[ { "path": ".clang-format", "content": "# http://clang.llvm.org/docs/ClangFormatStyleOptions.html\n# Defines the Google C++ style for automatic reformatting.\nBasedOnStyle: Google\nMaxEmptyLinesToKeep: 2\nDerivePointerAlignment: false\nPointerAlignment: Right\nAllowShortFunctionsOnASingleLine: Empty\nIncludeBlocks: Merge\nIncludeCategories:\n - Regex: '^$'\n Priority: 200\n - Regex: '^<[0-9A-Za-z_]+>$'\n Priority: 201\n - Regex: '^<[0-9A-Za-z_]+\\.[0-9A-Za-z]+>$'\n Priority: 202\n - Regex: '^<[0-9A-Za-z_]+/[0-9A-Za-z]+'\n Priority: 203\n - Regex: '^\\\"[0-9A-Za-z_]+/[0-9A-Za-z]+'\n Priority: 300\n - Regex: '^\\\"[0-9A-Za-z_]+\\.[0-9A-Za-z]+\\\"$'\n Priority: 301\n - Regex: '.*'\n Priority: 1000\n" }, { "path": ".git/HEAD", "content": "ref: refs/heads/main\n" }, { "path": ".git/config", "content": "[core]\n\trepositoryformatversion = 1\n\tfilemode = true\n\tbare = false\n\tlogallrefupdates = true\n[remote \"origin\"]\n\turl = https://github.com/alibaba/zvec\n\ttagOpt = --no-tags\n\tfetch = +refs/heads/main:refs/remotes/origin/main\n\tpromisor = true\n\tpartialclonefilter = blob:limit=1048576\n[branch \"main\"]\n\tremote = origin\n\tmerge = refs/heads/main\n" }, { "path": ".git/description", "content": "Unnamed repository; edit this file 'description' to name the repository.\n" }, { "path": ".git/hooks/applypatch-msg.sample", "content": "#!/bin/sh\n#\n# An example hook script to check the commit log message taken by\n# applypatch from an e-mail message.\n#\n# The hook should exit with non-zero status after issuing an\n# appropriate message if it wants to stop the commit. The hook is\n# allowed to edit the commit message file.\n#\n# To enable this hook, rename this file to \"applypatch-msg\".\n\n. git-sh-setup\ncommitmsg=\"$(git rev-parse --git-path hooks/commit-msg)\"\ntest -x \"$commitmsg\" && exec \"$commitmsg\" ${1+\"$@\"}\n:\n" }, { "path": ".git/hooks/commit-msg.sample", "content": "#!/bin/sh\n#\n# An example hook script to check the commit log message.\n# Called by \"git commit\" with one argument, the name of the file\n# that has the commit message. The hook should exit with non-zero\n# status after issuing an appropriate message if it wants to stop the\n# commit. The hook is allowed to edit the commit message file.\n#\n# To enable this hook, rename this file to \"commit-msg\".\n\n# Uncomment the below to add a Signed-off-by line to the message.\n# Doing this in a hook is a bad idea in general, but the prepare-commit-msg\n# hook is more suited to it.\n#\n# SOB=$(git var GIT_AUTHOR_IDENT | sed -n 's/^\\(.*>\\).*$/Signed-off-by: \\1/p')\n# grep -qs \"^$SOB\" \"$1\" || echo \"$SOB\" >> \"$1\"\n\n# This example catches duplicate Signed-off-by lines.\n\ntest \"\" = \"$(grep '^Signed-off-by: ' \"$1\" |\n\t sort | uniq -c | sed -e '/^[ \t]*1[ \t]/d')\" || {\n\techo >&2 Duplicate Signed-off-by lines.\n\texit 1\n}\n" }, { "path": ".git/hooks/fsmonitor-watchman.sample", "content": "#!/usr/bin/perl\n\nuse strict;\nuse warnings;\nuse IPC::Open2;\n\n# An example hook script to integrate Watchman\n# (https://facebook.github.io/watchman/) with git to speed up detecting\n# new and modified files.\n#\n# The hook is passed a version (currently 2) and last update token\n# formatted as a string and outputs to stdout a new update token and\n# all files that have been modified since the update token. Paths must\n# be relative to the root of the working tree and separated by a single NUL.\n#\n# To enable this hook, rename this file to \"query-watchman\" and set\n# 'git config core.fsmonitor .git/hooks/query-watchman'\n#\nmy ($version, $last_update_token) = @ARGV;\n\n# Uncomment for debugging\n# print STDERR \"$0 $version $last_update_token\\n\";\n\n# Check the hook interface version\nif ($version ne 2) {\n\tdie \"Unsupported query-fsmonitor hook version '$version'.\\n\" .\n\t \"Falling back to scanning...\\n\";\n}\n\nmy $git_work_tree = get_working_dir();\n\nmy $retry = 1;\n\nmy $json_pkg;\neval {\n\trequire JSON::XS;\n\t$json_pkg = \"JSON::XS\";\n\t1;\n} or do {\n\trequire JSON::PP;\n\t$json_pkg = \"JSON::PP\";\n};\n\nlaunch_watchman();\n\nsub launch_watchman {\n\tmy $o = watchman_query();\n\tif (is_work_tree_watched($o)) {\n\t\toutput_result($o->{clock}, @{$o->{files}});\n\t}\n}\n\nsub output_result {\n\tmy ($clockid, @files) = @_;\n\n\t# Uncomment for debugging watchman output\n\t# open (my $fh, \">\", \".git/watchman-output.out\");\n\t# binmode $fh, \":utf8\";\n\t# print $fh \"$clockid\\n@files\\n\";\n\t# close $fh;\n\n\tbinmode STDOUT, \":utf8\";\n\tprint $clockid;\n\tprint \"\\0\";\n\tlocal $, = \"\\0\";\n\tprint @files;\n}\n\nsub watchman_clock {\n\tmy $response = qx/watchman clock \"$git_work_tree\"/;\n\tdie \"Failed to get clock id on '$git_work_tree'.\\n\" .\n\t\t\"Falling back to scanning...\\n\" if $? != 0;\n\n\treturn $json_pkg->new->utf8->decode($response);\n}\n\nsub watchman_query {\n\tmy $pid = open2(\\*CHLD_OUT, \\*CHLD_IN, 'watchman -j --no-pretty')\n\tor die \"open2() failed: $!\\n\" .\n\t\"Falling back to scanning...\\n\";\n\n\t# In the query expression below we're asking for names of files that\n\t# changed since $last_update_token but not from the .git folder.\n\t#\n\t# To accomplish this, we're using the \"since\" generator to use the\n\t# recency index to select candidate nodes and \"fields\" to limit the\n\t# output to file names only. Then we're using the \"expression\" term to\n\t# further constrain the results.\n\tmy $last_update_line = \"\";\n\tif (substr($last_update_token, 0, 1) eq \"c\") {\n\t\t$last_update_token = \"\\\"$last_update_token\\\"\";\n\t\t$last_update_line = qq[\\n\"since\": $last_update_token,];\n\t}\n\tmy $query = <<\"\tEND\";\n\t\t[\"query\", \"$git_work_tree\", {$last_update_line\n\t\t\t\"fields\": [\"name\"],\n\t\t\t\"expression\": [\"not\", [\"dirname\", \".git\"]]\n\t\t}]\n\tEND\n\n\t# Uncomment for debugging the watchman query\n\t# open (my $fh, \">\", \".git/watchman-query.json\");\n\t# print $fh $query;\n\t# close $fh;\n\n\tprint CHLD_IN $query;\n\tclose CHLD_IN;\n\tmy $response = do {local $/; };\n\n\t# Uncomment for debugging the watch response\n\t# open ($fh, \">\", \".git/watchman-response.json\");\n\t# print $fh $response;\n\t# close $fh;\n\n\tdie \"Watchman: command returned no output.\\n\" .\n\t\"Falling back to scanning...\\n\" if $response eq \"\";\n\tdie \"Watchman: command returned invalid output: $response\\n\" .\n\t\"Falling back to scanning...\\n\" unless $response =~ /^\\{/;\n\n\treturn $json_pkg->new->utf8->decode($response);\n}\n\nsub is_work_tree_watched {\n\tmy ($output) = @_;\n\tmy $error = $output->{error};\n\tif ($retry > 0 and $error and $error =~ m/unable to resolve root .* directory (.*) is not watched/) {\n\t\t$retry--;\n\t\tmy $response = qx/watchman watch \"$git_work_tree\"/;\n\t\tdie \"Failed to make watchman watch '$git_work_tree'.\\n\" .\n\t\t \"Falling back to scanning...\\n\" if $? != 0;\n\t\t$output = $json_pkg->new->utf8->decode($response);\n\t\t$error = $output->{error};\n\t\tdie \"Watchman: $error.\\n\" .\n\t\t\"Falling back to scanning...\\n\" if $error;\n\n\t\t# Uncomment for debugging watchman output\n\t\t# open (my $fh, \">\", \".git/watchman-output.out\");\n\t\t# close $fh;\n\n\t\t# Watchman will always return all files on the first query so\n\t\t# return the fast \"everything is dirty\" flag to git and do the\n\t\t# Watchman query just to get it over with now so we won't pay\n\t\t# the cost in git to look up each individual file.\n\t\tmy $o = watchman_clock();\n\t\t$error = $output->{error};\n\n\t\tdie \"Watchman: $error.\\n\" .\n\t\t\"Falling back to scanning...\\n\" if $error;\n\n\t\toutput_result($o->{clock}, (\"/\"));\n\t\t$last_update_token = $o->{clock};\n\n\t\teval { launch_watchman() };\n\t\treturn 0;\n\t}\n\n\tdie \"Watchman: $error.\\n\" .\n\t\"Falling back to scanning...\\n\" if $error;\n\n\treturn 1;\n}\n\nsub get_working_dir {\n\tmy $working_dir;\n\tif ($^O =~ 'msys' || $^O =~ 'cygwin') {\n\t\t$working_dir = Win32::GetCwd();\n\t\t$working_dir =~ tr/\\\\/\\//;\n\t} else {\n\t\trequire Cwd;\n\t\t$working_dir = Cwd::cwd();\n\t}\n\n\treturn $working_dir;\n}\n" }, { "path": ".git/hooks/post-update.sample", "content": "#!/bin/sh\n#\n# An example hook script to prepare a packed repository for use over\n# dumb transports.\n#\n# To enable this hook, rename this file to \"post-update\".\n\nexec git update-server-info\n" }, { "path": ".git/hooks/pre-applypatch.sample", "content": "#!/bin/sh\n#\n# An example hook script to verify what is about to be committed\n# by applypatch from an e-mail message.\n#\n# The hook should exit with non-zero status after issuing an\n# appropriate message if it wants to stop the commit.\n#\n# To enable this hook, rename this file to \"pre-applypatch\".\n\n. git-sh-setup\nprecommit=\"$(git rev-parse --git-path hooks/pre-commit)\"\ntest -x \"$precommit\" && exec \"$precommit\" ${1+\"$@\"}\n:\n" }, { "path": ".git/hooks/pre-commit.sample", "content": "#!/bin/sh\n#\n# An example hook script to verify what is about to be committed.\n# Called by \"git commit\" with no arguments. The hook should\n# exit with non-zero status after issuing an appropriate message if\n# it wants to stop the commit.\n#\n# To enable this hook, rename this file to \"pre-commit\".\n\nif git rev-parse --verify HEAD >/dev/null 2>&1\nthen\n\tagainst=HEAD\nelse\n\t# Initial commit: diff against an empty tree object\n\tagainst=$(git hash-object -t tree /dev/null)\nfi\n\n# If you want to allow non-ASCII filenames set this variable to true.\nallownonascii=$(git config --type=bool hooks.allownonascii)\n\n# Redirect output to stderr.\nexec 1>&2\n\n# Cross platform projects tend to avoid non-ASCII filenames; prevent\n# them from being added to the repository. We exploit the fact that the\n# printable range starts at the space character and ends with tilde.\nif [ \"$allownonascii\" != \"true\" ] &&\n\t# Note that the use of brackets around a tr range is ok here, (it's\n\t# even required, for portability to Solaris 10's /usr/bin/tr), since\n\t# the square bracket bytes happen to fall in the designated range.\n\ttest $(git diff-index --cached --name-only --diff-filter=A -z $against |\n\t LC_ALL=C tr -d '[ -~]\\0' | wc -c) != 0\nthen\n\tcat <<\\EOF\nError: Attempt to add a non-ASCII file name.\n\nThis can cause problems if you want to work with people on other platforms.\n\nTo be portable it is advisable to rename the file.\n\nIf you know what you are doing you can disable this check using:\n\n git config hooks.allownonascii true\nEOF\n\texit 1\nfi\n\n# If there are whitespace errors, print the offending file names and fail.\nexec git diff-index --check --cached $against --\n" }, { "path": ".git/hooks/pre-merge-commit.sample", "content": "#!/bin/sh\n#\n# An example hook script to verify what is about to be committed.\n# Called by \"git merge\" with no arguments. The hook should\n# exit with non-zero status after issuing an appropriate message to\n# stderr if it wants to stop the merge commit.\n#\n# To enable this hook, rename this file to \"pre-merge-commit\".\n\n. git-sh-setup\ntest -x \"$GIT_DIR/hooks/pre-commit\" &&\n exec \"$GIT_DIR/hooks/pre-commit\"\n:\n" }, { "path": ".git/hooks/pre-push.sample", "content": "#!/bin/sh\n\n# An example hook script to verify what is about to be pushed. Called by \"git\n# push\" after it has checked the remote status, but before anything has been\n# pushed. If this script exits with a non-zero status nothing will be pushed.\n#\n# This hook is called with the following parameters:\n#\n# $1 -- Name of the remote to which the push is being done\n# $2 -- URL to which the push is being done\n#\n# If pushing without using a named remote those arguments will be equal.\n#\n# Information about the commits which are being pushed is supplied as lines to\n# the standard input in the form:\n#\n# \n#\n# This sample shows how to prevent push of commits where the log message starts\n# with \"WIP\" (work in progress).\n\nremote=\"$1\"\nurl=\"$2\"\n\nzero=$(git hash-object --stdin &2 \"Found WIP commit in $local_ref, not pushing\"\n\t\t\texit 1\n\t\tfi\n\tfi\ndone\n\nexit 0\n" }, { "path": ".git/hooks/pre-rebase.sample", "content": "#!/bin/sh\n#\n# Copyright (c) 2006, 2008 Junio C Hamano\n#\n# The \"pre-rebase\" hook is run just before \"git rebase\" starts doing\n# its job, and can prevent the command from running by exiting with\n# non-zero status.\n#\n# The hook is called with the following parameters:\n#\n# $1 -- the upstream the series was forked from.\n# $2 -- the branch being rebased (or empty when rebasing the current branch).\n#\n# This sample shows how to prevent topic branches that are already\n# merged to 'next' branch from getting rebased, because allowing it\n# would result in rebasing already published history.\n\npublish=next\nbasebranch=\"$1\"\nif test \"$#\" = 2\nthen\n\ttopic=\"refs/heads/$2\"\nelse\n\ttopic=`git symbolic-ref HEAD` ||\n\texit 0 ;# we do not interrupt rebasing detached HEAD\nfi\n\ncase \"$topic\" in\nrefs/heads/??/*)\n\t;;\n*)\n\texit 0 ;# we do not interrupt others.\n\t;;\nesac\n\n# Now we are dealing with a topic branch being rebased\n# on top of master. Is it OK to rebase it?\n\n# Does the topic really exist?\ngit show-ref -q \"$topic\" || {\n\techo >&2 \"No such branch $topic\"\n\texit 1\n}\n\n# Is topic fully merged to master?\nnot_in_master=`git rev-list --pretty=oneline ^master \"$topic\"`\nif test -z \"$not_in_master\"\nthen\n\techo >&2 \"$topic is fully merged to master; better remove it.\"\n\texit 1 ;# we could allow it, but there is no point.\nfi\n\n# Is topic ever merged to next? If so you should not be rebasing it.\nonly_next_1=`git rev-list ^master \"^$topic\" ${publish} | sort`\nonly_next_2=`git rev-list ^master ${publish} | sort`\nif test \"$only_next_1\" = \"$only_next_2\"\nthen\n\tnot_in_topic=`git rev-list \"^$topic\" master`\n\tif test -z \"$not_in_topic\"\n\tthen\n\t\techo >&2 \"$topic is already up to date with master\"\n\t\texit 1 ;# we could allow it, but there is no point.\n\telse\n\t\texit 0\n\tfi\nelse\n\tnot_in_next=`git rev-list --pretty=oneline ^${publish} \"$topic\"`\n\t/usr/bin/perl -e '\n\t\tmy $topic = $ARGV[0];\n\t\tmy $msg = \"* $topic has commits already merged to public branch:\\n\";\n\t\tmy (%not_in_next) = map {\n\t\t\t/^([0-9a-f]+) /;\n\t\t\t($1 => 1);\n\t\t} split(/\\n/, $ARGV[1]);\n\t\tfor my $elem (map {\n\t\t\t\t/^([0-9a-f]+) (.*)$/;\n\t\t\t\t[$1 => $2];\n\t\t\t} split(/\\n/, $ARGV[2])) {\n\t\t\tif (!exists $not_in_next{$elem->[0]}) {\n\t\t\t\tif ($msg) {\n\t\t\t\t\tprint STDERR $msg;\n\t\t\t\t\tundef $msg;\n\t\t\t\t}\n\t\t\t\tprint STDERR \" $elem->[1]\\n\";\n\t\t\t}\n\t\t}\n\t' \"$topic\" \"$not_in_next\" \"$not_in_master\"\n\texit 1\nfi\n\n<<\\DOC_END\n\nThis sample hook safeguards topic branches that have been\npublished from being rewound.\n\nThe workflow assumed here is:\n\n * Once a topic branch forks from \"master\", \"master\" is never\n merged into it again (either directly or indirectly).\n\n * Once a topic branch is fully cooked and merged into \"master\",\n it is deleted. If you need to build on top of it to correct\n earlier mistakes, a new topic branch is created by forking at\n the tip of the \"master\". This is not strictly necessary, but\n it makes it easier to keep your history simple.\n\n * Whenever you need to test or publish your changes to topic\n branches, merge them into \"next\" branch.\n\nThe script, being an example, hardcodes the publish branch name\nto be \"next\", but it is trivial to make it configurable via\n$GIT_DIR/config mechanism.\n\nWith this workflow, you would want to know:\n\n(1) ... if a topic branch has ever been merged to \"next\". Young\n topic branches can have stupid mistakes you would rather\n clean up before publishing, and things that have not been\n merged into other branches can be easily rebased without\n affecting other people. But once it is published, you would\n not want to rewind it.\n\n(2) ... if a topic branch has been fully merged to \"master\".\n Then you can delete it. More importantly, you should not\n build on top of it -- other people may already want to\n change things related to the topic as patches against your\n \"master\", so if you need further changes, it is better to\n fork the topic (perhaps with the same name) afresh from the\n tip of \"master\".\n\nLet's look at this example:\n\n\t\t o---o---o---o---o---o---o---o---o---o \"next\"\n\t\t / / / /\n\t\t / a---a---b A / /\n\t\t/ / / /\n\t / / c---c---c---c B /\n\t / / / \\ /\n\t / / / b---b C \\ /\n\t / / / / \\ /\n ---o---o---o---o---o---o---o---o---o---o---o \"master\"\n\n\nA, B and C are topic branches.\n\n * A has one fix since it was merged up to \"next\".\n\n * B has finished. It has been fully merged up to \"master\" and \"next\",\n and is ready to be deleted.\n\n * C has not merged to \"next\" at all.\n\nWe would want to allow C to be rebased, refuse A, and encourage\nB to be deleted.\n\nTo compute (1):\n\n\tgit rev-list ^master ^topic next\n\tgit rev-list ^master next\n\n\tif these match, topic has not merged in next at all.\n\nTo compute (2):\n\n\tgit rev-list master..topic\n\n\tif this is empty, it is fully merged to \"master\".\n\nDOC_END\n" }, { "path": ".git/hooks/pre-receive.sample", "content": "#!/bin/sh\n#\n# An example hook script to make use of push options.\n# The example simply echoes all push options that start with 'echoback='\n# and rejects all pushes when the \"reject\" push option is used.\n#\n# To enable this hook, rename this file to \"pre-receive\".\n\nif test -n \"$GIT_PUSH_OPTION_COUNT\"\nthen\n\ti=0\n\twhile test \"$i\" -lt \"$GIT_PUSH_OPTION_COUNT\"\n\tdo\n\t\teval \"value=\\$GIT_PUSH_OPTION_$i\"\n\t\tcase \"$value\" in\n\t\techoback=*)\n\t\t\techo \"echo from the pre-receive-hook: ${value#*=}\" >&2\n\t\t\t;;\n\t\treject)\n\t\t\texit 1\n\t\tesac\n\t\ti=$((i + 1))\n\tdone\nfi\n" }, { "path": ".git/hooks/prepare-commit-msg.sample", "content": "#!/bin/sh\n#\n# An example hook script to prepare the commit log message.\n# Called by \"git commit\" with the name of the file that has the\n# commit message, followed by the description of the commit\n# message's source. The hook's purpose is to edit the commit\n# message file. If the hook fails with a non-zero status,\n# the commit is aborted.\n#\n# To enable this hook, rename this file to \"prepare-commit-msg\".\n\n# This hook includes three examples. The first one removes the\n# \"# Please enter the commit message...\" help message.\n#\n# The second includes the output of \"git diff --name-status -r\"\n# into the message, just before the \"git status\" output. It is\n# commented because it doesn't cope with --amend or with squashed\n# commits.\n#\n# The third example adds a Signed-off-by line to the message, that can\n# still be edited. This is rarely a good idea.\n\nCOMMIT_MSG_FILE=$1\nCOMMIT_SOURCE=$2\nSHA1=$3\n\n/usr/bin/perl -i.bak -ne 'print unless(m/^. Please enter the commit message/..m/^#$/)' \"$COMMIT_MSG_FILE\"\n\n# case \"$COMMIT_SOURCE,$SHA1\" in\n# ,|template,)\n# /usr/bin/perl -i.bak -pe '\n# print \"\\n\" . `git diff --cached --name-status -r`\n# \t if /^#/ && $first++ == 0' \"$COMMIT_MSG_FILE\" ;;\n# *) ;;\n# esac\n\n# SOB=$(git var GIT_COMMITTER_IDENT | sed -n 's/^\\(.*>\\).*$/Signed-off-by: \\1/p')\n# git interpret-trailers --in-place --trailer \"$SOB\" \"$COMMIT_MSG_FILE\"\n# if test -z \"$COMMIT_SOURCE\"\n# then\n# /usr/bin/perl -i.bak -pe 'print \"\\n\" if !$first_line++' \"$COMMIT_MSG_FILE\"\n# fi\n" }, { "path": ".git/hooks/push-to-checkout.sample", "content": "#!/bin/sh\n\n# An example hook script to update a checked-out tree on a git push.\n#\n# This hook is invoked by git-receive-pack(1) when it reacts to git\n# push and updates reference(s) in its repository, and when the push\n# tries to update the branch that is currently checked out and the\n# receive.denyCurrentBranch configuration variable is set to\n# updateInstead.\n#\n# By default, such a push is refused if the working tree and the index\n# of the remote repository has any difference from the currently\n# checked out commit; when both the working tree and the index match\n# the current commit, they are updated to match the newly pushed tip\n# of the branch. This hook is to be used to override the default\n# behaviour; however the code below reimplements the default behaviour\n# as a starting point for convenient modification.\n#\n# The hook receives the commit with which the tip of the current\n# branch is going to be updated:\ncommit=$1\n\n# It can exit with a non-zero status to refuse the push (when it does\n# so, it must not modify the index or the working tree).\ndie () {\n\techo >&2 \"$*\"\n\texit 1\n}\n\n# Or it can make any necessary changes to the working tree and to the\n# index to bring them to the desired state when the tip of the current\n# branch is updated to the new commit, and exit with a zero status.\n#\n# For example, the hook can simply run git read-tree -u -m HEAD \"$1\"\n# in order to emulate git fetch that is run in the reverse direction\n# with git push, as the two-tree form of git read-tree -u -m is\n# essentially the same as git switch or git checkout that switches\n# branches while keeping the local changes in the working tree that do\n# not interfere with the difference between the branches.\n\n# The below is a more-or-less exact translation to shell of the C code\n# for the default behaviour for git's push-to-checkout hook defined in\n# the push_to_deploy() function in builtin/receive-pack.c.\n#\n# Note that the hook will be executed from the repository directory,\n# not from the working tree, so if you want to perform operations on\n# the working tree, you will have to adapt your code accordingly, e.g.\n# by adding \"cd ..\" or using relative paths.\n\nif ! git update-index -q --ignore-submodules --refresh\nthen\n\tdie \"Up-to-date check failed\"\nfi\n\nif ! git diff-files --quiet --ignore-submodules --\nthen\n\tdie \"Working directory has unstaged changes\"\nfi\n\n# This is a rough translation of:\n#\n# head_has_history() ? \"HEAD\" : EMPTY_TREE_SHA1_HEX\nif git cat-file -e HEAD 2>/dev/null\nthen\n\thead=HEAD\nelse\n\thead=$(git hash-object -t tree --stdin &2\n\texit 1\n}\n\nunset GIT_DIR GIT_WORK_TREE\ncd \"$worktree\" &&\n\nif grep -q \"^diff --git \" \"$1\"\nthen\n\tvalidate_patch \"$1\"\nelse\n\tvalidate_cover_letter \"$1\"\nfi &&\n\nif test \"$GIT_SENDEMAIL_FILE_COUNTER\" = \"$GIT_SENDEMAIL_FILE_TOTAL\"\nthen\n\tgit config --unset-all sendemail.validateWorktree &&\n\ttrap 'git worktree remove -ff \"$worktree\"' EXIT &&\n\tvalidate_series\nfi\n" }, { "path": ".git/hooks/update.sample", "content": "#!/bin/sh\n#\n# An example hook script to block unannotated tags from entering.\n# Called by \"git receive-pack\" with arguments: refname sha1-old sha1-new\n#\n# To enable this hook, rename this file to \"update\".\n#\n# Config\n# ------\n# hooks.allowunannotated\n# This boolean sets whether unannotated tags will be allowed into the\n# repository. By default they won't be.\n# hooks.allowdeletetag\n# This boolean sets whether deleting tags will be allowed in the\n# repository. By default they won't be.\n# hooks.allowmodifytag\n# This boolean sets whether a tag may be modified after creation. By default\n# it won't be.\n# hooks.allowdeletebranch\n# This boolean sets whether deleting branches will be allowed in the\n# repository. By default they won't be.\n# hooks.denycreatebranch\n# This boolean sets whether remotely creating branches will be denied\n# in the repository. By default this is allowed.\n#\n\n# --- Command line\nrefname=\"$1\"\noldrev=\"$2\"\nnewrev=\"$3\"\n\n# --- Safety check\nif [ -z \"$GIT_DIR\" ]; then\n\techo \"Don't run this script from the command line.\" >&2\n\techo \" (if you want, you could supply GIT_DIR then run\" >&2\n\techo \" $0 )\" >&2\n\texit 1\nfi\n\nif [ -z \"$refname\" -o -z \"$oldrev\" -o -z \"$newrev\" ]; then\n\techo \"usage: $0 \" >&2\n\texit 1\nfi\n\n# --- Config\nallowunannotated=$(git config --type=bool hooks.allowunannotated)\nallowdeletebranch=$(git config --type=bool hooks.allowdeletebranch)\ndenycreatebranch=$(git config --type=bool hooks.denycreatebranch)\nallowdeletetag=$(git config --type=bool hooks.allowdeletetag)\nallowmodifytag=$(git config --type=bool hooks.allowmodifytag)\n\n# check for no description\nprojectdesc=$(sed -e '1q' \"$GIT_DIR/description\")\ncase \"$projectdesc\" in\n\"Unnamed repository\"* | \"\")\n\techo \"*** Project description file hasn't been set\" >&2\n\texit 1\n\t;;\nesac\n\n# --- Check types\n# if $newrev is 0000...0000, it's a commit to delete a ref.\nzero=$(git hash-object --stdin &2\n\t\t\techo \"*** Use 'git tag [ -a | -s ]' for tags you want to propagate.\" >&2\n\t\t\texit 1\n\t\tfi\n\t\t;;\n\trefs/tags/*,delete)\n\t\t# delete tag\n\t\tif [ \"$allowdeletetag\" != \"true\" ]; then\n\t\t\techo \"*** Deleting a tag is not allowed in this repository\" >&2\n\t\t\texit 1\n\t\tfi\n\t\t;;\n\trefs/tags/*,tag)\n\t\t# annotated tag\n\t\tif [ \"$allowmodifytag\" != \"true\" ] && git rev-parse $refname > /dev/null 2>&1\n\t\tthen\n\t\t\techo \"*** Tag '$refname' already exists.\" >&2\n\t\t\techo \"*** Modifying a tag is not allowed in this repository.\" >&2\n\t\t\texit 1\n\t\tfi\n\t\t;;\n\trefs/heads/*,commit)\n\t\t# branch\n\t\tif [ \"$oldrev\" = \"$zero\" -a \"$denycreatebranch\" = \"true\" ]; then\n\t\t\techo \"*** Creating a branch is not allowed in this repository\" >&2\n\t\t\texit 1\n\t\tfi\n\t\t;;\n\trefs/heads/*,delete)\n\t\t# delete branch\n\t\tif [ \"$allowdeletebranch\" != \"true\" ]; then\n\t\t\techo \"*** Deleting a branch is not allowed in this repository\" >&2\n\t\t\texit 1\n\t\tfi\n\t\t;;\n\trefs/remotes/*,commit)\n\t\t# tracking branch\n\t\t;;\n\trefs/remotes/*,delete)\n\t\t# delete tracking branch\n\t\tif [ \"$allowdeletebranch\" != \"true\" ]; then\n\t\t\techo \"*** Deleting a tracking branch is not allowed in this repository\" >&2\n\t\t\texit 1\n\t\tfi\n\t\t;;\n\t*)\n\t\t# Anything else (is there anything else?)\n\t\techo \"*** Update hook: unknown type of update to ref $refname of type $newrev_type\" >&2\n\t\texit 1\n\t\t;;\nesac\n\n# --- Finished\nexit 0\n" }, { "path": ".git/info/exclude", "content": "# git ls-files --others --exclude-from=.git/info/exclude\n# Lines that start with '#' are comments.\n# For a project mostly in C, the following would be a good set of\n# exclude patterns (uncomment them if you want to use them):\n# *.[oa]\n# *~\n" }, { "path": ".git/logs/HEAD", "content": "0000000000000000000000000000000000000000 b49833bf56a0e102b8ac1ff95ed7766545f5bd1e appuser 1774064477 +0000\tclone: from https://github.com/alibaba/zvec\n" }, { "path": ".git/logs/refs/heads/main", "content": "0000000000000000000000000000000000000000 b49833bf56a0e102b8ac1ff95ed7766545f5bd1e appuser 1774064477 +0000\tclone: from https://github.com/alibaba/zvec\n" }, { "path": ".git/logs/refs/remotes/origin/HEAD", "content": "0000000000000000000000000000000000000000 b49833bf56a0e102b8ac1ff95ed7766545f5bd1e appuser 1774064477 +0000\tclone: from https://github.com/alibaba/zvec\n" }, { "path": ".git/objects/pack/pack-2b5e15ebe928a592991dc24c7ae7e8dc9e3500dc.promisor", "content": "b49833bf56a0e102b8ac1ff95ed7766545f5bd1e refs/heads/main\n" }, { "path": ".git/packed-refs", "content": "# pack-refs with: peeled fully-peeled sorted \nb49833bf56a0e102b8ac1ff95ed7766545f5bd1e refs/remotes/origin/main\n" }, { "path": ".git/refs/heads/main", "content": "b49833bf56a0e102b8ac1ff95ed7766545f5bd1e\n" }, { "path": ".git/refs/remotes/origin/HEAD", "content": "ref: refs/remotes/origin/main\n" }, { "path": ".git/shallow", "content": "b49833bf56a0e102b8ac1ff95ed7766545f5bd1e\n" }, { "path": ".github/ISSUE_TEMPLATE/benchmark.yml", "content": "name: Benchmarking\ndescription: Add, update, or fix benchmark cases for zvec\ntitle: \"[Benchmark]: \"\nlabels: [\"benchmark\"]\nbody:\n - type: markdown\n attributes:\n value: |\n Use this for benchmark-related work: new test cases, CI integration, or performance regression tracking.\n \n - type: input\n id: benchmark_type\n attributes:\n label: Benchmark Type\n description: e.g., filtered search, batch insert, recall@k, ARM64 vs x86\n validations:\n required: true\n \n - type: textarea\n id: goal\n attributes:\n label: Goal\n description: What performance aspect are you measuring or improving?\n validations:\n required: true\n \n - type: textarea\n id: methodology\n attributes:\n label: Methodology\n description: Dataset, query size, hardware, metrics (latency, throughput, memory)\n validations:\n required: true\n \n - type: textarea\n id: baseline\n attributes:\n label: Baseline (if applicable)\n description: Current performance numbers or competing systems for comparison.\n validations:\n required: false\n \n - type: textarea\n id: ci_integration\n attributes:\n label: CI Integration Plan\n description: Should this run in CI? How often?\n validations:\n required: false" }, { "path": ".github/ISSUE_TEMPLATE/bug_report.yml", "content": "name: Bug Report\ndescription: Report a bug or unexpected behavior (e.g., crash, incorrect vector query, memory leak)\ntitle: \"[Bug]: \"\nlabels: [\"bug\", \"triage\"]\nbody:\n - type: markdown\n attributes:\n value: |\n Thank you for reporting! Please provide detailed info so we can reproduce and fix it quickly.\n\n - type: textarea\n id: description\n attributes:\n label: Description\n description: What happened? What did you expect?\n placeholder: |\n e.g. \"Query with vector field crashes when using Zvec Python API\"\n validations:\n required: true\n\n - type: textarea\n id: steps_to_reproduce\n attributes:\n label: Steps to Reproduce\n description: Exact steps to trigger the issue (code snippets welcome)\n placeholder: |\n 1. Build Zvec with CMake (Debug/Release)\n 2. Run Python script: `python test.py`\n 3. Call `collection.query(VectorQuery())`\n 4. Process segfaults / hangs / returns wrong results\n render: python\n validations:\n required: true\n\n - type: textarea\n id: logs_or_trace\n attributes:\n label: Logs / Stack Trace\n description: Paste relevant logs, LLDB/GDB backtrace, or CI failures\n placeholder: |\n Thread 1 \"python\" received signal SIGSEGV, Segmentation fault.\n 0x0000000104a2c3f0 in std::__1::shared_ptr<...>::...\n render: shell\n validations:\n required: false\n\n - type: input\n id: os\n attributes:\n label: Operating System\n placeholder: macOS 14 (M1), Ubuntu 22.04, Windows 11 (WSL2)\n validations:\n required: true\n\n - type: input\n id: build_env\n attributes:\n label: Build & Runtime Environment\n description: Compiler, CMake, Python, key dependencies\n placeholder: |\n clang 15.0.0, CMake 4.1.2, Python 3.11.9, magic_enum v0.9.7 (via git submodule)\n validations:\n required: true\n\n - type: checkboxes\n id: additional_context\n attributes:\n label: Additional Context\n options:\n - label: I've checked `git status` — no uncommitted submodule changes\n - label: I built with `CMAKE_BUILD_TYPE=Debug`\n - label: This occurs with or without `COVERAGE=ON`\n - label: The issue involves Python ↔ C++ integration (pybind11)" }, { "path": ".github/ISSUE_TEMPLATE/config.yml", "content": "blank_issues_enabled: false\ncontact_links:\n - name: Documentation\n url: https://zvec.org/en/\n about: Check the quickstart, build guide, and API docs first.\n\n - name: Python API Examples\n url: https://zvec.org/en/docs/quickstart/\n about: See working usage examples.\n" }, { "path": ".github/ISSUE_TEMPLATE/enhancement.yml", "content": "name: Enhancement\ndescription: Improve an existing feature or component\ntitle: \"[Enhance]: \"\nlabels: [\"enhancement\"]\nbody:\n - type: markdown\n attributes:\n value: |\n This template is for improving existing functionality (e.g., performance, usability, robustness).\n \n - type: input\n id: component\n attributes:\n label: Affected Component\n description: e.g., HNSW index, buffer manager, Python API\n validations:\n required: true\n \n - type: textarea\n id: current\n attributes:\n label: Current Behavior\n description: What is the current state and its limitations?\n validations:\n required: true\n \n - type: textarea\n id: desired\n attributes:\n label: Desired Improvement\n description: What should be improved and how?\n validations:\n required: true\n \n - type: textarea\n id: impact\n attributes:\n label: Impact\n description: How will this benefit users? (e.g., faster queries, lower memory, easier integration)\n validations:\n required: true" }, { "path": ".github/ISSUE_TEMPLATE/feature_request.yml", "content": "name: Feature Request\ndescription: Suggest a new feature or improvement (e.g., better memory control, new query option)\ntitle: \"[Feature]: \"\nlabels: [\"feature\"]\nbody:\n - type: markdown\n attributes:\n value: |\n Thanks for your idea! Help us understand the motivation and scope.\n\n - type: textarea\n id: problem_or_motivation\n attributes:\n label: Problem / Motivation\n description: What problem does this solve? Why is it needed?\n placeholder: |\n e.g. \"Current vector queries don't allow filtering by metadata + distance threshold at once\"\n validations:\n required: true\n\n - type: textarea\n id: proposed_solution\n attributes:\n label: Proposed Solution\n description: How should it work? API sketch or pseudocode welcome.\n placeholder: |\n Add `filter=` and `max_distance=` args to `Zvec.query()`:\n ```python\n results = db.query(vector, filter=\"category == 'A'\", max_distance=0.5)\n ```\n render: python\n validations:\n required: false\n\n - type: textarea\n id: alternatives\n attributes:\n label: Alternatives Considered\n description: Are there workarounds? Why not use them?\n validations:\n required: false\n\n - type: dropdown\n id: impact_area\n attributes:\n label: Affected Area\n multiple: true\n options:\n - label: C++ Core (storage, indexing)\n - label: Python API / Bindings\n - label: Build System (CMake, Homebrew pkg)\n - label: Testing / CI / Coverage\n - label: Documentation\n validations:\n required: false" }, { "path": ".github/ISSUE_TEMPLATE/integration.yml", "content": "name: Ecosystem Integration\ndescription: Integrate zvec with external frameworks (e.g., LangChain, LlamaIndex)\ntitle: \"[Integration]: \"\nlabels: [\"integration\"]\nbody:\n - type: input\n id: framework\n attributes:\n label: Target Framework\n description: e.g., LangChain, LlamaIndex, Haystack\n validations:\n required: true\n \n - type: textarea\n id: motivation\n attributes:\n label: Motivation\n description: Why integrate with this framework? Who benefits?\n validations:\n required: true\n \n - type: textarea\n id: interface\n attributes:\n label: Required Interface\n description: What adapter or interface must be implemented? (e.g., VectorStore base class)\n validations:\n required: true\n \n - type: textarea\n id: reference\n attributes:\n label: Reference Implementations\n description: Links to similar integrations in other vector DBs.\n validations:\n required: false" }, { "path": ".github/ISSUE_TEMPLATE/profiling.yml", "content": "name: Profiling / Investigation\ndescription: Profile performance, compatibility, or behavior in a specific scenario\ntitle: \"[Profile]: \"\nlabels: [\"profile\"]\nbody:\n - type: markdown\n attributes:\n value: |\n Use this for tasks like performance profiling, architecture compatibility checks, or feasibility studies.\n \n - type: input\n id: scenario\n attributes:\n label: Target Scenario\n description: e.g., ARM64 deployment, high-concurrency load, large dataset ingestion\n validations:\n required: true\n \n - type: textarea\n id: objective\n attributes:\n label: Objective\n description: What do you want to learn or validate?\n validations:\n required: true\n \n - type: textarea\n id: methodology\n attributes:\n label: Proposed Methodology\n description: How will you conduct the investigation? (tools, metrics, test data)\n validations:\n required: true\n \n - type: textarea\n id: expected_outcome\n attributes:\n label: Expected Outcome\n description: What deliverables are expected? (e.g., report, optimization PR, benchmark results)\n validations:\n required: true" }, { "path": ".github/codecov.yml", "content": "codecov:\n require_ci_to_pass: true\n\ncoverage:\n precision: 2\n round: down\n range: \"60...75\"\n\n status:\n project:\n default: false\n patch:\n default: false\n\n\nparsers:\n gcov:\n branch_detection:\n conditional: true\n loop: true\n method: false\n macro: false\n\ncomment:\n require_changes: false\n layout: \"reach,diff,flags,tree\"\n behavior: default\n\n\nignore:\n - \"thirdparty/\"\n - \"tests/\"" }, { "path": ".github/dependabot.yml", "content": "version: 2\nupdates:\n # GitHub Actions dependencies\n - package-ecosystem: \"github-actions\"\n directory: \"/\"\n schedule:\n interval: \"weekly\"\n day: \"monday\"\n time: \"02:00\"\n timezone: \"Asia/Shanghai\"\n labels:\n - \"dependencies\"\n - \"github-actions\"\n commit-message:\n prefix: \"ci\"\n include: \"scope\"\n open-pull-requests-limit: 5\n" }, { "path": ".github/workflows/01-ci-pipeline.yml", "content": "name: Main\n\non:\n push:\n branches: [ \"main\" ]\n paths-ignore:\n - '**.md'\n merge_group:\n pull_request:\n branches: [ \"main\" ]\n paths-ignore:\n - '**.md'\n workflow_dispatch:\n\nconcurrency:\n group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref || '' }}-${{ github.base_ref || '' }}-${{ github.ref != 'refs/heads/main' || github.sha }}\n cancel-in-progress: true\n\npermissions:\n contents: read\n\njobs:\n # Code quality checks (fast, run first)\n lint:\n uses: ./.github/workflows/02-lint-check.yml\n\n # Main build and test matrix\n build-and-test-macos-arm64:\n name: Build & Test (macos-arm64)\n needs: lint\n uses: ./.github/workflows/03-macos-linux-build.yml\n with:\n platform: macos-arm64\n os: macos-15\n\n build-and-test-linux-arm64:\n name: Build & Test (linux-arm64)\n needs: lint\n uses: ./.github/workflows/03-macos-linux-build.yml\n with:\n platform: linux-arm64\n os: ubuntu-24.04-arm\n\n build-and-test-linux-x64:\n name: Build & Test (linux-x64)\n needs: lint\n uses: ./.github/workflows/03-macos-linux-build.yml\n with:\n platform: linux-x64\n os: ubuntu-24.04\n\n build-android:\n name: Build & Test (android)\n needs: lint\n uses: ./.github/workflows/04-android-build.yml\n" }, { "path": ".github/workflows/02-lint-check.yml", "content": "name: Lint\n\non:\n workflow_call:\n\njobs:\n lint:\n name: Code Quality Checks\n runs-on: ubuntu-24.04\n steps:\n - name: Checkout code\n uses: actions/checkout@v6\n\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: '3.10'\n cache: 'pip'\n cache-dependency-path: 'pyproject.toml'\n\n - name: Install linting tools\n run: |\n python -m pip install --upgrade pip \\\n ruff==v0.14.4 \\\n clang-format==18.1.8\n shell: bash\n\n - name: Run Ruff Linter\n run: python -m ruff check .\n shell: bash\n\n - name: Run Ruff Formatter Check\n run: python -m ruff format --check .\n shell: bash\n\n - name: Run clang-format Check\n run: |\n CPP_FILES=$(find . -type f \\( -name \"*.cpp\" -o -name \"*.h\" -o -name \"*.hpp\" -o -name \"*.cc\" -o -name \"*.cxx\" \\) \\\n ! -path \"./build/*\" \\\n ! -path \"./tests/*\" \\\n ! -path \"./scripts/*\" \\\n ! -path \"./python/*\" \\\n ! -path \"./thirdparty/*\" \\\n ! -path \"./.git/*\")\n\n if [ -z \"$CPP_FILES\" ]; then\n echo \"No C++ files found to check.\"\n exit 0\n fi\n\n clang-format --dry-run --Werror $CPP_FILES\n shell: bash\n" }, { "path": ".github/workflows/03-macos-linux-build.yml", "content": "name: MacOS & Linux Build\n\non:\n workflow_call:\n inputs:\n platform:\n description: 'Platform identifier'\n required: true\n type: string\n os:\n description: 'GitHub Actions runner OS'\n required: true\n type: string\n\npermissions:\n contents: read\n\njobs:\n # Build and test matrix (parallel execution)\n build-and-test:\n name: Build & Test (${{ inputs.platform }})\n runs-on: ${{ inputs.os }}\n \n strategy:\n fail-fast: false\n matrix:\n include:\n - os: ${{ inputs.os }}\n platform: ${{ inputs.platform }}\n arch_flag: \"\" # Use appropriate architecture\n\n steps:\n - name: Checkout code\n uses: actions/checkout@v6\n with:\n submodules: recursive\n\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: '3.10'\n cache: 'pip'\n cache-dependency-path: 'pyproject.toml'\n\n - name: Set up environment variables\n run: |\n # Set number of processors for parallel builds\n if [[ \"${{ matrix.platform }}\" == \"macos-arm64\" ]]; then\n NPROC=$(sysctl -n hw.ncpu 2>/dev/null || echo 2)\n else\n NPROC=$(nproc 2>/dev/null || echo 2)\n fi\n echo \"NPROC=$NPROC\" >> $GITHUB_ENV\n echo \"Using $NPROC parallel jobs for builds\"\n \n # Add Python user base bin to PATH for pip-installed CLI tools\n echo \"$(python -c 'import site; print(site.USER_BASE)')/bin\" >> $GITHUB_PATH\n shell: bash\n\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip \\\n pybind11==3.0 \\\n cmake==3.30.0 \\\n ninja==1.11.1 \\\n pytest \\\n scikit-build-core \\\n setuptools_scm\n shell: bash\n\n - name: Build from source\n run: |\n cd \"$GITHUB_WORKSPACE\"\n \n CMAKE_GENERATOR=\"Unix Makefiles\" \\\n CMAKE_BUILD_PARALLEL_LEVEL=\"$NPROC\" \\\n python -m pip install -v . \\\n --no-build-isolation \\\n --config-settings='cmake.define.BUILD_TOOLS=\"ON\"' \\\n ${{ matrix.arch_flag }}\n shell: bash\n\n - name: Run C++ Tests\n run: |\n cd \"$GITHUB_WORKSPACE/build\"\n make unittest -j$NPROC\n shell: bash\n\n - name: Run Python Tests\n run: |\n cd \"$GITHUB_WORKSPACE\"\n python -m pytest python/tests/\n shell: bash\n\n - name: Run C++ Examples\n run: |\n cd \"$GITHUB_WORKSPACE/examples/c++\"\n mkdir build && cd build\n cmake .. -DCMAKE_BUILD_TYPE=Release\n make -j $NPROC\n ./db-example\n ./core-example\n ./ailego-example\n shell: bash\n" }, { "path": ".github/workflows/04-android-build.yml", "content": "name: Android Cross Build\n\non:\n workflow_call:\n\npermissions:\n contents: read\n\njobs:\n build-android:\n # sdkmanager and other Android tools are x86‑only; ARM runners fail with exit code 1\n # switch back to an x86 image so the setup-android action can install the SDK\n runs-on: ubuntu-24.04\n strategy:\n fail-fast: false\n matrix:\n abi: [x86_64]\n api: [21]\n\n steps:\n - name: Checkout\n uses: actions/checkout@v6\n\n - name: Cache dependencies\n uses: actions/cache@v5\n with:\n path: |\n ~/.ccache\n key: ${{ runner.os }}-dependencies-cache-${{ hashFiles('**/CMakeLists.txt', 'thirdparty/**') }}-stl-fix\n\n - name: Install dependencies\n run: |\n sudo apt-get update\n sudo apt-get install -y --no-install-recommends \\\n cmake ninja-build git ca-certificates python3 \\\n build-essential make ccache\n\n - name: Setup Java 17\n uses: actions/setup-java@v5\n with:\n distribution: temurin\n java-version: '17'\n\n - name: Setup Android SDK\n uses: android-actions/setup-android@v3\n\n - name: Install NDK (side by side)\n shell: bash\n run: |\n sdkmanager \"ndk;26.1.10909125\"\n\n - name: Cache host protoc build\n uses: actions/cache@v5\n with:\n path: build-host\n key: ${{ runner.os }}-host-protoc-${{ hashFiles('src/**', 'CMakeLists.txt') }}-stl-fix\n restore-keys: |\n ${{ runner.os }}-host-protoc-\n\n - name: Use host env to compile protoc\n shell: bash\n run: |\n git submodule update --init\n if [ ! -d \"build-host\" ]; then\n export CCACHE_BASEDIR=\"$GITHUB_WORKSPACE\"\n export CCACHE_NOHASHDIR=1\n export CCACHE_SLOPPINESS=clang_index_store,file_stat_matches,include_file_mtime,locale,time_macros\n\n cmake -S . -B build-host -G Ninja \\\n -DCMAKE_C_COMPILER_LAUNCHER=ccache \\\n -DCMAKE_CXX_COMPILER_LAUNCHER=ccache\n cmake --build build-host --target protoc --parallel\n else\n echo \"Using cached host protoc build\"\n fi\n\n - name: Cache Android build\n uses: actions/cache@v5\n with:\n path: build-android-${{ matrix.abi }}\n key: ${{ runner.os }}-android-build-${{ matrix.abi }}-${{ hashFiles('src/**', 'CMakeLists.txt', 'cmake/**', 'thirdparty/**') }}-stl-fix-3\n\n - name: Configure and Build\n shell: bash\n run: |\n git submodule foreach --recursive 'git stash --include-untracked'\n\n export ANDROID_SDK_ROOT=\"$ANDROID_HOME\"\n export ANDROID_NDK_HOME=\"$ANDROID_SDK_ROOT/ndk/26.1.10909125\"\n\n export CCACHE_BASEDIR=\"$GITHUB_WORKSPACE\"\n export CCACHE_NOHASHDIR=1\n export CCACHE_SLOPPINESS=clang_index_store,file_stat_matches,include_file_mtime,locale,time_macros\n\n if [ ! -d \"build-android-${{ matrix.abi }}\" ]; then\n cmake -S . -B build-android-${{ matrix.abi }} -G Ninja \\\n -DCMAKE_BUILD_TYPE=Release \\\n -DCMAKE_TOOLCHAIN_FILE=\"$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake\" \\\n -DANDROID_ABI=${{ matrix.abi }} \\\n -DANDROID_PLATFORM=android-${{ matrix.api }} \\\n -DANDROID_STL=c++_static \\\n -DBUILD_PYTHON_BINDINGS=OFF \\\n -DENABLE_NATIVE=OFF \\\n -DAUTO_DETECT_ARCH=OFF \\\n -DBUILD_TOOLS=OFF \\\n -DGLOBAL_CC_PROTOBUF_PROTOC=\"$GITHUB_WORKSPACE/build-host/bin/protoc\" \\\n -DCMAKE_C_COMPILER_LAUNCHER=ccache \\\n -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \\\n -DCMAKE_VERBOSE_MAKEFILE=ON\n cmake --build build-android-${{ matrix.abi }} --parallel --verbose\n else\n echo \"Using cached Android build directory\"\n fi\n\n - name: Cache examples build\n uses: actions/cache@v5\n with:\n path: examples/c++/build-android-examples-${{ matrix.abi }}\n key: ${{ runner.os }}-examples-build-${{ matrix.abi }}-${{ hashFiles('examples/c++/**', 'CMakeLists.txt', 'src/**') }}-stl-fix-3\n\n - name: Build examples\n shell: bash\n run: |\n export ANDROID_SDK_ROOT=\"$ANDROID_HOME\"\n export ANDROID_NDK_HOME=\"$ANDROID_SDK_ROOT/ndk/26.1.10909125\"\n\n if [ ! -d \"examples/c++/build-android-examples-${{ matrix.abi }}\" ]; then\n cmake -S examples/c++ -B examples/c++/build-android-examples-${{ matrix.abi }} -G Ninja \\\n -DCMAKE_TOOLCHAIN_FILE=\"$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake\" \\\n -DANDROID_ABI=${{ matrix.abi }} \\\n -DANDROID_PLATFORM=android-${{ matrix.api }} \\\n -DANDROID_STL=c++_static \\\n -DCMAKE_BUILD_TYPE=Release \\\n -DCMAKE_INTERPROCEDURAL_OPTIMIZATION=ON \\\n -DHOST_BUILD_DIR=\"build-android-${{ matrix.abi }}\" \\\n -DCMAKE_C_COMPILER_LAUNCHER=ccache \\\n -DCMAKE_CXX_COMPILER_LAUNCHER=ccache\n cmake --build examples/c++/build-android-examples-${{ matrix.abi }} --parallel\n else\n echo \"Using cached examples build\"\n fi\n\n - name: Run on Android emulator (arm64) and verify\n uses: reactivecircus/android-emulator-runner@v2\n with:\n api-level: ${{ matrix.api }}\n arch: ${{ matrix.abi }}\n # target: google_apis\n # emulator-options: -no-window -gpu swiftshader_indirect -noaudio -no-boot-anim -netdelay none -netspeed full\n # disable-animations: true\n script: |\n adb wait-for-device\n\n echo \"Device ABI:\"\n adb shell getprop ro.product.cpu.abi\n adb shell getprop ro.product.cpu.abilist\n\n echo \"=== CPU ISA / Instruction Set Support ===\"\n echo \"--- /proc/cpuinfo flags ---\"\n adb shell 'cat /proc/cpuinfo | grep -E \"^(Features|flags)\"'\n\n echo \"Checking binary sizes:\"\n ls -lah examples/c++/build-android-examples-${{ matrix.abi }}/\n\n # Push executables to device\n adb push examples/c++/build-android-examples-${{ matrix.abi }}/ailego-example /data/local/tmp/\n adb push examples/c++/build-android-examples-${{ matrix.abi }}/core-example /data/local/tmp/\n adb push examples/c++/build-android-examples-${{ matrix.abi }}/db-example /data/local/tmp/\n\n adb shell chmod 755 /data/local/tmp/ailego-example\n adb shell chmod 755 /data/local/tmp/core-example\n adb shell chmod 755 /data/local/tmp/db-example\n\n echo \"File info on device:\"\n adb shell ls -la /data/local/tmp/ailego-example\n adb shell ls -la /data/local/tmp/core-example\n adb shell ls -la /data/local/tmp/db-example\n\n echo \"Running ailego example:\"\n adb shell 'cd /data/local/tmp && ./ailego-example'\n\n echo \"Running core example:\"\n adb shell 'cd /data/local/tmp && ./core-example'\n\n echo \"Running db example:\"\n adb shell 'cd /data/local/tmp && ./db-example'\n" }, { "path": ".github/workflows/_build_wheel_job.yml", "content": "name: \"(Reusable) Build, Publish and Smoke-test a Wheel\"\n\non:\n workflow_call:\n inputs:\n runner:\n description: \"GitHub Actions runner label\"\n required: true\n type: string\n pypi_repository_url:\n description: \"PyPI repository URL (empty string means official PyPI)\"\n required: false\n type: string\n default: \"\"\n secrets:\n PYPI_API_TOKEN:\n required: true\n\njobs:\n build_publish_test:\n name: Build / publish / smoke-test on ${{ inputs.runner }}\n runs-on: ${{ inputs.runner }}\n permissions:\n contents: read\n\n steps:\n - name: Checkout code\n uses: actions/checkout@v6\n with:\n submodules: recursive\n\n - name: Set up Python (for cibuildwheel controller)\n uses: actions/setup-python@v6\n with:\n python-version: '3.11'\n\n - name: Install cibuildwheel\n run: |\n pip install --upgrade pip\n pip install cibuildwheel==3.4.0\n\n - name: Build wheels using cibuildwheel\n run: |\n python -m cibuildwheel --output-dir wheelhouse\n # Save list of built wheels for publishing\n ls wheelhouse/*.whl | tee $GITHUB_STEP_SUMMARY\n echo \"wheels=$(ls wheelhouse/*.whl | tr '\\n' ' ')\" >> $GITHUB_ENV\n\n - name: Publish to PyPI\n if: success() && github.event_name == 'workflow_dispatch'\n env:\n TWINE_USERNAME: __token__\n TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}\n TWINE_REPOSITORY_URL: ${{ inputs.pypi_repository_url }}\n run: |\n pip install twine\n twine upload --skip-existing --verbose wheelhouse/*.whl\n\n - name: Smoke test from PyPI\n if: success() && github.event_name == 'workflow_dispatch'\n shell: bash\n env:\n PYPI_REPOSITORY_URL: ${{ inputs.pypi_repository_url }}\n run: |\n # Extract version from wheel filename (e.g. zvec-0.2.1.dev24-cp311-...whl -> 0.2.1.dev24)\n WHEEL_FILE=$(ls wheelhouse/zvec-*.whl | head -1)\n ZVEC_VERSION=$(basename \"$WHEEL_FILE\" | sed 's/zvec-\\([^-]*\\)-.*/\\1/')\n\n # Build index-url flags: use TestPyPI when repository URL is set, otherwise official PyPI\n if [ -n \"$PYPI_REPOSITORY_URL\" ]; then\n INDEX_FLAGS=\"--index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/\"\n echo \"Waiting for zvec==$ZVEC_VERSION to become available on TestPyPI...\"\n else\n INDEX_FLAGS=\"\"\n echo \"Waiting for zvec==$ZVEC_VERSION to become available on PyPI...\"\n fi\n # Poll until the version is available (max 5 minutes)\n FOUND=0\n for i in $(seq 1 30); do\n if pip install $INDEX_FLAGS --dry-run \"zvec==$ZVEC_VERSION\" > /dev/null 2>&1; then\n echo \"Version $ZVEC_VERSION is available.\"\n FOUND=1\n break\n fi\n echo \"Attempt $i/30: not yet available, retrying in 10s...\"\n sleep 10\n done\n\n if [ \"$FOUND\" -eq 0 ]; then\n echo \"ERROR: Timed out (5 min) waiting for zvec==$ZVEC_VERSION on PyPI. Aborting smoke test.\"\n exit 1\n fi\n\n # Create a clean venv and install\n python -m venv test_env\n source test_env/bin/activate\n pip install --upgrade pip\n pip install $INDEX_FLAGS \"zvec==$ZVEC_VERSION\"\n pip install --upgrade pip\n pip install $INDEX_FLAGS \"zvec==$ZVEC_VERSION\"\n # Run a simple smoke test\n python -c \"import zvec; print('Import OK:', zvec.__version__)\"\n" }, { "path": ".github/workflows/build_test_wheel.yml", "content": "name: Build Test PyPi Wheels\n\non:\n workflow_dispatch:\n\npermissions:\n contents: read\n\njobs:\n build_wheels_linux_x64:\n name: Build wheels on ubuntu-24.04 (x64) for TestPyPi\n uses: ./.github/workflows/_build_wheel_job.yml\n with:\n runner: ubuntu-24.04\n pypi_repository_url: https://test.pypi.org/legacy/\n secrets:\n PYPI_API_TOKEN: ${{ secrets.TEST_PYPI_API_TOKEN }}\n\n build_wheels_linux_arm64:\n name: Build wheels on ubuntu-24.04-arm (arm64) for TestPyPi\n uses: ./.github/workflows/_build_wheel_job.yml\n with:\n runner: ubuntu-24.04-arm\n pypi_repository_url: https://test.pypi.org/legacy/\n secrets:\n PYPI_API_TOKEN: ${{ secrets.TEST_PYPI_API_TOKEN }}\n\n build_wheels_macos_arm64:\n name: Build wheels on macos-15 (arm64) for TestPyPi\n uses: ./.github/workflows/_build_wheel_job.yml\n with:\n runner: macos-15\n pypi_repository_url: https://test.pypi.org/legacy/\n secrets:\n PYPI_API_TOKEN: ${{ secrets.TEST_PYPI_API_TOKEN }}\n" }, { "path": ".github/workflows/build_wheel.yml", "content": "name: Build Wheels\n\non:\n workflow_dispatch:\n\npermissions:\n contents: read\n\njobs:\n build_wheels_linux_x64:\n name: Build wheels on ubuntu-24.04 (x64) for PyPi\n uses: ./.github/workflows/_build_wheel_job.yml\n with:\n runner: ubuntu-24.04\n secrets:\n PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}\n\n build_wheels_linux_arm64:\n name: Build wheels on ubuntu-24.04-arm (arm64) for PyPi\n uses: ./.github/workflows/_build_wheel_job.yml\n with:\n runner: ubuntu-24.04-arm\n secrets:\n PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}\n\n build_wheels_macos_arm64:\n name: Build wheels on macos-15 (arm64) for PyPi\n uses: ./.github/workflows/_build_wheel_job.yml\n with:\n runner: macos-15\n secrets:\n PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}\n" }, { "path": ".github/workflows/continuous_bench.yml", "content": "name: Continuous Benchmark\non:\n push:\n branches: [ \"main\", \"ci/continuous_bench_squash\" ]\n paths-ignore:\n - '**.md'\n workflow_dispatch:\n\nconcurrency:\n group: cb-${{ github.workflow }}-${{ github.ref }}\n cancel-in-progress: true\n\npermissions:\n contents: read\n\njobs:\n benchmark:\n runs-on: vdbbench\n steps:\n - uses: actions/checkout@v6\n\n - name: Run VectorDBBench\n env:\n DATABASE_URL: ${{ secrets.DATABASE_URL }}\n run: |\n bash .github/workflows/scripts/run_vdb.sh" }, { "path": ".github/workflows/docker/Dockerfile.linux_x64_glibc228", "content": "# =============================================================================\n# Dockerfile.linux_x64_glibc228\n# Purpose: Ubuntu 18.10 gcc-9 + glibc 2.28 + CMake 3.30.0 + PyBind11 build environment\n# Warning: ubuntu:18.10 is EOL; use only for glibc 2.28 compatibility testing.\n# =============================================================================\n\n# Use official Ubuntu 18.10 (Cosmic Cuttlefish)\n# glibc version: 2.28 (confirmed via `ldd --version`)\nFROM ubuntu:18.10\n\n# Replace Ubuntu mirror with old-releases.ubuntu.com for older glibc compatibility\nRUN sed -i 's|http://\\(.*\\)/ubuntu|http://old-releases.ubuntu.com/ubuntu|g' /etc/apt/sources.list && \\\n sed -i 's|http://security.ubuntu.com/ubuntu|http://old-releases.ubuntu.com/ubuntu|g' /etc/apt/sources.list\n\n# Add Ubuntu 20.04 (focal) repo for GCC 9 ONLY\nRUN echo \"deb http://archive.ubuntu.com/ubuntu/ focal main universe\" >> /etc/apt/sources.list && \\\n echo \"deb http://security.ubuntu.com/ubuntu/ focal-security main universe\" >> /etc/apt/sources.list\n\n# Prevent interactive prompts & set non-root user\nENV DEBIAN_FRONTEND=noninteractive \\\n TZ=Etc/UTC\n\n# Create non-root user for safety (optional but recommended)\nRUN useradd -m -u 1000 builder && \\\n mkdir -p /workspace && chown builder:builder /workspace\n\n# Install base system dependencies\nRUN apt-get update && \\\n apt-get install -y --no-install-recommends \\\n build-essential \\\n gcc-9 g++-9 \\\n ninja-build git curl ca-certificates vim wget lcov gnupg clang-format-18\\\n rsync lsb-release \\\n uuid-dev zlib1g-dev libssl-dev libffi-dev \\\n pybind11-dev && \\\n update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-9 90 \\\n --slave /usr/bin/g++ g++ /usr/bin/g++-9 && \\\n rm -rf /var/lib/apt/lists/*\n\n# Install Miniforge (Conda) as root, then assign to builder\nENV MINIFORGE_VERSION=\"latest\"\nENV MINIFORGE_HOME=\"/opt/miniforge3\"\n\nRUN curl -sSL \"https://github.com/conda-forge/miniforge/releases/${MINIFORGE_VERSION}/download/Miniforge3-Linux-x86_64.sh\" -o miniforge.sh && \\\n bash miniforge.sh -b -p ${MINIFORGE_HOME} && \\\n rm miniforge.sh && \\\n chown -R builder:builder ${MINIFORGE_HOME}\n\n# Switch to non-root user\nUSER builder\nENV PATH=\"${MINIFORGE_HOME}/bin:${PATH}\"\nWORKDIR /workspace\n\n# Create conda envs for supported Python versions\nRUN conda create -n py310 python=3.10 -y && \\\n conda create -n py311 python=3.11 -y && \\\n conda create -n py312 python=3.12 -y\nRUN conda clean --all -f -y\n\n# Install CMake 3.30.0 from Kitware official binary\n# Ref: https://github.com/Kitware/CMake/releases/tag/v3.30.0\nRUN mkdir -p /tmp/cmake && cd /tmp/cmake && \\\n curl -sSL -o cmake.tar.gz \\\n \"https://github.com/Kitware/CMake/releases/download/v3.30.0/cmake-3.30.0-linux-x86_64.tar.gz\" && \\\n tar -xzf cmake.tar.gz --strip-components=1 -C /tmp/cmake && \\\n mkdir -p /home/builder/.local && \\\n mv * /home/builder/.local/ && \\\n chown -R builder:builder /home/builder/.local && \\\n rm -rf /tmp/cmake\n\n# Add CMake to PATH\nENV PATH=\"/home/builder/.local/bin:${PATH}\"\n\n# Verify installations\nRUN cmake --version && \\\n conda info && \\\n conda env list && \\\n python --version && \\\n gcc --version && \\\n ldd --version | head -n1\n\n# Final setup\nWORKDIR /workspace" }, { "path": ".github/workflows/nightly_coverage.yml", "content": "name: Nightly Coverage Report\n\non:\n schedule:\n # Runs daily at 00:00 CST (China Standard Time) = 16:00 UTC\n - cron: '0 16 * * *'\n\n workflow_dispatch:\n\npermissions:\n contents: read\n\njobs:\n coverage:\n name: Nightly Coverage Report\n runs-on: ubuntu-24.04\n\n strategy:\n matrix:\n python-version: ['3.10']\n fail-fast: false\n\n steps:\n - name: Checkout code\n uses: actions/checkout@v6\n with:\n ref: main # Always use main for nightly\n submodules: recursive\n\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: ${{ matrix.python-version }}\n cache: 'pip'\n cache-dependency-path: 'pyproject.toml'\n\n - name: Set up environment variables\n run: |\n # Set number of processors for parallel builds\n NPROC=$(nproc 2>/dev/null || echo 2)\n echo \"NPROC=$NPROC\" >> $GITHUB_ENV\n echo \"Using $NPROC parallel jobs for builds\"\n \n # Add Python user base bin to PATH for pip-installed CLI tools\n echo \"$(python -c 'import site; print(site.USER_BASE)')/bin\" >> $GITHUB_PATH\n shell: bash\n\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip \\\n cmake==3.30.0 \\\n ninja==1.11.1 \\\n pytest \\\n pytest-cov \\\n scikit-build-core \\\n setuptools_scm\n shell: bash\n\n - name: Build with COVERAGE config\n run: |\n cd \"$GITHUB_WORKSPACE\"\n \n CMAKE_GENERATOR=\"Unix Makefiles\" \\\n CMAKE_BUILD_PARALLEL_LEVEL=\"$NPROC\" \\\n python -m pip install -v . \\\n --no-build-isolation \\\n --config-settings=\"cmake.build-type=COVERAGE\" \\\n --config-settings='cmake.define.ENABLE_ZEN3=\"ON\"'\n shell: bash\n\n - name: Run Python Tests with Coverage\n run: |\n cd \"$GITHUB_WORKSPACE\"\n python -m pytest python/tests/ --cov=zvec --cov-report=xml\n shell: bash\n\n - name: Run C++ Tests and Generate Coverage\n run: |\n cd \"$GITHUB_WORKSPACE/build\"\n make unittest -j$NPROC\n cd \"$GITHUB_WORKSPACE\"\n # Ensure gcov.sh is executable\n chmod +x scripts/gcov.sh\n bash scripts/gcov.sh -k\n shell: bash\n\n - name: Upload Coverage to Codecov\n uses: codecov/codecov-action@v5\n with:\n files: ./proxima-zvec-filtered.lcov.info,./coverage.xml\n flags: python,cpp,nightly\n name: nightly-linux-py${{ matrix.python-version }}\n token: ${{ secrets.CODECOV_TOKEN }}\n" }, { "path": ".github/workflows/scripts/run_vdb.sh", "content": "set -e\n\nQUANTIZE_TYPE_LIST=\"int8 int4 fp16 fp32\"\nCASE_TYPE_LIST=\"Performance768D1M Performance768D10M Performance1536D500K\" # respectively test cosine, ip # Performance960D1M l2 metrics\nLOG_FILE=\"bench.log\"\nDATE=$(date +%Y-%m-%d_%H-%M-%S)\nNPROC=$(nproc 2>/dev/null || getconf _NPROCESSORS_ONLN 2>/dev/null || echo 2)\n\n# COMMIT_ID = branch-date-sha\nCOMMIT_ID=${GITHUB_REF_NAME}-\"$DATE\"-$(echo ${GITHUB_WORKFLOW_SHA} | cut -c1-8)\nCOMMIT_ID=$(echo \"$COMMIT_ID\" | sed 's/\\//_/g')\necho \"COMMIT_ID: $COMMIT_ID\"\necho \"GITHUB_WORKFLOW_SHA: $GITHUB_WORKFLOW_SHA\"\necho \"workspace: $GITHUB_WORKSPACE\"\nDB_LABEL_PREFIX=\"Zvec16c64g-$COMMIT_ID\"\n\n# install zvec\ngit submodule update --init\n\n# for debug\n#cd ..\n#export SKBUILD_BUILD_DIR=\"$GITHUB_WORKSPACE/../build\"\npwd\n\npython3 -m venv .venv\nsource .venv/bin/activate\npip install cmake ninja psycopg2-binary loguru fire\npip install -e /opt/VectorDBBench\n\nCMAKE_GENERATOR=\"Unix Makefiles\" \\\nCMAKE_BUILD_PARALLEL_LEVEL=\"$NPROC\" \\\npip install -v \"$GITHUB_WORKSPACE\"\n\nfor CASE_TYPE in $CASE_TYPE_LIST; do\n echo \"Running VectorDBBench for $CASE_TYPE\"\n DATASET_DESC=\"\"\n if [ \"$CASE_TYPE\" == \"Performance768D1M\" ]; then\n DATASET_DESC=\"Performance768D1M - Cohere Cosine\"\n elif [ \"$CASE_TYPE\" == \"Performance768D10M\" ]; then\n DATASET_DESC=\"Performance768D10M - Cohere Cosine\"\n else\n DATASET_DESC=\"Performance1536D500K - OpenAI IP\"\n fi\n\n for QUANTIZE_TYPE in $QUANTIZE_TYPE_LIST; do\n DB_LABEL=\"$DB_LABEL_PREFIX-$CASE_TYPE-$QUANTIZE_TYPE\"\n echo \"Running VectorDBBench for $DB_LABEL\"\n\n VDB_PARAMS=\"--path ${DB_LABEL} --db-label ${DB_LABEL} --case-type ${CASE_TYPE} --num-concurrency 12,14,16,18,20\"\n if [ \"$CASE_TYPE\" == \"Performance768D1M\" ]; then\n VDB_PARAMS=\"${VDB_PARAMS} --m 15 --ef-search 180\"\n elif [ \"$CASE_TYPE\" == \"Performance768D10M\" ]; then\n VDB_PARAMS=\"${VDB_PARAMS} --m 50 --ef-search 118 --is-using-refiner\"\n else #Performance1536D500K using default params + refiner to monitor performance degradation\n VDB_PARAMS=\"${VDB_PARAMS} --m 50 --ef-search 100 --is-using-refiner\"\n fi\n\n if [ \"$QUANTIZE_TYPE\" == \"fp32\" ]; then\n vectordbbench zvec ${VDB_PARAMS} 2>&1 | tee $LOG_FILE\n else\n vectordbbench zvec ${VDB_PARAMS} --quantize-type \"${QUANTIZE_TYPE}\" 2>&1 | tee $LOG_FILE\n fi\n\n RESULT_JSON_PATH=$(grep -o \"/opt/VectorDBBench/.*\\.json\" $LOG_FILE)\n QPS=$(jq -r '.results[0].metrics.qps' \"$RESULT_JSON_PATH\")\n RECALL=$(jq -r '.results[0].metrics.recall' \"$RESULT_JSON_PATH\")\n LATENCY_P99=$(jq -r '.results[0].metrics.serial_latency_p99' \"$RESULT_JSON_PATH\")\n LOAD_DURATION=$(jq -r '.results[0].metrics.load_duration' \"$RESULT_JSON_PATH\")\n\n #quote the var to avoid space in the label\n label_list=\"case_type=\\\"${CASE_TYPE}\\\",dataset_desc=\\\"${DATASET_DESC}\\\",db_label=\\\"${DB_LABEL}\\\",commit=\\\"${COMMIT_ID}\\\",date=\\\"${DATE}\\\",quantize_type=\\\"${QUANTIZE_TYPE}\\\"\"\n # replace `/` with `_` in label_list\n label_list=$(echo \"$label_list\" | sed 's/\\//_/g')\n cat < prom_metrics.txt\n # TYPE vdb_bench_qps gauge\n vdb_bench_qps{$label_list} $QPS\n # TYPE vdb_bench_recall gauge\n vdb_bench_recall{$label_list} $RECALL\n # TYPE vdb_bench_latency_p99 gauge\n vdb_bench_latency_p99{$label_list} $LATENCY_P99\n # TYPE vdb_bench_load_duration gauge\n vdb_bench_load_duration{$label_list} $LOAD_DURATION\nEOF\n echo \"prom_metrics:\"\n cat prom_metrics.txt\n curl --data-binary @prom_metrics.txt \"http://47.93.34.27:9091/metrics/job/benchmarks-${CASE_TYPE}/case_type/${CASE_TYPE}/quantize_type/${QUANTIZE_TYPE}\" -v\n done\ndone" }, { "path": ".gitignore", "content": ".*\n*~\nbazel-*\nbuild*\nbin/*\nlib/*\nvar/*\nvenv*\ntests/integration/conf/*\ntests/de_integration/conf/*\n**/__pycache__/*\ntests/bench/log/*\ntests/integration/integration\ntests/integration/log\ntests/integration/*.log\ntests/de_integration/log\ntests/de_integration/*.log\n!.git*\n!.clang-format\n!.circleci\n!.drone.yml\nsdk/python/dist/\ncompile_commands.json\ndist\nhtml\n*.lcov.info\n\n# Dependencies\n/node_modules\n\n# Production\n/build\n\n# Generated files\n.docusaurus\n.cache-loader\n\n# Misc\n.DS_Store\n.env.local\n.env.development.local\n.env.test.local\n.env.production.local\n\nnpm-debug.log*\nyarn-debug.log*\nyarn-error.log*\n\nallure-*\n\n!build_android.sh" }, { "path": ".gitmodules", "content": "[submodule \"thirdparty/googletest/googletest-1.10.0\"]\n\tpath = thirdparty/googletest/googletest-1.10.0\n\turl = https://github.com/google/googletest.git\n[submodule \"thirdparty/sparsehash/sparsehash-2.0.4\"]\n\tpath = thirdparty/sparsehash/sparsehash-2.0.4\n\turl = https://github.com/sparsehash/sparsehash.git\n\tignore = untracked\n[submodule \"thirdparty/gflags/gflags-2.2.2\"]\n\tpath = thirdparty/gflags/gflags-2.2.2\n\turl = https://github.com/gflags/gflags.git\n[submodule \"thirdparty/rocksdb/rocksdb-8.1.1\"]\n\tpath = thirdparty/rocksdb/rocksdb-8.1.1\n\turl = https://github.com/facebook/rocksdb.git\n\tignore = all\n[submodule \"thirdparty/yaml-cpp/yaml-cpp-0.6.3\"]\n\tpath = thirdparty/yaml-cpp/yaml-cpp-0.6.3\n\turl = https://github.com/jbeder/yaml-cpp.git\n[submodule \"thirdparty/arrow/apache-arrow-21.0.0\"]\n\tpath = thirdparty/arrow/apache-arrow-21.0.0\n\turl = https://github.com/apache/arrow.git\n\tignore = all\n[submodule \"thirdparty/CRoaring/CRoaring-2.0.4\"]\n\tpath = thirdparty/CRoaring/CRoaring-2.0.4\n\turl = https://github.com/RoaringBitmap/CRoaring.git\n[submodule \"thirdparty/glog/glog-0.5.0\"]\n\tpath = thirdparty/glog/glog-0.5.0\n\turl = https://github.com/google/glog.git\n\tignore = all\n[submodule \"thirdparty/protobuf/protobuf-3.21.12\"]\n\tpath = thirdparty/protobuf/protobuf-3.21.12\n\turl = https://github.com/protocolbuffers/protobuf.git\n[submodule \"thirdparty/lz4/lz4-1.9.4\"]\n\tpath = thirdparty/lz4/lz4-1.9.4\n\turl = https://github.com/lz4/lz4.git\n[submodule \"thirdparty/antlr/antlr4\"]\n\tpath = thirdparty/antlr/antlr4\n\turl = https://github.com/antlr/antlr4.git\n\tignore = all\n[submodule \"thirdparty/magic_enum/magic_enum-0.9.7\"]\n\tpath = thirdparty/magic_enum/magic_enum-0.9.7\n\turl = https://github.com/Neargye/magic_enum.git\n\tignore = all\n[submodule \"thirdparty/RaBitQ-Library/RaBitQ-Library-0.1\"]\n\tpath = thirdparty/RaBitQ-Library/RaBitQ-Library-0.1\n\turl = https://github.com/VectorDB-NTU/RaBitQ-Library.git\n" }, { "path": "CMakeLists.txt", "content": "cmake_minimum_required(VERSION 3.13)\ncmake_policy(SET CMP0077 NEW)\nproject(zvec)\nset(CC_CXX_STANDARD 17)\n\nset(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS} -Wall -Werror=return-type\")\nset(CMAKE_C_FLAGS \"${CMAKE_C_FLAGS} -Wall -Werror=return-type\")\n\nif(CMAKE_CXX_COMPILER_ID STREQUAL \"GNU\")\n set(CMAKE_EXE_LINKER_FLAGS \"${CMAKE_EXE_LINKER_FLAGS} -Wl,--no-as-needed\")\n set(CMAKE_SHARED_LINKER_FLAGS \"${CMAKE_SHARED_LINKER_FLAGS} -Wl,--no-as-needed\")\nendif()\n\nif(NOT DEFINED PROJECT_ROOT_DIR OR NOT PROJECT_ROOT_DIR)\n set(PROJECT_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR} CACHE PATH \"Root directory of the project\" FORCE)\nendif()\n\nmessage(STATUS \"PROJECT_ROOT_DIR = ${PROJECT_ROOT_DIR}\")\n\ninclude(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\nif(NOT ANDROID AND AUTO_DETECT_ARCH AND CMAKE_SYSTEM_PROCESSOR MATCHES \"x86_64|i686|i386|x64\")\n setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512 MATH_MARCH_FLAG_AVX512FP16)\n message(STATUS \"best compiler march, sse: \" ${MATH_MARCH_FLAG_SSE} \", avx2: \" ${MATH_MARCH_FLAG_AVX2} \", avx512: \" ${MATH_MARCH_FLAG_AVX512} \", avx512fp16: \" ${MATH_MARCH_FLAG_AVX512FP16})\nendif()\n\ninclude_directories(${PROJECT_ROOT_DIR}/src/include)\ninclude_directories(${PROJECT_ROOT_DIR}/src)\n\noption(BUILD_PYTHON_BINDINGS \"Build Python bindings using pybind11\" OFF)\nmessage(STATUS \"BUILD_PYTHON_BINDINGS:${BUILD_PYTHON_BINDINGS}\")\n\noption(BUILD_TOOLS \"Build tools\" ON)\nmessage(STATUS \"BUILD_TOOLS:${BUILD_TOOLS}\")\n\noption(RABITQ_ENABLE_AVX512 \"Compile RaBitQ with AVX-512 support\" OFF)\n\nif(CMAKE_SYSTEM_NAME STREQUAL \"Linux\" AND CMAKE_SYSTEM_PROCESSOR MATCHES \"x86_64|amd64|AMD64\" AND NOT ANDROID)\n include(CheckCCompilerFlag)\n\n check_c_compiler_flag(\"-mavx2\" COMPILER_SUPPORTS_AVX2)\n check_c_compiler_flag(\"-mavx512f -mavx512bw -mavx512vl\" COMPILER_SUPPORTS_AVX512)\n\n if(COMPILER_SUPPORTS_AVX2 OR COMPILER_SUPPORTS_AVX512)\n set(RABITQ_SUPPORTED ON)\n add_definitions(-DRABITQ_SUPPORTED=1)\n if(RABITQ_ENABLE_AVX512 AND COMPILER_SUPPORTS_AVX512)\n add_definitions(-DRABITQ_COMPILED_AVX512=1)\n set(RABITQ_ARCH_FLAG \"${MATH_MARCH_FLAG_AVX512}\")\n else()\n set(RABITQ_ARCH_FLAG \"${MATH_MARCH_FLAG_AVX2}\")\n endif()\n else()\n set(RABITQ_SUPPORTED OFF)\n add_definitions(-DRABITQ_SUPPORTED=0)\n message(STATUS \"RaBitQ support disabled - compiler does not support AVX2 or AVX-512\")\n endif()\nelse()\n set(RABITQ_SUPPORTED OFF)\n add_definitions(-DRABITQ_SUPPORTED=0)\n message(STATUS \"RaBitQ support disabled - only supported on Linux x86_64\")\nendif()\nmessage(STATUS \"RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}\")\n\noption(USE_OSS_MIRROR \"Use OSS mirror for faster third-party downloads\" OFF)\nif(DEFINED ENV{USE_OSS_MIRROR} AND NOT \"$ENV{USE_OSS_MIRROR}\" STREQUAL \"\")\n set(USE_OSS_MIRROR \"$ENV{USE_OSS_MIRROR}\" CACHE BOOL \"Use OSS mirror for faster third-party downloads\" FORCE)\nendif()\nmessage(STATUS \"USE_OSS_MIRROR:${USE_OSS_MIRROR}\")\n\ncc_directory(thirdparty)\ncc_directories(src)\ncc_directories(tests)\n\nif(BUILD_TOOLS)\n cc_directories(tools)\nendif()\n\ngit_version(GIT_SRCS_VER ${PROJECT_ROOT_DIR})\nset(CPACK_PACKAGE_VERSION ${GIT_SRCS_VER})\nset(CPACK_PACKAGE_NAME zvec)\ninclude(CPack)\n\nif(BUILD_PYTHON_BINDINGS)\n if(APPLE)\n set(CMAKE_STRIP \"\")\n message(STATUS \"Disabled strip on macOS to preserve code signature\")\n endif()\n\n include(GNUInstallDirs)\n if(DEFINED SKBUILD_PLATLIB_DIR)\n set(ZVEC_PY_INSTALL_DIR \"${SKBUILD_PLATLIB_DIR}\")\n elseif(DEFINED Python_SITEARCH)\n set(ZVEC_PY_INSTALL_DIR \"${Python_SITEARCH}\")\n else()\n set(ZVEC_PY_INSTALL_DIR \"${CMAKE_INSTALL_LIBDIR}\")\n endif()\n\n message(STATUS \"Zvec install path: ${ZVEC_PY_INSTALL_DIR}\")\n install(TARGETS _zvec LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR})\nendif()\n" }, { "path": "CODE_OF_CONDUCT.md", "content": "# Zvec Code of Conduct\n\n## Our Pledge\n\nWe 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.\n\n## Expected Behavior\n\n- Use welcoming and inclusive language\n- Respect differing viewpoints and experiences\n- Gracefully accept constructive criticism\n- Focus on what is best for the community\n- Show empathy and kindness toward others\n\n## Unacceptable Behavior\n\n- Harassment, intimidation, or discriminatory conduct\n- Trolling, insulting, or derogatory comments\n- Public or private harassment\n- Publishing others’ private information without consent\n- Any conduct that would reasonably be considered inappropriate in a professional setting\n\n## Enforcement\n\nInstances of abusive, harassing, or otherwise unacceptable behavior may be reported by contacting the project team \nat **zvec@alibaba-inc.com** (replace with your contact). All complaints will be reviewed and investigated promptly \nand fairly.\n\nThe project team is obligated to respect the privacy and security of the reporter.\n\nConsequences may include:\n- A formal warning\n- Temporary or permanent ban from project spaces\n- Removal of contributions (e.g. comments, PRs)\n\n## Attribution\n\nThis Code of Conduct is adapted from the [Contributor Covenant][homepage], version 2.1, available at \nhttps://www.contributor-covenant.org/version/2/1/code_of_conduct.html\n\n[homepage]: https://www.contributor-covenant.org" }, { "path": "CONTRIBUTING.md", "content": "# Contributing to Zvec\n\nFirst off, thank you for considering contributing to Zvec! 🙌 \nWhether you're reporting a bug, proposing a feature, improving documentation, or submitting code — every contribution helps make Zvec better.\n\n## Code of Conduct\n\nBy participating, you agree to abide by our [Code of Conduct](CODE_OF_CONDUCT.md). Please be respectful, collaborative, and inclusive.\n\n---\n\n## Development Setup\n\n### Prerequisites\n- Python 3.10 - 3.12\n- CMake ≥ 3.26, < 4.0 (`cmake --version`)\n- A C++17-compatible compiler (e.g., `g++-11+`, `clang++`, Apple Clang on macOS)\n\n### Clone & Initialize\n\n```bash\ngit clone --recursive https://github.com/alibaba/zvec.git\ncd zvec\n```\n\n> 💡 **Tip** \n> - Forgot `--recursive`? Run: \n> ```bash\n> git submodule update --init --recursive\n> ```\n> - Set up pre-commit hooks: \n> ```bash\n> pip install pre-commit && pre-commit install\n> ```\n\n### Build from Source (Editable Install)\n```bash\npip install -e \".[dev]\"\n# This installs dev dependencies (pytest, ruff, etc.) and builds the C++ extension in-place\n```\n\n> ✅ Verify:\n> ```bash\n> python -c \"import zvec; print('Success!')\"\n> ```\n\n---\n\n## Testing\n\n### Run All Tests\n```bash\npytest python/tests/ -v\n```\n\n### Run with Coverage (Debug/CI)\n```bash\npytest python/tests/ --cov=zvec --cov-report=term-missing\n```\n\n> 🔎 See full rules in `[tool.ruff]` section of `pyproject.toml`.\n\n---\n\n## Build Customization\n\nYou can control build behavior via environment variables or `pyproject.toml`:\n\n| Option | How to Set | Description |\n|--------|------------|-------------|\n| **Build Type** | `CMAKE_BUILD_TYPE=Debug` | `Debug`, `Release`, or `Coverage` (for gcov/lcov) |\n| **Generator** | `CMAKE_GENERATOR=\"Unix Makefiles\"` | Default: `Ninja`; use Make if preferred |\n| **AVX-512** | `ENABLE_SKYLAKE_AVX512=ON` | Enable AVX-512 optimizations (x86_64 only) |\n\nExample (Debug + Make):\n```bash\nCMAKE_BUILD_TYPE=Debug CMAKE_GENERATOR=\"Unix Makefiles\" pip install -v .\n```\n\n---\n\n## Submitting Changes\n\n1. Fork the repo and create a feature branch (`feat/...`, `fix/...`, `docs/...`)\n2. Write clear commit messages (e.g., `fix(query): handle null vector in dense_fp32`)\n3. Ensure tests pass & linter is clean\n4. Open a Pull Request to `main`\n5. Link related issue (e.g., `Closes #123`)\n\n✅ **PRs should include**:\n- Test coverage for new behavior\n- Updates to documentation (if applicable)\n- Reasoning behind non-obvious design choices\n\n---\n\n## Documentation\n\n- User guides: `docs/` (built with MkDocs)\n- API reference: generated from docstrings (follow [Google style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings))\n- Build & deploy: `mkdocs serve` / `mkdocs build`\n\n---\n\n## Need Help\n\n- Browse [existing issues](https://github.com/alibaba/zvec/issues)\n- For sensitive/security issues: email `zvec@alibaba-inc.com`\n\n---\n\n✨ Thanks again for being part of Zvec!\n" }, { "path": "LICENSE", "content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License." }, { "path": "README.md", "content": "
\n \n \n \"zvec\n \n
\n\n

\n \"Code\n \"Main\"/\n \"License\"/\n \"PyPI\n \"Python\n \"npm\n

\n\n

\n \"alibaba%2Fzvec\n

\n\n

\n 🚀 Quickstart |\n 🏠 Home |\n 📚 Docs |\n 📊 Benchmarks |\n 🔎 DeepWiki |\n 🎮 Discord \n

\n\n**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.\n\n## 💫 Features\n\n- **Blazing Fast**: Searches billions of vectors in milliseconds.\n- **Simple, Just Works**: [Install](#-installation) and start searching in seconds. No servers, no config, no fuss.\n- **Dense + Sparse Vectors**: Work with both dense and sparse embeddings, with native support for multi-vector queries in a single call.\n- **Hybrid Search**: Combine semantic similarity with structured filters for precise results.\n- **Runs Anywhere**: As an in-process library, Zvec runs wherever your code runs — notebooks, servers, CLI tools, or even edge devices.\n\n## 📦 Installation\n\n### [Python](https://pypi.org/project/zvec/)\n\n**Requirements**: Python 3.10 - 3.12\n\n```bash\npip install zvec\n```\n\n### [Node.js](https://www.npmjs.com/package/@zvec/zvec)\n\n```bash\nnpm install @zvec/zvec\n```\n\n### ✅ Supported Platforms\n\n- Linux (x86_64, ARM64)\n- macOS (ARM64)\n\n### 🛠️ Building from Source\n\nIf you prefer to build Zvec from source, please check the [Building from Source](https://zvec.org/en/docs/build/) guide.\n\n## ⚡ One-Minute Example\n\n```python\nimport zvec\n\n# Define collection schema\nschema = zvec.CollectionSchema(\n name=\"example\",\n vectors=zvec.VectorSchema(\"embedding\", zvec.DataType.VECTOR_FP32, 4),\n)\n\n# Create collection\ncollection = zvec.create_and_open(path=\"./zvec_example\", schema=schema)\n\n# Insert documents\ncollection.insert([\n zvec.Doc(id=\"doc_1\", vectors={\"embedding\": [0.1, 0.2, 0.3, 0.4]}),\n zvec.Doc(id=\"doc_2\", vectors={\"embedding\": [0.2, 0.3, 0.4, 0.1]}),\n])\n\n# Search by vector similarity\nresults = collection.query(\n zvec.VectorQuery(\"embedding\", vector=[0.4, 0.3, 0.3, 0.1]),\n topk=10\n)\n\n# Results: list of {'id': str, 'score': float, ...}, sorted by relevance\nprint(results)\n```\n\n## 📈 Performance at Scale\n\nZvec delivers exceptional speed and efficiency, making it ideal for demanding production workloads.\n\n\"Zvec\n\nFor detailed benchmark methodology, configurations, and complete results, please see our [Benchmarks documentation](https://zvec.org/en/docs/benchmarks/).\n\n## 🤝 Join Our Community\n\n
\n\nStay updated and get support — scan or click:\n\n
\n\n| 💬 DingTalk | 📱 WeChat | 🎮 Discord |\n|:---:|:---:|:---:|\n| | | [![Discord](https://img.shields.io/badge/Discord-Join%20Server-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/rKddFBBu9z) |\n| Scan to join | Scan to join | Click to join |\n\n
\n\n
\n\n## ❤️ Contributing\n\nWe 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.\n\nCheck out our [Contributing Guide](./CONTRIBUTING.md) to get started!\n" }, { "path": "cmake/bazel.cmake", "content": "##\n## The following functions used by user's CMakeLists.txt:\n##\n\n## 1. Functions for C/C++\n##\n## 1.1. Add a subdirectory to the build\n## cc_directory( [binary_dir])\n##\n## 1.2. Add subdirectories to the build\n## cc_directories( [source_dir2 ...])\n##\n## 1.3. Build a C/C++ static or shared library\n## cc_library(\n## NAME \n## [STATIC] [SHARED] [STRICT] [ALWAYS_LINK] [EXCLUDE] [PACKED] [SRCS_NO_GLOB]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [PUBINCS public_dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [PACKED_EXCLUDES pattern1 ...]\n## [VERSION ]\n## )\n##\n## 1.4. Build a C/C++ executable program\n## cc_binary(\n## NAME \n## [STRICT] [PACKED]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [VERSION ]\n## )\n##\n## 1.5. Build a C/C++ executable test program\n## cc_test(\n## NAME \n## [STRICT]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [ARGS args1 ...]\n## [VERSION ]\n## )\n##\n## 1.6. Add existing test cases to a test suite\n## cc_test_suite( [test_name ...])\n##\n## 1.7. Import a C/C++ static or shared library\n## cc_import(\n## NAME \n## [STATIC | SHARED] [PACKED]\n## PATH \n## [INCS dir1 ...]\n## [PUBINCS public_dir1 ...]\n## [DEPS target1 ...]\n## [IMPLIB ]\n## [PACKED_EXCLUDES pattern1 ...]\n## )\n##\n## 1.8. Import a C/C++ interface library\n## cc_interface(\n## NAME \n## [PACKED]\n## [INCS dir1 ...]\n## [PUBINCS public_dir1 ...]\n## [DEPS target1 ...]\n## [PACKED_EXCLUDES pattern1 ...]\n## )\n##\n## 1.9. Build a C/C++ executable google test program\n## cc_gtest(\n## NAME \n## [STRICT]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [ARGS args1 ...]\n## [VERSION ]\n## )\n##\n## 1.10. Build a C/C++ executable google mock program\n## cc_gmock(\n## NAME \n## [STRICT]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [ARGS args1 ...]\n## [VERSION ]\n## )\n##\n## 1.11. Build a C++ protobuf static or shared library\n## cc_proto_library(\n## NAME \n## [STATIC] [SHARED] [STRICT] [EXCLUDE] [PACKED]\n## SRCS [file2.proto ...]\n## [PROTOROOT path]\n## [CXXFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [VERSION ]\n## [PROTOBUF_VERSION ]\n## )\n##\n\n## 2. Functions for CUDA\n##\n## 2.1. Add a subdirectory to the build\n## cuda_directory( [binary_dir])\n##\n## 2.2. Add subdirectories to the build\n## cuda_directories( [source_dir2 ...])\n##\n## 2.3. Build a CUDA static or shared library\n## cuda_library(\n## NAME \n## [STATIC] [SHARED] [STRICT] [ALWAYS_LINK] [EXCLUDE] [PACKED]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [PUBINCS public_dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [CUDAFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [PACKED_EXCLUDES pattern1 ...]\n## [VERSION ]\n## )\n##\n## 2.4. Build a CUDA executable program\n## cuda_binary(\n## NAME \n## [STRICT] [PACKED]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [CUDAFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [VERSION ]\n## )\n##\n## 2.5. Build a CUDA executable test program\n## cuda_test(\n## NAME \n## [STRICT]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [CUDAFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [ARGS args1 ...]\n## [VERSION ]\n## )\n##\n## 2.6. Add existing test cases to a test suite\n## cuda_test_suite( [test_name ...])\n##\n## 2.7. Import a C/C++/CUDA static or shared library\n## cuda_import(\n## NAME \n## [STATIC | SHARED] [PACKED]\n## PATH \n## [INCS dir1 ...]\n## [PUBINCS public_dir1 ...]\n## [DEPS target1 ...]\n## [IMPLIB ]\n## [PACKED_EXCLUDES pattern1 ...]\n## )\n##\n## 2.8. Import a C/C++/CUDA interface library\n## cuda_interface(\n## NAME \n## [PACKED]\n## [INCS dir1 ...]\n## [PUBINCS public_dir1 ...]\n## [DEPS target1 ...]\n## [PACKED_EXCLUDES pattern1 ...]\n## )\n##\n## 2.9. Build a CUDA executable google test program\n## cuda_gtest(\n## NAME \n## [STRICT]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [CUDAFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [ARGS args1 ...]\n## [VERSION ]\n## )\n##\n## 2.10. Build a CUDA executable google mock program\n## cuda_gmock(\n## NAME \n## [STRICT]\n## SRCS [file2 ...]\n## [INCS dir1 ...]\n## [DEFS DEF1=1 ...]\n## [LIBS lib1 ...]\n## [CFLAGS flag1 ...]\n## [CXXFLAGS flag1 ...]\n## [CUDAFLAGS flag1 ...]\n## [LDFLAGS flag1 ...]\n## [DEPS target1 ...]\n## [ARGS args1 ...]\n## [VERSION ]\n## )\n##\n\n## 3. Utility functions\n##\n## 3.1. Download a git repository\n## git_repository(\n## NAME \n## URL \n## [TAG ]\n## [PATH ]\n## )\n##\n## 3.2. Download a hg repository\n## hg_repository(\n## NAME \n## URL \n## [TAG ]\n## [PATH ]\n## )\n##\n## 3.3. Download a svn repository\n## svn_repository(\n## NAME \n## URL \n## [REV ]\n## [PATH ]\n## )\n##\n## 3.4. Download a http archive\n## http_archive(\n## NAME \n## URL \n## [SHA256 | SHA1 | MD5 ]\n## [PATH ]\n## )\n##\n## 3.5. Retrieve a version string from GIT\n## git_version(\n## \n## \n## )\n##\n## 3.6. Retrieve a version string from HG\n## hg_version(\n## \n## \n## )\n##\n## 3.7. Retrieve a version string from SVN\n## svn_version(\n## \n## \n## )\n##\n\ncmake_minimum_required(VERSION 3.1 FATAL_ERROR)\ninclude(CMakeParseArguments)\n\n# Using AppleClang instead of Clang (Compiler id)\nif(POLICY CMP0025)\n cmake_policy(SET CMP0025 NEW)\nendif()\n\n# Enable unit testing\nenable_testing()\n\n# Add unittest target\nif(NOT TARGET unittest)\n add_custom_target(\n unittest\n COMMAND ${CMAKE_CTEST_COMMAND} --output-on-failure\n --build-config $\n )\nendif()\n\n# Directories of target output\nif(NOT CMAKE_ARCHIVE_OUTPUT_DIRECTORY)\n set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/lib)\nendif()\nif(NOT CMAKE_LIBRARY_OUTPUT_DIRECTORY)\n set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/lib)\nendif()\nif(NOT CMAKE_RUNTIME_OUTPUT_DIRECTORY)\n set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/bin)\nendif()\n\n# RPATH settings\nset(CMAKE_MACOSX_RPATH ON)\nif(NOT ${CMAKE_SYSTEM_NAME} MATCHES \"Darwin\")\n set(CMAKE_SKIP_BUILD_RPATH ON)\n set(CMAKE_BUILD_WITH_INSTALL_RPATH ON)\n if(${CMAKE_SIZEOF_VOID_P} EQUAL \"8\")\n set(CMAKE_INSTALL_RPATH \"$ORIGIN/../lib64:$ORIGIN/../lib:$ORIGIN\")\n else()\n set(CMAKE_INSTALL_RPATH \"$ORIGIN/../lib:$ORIGIN\")\n endif()\nelse()\n set(CMAKE_INSTALL_RPATH \"@loader_path/../lib:@loader_path\")\nendif()\n\n# Define standard installation directories\nif(NOT CMAKE_INSTALL_LIBDIR)\n set(CMAKE_INSTALL_LIBDIR lib)\nendif()\nif(NOT CMAKE_INSTALL_BINDIR)\n set(CMAKE_INSTALL_BINDIR bin)\nendif()\nif(NOT CMAKE_INSTALL_INCDIR)\n set(CMAKE_INSTALL_INCDIR include)\nendif()\nif(NOT CMAKE_INSTALL_ETCDIR)\n set(CMAKE_INSTALL_ETCDIR etc)\nendif()\n\n# Generates a compile_commands.json\nset(CMAKE_EXPORT_COMPILE_COMMANDS ON)\nset(CMAKE_POSITION_INDEPENDENT_CODE ON)\nset(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON)\n\nif(APPLE OR ANDROID)\n option(CLANG_USE_LIBCXX \"Use libc++ instead of libstdc++\" ON)\nelse()\n option(CLANG_USE_LIBCXX \"Use libc++ instead of libstdc++\" OFF)\nendif()\n\nset(CLANG_STDLIB_OPTION \"\")\nif(CLANG_USE_LIBCXX)\n set(CLANG_STDLIB_OPTION \"-stdlib=libc++\")\nelse()\n set(CLANG_STDLIB_OPTION \"-stdlib=libstdc++\")\nendif()\n\nif(NOT MSVC)\n # Use color in diagnostics\n set(\n _COMPILER_FLAGS\n \"$<$:-fcolor-diagnostics;${CLANG_STDLIB_OPTION}>\"\n \"$<$:-fcolor-diagnostics>\"\n \"$<$:-fdiagnostics-color=always>\"\n )\n add_compile_options(\n \"$<$:${_COMPILER_FLAGS}>\"\n \"$<$:${_COMPILER_FLAGS}>\"\n )\n unset(_COMPILER_FLAGS)\nelse()\n # Replace the default compiling flags\n set(\n _COMPILER_FLAGS\n CMAKE_CXX_FLAGS\n CMAKE_CXX_FLAGS_DEBUG\n CMAKE_CXX_FLAGS_RELEASE\n CMAKE_CXX_FLAGS_RELWITHDEBINFO\n CMAKE_CXX_FLAGS_MINSIZEREL\n CMAKE_C_FLAGS\n CMAKE_C_FLAGS_DEBUG\n CMAKE_C_FLAGS_RELEASE\n CMAKE_C_FLAGS_RELWITHDEBINFO\n CMAKE_C_FLAGS_MINSIZEREL\n )\n foreach(COMPILER_FLAG ${_COMPILER_FLAGS})\n string(REPLACE \"/MT\" \"/MD\" ${COMPILER_FLAG} \"${${COMPILER_FLAG}}\")\n string(REGEX REPLACE \"/W[0-9]\" \"\" ${COMPILER_FLAG} \"${${COMPILER_FLAG}}\")\n endforeach()\n unset(_COMPILER_FLAGS)\n add_definitions(-D_CRT_SECURE_NO_WARNINGS)\n # Build shared library as default\n set(BUILD_SHARED_LIBS ON)\nendif()\n\nset(CMAKE_C_FLAGS_ASAN ${CMAKE_C_FLAGS_DEBUG})\nset(CMAKE_CXX_FLAGS_ASAN ${CMAKE_CXX_FLAGS_DEBUG})\nset(CMAKE_EXE_LINKER_FLAGS_ASAN ${CMAKE_EXE_LINKER_FLAGS_DEBUG})\nset(CMAKE_SHARED_LINKER_FLAGS_ASAN ${CMAKE_SHARED_LINKER_FLAGS_DEBUG})\nset(CMAKE_STATIC_LINKER_FLAGS_ASAN ${CMAKE_STATIC_LINKER_FLAGS_DEBUG})\nset(CMAKE_MODULE_LINKER_FLAGS_ASAN ${CMAKE_MODULE_LINKER_FLAGS_DEBUG})\nset(CMAKE_C_FLAGS_COVERAGE ${CMAKE_C_FLAGS_DEBUG})\nset(CMAKE_CXX_FLAGS_COVERAGE ${CMAKE_CXX_FLAGS_DEBUG})\nset(CMAKE_EXE_LINKER_FLAGS_COVERAGE ${CMAKE_EXE_LINKER_FLAGS_DEBUG})\nset(CMAKE_SHARED_LINKER_FLAGS_COVERAGE ${CMAKE_SHARED_LINKER_FLAGS_DEBUG})\nset(CMAKE_STATIC_LINKER_FLAGS_COVERAGE ${CMAKE_STATIC_LINKER_FLAGS_DEBUG})\nset(CMAKE_MODULE_LINKER_FLAGS_COVERAGE ${CMAKE_MODULE_LINKER_FLAGS_DEBUG})\n\n# C/C++ ASAN compile flags\nset(\n BAZEL_CC_ASAN_COMPILE_FLAGS\n \"$<$:$<$:-fsanitize=address>>\"\n \"$<$:$<$:-fsanitize=address>>\"\n \"$<$:$<$:-fsanitize=address>>\"\n \"$<$:$<$:/fsanitize=address>>\"\n )\n\n# C/C++ COVERAGE compile flags\nset(\n BAZEL_CC_COVERAGE_COMPILE_FLAGS\n \"$<$:$<$:--coverage>>\"\n \"$<$:$<$:--coverage>>\"\n \"$<$:$<$:--coverage>>\"\n )\n\n# C/C++ strict compile flags\nif(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 7.0)\n set(\n BAZEL_CC_STRICT_COMPILE_FLAGS\n \"$<$:-Wall;-Wextra;-Wshadow>\"\n \"$<$:-Wall;-Wextra;-Wshadow>\"\n \"$<$:-Wall;-Wextra;-Wshadow-local;-Wno-misleading-indentation>\"\n \"$<$:/W4>\"\n ${BAZEL_CC_ASAN_COMPILE_FLAGS}\n ${BAZEL_CC_COVERAGE_COMPILE_FLAGS}\n )\nelse()\n set(\n BAZEL_CC_STRICT_COMPILE_FLAGS\n \"$<$:-Wall;-Wextra;-Wshadow>\"\n \"$<$:-Wall;-Wextra;-Wshadow>\"\n \"$<$:-Wall;-Wextra;-Wshadow;-Wno-misleading-indentation>\"\n \"$<$:/W4>\"\n ${BAZEL_CC_ASAN_COMPILE_FLAGS}\n ${BAZEL_CC_COVERAGE_COMPILE_FLAGS}\n )\nendif()\n\n\n# C/C++ strict link flags\nset(\n BAZEL_CC_STRICT_LINK_FLAGS\n \"$<$:${CLANG_STDLIB_OPTION}>\"\n ${BAZEL_CC_ASAN_COMPILE_FLAGS}\n ${BAZEL_CC_COVERAGE_COMPILE_FLAGS}\n )\n\n# C/C++ unstrict compile flags\nset(\n BAZEL_CC_UNSTRICT_COMPILE_FLAGS\n \"$<$:-Wall>\"\n \"$<$:-Wall>\"\n \"$<$:-Wall>\"\n \"$<$:/W3>\"\n ${BAZEL_CC_ASAN_COMPILE_FLAGS}\n ${BAZEL_CC_COVERAGE_COMPILE_FLAGS}\n )\n\n# C/C++ unstrict link flags\nset(\n BAZEL_CC_UNSTRICT_LINK_FLAGS\n \"$<$:${CLANG_STDLIB_OPTION}>\"\n ${BAZEL_CC_ASAN_COMPILE_FLAGS}\n ${BAZEL_CC_COVERAGE_COMPILE_FLAGS}\n )\n\n# CUDA strict compile flags\nset(\n BAZEL_CUDA_STRICT_COMPILE_FLAGS\n \"$<$:$<$:-Wall;-Wextra;-Wshadow>>\"\n \"$<$:$<$:-Wall;-Wextra;-Wshadow>>\"\n \"$<$:$<$:-Wall;-Wextra;-Wshadow>>\"\n \"$<$:$<$:/W4>>\"\n \"$<$:$<$:-Wall;-Wextra;-Wshadow>>\"\n \"$<$:$<$:-Wall;-Wextra;-Wshadow>>\"\n \"$<$:$<$:-Wall;-Wextra;-Wshadow>>\"\n \"$<$:$<$:/W4>>\"\n \"$<$:$<$:-G>>\"\n )\n\n# CUDA strict link flags\nset(BAZEL_CUDA_STRICT_LINK_FLAGS \"\")\n\n# CUDA unstrict compile flags\nset(\n BAZEL_CUDA_UNSTRICT_COMPILE_FLAGS\n \"$<$:$<$:-Wall>>\"\n \"$<$:$<$:-Wall>>\"\n \"$<$:$<$:-Wall>>\"\n \"$<$:$<$:/W3>>\"\n \"$<$:$<$:-Wall>>\"\n \"$<$:$<$:-Wall>>\"\n \"$<$:$<$:-Wall>>\"\n \"$<$:$<$:/W3>>\"\n \"$<$:$<$:-G>>\"\n )\n\n# CUDA unstrict link flags\nset(BAZEL_CUDA_UNSTRICT_LINK_FLAGS \"\")\n\n## Find workspace directory\nfunction(_find_workspace_directory _RESULT)\n # Find Workspace.cmake folder\n set(_CURRENT_WORKSPACE_DIR ${CMAKE_CURRENT_SOURCE_DIR})\n get_filename_component(\n _PARENT_WORKSPACE_DIR ${_CURRENT_WORKSPACE_DIR} DIRECTORY\n )\n while(NOT (\"${_CURRENT_WORKSPACE_DIR}\" STREQUAL \"${_PARENT_WORKSPACE_DIR}\"))\n if(EXISTS \"${_CURRENT_WORKSPACE_DIR}/Workspace.cmake\")\n set(${_RESULT} ${_CURRENT_WORKSPACE_DIR} PARENT_SCOPE)\n message(STATUS \"Found workspace at ${${_RESULT}}\")\n break()\n endif()\n\n # Find next parent folder\n set(_CURRENT_WORKSPACE_DIR ${_PARENT_WORKSPACE_DIR})\n get_filename_component(\n _PARENT_WORKSPACE_DIR ${_CURRENT_WORKSPACE_DIR} DIRECTORY\n )\n endwhile()\nendfunction()\n\n## Retrieve absolute paths\nfunction(_absolute_paths _RESULT)\n foreach(FILEPATH ${ARGN})\n if(NOT IS_ABSOLUTE ${FILEPATH})\n get_filename_component(FILEPATH ${FILEPATH} ABSOLUTE)\n endif()\n list(APPEND FILEPATHS ${FILEPATH})\n endforeach()\n set(${_RESULT} \"${FILEPATHS}\" PARENT_SCOPE)\nendfunction()\n\n## Add both shared and static library\nmacro(_add_library _NAME _OPTION)\n add_library(${_NAME}_objects OBJECT ${_OPTION} ${ARGN})\n add_library(\n ${_NAME}_static STATIC ${_OPTION} $\n )\n add_library(\n ${_NAME} SHARED ${_OPTION} $\n )\n add_dependencies(${_NAME} ${_NAME}_static)\n if(NOT MSVC)\n set_property(TARGET ${_NAME}_static PROPERTY OUTPUT_NAME ${_NAME})\n endif()\nendmacro()\n\n## Link dependencies\nfunction(_targets_link_dependencies _NAME)\n foreach(LIB ${ARGN})\n if(TARGET ${LIB})\n list(APPEND LIBS_DEPS ${LIB})\n list(\n APPEND LIBS_INCS\n \"$\"\n )\n endif()\n endforeach()\n\n if(LIBS_DEPS)\n add_dependencies(${_NAME} ${LIBS_DEPS})\n target_include_directories(${_NAME} PRIVATE \"${LIBS_INCS}\")\n endif()\nendfunction()\n\n## Link libraries\nfunction(_target_link_libraries _NAME)\n function(_collect_always_link_libs LIB_LIST RESULT_VAR)\n if(NOT _COLLECT_ALWAYS_LINK_VISITED)\n set(_COLLECT_ALWAYS_LINK_VISITED \"\" PARENT_SCOPE)\n endif()\n\n set(LOCAL_RESULT \"\")\n foreach(LIB ${LIB_LIST})\n if(NOT TARGET ${LIB})\n continue()\n endif()\n\n list(FIND _COLLECT_ALWAYS_LINK_VISITED ${LIB} ALREADY_VISITED)\n if(NOT ALREADY_VISITED EQUAL -1)\n continue()\n endif()\n\n list(APPEND _COLLECT_ALWAYS_LINK_VISITED ${LIB})\n set(_COLLECT_ALWAYS_LINK_VISITED \"${_COLLECT_ALWAYS_LINK_VISITED}\" PARENT_SCOPE)\n\n get_target_property(ALWAYS_LINK ${LIB} ALWAYS_LINK)\n if(ALWAYS_LINK)\n list(APPEND LOCAL_RESULT ${LIB})\n endif()\n\n get_target_property(DEP_LIBS ${LIB} INTERFACE_LINK_LIBRARIES)\n if(DEP_LIBS)\n _collect_always_link_libs(\"${DEP_LIBS}\" DEP_ALWAYS_LINK_LIBS)\n list(APPEND LOCAL_RESULT ${DEP_ALWAYS_LINK_LIBS})\n endif()\n\n get_target_property(LINK_LIBS ${LIB} LINK_LIBRARIES)\n if(LINK_LIBS)\n _collect_always_link_libs(\"${LINK_LIBS}\" LINK_ALWAYS_LINK_LIBS)\n list(APPEND LOCAL_RESULT ${LINK_ALWAYS_LINK_LIBS})\n endif()\n endforeach()\n\n list(REMOVE_DUPLICATES LOCAL_RESULT)\n set(${RESULT_VAR} \"${LOCAL_RESULT}\" PARENT_SCOPE)\n endfunction()\n\n _collect_always_link_libs(\"${ARGN}\" ALL_ALWAYS_LINK_LIBS)\n\n set(ALL_LIBS_TO_PROCESS ${ARGN})\n foreach(ALWAYS_LIB ${ALL_ALWAYS_LINK_LIBS})\n list(FIND ARGN ${ALWAYS_LIB} FOUND_INDEX)\n if(FOUND_INDEX EQUAL -1)\n list(APPEND ALL_LIBS_TO_PROCESS ${ALWAYS_LIB})\n endif()\n endforeach()\n\n list(REMOVE_DUPLICATES ALL_LIBS_TO_PROCESS)\n\n foreach(LIB ${ALL_LIBS_TO_PROCESS})\n if(NOT TARGET ${LIB})\n list(APPEND LINK_LIBS ${LIB})\n continue()\n endif()\n\n list(FIND ALL_ALWAYS_LINK_LIBS ${LIB} IS_ALWAYS_LINK)\n if(IS_ALWAYS_LINK EQUAL -1)\n list(APPEND LINK_LIBS ${LIB})\n continue()\n endif()\n\n if(NOT MSVC)\n if(NOT ${CMAKE_SYSTEM_NAME} MATCHES \"Darwin\")\n list(APPEND LINK_LIBS -Wl,--whole-archive ${LIB} -Wl,--no-whole-archive)\n else()\n list(APPEND LINK_LIBS -Wl,-force_load ${LIB})\n endif()\n else()\n # Microsoft Visual C++\n list(APPEND LINK_LIBS /WHOLEARCHIVE:$)\n get_target_property(OTHER_LINK_LIBS ${LIB} INTERFACE_LINK_LIBRARIES)\n if(OTHER_LINK_LIBS)\n foreach(OTHER_LIB ${OTHER_LINK_LIBS})\n list(FIND ALL_LIBS_TO_PROCESS ${OTHER_LIB} FOUND_INDEX)\n if(FOUND_INDEX EQUAL -1)\n list(APPEND LINK_LIBS ${OTHER_LIB})\n endif()\n endforeach()\n endif()\n list(APPEND LIBS_DEPS ${LIB})\n list(\n APPEND LIBS_INCS\n \"$\"\n )\n endif()\n endforeach()\n\n target_link_libraries(${_NAME} ${LINK_LIBS})\n if(LIBS_DEPS)\n add_dependencies(${_NAME} ${LIBS_DEPS})\n target_include_directories(${_NAME} PRIVATE \"${LIBS_INCS}\")\n endif()\nendfunction()\n\n## Add a subdirectory to the build\nfunction(cc_directory)\n add_subdirectory(${ARGN})\nendfunction()\n\n## Add subdirectories to the build\nfunction(cc_directories)\n foreach(SRC_DIR ${ARGN})\n add_subdirectory(${SRC_DIR})\n endforeach()\nendfunction()\n\n## Set the properties of target\nfunction(_cc_target_properties)\n cmake_parse_arguments(\n CC_ARGS \"STRICT;ALWAYS_LINK\" \"NAME;VERSION;C_STANDARD;CXX_STANDARD\"\n \"INCS;PUBINCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS\" ${ARGN}\n )\n\n if(NOT CC_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n get_target_property(TARGET_TYPE ${CC_ARGS_NAME} TYPE)\n if((\"${TARGET_TYPE}\" STREQUAL \"SHARED_LIBRARY\") OR\n (\"${TARGET_TYPE}\" STREQUAL \"STATIC_LIBRARY\") OR\n (\"${TARGET_TYPE}\" STREQUAL \"EXECUTABLE\"))\n set(TARGET_LINKABLE TRUE)\n endif()\n\n if(CC_ARGS_ALWAYS_LINK)\n if((\"${TARGET_TYPE}\" STREQUAL \"STATIC_LIBRARY\") OR\n (\"${TARGET_TYPE}\" STREQUAL \"OBJECT_LIBRARY\"))\n set_property(TARGET ${CC_ARGS_NAME} PROPERTY ALWAYS_LINK TRUE)\n endif()\n endif()\n\n # Set the warning level of compiling\n if(CC_ARGS_STRICT)\n target_compile_options(\n ${CC_ARGS_NAME} PRIVATE \"${BAZEL_CC_STRICT_COMPILE_FLAGS}\"\n )\n if(TARGET_LINKABLE)\n target_link_libraries(${CC_ARGS_NAME} \"${BAZEL_CC_STRICT_LINK_FLAGS}\")\n endif()\n else()\n target_compile_options(\n ${CC_ARGS_NAME} PRIVATE \"${BAZEL_CC_UNSTRICT_COMPILE_FLAGS}\"\n )\n if(TARGET_LINKABLE)\n target_link_libraries(${CC_ARGS_NAME} \"${BAZEL_CC_UNSTRICT_LINK_FLAGS}\")\n endif()\n endif()\n\n if(CC_ARGS_DEFS)\n target_compile_definitions(${CC_ARGS_NAME} PRIVATE \"${CC_ARGS_DEFS}\")\n endif()\n\n if(CC_ARGS_CFLAGS OR CC_ARGS_CXXFLAGS)\n target_compile_options(\n ${CC_ARGS_NAME} PRIVATE\n \"$<$:${CC_ARGS_CFLAGS}>\"\n \"$<$:${CC_ARGS_CXXFLAGS}>\"\n )\n endif()\n\n if(CC_ARGS_LDFLAGS)\n string(REPLACE \";\" \" \" CC_ARGS_LDFLAGS \"${CC_ARGS_LDFLAGS}\")\n set_property(\n TARGET ${CC_ARGS_NAME} PROPERTY LINK_FLAGS \"${CC_ARGS_LDFLAGS}\"\n )\n endif()\n\n if(CC_ARGS_INCS)\n _absolute_paths(INC_DIRS ${CC_ARGS_INCS})\n target_include_directories(${CC_ARGS_NAME} PRIVATE \"${INC_DIRS}\")\n endif()\n\n if(BAZEL_WORKSPACE_DIR)\n target_include_directories(${CC_ARGS_NAME} PRIVATE \"${BAZEL_WORKSPACE_DIR}\")\n endif()\n\n if(CC_ARGS_PUBINCS)\n _absolute_paths(INC_DIRS ${CC_ARGS_PUBINCS})\n target_include_directories(${CC_ARGS_NAME} PUBLIC \"${INC_DIRS}\")\n endif()\n\n if(CC_ARGS_LIBS)\n if(NOT TARGET_LINKABLE)\n _targets_link_dependencies(${CC_ARGS_NAME} ${CC_ARGS_LIBS})\n else()\n if (\"${TARGET_TYPE}\" STREQUAL \"EXECUTABLE\")\n _target_link_libraries(${CC_ARGS_NAME} \"${CC_ARGS_LIBS}\")\n else()\n target_link_libraries(${CC_ARGS_NAME} \"${CC_ARGS_LIBS}\")\n endif()\n endif()\n endif()\n\n if(CC_ARGS_DEPS)\n add_dependencies(${CC_ARGS_NAME} \"${CC_ARGS_DEPS}\")\n endif()\n\n if(CC_ARGS_VERSION)\n set_property(\n TARGET ${CC_ARGS_NAME} PROPERTY VERSION \"${CC_ARGS_VERSION}\"\n )\n endif()\n\n if(NOT CC_C_STANDARD)\n set(CC_C_STANDARD 99)\n endif()\n\n if(NOT CC_CXX_STANDARD)\n set(CC_CXX_STANDARD 11)\n endif()\n\n set_target_properties(\n ${CC_ARGS_NAME} PROPERTIES DEFINE_SYMBOL \"\"\n C_STANDARD ${CC_C_STANDARD} CXX_STANDARD ${CC_CXX_STANDARD}\n C_STANDARD_REQUIRED ON C_EXTENSIONS ON\n CXX_STANDARD_REQUIRED ON CXX_EXTENSIONS OFF\n WINDOWS_EXPORT_ALL_SYMBOLS ON\n )\nendfunction()\n\n## Build a C/C++ static or shared library\nfunction(cc_library)\n cmake_parse_arguments(\n CC_ARGS\n \"STATIC;SHARED;EXCLUDE;PACKED;SRCS_NO_GLOB\"\n \"NAME;VERSION\"\n \"SRCS;INCS;PUBINCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS;PACKED_EXCLUDES\"\n ${ARGN}\n )\n\n if(NOT CC_ARGS_NAME)\n message(FATAL_ERROR \"No target name provided.\")\n endif()\n\n if(CC_ARGS_SRCS_NO_GLOB)\n set(SOURCE_FILES ${CC_ARGS_SRCS})\n if(NOT SOURCE_FILES)\n message(FATAL_ERROR \"No source files provided for ${CC_ARGS_NAME} (SRCS_NO_GLOB mode).\")\n endif()\n else()\n set(SOURCE_FILES \"\")\n foreach(_src IN LISTS CC_ARGS_SRCS)\n if(IS_ABSOLUTE \"${_src}\" OR NOT \"${_src}\" MATCHES \"[*?]\")\n list(APPEND SOURCE_FILES \"${_src}\")\n else()\n file(GLOB _globbed_srcs ${_src})\n list(APPEND SOURCE_FILES ${_globbed_srcs})\n endif()\n endforeach()\n if(NOT SOURCE_FILES)\n message(FATAL_ERROR \"No source files found for ${CC_ARGS_NAME} after globbing.\")\n endif()\n endif()\n\n if(CC_ARGS_VERSION)\n string(REPLACE \"-\" \"_\" MACRO_PREFIX \"${CC_ARGS_NAME}\")\n list(APPEND CC_ARGS_DEFS ${MACRO_PREFIX}_VERSION=\"${CC_ARGS_VERSION}\")\n endif()\n\n if(CC_ARGS_EXCLUDE)\n set(EXCLUDE_OPTION EXCLUDE_FROM_ALL)\n endif()\n\n if(CC_ARGS_SHARED AND CC_ARGS_STATIC)\n _add_library(${CC_ARGS_NAME} \"${EXCLUDE_OPTION}\" ${SOURCE_FILES})\n elseif(CC_ARGS_SHARED)\n add_library(${CC_ARGS_NAME} SHARED ${EXCLUDE_OPTION} ${SOURCE_FILES})\n elseif(CC_ARGS_STATIC)\n add_library(${CC_ARGS_NAME} STATIC ${EXCLUDE_OPTION} ${SOURCE_FILES})\n else()\n add_library(${CC_ARGS_NAME} ${EXCLUDE_OPTION} ${SOURCE_FILES})\n endif()\n\n if(TARGET ${CC_ARGS_NAME}_objects)\n _cc_target_properties(\n NAME \"${CC_ARGS_NAME}_objects\"\n INCS \"${CC_ARGS_INCS};${CC_ARGS_PUBINCS}\"\n DEFS \"${CC_ARGS_DEFS}\"\n LIBS \"${CC_ARGS_LIBS}\"\n CFLAGS \"${CC_ARGS_CFLAGS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n \"${CC_ARGS_UNPARSED_ARGUMENTS}\"\n )\n endif()\n\n if(TARGET ${CC_ARGS_NAME}_static)\n _cc_target_properties(\n NAME \"${CC_ARGS_NAME}_static\"\n INCS \"${CC_ARGS_INCS}\"\n PUBINCS \"${CC_ARGS_PUBINCS}\"\n DEFS \"${CC_ARGS_DEFS}\"\n LIBS \"${CC_ARGS_LIBS}\"\n CFLAGS \"${CC_ARGS_CFLAGS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n \"${CC_ARGS_UNPARSED_ARGUMENTS}\"\n )\n if(CC_ARGS_PACKED)\n install(\n TARGETS ${CC_ARGS_NAME}_static\n ARCHIVE DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n )\n endif()\n endif()\n\n _cc_target_properties(\n NAME \"${CC_ARGS_NAME}\"\n INCS \"${CC_ARGS_INCS}\"\n PUBINCS \"${CC_ARGS_PUBINCS}\"\n DEFS \"${CC_ARGS_DEFS}\"\n LIBS \"${CC_ARGS_LIBS}\"\n CFLAGS \"${CC_ARGS_CFLAGS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n VERSION \"${CC_ARGS_VERSION}\"\n \"${CC_ARGS_UNPARSED_ARGUMENTS}\"\n )\n if(CC_ARGS_PACKED)\n install(\n TARGETS ${CC_ARGS_NAME}\n ARCHIVE DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n LIBRARY DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n )\n if(CC_ARGS_PUBINCS)\n foreach(PACKED_EXCLUDE ${CC_ARGS_PACKED_EXCLUDES})\n list(APPEND PATTERN_EXCLUDES \"PATTERN;${PACKED_EXCLUDE};EXCLUDE\")\n endforeach()\n install(\n DIRECTORY ${CC_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR}\n FILES_MATCHING PATTERN \"*.h\" PATTERN \"*.hpp\" PATTERN \"*.hxx\"\n ${PATTERN_EXCLUDES}\n )\n endif()\n endif()\nendfunction()\n\n## Build a C/C++ executable program\nfunction(cc_binary)\n cmake_parse_arguments(\n CC_ARGS \"PACKED\" \"NAME;VERSION\"\n \"SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS\" ${ARGN}\n )\n\n if(NOT CC_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n file(GLOB CC_ARGS_SRCS ${CC_ARGS_SRCS})\n if(NOT CC_ARGS_SRCS)\n message(FATAL_ERROR \"No source files found of ${CC_ARGS_NAME}.\")\n endif()\n\n if(CC_ARGS_VERSION)\n string(REPLACE \"-\" \"_\" MACRO_PREFIX \"${CC_ARGS_NAME}\")\n list(APPEND CC_ARGS_DEFS ${MACRO_PREFIX}_VERSION=\"${CC_ARGS_VERSION}\")\n endif()\n add_executable(${CC_ARGS_NAME} ${CC_ARGS_SRCS})\n\n if(CC_ARGS_PACKED)\n install(\n TARGETS ${CC_ARGS_NAME} RUNTIME DESTINATION \"${CMAKE_INSTALL_BINDIR}\"\n )\n endif()\n\n _cc_target_properties(\n NAME \"${CC_ARGS_NAME}\"\n INCS \"${CC_ARGS_INCS}\"\n DEFS \"${CC_ARGS_DEFS}\"\n LIBS \"${CC_ARGS_LIBS}\"\n CFLAGS \"${CC_ARGS_CFLAGS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n VERSION \"${CC_ARGS_VERSION}\"\n \"${CC_ARGS_UNPARSED_ARGUMENTS}\"\n )\nendfunction()\n\n## Build a C/C++ executable test program\nfunction(cc_test)\n cmake_parse_arguments(\n CC_ARGS \"\" \"NAME;VERSION\"\n \"SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS;ARGS\" ${ARGN}\n )\n\n if(NOT CC_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n file(GLOB CC_ARGS_SRCS ${CC_ARGS_SRCS})\n if(NOT CC_ARGS_SRCS)\n message(FATAL_ERROR \"No source files found of ${CC_ARGS_NAME}.\")\n endif()\n\n if(CC_ARGS_VERSION)\n string(REPLACE \"-\" \"_\" MACRO_PREFIX \"${CC_ARGS_NAME}\")\n list(APPEND CC_ARGS_DEFS ${MACRO_PREFIX}_VERSION=\"${CC_ARGS_VERSION}\")\n endif()\n add_executable(${CC_ARGS_NAME} EXCLUDE_FROM_ALL ${CC_ARGS_SRCS})\n\n _cc_target_properties(\n NAME \"${CC_ARGS_NAME}\"\n INCS \"${CC_ARGS_INCS}\"\n DEFS \"${CC_ARGS_DEFS}\"\n LIBS \"${CC_ARGS_LIBS}\"\n CFLAGS \"${CC_ARGS_CFLAGS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n \"${CC_ARGS_UNPARSED_ARGUMENTS}\"\n )\n add_dependencies(unittest ${CC_ARGS_NAME})\n add_custom_target(\n unittest.${CC_ARGS_NAME}\n COMMAND $ \"${CC_ARGS_ARGS}\"\n WORKING_DIRECTORY ${PROJECT_BINARY_DIR}\n DEPENDS ${CC_ARGS_NAME}\n )\n add_test(\n NAME ${CC_ARGS_NAME}\n COMMAND $ \"${CC_ARGS_ARGS}\"\n WORKING_DIRECTORY ${PROJECT_BINARY_DIR}\n )\nendfunction()\n\n## Add existing test cases to a test suite\nfunction(cc_test_suite _NAME)\n if(NOT TARGET unittest.${_NAME})\n add_custom_target(unittest.${_NAME} COMMAND \"\")\n endif()\n foreach(TEST_TARGET ${ARGN})\n list(APPEND TEST_TARGETS unittest.${TEST_TARGET})\n endforeach()\n if(TEST_TARGETS)\n add_dependencies(unittest.${_NAME} ${TEST_TARGETS})\n endif()\nendfunction()\n\n## Import a C/C++ static or shared library\nfunction(cc_import)\n cmake_parse_arguments(\n CC_ARGS \"STATIC;SHARED;PACKED\"\n \"NAME;PATH;IMPLIB\" \"INCS;PUBINCS;DEPS;PACKED_EXCLUDES\" ${ARGN}\n )\n\n if(NOT CC_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n file(GLOB CC_ARGS_PATH ${CC_ARGS_PATH})\n if(NOT CC_ARGS_PATH)\n message(FATAL_ERROR \"No imported target file found of ${CC_ARGS_NAME}.\")\n endif()\n if(MSVC AND CC_ARGS_SHARED AND NOT CC_ARGS_IMPLIB)\n string(REGEX REPLACE\n \".[Dd][Ll][Ll]$\" \".lib\" CC_ARGS_IMPLIB ${CC_ARGS_PATH}\n )\n endif()\n\n if(CC_ARGS_SHARED)\n add_library(${CC_ARGS_NAME} SHARED IMPORTED GLOBAL)\n elseif(CC_ARGS_STATIC)\n add_library(${CC_ARGS_NAME} STATIC IMPORTED GLOBAL)\n else()\n add_library(${CC_ARGS_NAME} UNKNOWN IMPORTED GLOBAL)\n endif()\n\n set_property(\n TARGET ${CC_ARGS_NAME} PROPERTY IMPORTED_LOCATION ${CC_ARGS_PATH}\n )\n if(MSVC AND CC_ARGS_SHARED)\n set_property(\n TARGET ${CC_ARGS_NAME} PROPERTY IMPORTED_IMPLIB ${CC_ARGS_IMPLIB}\n )\n endif()\n\n if(CC_ARGS_INCS)\n _absolute_paths(INC_DIRS ${CC_ARGS_INCS})\n foreach(INC_DIR ${INC_DIRS})\n set_property(\n TARGET ${CC_ARGS_NAME} APPEND PROPERTY\n INTERFACE_INCLUDE_DIRECTORIES \"${INC_DIR}\"\n )\n endforeach()\n endif()\n\n if(CC_ARGS_PUBINCS)\n _absolute_paths(INC_DIRS ${CC_ARGS_PUBINCS})\n foreach(INC_DIR ${INC_DIRS})\n set_property(\n TARGET ${CC_ARGS_NAME} APPEND PROPERTY\n INTERFACE_INCLUDE_DIRECTORIES \"${INC_DIR}\"\n )\n endforeach()\n endif()\n\n if(CC_ARGS_DEPS)\n add_dependencies(${CC_ARGS_NAME} \"${CC_ARGS_DEPS}\")\n endif()\n\n if(CC_ARGS_PACKED)\n install(\n TARGETS ${CC_ARGS_NAME}\n ARCHIVE DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n LIBRARY DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n )\n if(CC_ARGS_PUBINCS)\n foreach(PACKED_EXCLUDE ${CC_ARGS_PACKED_EXCLUDES})\n list(APPEND PATTERN_EXCLUDES \"PATTERN;${PACKED_EXCLUDE};EXCLUDE\")\n endforeach()\n install(\n DIRECTORY ${CC_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR}\n FILES_MATCHING PATTERN \"*.h\" PATTERN \"*.hpp\" PATTERN \"*.hxx\"\n ${PATTERN_EXCLUDES}\n )\n endif()\n endif()\nendfunction()\n\n## Import a C/C++ interface library\nfunction(cc_interface)\n cmake_parse_arguments(\n CC_ARGS \"PACKED\" \"NAME\" \"INCS;PUBINCS;DEPS;PACKED_EXCLUDES\" ${ARGN}\n )\n\n if(NOT CC_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n add_library(${CC_ARGS_NAME} INTERFACE GLOBAL)\n if(CC_ARGS_INCS)\n _absolute_paths(INC_DIRS ${CC_ARGS_INCS})\n target_include_directories(${CC_ARGS_NAME} INTERFACE \"${INC_DIRS}\")\n endif()\n\n if(CC_ARGS_PUBINCS)\n _absolute_paths(INC_DIRS ${CC_ARGS_PUBINCS})\n target_include_directories(${CC_ARGS_NAME} INTERFACE \"${INC_DIRS}\")\n endif()\n\n if(CC_ARGS_DEPS)\n add_dependencies(${CC_ARGS_NAME} \"${CC_ARGS_DEPS}\")\n endif()\n\n if(CC_ARGS_PACKED AND CC_ARGS_PUBINCS)\n foreach(PACKED_EXCLUDE ${CC_ARGS_PACKED_EXCLUDES})\n list(APPEND PATTERN_EXCLUDES \"PATTERN;${PACKED_EXCLUDE};EXCLUDE\")\n endforeach()\n install(\n DIRECTORY ${CC_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR}\n FILES_MATCHING PATTERN \"*.h\" PATTERN \"*.hpp\" PATTERN \"*.hxx\"\n ${PATTERN_EXCLUDES}\n )\n endif()\nendfunction()\n\n## Find gtest library\nfunction(_find_gtest)\n if(DEFINED FIND_GTEST_LIBS AND DEFINED FIND_GTEST_INCS)\n return()\n endif()\n\n if(NOT TARGET gtest OR NOT TARGET gtest_main)\n # Find gtest using 'find_package'\n find_package(GTest REQUIRED)\n set(\n FIND_GTEST_INCS \"${GTEST_INCLUDE_DIRS}\"\n CACHE STRING \"GTest includes\"\n )\n set(\n FIND_GTEST_LIBS \"${GTEST_BOTH_LIBRARIES}\"\n CACHE STRING \"GTest libraries\"\n )\n else()\n # Find gtest using target names\n set(FIND_GTEST_INCS \"\" CACHE STRING \"GTest includes\")\n set(FIND_GTEST_LIBS \"gtest;gtest_main\" CACHE STRING \"GTest libraries\")\n endif()\nendfunction()\n\n## Build a C/C++ executable google test program\nfunction(cc_gtest)\n cmake_parse_arguments(\n CC_ARGS \"\" \"NAME;VERSION\"\n \"SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS;ARGS\" ${ARGN}\n )\n _find_gtest()\n cc_test(\n NAME \"${CC_ARGS_NAME}\"\n VERSION \"${CC_ARGS_VERSION}\"\n SRCS \"${CC_ARGS_SRCS}\"\n INCS \"${CC_ARGS_INCS};${FIND_GTEST_INCS}\"\n DEFS \"${CC_ARGS_DEFS}\"\n LIBS \"${CC_ARGS_LIBS};${FIND_GTEST_LIBS}\"\n CFLAGS \"${CC_ARGS_CFLAGS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n ARGS \"${CC_ARGS_ARGS}\"\n )\nendfunction()\n\n## Find gmock library\nfunction(_find_gmock)\n if(DEFINED FIND_GMOCK_LIBS AND DEFINED FIND_GMOCK_INCS)\n return()\n endif()\n\n if(NOT TARGET gmock OR NOT TARGET gmock_main)\n # Find gmock/gtest using 'find_package'\n find_package(GMock REQUIRED)\n find_package(GTest REQUIRED)\n set(\n FIND_GMOCK_INCS \"${GMOCK_INCLUDE_DIRS};${GTEST_INCLUDE_DIRS}\"\n CACHE STRING \"GMock includes\"\n )\n set(\n FIND_GMOCK_LIBS \"${GMOCK_BOTH_LIBRARIES};${GTEST_LIBRARIES}\"\n CACHE STRING \"GMock libraries\"\n )\n else()\n # Find gmock using target names\n set(FIND_GMOCK_INCS \"\" CACHE STRING \"GMock includes\")\n set(FIND_GMOCK_LIBS \"gmock;gmock_main\" CACHE STRING \"GMock libraries\")\n endif()\nendfunction()\n\n## Build a C/C++ executable google mock program\nfunction(cc_gmock)\n cmake_parse_arguments(\n CC_ARGS \"\" \"NAME;VERSION\"\n \"SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;LDFLAGS;DEPS;ARGS\" ${ARGN}\n )\n _find_gmock()\n cc_test(\n NAME \"${CC_ARGS_NAME}\"\n VERSION \"${CC_ARGS_VERSION}\"\n SRCS \"${CC_ARGS_SRCS}\"\n INCS \"${CC_ARGS_INCS};${FIND_GMOCK_INCS}\"\n DEFS \"${CC_ARGS_DEFS}\"\n LIBS \"${CC_ARGS_LIBS};${FIND_GMOCK_LIBS}\"\n CFLAGS \"${CC_ARGS_CFLAGS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n ARGS \"${CC_ARGS_ARGS}\"\n )\nendfunction()\n\n## Find protobuf library\nfunction(_find_protobuf _VERSION)\n if(DEFINED CC_PROTOBUF_PROTOC_${_VERSION})\n return()\n endif()\n\n # Find protobuf using 'find_package'\n if(NOT TARGET protoc OR NOT TARGET libprotobuf)\n find_package(Protobuf ${_VERSION} REQUIRED)\n set(\n CC_PROTOBUF_PROTOC_${_VERSION}\n \"${PROTOBUF_PROTOC_EXECUTABLE}\" CACHE PATH \"Protobuf compiler\"\n )\n set(\n CC_PROTOBUF_INCS_${_VERSION}\n \"${PROTOBUF_INCLUDE_DIRS}\" CACHE STRING \"Protobuf includes\"\n )\n set(\n CC_PROTOBUF_LIBS_${_VERSION}\n \"${PROTOBUF_LIBRARIES}\" CACHE STRING \"Protobuf libraries\"\n )\n return()\n endif()\n\n # Find protobuf using target names\n get_target_property(protoc_VERSION protoc VERSION)\n get_target_property(libprotobuf_VERSION libprotobuf VERSION)\n if(_VERSION)\n if(${protoc_VERSION} VERSION_LESS ${_VERSION})\n message(\n FATAL_ERROR\n \"The 'protoc' version is ${protoc_VERSION}, less than ${_VERSION}.\"\n )\n endif()\n if(${libprotobuf_VERSION} VERSION_LESS ${_VERSION})\n message(\n FATAL_ERROR\n \"The 'libprotobuf' version is ${libprotobuf_VERSION}, \"\n \"less than ${_VERSION}.\"\n )\n endif()\n endif()\n\n message(STATUS \"Found binary 'protoc ${protoc_VERSION}'\")\n message(STATUS \"Found library 'libprotobuf ${libprotobuf_VERSION}'\")\n set(\n CC_PROTOBUF_PROTOC_${_VERSION}\n \"$\" CACHE PATH \"Protobuf compiler\"\n )\n get_target_property(protoc_SOURCE_DIR protoc SOURCE_DIR)\n get_filename_component(protoc_INCLUDE_DIR ${protoc_SOURCE_DIR}/../src ABSOLUTE)\n set(\n CC_PROTOBUF_INCS_${_VERSION}\n \"${protoc_INCLUDE_DIR}\" CACHE STRING \"Protobuf includes\"\n )\n set(\n CC_PROTOBUF_LIBS_${_VERSION} libprotobuf CACHE STRING \"Protobuf libraries\"\n )\nendfunction()\n\n## Build a C++ protobuf static or shared library\nfunction(cc_proto_library)\n cmake_parse_arguments(\n CC_ARGS \"STATIC;SHARED;EXCLUDE;PACKED\"\n \"NAME;VERSION;PROTOROOT;PROTOBUF_VERSION\"\n \"SRCS;CXXFLAGS;LDFLAGS;DEPS\" ${ARGN}\n )\n\n _find_protobuf(\"${CC_ARGS_PROTOBUF_VERSION}\")\n set(CC_PROTOBUF_PROTOC ${CC_PROTOBUF_PROTOC_${CC_ARGS_PROTOBUF_VERSION}})\n if(DEFINED GLOBAL_CC_PROTOBUF_PROTOC)\n set(CC_PROTOBUF_PROTOC ${GLOBAL_CC_PROTOBUF_PROTOC})\n endif()\n set(CC_PROTOBUF_INCS ${CC_PROTOBUF_INCS_${CC_ARGS_PROTOBUF_VERSION}})\n set(CC_PROTOBUF_LIBS ${CC_PROTOBUF_LIBS_${CC_ARGS_PROTOBUF_VERSION}})\n\n if(NOT CC_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n file(GLOB CC_ARGS_SRCS ${CC_ARGS_SRCS})\n if(NOT CC_ARGS_SRCS)\n message(FATAL_ERROR \"No source files found of ${CC_ARGS_NAME}.\")\n endif()\n\n if(CC_ARGS_VERSION)\n string(REPLACE \"-\" \"_\" MACRO_PREFIX \"${CC_ARGS_NAME}\")\n list(APPEND CC_ARGS_DEFS ${MACRO_PREFIX}_VERSION=\"${CC_ARGS_VERSION}\")\n endif()\n\n if(CC_ARGS_EXCLUDE)\n set(EXCLUDE_OPTION EXCLUDE_FROM_ALL)\n endif()\n\n set(PROTO_ROOT ${CMAKE_CURRENT_SOURCE_DIR})\n if(CC_ARGS_PROTOROOT)\n get_filename_component(PROTO_ROOT ${CC_ARGS_PROTOROOT} ABSOLUTE)\n endif()\n\n # Compile proto files to C++ sources\n set(CPP_OUTPATH \"${CMAKE_CURRENT_BINARY_DIR}\")\n foreach(PROTO_FILE ${CC_ARGS_SRCS})\n get_filename_component(PROTO_FILE ${PROTO_FILE} ABSOLUTE)\n\n if(NOT ${PROTO_FILE} MATCHES \"\\\\.proto$$\")\n message(FATAL_ERROR \"Unrecognized proto file ${PROTOFILE}\")\n endif()\n if(NOT ${PROTO_FILE} MATCHES \"^${PROTO_ROOT}\")\n message(FATAL_ERROR \"'${PROTO_FILE}' NOT IN '${PROTO_ROOT}'\")\n endif()\n\n string(\n REGEX REPLACE \"^${PROTO_ROOT}(/?)\" \"\" ROOT_CLEANED_FILE ${PROTO_FILE}\n )\n string(REGEX REPLACE \"\\\\.proto$$\" \"\" EXT_CLEANED_FILE ${ROOT_CLEANED_FILE})\n set(CPP_FILE \"${CPP_OUTPATH}/${EXT_CLEANED_FILE}.pb.cc\")\n set(HDR_FILE \"${CPP_OUTPATH}/${EXT_CLEANED_FILE}.pb.h\")\n set(INJ_FILE \"${CPP_OUTPATH}/${EXT_CLEANED_FILE}.pb.cmake\")\n file(RELATIVE_PATH REL_CPP_FILE ${CMAKE_BINARY_DIR} ${CPP_FILE})\n\n set(INJECTED_SCRIPT\n \"foreach(SRC ${EXT_CLEANED_FILE}.pb.cc ${EXT_CLEANED_FILE}.pb.h)\\n\"\n \" file(READ \\$\\{SRC\\} SRC_CODE)\\n\"\n \" file(REMOVE \\$\\{SRC\\})\\n\"\n \" file(APPEND \\$\\{SRC\\} \\\"#ifdef __GNUC__\\\\n\\\")\\n\"\n \" file(APPEND \\$\\{SRC\\} \\\"#pragma GCC diagnostic push\\\\n\\\")\\n\"\n \" file(APPEND \\$\\{SRC\\} \\\"#pragma GCC diagnostic ignored \\\\\\\"-Wshadow\\\\\\\"\\\\n\\\")\\n\"\n \" file(APPEND \\$\\{SRC\\} \\\"#pragma GCC diagnostic ignored \\\\\\\"-Wunused-parameter\\\\\\\"\\\\n\\\")\\n\"\n \" file(APPEND \\$\\{SRC\\} \\\"#endif\\\\n\\\\n\\\")\\n\"\n \" file(APPEND \\$\\{SRC\\} \\\"\\$\\{SRC_CODE\\}\\\")\\n\"\n \" file(APPEND \\$\\{SRC\\} \\\"\\\\n#ifdef __GNUC__\\\\n\\\")\\n\"\n \" file(APPEND \\$\\{SRC\\} \\\"#pragma GCC diagnostic pop\\\\n\\\")\\n\"\n \" file(APPEND \\$\\{SRC\\} \\\"#endif\\\\n\\\")\\n\"\n \"endforeach()\\n\"\n )\n file(WRITE \"${INJ_FILE}\" ${INJECTED_SCRIPT})\n\n add_custom_command(\n OUTPUT \"${CPP_FILE}\" \"${HDR_FILE}\"\n # COMMAND ${CMAKE_COMMAND} -E make_directory ${CPP_OUTPATH}\n COMMAND ${CC_PROTOBUF_PROTOC}\n --cpp_out \"${CPP_OUTPATH}\" --python_out \"${CPP_OUTPATH}\"\n --proto_path \"${PROTO_ROOT}\" --proto_path \"${CC_PROTOBUF_INCS}\" \"${PROTO_FILE}\"\n\n COMMAND ${CMAKE_COMMAND} -P \"${INJ_FILE}\"\n DEPENDS \"${PROTO_FILE}\"\n COMMENT \"Generating CXX source ${REL_CPP_FILE}\"\n VERBATIM\n )\n list(APPEND CC_SRCS \"${CPP_FILE}\" \"${HDR_FILE}\")\n endforeach()\n\n # Compile C++ sources\n if(CC_ARGS_SHARED AND CC_ARGS_STATIC)\n _add_library(${CC_ARGS_NAME} \"${EXCLUDE_OPTION}\" \"${CC_SRCS}\")\n elseif(CC_ARGS_SHARED)\n add_library(${CC_ARGS_NAME} SHARED ${EXCLUDE_OPTION} ${CC_SRCS})\n elseif(CC_ARGS_STATIC)\n add_library(${CC_ARGS_NAME} STATIC ${EXCLUDE_OPTION} ${CC_SRCS})\n else()\n add_library(${CC_ARGS_NAME} ${EXCLUDE_OPTION} ${CC_SRCS})\n endif()\n\n if(TARGET ${CC_ARGS_NAME}_objects)\n _cc_target_properties(\n NAME \"${CC_ARGS_NAME}_objects\"\n INCS \"${CPP_OUTPATH};${CC_PROTOBUF_INCS}\"\n LIBS \"${CC_PROTOBUF_LIBS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n \"${CC_ARGS_UNPARSED_ARGUMENTS}\"\n )\n endif()\n\n if(TARGET ${CC_ARGS_NAME}_static)\n _cc_target_properties(\n NAME \"${CC_ARGS_NAME}_static\"\n PUBINCS \"${CPP_OUTPATH};${CC_PROTOBUF_INCS}\"\n LIBS \"${CC_PROTOBUF_LIBS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n \"${CC_ARGS_UNPARSED_ARGUMENTS}\"\n )\n if(CC_ARGS_PACKED)\n install(\n TARGETS ${CC_ARGS_NAME}_static\n ARCHIVE DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n )\n endif()\n endif()\n\n _cc_target_properties(\n NAME \"${CC_ARGS_NAME}\"\n PUBINCS \"${CPP_OUTPATH};${CC_PROTOBUF_INCS}\"\n LIBS \"${CC_PROTOBUF_LIBS}\"\n CXXFLAGS \"${CC_ARGS_CXXFLAGS}\"\n LDFLAGS \"${CC_ARGS_LDFLAGS}\"\n DEPS \"${CC_ARGS_DEPS}\"\n VERSION \"${CC_ARGS_VERSION}\"\n \"${CC_ARGS_UNPARSED_ARGUMENTS}\"\n )\n if(CC_ARGS_PACKED)\n install(\n TARGETS ${CC_ARGS_NAME}\n ARCHIVE DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n LIBRARY DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n )\n endif()\nendfunction()\n\n## Add a subdirectory to the build\nfunction(cuda_directory)\n if(NOT CMAKE_CUDA_COMPILER)\n message(FATAL_ERROR \"No CUDA language supported.\")\n endif()\n cc_directory(${ARGN})\nendfunction()\n\n## Add subdirectories to the build\nfunction(cuda_directories)\n if(NOT CMAKE_CUDA_COMPILER)\n message(FATAL_ERROR \"No CUDA language supported.\")\n endif()\n cc_directories(${ARGN})\nendfunction()\n\n## Set the properties of cuda target\nfunction(_cuda_target_properties)\n cmake_parse_arguments(\n CUDA_ARGS \"STRICT;ALWAYS_LINK\" \"NAME;VERSION;C_STANDARD;CXX_STANDARD\"\n \"INCS;PUBINCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS\" ${ARGN}\n )\n\n if(NOT CUDA_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n get_target_property(TARGET_TYPE ${CUDA_ARGS_NAME} TYPE)\n if((\"${TARGET_TYPE}\" STREQUAL \"SHARED_LIBRARY\") OR\n (\"${TARGET_TYPE}\" STREQUAL \"STATIC_LIBRARY\") OR\n (\"${TARGET_TYPE}\" STREQUAL \"EXECUTABLE\"))\n set(TARGET_LINKABLE TRUE)\n endif()\n\n if(CUDA_ARGS_ALWAYS_LINK)\n if((\"${TARGET_TYPE}\" STREQUAL \"STATIC_LIBRARY\") OR\n (\"${TARGET_TYPE}\" STREQUAL \"OBJECT_LIBRARY\"))\n set_property(TARGET ${CUDA_ARGS_NAME} PROPERTY ALWAYS_LINK TRUE)\n endif()\n endif()\n\n # Set the warning level of compiling\n if(CUDA_ARGS_STRICT)\n target_compile_options(\n ${CUDA_ARGS_NAME} PRIVATE \"${BAZEL_CUDA_STRICT_COMPILE_FLAGS}\"\n )\n if(TARGET_LINKABLE)\n target_link_libraries(\n ${CUDA_ARGS_NAME} \"${BAZEL_CUDA_STRICT_LINK_FLAGS}\"\n )\n endif()\n else()\n target_compile_options(\n ${CUDA_ARGS_NAME} PRIVATE \"${BAZEL_CUDA_UNSTRICT_COMPILE_FLAGS}\"\n )\n if(TARGET_LINKABLE)\n target_link_libraries(\n ${CUDA_ARGS_NAME} \"${BAZEL_CUDA_UNSTRICT_LINK_FLAGS}\"\n )\n endif()\n endif()\n\n target_compile_options(\n ${CUDA_ARGS_NAME} PRIVATE\n \"$<$:-ccbin=${CMAKE_CXX_COMPILER}>\"\n )\n\n if(CUDA_ARGS_DEFS)\n target_compile_definitions(${CUDA_ARGS_NAME} PRIVATE \"${CUDA_ARGS_DEFS}\")\n endif()\n\n if(CUDA_ARGS_CFLAGS OR CUDA_ARGS_CXXFLAGS OR CUDA_ARGS_CUDAFLAGS)\n target_compile_options(\n ${CUDA_ARGS_NAME} PRIVATE\n \"$<$:${CUDA_ARGS_CFLAGS}>\"\n \"$<$:${CUDA_ARGS_CXXFLAGS}>\"\n \"$<$:${CUDA_ARGS_CUDAFLAGS}>\"\n )\n endif()\n\n if(CUDA_ARGS_LDFLAGS)\n string(REPLACE \";\" \" \" CUDA_ARGS_LDFLAGS \"${CUDA_ARGS_LDFLAGS}\")\n set_property(\n TARGET ${CUDA_ARGS_NAME} PROPERTY LINK_FLAGS \"${CUDA_ARGS_LDFLAGS}\"\n )\n endif()\n\n if(CUDA_ARGS_INCS)\n _absolute_paths(INC_DIRS ${CUDA_ARGS_INCS})\n target_include_directories(${CUDA_ARGS_NAME} PRIVATE \"${INC_DIRS}\")\n endif()\n\n target_include_directories(\n ${CUDA_ARGS_NAME} PRIVATE \"${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}\"\n )\n\n if(BAZEL_WORKSPACE_DIR)\n target_include_directories(\n ${CUDA_ARGS_NAME} PRIVATE \"${BAZEL_WORKSPACE_DIR}\"\n )\n endif()\n\n if(CUDA_ARGS_PUBINCS)\n _absolute_paths(INC_DIRS ${CUDA_ARGS_PUBINCS})\n target_include_directories(${CUDA_ARGS_NAME} PUBLIC \"${INC_DIRS}\")\n endif()\n\n if(CUDA_ARGS_LIBS)\n if(NOT TARGET_LINKABLE)\n _targets_link_dependencies(${CUDA_ARGS_NAME} ${CUDA_ARGS_LIBS})\n else()\n if (\"${TARGET_TYPE}\" STREQUAL \"EXECUTABLE\")\n _target_link_libraries(${CUDA_ARGS_NAME} \"${CUDA_ARGS_LIBS}\")\n else()\n target_link_libraries(${CUDA_ARGS_NAME} \"${CUDA_ARGS_LIBS}\")\n endif()\n endif()\n endif()\n\n if(CUDA_ARGS_DEPS)\n add_dependencies(${CUDA_ARGS_NAME} \"${CUDA_ARGS_DEPS}\")\n endif()\n\n if(CUDA_ARGS_VERSION)\n set_property(\n TARGET ${CUDA_ARGS_NAME} PROPERTY VERSION \"${CUDA_ARGS_VERSION}\"\n )\n endif()\n\n if(NOT CUDA_C_STANDARD)\n set(CUDA_C_STANDARD 99)\n endif()\n\n if(NOT CUDA_CXX_STANDARD)\n set(CUDA_CXX_STANDARD 11)\n endif()\n\n set_target_properties(\n ${CUDA_ARGS_NAME} PROPERTIES DEFINE_SYMBOL \"\"\n C_STANDARD ${CUDA_C_STANDARD} CXX_STANDARD ${CUDA_CXX_STANDARD}\n C_STANDARD_REQUIRED ON C_EXTENSIONS ON\n CXX_STANDARD_REQUIRED ON CXX_EXTENSIONS OFF\n CUDA_STANDARD 11 CUDA_STANDARD_REQUIRED ON CUDA_EXTENSIONS OFF\n WINDOWS_EXPORT_ALL_SYMBOLS ON\n )\nendfunction()\n\n## Build a CUDA static or shared library\nfunction(cuda_library)\n if(NOT CMAKE_CUDA_COMPILER)\n message(FATAL_ERROR \"No CUDA language supported.\")\n endif()\n\n cmake_parse_arguments(\n CUDA_ARGS \"STATIC;SHARED;EXCLUDE;PACKED\" \"NAME;VERSION\"\n \"SRCS;INCS;PUBINCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS;PACKED_EXCS\"\n ${ARGN}\n )\n\n if(NOT CUDA_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n file(GLOB CUDA_ARGS_SRCS ${CUDA_ARGS_SRCS})\n if(NOT CUDA_ARGS_SRCS)\n message(FATAL_ERROR \"No source files found of ${CUDA_ARGS_NAME}.\")\n endif()\n\n if(CUDA_ARGS_VERSION)\n string(REPLACE \"-\" \"_\" MACRO_PREFIX \"${CUDA_ARGS_NAME}\")\n list(APPEND CUDA_ARGS_DEFS ${MACRO_PREFIX}_VERSION=\"${CUDA_ARGS_VERSION}\")\n endif()\n\n if(CUDA_ARGS_EXCLUDE)\n set(EXCLUDE_OPTION EXCLUDE_FROM_ALL)\n endif()\n\n if(CUDA_ARGS_SHARED AND CUDA_ARGS_STATIC)\n _add_library(${CUDA_ARGS_NAME} \"${EXCLUDE_OPTION}\" \"${CUDA_ARGS_SRCS}\")\n elseif(CUDA_ARGS_SHARED)\n add_library(${CUDA_ARGS_NAME} SHARED ${EXCLUDE_OPTION} ${CUDA_ARGS_SRCS})\n elseif(CUDA_ARGS_STATIC)\n add_library(${CUDA_ARGS_NAME} STATIC ${EXCLUDE_OPTION} ${CUDA_ARGS_SRCS})\n else()\n add_library(${CUDA_ARGS_NAME} ${EXCLUDE_OPTION} ${CUDA_ARGS_SRCS})\n endif()\n\n if(TARGET ${CUDA_ARGS_NAME}_objects)\n _cuda_target_properties(\n NAME \"${CUDA_ARGS_NAME}_objects\"\n INCS \"${CUDA_ARGS_INCS};${CUDA_ARGS_PUBINCS}\"\n DEFS \"${CUDA_ARGS_DEFS}\"\n LIBS \"${CUDA_ARGS_LIBS}\"\n CFLAGS \"${CUDA_ARGS_CFLAGS}\"\n CXXFLAGS \"${CUDA_ARGS_CXXFLAGS}\"\n CUDAFLAGS \"${CUDA_ARGS_CUDAFLAGS}\"\n LDFLAGS \"${CUDA_ARGS_LDFLAGS}\"\n DEPS \"${CUDA_ARGS_DEPS}\"\n \"${CUDA_ARGS_UNPARSED_ARGUMENTS}\"\n )\n endif()\n\n if(TARGET ${CUDA_ARGS_NAME}_static)\n _cuda_target_properties(\n NAME \"${CUDA_ARGS_NAME}_static\"\n INCS \"${CUDA_ARGS_INCS}\"\n PUBINCS \"${CUDA_ARGS_PUBINCS}\"\n DEFS \"${CUDA_ARGS_DEFS}\"\n LIBS \"${CUDA_ARGS_LIBS}\"\n CFLAGS \"${CUDA_ARGS_CFLAGS}\"\n CXXFLAGS \"${CUDA_ARGS_CXXFLAGS}\"\n CUDAFLAGS \"${CUDA_ARGS_CUDAFLAGS}\"\n LDFLAGS \"${CUDA_ARGS_LDFLAGS}\"\n DEPS \"${CUDA_ARGS_DEPS}\"\n \"${CUDA_ARGS_UNPARSED_ARGUMENTS}\"\n )\n if(CUDA_ARGS_PACKED)\n install(\n TARGETS ${CUDA_ARGS_NAME}_static\n ARCHIVE DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n )\n endif()\n endif()\n\n _cuda_target_properties(\n NAME \"${CUDA_ARGS_NAME}\"\n INCS \"${CUDA_ARGS_INCS}\"\n PUBINCS \"${CUDA_ARGS_PUBINCS}\"\n DEFS \"${CUDA_ARGS_DEFS}\"\n LIBS \"${CUDA_ARGS_LIBS}\"\n CFLAGS \"${CUDA_ARGS_CFLAGS}\"\n CXXFLAGS \"${CUDA_ARGS_CXXFLAGS}\"\n CUDAFLAGS \"${CUDA_ARGS_CUDAFLAGS}\"\n LDFLAGS \"${CUDA_ARGS_LDFLAGS}\"\n DEPS \"${CUDA_ARGS_DEPS}\"\n VERSION \"${CUDA_ARGS_VERSION}\"\n \"${CUDA_ARGS_UNPARSED_ARGUMENTS}\"\n )\n if(CUDA_ARGS_PACKED)\n install(\n TARGETS ${CUDA_ARGS_NAME}\n ARCHIVE DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n LIBRARY DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n )\n if(CUDA_ARGS_PUBINCS)\n foreach(PACKED_EXCLUDE ${CUDA_ARGS_PACKED_IGORNES})\n list(APPEND PATTERN_EXCLUDES \"PATTERN;${PACKED_EXCLUDE};EXCLUDE\")\n endforeach()\n install(\n DIRECTORY ${CUDA_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR}\n FILES_MATCHING PATTERN \"*.h\" PATTERN \"*.hpp\"\n PATTERN \"*.hxx\" PATTERN \"*.cuh\"\n ${PATTERN_EXCLUDES}\n )\n endif()\n endif()\nendfunction()\n\n## Build a CUDA executable program\nfunction(cuda_binary)\n if(NOT CMAKE_CUDA_COMPILER)\n message(FATAL_ERROR \"No CUDA language supported.\")\n endif()\n\n cmake_parse_arguments(\n CUDA_ARGS \"PACKED\" \"NAME;VERSION\"\n \"SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS\" ${ARGN}\n )\n\n if(NOT CUDA_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n file(GLOB CUDA_ARGS_SRCS ${CUDA_ARGS_SRCS})\n if(NOT CUDA_ARGS_SRCS)\n message(FATAL_ERROR \"No source files found of ${CUDA_ARGS_NAME}.\")\n endif()\n\n if(CUDA_ARGS_VERSION)\n string(REPLACE \"-\" \"_\" MACRO_PREFIX \"${CUDA_ARGS_NAME}\")\n list(APPEND CUDA_ARGS_DEFS ${MACRO_PREFIX}_VERSION=\"${CUDA_ARGS_VERSION}\")\n endif()\n add_executable(${CUDA_ARGS_NAME} ${CUDA_ARGS_SRCS})\n\n if(CUDA_ARGS_PACKED)\n install(\n TARGETS ${CUDA_ARGS_NAME} RUNTIME DESTINATION \"${CMAKE_INSTALL_BINDIR}\"\n )\n endif()\n\n _cuda_target_properties(\n NAME \"${CUDA_ARGS_NAME}\"\n INCS \"${CUDA_ARGS_INCS}\"\n DEFS \"${CUDA_ARGS_DEFS}\"\n LIBS \"${CUDA_ARGS_LIBS}\"\n CFLAGS \"${CUDA_ARGS_CFLAGS}\"\n CXXFLAGS \"${CUDA_ARGS_CXXFLAGS}\"\n CUDAFLAGS \"${CUDA_ARGS_CUDAFLAGS}\"\n LDFLAGS \"${CUDA_ARGS_LDFLAGS}\"\n DEPS \"${CUDA_ARGS_DEPS}\"\n VERSION \"${CUDA_ARGS_VERSION}\"\n \"${CUDA_ARGS_UNPARSED_ARGUMENTS}\"\n )\nendfunction()\n\n## Build a CUDA executable test program\nfunction(cuda_test)\n if(NOT CMAKE_CUDA_COMPILER)\n message(FATAL_ERROR \"No CUDA language supported.\")\n endif()\n\n cmake_parse_arguments(\n CUDA_ARGS \"\" \"NAME;VERSION\"\n \"SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS;ARGS\" ${ARGN}\n )\n\n if(NOT CUDA_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n file(GLOB CUDA_ARGS_SRCS ${CUDA_ARGS_SRCS})\n if(NOT CUDA_ARGS_SRCS)\n message(FATAL_ERROR \"No source files found of ${CUDA_ARGS_NAME}.\")\n endif()\n\n if(CUDA_ARGS_VERSION)\n string(REPLACE \"-\" \"_\" MACRO_PREFIX \"${CUDA_ARGS_NAME}\")\n list(APPEND CUDA_ARGS_DEFS ${MACRO_PREFIX}_VERSION=\"${CUDA_ARGS_VERSION}\")\n endif()\n add_executable(${CUDA_ARGS_NAME} EXCLUDE_FROM_ALL ${CUDA_ARGS_SRCS})\n\n _cuda_target_properties(\n NAME \"${CUDA_ARGS_NAME}\"\n INCS \"${CUDA_ARGS_INCS}\"\n DEFS \"${CUDA_ARGS_DEFS}\"\n LIBS \"${CUDA_ARGS_LIBS}\"\n CFLAGS \"${CUDA_ARGS_CFLAGS}\"\n CXXFLAGS \"${CUDA_ARGS_CXXFLAGS}\"\n CUDAFLAGS \"${CUDA_ARGS_CUDAFLAGS}\"\n LDFLAGS \"${CUDA_ARGS_LDFLAGS}\"\n DEPS \"${CUDA_ARGS_DEPS}\"\n \"${CUDA_ARGS_UNPARSED_ARGUMENTS}\"\n )\n add_dependencies(unittest ${CUDA_ARGS_NAME})\n add_custom_target(\n unittest.${CUDA_ARGS_NAME}\n COMMAND $ \"${CUDA_ARGS_ARGS}\"\n WORKING_DIRECTORY ${PROJECT_BINARY_DIR}\n DEPENDS ${CUDA_ARGS_NAME}\n )\n add_test(\n NAME ${CUDA_ARGS_NAME}\n COMMAND $ \"${CUDA_ARGS_ARGS}\"\n WORKING_DIRECTORY ${PROJECT_BINARY_DIR}\n )\nendfunction()\n\n## Add existing test cases to a test suite\nfunction(cuda_test_suite)\n if(NOT CMAKE_CUDA_COMPILER)\n message(FATAL_ERROR \"No CUDA language supported.\")\n endif()\n cc_test_suite(${ARGN})\nendfunction()\n\n## Import a C/C++/CUDA static or shared library\nfunction(cuda_import)\n if(NOT CMAKE_CUDA_COMPILER)\n message(FATAL_ERROR \"No CUDA language supported.\")\n endif()\n\n cmake_parse_arguments(\n CUDA_ARGS \"STATIC;SHARED;PACKED\"\n \"NAME;PATH;IMPLIB\" \"INCS;PUBINCS;DEPS;PACKED_EXCLUDES\" ${ARGN}\n )\n\n if(NOT CUDA_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n file(GLOB CUDA_ARGS_PATH ${CUDA_ARGS_PATH})\n if(NOT CUDA_ARGS_PATH)\n message(FATAL_ERROR \"No imported target file found of ${CUDA_ARGS_NAME}.\")\n endif()\n if(MSVC AND CUDA_ARGS_SHARED AND NOT CUDA_ARGS_IMPLIB)\n string(REGEX REPLACE\n \".[Dd][Ll][Ll]$\" \".lib\" CUDA_ARGS_IMPLIB ${CUDA_ARGS_PATH}\n )\n endif()\n\n if(CUDA_ARGS_SHARED)\n add_library(${CUDA_ARGS_NAME} SHARED IMPORTED GLOBAL)\n elseif(CUDA_ARGS_STATIC)\n add_library(${CUDA_ARGS_NAME} STATIC IMPORTED GLOBAL)\n else()\n add_library(${CUDA_ARGS_NAME} UNKNOWN IMPORTED GLOBAL)\n endif()\n\n set_property(\n TARGET ${CUDA_ARGS_NAME} PROPERTY IMPORTED_LOCATION ${CUDA_ARGS_PATH}\n )\n if(MSVC AND CUDA_ARGS_SHARED)\n set_property(\n TARGET ${CUDA_ARGS_NAME} PROPERTY IMPORTED_IMPLIB ${CUDA_ARGS_IMPLIB}\n )\n endif()\n\n if(CUDA_ARGS_INCS)\n _absolute_paths(INC_DIRS ${CUDA_ARGS_INCS})\n foreach(INC_DIR ${INC_DIRS})\n set_property(\n TARGET ${CUDA_ARGS_NAME} APPEND PROPERTY\n INTERFACE_INCLUDE_DIRECTORIES \"${INC_DIR}\"\n )\n endforeach()\n endif()\n\n if(CUDA_ARGS_PUBINCS)\n _absolute_paths(INC_DIRS ${CUDA_ARGS_PUBINCS})\n foreach(INC_DIR ${INC_DIRS})\n set_property(\n TARGET ${CUDA_ARGS_NAME} APPEND PROPERTY\n INTERFACE_INCLUDE_DIRECTORIES \"${INC_DIR}\"\n )\n endforeach()\n endif()\n\n if(CUDA_ARGS_DEPS)\n add_dependencies(${CUDA_ARGS_NAME} \"${CUDA_ARGS_DEPS}\")\n endif()\n\n if(CUDA_ARGS_PACKED)\n install(\n TARGETS ${CUDA_ARGS_NAME}\n ARCHIVE DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n LIBRARY DESTINATION \"${CMAKE_INSTALL_LIBDIR}\"\n )\n if(CUDA_ARGS_PUBINCS)\n foreach(PACKED_EXCLUDE ${CUDA_ARGS_PACKED_EXCLUDES})\n list(APPEND PATTERN_EXCLUDES \"PATTERN;${PACKED_EXCLUDE};EXCLUDE\")\n endforeach()\n install(\n DIRECTORY ${CUDA_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR}\n FILES_MATCHING PATTERN \"*.h\" PATTERN \"*.hpp\"\n PATTERN \"*.hxx\" PATTERN \"*.cuh\"\n ${PATTERN_EXCLUDES}\n )\n endif()\n endif()\nendfunction()\n\n## Import a C/C++/CUDA interface library\nfunction(cuda_interface)\n if(NOT CMAKE_CUDA_COMPILER)\n message(FATAL_ERROR \"No CUDA language supported.\")\n endif()\n\n cmake_parse_arguments(\n CUDA_ARGS \"PACKED\" \"NAME\" \"INCS;PUBINCS;DEPS;PACKED_EXCLUDES\" ${ARGN}\n )\n\n if(NOT CUDA_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n add_library(${CUDA_ARGS_NAME} INTERFACE GLOBAL)\n if(CUDA_ARGS_INCS)\n _absolute_paths(INC_DIRS ${CUDA_ARGS_INCS})\n target_include_directories(${CUDA_ARGS_NAME} INTERFACE \"${INC_DIRS}\")\n endif()\n\n if(CUDA_ARGS_PUBINCS)\n _absolute_paths(INC_DIRS ${CUDA_ARGS_PUBINCS})\n target_include_directories(${CUDA_ARGS_NAME} INTERFACE \"${INC_DIRS}\")\n endif()\n\n if(CUDA_ARGS_DEPS)\n add_dependencies(${CUDA_ARGS_NAME} \"${CUDA_ARGS_DEPS}\")\n endif()\n\n if(CUDA_ARGS_PACKED AND CUDA_ARGS_PUBINCS)\n foreach(PACKED_EXCLUDE ${CUDA_ARGS_PACKED_EXCLUDES})\n list(APPEND PATTERN_EXCLUDES \"PATTERN;${PACKED_EXCLUDE};EXCLUDE\")\n endforeach()\n install(\n DIRECTORY ${CUDA_ARGS_PUBINCS} DESTINATION ${CMAKE_INSTALL_INCDIR}\n FILES_MATCHING PATTERN \"*.h\" PATTERN \"*.hpp\"\n PATTERN \"*.hxx\" PATTERN \"*.cuh\"\n ${PATTERN_EXCLUDES}\n )\n endif()\nendfunction()\n\n## Build a C/C++/CUDA executable google test program\nfunction(cuda_gtest)\n cmake_parse_arguments(\n CUDA_ARGS \"\" \"NAME;VERSION\"\n \"SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS;ARGS\" ${ARGN}\n )\n _find_gtest()\n cuda_test(\n NAME \"${CUDA_ARGS_NAME}\"\n VERSION \"${CUDA_ARGS_VERSION}\"\n SRCS \"${CUDA_ARGS_SRCS}\"\n INCS \"${CUDA_ARGS_INCS};${FIND_GTEST_INCS}\"\n DEFS \"${CUDA_ARGS_DEFS}\"\n LIBS \"${CUDA_ARGS_LIBS};${FIND_GTEST_LIBS}\"\n CFLAGS \"${CUDA_ARGS_CFLAGS}\"\n CXXFLAGS \"${CUDA_ARGS_CXXFLAGS}\"\n CUDAFLAGS \"${CUDA_ARGS_CUDAFLAGS}\"\n LDFLAGS \"${CUDA_ARGS_LDFLAGS}\"\n DEPS \"${CUDA_ARGS_DEPS}\"\n ARGS \"${CUDA_ARGS_ARGS}\"\n )\nendfunction()\n\n## Build a C/C++/CUDA executable google mock program\nfunction(cuda_gmock)\n cmake_parse_arguments(\n CUDA_ARGS \"\" \"NAME;VERSION\"\n \"SRCS;INCS;DEFS;LIBS;CFLAGS;CXXFLAGS;CUDAFLAGS;LDFLAGS;DEPS;ARGS\" ${ARGN}\n )\n _find_gmock()\n cuda_test(\n NAME \"${CUDA_ARGS_NAME}\"\n VERSION \"${CUDA_ARGS_VERSION}\"\n SRCS \"${CUDA_ARGS_SRCS}\"\n INCS \"${CUDA_ARGS_INCS};${FIND_GMOCK_INCS}\"\n DEFS \"${CUDA_ARGS_DEFS}\"\n LIBS \"${CUDA_ARGS_LIBS};${FIND_GMOCK_LIBS}\"\n CFLAGS \"${CUDA_ARGS_CFLAGS}\"\n CXXFLAGS \"${CUDA_ARGS_CXXFLAGS}\"\n CUDAFLAGS \"${CUDA_ARGS_CUDAFLAGS}\"\n LDFLAGS \"${CUDA_ARGS_LDFLAGS}\"\n DEPS \"${CUDA_ARGS_DEPS}\"\n ARGS \"${CUDA_ARGS_ARGS}\"\n )\nendfunction()\n\n## Add a subdirectory to the build\nfunction(go_directory)\n add_subdirectory(${ARGN})\nendfunction()\n\n## Add subdirectories to the build\nfunction(go_directories)\n foreach(SRC_DIR ${ARGN})\n add_subdirectory(${SRC_DIR})\n endforeach()\nendfunction()\n\n## Build a go executable program\nfunction(go_binary)\n find_program(\n GO_EXECUTABLE go PATHS $ENV{HOME}/go ENV GOROOT GOPATH PATH_SUFFIXES bin\n )\n if(NOT GO_EXECUTABLE)\n message(FATAL_ERROR \"No go language compiler found.\")\n endif()\n\n cmake_parse_arguments(\n GO_ARGS \"PACKED\" \"NAME\"\n \"GOPATH;SRCS;ASMFLAGS;GCFLAGS;LDFLAGS;DEPS\" ${ARGN}\n )\n if(NOT GO_ARGS_NAME)\n message(FATAL_ERROR \"No target name privated.\")\n endif()\n\n file(GLOB GO_ARGS_SRCS ${GO_ARGS_SRCS})\n if(NOT GO_ARGS_SRCS)\n message(FATAL_ERROR \"No source files/directories found of ${GO_ARGS_NAME}.\")\n endif()\n\n if(${CMAKE_SYSTEM_NAME} MATCHES \"Windows\")\n string(REPLACE \";\" \"\\;\" GO_ARGS_GOPATH \"${GO_ARGS_GOPATH}\")\n else()\n string(REPLACE \";\" \":\" GO_ARGS_GOPATH \"${GO_ARGS_GOPATH}\")\n endif()\n\n set(\n GO_OUTPUT_FILE\n ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${GO_ARGS_NAME}${CMAKE_EXECUTABLE_SUFFIX}\n )\n file(RELATIVE_PATH GO_OUTPUT_REL_FILE ${CMAKE_BINARY_DIR} ${GO_OUTPUT_FILE})\n add_custom_target(\n ${GO_ARGS_NAME}\n COMMAND ${CMAKE_COMMAND} -E env GOPATH=\"${GO_ARGS_GOPATH}\"\n \"${GO_EXECUTABLE}\" build -v -buildmode=exe\n -compiler=gc -gcflags=\"${GO_ARGS_GCFLAGS}\" -asmflags=\"${GO_ARGS_ASMFLAGS}\"\n -ldflags=\"${GO_ARGS_LDFLAGS}\"\n -o \"${GO_OUTPUT_FILE}\" \"${GO_ARGS_SRCS}\"\n WORKING_DIRECTORY \"${CMAKE_CURRENT_BINARY_DIR}\"\n DEPENDS \"${GO_ARGS_DEPS}\"\n COMMENT \"Building GO executable ${GO_OUTPUT_REL_FILE}\"\n )\n if(GO_ARGS_PACKED)\n install(PROGRAMS ${GO_OUTPUT_FILE} DESTINATION \"${CMAKE_INSTALL_BINDIR}\")\n endif()\nendfunction()\n\n## Fetch content\nfunction(_fetch_content)\n cmake_parse_arguments(\n DL_ARGS \"\"\n \"NAME;PATH;GIT_URL;GIT_TAG;HG_URL;HG_TAG;SVN_URL;SVN_REV;URL;URL_HASH\"\n \"\" ${ARGN}\n )\n\n if(NOT DL_ARGS_NAME)\n message(FATAL_ERROR \"No fetch name privated.\")\n endif()\n\n if(NOT DL_ARGS_PATH)\n # Download to current source directory\n set(DL_ARGS_PATH \"${CMAKE_CURRENT_SOURCE_DIR}/${DL_ARGS_NAME}\")\n endif()\n\n set(\n CMAKELISTS_CONTENT\n \"cmake_minimum_required(VERSION 3.1)\\n\"\n \"project(${DL_ARGS_NAME})\\n\"\n \"include(ExternalProject)\\n\"\n \"ExternalProject_Add(\\n\"\n \" ${DL_ARGS_NAME}\\n\"\n \" PREFIX \\\"external\\\"\\n\"\n \" GIT_REPOSITORY \\\"${DL_ARGS_GIT_URL}\\\"\\n\"\n \" GIT_TAG \\\"${DL_ARGS_GIT_TAG}\\\"\\n\"\n \" HG_REPOSITORY \\\"${DL_ARGS_HG_URL}\\\"\\n\"\n \" HG_TAG \\\"${DL_ARGS_HG_TAG}\\\"\\n\"\n \" SVN_REPOSITORY \\\"${DL_ARGS_SVN_URL}\\\"\\n\"\n \" SVN_REVISION \\\"${DL_ARGS_SVN_REV}\\\"\\n\"\n \" URL \\\"${DL_ARGS_URL}\\\"\\n\"\n \" URL_HASH \\\"${DL_ARGS_URL_HASH}\\\"\\n\"\n \" SOURCE_DIR \\\"${DL_ARGS_PATH}\\\"\\n\"\n \" BINARY_DIR \\\"\\\"\\n\"\n \" CONFIGURE_COMMAND \\\"\\\"\\n\"\n \" BUILD_COMMAND \\\"\\\"\\n\"\n \" INSTALL_COMMAND \\\"\\\"\\n\"\n \" TEST_COMMAND \\\"\\\"\\n\"\n \" LOG_DOWNLOAD ON\\n\"\n \" )\\n\"\n )\n set(\n CMAKELISTS_DIRECTORY\n \"${PROJECT_BINARY_DIR}/downloads/${DL_ARGS_NAME}\"\n )\n add_custom_target(\n external.${DL_ARGS_NAME}\n COMMAND \"${CMAKE_COMMAND}\" -G \"${CMAKE_GENERATOR}\" . &&\n \"${CMAKE_COMMAND}\" --build .\n WORKING_DIRECTORY \"${CMAKELISTS_DIRECTORY}\"\n )\n\n # Write a cmake script into folder\n file(WRITE \"${CMAKELISTS_DIRECTORY}/CMakeLists.txt\" ${CMAKELISTS_CONTENT})\n\n execute_process(\n COMMAND \"${CMAKE_COMMAND}\" -G \"${CMAKE_GENERATOR}\" .\n WORKING_DIRECTORY \"${CMAKELISTS_DIRECTORY}\"\n )\n execute_process(\n COMMAND \"${CMAKE_COMMAND}\" --build .\n WORKING_DIRECTORY \"${CMAKELISTS_DIRECTORY}\"\n )\nendfunction()\n\n## Download a git repository\nfunction(git_repository)\n cmake_parse_arguments(GIT_ARGS \"\" \"NAME;PATH;URL;TAG\" \"\" ${ARGN})\n\n if(NOT GIT_ARGS_NAME)\n message(FATAL_ERROR \"No repository name privated.\")\n endif()\n if(NOT GIT_ARGS_URL)\n message(FATAL_ERROR \"No repository URL privated.\")\n endif()\n\n if(GIT_ARGS_PATH AND NOT IS_ABSOLUTE ${GIT_ARGS_PATH})\n get_filename_component(GIT_ARGS_PATH ${GIT_ARGS_PATH} ABSOLUTE)\n endif()\n\n _fetch_content(\n NAME \"${GIT_ARGS_NAME}\"\n PATH \"${GIT_ARGS_PATH}\"\n GIT_URL \"${GIT_ARGS_URL}\"\n GIT_TAG \"${GIT_ARGS_TAG}\"\n )\nendfunction()\n\n## Download a hg repository\nfunction(hg_repository)\n cmake_parse_arguments(HG_ARGS \"\" \"NAME;PATH;URL;TAG\" \"\" ${ARGN})\n\n if(NOT HG_ARGS_NAME)\n message(FATAL_ERROR \"No repository name privated.\")\n endif()\n if(NOT HG_ARGS_URL)\n message(FATAL_ERROR \"No repository URL privated.\")\n endif()\n\n if(HG_ARGS_PATH AND NOT IS_ABSOLUTE ${HG_ARGS_PATH})\n get_filename_component(HG_ARGS_PATH ${HG_ARGS_PATH} ABSOLUTE)\n endif()\n\n _fetch_content(\n NAME \"${HG_ARGS_NAME}\"\n PATH \"${HG_ARGS_PATH}\"\n HG_URL \"${HG_ARGS_URL}\"\n HG_TAG \"${HG_ARGS_TAG}\"\n )\nendfunction()\n\n## Download a svn repository\nfunction(svn_repository)\n cmake_parse_arguments(SVN_ARGS \"\" \"NAME;PATH;URL;REV\" \"\" ${ARGN})\n\n if(NOT SVN_ARGS_NAME)\n message(FATAL_ERROR \"No repository name privated.\")\n endif()\n if(NOT SVN_ARGS_URL)\n message(FATAL_ERROR \"No repository URL privated.\")\n endif()\n\n if(SVN_ARGS_PATH AND NOT IS_ABSOLUTE ${SVN_ARGS_PATH})\n get_filename_component(SVN_ARGS_PATH ${SVN_ARGS_PATH} ABSOLUTE)\n endif()\n\n _fetch_content(\n NAME \"${SVN_ARGS_NAME}\"\n PATH \"${SVN_ARGS_PATH}\"\n SVN_URL \"${SVN_ARGS_URL}\"\n SVN_REV \"${SVN_ARGS_REV}\"\n )\nendfunction()\n\n## Download a http archive\nfunction(http_archive)\n cmake_parse_arguments(HTTP_ARGS \"\" \"NAME;PATH;URL;SHA256;SHA1;MD5\" \"\" ${ARGN})\n\n if(NOT HTTP_ARGS_NAME)\n message(FATAL_ERROR \"No archive name privated.\")\n endif()\n if(NOT HTTP_ARGS_URL)\n message(FATAL_ERROR \"No archive URL privated.\")\n endif()\n\n if(HTTP_ARGS_PATH AND NOT IS_ABSOLUTE ${HTTP_ARGS_PATH})\n get_filename_component(HTTP_ARGS_PATH ${HTTP_ARGS_PATH} ABSOLUTE)\n endif()\n\n if(HTTP_ARGS_SHA256)\n set(HTTP_URL_HASH \"SHA256=${HTTP_ARGS_SHA256}\")\n elseif(HTTP_ARGS_SHA1)\n set(HTTP_URL_HASH \"SHA1=${HTTP_ARGS_SHA1}\")\n elseif(HTTP_ARGS_MD5)\n set(HTTP_URL_HASH \"MD5=${HTTP_ARGS_MD5}\")\n else()\n set(HTTP_URL_HASH \"\")\n endif()\n\n _fetch_content(\n NAME \"${HTTP_ARGS_NAME}\"\n PATH \"${HTTP_ARGS_PATH}\"\n URL \"${HTTP_ARGS_URL}\"\n URL_HASH \"${HTTP_URL_HASH}\"\n )\nendfunction()\n\n## Retrieve a version string from GIT\nfunction(git_version _RESULT _SOURCES_DIR)\n find_package(Git REQUIRED)\n\n if(NOT IS_ABSOLUTE ${_SOURCES_DIR})\n get_filename_component(_SOURCES_DIR ${_SOURCES_DIR} ABSOLUTE)\n endif()\n\n # git describe --tags\n execute_process(\n COMMAND \"${GIT_EXECUTABLE}\" describe --tags\n WORKING_DIRECTORY \"${_SOURCES_DIR}\"\n RESULT_VARIABLE GIT_VER_RESULT\n OUTPUT_VARIABLE GIT_VER_OUTPUT\n ERROR_VARIABLE GIT_VER_ERROR\n )\n if(GIT_VER_RESULT EQUAL 0)\n string(STRIP ${GIT_VER_OUTPUT} GIT_VER_OUTPUT)\n set(${_RESULT} \"${GIT_VER_OUTPUT}\" PARENT_SCOPE)\n return()\n endif()\n\n # git rev-parse --short HEAD\n execute_process(\n COMMAND \"${GIT_EXECUTABLE}\" rev-parse --short HEAD\n WORKING_DIRECTORY \"${_SOURCES_DIR}\"\n RESULT_VARIABLE GIT_VER_RESULT\n OUTPUT_VARIABLE GIT_VER_OUTPUT\n ERROR_VARIABLE GIT_VER_ERROR\n )\n if(GIT_VER_RESULT EQUAL 0)\n string(STRIP ${GIT_VER_OUTPUT} GIT_VER_OUTPUT)\n set(${_RESULT} \"g${GIT_VER_OUTPUT}\" PARENT_SCOPE)\n return()\n endif()\n\n set(${_RESULT} \"\" PARENT_SCOPE)\nendfunction()\n\n## Retrieve a version string from HG\nfunction(hg_version _RESULT _SOURCES_DIR)\n find_package(Hg REQUIRED)\n\n if(NOT IS_ABSOLUTE ${_SOURCES_DIR})\n get_filename_component(_SOURCES_DIR ${_SOURCES_DIR} ABSOLUTE)\n endif()\n\n # hg log -T \"{latesttagdistance}\" -r .\n execute_process(\n COMMAND \"${HG_EXECUTABLE}\" log -T \"{latesttagdistance}\" -r .\n WORKING_DIRECTORY \"${_SOURCES_DIR}\"\n RESULT_VARIABLE HG_VER_RESULT\n OUTPUT_VARIABLE HG_VER_OUTPUT\n ERROR_VARIABLE HG_VER_ERROR\n )\n if(HG_VER_RESULT EQUAL 0)\n string(STRIP ${HG_VER_OUTPUT} HG_VER_OUTPUT)\n if(HG_VER_OUTPUT STREQUAL \"0\")\n # hg log -T \"{latesttag}\" -r .\n execute_process(\n COMMAND \"${HG_EXECUTABLE}\" log -T \"{latesttag}\" -r .\n WORKING_DIRECTORY \"${_SOURCES_DIR}\"\n RESULT_VARIABLE HG_VER_RESULT\n OUTPUT_VARIABLE HG_VER_OUTPUT\n ERROR_VARIABLE HG_VER_ERROR\n )\n else()\n # hg log -T \"{latesttag}-{latesttagdistance}-h{node|short}\" -r .\n execute_process(\n COMMAND \"${HG_EXECUTABLE}\" log\n -T \"{latesttag}-{latesttagdistance}-h{node|short}\" -r .\n WORKING_DIRECTORY \"${_SOURCES_DIR}\"\n RESULT_VARIABLE HG_VER_RESULT\n OUTPUT_VARIABLE HG_VER_OUTPUT\n ERROR_VARIABLE HG_VER_ERROR\n )\n endif()\n\n if(HG_VER_RESULT EQUAL 0)\n string(STRIP ${HG_VER_OUTPUT} HG_VER_OUTPUT)\n if(NOT HG_VER_OUTPUT MATCHES \"^null.*\")\n set(${_RESULT} \"${HG_VER_OUTPUT}\" PARENT_SCOPE)\n return()\n endif()\n endif()\n endif()\n\n # hg log -T \"h{node|short}\" -r .\n execute_process(\n COMMAND \"${HG_EXECUTABLE}\" log -T \"h{node|short}\" -r .\n WORKING_DIRECTORY \"${_SOURCES_DIR}\"\n RESULT_VARIABLE HG_VER_RESULT\n OUTPUT_VARIABLE HG_VER_OUTPUT\n ERROR_VARIABLE HG_VER_ERROR\n )\n if(HG_VER_RESULT EQUAL 0)\n string(STRIP ${HG_VER_OUTPUT} HG_VER_OUTPUT)\n set(${_RESULT} \"${HG_VER_OUTPUT}\" PARENT_SCOPE)\n return()\n endif()\n\n set(${_RESULT} \"\" PARENT_SCOPE)\nendfunction()\n\n## Retrieve a version string from SVN\nfunction(svn_version _RESULT _SOURCES_DIR)\n find_package(Subversion REQUIRED)\n\n if(NOT IS_ABSOLUTE ${_SOURCES_DIR})\n get_filename_component(_SOURCES_DIR ${_SOURCES_DIR} ABSOLUTE)\n endif()\n\n # svn info --show-item revision\n execute_process(\n COMMAND \"${Subversion_SVN_EXECUTABLE}\" info --show-item revision\n WORKING_DIRECTORY \"${_SOURCES_DIR}\"\n RESULT_VARIABLE SVN_VER_RESULT\n OUTPUT_VARIABLE SVN_VER_OUTPUT\n ERROR_VARIABLE SVN_VER_ERROR\n )\n if(SVN_VER_RESULT EQUAL 0)\n string(STRIP ${SVN_VER_OUTPUT} SVN_VER_OUTPUT)\n set(${_RESULT} \"r${SVN_VER_OUTPUT}\" PARENT_SCOPE)\n return()\n endif()\n\n set(${_RESULT} \"\" PARENT_SCOPE)\nendfunction()\n\n_find_workspace_directory(BAZEL_WORKSPACE_DIR)\nif(BAZEL_WORKSPACE_DIR)\n include(\"${BAZEL_WORKSPACE_DIR}/Workspace.cmake\")\nendif()\n" }, { "path": "cmake/option.cmake", "content": "## https://en.wikipedia.org/wiki/List_of_Intel_CPU_microarchitectures \n## https://en.wikipedia.org/wiki/List_of_AMD_CPU_microarchitectures \n## https://gcc.gnu.org/onlinedocs/gcc/x86-Options.html \n\n## Intel Microarchitectures\noption(ENABLE_NEHALEM \"Enable Intel Nehalem CPU microarchitecture\" OFF)\noption(ENABLE_SANDYBRIDGE \"Enable Intel Sandy Bridge CPU microarchitecture\" OFF)\noption(ENABLE_HASWELL \"Enable Intel Haswell CPU microarchitecture\" OFF)\noption(ENABLE_BROADWELL \"Enable Intel Broadwell CPU microarchitecture\" OFF)\noption(ENABLE_SKYLAKE \"Enable Intel Skylake CPU microarchitecture\" OFF)\noption(ENABLE_SKYLAKE_AVX512 \"Enable Intel Skylake Server CPU microarchitecture\" OFF)\noption(ENABLE_ICELAKE \"Enable Intel Icelake CPU microarchitecture\" OFF)\noption(ENABLE_SAPPHIRERAPIDS \"Enable Intel Sapphire Rapids Server CPU microarchitecture\" OFF)\noption(ENABLE_EMERALDRAPIDS \"Enable Intel Emerald Rapids Server CPU microarchitecture\" OFF)\noption(ENABLE_GRANITERAPIDS \"Enable Intel Granite Rapids Server CPU microarchitecture\" OFF)\n\noption(ENABLE_NATIVE \"Enable native CPU microarchitecture\" OFF)\n\n## AMD Microarchitectures\noption(ENABLE_ZEN1 \"Enable AMD Zen+ Family 17h CPU microarchitecture\" OFF)\noption(ENABLE_ZEN2 \"Enable AMD Zen 2 Family 17h CPU microarchitecture\" OFF)\noption(ENABLE_ZEN3 \"Enable AMD Zen 3 Family 19h CPU microarchitecture\" OFF)\n\n## ARM architectures\noption(ENABLE_ARMV8A \"Enable ARMv8-a architecture\" OFF)\noption(ENABLE_ARMV8.1A \"Enable ARMv8.1-a architecture\" OFF)\noption(ENABLE_ARMV8.2A \"Enable ARMv8.2-a architecture\" OFF)\noption(ENABLE_ARMV8.3A \"Enable ARMv8.3-a architecture\" OFF)\noption(ENABLE_ARMV8.4A \"Enable ARMv8.4-a architecture\" OFF)\noption(ENABLE_ARMV8.5A \"Enable ARMv8.5-a architecture\" OFF)\noption(ENABLE_ARMV8.6A \"Enable ARMv8.6-a architecture\" OFF)\n\n## OpenMP option\noption(ENABLE_OPENMP \"Enable OpenMP support\" OFF)\n\nset(ARCH_OPTIONS\n ENABLE_NEHALEM ENABLE_SANDYBRIDGE ENABLE_HASWELL ENABLE_BROADWELL ENABLE_SKYLAKE\n ENABLE_SKYLAKE_AVX512 ENABLE_ICELAKE ENABLE_SAPPHIRERAPIDS ENABLE_EMERALDRAPIDS ENABLE_GRANITERAPIDS\n ENABLE_ZEN1 ENABLE_ZEN2 ENABLE_ZEN3\n ENABLE_ARMV8A ENABLE_ARMV8.1A ENABLE_ARMV8.2A ENABLE_ARMV8.3A ENABLE_ARMV8.4A\n ENABLE_ARMV8.5A ENABLE_ARMV8.6A\n ENABLE_NATIVE\n)\n\noption(AUTO_DETECT_ARCH \"Auto detect CPU microarchitecture\" ON)\nforeach(opt IN LISTS ARCH_OPTIONS)\n if(${opt})\n set(AUTO_DETECT_ARCH OFF)\n break()\n endif()\nendforeach()\n\ninclude(CheckCCompilerFlag)\n\nfunction(_AppendFlags _RESULT _FLAG)\n if(${_RESULT} AND NOT \"${${_RESULT}}\" MATCHES \"${_FLAG}\")\n set(${_RESULT} \"${${_RESULT}} ${_FLAG}\" PARENT_SCOPE)\n else()\n set(${_RESULT} \"${_FLAG}\" PARENT_SCOPE)\n endif()\nendfunction()\n\nmacro(add_arch_flag FLAG VAR_NAME OPTION_NAME)\n check_c_compiler_flag(\"${FLAG}\" COMPILER_SUPPORT_${VAR_NAME})\n if(COMPILER_SUPPORT_${VAR_NAME})\n _AppendFlags(CMAKE_C_FLAGS \"${FLAG}\")\n _AppendFlags(CMAKE_CXX_FLAGS \"${FLAG}\")\n set(${VAR_NAME}_ENABLED ON)\n else()\n if(${OPTION_NAME})\n message(FATAL_ERROR \"Compiler does not support required flag: '${FLAG}' for ${OPTION_NAME}\")\n else()\n set(${VAR_NAME}_ENABLED OFF)\n endif()\n endif()\nendmacro()\n\nfunction(_setup_armv8_march)\n set(_arch \"armv8\")\n check_c_compiler_flag(\"-march=${_arch}\" _COMP_SUPP_${_arch})\n if(_COMP_SUPP_${_arch})\n _AppendFlags(CMAKE_C_FLAGS \"-march=${_arch}\")\n _AppendFlags(CMAKE_CXX_FLAGS \"-march=${_arch}\")\n set(CMAKE_C_FLAGS \"${CMAKE_C_FLAGS}\" PARENT_SCOPE)\n set(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS}\" PARENT_SCOPE)\n return()\n else()\n message(WARNING \"No ARMv8 march flag supported by compiler.\")\n endif()\nendfunction()\n\nfunction(_setup_x86_march)\n set(_arch \"x86-64\")\n check_c_compiler_flag(\"-march=${_arch}\" _COMP_SUPP_${_arch})\n if(_COMP_SUPP_${_arch})\n _AppendFlags(CMAKE_C_FLAGS \"-march=${_arch}\")\n _AppendFlags(CMAKE_CXX_FLAGS \"-march=${_arch}\")\n set(CMAKE_C_FLAGS \"${CMAKE_C_FLAGS}\" PARENT_SCOPE)\n set(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS}\" PARENT_SCOPE)\n return()\n else()\n message(WARNING \"No known x86 march flag supported; falling back to generic.\")\n endif()\nendfunction()\n\nfunction(setup_compiler_march_for_x86 VAR_NAME_SSE VAR_NAME_AVX2 VAR_NAME_AVX512 VAR_NAME_AVX512FP16)\n #sse\n set(${VAR_NAME_SSE} \"-march=corei7\" PARENT_SCOPE)\n\n #avx 2\n set(${VAR_NAME_AVX2} \"-march=core-avx2\" PARENT_SCOPE)\n\n #avx512\n set(_x86_flags_avx512 \"icelake-server\" \"skylake-avx512\" \"core-avx2\" \"x86-64\")\n foreach(_arch_avx512 IN LISTS _x86_flags_avx512)\n check_c_compiler_flag(\"-march=${_arch_avx512}\" _COMP_SUPP_${_arch_avx512})\n if(_COMP_SUPP_${_arch_avx512})\n set(${VAR_NAME_AVX512} \"-march=${_arch_avx512}\" PARENT_SCOPE)\n break()\n endif()\n endforeach()\n\n #avx512fp16\n set(_x86_flags_avx512fp16\n \"sapphirerapids\" \"icelake-server\" \"skylake-avx512\" \"core-avx2\" \"x86-64\"\n )\n foreach(_arch_avx512fp16 IN LISTS _x86_flags_avx512fp16)\n check_c_compiler_flag(\"-march=${_arch_avx512fp16}\" _COMP_SUPP_${_arch_avx512fp16})\n if(_COMP_SUPP_${_arch_avx512fp16})\n set(${VAR_NAME_AVX512FP16} \"-march=${_arch_avx512fp16}\" PARENT_SCOPE)\n break()\n endif()\n endforeach()\nendfunction()\n\nif(MSVC)\n # Prefer higher ISAs\n foreach(_isa IN ITEMS \"AVX512\" \"AVX2\" \"AVX\" \"SSE2\")\n check_c_compiler_flag(\"/arch:${_isa}\" _COMP_SUPP_${_isa})\n if(_COMP_SUPP_${_isa})\n _AppendFlags(CMAKE_C_FLAGS \"/arch:${_isa}\")\n _AppendFlags(CMAKE_CXX_FLAGS \"/arch:${_isa}\")\n message(STATUS \"MSVC: enabled /arch:${_isa}\")\n break()\n endif()\n endforeach()\n return()\nendif()\n\nif(NOT AUTO_DETECT_ARCH)\n if(ENABLE_NATIVE)\n add_arch_flag(\"-march=native\" NATIVE ENABLE_NATIVE)\n endif()\n\n if(ENABLE_ZEN3)\n add_arch_flag(\"-march=znver3\" ZNVER3 ENABLE_ZEN3)\n endif()\n\n if(ENABLE_ZEN2)\n add_arch_flag(\"-march=znver2\" ZNVER2 ENABLE_ZEN2)\n endif()\n\n if(ENABLE_ZEN1)\n add_arch_flag(\"-march=znver1\" ZNVER1 ENABLE_ZEN1)\n endif()\n\n if(ENABLE_GRANITERAPIDS)\n add_arch_flag(\"-march=graniterapids\" GRANITERAPIDS ENABLE_GRANITERAPIDS)\n endif()\n\n if(ENABLE_EMERALDRAPIDS)\n add_arch_flag(\"-march=emeraldrapids\" EMERALDRAPIDS ENABLE_EMERALDRAPIDS)\n endif()\n\n if(ENABLE_SAPPHIRERAPIDS)\n add_arch_flag(\"-march=sapphirerapids\" SAPPHIRERAPIDS ENABLE_SAPPHIRERAPIDS)\n endif()\n\n if(ENABLE_ICELAKE)\n add_arch_flag(\"-march=icelake-server\" ICELAKE ENABLE_ICELAKE)\n endif()\n\n if(ENABLE_SKYLAKE_AVX512)\n add_arch_flag(\"-march=skylake-avx512\" SKYLAKE_AVX512 ENABLE_SKYLAKE_AVX512)\n endif()\n\n if(ENABLE_SKYLAKE)\n add_arch_flag(\"-march=skylake\" SKYLAKE ENABLE_SKYLAKE)\n endif()\n\n if(ENABLE_BROADWELL)\n add_arch_flag(\"-march=broadwell\" BROADWELL ENABLE_BROADWELL)\n endif()\n\n if(ENABLE_HASWELL)\n add_arch_flag(\"-march=haswell\" HASWELL ENABLE_HASWELL)\n endif()\n\n if(ENABLE_SANDYBRIDGE)\n add_arch_flag(\"-march=sandybridge\" SANDYBRIDGE ENABLE_SANDYBRIDGE)\n endif()\n\n if(ENABLE_NEHALEM)\n add_arch_flag(\"-march=nehalem\" NEHALEM ENABLE_NEHALEM)\n endif()\n\n # ARM (newest first — allow multiple? usually only one)\n # But GCC allows only one -march=, so honor highest enabled\n if(ENABLE_ARMV8.6A)\n add_arch_flag(\"-march=armv8.6-a\" ARMV86A ENABLE_ARMV8.6A)\n endif()\n if(ENABLE_ARMV8.5A)\n add_arch_flag(\"-march=armv8.5-a\" ARMV85A ENABLE_ARMV8.5A)\n endif()\n if(ENABLE_ARMV8.4A)\n add_arch_flag(\"-march=armv8.4-a\" ARMV84A ENABLE_ARMV8.4A)\n endif()\n if(ENABLE_ARMV8.3A)\n add_arch_flag(\"-march=armv8.3-a\" ARMV83A ENABLE_ARMV8.3A)\n endif()\n if(ENABLE_ARMV8.2A)\n add_arch_flag(\"-march=armv8.2-a\" ARMV82A ENABLE_ARMV8.2A)\n endif()\n if(ENABLE_ARMV8.1A)\n add_arch_flag(\"-march=armv8.1-a\" ARMV81A ENABLE_ARMV8.1A)\n endif()\n if(ENABLE_ARMV8A)\n add_arch_flag(\"-march=armv8-a\" ARMV8A ENABLE_ARMV8A)\n endif()\n\nelse()\n # AUTO DETECT\n # Heuristic: detect host architecture and probe appropriate flags\n if(CMAKE_SYSTEM_PROCESSOR MATCHES \"aarch64|arm64|ARM64\")\n _setup_armv8_march()\n elseif(CMAKE_SYSTEM_PROCESSOR MATCHES \"x86_64|i686|i386|x64\")\n _setup_x86_march()\n else()\n message(WARNING \"Unknown host architecture: ${CMAKE_SYSTEM_PROCESSOR}; no -march= set.\")\n endif()\nendif()\n\n# -----------------------------\n# OpenMP\n# -----------------------------\nif(ENABLE_OPENMP)\n find_package(OpenMP REQUIRED)\n if(OpenMP_C_FLAGS)\n _AppendFlags(CMAKE_C_FLAGS \"${OpenMP_C_FLAGS}\")\n endif()\n if(OpenMP_CXX_FLAGS)\n _AppendFlags(CMAKE_CXX_FLAGS \"${OpenMP_CXX_FLAGS}\")\n endif()\nendif()\n" }, { "path": "cmake/utils.cmake", "content": "function(apply_patch_once patch_name target_dir patch_file)\n set(mark_file \"${target_dir}/.${patch_name}_patched\")\n\n if(EXISTS \"${mark_file}\")\n #message(STATUS \"Patch '${patch_name}' already applied to ${target_dir}, skipping.\")\n return()\n endif()\n\n if(NOT EXISTS \"${patch_file}\")\n message(FATAL_ERROR \"Patch file '${patch_file}' not found!\")\n endif()\n\n #message(STATUS \"Applying patch '${patch_name}' to ${target_dir} ...\")\n execute_process(\n COMMAND patch -p1 -i \"${patch_file}\"\n WORKING_DIRECTORY \"${target_dir}\"\n RESULT_VARIABLE patch_result\n OUTPUT_VARIABLE patch_stdout\n ERROR_VARIABLE patch_stderr\n )\n\n if(NOT patch_result EQUAL 0)\n message(FATAL_ERROR \"Failed to apply patch '${patch_name}' to ${target_dir}:\\n${patch_stderr}\")\n else()\n #message(STATUS \"Patch '${patch_name}' applied successfully:\\n${patch_stdout}\")\n file(WRITE \"${mark_file}\" \"patched\")\n endif()\nendfunction()\n" }, { "path": "examples/c++/CMakeLists.txt", "content": "cmake_minimum_required(VERSION 3.13)\ncmake_policy(SET CMP0077 NEW)\nproject(zvec-example-c++)\nset(CMAKE_CXX_STANDARD 17)\n\n# Enable compile_commands.json\nset(CMAKE_EXPORT_COMPILE_COMMANDS ON)\n\n# --- Paths to Zvec and dependencies ---\n# Allow custom host build directory, default to \"build\"\nif(NOT DEFINED HOST_BUILD_DIR)\n set(HOST_BUILD_DIR \"build\")\nendif()\n\nset(ZVEC_INCLUDE_DIR ${CMAKE_BINARY_DIR}/../../../src/include)\nset(ZVEC_LIB_DIR ${CMAKE_BINARY_DIR}/../../../${HOST_BUILD_DIR}/lib)\nset(ZVEC_DEPENDENCY_LIB_DIR ${CMAKE_BINARY_DIR}/../../../${HOST_BUILD_DIR}/external/usr/local/lib)\n\n# Add include and library search paths\ninclude_directories(${ZVEC_INCLUDE_DIR})\nlink_directories(${ZVEC_LIB_DIR} ${ZVEC_DEPENDENCY_LIB_DIR})\n\n# --- Determine debug/release library names ---\nif(CMAKE_BUILD_TYPE STREQUAL \"Debug\")\n set(GLOG_LIB glogd)\n set(GFLAGS_LIB gflags_nothreads_debug)\n set(PROTOBUF_LIB protobufd)\nelse()\n set(GLOG_LIB glog)\n set(GFLAGS_LIB gflags_nothreads)\n set(PROTOBUF_LIB protobuf)\nendif()\n\n# --- Dependency groups ---\nfind_package(Threads REQUIRED)\n\nset(zvec_ailego_deps\n arrow\n parquet\n arrow_bundled_dependencies\n ${CMAKE_THREAD_LIBS_INIT}\n ${CMAKE_DL_LIBS}\n)\n\nset(zvec_core_deps\n zvec_turbo\n)\n\nset(zvec_db_deps\n roaring\n rocksdb\n arrow\n arrow_acero\n arrow_bundled_dependencies\n arrow_compute\n arrow_dataset\n parquet\n antlr4-runtime\n ${GLOG_LIB}\n ${GFLAGS_LIB}\n ${PROTOBUF_LIB}\n lz4\n)\n\n# --- Create INTERFACE targets for Zvec components ---\n\n# zvec_ailego: links libzvec_ailego.a + its deps\nadd_library(zvec-ailego INTERFACE)\ntarget_link_libraries(zvec-ailego INTERFACE\n -lzvec_ailego\n ${zvec_ailego_deps}\n)\n\n# zvec_core: links libzvec_core.a via special flags (handled externally), but declare logical deps\nadd_library(zvec-core INTERFACE)\nif(CMAKE_SYSTEM_NAME STREQUAL \"Linux\")\n target_link_libraries(zvec-core INTERFACE\n -Wl,--whole-archive\n zvec_core\n -Wl,--no-whole-archive\n -Wl,--start-group\n zvec-ailego\n ${zvec_core_deps}\n -Wl,--end-group\n )\nelseif(APPLE)\n target_link_libraries(zvec-core INTERFACE\n -Wl,-force_load ${ZVEC_LIB_DIR}/libzvec_core.a\n zvec-ailego\n ${zvec_core_deps}\n )\nelseif(ANDROID)\n target_link_libraries(zvec-core INTERFACE\n -Wl,--whole-archive\n zvec_core\n -Wl,--no-whole-archive\n -Wl,--start-group\n zvec-ailego\n ${zvec_core_deps}\n -Wl,--end-group\n )\nelse()\n message(FATAL_ERROR \"Unsupported platform: ${CMAKE_SYSTEM_NAME}\")\nendif()\n\n# zvec_db: links libzvec_db.a + all deps\nadd_library(zvec-db INTERFACE)\nif(CMAKE_SYSTEM_NAME STREQUAL \"Linux\")\n target_link_libraries(zvec-db INTERFACE\n zvec_db\n zvec-core\n zvec-ailego\n -Wl,--start-group\n ${zvec_db_deps}\n -Wl,--end-group\n )\nelseif(APPLE)\n target_link_libraries(zvec-db INTERFACE\n zvec_db\n zvec-core\n zvec-ailego\n ${zvec_db_deps}\n )\nelseif(ANDROID)\n target_link_libraries(zvec-db INTERFACE\n zvec_db\n zvec-core\n zvec-ailego\n -Wl,--start-group\n ${zvec_db_deps}\n -Wl,--end-group\n )\nelse()\n message(FATAL_ERROR \"Unsupported platform: ${CMAKE_SYSTEM_NAME}\")\nendif()\n\n\n# --- Main executable ---\nadd_executable(db-example db/main.cc)\ntarget_link_libraries(db-example PRIVATE\n zvec-db\n)\nif(ANDROID)\n target_link_libraries(db-example PRIVATE\n log\n )\nendif()\n\nadd_executable(core-example core/main.cc)\ntarget_link_libraries(core-example PRIVATE\n zvec-core\n)\n\nadd_executable(ailego-example ailego/main.cc)\ntarget_link_libraries(ailego-example PRIVATE\n zvec-ailego\n)\n\n# Strip symbols to reduce executable size\nif(CMAKE_BUILD_TYPE STREQUAL \"Release\" AND ANDROID)\n add_custom_command(TARGET db-example POST_BUILD\n COMMAND ${CMAKE_STRIP} \"$\"\n COMMENT \"Stripping symbols from db-example\")\n add_custom_command(TARGET core-example POST_BUILD\n COMMAND ${CMAKE_STRIP} \"$\"\n COMMENT \"Stripping symbols from core-example\")\n add_custom_command(TARGET ailego-example POST_BUILD\n COMMAND ${CMAKE_STRIP} \"$\"\n COMMENT \"Stripping symbols from ailego-example\")\nendif()\n\n# Optimize for size\nif(CMAKE_BUILD_TYPE STREQUAL \"Release\" AND ANDROID)\n set_property(TARGET db-example core-example ailego-example\n PROPERTY COMPILE_FLAGS \"-Os\")\n set_property(TARGET db-example core-example ailego-example\n PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)\nendif()\n" }, { "path": "examples/c++/ailego/main.cc", "content": "#include \n#include \n#include \n\nusing namespace zvec;\n\nint main() {\n std::string a{\"hello world\"};\n\n std::cout << ailego::StringHelper::StartsWith(a, \"hello\") << std::endl;\n}" }, { "path": "examples/c++/core/main.cc", "content": "#include \n#include \n#include \n#include \n#include \n#include \n\nusing namespace zvec::core_interface;\n\nconstexpr uint32_t kDimension = 64;\nconst std::string index_name{\"test.index\"};\n\nIndex::Pointer create_index(const BaseIndexParam::Pointer ¶m,\n int doc_num = 10) {\n auto index = IndexFactory::CreateAndInitIndex(*param);\n if (!index) {\n std::cout << \"Failed to create index.\" << std::endl;\n return nullptr;\n }\n\n int ret = index->Open(\n index_name, StorageOptions{StorageOptions::StorageType::kMMAP, true});\n if (ret != 0) {\n std::cout << \"Failed to open index.\" << std::endl;\n return nullptr;\n }\n\n for (int i = 0; i < doc_num; ++i) {\n std::vector vector(kDimension, i / 10.0f + 0.1f);\n VectorData vector_data;\n vector_data.vector = DenseVector{vector.data()};\n ret = index->Add(vector_data, i);\n if (ret != 0) {\n std::cout << \"Failed to add to index.\" << std::endl;\n return nullptr;\n }\n }\n\n ret = index->Train();\n if (ret != 0) {\n std::cout << \"Failed to train index.\" << std::endl;\n return nullptr;\n }\n\n return index;\n}\n\nint main() {\n char cmd_buf[100];\n snprintf(cmd_buf, 100, \"rm -f %s\", index_name.c_str());\n system(cmd_buf);\n\n auto param = HNSWIndexParamBuilder()\n .WithMetricType(MetricType::kInnerProduct)\n .WithDataType(DataType::DT_FP32)\n .WithDimension(kDimension)\n .WithIsSparse(false)\n .Build();\n auto index = create_index(param, 1);\n std::cout << \"index stats: \" << index->GetDocCount() << std::endl;\n\n // query\n auto query_param = HNSWQueryParamBuilder()\n .with_topk(10)\n .with_fetch_vector(true)\n .with_ef_search(20)\n .build();\n\n SearchResult result;\n VectorData query;\n std::vector vector(kDimension, 0.1f);\n query.vector = DenseVector{vector.data()};\n int ret = index->Search(query, query_param, &result);\n if (ret != 0) {\n std::cout << \"Failed to search index.\" << std::endl;\n return -1;\n }\n\n std::cout << \"query results: \" << result.doc_list_.size() << std::endl;\n if (result.doc_list_.size() == 0) {\n std::cout << \"No results found.\" << std::endl;\n return -1;\n }\n\n std::cout << \"key: \" << result.doc_list_[0].key()\n << \", score: \" << result.doc_list_[0].score() << std::endl;\n\n return 0;\n}" }, { "path": "examples/c++/db/main.cc", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nusing namespace zvec;\n\nDoc create_doc(const uint64_t doc_id, const CollectionSchema &schema,\n std::string pk = \"\") {\n Doc new_doc;\n if (pk.empty()) {\n pk = \"pk_\" + std::to_string(doc_id);\n }\n new_doc.set_pk(pk);\n\n for (auto &field : schema.fields()) {\n switch (field->data_type()) {\n case DataType::BINARY: {\n std::string binary_str(\"binary_\" + std::to_string(doc_id));\n new_doc.set(field->name(), binary_str);\n break;\n }\n case DataType::BOOL:\n new_doc.set(field->name(), doc_id % 10 == 0);\n break;\n case DataType::INT32:\n new_doc.set(field->name(), (int32_t)doc_id);\n break;\n case DataType::INT64:\n new_doc.set(field->name(), (int64_t)doc_id);\n break;\n case DataType::UINT32:\n new_doc.set(field->name(), (uint32_t)doc_id);\n break;\n case DataType::UINT64:\n new_doc.set(field->name(), (uint64_t)doc_id);\n break;\n case DataType::FLOAT:\n new_doc.set(field->name(), (float)doc_id);\n break;\n case DataType::DOUBLE:\n new_doc.set(field->name(), (double)doc_id);\n break;\n case DataType::STRING:\n new_doc.set(field->name(),\n \"value_\" + std::to_string(doc_id));\n break;\n case DataType::ARRAY_BINARY: {\n std::vector bin_vec;\n for (size_t i = 0; i < (doc_id % 10); i++) {\n bin_vec.push_back(\"bin_\" + std::to_string(i));\n }\n new_doc.set>(field->name(), bin_vec);\n break;\n }\n case DataType::ARRAY_BOOL:\n new_doc.set>(field->name(),\n std::vector(10, doc_id % 10 == 0));\n break;\n case DataType::ARRAY_INT32:\n new_doc.set>(\n field->name(), std::vector(10, (int32_t)doc_id));\n break;\n case DataType::ARRAY_INT64:\n new_doc.set>(\n field->name(), std::vector(10, (int64_t)doc_id));\n break;\n case DataType::ARRAY_UINT32:\n new_doc.set>(\n field->name(), std::vector(10, (uint32_t)doc_id));\n break;\n case DataType::ARRAY_UINT64:\n new_doc.set>(\n field->name(), std::vector(10, (uint64_t)doc_id));\n break;\n case DataType::ARRAY_FLOAT:\n new_doc.set>(field->name(),\n std::vector(10, (float)doc_id));\n break;\n case DataType::ARRAY_DOUBLE:\n new_doc.set>(\n field->name(), std::vector(10, (double)doc_id));\n break;\n case DataType::ARRAY_STRING:\n new_doc.set>(\n field->name(),\n std::vector(10, \"value_\" + std::to_string(doc_id)));\n break;\n case DataType::VECTOR_BINARY32:\n new_doc.set>(\n field->name(),\n std::vector(field->dimension(), uint32_t(doc_id + 0.1)));\n break;\n case DataType::VECTOR_BINARY64:\n new_doc.set>(\n field->name(),\n std::vector(field->dimension(), uint64_t(doc_id + 0.1)));\n break;\n case DataType::VECTOR_FP32:\n new_doc.set>(\n field->name(),\n std::vector(field->dimension(), float(doc_id + 0.1)));\n break;\n case DataType::VECTOR_FP64:\n new_doc.set>(\n field->name(),\n std::vector(field->dimension(), double(doc_id + 0.1)));\n break;\n case DataType::VECTOR_FP16:\n new_doc.set>(\n field->name(), std::vector(\n field->dimension(), static_cast(\n float(doc_id + 0.1))));\n break;\n case DataType::VECTOR_INT8:\n new_doc.set>(\n field->name(),\n std::vector(field->dimension(), (int8_t)doc_id));\n break;\n case DataType::VECTOR_INT16:\n new_doc.set>(\n field->name(),\n std::vector(field->dimension(), (int16_t)doc_id));\n break;\n case DataType::SPARSE_VECTOR_FP16: {\n std::vector indices;\n std::vector values;\n for (uint32_t i = 0; i < 100; i++) {\n indices.push_back(i);\n values.push_back(zvec::float16_t(float(doc_id + 0.1)));\n }\n std::pair, std::vector>\n sparse_float_vec;\n sparse_float_vec.first = indices;\n sparse_float_vec.second = values;\n new_doc.set<\n std::pair, std::vector>>(\n field->name(), sparse_float_vec);\n break;\n }\n case DataType::SPARSE_VECTOR_FP32: {\n std::vector indices;\n std::vector values;\n for (uint32_t i = 0; i < 100; i++) {\n indices.push_back(i);\n values.push_back(float(doc_id + 0.1));\n }\n std::pair, std::vector> sparse_float_vec;\n sparse_float_vec.first = indices;\n sparse_float_vec.second = values;\n new_doc.set, std::vector>>(\n field->name(), sparse_float_vec);\n break;\n }\n default:\n std::cout << \"Unsupported data type: \" << field->name() << std::endl;\n throw std::runtime_error(\"Unsupported vector data type\");\n }\n }\n\n return new_doc;\n}\n\nCollectionSchema::Ptr create_schema() {\n auto schema = std::make_shared(\"demo\");\n schema->set_max_doc_count_per_segment(1000);\n\n schema->add_field(std::make_shared(\n \"id\", DataType::INT64, false, std::make_shared(true)));\n schema->add_field(std::make_shared(\n \"name\", DataType::STRING, false,\n std::make_shared(false)));\n schema->add_field(\n std::make_shared(\"weight\", DataType::FLOAT, true));\n\n schema->add_field(std::make_shared(\n \"dense\", DataType::VECTOR_FP32, 128, false,\n std::make_shared(MetricType::IP)));\n schema->add_field(std::make_shared(\n \"sparse\", DataType::SPARSE_VECTOR_FP32, 0, false,\n std::make_shared(MetricType::IP)));\n\n return schema;\n}\n\nint main() {\n std::string path = \"./demo\";\n std::string rm_cmd = \"rm -rf \" + path;\n system(rm_cmd.c_str());\n\n auto schema = create_schema();\n CollectionOptions options{false, true};\n\n auto result = Collection::CreateAndOpen(path, *schema, options);\n if (!result.has_value()) {\n std::cout << result.error().message() << std::endl;\n return -1;\n }\n\n std::cout << \"init stats: \" << result.value()->Stats().value().to_string()\n << std::endl;\n\n auto coll = std::move(result).value();\n\n // insert docs\n {\n auto doc1 = create_doc(0, *schema);\n std::vector docs{doc1};\n auto res = coll->Insert(docs);\n if (!res.has_value()) {\n std::cout << res.error().message() << std::endl;\n return -1;\n }\n std::cout << \"after insert stats \" << coll->Stats().value().to_string()\n << std::endl;\n }\n\n // optimize\n {\n auto res = coll->Optimize();\n if (!res.ok()) {\n std::cout << res.message() << std::endl;\n return -1;\n }\n std::cout << \"after optimize stats \" << coll->Stats().value().to_string()\n << std::endl;\n }\n\n // query\n {\n VectorQuery query;\n query.topk_ = 10;\n query.field_name_ = \"dense\";\n query.include_vector_ = true;\n std::vector query_vector = std::vector(128, 0.1);\n query.query_vector_.assign((char *)query_vector.data(),\n query_vector.size() * sizeof(float));\n auto res = coll->Query(query);\n if (!res.has_value()) {\n std::cout << res.error().message() << std::endl;\n return -1;\n }\n std::cout << \"query result: doc_count[\" << res.value().size() << \"]\"\n << std::endl;\n std::cout << \"first doc: \" << res.value()[0]->to_detail_string()\n << std::endl;\n }\n\n // close and reopen\n coll.reset();\n options.read_only_ = true;\n result = Collection::Open(path, options);\n if (!result.has_value()) {\n std::cout << result.error().message() << std::endl;\n return -1;\n }\n std::cout << \"reopen stats: \" << result.value()->Stats().value().to_string()\n << std::endl;\n\n return 0;\n}" }, { "path": "pyproject.toml", "content": "######################################################################################################\n# Zvec: High-Performance Vector Database with PyBind11 & C++ Backend\n######################################################################################################\n[project]\nname = \"zvec\"\ndynamic = [\"version\"]\ndescription = \"A high-performance vector database engine with native C++ backend and Python bindings\"\nreadme = \"README.md\"\nlicense = { text = \"Apache-2.0\" }\nauthors = [\n { name = \"zvec\", email = \"zvec@alibaba-inc.com\" },\n]\nmaintainers = [\n { name = \"Zvec Core Team\", email = \"zvec@alibaba-inc.com\" },\n]\nrequires-python = \">=3.9\"\nclassifiers = [\n \"Development Status :: 3 - Alpha\",\n \"Intended Audience :: Developers\",\n \"Intended Audience :: Education\",\n \"Intended Audience :: Science/Research\",\n \"License :: OSI Approved :: Apache Software License\",\n \"Operating System :: POSIX :: Linux\",\n \"Operating System :: MacOS\",\n \"Programming Language :: C++\",\n \"Programming Language :: Python :: 3 :: Only\",\n \"Programming Language :: Python :: 3.10\",\n \"Programming Language :: Python :: 3.11\",\n \"Programming Language :: Python :: 3.12\",\n \"Programming Language :: Python :: 3.13\",\n \"Programming Language :: Python :: 3.14\",\n \"Topic :: Database\",\n \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n \"Topic :: Software Development :: Libraries :: Python Modules\",\n]\nkeywords = [\n \"vector-database\", \"ann\", \"nearest-neighbor\"\n]\ndependencies = [\n \"numpy >=1.23\",\n]\n\n[project.urls]\nHomepage = \"https://github.com/alibaba/zvec\"\nRepository = \"https://github.com/alibaba/zvec\"\n\"Bug Tracker\" = \"https://github.com/alibaba/zvec/issues\"\n\"Documentation\" = \"https://zvec.org\"\n\n[project.optional-dependencies]\ntest = [\n \"pytest >=8.0\",\n \"pytest-cov >=4.1\",\n \"pytest-mock >=3.12\",\n \"cibuildwheel == 3.4.0\",\n]\ndocs = [\n \"mkdocs >=1.5\",\n \"mkdocs-material >=9.5\",\n \"mkdocstrings[python] >=0.24\",\n]\ndev = [\n \"ruff >=0.4\",\n \"black >=24.0\",\n \"mypy >=1.8\",\n \"pre-commit >=3.6\",\n \"build >=1.0\",\n \"twine >=4.0\",\n \"numpy >=1.23\",\n # Inherit test deps\n \"pytest >=8.0\",\n \"pytest-cov >=4.1\",\n \"pytest-mock >=3.12\",\n \"cibuildwheel == 3.4.0\",\n # Inherit docs deps\n \"mkdocs >=1.5\",\n \"mkdocs-material >=9.5\",\n \"mkdocstrings[python] >=0.24\",\n \"pybind11-stubgen>=2.5.5\",\n \"pybind11 >=3.0\",\n]\n######################################################################################################\n# BUILD SYSTEM CONFIGURATION (scikit-build-core)\n######################################################################################################\n[build-system]\nrequires = [\n \"scikit-build-core >=0.11\",\n \"pybind11 >=3.0\",\n \"setuptools_scm>=8.0\",\n \"cmake>=3.26,<4.0\",\n \"ninja>=1.11\",\n]\nbuild-backend = \"scikit_build_core.build\"\n\n[tool.scikit-build]\n# Core settings\nminimum-version = \"0.11\"\nmetadata.version.provider = \"scikit_build_core.metadata.setuptools_scm\"\n\n# CMake configuration\ncmake.version = \">=3.26,<4.0\"\nninja.version = \">=1.11\"\ncmake.build-type = \"Release\"\ninstall.strip = true # Strip symbols in release builds to reduce wheel size\n\n# Build directory\nbuild-dir = \"build\"\n\n# Platform support\nwheel.expand-macos-universal-tags = true\nwheel.packages = [\"python/zvec\"]\n\n# Source distribution\nsdist.include = [\n \"README.md\",\n \"LICENSE\",\n \"pyproject.toml\",\n \"CMakeLists.txt\",\n \"src/**/*\",\n \"stub/zvec/**/*\",\n \"python/zvec/py.typed\",\n]\n\n# CMake defines (env-overridable)\n[tool.scikit-build.cmake.define]\nBUILD_TOOLS = \"OFF\"\nBUILD_PYTHON_BINDINGS = \"ON\"\n#CMAKE_VERBOSE_MAKEFILE = \"ON\"\n\n# Setuptools config for test pypi\n[tool.setuptools_scm]\nlocal_scheme = \"no-local-version\"\nversion_scheme = \"guess-next-dev\"\nfallback_version = \"0.2.1b1\"\n######################################################################################################\n# TESTING & QUALITY\n######################################################################################################\n[tool.pytest.ini_options]\nminversion = \"8.0\"\naddopts = [\n \"-ra\",\n \"--showlocals\",\n \"--strict-markers\",\n \"--strict-config\",\n \"--tb=short\",\n]\nxfail_strict = true\nlog_cli_level = \"INFO\"\nfilterwarnings = [\n \"error\",\n \"ignore::pytest.PytestCacheWarning\",\n # Ignore numpy deprecation warnings in tests (if any)\n \"ignore:.*numpy.*:DeprecationWarning\",\n]\ntestpaths = [\"python/tests\"]\nmarkers = [\n \"title: Custom marker for test title/description\",\n # \"slow: marks tests as slow\",\n]\n\n######################################################################################################\n# BUILD WHEEL\n######################################################################################################\n[tool.cibuildwheel]\nbuild = [\n \"cp310-*\",\n \"cp311-*\",\n \"cp312-*\",\n \"cp313-*\",\n \"cp314-*\",\n]\nbuild-frontend = \"build\"\ntest-requires = [\"pytest\", \"numpy\"]\ntest-command = \"cd {project} && pytest python/tests -v --tb=short\"\nbuild-verbosity = 1\n\n[tool.cibuildwheel.linux]\narchs = [\"auto\"]\nenvironment = { CMAKE_GENERATOR = \"Unix Makefiles\", CMAKE_BUILD_PARALLEL_LEVEL = \"16\" }\nmanylinux-x86_64-image = \"manylinux_2_28\"\nmanylinux-aarch64-image = \"manylinux_2_28\"\n# Skip 32-bit builds and musllinux\nskip = [\"*-manylinux_i686\", \"*-musllinux*\"]\n\n[tool.cibuildwheel.macos]\narchs = [\"arm64\"]\n# Inherits CMAKE_GENERATOR and CMAKE_BUILD_PARALLEL_LEVEL from [tool.cibuildwheel] won't work;\n# platform-level environment overrides the top-level entirely, so all vars must be listed here\nenvironment = { CMAKE_GENERATOR = \"Unix Makefiles\", CMAKE_BUILD_PARALLEL_LEVEL = \"16\", MACOSX_DEPLOYMENT_TARGET = \"11.0\" }\n######################################################################################################\n# CODE QUALITY & FORMATTING (Ruff)\n######################################################################################################\n[tool.ruff]\ntarget-version = \"py310\"\nline-length = 88\nexclude = [\n \"build/\",\n \"dist/\",\n \".git/\",\n \".venv/\",\n \"venv/\",\n \"thirdparty\",\n]\n\n[tool.ruff.lint]\nextend-select = [\n \"B\", # flake8-bugbear\n \"I\", # isort\n \"ARG\", # flake8-unused-arguments\n \"C4\", # flake8-comprehensions\n \"EM\", # flake8-errmsg\n \"ICN\", # flake8-import-conventions\n \"G\", # flake8-logging-format\n \"PGH\", # pygrep-hooks\n \"PIE\", # flake8-pie\n \"PL\", # pylint\n \"PT\", # flake8-pytest-style\n \"PTH\", # flake8-use-pathlib\n \"RET\", # flake8-return\n \"RUF\", # Ruff-specific\n \"SIM\", # flake8-simplify\n \"T20\", # flake8-print\n \"UP\", # pyupgrade\n \"YTT\", # flake8-2020\n \"EXE\", # flake8-executable\n \"NPY\", # NumPy-specific\n \"PD\", # pandas-vet\n]\nignore = [\n \"PLR0913\", # Too many arguments (common in bindings)\n \"PLR2004\", # Magic value used in comparison\n \"UP045\", \"UP007\", # Use list() instead of [] (breaks C++ init)\n \"EM101\", \"EM102\", # Exception messages as literals (ok in tests/utils)\n \"B008\", # Mutable default args (cautiously allowed in config)\n \"E731\", # Lambda assignment (used in callbacks)\n \"B019\", # `functools.lru_cache` on methods (handled manually)\n \"PLR0912\", # Too many branches\n \"PLC0105\", # Ignore contravariant\n \"RUF002\", # Ignore Unicode\n]\nfixable = [\"ALL\"]\nunfixable = []\n\n# Ignore all errors in docstrings\n[tool.ruff.lint.pydocstyle]\nconvention = \"google\" # or \"numpy\", \"pep257\"\nignore-decorators = [\"typing.overload\"]\n\n[tool.ruff.lint.flake8-type-checking]\n# Don't check code examples in docstrings\nquote-annotations = true\n\n[tool.ruff.lint.isort]\nrequired-imports = [\"from __future__ import annotations\"]\nknown-first-party = [\"zvec\"]\n\n[tool.ruff.lint.per-file-ignores]\n\"python/tests/**\" = [\"ALL\"]\n\"bench/core/**\" = [\"ALL\"]\n\"python/zvec/__init__.py\" = [\n \"F401\", # Unused import (for __all__)\n \"E402\", # Module level import not at top (C++ module init order)\n \"PLE0605\", # Invalid format for __all__\n \"RUF022\", # __all__ is not sorted\n]\n\"python/zvec/model/doc.py\" = [\n \"RUF023\", # Unused sort (for __slot__)\n]\n\"python/zvec/extension/**\" = [\n \"PLC0415\", # Import outside top-level (dynamic imports in _get_model)\n]\n\n[tool.ruff.format]\nindent-style = \"space\"\nquote-style = \"double\"\nline-ending = \"lf\"\nskip-magic-trailing-comma = false\n" }, { "path": "python/tests/detail/distance_helper.py", "content": "import logging\nimport math\nimport numpy as np\n\nfrom zvec import (\n MetricType,\n DataType,\n QuantizeType,\n Doc,\n CollectionSchema,\n FieldSchema,\n VectorSchema,\n)\n\nfrom typing import Dict\n\n\ndef is_float_equal(actual, expected, rel_tol=1e-5, abs_tol=1e-8):\n if actual is None and expected is None:\n return True\n return math.isclose(actual, expected, rel_tol=rel_tol, abs_tol=abs_tol)\n\n\ndef is_dense_vector_equal(vec1, vec2, rtol=1e-5, atol=1e-8):\n \"\"\"Compare two dense vectors with tolerance.\"\"\"\n return np.allclose(vec1, vec2, rtol=rtol, atol=atol)\n\n\ndef is_sparse_vector_equal(vec1, vec2, rtol=1e-5, atol=1e-8):\n \"\"\"Compare two sparse vectors with tolerance.\"\"\"\n # Check if they have the same keys\n if set(vec1.keys()) != set(vec2.keys()):\n return False\n\n # Check if all values are close\n for key in vec1:\n if not math.isclose(vec1[key], vec2[key], rel_tol=rtol, abs_tol=atol):\n return False\n\n return True\n\n\ndef is_float_array_equal(arr1, arr2, rtol=1e-5, atol=1e-8):\n \"\"\"Compare two float arrays with tolerance.\"\"\"\n return np.allclose(arr1, arr2, rtol=rtol, atol=atol)\n\n\ndef is_double_array_equal(arr1, arr2, rtol=1e-9, atol=1e-12):\n \"\"\"Compare two double arrays with tolerance.\"\"\"\n return np.allclose(arr1, arr2, rtol=rtol, atol=atol)\n\n\ndef is_int_array_equal(arr1, arr2):\n \"\"\"Compare two integer arrays with exact equality.\"\"\"\n return np.array_equal(arr1, arr2)\n\n\ndef cosine_distance_dense(\n vec1,\n vec2,\n dtype: DataType = DataType.VECTOR_FP32,\n quantize_type: QuantizeType = QuantizeType.UNDEFINED,\n):\n if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16:\n # More stable conversion to float16 to avoid numerical issues\n vec1 = [float(np.float16(a)) for a in vec1]\n vec2 = [float(np.float16(b)) for b in vec2]\n elif dtype == DataType.VECTOR_INT8:\n # For INT8 vectors, convert to integers for proper calculation\n vec1 = [\n int(round(min(max(val, -128), 127))) for val in vec1\n ] # Clamp to valid INT8 range\n vec2 = [\n int(round(min(max(val, -128), 127))) for val in vec2\n ] # Clamp to valid INT8 range\n\n dot_product = sum(a * b for a, b in zip(vec1, vec2))\n\n magnitude1 = math.sqrt(sum(a * a for a in vec1))\n magnitude2 = math.sqrt(sum(b * b for b in vec2))\n\n if magnitude1 == 0 or magnitude2 == 0:\n return 1.0 # Zero vector case - maximum distance\n\n cosine_similarity = dot_product / (magnitude1 * magnitude2)\n\n # Clamp to [-1, 1] range to handle floating-point precision errors\n cosine_similarity = max(-1.0, min(1.0, cosine_similarity))\n\n # For identical vectors (within floating point precision), ensure cosine distance is 0.0\n # This is especially important for low-precision types which have limited precision\n if (\n dtype == DataType.VECTOR_FP16\n or quantize_type == QuantizeType.FP16\n or dtype == DataType.VECTOR_INT8\n ):\n if (\n abs(cosine_similarity - 1.0) < 1e-3\n ): # Handle precision issues for low-precision types\n cosine_similarity = 1.0\n\n # Return cosine distance (1 - cosine similarity) to maintain compatibility\n # with system internal processing and existing test expectations\n return 1.0 - cosine_similarity\n\n\ndef dp_distance_dense(\n vec1,\n vec2,\n dtype: DataType = DataType.VECTOR_FP32,\n quantize_type: QuantizeType = QuantizeType.UNDEFINED,\n):\n if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16:\n # More stable computation to avoid numerical issues\n products = [\n float(np.float16(a)) * float(np.float16(b)) for a, b in zip(vec1, vec2)\n ]\n return sum(products)\n elif dtype == DataType.VECTOR_INT8:\n # For INT8 vectors, convert to integers for proper calculation\n products = [\n int(round(min(max(a, -128), 127))) * int(round(min(max(b, -128), 127)))\n for a, b in zip(vec1, vec2)\n ]\n return sum(products)\n return sum(a * b for a, b in zip(vec1, vec2))\n\n\ndef euclidean_distance_dense(\n vec1,\n vec2,\n dtype: DataType = DataType.VECTOR_FP32,\n quantize_type: QuantizeType = QuantizeType.UNDEFINED,\n):\n if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16:\n # Convert to float16 and compute squared differences safely\n # Use a more stable computation to avoid overflow\n squared_diffs = []\n for a, b in zip(vec1, vec2):\n diff = np.float16(a) - np.float16(b)\n squared_diff = float(diff) * float(\n diff\n ) # Convert to float for multiplication\n squared_diffs.append(squared_diff)\n squared_distance = sum(squared_diffs)\n elif dtype == DataType.VECTOR_INT8:\n # For INT8 vectors, convert to integers and handle potential scaling\n # INT8 values might be treated differently in the library implementation\n vec1_int = [\n int(round(min(max(val, -128), 127))) for val in vec1\n ] # Clamp to valid INT8 range\n vec2_int = [\n int(round(min(max(val, -128), 127))) for val in vec2\n ] # Clamp to valid INT8 range\n # Use float type to prevent overflow when summing large squared differences\n squared_distance = sum(float(a - b) ** 2 for a, b in zip(vec1_int, vec2_int))\n else:\n squared_distance = sum((a - b) ** 2 for a, b in zip(vec1, vec2))\n\n return squared_distance # Return squared distance for INT8\n\n\ndef distance_dense(\n vec1,\n vec2,\n metric: MetricType,\n data_type: DataType = DataType.VECTOR_FP32,\n quantize_type: QuantizeType = QuantizeType.UNDEFINED,\n):\n if metric == MetricType.COSINE:\n return cosine_distance_dense(vec1, vec2, data_type, quantize_type)\n elif metric == MetricType.L2:\n return euclidean_distance_dense(vec1, vec2, data_type, quantize_type)\n elif metric == MetricType.IP:\n return dp_distance_dense(vec1, vec2, data_type, quantize_type)\n else:\n raise ValueError(\"Unsupported metric type\")\n\n\ndef dp_distance_sparse(\n vec1,\n vec2,\n data_type: DataType = DataType.SPARSE_VECTOR_FP32,\n quantize_type: QuantizeType = QuantizeType.UNDEFINED,\n):\n dot_product = 0.0\n for dim in set(vec1.keys()) & set(vec2.keys()):\n print(\"dim,vec1,vec2:\\n\")\n print(dim, vec1, vec2)\n if (\n data_type == DataType.SPARSE_VECTOR_FP16\n or quantize_type == QuantizeType.FP16\n ):\n vec1[dim] = np.float16(vec1[dim])\n vec2[dim] = np.float16(vec2[dim])\n dot_product += vec1[dim] * vec2[dim]\n return dot_product\n\n\ndef distance(\n vec1,\n vec2,\n metric: MetricType,\n data_type: DataType,\n quantize_type: QuantizeType = QuantizeType.UNDEFINED,\n):\n is_sparse = (\n data_type == DataType.SPARSE_VECTOR_FP32\n or data_type == DataType.SPARSE_VECTOR_FP16\n )\n\n if is_sparse:\n if metric != MetricType.IP:\n raise ValueError(\"Unsupported metric type for sparse vectors\")\n\n if is_sparse:\n return dp_distance_sparse(vec1, vec2, data_type, quantize_type)\n else:\n return distance_dense(vec1, vec2, metric, data_type, quantize_type)\n\n\ndef distance_recall(\n vec1,\n vec2,\n metric: MetricType,\n data_type: DataType,\n quantize_type: QuantizeType = QuantizeType.UNDEFINED,\n):\n is_sparse = (\n data_type == DataType.SPARSE_VECTOR_FP32\n or data_type == DataType.SPARSE_VECTOR_FP16\n )\n\n if is_sparse:\n return dp_distance_sparse(vec1, vec2, data_type, quantize_type)\n else:\n if data_type in [DataType.VECTOR_FP32, DataType.VECTOR_FP16]:\n return distance_dense(vec1, vec2, metric, data_type, quantize_type)\n elif data_type in [DataType.VECTOR_INT8] and metric in [\n MetricType.L2,\n MetricType.IP,\n ]:\n return distance_dense(vec1, vec2, metric, data_type, quantize_type)\n else:\n return dp_distance_dense(vec1, vec2, data_type, quantize_type)\n\n\ndef calculate_rrf_score(rank, k=60):\n return 1.0 / (k + rank + 1)\n\n\ndef calculate_multi_vector_rrf_scores(query_results: Dict[str, Doc], k=60):\n rrf_scores = {}\n\n for vector_name, docs in query_results.items():\n for rank, doc in enumerate(docs):\n doc_id = doc.id\n rrf_score = calculate_rrf_score(rank, k)\n if doc_id in rrf_scores:\n rrf_scores[doc_id] += rrf_score\n else:\n rrf_scores[doc_id] = rrf_score\n\n return rrf_scores\n\n\ndef calculate_multi_vector_weighted_scores(\n query_results: Dict[str, Doc], weights: Dict[str, float], metric: MetricType\n):\n def _normalize_score(score: float, metric: MetricType) -> float:\n if metric == MetricType.L2:\n return 1.0 - 2 * math.atan(score) / math.pi\n if metric == MetricType.IP:\n return 0.5 + math.atan(score) / math.pi\n if metric == MetricType.COSINE:\n return 1.0 - score / 2.0\n raise ValueError(\"Unsupported metric type\")\n\n weighted_scores = {}\n\n for vector_name, docs in query_results.items():\n weight = weights.get(vector_name, 1.0)\n\n for doc in docs:\n doc_id = doc.id\n weighted_score = (_normalize_score(doc.score, metric)) * weight\n if doc_id in weighted_scores:\n weighted_scores[doc_id] += weighted_score\n else:\n weighted_scores[doc_id] = weighted_score\n\n return weighted_scores\n\n\ndef is_field_equal(field1, field2, schema: FieldSchema) -> bool:\n if field1 is None and field2 is None:\n return True\n if field1 is None or field2 is None:\n return False\n\n if schema.data_type == DataType.ARRAY_FLOAT:\n return is_float_array_equal(field1, field2)\n elif schema.data_type == DataType.ARRAY_DOUBLE:\n return is_double_array_equal(field1, field2)\n elif schema.data_type in [\n DataType.ARRAY_INT32,\n DataType.ARRAY_INT64,\n DataType.ARRAY_BOOL,\n DataType.ARRAY_STRING,\n DataType.ARRAY_UINT32,\n DataType.ARRAY_UINT64,\n DataType.ARRAY_INT64,\n ]:\n return is_int_array_equal(field1, field2)\n elif schema.data_type in [DataType.FLOAT, DataType.DOUBLE]:\n return is_float_equal(field1, field2)\n\n return field1 == field2\n\n\ndef is_vector_equal(vec1, vec2, schema: VectorSchema) -> bool:\n if (\n schema.data_type == DataType.SPARSE_VECTOR_FP16\n or schema.data_type == DataType.VECTOR_FP16\n ):\n # skip fp16 vector equal\n return True\n\n is_sparse = (\n schema.data_type == DataType.SPARSE_VECTOR_FP32\n or schema.data_type == DataType.SPARSE_VECTOR_FP16\n )\n\n if is_sparse:\n return is_sparse_vector_equal(vec1, vec2)\n else:\n return is_dense_vector_equal(vec1, vec2)\n\n\ndef is_doc_equal(\n doc1: Doc,\n doc2: Doc,\n schema: CollectionSchema,\n except_score: bool = True,\n include_vector: bool = True,\n):\n if doc1.id != doc2.id:\n logging.error(\"doc ids are not equal\")\n return False\n\n reduce_field_names = set(doc1.field_names() + doc2.field_names())\n reduce_vector_names = set(doc1.vector_names() + doc2.vector_names())\n\n is_doc1_fields_empty = doc1.fields is None or doc1.fields == {}\n is_doc2_fields_empty = doc2.fields is None or doc2.fields == {}\n\n if is_doc1_fields_empty or is_doc2_fields_empty:\n if is_doc1_fields_empty != is_doc2_fields_empty:\n return False\n else:\n for field_name in reduce_field_names:\n field_schema = schema.field(field_name)\n if field_schema is None:\n return False\n if is_field_equal(\n doc1.field(field_name), doc2.field(field_name), field_schema\n ):\n continue\n else:\n logging.error(f\"{field_name} are not equal\")\n return False\n\n if include_vector:\n is_doc1_vectors_empty = doc1.vectors is None or doc1.vectors == {}\n is_doc2_vectors_empty = doc2.vectors is None or doc2.vectors == {}\n\n if is_doc1_vectors_empty or is_doc2_vectors_empty:\n if is_doc1_fields_empty != is_doc2_vectors_empty:\n return False\n else:\n for vector_name in reduce_vector_names:\n vector_schema = schema.vector(vector_name)\n if vector_schema is None:\n return False\n if is_vector_equal(\n doc1.vector(vector_name), doc2.vector(vector_name), vector_schema\n ):\n continue\n else:\n return False\n\n return True\n" }, { "path": "python/tests/detail/doc_helper.py", "content": "from zvec import CollectionSchema, Doc\n\nfrom support_helper import *\n\nimport numpy as np\nfrom typing import Literal, Optional, Union, Tuple\n\nimport random\nimport string\nimport math\n\n\ndef generate_constant_vector(\n i: int, dimension: int, dtype: Literal[\"int8\", \"float16\", \"float32\"] = \"float32\"\n):\n if dtype == \"int8\":\n vec = [(i % 127)] * dimension\n vec[i % dimension] = (i + 1) % 127\n else:\n base_val = (i % 1000) / 256.0\n special_val = ((i + 1) % 1000) / 256.0\n vec = [base_val] * dimension\n vec[i % dimension] = special_val\n\n return vec\n\n\ndef generate_constant_vector_recall(\n i: int, dimension: int, dtype: Literal[\"int8\", \"float16\", \"float32\"] = \"float32\"\n):\n if dtype == \"int8\":\n vec = [(i % 127)] * dimension\n vec[i % dimension] = (i + 1) % 127\n else:\n base_val = math.sin((i) * 1000) / 256.0\n special_val = math.sin((i + 1) * 1000) / 256.0\n vec = [base_val] * dimension\n vec[i % dimension] = special_val\n\n return vec\n\n\ndef generate_sparse_vector(i: int):\n return {i: i + 0.1}\n\n\ndef generate_vectordict(i: int, schema: CollectionSchema) -> Doc:\n doc_fields = {}\n doc_vectors = {}\n doc_fields = {}\n doc_vectors = {}\n for field in schema.fields:\n if field.data_type == DataType.BOOL:\n doc_fields[field.name] = i % 2 == 0\n elif field.data_type == DataType.INT32:\n doc_fields[field.name] = i\n elif field.data_type == DataType.UINT32:\n doc_fields[field.name] = i\n elif field.data_type == DataType.INT64:\n doc_fields[field.name] = i\n elif field.data_type == DataType.UINT64:\n doc_fields[field.name] = i\n elif field.data_type == DataType.FLOAT:\n doc_fields[field.name] = float(i) + 0.1\n elif field.data_type == DataType.DOUBLE:\n doc_fields[field.name] = float(i) + 0.11\n elif field.data_type == DataType.STRING:\n doc_fields[field.name] = f\"test_{i}\"\n elif field.data_type == DataType.ARRAY_BOOL:\n doc_fields[field.name] = [i % 2 == 0, i % 3 == 0]\n elif field.data_type == DataType.ARRAY_INT32:\n doc_fields[field.name] = [i, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_UINT32:\n doc_fields[field.name] = [i, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_INT64:\n doc_fields[field.name] = [i, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_UINT64:\n doc_fields[field.name] = [i, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_FLOAT:\n doc_fields[field.name] = [float(i + 0.1), float(i + 1.1), float(i + 2.1)]\n elif field.data_type == DataType.ARRAY_DOUBLE:\n doc_fields[field.name] = [float(i + 0.11), float(i + 1.11), float(i + 2.11)]\n elif field.data_type == DataType.ARRAY_STRING:\n doc_fields[field.name] = [f\"test_{i}\", f\"test_{i + 1}\", f\"test_{i + 2}\"]\n else:\n raise ValueError(f\"Unsupported field type: {field.data_type}\")\n for vector in schema.vectors:\n if vector.data_type == DataType.VECTOR_FP16:\n doc_vectors[vector.name] = generate_constant_vector(\n i, vector.dimension, \"float16\"\n )\n elif vector.data_type == DataType.VECTOR_FP32:\n doc_vectors[vector.name] = generate_constant_vector(\n i, vector.dimension, \"float32\"\n )\n elif vector.data_type == DataType.VECTOR_INT8:\n doc_vectors[vector.name] = generate_constant_vector(\n i,\n vector.dimension,\n \"int8\",\n )\n elif vector.data_type == DataType.SPARSE_VECTOR_FP32:\n doc_vectors[vector.name] = generate_sparse_vector(i)\n elif vector.data_type == DataType.SPARSE_VECTOR_FP16:\n doc_vectors[vector.name] = generate_sparse_vector(i)\n else:\n raise ValueError(f\"Unsupported vector type: {vector.data_type}\")\n return doc_fields, doc_vectors\n\n\ndef generate_vectordict_recall(i: int, schema: CollectionSchema) -> Doc:\n doc_fields = {}\n doc_vectors = {}\n doc_fields = {}\n doc_vectors = {}\n for field in schema.fields:\n if field.data_type == DataType.BOOL:\n doc_fields[field.name] = i % 2 == 0\n elif field.data_type == DataType.INT32:\n doc_fields[field.name] = i\n elif field.data_type == DataType.UINT32:\n doc_fields[field.name] = i\n elif field.data_type == DataType.INT64:\n doc_fields[field.name] = i\n elif field.data_type == DataType.UINT64:\n doc_fields[field.name] = i\n elif field.data_type == DataType.FLOAT:\n doc_fields[field.name] = float(i) + 0.1\n elif field.data_type == DataType.DOUBLE:\n doc_fields[field.name] = float(i) + 0.11\n elif field.data_type == DataType.STRING:\n doc_fields[field.name] = f\"test_{i}\"\n elif field.data_type == DataType.ARRAY_BOOL:\n doc_fields[field.name] = [i % 2 == 0, i % 3 == 0]\n elif field.data_type == DataType.ARRAY_INT32:\n doc_fields[field.name] = [i, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_UINT32:\n doc_fields[field.name] = [i, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_INT64:\n doc_fields[field.name] = [i, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_UINT64:\n doc_fields[field.name] = [i, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_FLOAT:\n doc_fields[field.name] = [float(i + 0.1), float(i + 1.1), float(i + 2.1)]\n elif field.data_type == DataType.ARRAY_DOUBLE:\n doc_fields[field.name] = [float(i + 0.11), float(i + 1.11), float(i + 2.11)]\n elif field.data_type == DataType.ARRAY_STRING:\n doc_fields[field.name] = [f\"test_{i}\", f\"test_{i + 1}\", f\"test_{i + 2}\"]\n else:\n raise ValueError(f\"Unsupported field type: {field.data_type}\")\n for vector in schema.vectors:\n if vector.data_type == DataType.VECTOR_FP16:\n doc_vectors[vector.name] = generate_constant_vector_recall(\n i, vector.dimension, \"float16\"\n )\n elif vector.data_type == DataType.VECTOR_FP32:\n doc_vectors[vector.name] = generate_constant_vector_recall(\n i, vector.dimension, \"float32\"\n )\n elif vector.data_type == DataType.VECTOR_INT8:\n doc_vectors[vector.name] = generate_constant_vector_recall(\n i,\n vector.dimension,\n \"int8\",\n )\n elif vector.data_type == DataType.SPARSE_VECTOR_FP32:\n doc_vectors[vector.name] = generate_sparse_vector(i)\n elif vector.data_type == DataType.SPARSE_VECTOR_FP16:\n doc_vectors[vector.name] = generate_sparse_vector(i)\n else:\n raise ValueError(f\"Unsupported vector type: {vector.data_type}\")\n return doc_fields, doc_vectors\n\n\ndef generate_vectordict_update(i: int, schema: CollectionSchema) -> Doc:\n doc_fields = {}\n doc_vectors = {}\n doc_fields = {}\n doc_vectors = {}\n for field in schema.fields:\n if field.data_type == DataType.BOOL:\n doc_fields[field.name] = (i + 1) % 2 == 0\n elif field.data_type == DataType.INT32:\n doc_fields[field.name] = i + 1\n elif field.data_type == DataType.UINT32:\n doc_fields[field.name] = i + 1\n elif field.data_type == DataType.INT64:\n doc_fields[field.name] = i + 1\n elif field.data_type == DataType.UINT64:\n doc_fields[field.name] = i + 1\n elif field.data_type == DataType.FLOAT:\n doc_fields[field.name] = float(i + 1) + 0.1\n elif field.data_type == DataType.DOUBLE:\n doc_fields[field.name] = float(i + 1) + 0.11\n elif field.data_type == DataType.STRING:\n doc_fields[field.name] = f\"test_{i + 1}\"\n elif field.data_type == DataType.ARRAY_BOOL:\n doc_fields[field.name] = [(i + 1) % 2 == 0, (i + 1) % 3 == 0]\n elif field.data_type == DataType.ARRAY_INT32:\n doc_fields[field.name] = [i + 1, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_UINT32:\n doc_fields[field.name] = [i + 1, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_INT64:\n doc_fields[field.name] = [i + 1, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_UINT64:\n doc_fields[field.name] = [i + 1, i + 1, i + 2]\n elif field.data_type == DataType.ARRAY_FLOAT:\n doc_fields[field.name] = [float(i + 1.1), float(i + 2.1), float(i + 3.1)]\n elif field.data_type == DataType.ARRAY_DOUBLE:\n doc_fields[field.name] = [float(i + 1.11), float(i + 2.11), float(i + 3.11)]\n elif field.data_type == DataType.ARRAY_STRING:\n doc_fields[field.name] = [f\"test_{i + 1}\", f\"test_{i + 2}\", f\"test_{i + 3}\"]\n else:\n raise ValueError(f\"Unsupported field type: {field.data_type}\")\n for vector in schema.vectors:\n if vector.data_type == DataType.VECTOR_FP16:\n doc_vectors[vector.name] = generate_constant_vector(\n i + 1, vector.dimension, \"float16\"\n )\n elif vector.data_type == DataType.VECTOR_FP32:\n doc_vectors[vector.name] = generate_constant_vector(\n i + 1, vector.dimension, \"float32\"\n )\n elif vector.data_type == DataType.VECTOR_INT8:\n doc_vectors[vector.name] = generate_constant_vector(\n i + 1,\n vector.dimension,\n \"int8\",\n )\n elif vector.data_type == DataType.SPARSE_VECTOR_FP32:\n doc_vectors[vector.name] = generate_sparse_vector(i + 1)\n elif vector.data_type == DataType.SPARSE_VECTOR_FP16:\n doc_vectors[vector.name] = generate_sparse_vector(i + 1)\n else:\n raise ValueError(f\"Unsupported vector type: {vector.data_type}\")\n return doc_fields, doc_vectors\n\n\ndef generate_doc(i: int, schema: CollectionSchema) -> Doc:\n doc_fields = {}\n doc_vectors = {}\n doc_fields, doc_vectors = generate_vectordict(i, schema)\n doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)\n return doc\n\n\ndef generate_doc_recall(i: int, schema: CollectionSchema) -> Doc:\n doc_fields = {}\n doc_vectors = {}\n doc_fields, doc_vectors = generate_vectordict_recall(i, schema)\n doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)\n return doc\n\n\ndef generate_update_doc(i: int, schema: CollectionSchema) -> Doc:\n doc_fields = {}\n doc_vectors = {}\n doc_fields, doc_vectors = generate_vectordict_update(i, schema)\n doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)\n return doc\n\n\ndef generate_doc_random(i, schema: CollectionSchema) -> Doc:\n doc_fields = {}\n doc_vectors = {}\n\n random.seed(i)\n\n for field in schema.fields:\n if field.data_type == DataType.BOOL:\n doc_fields[field.name] = random.choice([True, False])\n elif field.data_type == DataType.INT32:\n doc_fields[field.name] = random.randint(-2147483648, 2147483647)\n elif field.data_type == DataType.UINT32:\n doc_fields[field.name] = random.randint(0, 4294967295)\n elif field.data_type == DataType.INT64:\n doc_fields[field.name] = random.randint(\n -9223372036854775808, 9223372036854775807\n )\n elif field.data_type == DataType.UINT64:\n doc_fields[field.name] = random.randint(0, 18446744073709551615)\n elif field.data_type == DataType.FLOAT:\n doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38)\n elif field.data_type == DataType.DOUBLE:\n doc_fields[field.name] = random.uniform(\n -1.7976931348623157e308, 1.7976931348623157e308\n )\n elif field.data_type == DataType.STRING:\n length = random.randint(1, 999)\n doc_fields[field.name] = \"\".join(\n random.choices(string.ascii_letters + string.digits, k=length)\n )\n elif field.data_type == DataType.ARRAY_BOOL:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.choice([True, False]) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_INT32:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.randint(-2147483648, 2147483647) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_UINT32:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.randint(0, 4294967295) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_INT64:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.randint(-9223372036854775808, 9223372036854775807)\n for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_UINT64:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.randint(0, 18446744073709551615) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_FLOAT:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_DOUBLE:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.uniform(-1.7976931348623157e308, 1.7976931348623157e308)\n for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_STRING:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n \"\".join(\n random.choices(\n string.ascii_letters + string.digits, k=random.randint(1, 100)\n )\n )\n for _ in range(array_length)\n ]\n else:\n raise ValueError(f\"Unsupported field type: {field.data_type}\")\n\n for vector in schema.vectors:\n if vector.data_type == DataType.VECTOR_FP16:\n doc_vectors[vector.name] = generate_constant_vector(\n random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, \"float16\"\n )\n elif vector.data_type == DataType.VECTOR_FP32:\n doc_vectors[vector.name] = generate_constant_vector(\n random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, \"float32\"\n )\n elif vector.data_type == DataType.VECTOR_INT8:\n doc_vectors[vector.name] = generate_constant_vector(\n random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, \"int8\"\n )\n elif vector.data_type == DataType.SPARSE_VECTOR_FP32:\n doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))\n elif vector.data_type == DataType.SPARSE_VECTOR_FP16:\n doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))\n else:\n raise ValueError(f\"Unsupported vector type: {vector.data_type}\")\n\n doc = Doc(id=i, fields=doc_fields, vectors=doc_vectors)\n return doc\n\n\ndef generate_vectordict_random(schema: CollectionSchema):\n doc_fields = {}\n doc_vectors = {}\n for field in schema.fields:\n if field.data_type == DataType.BOOL:\n doc_fields[field.name] = random.choice([True, False])\n elif field.data_type == DataType.INT32:\n doc_fields[field.name] = random.randint(-2147483648, 2147483647)\n elif field.data_type == DataType.UINT32:\n doc_fields[field.name] = random.randint(0, 4294967295)\n elif field.data_type == DataType.INT64:\n doc_fields[field.name] = random.randint(\n -9223372036854775808, 9223372036854775807\n )\n elif field.data_type == DataType.UINT64:\n doc_fields[field.name] = random.randint(0, 18446744073709551615)\n elif field.data_type == DataType.FLOAT:\n doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38)\n elif field.data_type == DataType.DOUBLE:\n doc_fields[field.name] = random.uniform(\n -1.7976931348623157e308, 1.7976931348623157e308\n )\n elif field.data_type == DataType.STRING:\n length = random.randint(1, 999)\n doc_fields[field.name] = \"\".join(\n random.choices(string.ascii_letters + string.digits, k=length)\n )\n elif field.data_type == DataType.ARRAY_BOOL:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.choice([True, False]) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_INT32:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.randint(-2147483648, 2147483647) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_UINT32:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.randint(0, 4294967295) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_INT64:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.randint(-9223372036854775808, 9223372036854775807)\n for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_UINT64:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.randint(0, 18446744073709551615) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_FLOAT:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_DOUBLE:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n random.uniform(-1.7976931348623157e308, 1.7976931348623157e308)\n for _ in range(array_length)\n ]\n elif field.data_type == DataType.ARRAY_STRING:\n array_length = random.randint(0, 10)\n doc_fields[field.name] = [\n \"\".join(\n random.choices(\n string.ascii_letters + string.digits, k=random.randint(1, 100)\n )\n )\n for _ in range(array_length)\n ]\n else:\n raise ValueError(f\"Unsupported field type: {field.data_type}\")\n\n for vector in schema.vectors:\n if vector.data_type == DataType.VECTOR_FP16:\n doc_vectors[vector.name] = generate_constant_vector(\n random.randint(1, 100), vector.dimension, \"float16\"\n )\n elif vector.data_type == DataType.VECTOR_FP32:\n doc_vectors[vector.name] = generate_constant_vector(\n random.randint(1, 100), vector.dimension, \"float32\"\n )\n elif vector.data_type == DataType.VECTOR_INT8:\n doc_vectors[vector.name] = generate_constant_vector(\n random.randint(1, 100), vector.dimension, \"int8\"\n )\n elif vector.data_type == DataType.SPARSE_VECTOR_FP32:\n doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))\n elif vector.data_type == DataType.SPARSE_VECTOR_FP16:\n doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))\n else:\n raise ValueError(f\"Unsupported vector type: {vector.data_type}\")\n\n return doc_fields, doc_vectors\n" }, { "path": "python/tests/detail/fixture_helper.py", "content": "import pytest\nimport logging\n\nfrom typing import Any, Generator\nfrom zvec.typing import DataType, StatusCode, MetricType, QuantizeType\nimport zvec\nfrom zvec import (\n CollectionOption,\n InvertIndexParam,\n HnswIndexParam,\n FlatIndexParam,\n IVFIndexParam,\n FieldSchema,\n VectorSchema,\n CollectionSchema,\n Collection,\n Doc,\n VectorQuery,\n)\n\nfrom support_helper import *\n\n\n@pytest.fixture(scope=\"session\")\ndef basic_schema(collection_name=\"test_collection\") -> CollectionSchema:\n return CollectionSchema(\n name=collection_name if len(collection_name) > 0 else \"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\", DataType.STRING, nullable=False, index_param=InvertIndexParam()\n ),\n FieldSchema(\"weight\", DataType.FLOAT, nullable=True),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n VectorSchema(\n \"sparse\", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()\n ),\n ],\n )\n\n\n@pytest.fixture(scope=\"session\")\ndef full_schema(\n nullable: bool = False,\n has_index: bool = False,\n) -> CollectionSchema:\n scalar_index_param = None\n vector_index_param = None\n if has_index:\n scalar_index_param = InvertIndexParam(enable_range_optimization=True)\n vector_index_param = HnswIndexParam()\n\n fields = []\n for k, v in DEFAULT_SCALAR_FIELD_NAME.items():\n fields.append(\n FieldSchema(\n v,\n k,\n nullable=nullable,\n index_param=scalar_index_param,\n )\n )\n vetors = []\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n vetors.append(\n VectorSchema(\n v,\n k,\n dimension=DEFAULT_VECTOR_DIMENSION,\n index_param=vector_index_param,\n )\n )\n\n return CollectionSchema(\n name=\"full_collection\",\n fields=fields,\n vectors=vetors,\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef full_schema_new(request) -> CollectionSchema:\n if hasattr(request, \"param\"):\n nullable, has_index, vector_index = request.param\n else:\n nullable, has_index, vector_index = True, False, HnswIndexParam()\n\n scalar_index_param = None\n vector_index_param = None\n if has_index:\n scalar_index_param = InvertIndexParam(enable_range_optimization=True)\n vector_index_param = vector_index\n\n fields = []\n for k, v in DEFAULT_SCALAR_FIELD_NAME.items():\n fields.append(\n FieldSchema(\n v,\n k,\n nullable=nullable,\n index_param=scalar_index_param,\n )\n )\n vectors = []\n\n if vector_index_param in [\n HnswIndexParam(),\n FlatIndexParam(),\n HnswIndexParam(\n metric_type=MetricType.IP,\n m=16,\n ef_construction=100,\n ),\n FlatIndexParam(\n metric_type=MetricType.IP,\n ),\n ]:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=DEFAULT_VECTOR_DIMENSION,\n index_param=vector_index_param,\n )\n )\n elif vector_index_param in [\n IVFIndexParam(),\n IVFIndexParam(\n metric_type=MetricType.IP,\n n_list=100,\n n_iters=10,\n use_soar=False,\n ),\n IVFIndexParam(\n metric_type=MetricType.L2,\n n_list=200,\n n_iters=20,\n use_soar=True,\n ),\n (\n IVFIndexParam(\n metric_type=MetricType.COSINE,\n n_list=150,\n n_iters=15,\n use_soar=False,\n )\n ),\n (\n HnswIndexParam(\n metric_type=MetricType.COSINE,\n m=24,\n ef_construction=150,\n )\n ),\n (\n HnswIndexParam(\n metric_type=MetricType.L2,\n m=32,\n ef_construction=200,\n )\n ),\n (\n FlatIndexParam(\n metric_type=MetricType.COSINE,\n )\n ),\n (\n FlatIndexParam(\n metric_type=MetricType.L2,\n )\n ),\n ]:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v in [\"vector_fp16_field\", \"vector_fp32_field\"]:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=DEFAULT_VECTOR_DIMENSION,\n index_param=vector_index_param,\n )\n )\n elif v in [\"vector_int8_field\"] and vector_index_param in [\n IVFIndexParam(\n metric_type=MetricType.L2,\n n_list=200,\n n_iters=20,\n use_soar=True,\n ),\n (\n HnswIndexParam(\n metric_type=MetricType.L2,\n m=32,\n ef_construction=200,\n )\n ),\n (\n FlatIndexParam(\n metric_type=MetricType.L2,\n )\n ),\n ]:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=DEFAULT_VECTOR_DIMENSION,\n index_param=vector_index_param,\n )\n )\n else:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=DEFAULT_VECTOR_DIMENSION,\n index_param=HnswIndexParam(),\n )\n )\n else:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v in [\"vector_fp16_field\", \"vector_fp32_field\"]:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=DEFAULT_VECTOR_DIMENSION,\n index_param=vector_index_param,\n )\n )\n else:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=DEFAULT_VECTOR_DIMENSION,\n index_param=HnswIndexParam(),\n )\n )\n\n return CollectionSchema(\n name=\"full_collection_new\",\n fields=fields,\n vectors=vectors,\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef full_schema_ivf(request) -> CollectionSchema:\n if hasattr(request, \"param\"):\n nullable, has_index, vector_index = request.param\n else:\n nullable, has_index, vector_index = True, False, IVFIndexParam()\n\n scalar_index_param = None\n vector_index_param = None\n if has_index:\n scalar_index_param = InvertIndexParam(enable_range_optimization=True)\n vector_index_param = vector_index\n\n fields = []\n for k, v in DEFAULT_SCALAR_FIELD_NAME.items():\n fields.append(\n FieldSchema(\n v,\n k,\n nullable=nullable,\n index_param=scalar_index_param,\n )\n )\n vectors = []\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v in [\"vector_fp16_field\", \"vector_fp32_field\"]:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=DEFAULT_VECTOR_DIMENSION,\n index_param=vector_index_param,\n )\n )\n\n return CollectionSchema(\n name=\"full_collection_ivf\",\n fields=fields,\n vectors=vectors,\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef full_schema_1024(request) -> CollectionSchema:\n if hasattr(request, \"param\"):\n nullable, has_index, vector_index = request.param\n else:\n nullable, has_index, vector_index = True, False, HnswIndexParam()\n\n scalar_index_param = None\n vector_index_param = None\n if has_index:\n scalar_index_param = InvertIndexParam(enable_range_optimization=True)\n vector_index_param = vector_index\n\n fields = []\n for k, v in DEFAULT_SCALAR_FIELD_NAME.items():\n fields.append(\n FieldSchema(\n v,\n k,\n nullable=nullable,\n index_param=scalar_index_param,\n )\n )\n vectors = []\n\n if vector_index_param in [\n HnswIndexParam(),\n FlatIndexParam(),\n HnswIndexParam(\n metric_type=MetricType.IP,\n m=16,\n ef_construction=100,\n ),\n FlatIndexParam(\n metric_type=MetricType.IP,\n ),\n ]:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=VECTOR_DIMENSION_1024,\n index_param=vector_index_param,\n )\n )\n elif vector_index_param in [\n IVFIndexParam(),\n IVFIndexParam(\n metric_type=MetricType.IP,\n n_list=100,\n n_iters=10,\n use_soar=False,\n ),\n IVFIndexParam(\n metric_type=MetricType.L2,\n n_list=200,\n n_iters=20,\n use_soar=True,\n ),\n IVFIndexParam(\n metric_type=MetricType.COSINE,\n n_list=150,\n n_iters=15,\n use_soar=False,\n ),\n ]:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v in [\"vector_fp16_field\", \"vector_fp32_field\"]:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=VECTOR_DIMENSION_1024,\n index_param=vector_index_param,\n )\n )\n elif v in [\"vector_int8_field\"] and vector_index_param in [\n IVFIndexParam(\n metric_type=MetricType.L2,\n n_list=200,\n n_iters=20,\n use_soar=True,\n ),\n IVFIndexParam(\n metric_type=MetricType.COSINE,\n n_list=150,\n n_iters=15,\n use_soar=False,\n ),\n ]:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=DVECTOR_DIMENSION_1024,\n index_param=vector_index_param,\n )\n )\n else:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=VECTOR_DIMENSION_1024,\n index_param=HnswIndexParam(),\n )\n )\n else:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v in [\"vector_fp16_field\", \"vector_fp32_field\", \"vector_int8_field\"]:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=VECTOR_DIMENSION_1024,\n index_param=vector_index_param,\n )\n )\n else:\n vectors.append(\n VectorSchema(\n v,\n k,\n dimension=VECTOR_DIMENSION_1024,\n index_param=HnswIndexParam(),\n )\n )\n\n return CollectionSchema(\n name=\"full_collection_new\",\n fields=fields,\n vectors=vectors,\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef single_vector_schema(\n data_type: DataType,\n) -> CollectionSchema:\n vector_schema = [\n VectorSchema(\n DEFAULT_VECTOR_FIELD_NAME[data_type],\n data_type,\n DEFAULT_VECTOR_DIMENSION,\n )\n ]\n\n return CollectionSchema(\n name=\"full_collection\",\n vectors=vector_schema,\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef single_vector_schema_with_index_param(\n data_type: DataType, index_param\n) -> CollectionSchema:\n vector_schema = [\n VectorSchema(\n DEFAULT_VECTOR_FIELD_NAME[data_type],\n data_type,\n DEFAULT_VECTOR_DIMENSION,\n index_param,\n )\n ]\n\n return CollectionSchema(\n name=\"full_collection\",\n vectors=vector_schema,\n )\n\n\ndef create_collection_fixture(\n collection_temp_dir, schema: CollectionSchema, collection_option: CollectionOption\n) -> Generator[Any, Any, Collection]:\n \"\"\"Common helper function to create and manage collection fixtures.\"\"\"\n coll = zvec.create_and_open(\n path=str(collection_temp_dir),\n schema=schema,\n option=collection_option,\n )\n\n assert coll is not None, \"Failed to create and open collection\"\n assert coll.path == str(collection_temp_dir)\n assert coll.schema.name == schema.name\n assert list(coll.schema.fields) == list(schema.fields)\n assert list(coll.schema.vectors) == list(schema.vectors)\n assert coll.option.read_only == collection_option.read_only\n assert coll.option.enable_mmap == collection_option.enable_mmap\n\n try:\n yield coll\n finally:\n if hasattr(coll, \"destroy\") and coll is not None:\n try:\n coll.destroy()\n except Exception as e:\n logging.warning(f\"Warning: failed to destroy collection: {e}\")\n\n\n@pytest.fixture(scope=\"function\")\ndef basic_collection(\n collection_temp_dir, basic_schema, collection_option\n) -> Generator[Any, Any, Collection]:\n yield from create_collection_fixture(\n collection_temp_dir, basic_schema, collection_option\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef collection_option():\n return CollectionOption(read_only=False, enable_mmap=True)\n\n\n@pytest.fixture(scope=\"function\")\ndef collection_temp_dir(tmp_path_factory):\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection_path\"\n return str(collection_path)\n\n\n@pytest.fixture(scope=\"function\")\ndef full_collection(\n collection_temp_dir,\n full_schema,\n collection_option,\n nullable: bool = True,\n has_index: bool = False,\n) -> Generator[Any, Any, Collection]:\n yield from create_collection_fixture(\n collection_temp_dir, full_schema, collection_option\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef full_collection_new(\n collection_temp_dir, full_schema_new, collection_option\n) -> Generator[Any, Any, Collection]:\n yield from create_collection_fixture(\n collection_temp_dir, full_schema_new, collection_option\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef full_collection_ivf(\n collection_temp_dir, full_schema_ivf, collection_option\n) -> Generator[Any, Any, Collection]:\n yield from create_collection_fixture(\n collection_temp_dir, full_schema_ivf, collection_option\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef full_collection_1024(\n collection_temp_dir, full_schema_1024, collection_option\n) -> Generator[Any, Any, Collection]:\n yield from create_collection_fixture(\n collection_temp_dir, full_schema_1024, collection_option\n )\n\n\n@pytest.fixture\ndef sample_field_list(nullable: bool = True, scalar_index_param=None, name_prefix=\"\"):\n field_list = []\n for k, v in DEFAULT_SCALAR_FIELD_NAME.items():\n field_list.append(\n FieldSchema(\n f\"{name_prefix}_{v}\" if len(name_prefix) > 0 else v,\n k,\n nullable=nullable,\n index_param=scalar_index_param,\n )\n )\n return field_list\n\n\n@pytest.fixture\ndef sample_vector_list(vector_index_param=None, name_prefix=\"\"):\n vector_list = []\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n vector_list.append(\n VectorSchema(\n f\"{name_prefix}_{v}\" if len(name_prefix) > 0 else v,\n k,\n dimension=DEFAULT_VECTOR_DIMENSION,\n index_param=vector_index_param,\n )\n )\n return vector_list\n" }, { "path": "python/tests/detail/params_helper.py", "content": "from zvec import (\n CollectionOption,\n IndexOption,\n OptimizeOption,\n InvertIndexParam,\n HnswIndexParam,\n IVFIndexParam,\n FlatIndexParam,\n AlterColumnOption,\n AddColumnOption,\n DataType,\n MetricType,\n QuantizeType,\n)\n\n\nVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP = {\n DataType.VECTOR_FP32: [\n HnswIndexParam(),\n HnswIndexParam(\n metric_type=MetricType.IP,\n m=16,\n ef_construction=100,\n quantize_type=QuantizeType.INT8,\n ),\n HnswIndexParam(\n metric_type=MetricType.COSINE,\n m=24,\n ef_construction=150,\n quantize_type=QuantizeType.INT4,\n ),\n HnswIndexParam(\n metric_type=MetricType.L2,\n m=32,\n ef_construction=200,\n quantize_type=QuantizeType.FP16,\n ),\n FlatIndexParam(),\n FlatIndexParam(metric_type=MetricType.IP, quantize_type=QuantizeType.INT4),\n FlatIndexParam(metric_type=MetricType.L2, quantize_type=QuantizeType.INT8),\n FlatIndexParam(metric_type=MetricType.COSINE, quantize_type=QuantizeType.FP16),\n IVFIndexParam(),\n IVFIndexParam(\n metric_type=MetricType.IP,\n quantize_type=QuantizeType.INT4,\n n_list=100,\n n_iters=10,\n use_soar=False,\n ),\n IVFIndexParam(\n metric_type=MetricType.L2,\n quantize_type=QuantizeType.INT8,\n n_list=200,\n n_iters=20,\n use_soar=True,\n ),\n IVFIndexParam(\n metric_type=MetricType.COSINE,\n quantize_type=QuantizeType.FP16,\n n_list=150,\n n_iters=15,\n use_soar=False,\n ),\n ],\n DataType.VECTOR_FP16: [\n HnswIndexParam(),\n FlatIndexParam(),\n # IVFIndexParam(),\n ],\n DataType.VECTOR_INT8: [\n HnswIndexParam(),\n FlatIndexParam(),\n # IVFIndexParam(),\n ],\n DataType.SPARSE_VECTOR_FP32: [\n HnswIndexParam(),\n FlatIndexParam(),\n HnswIndexParam(\n metric_type=MetricType.IP,\n m=16,\n ef_construction=100,\n quantize_type=QuantizeType.FP16,\n ),\n ],\n DataType.SPARSE_VECTOR_FP16: [\n HnswIndexParam(),\n FlatIndexParam(),\n HnswIndexParam(\n metric_type=MetricType.IP,\n m=16,\n ef_construction=100,\n ),\n ],\n}\n\nVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS = [\n (data_type, param)\n for data_type, params in VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP.items()\n for param in params\n]\n\nINVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP = {\n DataType.VECTOR_FP32: [\n InvertIndexParam(),\n ],\n DataType.VECTOR_FP16: [\n InvertIndexParam(),\n ],\n DataType.VECTOR_INT8: [\n InvertIndexParam(),\n ],\n DataType.SPARSE_VECTOR_FP32: [\n HnswIndexParam(metric_type=MetricType.L2),\n FlatIndexParam(metric_type=MetricType.COSINE),\n IVFIndexParam(),\n InvertIndexParam(),\n ],\n DataType.SPARSE_VECTOR_FP16: [\n HnswIndexParam(metric_type=MetricType.L2),\n FlatIndexParam(metric_type=MetricType.COSINE),\n IVFIndexParam(),\n InvertIndexParam(),\n ],\n}\n\nINVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS = [\n (data_type, param)\n for data_type, params in INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP.items()\n for param in params\n]\n\nCOLLECTION_NAME_MAX_LENGTH = 64\n\nCOLLECTION_NAME_VALID_LIST = [\n \"col\",\n \"C0llECTION\",\n \"Collection1\",\n \"collection_2\",\n \"123collection-\",\n \"a\" * COLLECTION_NAME_MAX_LENGTH,\n]\n\nCOLLECTION_NAME_INVALID_LIST = [\n \"l\",\n \"1C\",\n \"\",\n \" \",\n None,\n \"abcdefghijklmnopqrstuvwxzy123456abcdefghijklmnopqrstuvwxzy1234561\",\n \"test/\",\n \"!@#$%^&*()test\",\n]\n\nFIELD_NAME_VALID_LIST = [\n \"1\",\n \"12\",\n \"col\",\n \"ID\",\n \"name1\",\n \"Weigt_12-\",\n \"123age\",\n \"name_with_underscores\",\n \"123numeric_start\",\n \"name-with-dashes\",\n]\n\nFIELD_NAME_INVALID_LIST = [\n \"\",\n \" \",\n None,\n \"abcdefghijklmnopqrstuvwxzy1234561\",\n \"test/\",\n \"!@#$%^&*()test\",\n \"name@with#special$chars\",\n \"name with spaces\",\n]\n\nFIELD_LIST_MAX_LENGTH = 1024\nVECTOR_LIST_MAX_LENGTH = 5\nDENSE_VECTOR_MAX_DIMENSION = 20000\nSPARSE_VECTOR_MAX_DIMENSION = 4096\n\nFIELD_VECTOR_LIST_DIMENSION_VALID_LIST = [\n # field_list_len, vector_list_len, dimension\n (1, 1, 1),\n (2, 2, 512),\n (512, 3, 1024),\n (1024, 4, 20000),\n]\n\nFIELD_VECTOR_LIST_DIMENSION_INVALID_LIST = [\n # field_list_len, vector_list_len, dimension\n (1, 1, 0),\n (1, 1, -1),\n (1, 1, \"1\"),\n (1, 1, 20001),\n]\n\n\nINCOMPATIBLE_CONSTRUCTOR_ERROR_MSG = \"incompatible constructor arguments\"\nSCHEMA_VALIDATE_ERROR_MSG = \"schema validate failed\"\nCREATE_READ_ONLY_ERROR_MSG = \"Unable to create collection with read-only mode\"\nINCOMPATIBLE_FUNCTION_ERROR_MSG = \"incompatible function arguments\"\nINVALID_PATH_ERROR_MSG = \"path validate failed\"\nINDEX_NON_EXISTENT_COLUMN_ERROR_MSG = \"not found in schema\"\nACCESS_DESTROYED_COLLECTION_ERROR_MSG = \"is already destroyed\"\nCOLLECTION_PATH_NOT_EXIST_ERROR_MSG = \"not exist\"\nNOT_SUPPORT_ADD_COLUMN_ERROR_MSG = \"Only support basic numeric data type\"\nNOT_EXIST_COLUMN_TO_DROP_ERROR_MSG = \"Column not exists\"\n" }, { "path": "python/tests/detail/support_helper.py", "content": "from zvec import (\n CollectionOption,\n IndexOption,\n OptimizeOption,\n InvertIndexParam,\n HnswIndexParam,\n IVFIndexParam,\n FlatIndexParam,\n DataType,\n IndexType,\n QuantizeType,\n)\n\nSUPPORT_SCALAR_DATA_TYPES = [\n DataType.BOOL,\n DataType.FLOAT,\n DataType.DOUBLE,\n DataType.INT32,\n DataType.INT64,\n DataType.UINT32,\n DataType.UINT64,\n DataType.STRING,\n DataType.ARRAY_BOOL,\n DataType.ARRAY_FLOAT,\n DataType.ARRAY_DOUBLE,\n DataType.ARRAY_INT32,\n DataType.ARRAY_INT64,\n DataType.ARRAY_UINT32,\n DataType.ARRAY_UINT64,\n DataType.ARRAY_STRING,\n]\n\nDEFAULT_SCALAR_FIELD_NAME = {\n DataType.BOOL: \"bool_field\",\n DataType.FLOAT: \"float_field\",\n DataType.DOUBLE: \"double_field\",\n DataType.INT32: \"int32_field\",\n DataType.INT64: \"int64_field\",\n DataType.UINT32: \"uint32_field\",\n DataType.UINT64: \"uint64_field\",\n DataType.STRING: \"string_field\",\n DataType.ARRAY_BOOL: \"array_bool_field\",\n DataType.ARRAY_FLOAT: \"array_float_field\",\n DataType.ARRAY_DOUBLE: \"array_double_field\",\n DataType.ARRAY_INT32: \"array_int32_field\",\n DataType.ARRAY_INT64: \"array_int64_field\",\n DataType.ARRAY_UINT32: \"array_uint32_field\",\n DataType.ARRAY_UINT64: \"array_uint64_field\",\n DataType.ARRAY_STRING: \"array_string_field\",\n}\n\nSUPPORT_SCALAR_INDEX_TYPES = [\n IndexType.INVERT,\n]\n\nSUPPORT_VECTOR_DATA_TYPES = [\n DataType.VECTOR_FP16,\n DataType.VECTOR_FP32,\n DataType.VECTOR_INT8,\n DataType.SPARSE_VECTOR_FP32,\n DataType.SPARSE_VECTOR_FP16,\n]\n\nSUPPORT_VECTOR_INDEX_TYPES = [\n IndexType.FLAT,\n IndexType.HNSW,\n IndexType.IVF,\n]\n\nDEFAULT_VECTOR_FIELD_NAME = {\n DataType.VECTOR_FP16: \"vector_fp16_field\",\n DataType.VECTOR_FP32: \"vector_fp32_field\",\n DataType.VECTOR_INT8: \"vector_int8_field\",\n DataType.SPARSE_VECTOR_FP32: \"sparse_vector_fp32_field\",\n DataType.SPARSE_VECTOR_FP16: \"sparse_vector_fp16_field\",\n}\n\nDEFAULT_VECTOR_DIMENSION = 128\nVECTOR_DIMENSION_1024 = 4\nSUPPORT_VECTOR_DATA_TYPE_INDEX_MAP = {\n DataType.VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF],\n DataType.VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF],\n DataType.VECTOR_INT8: [IndexType.FLAT, IndexType.HNSW],\n DataType.SPARSE_VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW],\n DataType.SPARSE_VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW],\n}\n\nSUPPORT_VECTOR_DATA_TYPE_INDEX_MAP_PARAMS = [\n (data_type, index_type)\n for data_type, index_types in SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP.items()\n for index_type in index_types\n]\n\nDEFAULT_INDEX_PARAMS = {\n IndexType.FLAT: FlatIndexParam(),\n IndexType.HNSW: HnswIndexParam(),\n IndexType.IVF: IVFIndexParam(),\n IndexType.INVERT: InvertIndexParam(),\n}\n\nSUPPORT_VECTOR_DATA_TYPE_QUANT_MAP = {\n DataType.VECTOR_FP32: [QuantizeType.FP16, QuantizeType.INT8, QuantizeType.INT4],\n DataType.SPARSE_VECTOR_FP32: [QuantizeType.FP16],\n}\n\nSUPPORT_ADD_COLUMN_DATA_TYPE = [\n DataType.INT32,\n DataType.UINT32,\n DataType.INT64,\n DataType.UINT64,\n DataType.FLOAT,\n DataType.DOUBLE,\n]\n\nNOT_SUPPORT_ADD_COLUMN_DATA_TYPE = [\n DataType.BOOL,\n DataType.STRING,\n DataType.ARRAY_BOOL,\n DataType.ARRAY_INT32,\n DataType.ARRAY_INT64,\n DataType.ARRAY_UINT32,\n DataType.ARRAY_UINT64,\n DataType.ARRAY_FLOAT,\n DataType.ARRAY_DOUBLE,\n DataType.ARRAY_STRING,\n]\n" }, { "path": "python/tests/detail/test_collection_concurrency.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport pytest\nimport threading\nimport numpy as np\nimport zvec\n\nfrom zvec import (\n CollectionOption,\n InvertIndexParam,\n HnswIndexParam,\n Collection,\n Doc,\n DataType,\n FieldSchema,\n VectorSchema,\n)\n\n\nclass TestCollectionConcurrency:\n @pytest.fixture(scope=\"function\")\n def test_collection(self, tmp_path_factory):\n \"\"\"Fixture to create a test collection\"\"\"\n collection_schema = zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\",\n DataType.STRING,\n nullable=False,\n index_param=InvertIndexParam(),\n ),\n FieldSchema(\"weight\", DataType.FLOAT, nullable=True),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n VectorSchema(\n \"sparse\", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()\n ),\n ],\n )\n\n collection_option = CollectionOption(read_only=False, enable_mmap=True)\n\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n coll = zvec.create_and_open(\n path=str(collection_path),\n schema=collection_schema,\n option=collection_option,\n )\n\n assert coll is not None, \"Failed to create and open collection\"\n\n yield coll\n\n # Clean up\n if hasattr(coll, \"destroy\") and coll is not None:\n try:\n coll.destroy()\n except Exception as e:\n print(f\"Warning: failed to destroy collection: {e}\")\n\n def test_concurrent_read_write(self, test_collection: Collection):\n results = []\n\n def insert_docs(thread_id):\n try:\n docs = [\n Doc(\n id=f\"{thread_id}_{i}\",\n fields={\n \"id\": int(f\"{thread_id}{i}\"),\n \"name\": f\"thread_{thread_id}_doc_{i}\",\n \"weight\": float(i),\n },\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: float(i), 2: float(i * 2)},\n },\n )\n for i in range(5)\n ]\n\n result = test_collection.insert(docs)\n results.append((thread_id, \"insert\", len(result)))\n except Exception as e:\n results.append((thread_id, \"insert_exception\", str(e)))\n\n def query_docs(thread_id):\n try:\n result = test_collection.query(filter=\"id > 0\", topk=10)\n results.append((thread_id, \"query\", len(result)))\n except Exception as e:\n results.append((thread_id, \"query_exception\", str(e)))\n\n # Create threads for concurrent operations\n threads = []\n\n # Start insert threads\n for i in range(3):\n thread = threading.Thread(target=insert_docs, args=(i,))\n threads.append(thread)\n thread.start()\n\n # Start query threads\n for i in range(3):\n thread = threading.Thread(target=query_docs, args=(i,))\n threads.append(thread)\n thread.start()\n\n # Wait for all threads to complete\n for thread in threads:\n thread.join()\n\n # Analyze results\n insert_results = [r for r in results if r[1] == \"insert\"]\n query_results = [r for r in results if r[1] == \"query\"]\n\n logging.info(\n f\"Concurrent read/write results - Inserts: {len(insert_results)}, Queries: {len(query_results)}\"\n )\n\n # At least some operations should succeed\n assert len(insert_results) + len(query_results) > 0\n\n def test_concurrent_query(self, test_collection: Collection):\n # First insert some data\n docs = [\n Doc(\n id=f\"{i}\",\n fields={\"id\": i, \"name\": f\"test_{i}\", \"weight\": float(i)},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: float(i), 2: float(i * 2)},\n },\n )\n for i in range(20)\n ]\n\n insert_result = test_collection.insert(docs)\n assert len(insert_result) == 20\n\n results = []\n\n def query_operation(thread_id):\n \"\"\"Perform query operation from a thread\"\"\"\n try:\n result = test_collection.query(filter=f\"id > {thread_id}\", topk=5)\n results.append((thread_id, \"query\", len(result)))\n except Exception as e:\n results.append((thread_id, \"query_exception\", str(e)))\n\n # Create multiple threads for concurrent queries\n threads = []\n for i in range(5):\n thread = threading.Thread(target=query_operation, args=(i,))\n threads.append(thread)\n thread.start()\n\n # Wait for all threads to complete\n for thread in threads:\n thread.join()\n\n # Analyze results\n query_results = [r for r in results if r[1] == \"query\"]\n logging.info(f\"Concurrent query results - Queries: {len(query_results)}\")\n\n # All query operations should succeed\n assert len(query_results) == 5\n\n def test_concurrent_modifications(self, test_collection: Collection):\n # First insert some data\n docs = [\n Doc(\n id=f\"{i}\",\n fields={\"id\": i, \"name\": f\"test_{i}\", \"weight\": float(i)},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: float(i), 2: float(i * 2)},\n },\n )\n for i in range(10)\n ]\n\n insert_result = test_collection.insert(docs)\n assert len(insert_result) == 10\n\n results = []\n\n def update_operation(thread_id):\n \"\"\"Perform update operation from a thread\"\"\"\n try:\n # Each thread updates different documents\n update_docs = [\n Doc(\n id=f\"{i}\",\n fields={\n \"id\": i,\n \"name\": f\"updated_by_thread_{thread_id}\",\n \"weight\": float(i + thread_id),\n },\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: float(i) + 0.5, 2: float(i * 2) + 0.5},\n },\n )\n for i in range(thread_id * 2, thread_id * 2 + 2)\n ]\n\n result = test_collection.update(update_docs)\n results.append((thread_id, \"update\", len(result)))\n except Exception as e:\n results.append((thread_id, \"update_exception\", str(e)))\n\n def delete_operation(thread_id):\n \"\"\"Perform delete operation from a thread\"\"\"\n try:\n # Each thread deletes different documents\n delete_ids = [f\"{thread_id * 2 + 2}\", f\"{thread_id * 2 + 3}\"]\n result = test_collection.delete(delete_ids)\n results.append((thread_id, \"delete\", len(result)))\n except Exception as e:\n results.append((thread_id, \"delete_exception\", str(e)))\n\n # Create threads for concurrent operations\n threads = []\n\n # Start update threads\n for i in range(3):\n thread = threading.Thread(target=update_operation, args=(i,))\n threads.append(thread)\n thread.start()\n\n # Start delete threads\n for i in range(2):\n thread = threading.Thread(target=delete_operation, args=(i,))\n threads.append(thread)\n thread.start()\n\n # Wait for all threads to complete\n for thread in threads:\n thread.join()\n\n # Analyze results\n update_results = [r for r in results if r[1] == \"update\"]\n delete_results = [r for r in results if r[1] == \"delete\"]\n\n logging.info(\n f\"Concurrent modification results - Updates: {len(update_results)}, Deletes: {len(delete_results)}\"\n )\n\n # At least some operations should succeed\n assert len(update_results) + len(delete_results) > 0\n\n def test_read_write_locking(self, test_collection: Collection):\n # Perform operations that should be thread-safe\n docs = [\n Doc(\n id=f\"{i}\",\n fields={\"id\": i, \"name\": f\"test_{i}\", \"weight\": float(i)},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: float(i), 2: float(i * 2)},\n },\n )\n for i in range(5)\n ]\n\n # Insert data\n insert_result = test_collection.insert(docs)\n assert len(insert_result) == 5\n\n # Concurrent operations should not cause data corruption\n results = []\n\n def mixed_operation(thread_id):\n \"\"\"Perform mixed operations from a thread\"\"\"\n try:\n # Mix of read and write operations\n if thread_id % 2 == 0:\n # Read operation\n result = test_collection.fetch([f\"{thread_id % 5}\"])\n results.append((thread_id, \"read\", len(result)))\n else:\n # Write operation\n doc = Doc(\n id=f\"{thread_id % 5}\",\n fields={\n \"id\": thread_id % 5,\n \"name\": f\"mixed_op_{thread_id}\",\n \"weight\": float(thread_id),\n },\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: float(thread_id), 2: float(thread_id * 2)},\n },\n )\n result = test_collection.upsert(doc)\n results.append((thread_id, \"write\", len(result)))\n except Exception as e:\n results.append((thread_id, \"exception\", str(e)))\n\n # Create multiple threads\n threads = []\n for i in range(10):\n thread = threading.Thread(target=mixed_operation, args=(i,))\n threads.append(thread)\n thread.start()\n\n # Wait for all threads to complete\n for thread in threads:\n thread.join()\n\n # Verify that the collection is still in a consistent state\n final_result = test_collection.query()\n assert len(final_result) >= 0 # Should not crash or return corrupted data\n\n def test_race_condition_detection(self, test_collection: Collection):\n # Insert initial data\n docs = [\n Doc(\n id=f\"{i}\",\n fields={\"id\": i, \"name\": f\"initial_{i}\", \"weight\": float(i)},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: float(i), 2: float(i * 2)},\n },\n )\n for i in range(10)\n ]\n\n insert_result = test_collection.insert(docs)\n assert len(insert_result) == 10\n\n # Perform many rapid concurrent operations\n operation_count = 100\n results = []\n\n def rapid_operation(op_id):\n \"\"\"Perform rapid operations\"\"\"\n try:\n # Alternate between different types of operations\n if op_id % 4 == 0:\n # Insert\n doc = Doc(\n id=f\"rapid_{op_id}\",\n fields={\n \"id\": op_id,\n \"name\": f\"rapid_{op_id}\",\n \"weight\": float(op_id),\n },\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: float(op_id), 2: float(op_id * 2)},\n },\n )\n result = test_collection.insert(doc)\n results.append((\"insert\", len(result)))\n elif op_id % 4 == 1:\n # Update\n doc = Doc(\n id=f\"{op_id % 10}\",\n fields={\n \"id\": op_id % 10,\n \"name\": f\"rapid_update_{op_id}\",\n \"weight\": float(op_id),\n },\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: float(op_id), 2: float(op_id * 2)},\n },\n )\n result = test_collection.update(doc)\n results.append((\"update\", len(result)))\n elif op_id % 4 == 2:\n # Query\n result = test_collection.query(filter=f\"id > {op_id % 5}\", topk=3)\n results.append((\"query\", len(result)))\n else:\n # Fetch\n result = test_collection.fetch([f\"{op_id % 10}\"])\n results.append((\"fetch\", len(result)))\n except Exception as e:\n results.append((\"exception\", str(e)))\n\n # Create many threads for rapid concurrent operations\n threads = []\n for i in range(operation_count):\n thread = threading.Thread(target=rapid_operation, args=(i,))\n threads.append(thread)\n thread.start()\n\n # Wait for all threads to complete\n for thread in threads:\n thread.join()\n\n # Verify collection is still functional\n final_query = test_collection.query()\n assert len(final_query) >= 0 # Should not be corrupted\n\n logging.info(\n f\"Rapid concurrent operations completed - Total operations: {len(results)}\"\n )\n" }, { "path": "python/tests/detail/test_collection_create_and_open.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport threading\nimport os\n\nfrom distance_helper import *\nfrom fixture_helper import *\nfrom doc_helper import *\nfrom params_helper import *\n\n\ndef check_collection_info(\n coll: Collection, schema: CollectionSchema, option: CollectionOption, path: str\n):\n assert coll is not None, \"Failed to create and open collection\"\n assert coll.path == path\n assert coll.schema.name == schema.name\n assert list(coll.schema.fields) == list(schema.fields)\n assert list(coll.schema.vectors) == list(schema.vectors)\n assert coll.option.read_only == option.read_only\n assert coll.option.enable_mmap == option.enable_mmap\n\n\ndef check_collection_basic(coll: Collection, optimize: bool = False):\n schema = coll.schema\n\n docs = [generate_doc(i, schema) for i in range(10)]\n\n results = coll.insert(docs=docs)\n assert len(results) == len(docs)\n for result in results:\n assert result.ok()\n\n assert coll.stats.doc_count == len(docs)\n\n def check_fetch_query():\n results = coll.fetch([str(i) for i in range(len(docs))])\n assert len(results) == len(docs)\n for i in range(len(docs)):\n assert str(i) in results\n\n results = coll.query()\n assert len(results) == len(docs)\n\n check_fetch_query()\n\n if optimize:\n coll.optimize()\n check_fetch_query()\n\n\ndef check_collection_full(coll: Collection):\n test_doc = generate_doc(1, coll.schema)\n\n insert_result = coll.insert(test_doc)\n assert insert_result.ok()\n\n stats = coll.stats\n assert stats.doc_count == 1\n\n fetched_docs = coll.fetch(ids=[\"1\"])\n assert len(fetched_docs) == 1\n assert \"1\" in fetched_docs\n assert fetched_docs[\"1\"] is not None\n assert is_doc_equal(fetched_docs[\"1\"], test_doc, coll.schema)\n\n query_result = coll.query()\n assert len(query_result) == 1\n\n updated_doc = Doc(\n id=\"1\",\n fields={\"int32_field\": 1},\n vectors={\"vector_fp32_field\": [0.2] * 128},\n )\n update_result = coll.update(updated_doc)\n assert update_result.ok()\n\n upserted_doc = generate_doc(1, coll.schema)\n upsert_result = coll.upsert(upserted_doc)\n assert upsert_result.ok()\n\n # 8. Delete document\n delete_result = coll.delete(\"1\")\n assert delete_result.ok()\n\n # Verify document was deleted\n stats = coll.stats\n assert stats.doc_count == 0\n\n\nvalid_collection_options = [\n # (read_only, enable_mmap)\n (False, True),\n (False, False),\n]\ninvalid_collection_options = [\n # (read_only, enable_mmap)\n (True, True),\n (True, False),\n]\nduplicate_names_test = [\n (\"field1\", \"field1\", \"vector1\", \"vector2\"),\n (\"field1\", \"field2\", \"vector1\", \"vector1\"),\n (\n \"shared_name1\",\n \"shared_name2\",\n \"shared_name1\",\n \"shared_name2\",\n ),\n]\nlong_names = [\n \"a\" * 100, # 100 characters\n \"b\" * 200, # 200 characters\n]\n\nvalid_path_list = [\n \"/tmp/nonexistent/directory/test_collection\",\n \"test/collection/with/slashes\",\n]\ninvalid_path_list = [\n \"invalid:path\",\n \"\",\n \"test_collection_with_spaces \",\n \"test@#$%collection\",\n]\n\n\nclass TestCreateAndOpen:\n @pytest.mark.parametrize(\"collection_name\", COLLECTION_NAME_VALID_LIST)\n def test_valid_collection_name(\n self,\n collection_temp_dir,\n collection_name,\n collection_option,\n sample_field_list,\n sample_vector_list,\n ):\n collection_schema = zvec.CollectionSchema(\n name=collection_name,\n fields=sample_field_list,\n vectors=sample_vector_list,\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n check_collection_info(\n coll, collection_schema, collection_option, collection_temp_dir\n )\n check_collection_basic(coll)\n\n coll.destroy()\n\n @pytest.mark.parametrize(\"collection_name\", COLLECTION_NAME_INVALID_LIST)\n def test_invalid_collection_name(\n self,\n collection_temp_dir,\n collection_name,\n collection_option,\n sample_field_list,\n sample_vector_list,\n ):\n with pytest.raises(Exception) as exc_info:\n collection_schema = zvec.CollectionSchema(\n name=collection_name,\n fields=sample_field_list,\n vectors=sample_vector_list,\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n\n @pytest.mark.parametrize(\"name_prefix\", FIELD_NAME_VALID_LIST)\n def test_valid_field_vector_name(\n self,\n collection_temp_dir,\n collection_option,\n name_prefix,\n sample_field_list,\n sample_vector_list,\n ):\n collection_schema = zvec.CollectionSchema(\n name=\"test_collection\",\n fields=sample_field_list,\n vectors=sample_vector_list,\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n check_collection_info(\n coll, collection_schema, collection_option, collection_temp_dir\n )\n check_collection_basic(coll)\n\n coll.destroy()\n\n @pytest.mark.parametrize(\"field_name\", FIELD_NAME_INVALID_LIST)\n def test_invalid_field_name(\n self, collection_temp_dir, collection_option, field_name\n ):\n with pytest.raises(Exception) as exc_info:\n field_list = [FieldSchema(field_name, DataType.STRING)]\n vector_list = [\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n ]\n\n collection_schema = zvec.CollectionSchema(\n name=\"collection_name\", fields=field_list, vectors=vector_list\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n\n @pytest.mark.parametrize(\"vector_name\", FIELD_NAME_INVALID_LIST)\n def test_invalid_vector_name(\n self, collection_temp_dir, collection_option, vector_name\n ):\n with pytest.raises(Exception) as exc_info:\n field_list = [\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n )\n ]\n vector_list = [\n VectorSchema(vector_name, DataType.VECTOR_FP32, dimension=128)\n ]\n\n collection_schema = zvec.CollectionSchema(\n name=\"collection_name\", fields=field_list, vectors=vector_list\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n\n @pytest.mark.parametrize(\n \"field_list_len,vector_list_len,dimension\",\n FIELD_VECTOR_LIST_DIMENSION_VALID_LIST,\n )\n def test_valid_field_vector_size_dimension(\n self,\n collection_temp_dir,\n collection_option,\n field_list_len,\n vector_list_len,\n dimension,\n ):\n field_list = []\n vector_list = []\n for i in range(0, field_list_len):\n field_list.append(\n FieldSchema(\"id_\" + str(i), DataType.INT64, nullable=True)\n )\n\n for i in range(0, vector_list_len):\n vector_list.append(\n VectorSchema(\n \"dense_vector_\" + str(i),\n DataType.VECTOR_FP32,\n dimension=dimension,\n index_param=HnswIndexParam(),\n )\n )\n\n collection_schema = zvec.CollectionSchema(\n name=\"test_dense_vector_list\", fields=field_list, vectors=vector_list\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n check_collection_info(\n coll, collection_schema, collection_option, collection_temp_dir\n )\n check_collection_basic(coll)\n\n coll.destroy()\n\n @pytest.mark.parametrize(\n \"field_list_len,vector_list_len,dimension\",\n FIELD_VECTOR_LIST_DIMENSION_INVALID_LIST,\n )\n def test_invalid_field_vector_size_dimension(\n self,\n collection_temp_dir,\n collection_option,\n vector_list_len,\n field_list_len,\n dimension,\n ):\n with pytest.raises(Exception) as exc_info:\n field_list = []\n vector_list = []\n for i in range(0, field_list_len):\n field_list.append(\n FieldSchema(\n \"id_\" + str(i),\n DataType.INT64,\n nullable=False,\n )\n )\n\n for i in range(0, vector_list_len):\n vector_list.append(\n VectorSchema(\n \"dense_vector_\" + str(i),\n DataType.VECTOR_FP32,\n dimension=dimension,\n index_param=HnswIndexParam(),\n )\n )\n\n collection_schema = zvec.CollectionSchema(\n name=\"test_dense_vector_list\", fields=field_list, vectors=vector_list\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n\n def test_valid_single_vector_field_construction(\n self, collection_temp_dir, collection_option\n ):\n field = FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=True,\n index_param=InvertIndexParam(enable_range_optimization=True),\n )\n\n vector = VectorSchema(\n \"dense_vector\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n\n collection_schema = zvec.CollectionSchema(\n name=\"test_single_dense_vector_non_list\",\n fields=field,\n vectors=vector, # Non-list form\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n check_collection_info(\n coll, collection_schema, collection_option, collection_temp_dir\n )\n check_collection_basic(coll)\n coll.destroy()\n\n def test_collection_concurrent_create(\n self, collection_temp_dir, basic_schema, collection_option\n ):\n results = []\n errors = []\n lock = threading.Lock()\n\n # Function to be executed by each thread\n def create_collection_thread(thread_id):\n try:\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=basic_schema,\n option=collection_option,\n )\n with lock:\n results.append((thread_id, coll))\n except Exception as e:\n with lock:\n errors.append((thread_id, str(e)))\n\n threads = []\n for i in range(5):\n thread = threading.Thread(target=create_collection_thread, args=(i,))\n threads.append(thread)\n thread.start()\n\n for thread in threads:\n thread.join()\n assert len(results) == 1, (\n f\"Expected exactly one successful creation, but got {len(results)}\"\n )\n assert len(errors) == 4, (\n f\"Expected exactly four failures, but got {len(errors)}\"\n )\n\n successful_thread_id, successful_collection = results[0]\n assert successful_collection is not None, (\n \"Successful creation should return a valid collection\"\n )\n assert successful_collection.path == collection_temp_dir, (\n \"Collection path mismatch\"\n )\n\n def test_create_open_loop(\n self, collection_temp_dir, collection_option, full_schema\n ):\n for cycle in range(10):\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=full_schema,\n option=collection_option,\n )\n assert coll is not None, (\n f\"Failed to create and open collection in cycle {cycle}\"\n )\n assert coll.path == collection_temp_dir, (\n f\"Collection path mismatch in cycle {cycle}\"\n )\n\n del coll\n\n reopened_coll = zvec.open(\n path=collection_temp_dir, option=collection_option\n )\n assert reopened_coll is not None, (\n f\"Failed to reopen collection in cycle {cycle}\"\n )\n assert reopened_coll.path == collection_temp_dir, (\n f\"Reopened collection path mismatch in cycle {cycle}\"\n )\n\n check_collection_full(reopened_coll)\n\n reopened_coll.destroy()\n\n @pytest.mark.parametrize(\n \"data_type, index_param\", VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS\n )\n def test_valid_vector_index_params(\n self,\n data_type,\n index_param,\n single_vector_schema_with_index_param,\n collection_temp_dir,\n collection_option,\n ):\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=single_vector_schema_with_index_param,\n option=collection_option,\n )\n\n check_collection_info(\n coll,\n single_vector_schema_with_index_param,\n collection_option,\n collection_temp_dir,\n )\n\n check_collection_basic(coll, True)\n\n @pytest.mark.parametrize(\n \"data_type, index_param\", INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS\n )\n def test_invalid_vector_index_params(\n self,\n data_type,\n index_param,\n single_vector_schema_with_index_param,\n collection_temp_dir,\n collection_option,\n ):\n with pytest.raises(Exception) as exc_info:\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=single_vector_schema_with_index_param,\n option=collection_option,\n )\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n\n def test_open_concurrent_same_path(self, tmp_path_factory, collection_option):\n \"\"\"Test concurrent opening of the same collection path.\n\n - Multi-threading concurrency: 5 threads simultaneously open the same collection\n - Result verification: Verify that only one can open successfully, others must fail\n \"\"\"\n # Create a temporary directory and path for the collection\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"concurrent_open_test_collection\"\n\n # First, create a collection that we'll try to open concurrently\n field_list = [\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\", DataType.STRING, nullable=False, index_param=InvertIndexParam()\n ),\n ]\n\n vector_list = [\n VectorSchema(\n \"dense_vector\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n ]\n\n collection_schema = zvec.CollectionSchema(\n name=\"concurrent_open_test_collection\",\n fields=field_list,\n vectors=vector_list,\n )\n\n # Create the collection first\n coll = zvec.create_and_open(\n path=str(collection_path),\n schema=collection_schema,\n option=collection_option,\n )\n\n # Close the collection so we can test opening it\n if hasattr(coll, \"close\") and coll is not None:\n coll.close()\n\n # Shared variables to collect results from threads\n results = []\n errors = []\n\n # Lock for thread-safe operations\n lock = threading.Lock()\n # Clean up the created collection reference\n del coll\n\n # Function to be executed by each thread\n def open_collection_thread(thread_id):\n try:\n reopened_coll = zvec.open(\n path=str(collection_path), option=collection_option\n )\n with lock:\n results.append((thread_id, reopened_coll))\n # Clean up the collection if opened successfully\n if hasattr(reopened_coll, \"close\") and reopened_coll is not None:\n reopened_coll.close()\n except Exception as e:\n with lock:\n errors.append((thread_id, str(e)))\n\n # Create and start 5 threads\n threads = []\n for i in range(5):\n thread = threading.Thread(target=open_collection_thread, args=(i,))\n threads.append(thread)\n thread.start()\n\n # Wait for all threads to complete\n for thread in threads:\n thread.join()\n\n # Verify results:\n # 1. Only one open should succeed (exactly one collection in results)\n # 2. Others should fail (4 errors in errors)\n assert len(results) == 1, (\n f\"Expected exactly one successful open, but got {len(results)}\"\n )\n assert len(errors) == 4, (\n f\"Expected exactly four failures, but got {len(errors)}\"\n )\n\n # Additional verification: check that the successful open has a valid collection\n successful_thread_id, successful_collection = results[0]\n assert successful_collection is not None, (\n \"Successful open should return a valid collection\"\n )\n assert successful_collection.path == str(collection_path), (\n \"Collection path mismatch\"\n )\n\n @pytest.mark.parametrize(\"read_only,enable_mmap\", valid_collection_options)\n def test_valid_option(\n self, collection_temp_dir, basic_schema, read_only, enable_mmap\n ):\n option = CollectionOption(read_only=read_only, enable_mmap=enable_mmap)\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=basic_schema,\n option=option,\n )\n\n check_collection_info(coll, basic_schema, option, collection_temp_dir)\n check_collection_basic(coll)\n\n coll.destroy()\n\n def test_valid_none_option(self, collection_temp_dir, basic_schema):\n zvec.create_and_open(\n path=collection_temp_dir,\n schema=basic_schema,\n option=None,\n )\n\n @pytest.mark.parametrize(\"read_only,enable_mmap\", invalid_collection_options)\n def test_invalid_option(\n self, collection_temp_dir, basic_schema, read_only, enable_mmap\n ):\n with pytest.raises(Exception) as exc_info:\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=basic_schema,\n option=CollectionOption(read_only=read_only, enable_mmap=enable_mmap),\n )\n\n assert CREATE_READ_ONLY_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n\n @pytest.mark.parametrize(\n \"field_name1,field_name2,vector_name1,vector_name2\",\n duplicate_names_test,\n )\n def test_duplicate_field_names(\n self,\n collection_temp_dir,\n collection_option,\n field_name1,\n field_name2,\n vector_name1,\n vector_name2,\n ):\n with pytest.raises(Exception) as exc_info:\n collection_schema = zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n field_name1,\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n field_name2,\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n ],\n vectors=[\n VectorSchema(\n vector_name1,\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n VectorSchema(\n vector_name2,\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n ],\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n\n @pytest.mark.parametrize(\"long_name\", long_names)\n def test_invalid_long_field_names(\n self, collection_option, collection_temp_dir, long_name\n ):\n collection_schema = zvec.CollectionSchema(\n name=long_name,\n fields=[\n FieldSchema(\n long_name + \"_field\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n ],\n vectors=[\n VectorSchema(\n long_name + \"_vector\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n ],\n )\n\n with pytest.raises(Exception) as exc_info:\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n\n def test_invalid_empty_fields_and_vectors(\n self, collection_temp_dir, collection_option\n ):\n collection_schema = zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[], # Empty fields\n vectors=[], # Empty vectors\n )\n\n with pytest.raises(Exception) as exc_info:\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=collection_schema,\n option=collection_option,\n )\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n\n @pytest.mark.parametrize(\"valid_path\", valid_path_list)\n def test_valid_path(self, basic_schema, collection_option, valid_path):\n if os.path.exists(valid_path):\n import shutil\n\n shutil.rmtree(valid_path)\n\n coll = zvec.create_and_open(\n path=valid_path, schema=basic_schema, option=collection_option\n )\n\n check_collection_info(coll, basic_schema, collection_option, valid_path)\n\n coll.destroy()\n\n @pytest.mark.parametrize(\"invalid_path\", invalid_path_list)\n def test_invalid_path(self, basic_schema, collection_option, invalid_path):\n with pytest.raises(Exception) as exc_info:\n coll = zvec.create_and_open(\n path=invalid_path, schema=basic_schema, option=collection_option\n )\n\n assert INVALID_PATH_ERROR_MSG in str(exc_info.value), str(exc_info.value)\n" }, { "path": "python/tests/detail/test_collection_ddl.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom distance_helper import *\nfrom fixture_helper import *\nfrom doc_helper import *\nfrom params_helper import *\n\n\nclass TestDDL:\n def test_collection_stats(self, basic_collection: Collection):\n assert basic_collection.stats is not None\n stats = basic_collection.stats\n assert stats.doc_count == 0\n assert len(stats.index_completeness) == 2\n assert stats.index_completeness[\"dense\"] == 1\n assert stats.index_completeness[\"sparse\"] == 1\n\n def test_collection_destroy(\n self, basic_collection: Collection, collection_temp_dir, collection_option\n ):\n doc = generate_doc(1, basic_collection.schema)\n\n result = basic_collection.insert(doc)\n assert bool(result)\n assert result.ok()\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n basic_collection.destroy()\n\n with pytest.raises(Exception) as exc_info:\n stats = basic_collection.stats\n assert ACCESS_DESTROYED_COLLECTION_ERROR_MSG in str(exc_info.value)\n\n with pytest.raises(Exception) as exc_info:\n zvec.open(path=collection_temp_dir, option=collection_option)\n assert COLLECTION_PATH_NOT_EXIST_ERROR_MSG in str(exc_info.value)\n\n def test_collection_flush(self, basic_collection: Collection):\n doc = generate_doc(1, basic_collection.schema)\n\n result = basic_collection.insert(doc)\n assert bool(result)\n assert result.ok()\n\n basic_collection.flush()\n\n fetched_docs = basic_collection.fetch([\"1\"])\n assert \"1\" in fetched_docs\n assert fetched_docs[\"1\"].id == \"1\"\n\n\nclass TestIndexDDL:\n @pytest.mark.parametrize(\"field_name\", DEFAULT_SCALAR_FIELD_NAME.values())\n @pytest.mark.parametrize(\"index_type\", SUPPORT_SCALAR_INDEX_TYPES)\n def test_scalar_index_operation(\n self,\n full_collection: Collection,\n field_name: str,\n index_type: IndexType,\n ):\n # INSERT 0~5 Doc\n docs = [generate_doc(i, full_collection.schema) for i in range(5)]\n\n result = full_collection.insert(docs)\n assert len(result) == 5\n for item in result:\n assert item.ok()\n\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 5\n\n if field_name in [\"bool_field\"]:\n query_filter = f\"{field_name} = true\"\n elif field_name in [\"double_field\", \"float_field\"]:\n query_filter = f\"{field_name} >= 3.0\"\n elif field_name in [\n \"int32_field\",\n \"int64_field\",\n \"uint32_field\",\n \"uint64_field\",\n ]:\n query_filter = f\"{field_name} >= 30\"\n elif field_name in [\"string_field\"]:\n query_filter = f\"{field_name} >= 'test_3'\"\n elif field_name in [\"array_bool_field\"]:\n query_filter = f\"{field_name} contain_any (false)\"\n elif field_name in [\"array_double_field\", \"array_float_field\"]:\n query_filter = f\"{field_name} contain_any (3.0, 4.0)\"\n elif field_name in [\n \"array_int64_field\",\n \"array_int32_field\",\n \"array_uint64_field\",\n \"array_uint32_field\",\n ]:\n query_filter = f\"{field_name} contain_any (3, 4)\"\n elif field_name == \"array_string_field\":\n query_filter = f\"{field_name} contain_any ('test_3', 'test_4')\"\n else:\n assert False, f\"Unsupported field type for index creation: {field_name}\"\n\n query_result_before = full_collection.query(filter=query_filter, topk=10)\n\n if index_type not in DEFAULT_INDEX_PARAMS:\n pytest.fail(f\"Unsupported index type for index creation: {index_type}\")\n index_param = DEFAULT_INDEX_PARAMS[index_type]\n\n full_collection.create_index(\n field_name=field_name, index_param=index_param, option=IndexOption()\n )\n stats_after_create = full_collection.stats\n assert stats_after_create is not None\n assert stats_after_create.doc_count == 5\n\n query_result_after = full_collection.query(filter=query_filter, topk=10)\n\n assert len(query_result_before) == len(query_result_after), (\n f\"Query result count mismatch for {field_name} with index type {index_type}: before={len(query_result_before)}, after={len(query_result_after)}\"\n )\n\n before_ids = set(doc.id for doc in query_result_before)\n after_ids = set(doc.id for doc in query_result_after)\n assert before_ids == after_ids, (\n f\"Query result IDs mismatch for {field_name} with index type {index_type}: before={before_ids}, after={after_ids}\"\n )\n\n # INSERT 5~8 Doc\n new_docs = [generate_doc(i, full_collection.schema) for i in range(5, 8)]\n\n result = full_collection.insert(new_docs)\n assert len(result) == 3\n for item in result:\n assert item.ok()\n\n stats_after_insert1 = full_collection.stats\n assert stats_after_insert1 is not None\n assert stats_after_insert1.doc_count == 8\n\n fetched_docs = full_collection.fetch([f\"{i}\" for i in range(5, 8)])\n assert len(fetched_docs) == 3\n\n for i in range(5, 8):\n doc_id = f\"{i}\"\n assert doc_id in fetched_docs\n\n query_result = full_collection.query(filter=query_filter, topk=20)\n assert len(query_result) >= len(query_result_before)\n\n full_collection.drop_index(field_name=field_name)\n\n # Insert 8~10 Doc\n more_docs = [generate_doc(i, full_collection.schema) for i in range(8, 10)]\n\n result = full_collection.insert(more_docs)\n assert len(result) == 2\n for item in result:\n assert item.ok()\n\n stats_after_insert2 = full_collection.stats\n assert stats_after_insert2 is not None\n assert stats_after_insert2.doc_count == 10\n\n fetched_docs = full_collection.fetch([f\"{i}\" for i in range(8, 10)])\n assert len(fetched_docs) == 2\n\n for i in range(8, 10):\n doc_id = f\"{i}\"\n assert doc_id in fetched_docs\n\n query_result = full_collection.query(filter=query_filter, topk=20)\n assert len(query_result) >= len(query_result_before)\n\n final_stats = full_collection.stats\n assert final_stats is not None\n assert final_stats.doc_count == 10\n full_collection.destroy()\n\n @pytest.mark.parametrize(\"field_name\", DEFAULT_SCALAR_FIELD_NAME.values())\n @pytest.mark.parametrize(\"index_type\", SUPPORT_SCALAR_INDEX_TYPES)\n def test_duplicate_create_index(\n self, full_collection: Collection, field_name: str, index_type: IndexType\n ):\n docs = [generate_doc(i, full_collection.schema) for i in range(10)]\n\n result = full_collection.insert(docs)\n assert bool(result)\n for item in result:\n assert item.ok()\n\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 10\n\n if field_name in [\"bool_field\"]:\n query_filter = f\"{field_name} = true\"\n elif field_name in [\"double_field\", \"float_field\"]:\n query_filter = f\"{field_name} >= 3.0\"\n elif field_name in [\n \"int32_field\",\n \"int64_field\",\n \"uint32_field\",\n \"uint64_field\",\n ]:\n query_filter = f\"{field_name} >= 30\"\n elif field_name in [\"string_field\"]:\n query_filter = f\"{field_name} >= 'test_3'\"\n elif field_name in [\"array_bool_field\"]:\n query_filter = f\"{field_name} contain_any (false)\"\n elif field_name in [\"array_double_field\", \"array_float_field\"]:\n query_filter = f\"{field_name} contain_any (3.0, 4.0)\"\n elif field_name in [\n \"array_int64_field\",\n \"array_int32_field\",\n \"array_uint64_field\",\n \"array_uint32_field\",\n ]:\n query_filter = f\"{field_name} contain_any (3, 4)\"\n elif field_name == \"array_string_field\":\n query_filter = f\"{field_name} contain_any ('test_3', 'test_4')\"\n else:\n assert False, f\"Unsupported field type for index creation: {field_name}\"\n\n query_result_before = full_collection.query(filter=query_filter, topk=5)\n\n if index_type not in DEFAULT_INDEX_PARAMS:\n pytest.fail(f\"Unsupported index type for index creation: {index_type}\")\n index_param = DEFAULT_INDEX_PARAMS[index_type]\n\n full_collection.create_index(\n field_name=field_name, index_param=index_param, option=IndexOption()\n )\n\n query_result_after = full_collection.query(filter=query_filter, topk=5)\n\n assert len(query_result_before) == len(query_result_after), (\n f\"Query result count mismatch: before={len(query_result_before)}, after={len(query_result_after)}\"\n )\n\n before_ids = set(doc.id for doc in query_result_before)\n after_ids = set(doc.id for doc in query_result_after)\n assert before_ids == after_ids, (\n f\"Query result IDs mismatch: before={before_ids}, after={after_ids}\"\n )\n\n full_collection.create_index(\n field_name=field_name, index_param=index_param, option=IndexOption()\n )\n\n def test_optimize(self, full_collection: Collection):\n docs = [generate_doc(i, full_collection.schema) for i in range(10)]\n\n result = full_collection.insert(docs)\n assert bool(result)\n for item in result:\n assert item.ok()\n\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 10\n\n full_collection.optimize(option=OptimizeOption())\n\n fetched_docs = full_collection.fetch([\"1\"])\n assert \"1\" in fetched_docs\n assert fetched_docs[\"1\"].id == \"1\"\n\n @pytest.mark.parametrize(\n \"vector_type, index_type\", SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP_PARAMS\n )\n def test_vector_index_operation(\n self,\n full_collection: Collection,\n vector_type: DataType,\n index_type: IndexType,\n ):\n vector_field_name = DEFAULT_VECTOR_FIELD_NAME[vector_type]\n\n docs = [generate_doc(i, full_collection.schema) for i in range(5)]\n\n result = full_collection.insert(docs)\n assert len(result) == 5, (\n f\"Expected 5 insertion results, got {len(result)} for vector type {vector_type} and index type {index_type}\"\n )\n for i, item in enumerate(result):\n assert item.ok(), (\n 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}\"\n )\n\n stats = full_collection.stats\n assert stats is not None, (\n f\"stats is None for vector type {vector_type} and index type {index_type}\"\n )\n assert stats.doc_count == 5, (\n f\"doc_count!=5 for vector type {vector_type} and index type {index_type}\"\n )\n\n if index_type not in DEFAULT_INDEX_PARAMS:\n pytest.fail(\n f\"Unsupported index type {index_type} for vector type {vector_type} in test_vector_all_data_types_index_create_drop_validation\"\n )\n index_param = DEFAULT_INDEX_PARAMS[index_type]\n\n full_collection.create_index(\n field_name=vector_field_name,\n index_param=index_param,\n option=IndexOption(),\n )\n\n stats_after_create = full_collection.stats\n assert stats_after_create is not None, (\n f\"stats_after_create_index is None for vector type {vector_type} and index type {index_type}\"\n )\n\n new_docs = [generate_doc(i, full_collection.schema) for i in range(5, 8)]\n\n result = full_collection.insert(new_docs)\n assert len(result) == 3, (\n f\"Expected 3 insertion results, got {len(result)} for vector type {vector_type} and index type {index_type}\"\n )\n for i, item in enumerate(result):\n assert item.ok(), (\n f\"Before drop_index,result={result},BInsertion result {i} is not OK for vector type {vector_type} and index type {index_type} and \"\n )\n\n stats_after_insert1 = full_collection.stats\n assert stats_after_insert1 is not None, (\n f\"stats_after_insert1 is None for vector type {vector_type} and index type {index_type}\"\n )\n assert stats_after_insert1.doc_count == 8, (\n f\"Expected 8 documents, got {stats_after_insert1.doc_count} for vector type {vector_type} and index type {index_type}\"\n )\n\n fetched_docs = full_collection.fetch([f\"{i}\" for i in range(5, 8)])\n assert len(fetched_docs) == 3, (\n f\"Expected 3 fetched documents, got {len(fetched_docs)} for vector type {vector_type} and index type {index_type}\"\n )\n\n for i in range(5, 8):\n doc_id = f\"{i}\"\n assert doc_id in fetched_docs, (\n f\"Document ID {doc_id} not found in fetched results for vector type {vector_type} and index type {index_type}\"\n )\n assert fetched_docs[doc_id].id == doc_id, (\n f\"Document {doc_id} has incorrect ID field value for vector type {vector_type} and index type {index_type}\"\n )\n\n full_collection.drop_index(field_name=vector_field_name)\n\n more_docs = [generate_doc(i, full_collection.schema) for i in range(8, 10)]\n result = full_collection.insert(more_docs)\n assert len(result) == 2, (\n f\"Expected 2 insertion results, got {len(result)} for vector type {vector_type} and index type {index_type}\"\n )\n for i, item in enumerate(result):\n assert item.ok(), (\n f\"After drop_index,Insertion result {i} is not OK for vector type {vector_type} and index type {index_type} and result={result}\"\n )\n\n # Verify document count after second insertion\n stats_after_insert2 = full_collection.stats\n assert stats_after_insert2 is not None, (\n f\"stats_after_insert2 is None for vector type {vector_type} and index type {index_type}\"\n )\n assert stats_after_insert2.doc_count == 10, (\n f\"Expected 10 documents, got {stats_after_insert2.doc_count} for vector type {vector_type} and index type {index_type}\"\n )\n\n # Fetch data\n fetched_docs = full_collection.fetch([f\"{i}\" for i in range(8, 10)])\n assert len(fetched_docs) == 2, (\n f\"Expected 2 fetched documents, got {len(fetched_docs)} for vector type {vector_type} and index type {index_type}\"\n )\n\n # Verify fetched documents have correct data\n for i in range(8, 10):\n doc_id = f\"{i}\"\n assert doc_id in fetched_docs, (\n f\"Document ID {doc_id} not found in fetched results for vector type {vector_type} and index type {index_type}\"\n )\n assert fetched_docs[doc_id].id == doc_id, (\n f\"Document {doc_id} has incorrect ID field value for vector type {vector_type} and index type {index_type}\"\n )\n\n # Final verification\n final_stats = full_collection.stats\n assert final_stats is not None, (\n f\"final_stats is None for vector type {vector_type} and index type {index_type}\"\n )\n assert final_stats.doc_count == 10, (\n f\"Expected 10 documents, got {final_stats.doc_count} for vector type {vector_type} and index type {index_type}\"\n )\n full_collection.destroy()\n\n @staticmethod\n def create_collection(\n collection_path, collection_option: CollectionOption\n ) -> Collection:\n schema = CollectionSchema(\n name=\"test_collection_invalid_vector_index\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\",\n DataType.STRING,\n nullable=True,\n index_param=InvertIndexParam(),\n ),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n ],\n )\n coll = zvec.create_and_open(\n path=collection_path, schema=schema, option=collection_option\n )\n assert coll is not None, \"Failed to create and open collection\"\n return coll\n\n @staticmethod\n def check_error_message(exc_info, invalid_name):\n if type(invalid_name) is str:\n assert INDEX_NON_EXISTENT_COLUMN_ERROR_MSG in str(exc_info.value), (\n \"Error message is unreasonable: e=\" + str(exc_info.value)\n )\n else:\n assert INCOMPATIBLE_FUNCTION_ERROR_MSG in str(exc_info.value), (\n \"Error message is unreasonable: e=\" + str(exc_info.value)\n )\n\n @pytest.mark.parametrize(\n \"invalid_field_name,invalid_vector_name\",\n [\n (\"\", \"\"), # Empty string\n (\" \", \" \"), # Space only\n (\"v\" * 33, \"v\" * 33), # Too long (33 characters, exceeds 32)\n (\"vector name\", \"vector_name\"), # Contains space\n (\"vector@name\", \"vector@name\"), # Contains special character\n (\"vector/name\", \"vector/name\"), # Contains slash\n (\"vector\\\\name\", \"vector\\\\name\"), # Contains backslash\n (\"vector.name\", \"vector.name\"), # Contains dot\n (\"vector$data\", \"vector$data\"), # Contains dollar sign\n (\"vector+name\", \"vector+name\"), # Contains plus sign\n (\"vector=name\", \"vector=name\"), # Contains equals sign\n (None, None), # None value,\n (1, 1),\n (1.1, 1.1),\n ],\n )\n def test_invalid_field_and_vector_name(\n self,\n collection_temp_dir,\n collection_option: CollectionOption,\n invalid_field_name: Any,\n invalid_vector_name: Any,\n ):\n coll = self.create_collection(collection_temp_dir, collection_option)\n with pytest.raises(Exception) as exc_info:\n coll.create_index(\n field_name=invalid_vector_name,\n index_param=HnswIndexParam(),\n option=IndexOption(),\n )\n self.check_error_message(exc_info, invalid_vector_name)\n with pytest.raises(Exception) as exc_info:\n coll.create_index(\n field_name=invalid_field_name,\n index_param=InvertIndexParam(),\n option=IndexOption(),\n )\n self.check_error_message(exc_info, invalid_field_name)\n coll.destroy()\n coll = self.create_collection(collection_temp_dir, collection_option)\n with pytest.raises(Exception) as exc_info:\n coll.drop_index(field_name=invalid_vector_name)\n self.check_error_message(exc_info, invalid_vector_name)\n with pytest.raises(Exception) as exc_info:\n coll.drop_index(field_name=invalid_field_name)\n self.check_error_message(exc_info, invalid_field_name)\n coll.destroy()\n\n @pytest.mark.parametrize(\n \"field_name,vector_name\",\n [\n (\"2\", \"3\"),\n (\"col\", \"co1\"),\n (\"ID\", \"IM\"),\n (\"name-1\", \"name2\"),\n (\"Weigt_12\", \"Weigt_13\"),\n (\"123age\", \"123agl\"),\n ],\n )\n def test_valid_field_and_vector_name(\n self,\n collection_temp_dir,\n collection_option: CollectionOption,\n field_name: str,\n vector_name: str,\n ):\n schema = zvec.CollectionSchema(\n name=\"test_index_names\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(field_name, DataType.STRING, nullable=True),\n ],\n vectors=[\n VectorSchema(\n vector_name,\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n ],\n )\n\n coll = zvec.create_and_open(\n path=collection_temp_dir, schema=schema, option=collection_option\n )\n\n assert coll is not None, (\n f\"Failed to create and open collection with field_name={field_name}, vector_name={vector_name}\"\n )\n\n # Insert some data\n docs = [\n Doc(\n id=f\"{i}\",\n fields={\"id\": i, field_name: f\"value_{i}\"},\n vectors={vector_name: [float(j % 10) for j in range(128)]},\n )\n for i in range(5)\n ]\n\n result = coll.insert(docs)\n assert len(result) == 5, (\n f\"Expected 5 insertion results, got {len(result)} for field_name={field_name}, vector_name={vector_name}\"\n )\n for item in result:\n assert item.ok(), (\n f\"Insertion failed for field_name={field_name}, vector_name={vector_name}: {item}\"\n )\n\n # Create index on field\n coll.create_index(\n field_name=field_name,\n index_param=InvertIndexParam(),\n option=IndexOption(),\n )\n\n # Create index on vector\n coll.create_index(\n field_name=vector_name,\n index_param=HnswIndexParam(),\n option=IndexOption(),\n )\n\n # Verify indexes were created successfully\n stats = coll.stats\n assert stats is not None, (\n f\"Stats is None for field_name={field_name}, vector_name={vector_name}\"\n )\n\n coll.destroy()\n\n def test_compicated_workflow(\n self,\n collection_temp_dir,\n basic_schema: CollectionSchema,\n collection_option: CollectionOption,\n ):\n \"\"\"\n Test the complete workflow:\n 1. Create collection\n 2. Create index\n 3. Insert doc\n 4. Upsert\n 5. Update doc\n 6. Fetch doc\n 7. Query doc\n 8. Drop index\n 9. Insert doc\n 10. Update doc\n 11. Upsert doc\n 12. Fetch doc\n 13. Query doc\n 14. Flush\n 15. Destroy\n \"\"\"\n # Step 1: Create collection\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=basic_schema,\n option=collection_option,\n )\n\n assert coll is not None, \"Failed to create and open collection\"\n assert coll.path == collection_temp_dir\n assert coll.schema.name == basic_schema.name\n assert coll.stats.doc_count == 0\n\n # Step 2: Create index\n coll.create_index(\n field_name=\"name\", index_param=InvertIndexParam(), option=IndexOption()\n )\n # Verify index was created\n stats = coll.stats\n assert stats is not None, \"coll.stats is None!\"\n\n # Step 3: Insert doc\n doc1 = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test1\", \"weight\": 80.5},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n\n result = coll.insert(doc1)\n assert bool(result)\n assert result.ok()\n assert coll.stats.doc_count == 1\n\n # Step 4: Upsert (existing doc)\n doc1_updated = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test1_updated\", \"weight\": 85.0},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.5, 2: 2.5},\n },\n )\n\n result = coll.upsert(doc1_updated)\n assert bool(result)\n assert result.ok()\n assert coll.stats.doc_count == 1\n\n # Step 5: Update doc\n doc2 = Doc(\n id=\"2\",\n fields={\"id\": 2, \"name\": \"test2\", \"weight\": 90.0},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 3.0, 2: 4.0},\n },\n )\n\n # First insert doc2\n result = coll.insert(doc2)\n assert bool(result)\n assert result.ok()\n assert coll.stats.doc_count == 2\n\n # Then update it\n doc2_updated = Doc(\n id=\"2\",\n fields={\"id\": 2, \"name\": \"test2_updated\", \"weight\": 95.0},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 3.5, 2: 4.5},\n },\n )\n\n result = coll.update(doc2_updated)\n assert bool(result)\n assert result.ok()\n assert coll.stats.doc_count == 2\n\n # Step 6: Fetch doc\n fetched_docs = coll.fetch([\"1\", \"2\"])\n assert len(fetched_docs) == 2\n assert \"1\" in fetched_docs\n assert \"2\" in fetched_docs\n assert fetched_docs[\"1\"].field(\"name\") == \"test1_updated\"\n assert fetched_docs[\"2\"].field(\"name\") == \"test2_updated\"\n\n # Step 7: Query doc\n query_result = coll.query(filter=\"id >= 1\", topk=10)\n assert len(query_result) == 2\n\n # Step 8: Drop index\n coll.drop_index(field_name=\"name\")\n\n # Step 9: Insert doc\n doc3 = Doc(\n id=\"3\",\n fields={\"id\": 3, \"name\": \"test3\", \"weight\": 100.0},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 5.0, 2: 6.0},\n },\n )\n\n result = coll.insert(doc3)\n assert bool(result)\n assert result.ok()\n assert coll.stats.doc_count == 3\n\n # Step 10: Update doc\n doc3_updated = Doc(\n id=\"3\",\n fields={\"id\": 3, \"name\": \"test3_updated\", \"weight\": 105.0},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 5.5, 2: 6.5},\n },\n )\n\n result = coll.update(doc3_updated)\n assert bool(result)\n assert result.ok()\n assert coll.stats.doc_count == 3\n\n # Step 11: Upsert doc\n doc4 = Doc(\n id=\"4\",\n fields={\"id\": 4, \"name\": \"test4\", \"weight\": 110.0},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 7.0, 2: 8.0},\n },\n )\n\n result = coll.upsert(doc4)\n assert bool(result)\n assert result.ok()\n assert coll.stats.doc_count == 4\n\n # Step 12: Fetch doc\n fetched_docs = coll.fetch([\"3\", \"4\"])\n assert len(fetched_docs) == 2\n assert \"3\" in fetched_docs\n assert \"4\" in fetched_docs\n assert fetched_docs[\"3\"].field(\"name\") == \"test3_updated\"\n assert fetched_docs[\"4\"].field(\"name\") == \"test4\"\n\n # Step 13: Query doc\n query_result = coll.query(filter=\"id >= 3\", topk=10)\n assert len(query_result) == 2\n\n # Step 14: Flush\n coll.flush()\n\n # Verify data is still accessible after flush\n fetched_docs = coll.fetch([\"1\", \"2\", \"3\", \"4\"])\n assert len(fetched_docs) == 4\n\n # Step 15: Destroy\n coll.destroy()\n\n @pytest.mark.parametrize(\n \"data_type, index_param\", VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS\n )\n def test_vector_index_params(\n self,\n collection_temp_dir,\n collection_option: CollectionOption,\n data_type: DataType,\n index_param,\n single_vector_schema,\n ):\n vector_name = DEFAULT_VECTOR_FIELD_NAME[data_type]\n dimension = DEFAULT_VECTOR_DIMENSION\n\n coll = zvec.create_and_open(\n path=collection_temp_dir,\n schema=single_vector_schema,\n option=collection_option,\n )\n\n assert coll is not None, (\n f\"Failed to create and open collection, {data_type}, {index_param}\"\n )\n\n docs = {str(i): generate_doc(i, single_vector_schema) for i in range(5)}\n result = coll.insert(docs.values())\n assert len(result) == len(docs), (\n f\"Expected 5 results, got {len(result)}, {data_type}, {index_param}\"\n )\n for item in result:\n assert item.ok(), f\"Insertion failed for, {data_type}, {index_param}\"\n\n def check_result(\n label: str, metric_type: MetricType, quantize_type: QuantizeType\n ):\n query_vector = [1] * dimension\n if data_type in [DataType.SPARSE_VECTOR_FP16, DataType.SPARSE_VECTOR_FP32]:\n query_vector = {1: 1}\n\n fetch_result = coll.fetch([str(i) for i in range(len(docs))])\n assert len(fetch_result) == len(docs), (\n f\"{label}, Expected 5 fetched docs, got {len(fetch_result)}, {data_type}, {index_param}\"\n )\n for i in range(len(docs)):\n doc_id = str(i)\n assert doc_id in fetch_result, (\n f\"{label}, Document ID '{doc_id}' not found, {data_type}, {index_param}\"\n )\n fetched_doc = fetch_result[doc_id]\n # Verify doc equal\n assert is_doc_equal(fetched_doc, docs[doc_id], single_vector_schema), (\n f\"{label}, doc not equal, insert: {docs[doc_id]}, fetched: {fetched_doc}, {data_type}, {index_param}\"\n )\n\n query_result: list[Doc] = coll.query(\n VectorQuery(field_name=vector_name, vector=query_vector),\n include_vector=False,\n topk=len(docs),\n )\n assert len(query_result) == len(docs), (\n f\"{label}, Expected {len(docs)} result, got {len(query_result)}, {data_type}, {index_param}\"\n )\n inserted_ids = [str(i) for i in range(len(docs))]\n queried_ids = [doc.id for doc in query_result]\n assert set(inserted_ids) == set(queried_ids), (\n f\"{label}, inserted_ids != queried_ids, insert: {inserted_ids}, query: {queried_ids}, {data_type}, {index_param}\"\n )\n\n last_score = None\n for i, doc in enumerate(query_result):\n # Get the document's vector for comparison\n expect_doc = generate_doc(int(doc.id), single_vector_schema)\n doc_vector = expect_doc.vector(vector_name)\n expected_score = distance(\n doc_vector,\n query_vector,\n metric_type,\n data_type,\n quantize_type,\n )\n print(f\"query: {doc}, expect_core: {expected_score}\")\n if quantize_type is QuantizeType.UNDEFINED:\n assert is_float_equal(doc.score, expected_score), (\n f\"{label} top{i} pk{doc.id} score {doc.score:6f} expected:{expected_score:6f}, {data_type}, {index_param}\"\n )\n if last_score is not None:\n if metric_type == MetricType.IP:\n assert last_score >= doc.score, (\n f\"{label}, score not sorted, last_score: {last_score}, current_score: {doc.score}, {data_type}, {index_param}\"\n )\n else:\n assert last_score <= doc.score, (\n f\"{label}, score not sorted, last_score: {last_score}, current_score: {doc.score}, {data_type}, {index_param}\"\n )\n last_score = doc.score\n\n # default metric_type=IP, quantize_type=None\n check_result(\"pre_create_index\", MetricType.IP, QuantizeType.UNDEFINED)\n\n # create index\n coll.create_index(\n field_name=vector_name,\n index_param=index_param,\n option=IndexOption(),\n )\n check_result(\n \"post_create_index\", index_param.metric_type, index_param.quantize_type\n )\n\n coll.drop_index(field_name=vector_name)\n check_result(\"post_drop_index\", MetricType.IP, QuantizeType.UNDEFINED)\n\n new_docs = {str(i): generate_doc(i, single_vector_schema) for i in range(5, 8)}\n new_result = coll.insert(new_docs.values())\n assert len(new_result) == len(new_docs), (\n f\"Expected {len(new_docs)} insertion results for new docs, got {len(new_result)} for vector {vector_name}\"\n )\n for item in new_result:\n assert item.ok(), (\n f\"New document insertion failed for vector {vector_name}: {item}\"\n )\n docs |= new_docs\n coll.create_index(\n field_name=vector_name,\n index_param=index_param,\n option=IndexOption(),\n )\n\n check_result(\n \"post_create_index2\", index_param.metric_type, index_param.quantize_type\n )\n coll.destroy()\n\n\nclass TestColumnDDL:\n def test_add_column(self, basic_collection: Collection):\n basic_collection.add_column(\n field_schema=FieldSchema(\"income\", DataType.INT32),\n expression=\"'weight' * 2\", # Simple expression\n )\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, \"income\": 1},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_add_column_with_default_option(self, basic_collection: Collection):\n # Add a new column with default option\n basic_collection.add_column(\n field_schema=FieldSchema(\"test_column_default\", DataType.INT32),\n expression=\"100\",\n option=AddColumnOption(), # Default option\n )\n # Verify column was added by inserting data\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, \"test_column_default\": 1},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n # Verify document was inserted\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n @pytest.mark.parametrize(\"concurrency\", [0, 1, 4, 8])\n def test_add_column_with_various_concurrency_options(\n self, basic_collection: Collection, concurrency\n ):\n field_name = f\"test_column_concurrent_{concurrency}\"\n basic_collection.add_column(\n field_schema=FieldSchema(field_name, DataType.INT32),\n expression=\"100\",\n option=AddColumnOption(concurrency=concurrency),\n )\n\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, field_name: 200},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n @pytest.mark.parametrize(\"data_type\", SUPPORT_ADD_COLUMN_DATA_TYPE)\n def test_add_column_valid_data_types(self, basic_collection: Collection, data_type):\n field_name = f\"test_field_{data_type.name.lower()}\"\n\n # Add a new column with specific data type\n basic_collection.add_column(\n field_schema=FieldSchema(field_name, data_type),\n expression=\"1\" if data_type != DataType.STRING else \"'test'\",\n )\n\n # Verify column was added by inserting data\n if data_type == DataType.STRING:\n field_value = \"test_value\"\n elif data_type in [DataType.ARRAY_STRING]:\n field_value = [\"test_value\"]\n elif data_type in [DataType.ARRAY_INT32, DataType.ARRAY_INT64]:\n field_value = [1, 2, 3]\n elif data_type in [DataType.ARRAY_FLOAT, DataType.ARRAY_DOUBLE]:\n field_value = [1.1, 2.2, 3.3]\n elif data_type == DataType.ARRAY_BOOL:\n field_value = [True, False]\n elif data_type in [DataType.FLOAT, DataType.DOUBLE]:\n field_value = 1.5\n elif data_type in [DataType.INT32, DataType.INT64]:\n field_value = 100\n elif data_type == DataType.BOOL:\n field_value = True\n else:\n field_value = 1\n\n doc = Doc(\n id=\"1\",\n fields={\n \"id\": 1,\n \"name\": \"test\",\n \"weight\": 80.5,\n field_name: field_value,\n },\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n # Verify document was inserted\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n @pytest.mark.parametrize(\"data_type\", NOT_SUPPORT_ADD_COLUMN_DATA_TYPE)\n def test_add_column_invalid_data_types(\n self, basic_collection: Collection, data_type\n ):\n with pytest.raises(Exception) as exc_info:\n field_name = f\"test_field_{data_type.name.lower()}\"\n\n # Add a new column with specific data type\n basic_collection.add_column(\n field_schema=FieldSchema(field_name, data_type),\n expression=\"1\" if data_type != DataType.STRING else \"'test'\",\n )\n\n assert NOT_SUPPORT_ADD_COLUMN_ERROR_MSG in str(exc_info.value)\n\n @pytest.mark.parametrize(\"nullable\", [True, False])\n def test_add_column_with_nullable_options(\n self, basic_collection: Collection, nullable\n ):\n field_name = f\"test_field_nullable_{str(nullable).lower()}\"\n\n # Add a new column with specific nullable option\n basic_collection.add_column(\n field_schema=FieldSchema(field_name, DataType.INT32, nullable=nullable),\n expression=\"100\",\n )\n\n # Verify column was added by inserting data\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, field_name: 200},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n # Verify document was inserted\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n # Verify column was added by inserting data\n doc = Doc(\n id=\"2\",\n fields={\"id\": 2, \"name\": \"test\", \"weight\": 80.5, field_name: None},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n if nullable:\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n else:\n with pytest.raises(ValueError) as e:\n basic_collection.insert(doc)\n assert (\n \"Field 'test_field_nullable_false': expected non-nullable type\"\n in str(e.value)\n )\n\n # Verify document was inserted\n stats = basic_collection.stats\n assert stats is not None\n if nullable:\n assert stats.doc_count == 2\n else:\n assert stats.doc_count == 1\n\n @pytest.mark.parametrize(\n \"expression\",\n [\n \"1\", # Constant integer\n \"1.5\", # Constant float\n \"'test'\", # Constant string\n \"id\", # Reference to existing field\n \"weight * 2\", # Simple arithmetic\n \"weight + id\", # Complex arithmetic\n \"CASE WHEN weight > 50 THEN 1 ELSE 0 END\", # Conditional expression\n ],\n )\n def test_add_column_with_different_expressions(\n self, basic_collection: Collection, expression\n ):\n field_name = f\"test_field_expr_{abs(hash(expression)) % 1000}\"\n\n # Add a new column with specific expression\n basic_collection.add_column(\n field_schema=FieldSchema(field_name, DataType.INT32),\n expression=expression,\n )\n\n # Verify column was added by inserting data\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, field_name: 200},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n # Verify document was inserted\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_add_column_with_index_param(self, basic_collection: Collection):\n basic_collection.add_column(\n field_schema=FieldSchema(\n \"indexed_field\",\n DataType.INT32,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n expression=\"id * 2\",\n )\n\n # Verify column was added by inserting data\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, \"indexed_field\": 200},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n # Verify document was inserted\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n @pytest.mark.parametrize(\n \"field_name\",\n [\n \"a\", # Minimum length\n \"a\" * 32, # Maximum length (32 characters)\n \"valid_field_name_123\", # Alphanumeric with underscore\n \"Valid-Field-Name\", # With hyphens\n \"_underscore_start\", # Starting with underscore\n \"field_name_with_123_numbers\", # Numbers in middle\n \"FIELD_NAME_UPPERCASE\", # Uppercase\n # \"field_with_nums_123_and_hyphens-456\", # Complex valid name within limit\n ],\n )\n def test_add_column_with_valid_field_names(\n self, basic_collection: Collection, field_name\n ):\n basic_collection.add_column(\n field_schema=FieldSchema(field_name, DataType.INT32), expression=\"200\"\n )\n\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, field_name: 300},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n @pytest.mark.parametrize(\n \"invalid_field_name\",\n [\n \"\", # Empty string\n \" \", # Space only\n \"a\" * 33, # Too long (33 characters, exceeds 32)\n \"field name\", # Contains space\n \"field.name\", # Contains dot\n \"field@name\", # Contains special character\n \"field/name\", # Contains slash\n \"field\\\\name\", # Contains backslash\n \"field$name\", # Contains dollar sign\n \"field+name\", # Contains plus sign\n \"field=name\", # Contains equals sign\n None, # None value\n ],\n )\n def test_add_column_with_invalid_field_names(\n self, basic_collection: Collection, invalid_field_name\n ):\n with pytest.raises(Exception) as exc_info:\n basic_collection.add_column(\n field_schema=FieldSchema(invalid_field_name, DataType.INT32),\n expression=\"100\",\n )\n\n if invalid_field_name is None:\n assert \"validate failed\" in str(exc_info.value), (\n \"Error message is unreasonable: e=\" + str(exc_info.value)\n )\n else:\n assert (\n \"invalid\" in str(exc_info.value).lower()\n or \"name\" in str(exc_info.value).lower()\n )\n\n def test_alter_column_rename(self, basic_collection: Collection):\n basic_collection.alter_column(\n old_name=\"weight\",\n new_name=\"mass\",\n option=AlterColumnOption(),\n )\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"mass\": 80.5}, # Use new name\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_alter_column_non_exist(self, basic_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n basic_collection.alter_column(\n old_name=\"non_existing\",\n new_name=\"new_name\",\n field_schema=FieldSchema(\"new_name\", DataType.STRING),\n )\n assert \"column non_existing not found\" in str(exc_info.value), (\n \"Error message is unreasonable: e=\" + str(exc_info.value)\n )\n\n def test_alter_column_with_default_option(self, basic_collection: Collection):\n basic_collection.add_column(\n field_schema=FieldSchema(\"original_field\", DataType.INT32), expression=\"100\"\n )\n\n basic_collection.alter_column(\n old_name=\"original_field\",\n new_name=\"renamed_field\",\n option=AlterColumnOption(),\n )\n\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, \"renamed_field\": 200},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n @pytest.mark.parametrize(\"concurrency\", [0, 1, 4, 8])\n def test_alter_column_with_various_concurrency_options(\n self, basic_collection: Collection, concurrency\n ):\n old_field_name = f\"orig_field_{concurrency}\"\n new_field_name = f\"modified_field_{concurrency}\"\n\n basic_collection.add_column(\n field_schema=FieldSchema(old_field_name, DataType.INT32),\n expression=\"100\",\n )\n\n basic_collection.alter_column(\n old_name=old_field_name,\n new_name=new_field_name,\n option=AlterColumnOption(concurrency=concurrency),\n )\n\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, new_field_name: 200},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n @pytest.mark.parametrize(\n \"old_field_name,new_field_name\",\n [\n (\"a\", \"new_a\"), # Minimum length\n (\n \"abcdefghijklmnopqrstuvwxyz123456\",\n \"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\",\n ), # Maximum length (32 characters)\n (\"valid_field_name_123\", \"new_valid_field\"), # Alphanumeric with underscore\n (\"Valid-Field-Name\", \"New-Field-Name\"), # With hyphens\n (\"_underscore_start\", \"new_underscore\"), # Starting with underscore\n (\"field_name_with_123_numbers\", \"new_with_nums\"), # Numbers in middle\n (\"FIELD_NAME_UPPERCASE\", \"new_uppercase\"), # Uppercase\n (\n \"field_with_nums_3_and_hyphens-6\",\n \"new_field_hyphens\",\n ), # Complex valid name\n ],\n )\n def test_alter_column_field_name_valid(\n self, basic_collection: Collection, old_field_name, new_field_name\n ):\n basic_collection.add_column(\n field_schema=FieldSchema(old_field_name, DataType.INT32),\n expression=\"100\",\n )\n basic_collection.alter_column(\n old_name=old_field_name,\n new_name=new_field_name,\n option=AlterColumnOption(),\n )\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, new_field_name: 200},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n @pytest.mark.parametrize(\n \"valid_old_name,invalid_new_name\",\n [\n (\"temp_field\", \"\"), # Empty new name\n (\"temp_field\", \"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\"), # Too long new name\n (\"temp_field\", \"field name\"), # New name with space\n (\"temp_field\", \"field.name\"), # New name with dot\n (\"temp_field\", \"field@name\"), # New name with special character\n (\"temp_field\", \"field/name\"), # New name with slash\n (\"temp_field\", \"field\\\\name\"), # New name with backslash\n (\"temp_field\", \"field$name\"), # New name with dollar sign\n (\"temp_field\", \"field+name\"), # New name with plus sign\n (\"temp_field\", \"field=name\"), # New name with equals sign\n (\"temp_field\", None), # None new name\n ],\n )\n def test_alter_column_with_invalid_field_names(\n self, basic_collection: Collection, valid_old_name, invalid_new_name\n ):\n basic_collection.add_column(\n field_schema=FieldSchema(\"temp_field\", DataType.INT32), expression=\"100\"\n )\n with pytest.raises(Exception) as exc_info:\n basic_collection.alter_column(\n old_name=valid_old_name,\n new_name=invalid_new_name if invalid_new_name is not None else \"\",\n field_schema=FieldSchema(\n invalid_new_name if invalid_new_name is not None else \"\",\n DataType.INT32,\n ),\n )\n\n assert (\n \"invalid\" in str(exc_info.value).lower()\n or \"name\" in str(exc_info.value).lower()\n or \"incompatible\" in str(exc_info.value).lower()\n )\n\n def test_drop_column_exist(self, basic_collection: Collection):\n basic_collection.add_column(\n field_schema=FieldSchema(\"temp_field\", DataType.INT32), expression=\"100\"\n )\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, \"temp_field\": 1},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n basic_collection.drop_column(\"temp_field\")\n doc = Doc(\n id=\"2\",\n fields={\"id\": 2, \"name\": \"test\", \"weight\": 80.5, \"temp_field\": 1},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n with pytest.raises(Exception) as exc_info:\n result = basic_collection.insert(doc)\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value)\n\n def test_drop_column_non_exist(self, basic_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n basic_collection.drop_column(\"non_existing_column\")\n assert NOT_EXIST_COLUMN_TO_DROP_ERROR_MSG in str(exc_info.value)\n\n @pytest.mark.parametrize(\n \"field_name\",\n [\n \"a\", # Minimum length\n \"a\" * 32, # Maximum length (32 characters)\n \"valid_field_name_123\", # Alphanumeric with underscore\n \"Valid-Field-Name\", # With hyphens\n \"_underscore_start\", # Starting with underscore\n \"field_name_with_123_numbers\", # Numbers in middle\n \"FIELD_NAME_UPPERCASE\", # Uppercase\n \"field_with_nums_3_and_hyphens-6\", # Complex valid name within limit\n ],\n )\n def test_drop_column_field_name_valid(\n self, basic_collection: Collection, field_name\n ):\n basic_collection.add_column(\n field_schema=FieldSchema(field_name, DataType.INT32), expression=\"100\"\n )\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5, field_name: 200},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n\n result = basic_collection.insert(doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.ok(), (\n f\"result={result},Insert operation failed with code = {result.code()}\"\n )\n\n stats = basic_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n basic_collection.drop_column(field_name)\n\n doc = Doc(\n id=\"2\",\n fields={\"id\": 2, \"name\": \"test\", \"weight\": 80.5, field_name: 200},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n with pytest.raises(Exception) as exc_info:\n result = basic_collection.insert(doc)\n\n assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value)\n" }, { "path": "python/tests/detail/test_collection_dml.py", "content": "import logging\nimport pytest\n\n\nfrom zvec import (\n CollectionOption,\n InvertIndexParam,\n HnswIndexParam,\n FieldSchema,\n VectorSchema,\n CollectionSchema,\n Collection,\n Doc,\n VectorQuery,\n StatusCode,\n)\nfrom distance_helper import *\nfrom fixture_helper import *\nfrom doc_helper import *\n\nMaximum = 1024\n\nDOCID_VALID_LIST = [\n \"1valid_Id\",\n \"123.45\",\n \"123abc\",\n \"-!@#$%+=.123abc_+\",\n \"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ123456789012\",\n]\nDOCID_INVALID_LIST = [\n None,\n \"\",\n \"()qsd123\",\n \" \",\n \"/&AS12\",\n \"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890121\",\n]\n\nFIELD_VALUE_VALID_LIST = [\n (\n \"bool_field\",\n [\n None,\n True,\n False,\n ],\n ),\n (\n \"float_field\",\n [\n None,\n 0.0,\n -1.0,\n 1.0,\n 3.4028235e38,\n -3.4028235e38,\n 1.17549435e-38,\n -1.17549435e-38,\n float(\"inf\"),\n float(\"-inf\"),\n ],\n ),\n (\n \"double_field\",\n [\n None,\n 0.0,\n -1.0,\n 1.0,\n 1.7976931348623157e308,\n -1.7976931348623157e308,\n 2.2250738585072014e-308,\n -2.2250738585072014e-308,\n float(\"inf\"),\n float(\"-inf\"),\n ],\n ),\n (\n \"int32_field\",\n [\n None,\n 0,\n 1,\n -1,\n 2147483647,\n -2147483648,\n ],\n ),\n (\n \"int64_field\",\n [\n None,\n 0,\n 1,\n -1,\n 9223372036854775807,\n -9223372036854775808,\n ],\n ),\n (\n \"uint32_field\",\n [\n None,\n 0,\n 1,\n 4294967295,\n ],\n ),\n (\n \"uint64_field\",\n [\n None,\n 0,\n 1,\n 18446744073709551615,\n ],\n ),\n (\n \"string_field\",\n [\n None,\n \"\",\n \"a\",\n \"test_name\",\n \"这是一个中文名称测试\",\n \"a\" * 1000,\n ],\n ),\n (\n \"array_bool_field\",\n [\n None,\n [],\n [True],\n [False, True],\n [True, False, True, False] * 10,\n ],\n ),\n (\n \"array_float_field\",\n [\n None,\n [],\n [0.0],\n [1.0, 2.0, 3.0],\n [3.4028235e38, -3.4028235e38],\n ],\n ),\n (\n \"array_double_field\",\n [\n None,\n [],\n [0.0],\n [1.0, 2.0, 3.0],\n [1.7976931348623157e308, -1.7976931348623157e308],\n ],\n ),\n (\n \"array_int32_field\",\n [\n None,\n [],\n [0],\n [1, 2, 3],\n [2147483647, -2147483648],\n ],\n ),\n (\n \"array_int64_field\",\n [\n None,\n [],\n [0],\n [1, 2, 3],\n [9223372036854775807, -9223372036854775808],\n ],\n ),\n (\n \"array_uint32_field\",\n [\n None,\n [],\n [0],\n [1, 2, 3],\n [4294967295],\n ],\n ),\n (\n \"array_uint64_field\",\n [\n None,\n [],\n [0],\n [1, 2, 3],\n [18446744073709551615],\n ],\n ),\n (\n \"array_string_field\",\n [\n None,\n [],\n [\"\"],\n [\"a\", \"b\", \"c\"],\n [\"test_string\", \"测试字符串\"],\n [\"a\" * 100] * 5,\n ],\n ),\n]\nFIELD_VALUE_INVALID_LIST = [\n (\n \"bool_field\",\n [\n \"True\",\n \"False\",\n \"\",\n ],\n ),\n (\"float_field\", [\"invalid\", [1.0], {\"value\": 1.0}]),\n (\"double_field\", [\"invalid\", [1.0], {\"value\": 1.0}]),\n (\n \"int32_field\",\n [\n \"invalid\",\n [1],\n {\"value\": 1},\n 2147483648,\n -2147483649,\n ],\n ),\n (\n \"int64_field\",\n [\n \"invalid\",\n [1],\n {\"value\": 1},\n 9223372036854775808,\n -9223372036854775809,\n ],\n ),\n (\n \"uint32_field\",\n [\n \"invalid\",\n [1],\n {\"value\": 1},\n 4294967296,\n -1,\n ],\n ),\n (\n \"uint64_field\",\n [\n \"invalid\",\n [1],\n {\"value\": 1},\n 18446744073709551616,\n -1,\n ],\n ),\n (\n \"string_field\",\n [\n 123,\n 12.34,\n True,\n [\"array\"],\n {\"key\": \"value\"},\n ],\n ),\n (\n \"array_bool_field\",\n [\n True,\n False,\n [True, \"invalid\"],\n {\"key\": True},\n ],\n ),\n (\n \"array_float_field\",\n [\n [1.0, \"invalid\"],\n [1.0, None],\n \"invalid\",\n [1.0, [2.0]],\n 1.0,\n ],\n ),\n (\n \"array_double_field\",\n [\n [1.0, \"invalid\"],\n [1.0, None],\n \"invalid\",\n [1.0, [2.0]],\n 1.0,\n ],\n ),\n (\n \"array_int32_field\",\n [\n [1, \"invalid\"],\n [1, None],\n \"invalid\",\n [1, [2]],\n 1,\n ],\n ),\n (\n \"array_int64_field\",\n [\n [1, \"invalid\"],\n [1, None],\n \"invalid\",\n [1, [2]],\n 1,\n ],\n ),\n (\n \"array_uint32_field\",\n [\n [1, \"invalid\"],\n [1, None],\n [1, -1],\n \"invalid\",\n [1, [2]],\n 1,\n ],\n ),\n (\n \"array_uint64_field\",\n [\n [1, \"invalid\"],\n [1, None],\n [1, -1],\n \"invalid\",\n [1, [2]],\n 1,\n ],\n ),\n (\n \"array_string_field\",\n [\n [\"valid\", 123],\n [\"valid\", None],\n \"invalid\",\n [[\"nested\"]],\n 123,\n ],\n ),\n]\n\nVECTOR_VALUE_VALID_LIST = [\n (\n \"vector_fp32_field\",\n [\n [0.0] * 128,\n [1.0] * 128,\n [-1.0] * 128,\n [float(\"inf\")] * 128,\n [float(\"-inf\")] * 128,\n [i / 128.0 for i in range(128)],\n [-i / 128.0 for i in range(128)],\n ],\n ),\n (\n \"vector_fp16_field\",\n [\n [0.0] * 128,\n [1.0] * 128,\n [-1.0] * 128,\n [float(\"inf\")] * 128,\n [float(\"-inf\")] * 128,\n [i / 128.0 for i in range(128)],\n [-i / 128.0 for i in range(128)],\n ],\n ),\n (\"vector_int8_field\", [[100] * 128, [0] * 128, [-100] * 128]),\n (\n \"sparse_vector_fp32_field\",\n [\n {0: 1.0},\n {0: 0.0, 1: 1.0, 2: -1.0},\n {0: float(\"inf\"), 1: float(\"-inf\")},\n {i: float(i) for i in range(10)},\n {128: 1.0, 256: -1.0, 512: 0.5},\n ],\n ),\n (\n \"sparse_vector_fp16_field\",\n [\n {0: 1.0},\n {0: 0.0, 1: 1.0, 2: -1.0},\n {0: float(\"inf\"), 1: float(\"-inf\")},\n {i: float(i) for i in range(10)},\n {128: 1.0, 256: -1.0, 512: 0.5},\n ],\n ),\n]\nVECTOR_VALUE_INVALID_LIST = [\n (\n \"vector_fp32_field\",\n [\n None,\n [],\n [0.0] * 127,\n [0.0] * 129,\n [0.0] * 1000,\n [\"invalid\"],\n [0, 1, 2],\n [None] * 128,\n ],\n ),\n (\n \"vector_fp16_field\",\n [\n None,\n [],\n [0.0] * 127,\n [0.0] * 129,\n [0.0] * 1000,\n [\"invalid\"],\n [0, 1, 2],\n [None] * 128,\n ],\n ),\n (\n \"vector_int8_field\",\n [\n None,\n [],\n [1] * 127,\n [10] * 129,\n [0] * 1000,\n [\"invalid\"],\n [0, 1, 2],\n [None] * 128,\n ],\n ),\n (\n \"sparse_vector_fp32_field\",\n [\n None,\n \"invalid\",\n {None: 1.0},\n {\"0\": 1.0},\n {0: \"invalid\"},\n {0: None},\n {-1: 1.0},\n ],\n ),\n (\n \"sparse_vector_fp16_field\",\n [\n None,\n \"invalid\",\n {None: 1.0},\n {\"0\": 1.0},\n {0: \"invalid\"},\n {0: None},\n {-1: 1.0},\n ],\n ),\n]\n\nUPDATE_PARTIAL_VALUE = [\n (\n \"partial_fields\",\n {\"string_field\": \"partially_updated_test\", \"float_field\": 95.5},\n {},\n ),\n (\"dense_vector_only\", {}, {\"vector_fp32_field\": [0.3] * 128}),\n (\"dense_vector_only\", {}, {\"vector_fp16_field\": [0.6] * 128}),\n (\"dense_vector_only\", {}, {\"vector_int8_field\": [3] * 128}),\n (\"sparse_vector_only\", {}, {\"sparse_vector_fp32_field\": {1: 2.0, 2: 3.0, 4: 4.0}}),\n (\n \"sparse_vector_only\",\n {},\n {\"sparse_vector_fp16_field\": {10: 2.1, 20: 3.1, 40: 4.1}},\n ),\n (\n \"fields_and_vectors\",\n {\"string_field\": \"fully_updated_test\", \"bool_field\": False},\n {\n \"vector_fp32_field\": [0.4] * 128,\n \"sparse_vector_fp32_field\": {1: 3.0, 3: 5.0},\n },\n ),\n]\n\n\n# ==================== helper ====================\ndef singledoc_and_check(\n collection: Collection, insert_doc, operator=\"insert\", is_delete=1\n):\n if operator == \"insert\":\n result = collection.insert(insert_doc)\n elif operator == \"upsert\":\n result = collection.upsert(insert_doc)\n elif operator == \"update\":\n result = collection.update(insert_doc)\n else:\n logging.error(\"operator value is error!\")\n\n assert bool(result)\n assert result.ok()\n\n stats = collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n fetched_docs = collection.fetch([insert_doc.id])\n assert len(fetched_docs) == 1\n assert insert_doc.id in fetched_docs\n\n fetched_doc = fetched_docs[insert_doc.id]\n\n assert is_doc_equal(fetched_doc, insert_doc, collection.schema)\n assert hasattr(fetched_doc, \"score\"), \"Document should have a score attribute\"\n assert fetched_doc.score == 0.0, (\n \"Fetch operation should return default score of 0.0\"\n )\n\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v != {}:\n query_result = collection.query(\n VectorQuery(field_name=v, vector=insert_doc.vectors[v]),\n topk=10,\n )\n assert len(query_result) > 0, (\n f\"Expected at least 1 query result, but got {len(query_result)}\"\n )\n\n found_doc = None\n for doc in query_result:\n if doc.id == insert_doc.id:\n found_doc = doc\n break\n assert found_doc is not None, (\n f\"Inserted document {insert_doc.id} not found in query results\"\n )\n assert is_doc_equal(found_doc, insert_doc, collection.schema, True, False)\n if is_delete == 1:\n collection.delete(insert_doc.id)\n assert collection.stats.doc_count == 0, \"Document should be deleted\"\n\n\ndef updatedoc_partial_check(\n collection, update_doc_partial, update_doc_full, operator=\"update\", is_delete=1\n):\n if operator == \"upsert\":\n result = collection.upsert(update_doc_partial)\n elif operator == \"update\":\n result = collection.update(update_doc_partial)\n else:\n logging.error(\"operator value is error!\")\n\n assert bool(result)\n assert result.ok()\n\n stats = collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n fetched_docs = collection.fetch([update_doc_partial.id])\n assert len(fetched_docs) == 1, (\n f\"fetched_docs={fetched_docs},Expected 1 fetched document, but got {len(fetched_docs)}\"\n )\n assert update_doc_partial.id in fetched_docs, (\n f\"Expected document ID {update_doc_partial.id} in fetched documents\"\n )\n\n fetched_doc = fetched_docs[update_doc_partial.id]\n assert is_doc_equal(fetched_doc, update_doc_full, collection.schema)\n assert hasattr(fetched_doc, \"score\"), \"Document should have a score attribute\"\n assert fetched_doc.score == 0.0, (\n \"Fetch operation should return default score of 0.0\"\n )\n\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v != {}:\n query_result = collection.query(\n VectorQuery(field_name=v, vector=update_doc_full.vectors[v]),\n topk=10,\n )\n assert len(query_result) > 0, (\n f\"Expected at least 1 query result, but got {len(query_result)}\"\n )\n\n found_doc = None\n for doc in query_result:\n if doc.id == update_doc_partial.id:\n found_doc = doc\n break\n assert found_doc is not None, (\n f\"Inserted document {update_doc_partial.id} not found in query results\"\n )\n assert is_doc_equal(\n found_doc, update_doc_full, collection.schema, True, False\n )\n if is_delete == 1:\n collection.delete(update_doc_partial.id)\n assert collection.stats.doc_count == 0, \"Document should be deleted\"\n\n\ndef batchdoc_and_check(collection, multiple_docs, doc_num, operator=\"insert\"):\n if operator == \"insert\":\n result = collection.insert(multiple_docs)\n elif operator == \"upsert\":\n result = collection.upsert(multiple_docs)\n\n elif operator == \"update\":\n result = collection.update(multiple_docs)\n else:\n logging.error(\"operator value is error!\")\n\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok(), (\n f\"result={result},Insert operation failed with code {item.code()}\"\n )\n\n stats = collection.stats\n assert stats is not None, \"Collection stats should not be None\"\n assert stats.doc_count == len(multiple_docs), (\n f\"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}\"\n )\n\n doc_ids = [doc.id for doc in multiple_docs]\n fetched_docs = collection.fetch(doc_ids)\n assert len(fetched_docs) == len(multiple_docs), (\n f\"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}\"\n )\n\n for original_doc in multiple_docs:\n assert original_doc.id in fetched_docs, (\n f\"Expected document ID {original_doc.id} in fetched documents\"\n )\n fetched_doc = fetched_docs[original_doc.id]\n\n assert is_doc_equal(fetched_doc, original_doc, collection.schema)\n\n assert hasattr(fetched_doc, \"score\"), \"Document should have a score attribute\"\n assert fetched_doc.score == 0.0, (\n \"Fetch operation should return default score of 0.0\"\n )\n\n first_doc = multiple_docs[doc_num - 1]\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n query_result = collection.query(\n VectorQuery(field_name=v, vector=first_doc.vectors[v]),\n topk=1024,\n )\n assert len(query_result) > 0, (\n f\"Expected at least 1 query result, but got {len(query_result)}\"\n )\n\n found_doc = None\n\n for doc in query_result:\n if doc.id == first_doc.id:\n found_doc = doc\n break\n assert found_doc is not None, (\n f\"Inserted document {first_doc.id} not found in query results\"\n )\n\n assert is_doc_equal(found_doc, first_doc, collection.schema, True, False)\n\n\n# ==================== Tests ====================\n# ----------------------------\n# Collection Insert Test Case\n# ----------------------------\n\n\nclass TestCollectionInsert:\n def test_insert(self, full_collection: Collection):\n single_doc = generate_doc(1, full_collection.schema)\n singledoc_and_check(full_collection, single_doc)\n\n @pytest.mark.parametrize(\"doc_num\", [1, 5, Maximum])\n def test_insert_batch(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num)\n\n def test_insert_duplicate(self, full_collection: Collection):\n insert_doc = generate_doc(1, full_collection.schema)\n\n result = full_collection.insert(insert_doc)\n assert result.code().value == 0\n assert result.ok()\n\n # Verify documents were inserted\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n insert_doc_duplicate = full_collection.insert(insert_doc)\n assert bool(insert_doc_duplicate)\n assert insert_doc_duplicate.code() == StatusCode.ALREADY_EXISTS, (\n f\"Second insert operation should fail with ALREADY_EXISTS, but got code {insert_doc_duplicate.code()}\"\n )\n\n stats = full_collection.stats\n assert stats is not None, \"Collection stats should not be None\"\n assert stats.doc_count == 1, (\n f\"Document count should still be 1 after failed insert, but got {stats.doc_count}\"\n )\n\n @pytest.mark.parametrize(\"doc_id\", DOCID_VALID_LIST)\n def test_insert_docid_valid(self, full_collection: Collection, doc_id):\n insert_doc = generate_doc_random(doc_id, full_collection.schema)\n singledoc_and_check(full_collection, insert_doc)\n\n @pytest.mark.parametrize(\"doc_id\", DOCID_INVALID_LIST)\n def test_insert_docid_invalid(self, full_collection: Collection, doc_id):\n insert_doc = generate_doc_random(doc_id, full_collection.schema)\n\n with pytest.raises(Exception) as exc_info:\n full_collection.insert(insert_doc)\n\n assert exc_info.value is not None\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n @pytest.mark.parametrize(\"field_name, field_values\", FIELD_VALUE_VALID_LIST)\n @pytest.mark.parametrize(\n \"full_schema_new\",\n [(True, True, HnswIndexParam()), (False, True, HnswIndexParam())],\n indirect=True,\n )\n def test_insert_fields_valid(\n self, full_collection_new: Collection, field_name: str, field_values, request\n ):\n for i, field_value in enumerate(field_values):\n doc_id = str(field_value) if field_name == \"id\" else str(i)\n doc_fields, doc_vectors = generate_vectordict_random(\n full_collection_new.schema\n )\n full_schema_params = request.getfixturevalue(\"full_schema_new\")\n target_field = None\n for field in full_schema_params.fields:\n if field.name == field_name:\n target_field = field\n break\n doc_fields[field_name] = field_value\n insert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors)\n if target_field and not target_field.nullable and field_value is None:\n with pytest.raises(Exception) as exc_info:\n full_collection_new.insert(insert_doc)\n assert exc_info.value is not None\n else:\n singledoc_and_check(full_collection_new, insert_doc)\n\n @pytest.mark.parametrize(\"field_name, field_values\", FIELD_VALUE_INVALID_LIST)\n def test_insert_fields_invalid(\n self, full_collection: Collection, field_name: str, field_values\n ):\n for i, field_value in enumerate(field_values):\n doc_id = str(field_value) if field_name == \"id\" else str(i)\n doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)\n doc_fields[field_name] = field_value\n insert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors)\n\n with pytest.raises(Exception) as exc_info:\n full_collection.insert(insert_doc)\n assert exc_info.value is not None\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n @pytest.mark.parametrize(\"vector_field, vector_values\", VECTOR_VALUE_VALID_LIST)\n def test_insert_vector_valid(\n self, full_collection: Collection, vector_field: str, vector_values\n ):\n for i, vector_value in enumerate(vector_values):\n doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)\n\n doc_vectors[vector_field] = vector_value\n\n insert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)\n\n singledoc_and_check(full_collection, insert_doc)\n\n @pytest.mark.parametrize(\"vector_field, vector_values\", VECTOR_VALUE_INVALID_LIST)\n def test_insert_vector_invalid(\n self, full_collection: Collection, vector_field: str, vector_values\n ):\n for i, vector_value in enumerate(vector_values):\n doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)\n doc_vectors[vector_field] = vector_value\n insert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)\n with pytest.raises(Exception) as exc_info:\n full_collection.insert(insert_doc)\n\n assert exc_info.value is not None\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n\nclass TestCollectionUpdate:\n def test_update(self, full_collection: Collection):\n insert_doc = generate_doc(1, full_collection.schema)\n singledoc_and_check(full_collection, insert_doc, is_delete=0)\n updated_doc = generate_update_doc(1, full_collection.schema)\n singledoc_and_check(full_collection, updated_doc, operator=\"update\")\n\n @pytest.mark.parametrize(\"doc_num\", [1, 5, Maximum])\n def test_update_batch(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num)\n multiple_update_docs = [\n generate_update_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(\n full_collection, multiple_update_docs, doc_num, operator=\"update\"\n )\n\n def test_empty_collection_update(self, full_collection: Collection):\n updated_doc = generate_update_doc(1, full_collection.schema)\n result = full_collection.update(updated_doc)\n assert bool(result), f\"Expected 1 result, but got {len(result)}\"\n assert result.code() == StatusCode.NOT_FOUND, (\n f\"Update operation should fail with NOT_FOUND, but got code {result.code()}\"\n )\n fetched_docs = full_collection.fetch([updated_doc.id])\n assert len(fetched_docs) == 0\n\n stats = full_collection.stats\n assert stats is not None, \"Collection stats should not be None\"\n assert stats.doc_count == 0, (\n f\"Document count should be 0, but got {stats.doc_count}\"\n )\n\n @pytest.mark.parametrize(\"doc_num\", [1, 5, Maximum])\n def test_empty_collection_update_batch(self, full_collection: Collection, doc_num):\n multiple_update_docs = [\n generate_update_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n result = full_collection.update(multiple_update_docs)\n assert len(result) == len(multiple_update_docs), (\n f\"Expected {len(multiple_update_docs)} results, but got {len(result)}\"\n )\n for item in result:\n assert item.code() == StatusCode.NOT_FOUND, (\n f\"Update operation should fail with NOT_FOUND, but got code {item.code()}\"\n )\n\n stats = full_collection.stats\n assert stats is not None, \"Collection stats should not be None\"\n assert stats.doc_count == 0, (\n f\"Document count should be 0, but got {stats.doc_count}\"\n )\n\n doc_ids = [doc.id for doc in multiple_update_docs]\n fetched_docs = full_collection.fetch(doc_ids)\n assert len(fetched_docs) == 0\n\n @pytest.mark.parametrize(\"field_name, field_values\", FIELD_VALUE_VALID_LIST)\n @pytest.mark.parametrize(\n \"full_schema_new\",\n [(True, True, HnswIndexParam()), (False, True, HnswIndexParam())],\n indirect=True,\n )\n def test_update_fields_valid(\n self, full_collection_new: Collection, field_name: str, field_values, request\n ):\n for i, field_value in enumerate(field_values):\n insert_doc = generate_doc(i, full_collection_new.schema)\n singledoc_and_check(full_collection_new, insert_doc, is_delete=0)\n update_doc_fields, update_doc_vectors = generate_vectordict_random(\n full_collection_new.schema\n )\n full_schema_params = request.getfixturevalue(\"full_schema_new\")\n target_field = None\n for field in full_schema_params.fields:\n if field.name == field_name:\n target_field = field\n break\n update_doc_fields[field_name] = field_value\n update_doc = Doc(\n id=str(i), fields=update_doc_fields, vectors=update_doc_vectors\n )\n if target_field and not target_field.nullable and field_value is None:\n with pytest.raises(Exception) as exc_info:\n update_doc_fields[field_name] = field_value\n full_collection_new.update(update_doc)\n assert exc_info.value is not None\n full_collection_new.delete(insert_doc.id)\n else:\n singledoc_and_check(\n full_collection_new, update_doc, operator=\"update\", is_delete=1\n )\n\n @pytest.mark.parametrize(\"field_name, field_values\", FIELD_VALUE_INVALID_LIST)\n def test_update_fields_invalid(\n self, full_collection: Collection, field_name: str, field_values\n ):\n for i, field_value in enumerate(field_values):\n insert_doc = generate_doc(i, full_collection.schema)\n singledoc_and_check(full_collection, insert_doc, is_delete=0)\n update_doc_fields, update_doc_vectors = generate_vectordict_random(\n full_collection.schema\n )\n update_doc_fields[field_name] = field_value\n update_doc = Doc(\n id=str(i), fields=update_doc_fields, vectors=update_doc_vectors\n )\n\n with pytest.raises(Exception) as exc_info:\n full_collection.update(update_doc)\n\n assert exc_info.value is not None\n full_collection.delete(insert_doc.id)\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n @pytest.mark.parametrize(\"vector_field, vector_values\", VECTOR_VALUE_VALID_LIST)\n def test_update_doc_vector_valid(\n self,\n full_collection: Collection,\n collection_temp_dir,\n collection_option,\n vector_field: str,\n vector_values,\n ):\n for i, vector_value in enumerate(vector_values):\n insert_doc = generate_doc(i, full_collection.schema)\n singledoc_and_check(full_collection, insert_doc, is_delete=0)\n update_doc_fields, update_doc_vectors = generate_vectordict_random(\n full_collection.schema\n )\n update_doc_vectors[vector_field] = vector_value\n update_doc = Doc(\n id=str(i), fields=update_doc_fields, vectors=update_doc_vectors\n )\n singledoc_and_check(full_collection, update_doc, operator=\"update\")\n\n @pytest.mark.parametrize(\"vector_field, vector_values\", VECTOR_VALUE_INVALID_LIST)\n def test_update_doc_vector_invalid(\n self,\n full_collection: Collection,\n collection_temp_dir,\n collection_option,\n vector_field: str,\n vector_values,\n ):\n for i, vector_value in enumerate(vector_values):\n insert_doc = generate_doc(i, full_collection.schema)\n singledoc_and_check(full_collection, insert_doc, is_delete=0)\n update_doc_fields, update_doc_vectors = generate_vectordict_random(\n full_collection.schema\n )\n update_doc_vectors[vector_field] = vector_value\n update_doc = Doc(\n id=str(i), fields=update_doc_fields, vectors=update_doc_vectors\n )\n with pytest.raises(Exception) as exc_info:\n full_collection.update(update_doc)\n assert exc_info.value is not None\n full_collection.delete(insert_doc.id)\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n @pytest.mark.parametrize(\n \"update_type, fields_to_update, vectors_to_update\", UPDATE_PARTIAL_VALUE\n )\n def test_update_partial_fields(\n self,\n full_collection: Collection,\n collection_temp_dir,\n collection_option,\n update_type: str,\n fields_to_update: dict,\n vectors_to_update: dict,\n doc_id=1,\n ):\n insert_doc = generate_doc(doc_id, full_collection.schema)\n singledoc_and_check(full_collection, insert_doc, is_delete=0)\n\n update_doc_fields, update_doc_vectors = insert_doc.fields, insert_doc.vectors\n for k, v in fields_to_update.items():\n update_doc_fields[k] = v\n for k, v in vectors_to_update.items():\n update_doc_vectors[k] = v\n\n update_doc_full = Doc(\n id=str(doc_id), fields=update_doc_fields, vectors=update_doc_vectors\n )\n\n update_doc_partial = Doc(\n id=str(doc_id), fields=fields_to_update, vectors=vectors_to_update\n )\n\n updatedoc_partial_check(\n full_collection,\n update_doc_partial,\n update_doc_full,\n operator=\"update\",\n is_delete=1,\n )\n\n\nclass TestCollectionUpsert:\n def test_new_doc_upsert(self, full_collection: Collection):\n single_doc = generate_doc(1, full_collection.schema)\n singledoc_and_check(full_collection, single_doc, operator=\"upsert\", is_delete=1)\n\n @pytest.mark.parametrize(\"doc_num\", [1, 5, Maximum])\n def test_new_doc_upsert_batch(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"upsert\")\n\n def test_existing_doc_upsert(self, full_collection: Collection):\n insert_doc = generate_doc(1, full_collection.schema)\n singledoc_and_check(full_collection, insert_doc, is_delete=0)\n updated_doc = generate_update_doc(1, full_collection.schema)\n singledoc_and_check(full_collection, updated_doc, operator=\"upsert\")\n\n @pytest.mark.parametrize(\"doc_id\", DOCID_VALID_LIST)\n def test_upsert_docid_valid(self, full_collection: Collection, doc_id):\n upsert_doc = generate_doc_random(doc_id, full_collection.schema)\n singledoc_and_check(full_collection, upsert_doc, operator=\"upsert\", is_delete=1)\n\n @pytest.mark.parametrize(\"doc_id\", DOCID_INVALID_LIST)\n def test_upsert_docid_invalid(self, full_collection: Collection, doc_id):\n upsert_doc = generate_doc_random(doc_id, full_collection.schema)\n\n with pytest.raises(Exception) as exc_info:\n full_collection.upsert(upsert_doc)\n\n assert exc_info.value is not None\n\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n @pytest.mark.parametrize(\"field_name, field_values\", FIELD_VALUE_VALID_LIST)\n @pytest.mark.parametrize(\n \"full_schema_new\",\n [(True, True, HnswIndexParam()), (False, True, HnswIndexParam())],\n indirect=True,\n )\n def test_upsert_fields_valid(\n self, full_collection_new: Collection, field_name: str, field_values, request\n ):\n for i, field_value in enumerate(field_values):\n doc_id = str(field_value) if field_name == \"id\" else str(i)\n doc_fields, doc_vectors = generate_vectordict_random(\n full_collection_new.schema\n )\n\n full_schema_params = request.getfixturevalue(\"full_schema_new\")\n target_field = None\n for field in full_schema_params.fields:\n if field.name == field_name:\n target_field = field\n break\n doc_fields[field_name] = field_value\n upsert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors)\n if target_field and not target_field.nullable and field_value is None:\n with pytest.raises(Exception) as exc_info:\n full_collection_new.upsert(upsert_doc)\n assert exc_info.value is not None\n else:\n singledoc_and_check(\n full_collection_new, upsert_doc, operator=\"upsert\", is_delete=1\n )\n\n @pytest.mark.parametrize(\"field_name, field_values\", FIELD_VALUE_INVALID_LIST)\n def test_upsert_fields_invalid(\n self, full_collection: Collection, field_name: str, field_values\n ):\n for i, field_value in enumerate(field_values):\n doc_id = str(field_value) if field_name == \"id\" else str(i)\n doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)\n doc_fields[field_name] = field_value\n upsert_doc = Doc(id=doc_id, fields=doc_fields, vectors=doc_vectors)\n\n with pytest.raises(Exception) as exc_info:\n full_collection.upsert(upsert_doc)\n assert exc_info.value is not None\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n @pytest.mark.parametrize(\"vector_field, vector_values\", VECTOR_VALUE_VALID_LIST)\n def test_upsert_vector_valid(\n self, full_collection: Collection, vector_field: str, vector_values\n ):\n for i, vector_value in enumerate(vector_values):\n doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)\n\n doc_vectors[vector_field] = vector_value\n\n upsert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)\n\n singledoc_and_check(\n full_collection, upsert_doc, operator=\"upsert\", is_delete=1\n )\n\n @pytest.mark.parametrize(\"vector_field, vector_values\", VECTOR_VALUE_INVALID_LIST)\n def test_upsert_vector_invalid(\n self, full_collection: Collection, vector_field: str, vector_values\n ):\n for i, vector_value in enumerate(vector_values):\n doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)\n doc_vectors[vector_field] = vector_value\n upsert_doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)\n with pytest.raises(Exception) as exc_info:\n full_collection.upsert(upsert_doc)\n\n assert exc_info.value is not None\n stats = full_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n\nclass TestCollectionDelete:\n @pytest.mark.parametrize(\"doc_num\", [1, 5, Maximum])\n def test_delete_batch(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n\n doc_ids = [doc.id for doc in multiple_docs]\n result = full_collection.delete(doc_ids)\n assert len(result) == len(doc_ids)\n for item in result:\n assert item.ok()\n\n def test_delete_non_exist(self, full_collection: Collection):\n result = full_collection.delete(\"non_existing_id\")\n assert result.code().value == 1\n assert result.code() == StatusCode.NOT_FOUND\n\n @pytest.mark.parametrize(\"doc_num\", [5])\n def test_delete_batch_part_non_exist(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n doc_ids = [doc.id for doc in multiple_docs]\n doc_ids.extend([str(doc_num), str(doc_num + 1)])\n result = full_collection.delete(doc_ids)\n\n assert len(result) == len(doc_ids)\n for i in range(len(result)):\n if i < doc_num:\n assert result[i].ok()\n else:\n assert result[i].code().value == 1\n assert result[i].code() == StatusCode.NOT_FOUND\n\n @pytest.mark.parametrize(\"doc_num\", [5])\n def test_delete_by_filter(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n\n result = full_collection.delete_by_filter(\"int32_field > 0\")\n assert result is None\n\n def test_delete_empty_ids(self, full_collection: Collection):\n result = full_collection.delete([])\n assert len(result) == 0\n" }, { "path": "python/tests/detail/test_collection_dql.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom zvec.typing import DataType, StatusCode, MetricType, QuantizeType\nfrom zvec.model import Collection, Doc, VectorQuery\nfrom zvec.model.param import (\n CollectionOption,\n InvertIndexParam,\n HnswIndexParam,\n FlatIndexParam,\n IVFIndexParam,\n HnswQueryParam,\n IVFQueryParam,\n)\n\n\nfrom zvec.model.schema import FieldSchema, VectorSchema\nfrom zvec.extension import RrfReRanker, WeightedReRanker, QwenReRanker\nfrom distance_helper import *\n\nfrom zvec import StatusCode\nfrom distance_helper import *\nfrom fixture_helper import *\nfrom doc_helper import *\nfrom params_helper import *\n\n\n# ==================== helper ====================\ndef batchdoc_and_check(\n collection: Collection, multiple_docs, doc_num, operator=\"insert\"\n):\n if operator == \"insert\":\n result = collection.insert(multiple_docs)\n elif operator == \"upsert\":\n result = collection.upsert(multiple_docs)\n\n elif operator == \"update\":\n result = collection.update(multiple_docs)\n else:\n logging.error(\"operator value is error!\")\n\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok(), (\n f\"result={result},Insert operation failed with code {item.code()}\"\n )\n\n stats = collection.stats\n assert stats is not None, \"Collection stats should not be None\"\n assert stats.doc_count == len(multiple_docs), (\n f\"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}\"\n )\n\n doc_ids = [doc.id for doc in multiple_docs]\n fetched_docs = collection.fetch(doc_ids)\n assert len(fetched_docs) == len(multiple_docs), (\n f\"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}\"\n )\n\n for original_doc in multiple_docs:\n assert original_doc.id in fetched_docs, (\n f\"Expected document ID {original_doc.id} in fetched documents\"\n )\n fetched_doc = fetched_docs[original_doc.id]\n\n assert is_doc_equal(fetched_doc, original_doc, collection.schema)\n\n assert hasattr(fetched_doc, \"score\"), \"Document should have a score attribute\"\n assert fetched_doc.score == 0.0, (\n \"Fetch operation should return default score of 0.0\"\n )\n\n first_doc = multiple_docs[doc_num - 1]\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n query_result = collection.query(\n VectorQuery(field_name=v, vector=first_doc.vectors[v]),\n topk=1024,\n include_vector=True,\n )\n assert len(query_result) > 0, (\n f\"Expected at least 1 query result, but got {len(query_result)}\"\n )\n\n found_doc = None\n\n for doc in query_result:\n if doc.id == first_doc.id:\n found_doc = doc\n break\n assert found_doc is not None, (\n f\"Inserted document {first_doc.id} not found in query results\"\n )\n\n assert is_doc_equal(found_doc, first_doc, collection.schema)\n assert hasattr(found_doc, \"score\")\n assert isinstance(found_doc.score, (int, float))\n\n\ndef batchdoc_and_check_ivf(\n collection: Collection, multiple_docs, doc_num, operator=\"insert\"\n):\n if operator == \"insert\":\n result = collection.insert(multiple_docs)\n elif operator == \"upsert\":\n result = collection.upsert(multiple_docs)\n\n elif operator == \"update\":\n result = collection.update(multiple_docs)\n else:\n logging.error(\"operator value is error!\")\n\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok(), (\n f\"result={result},Insert operation failed with code {item.code()}\"\n )\n\n stats = collection.stats\n assert stats is not None, \"Collection stats should not be None\"\n assert stats.doc_count == len(multiple_docs), (\n f\"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}\"\n )\n\n doc_ids = [doc.id for doc in multiple_docs]\n fetched_docs = collection.fetch(doc_ids)\n assert len(fetched_docs) == len(multiple_docs), (\n f\"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}\"\n )\n\n for original_doc in multiple_docs:\n assert original_doc.id in fetched_docs, (\n f\"Expected document ID {original_doc.id} in fetched documents\"\n )\n fetched_doc = fetched_docs[original_doc.id]\n\n assert is_doc_equal(fetched_doc, original_doc, collection.schema)\n\n assert hasattr(fetched_doc, \"score\"), \"Document should have a score attribute\"\n assert fetched_doc.score == 0.0, (\n \"Fetch operation should return default score of 0.0\"\n )\n\n first_doc = multiple_docs[doc_num - 1]\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v in [\"vector_fp16_field\", \"vector_fp32_field\"]:\n query_result = collection.query(\n VectorQuery(field_name=v, vector=first_doc.vectors[v]),\n topk=1024,\n include_vector=True,\n )\n assert len(query_result) > 0, (\n f\"Expected at least 1 query result, but got {len(query_result)}\"\n )\n\n found_doc = None\n\n for doc in query_result:\n if doc.id == first_doc.id:\n found_doc = doc\n break\n assert found_doc is not None, (\n f\"Inserted document {first_doc.id} not found in query results\"\n )\n\n assert is_doc_equal(found_doc, first_doc, collection.schema)\n assert hasattr(found_doc, \"score\")\n assert isinstance(found_doc.score, (int, float))\n\n\ndef single_querydoc_check(\n multiple_docs,\n query_result,\n full_collection: Collection,\n is_by_vector=0,\n query_vector=None,\n data_type=None,\n vector_name=None,\n metric_type=MetricType.IP,\n id_include_vector: bool = False,\n is_output_fields=0,\n):\n for original_doc in multiple_docs:\n for doc in query_result:\n if doc.id == original_doc.id:\n found_doc = doc\n if is_output_fields == 0:\n assert is_doc_equal(\n found_doc,\n original_doc,\n full_collection.schema,\n True,\n id_include_vector,\n )\n assert hasattr(found_doc, \"score\")\n # assert found_doc.score >= 0.0\n if not id_include_vector:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n assert found_doc.vector(v) == {}\n else:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n assert found_doc.vector(v) != {}\n if is_by_vector:\n prev_score = float(\"inf\")\n for i, doc in enumerate(query_result):\n doc_vector = full_collection.fetch(doc.id)[doc.id].vector(\n vector_name\n )\n expected_score = distance(\n query_vector, doc_vector, metric_type, data_type, k\n )\n if (\n full_collection.schema.vector(vector_name).data_type\n != DataType.VECTOR_FP16\n ):\n assert abs(doc.score - expected_score) < 0.001, (\n f\"{data_type} {vector_name} :Expected score {expected_score:.6f}, but got {doc.score:.6f} for document {doc.id}\"\n )\n assert doc.score <= prev_score, (\n f\"{data_type} {vector_name} :Scores should be in descending order. Current: {doc.score}, Previous: {prev_score}\"\n )\n prev_score = doc.score\n\n\ndef multi_querydoc_check(multiple_docs, query_result, full_collection):\n for original_doc in multiple_docs:\n for doc in query_result:\n if doc.id == original_doc.id:\n found_doc = doc\n assert is_doc_equal(\n found_doc, original_doc, full_collection.schema, False, False\n )\n assert hasattr(found_doc, \"score\"), (\n \"Document should have a score attribute\"\n )\n assert found_doc.score >= 0.0, (\n \"Fetch operation should return default score of 0.0\"\n )\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n assert found_doc.vector(v) == {}\n\n\n# ==================== Tests ====================\nclass TestCollectionFetch:\n def test_fetch_non_existing(self, full_collection: Collection):\n result = full_collection.fetch(ids=[\"non_existing_id1\", \"non_existing_id2\"])\n assert len(result) == 0\n\n @pytest.mark.parametrize(\"doc_num\", [3])\n def test_fetch_partial_non_existing(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n\n fetch_id_list = [doc.id for doc in multiple_docs]\n fetch_id_list.append(\"non_existing_id\")\n result = full_collection.fetch(ids=fetch_id_list)\n\n assert len(result) == doc_num\n assert \"non_existing_id\" not in result.keys()\n\n def test_fetch_empty_ids(self, full_collection: Collection):\n result = full_collection.fetch(ids=[])\n assert len(result) == 0, (\n f\"Expected 0 results for empty ID list, but got {len(result)}\"\n )\n\n\nclass TestCollectionQuery:\n @pytest.mark.parametrize(\"doc_num\", [5])\n def test_query_with_no_condition(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n query_result = full_collection.query()\n assert len(query_result) == doc_num\n single_querydoc_check(multiple_docs, query_result, full_collection)\n\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_with_filter_empty(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n result1 = full_collection.query(filter=\"\")\n assert len(result1) == doc_num\n single_querydoc_check(multiple_docs, result1, full_collection)\n result2 = full_collection.query(filter=None)\n assert len(result2) == doc_num\n single_querydoc_check(multiple_docs, result2, full_collection)\n ids1 = set(doc.id for doc in result1)\n ids2 = set(doc.id for doc in result2)\n assert ids1 == ids2\n\n @pytest.mark.parametrize(\"field_name\", [\"int32_field\"])\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_with_filter_single_condition(\n self, full_collection: Collection, doc_num, field_name\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n filter = field_name + \" > 5\"\n query_result = full_collection.query(filter=filter)\n assert len(query_result) == doc_num - 6\n\n returned_doc_ids = set()\n for doc in query_result:\n returned_doc_ids.add(doc.id)\n\n expected_doc_ids = set(str(i) for i in range(6, doc_num))\n\n for doc in query_result:\n assert doc.id in expected_doc_ids\n assert int(doc.field(field_name)) > 5\n\n single_querydoc_check(multiple_docs, query_result, full_collection)\n\n @pytest.mark.parametrize(\"field_name\", [\"int32_field\"])\n @pytest.mark.parametrize(\n \"filter\",\n [\n \"int32_field > 3 and int32_field < 9\",\n \"int32_field >= 5 and int32_field <= 7\",\n ],\n )\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_with_filter_and(\n self, full_collection: Collection, doc_num, field_name, filter\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n filter = field_name + \" > 3 and \" + field_name + \" < 9\"\n query_result = full_collection.query(filter=filter)\n if filter == \"int32_field > 3 and int32_field < 9\":\n assert len(query_result) == doc_num - 4 - 1\n expected_doc_ids = set(str(i) for i in range(4, 9))\n\n for doc in query_result:\n assert doc.id in expected_doc_ids\n field_value = int(doc.field(field_name))\n assert field_value > 3 and field_value < 9\n else:\n assert len(query_result) == 3\n expected_doc_ids = set(str(i) for i in range(5, 8))\n\n for doc in query_result:\n assert doc.id in expected_doc_ids\n field_value = int(doc.field(field_name))\n assert field_value >= 5 and field_value <= 7\n\n single_querydoc_check(multiple_docs, query_result, full_collection)\n\n @pytest.mark.parametrize(\"field_name\", [\"int32_field\"])\n @pytest.mark.parametrize(\n \"filter\",\n [\n \"int32_field < 3 or int32_field > 8\",\n \"int32_field = 3 or int32_field = 7\",\n \"int32_field <= 3 or int32_field >= 8\",\n ],\n )\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_with_filter_or(\n self, full_collection: Collection, doc_num, field_name, filter\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n query_result = full_collection.query(filter=filter)\n if filter == \"int32_field < 3 or int32_field > 8\":\n assert len(query_result) == 4\n expected_doc_ids = set([str(0), str(1), str(2), str(9)])\n for doc in query_result:\n assert doc.id in expected_doc_ids\n field_value = int(doc.field(field_name))\n assert field_value < 3 or field_value > 8\n elif filter == \"int32_field = 3 or int32_field = 7\":\n assert len(query_result) == 2\n expected_doc_ids = set([str(3), str(7)])\n for doc in query_result:\n assert doc.id in expected_doc_ids\n field_value = int(doc.field(field_name))\n assert field_value == 3 or field_value == 7\n else:\n assert len(query_result) == 6\n expected_doc_ids = set([str(0), str(1), str(2), str(3), str(8), str(9)])\n for doc in query_result:\n assert doc.id in expected_doc_ids\n field_value = int(doc.field(field_name))\n assert field_value <= 3 or field_value >= 8\n\n single_querydoc_check(multiple_docs, query_result, full_collection)\n\n @pytest.mark.parametrize(\"field_names\", [(\"int32_field\", \"bool_field\")])\n @pytest.mark.parametrize(\n \"filter\",\n [\n \"(int32_field < 3 or int32_field > 8) and bool_field = false\",\n \"(int32_field > 2 and int32_field < 5) or (int32_field > 7 and bool_field = true)\",\n ],\n )\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_with_filter_parentheses(\n self, full_collection: Collection, doc_num, field_names, filter\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n query_result = full_collection.query(filter=filter)\n if filter == \"(int32_field < 3 or int32_field > 8) and bool_field = false\":\n assert len(query_result) == 2\n expected_doc_ids = set([str(1), str(9)])\n for doc in query_result:\n assert doc.id in expected_doc_ids\n assert (\n int(doc.field(field_names[0])) < 3\n or int(doc.field(field_names[0])) > 8\n ) and doc.field(field_names[1]) == False\n else:\n assert len(query_result) == 3\n expected_doc_ids = set([str(3), str(4), str(8)])\n for doc in query_result:\n assert doc.id in expected_doc_ids\n assert (\n (\n int(doc.field(field_names[0])) > 2\n and int(doc.field(field_names[0])) < 5\n )\n or (doc.field(field_names[0])) > 7\n and doc.field(field_names[1]) == True\n )\n single_querydoc_check(multiple_docs, query_result, full_collection)\n\n @pytest.mark.parametrize(\n \"filter\",\n [\n \"int32_field >\",\n \"int32_field = 'string'\",\n \"nonexistent_field = 5\",\n \"int32_field > 5 and\",\n \"int32_field > > 5\",\n ],\n )\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_filter_invalid(self, full_collection: Collection, doc_num, filter):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n with pytest.raises(Exception) as exc_info:\n full_collection.query(filter=filter)\n if filter in [\"int32_field = 'string'\", \"nonexistent_field = 5\"]:\n assert \"Analyze sql info failed\" in str(exc_info.value)\n else:\n assert \"Invalid filter\" in str(exc_info.value)\n\n @pytest.mark.parametrize(\"field_name\", [\"int32_field\"])\n @pytest.mark.parametrize(\"topk_value\", [1, 5, 10, 50, 100, 500, 1000, 1024])\n def test_query_with_filter_topk_valid(\n self, full_collection: Collection, topk_value: int, field_name\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(topk_value)\n ]\n batchdoc_and_check(\n full_collection, multiple_docs, topk_value, operator=\"insert\"\n )\n filter = (\n field_name + f\" >={topk_value - 1} and \" + field_name + f\" <={topk_value}\"\n )\n print(\"filter:\\n\")\n print(filter)\n query_result = full_collection.query(filter=filter, topk=topk_value)\n assert len(query_result) == 1\n expected_doc_ids = [str(topk_value - 1)]\n\n for doc in query_result:\n assert doc.id in expected_doc_ids\n field_value = int(doc.field(field_name))\n assert field_value >= topk_value - 1 and field_value <= topk_value\n single_querydoc_check(multiple_docs, query_result, full_collection)\n\n @pytest.mark.parametrize(\"field_name\", [\"int32_field\"])\n @pytest.mark.parametrize(\"topk_value\", [1, 5, 10, 50, 100, 500, 1000, 1024])\n def test_query_without_filter_topk_valid(\n self, full_collection: Collection, topk_value: int, field_name\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(topk_value)\n ]\n batchdoc_and_check(\n full_collection, multiple_docs, topk_value, operator=\"insert\"\n )\n\n query_result = full_collection.query(topk=topk_value)\n assert len(query_result) == topk_value\n single_querydoc_check(multiple_docs, query_result, full_collection)\n\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_with_include_vector(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n\n query_result = full_collection.query(include_vector=True)\n assert len(query_result) > 0\n single_querydoc_check(\n multiple_docs, query_result, full_collection, id_include_vector=1\n )\n\n @pytest.mark.parametrize(\"output_fields\", [[\"int32_field\", \"int64_field\"]])\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_with_output_fields(\n self, full_collection: Collection, doc_num, output_fields\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n query_result = full_collection.query(output_fields=output_fields)\n assert len(query_result) > 0\n for doc in query_result:\n field_names = doc.field_names()\n assert field_names == output_fields\n\n @pytest.mark.parametrize(\n \"filter\",\n [\n \"int32_field >= 10 and int32_field <= 20\",\n \"int32_field = 3 and int32_field = 8\",\n ],\n )\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_empty_result(self, full_collection: Collection, doc_num, filter):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n result = full_collection.query(filter=filter)\n assert len(result) == 0\n\n @pytest.mark.parametrize(\n \"full_schema_new\",\n [(True, True, HnswIndexParam()), (False, True, FlatIndexParam())],\n indirect=True,\n )\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_by_id(\n self, full_collection_new: Collection, doc_num, full_schema_new\n ):\n multiple_docs = [\n generate_doc(i, full_collection_new.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(\n full_collection_new, multiple_docs, doc_num, operator=\"insert\"\n )\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n query_result = full_collection_new.query(VectorQuery(field_name=v, id=\"1\"))\n assert len(query_result) > 0\n query_doc = full_collection_new.fetch(ids=[\"1\"])\n query_vector = query_doc[\"1\"].vector(v)\n single_querydoc_check(\n multiple_docs,\n query_result,\n full_collection_new,\n is_by_vector=1,\n query_vector=query_vector,\n data_type=k,\n vector_name=v,\n )\n\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_by_id_ivf(self, full_collection_ivf: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection_ivf.schema) for i in range(doc_num)\n ]\n batchdoc_and_check_ivf(\n full_collection_ivf, multiple_docs, doc_num, operator=\"insert\"\n )\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v in [\"vector_fp16_field\", \"vector_fp32_field\"]:\n query_result = full_collection_ivf.query(\n VectorQuery(field_name=v, id=\"1\")\n )\n assert len(query_result) > 0\n query_doc = full_collection_ivf.fetch(ids=[\"1\"])\n query_vector = query_doc[\"1\"].vector(v)\n single_querydoc_check(\n multiple_docs,\n query_result,\n full_collection_ivf,\n is_by_vector=1,\n query_vector=query_vector,\n data_type=k,\n vector_name=v,\n )\n\n @pytest.mark.parametrize(\n \"full_schema_new\",\n [(True, True, HnswIndexParam()), (False, True, FlatIndexParam())],\n indirect=True,\n )\n @pytest.mark.parametrize(\"doc_num\", [10])\n @pytest.mark.parametrize(\"topk\", [None, 1024])\n @pytest.mark.parametrize(\"filter\", [None, \"int32_field >= 3 and int32_field <= 7\"])\n def test_query_by_vector(\n self, full_collection_new: Collection, doc_num, full_schema_new, topk, filter\n ):\n multiple_docs = [\n generate_doc(i, full_collection_new.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(\n full_collection_new, multiple_docs, doc_num, operator=\"insert\"\n )\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n doc_fields, doc_vectors = generate_vectordict_random(\n full_collection_new.schema\n )\n query_vector = doc_vectors[v]\n if topk and filter:\n query_result = full_collection_new.query(\n filter=filter,\n vectors=VectorQuery(field_name=v, vector=query_vector),\n topk=topk,\n )\n elif topk and not filter:\n query_result = full_collection_new.query(\n VectorQuery(field_name=v, vector=query_vector), topk=topk\n )\n elif not topk and filter:\n query_result = full_collection_new.query(\n filter=filter,\n vectors=VectorQuery(field_name=v, vector=query_vector),\n )\n else:\n query_result = full_collection_new.query(\n VectorQuery(field_name=v, vector=query_vector)\n )\n assert len(query_result) > 0, (\n f\"Expected at least 1 query result, but got {len(query_result)}\"\n )\n single_querydoc_check(\n multiple_docs,\n query_result,\n full_collection_new,\n is_by_vector=1,\n query_vector=query_vector,\n data_type=k,\n vector_name=v,\n )\n\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_by_vector_ivf(self, full_collection_ivf: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection_ivf.schema) for i in range(doc_num)\n ]\n batchdoc_and_check_ivf(\n full_collection_ivf, multiple_docs, doc_num, operator=\"insert\"\n )\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n if v in [\"vector_fp16_field\", \"vector_fp32_field\"]:\n doc_fields, doc_vectors = generate_vectordict_random(\n full_collection_ivf.schema\n )\n query_vector = doc_vectors[v]\n query_result = full_collection_ivf.query(\n VectorQuery(field_name=v, vector=query_vector),\n topk=1024,\n )\n assert len(query_result) > 0, (\n f\"Expected at least 1 query result, but got {len(query_result)}\"\n )\n single_querydoc_check(\n multiple_docs,\n query_result,\n full_collection_ivf,\n is_by_vector=1,\n query_vector=query_vector,\n data_type=k,\n vector_name=v,\n )\n\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_multivector_rrf(self, full_collection: Collection, doc_num):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n\n doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)\n single_query_results = {}\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n single_query_results[v] = full_collection.query(\n VectorQuery(field_name=v, vector=doc_vectors[v])\n )\n expected_rrf_scores = calculate_multi_vector_rrf_scores(single_query_results)\n multi_query_vectors = []\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n multi_query_vectors.append(VectorQuery(field_name=v, vector=doc_vectors[v]))\n\n rrf_reranker = RrfReRanker(topn=3)\n multi_query_result = full_collection.query(\n vectors=multi_query_vectors,\n reranker=rrf_reranker,\n )\n assert len(multi_query_result) > 0, (\n f\"Expected at least 1 result, but got {len(multi_query_result)}\"\n )\n\n multi_querydoc_check(multiple_docs, multi_query_result, full_collection)\n\n prev_score = float(\"inf\")\n for i, doc in enumerate(multi_query_result):\n doc_id = doc.id\n assert doc_id in expected_rrf_scores, (\n f\"Document {doc_id} should be in expected RRF scores\"\n )\n expected_score = expected_rrf_scores[doc_id]\n actual_score = doc.score\n assert abs(actual_score - expected_score) < 1e-10, (\n f\"RRF score mismatch for document {doc_id}: expected {expected_score}, got {actual_score}\"\n )\n assert doc.score <= prev_score, (\n f\"Scores should be in descending order. Current: {doc.score}, Previous: {prev_score}\"\n )\n prev_score = doc.score\n\n @pytest.mark.parametrize(\n \"weights\",\n [\n {\n \"vector_fp32_field\": 0.3,\n \"vector_fp16_field\": 0.2,\n \"vector_int8_field\": 0.3,\n \"sparse_vector_fp32_field\": 0.1,\n \"sparse_vector_fp16_field\": 0.1,\n }\n ],\n )\n @pytest.mark.parametrize(\n \"metric_type\", [MetricType.L2, MetricType.IP, MetricType.COSINE]\n )\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_multivector_weighted(\n self, full_collection: Collection, doc_num, weights, metric_type\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)\n\n weighted_reranker = WeightedReRanker(\n topn=3, weights=weights, metric=MetricType.IP\n )\n\n single_query_results = {}\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n single_query_results[v] = full_collection.query(\n VectorQuery(field_name=v, vector=doc_vectors[v])\n )\n expected_weighted_scores = calculate_multi_vector_weighted_scores(\n single_query_results, weights, MetricType.IP\n )\n\n multi_query_vectors = []\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n multi_query_vectors.append(VectorQuery(field_name=v, vector=doc_vectors[v]))\n\n multi_query_result = full_collection.query(\n vectors=multi_query_vectors,\n reranker=weighted_reranker,\n )\n assert len(multi_query_result) > 0, (\n f\"Expected at least 1 result, but got {len(multi_query_result)}\"\n )\n\n multi_querydoc_check(multiple_docs, multi_query_result, full_collection)\n\n prev_score = float(\"inf\")\n for i, doc in enumerate(multi_query_result):\n doc_id = doc.id\n assert doc_id in expected_weighted_scores, (\n f\"Document {doc_id} should be in expected scores\"\n )\n expected_score = expected_weighted_scores[doc_id]\n actual_score = doc.score\n assert abs(actual_score - expected_score) < 1e-10, (\n f\"score mismatch for document {doc_id}: expected {expected_score}, got {actual_score}\"\n )\n assert doc.score <= prev_score, (\n f\"Scores should be in descending order. Current: {doc.score}, Previous: {prev_score}\"\n )\n prev_score = doc.score\n\n @pytest.mark.parametrize(\"topk\", [5])\n @pytest.mark.parametrize(\"doc_num\", [10])\n @pytest.mark.parametrize(\"filter\", [\"int32_field >= 3 and int32_field <= 7\"])\n def test_query_consistency(\n self, full_collection: Collection, filter, doc_num, topk\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n results = []\n for i in range(5):\n query_result = full_collection.query(filter=filter, topk=topk)\n single_querydoc_check(multiple_docs, query_result, full_collection)\n\n results.append(query_result)\n assert len(results) == 5\n expected_count = len(results[0])\n for i, result in enumerate(results):\n assert len(result) == expected_count\n\n expected_ids = set(doc.id for doc in results[0])\n for i, result in enumerate(results):\n result_ids = set(doc.id for doc in result)\n assert result_ids == expected_ids\n\n for i, result in enumerate(results):\n result_ids = [doc.id for doc in result]\n expected_sorted_ids = sorted(result_ids, key=lambda x: int(x))\n assert result_ids == expected_sorted_ids\n\n @pytest.mark.parametrize(\"ef\", [0, 100, 1024, 2048])\n @pytest.mark.parametrize(\"doc_num\", [10])\n @pytest.mark.parametrize(\"topk\", [1024])\n @pytest.mark.parametrize(\"filter\", [\"int32_field >= 3 and int32_field <= 7\"])\n def test_query_vector_with_HnswQueryParam_valid(\n self,\n full_collection_new: Collection,\n doc_num,\n full_schema_new,\n topk,\n filter,\n ef,\n ):\n multiple_docs = [\n generate_doc(i, full_collection_new.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(\n full_collection_new, multiple_docs, doc_num, operator=\"insert\"\n )\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n doc_fields, doc_vectors = generate_vectordict_random(\n full_collection_new.schema\n )\n query_vector = doc_vectors[v]\n query_result = full_collection_new.query(\n filter=filter,\n vectors=VectorQuery(\n field_name=v, vector=query_vector, param=HnswQueryParam(ef=ef)\n ),\n topk=topk,\n )\n assert len(query_result) > 0, (\n f\"Expected at least 1 query result, but got {len(query_result)}\"\n )\n single_querydoc_check(\n multiple_docs,\n query_result,\n full_collection_new,\n is_by_vector=1,\n query_vector=query_vector,\n data_type=k,\n vector_name=v,\n )\n\n @pytest.mark.parametrize(\"ef\", [None, \"invalid\", 10.5])\n @pytest.mark.parametrize(\"doc_num\", [10])\n @pytest.mark.parametrize(\"topk\", [10])\n @pytest.mark.parametrize(\"filter\", [\"int32_field >= 3 and int32_field <= 7\"])\n def test_query_vector_with_HnswQueryParam_invalid(\n self, full_collection: Collection, doc_num, topk, ef, filter\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)\n query_vector = doc_vectors[v]\n with pytest.raises(Exception) as exc_info:\n full_collection.query(\n filter=filter,\n vectors=VectorQuery(\n field_name=v, vector=query_vector, param=HnswQueryParam(ef=ef)\n ),\n topk=topk,\n )\n assert INCOMPATIBLE_CONSTRUCTOR_ERROR_MSG in str(exc_info.value)\n\n @pytest.mark.parametrize(\"nprobe\", [1, 10, 100, 2048])\n @pytest.mark.parametrize(\"doc_num\", [10])\n @pytest.mark.parametrize(\"topk\", [10])\n @pytest.mark.parametrize(\"filter\", [\"int32_field >= 3 and int32_field <= 7\"])\n def test_query_vector_with_IVFQueryParam_valid(\n self, full_collection_ivf: Collection, nprobe, doc_num, topk, filter\n ):\n multiple_docs = [\n generate_doc(i, full_collection_ivf.schema) for i in range(doc_num)\n ]\n batchdoc_and_check_ivf(\n full_collection_ivf, multiple_docs, doc_num, operator=\"insert\"\n )\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n doc_fields, doc_vectors = generate_vectordict_random(\n full_collection_ivf.schema\n )\n if v in [\"vector_fp32_field\"]:\n query_vector = doc_vectors[v]\n\n query_result = full_collection_ivf.query(\n filter=filter,\n vectors=VectorQuery(\n field_name=v,\n vector=query_vector,\n param=IVFQueryParam(nprobe=nprobe),\n ),\n topk=topk,\n )\n assert len(query_result) > 0\n single_querydoc_check(\n multiple_docs,\n query_result,\n full_collection_ivf,\n is_by_vector=1,\n query_vector=query_vector,\n data_type=k,\n vector_name=v,\n )\n\n @pytest.mark.parametrize(\"nprobe\", [None, 10.5])\n @pytest.mark.parametrize(\"doc_num\", [10])\n @pytest.mark.parametrize(\"topk\", [10])\n @pytest.mark.parametrize(\"filter\", [\"int32_field >= 3 and int32_field <= 7\"])\n def test_query_vector_with_IVFQueryParam_invalid(\n self, full_collection_ivf: Collection, nprobe, doc_num, topk, filter\n ):\n multiple_docs = [\n generate_doc(i, full_collection_ivf.schema) for i in range(doc_num)\n ]\n batchdoc_and_check_ivf(\n full_collection_ivf, multiple_docs, doc_num, operator=\"insert\"\n )\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n doc_fields, doc_vectors = generate_vectordict_random(\n full_collection_ivf.schema\n )\n if v in [\"vector_fp32_field\"]:\n print(\"v:\\n\")\n print(v)\n query_vector = doc_vectors[v]\n with pytest.raises(Exception) as exc_info:\n full_collection_ivf.query(\n # filter=filter,\n vectors=VectorQuery(\n field_name=v,\n vector=query_vector,\n param=IVFQueryParam(nprobe=nprobe),\n ),\n topk=topk,\n )\n assert INCOMPATIBLE_CONSTRUCTOR_ERROR_MSG in str(exc_info.value)\n\n @pytest.mark.parametrize(\"filter\", [\"int32_field >= 3 and int32_field <= 7\"])\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_vector_with_param_invalid(\n self, full_collection: Collection, doc_num, filter\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n\n with pytest.raises(Exception) as exc_info:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n doc_fields, doc_vectors = generate_vectordict_random(\n full_collection.schema\n )\n query_vector = doc_vectors[v]\n if v in [\"vector_fp16_field\", \"vector_fp32_field\"]:\n full_collection.query(\n filter=filter,\n vectors=VectorQuery(\n field_name=v, vector=query_vector, param=HnswIndexParam()\n ),\n )\n assert INCOMPATIBLE_FUNCTION_ERROR_MSG in str(exc_info.value)\n\n @pytest.mark.parametrize(\"doc_num\", [10])\n @pytest.mark.parametrize(\n \"test_case_name,vector_query,expected_error_msg\",\n [\n (\n \"Non-existent vector field name\",\n lambda ref_dense_vector: VectorQuery(\n field_name=\"nonexistent_vector\", vector=ref_dense_vector\n ),\n \"Expected exception for non-existent vector field name\",\n ),\n (\n \"Invalid vector data type for dense vector (string instead of list)\",\n lambda ref_dense_vector: VectorQuery(\n field_name=\"vector_fp32_field\", vector=\"invalid_vector_data\"\n ),\n \"Expected exception for invalid dense vector data type\",\n ),\n (\n \"Invalid vector data type for sparse vector (list instead of dict)\",\n lambda ref_dense_vector: VectorQuery(\n field_name=\"sparse_fp32\", vector=[1.0, 2.0, 3.0]\n ),\n \"Expected exception for invalid sparse vector data type\",\n ),\n (\n \"Empty vector data for dense vector\",\n lambda ref_dense_vector: VectorQuery(\n field_name=\"vector_fp32_field\", vector=[]\n ),\n \"Expected exception for empty dense vector data\",\n ),\n (\n \"Invalid dimension for dense vector\",\n lambda ref_dense_vector: VectorQuery(\n field_name=\"vector_fp32_field\", vector=[1.0, 2.0]\n ), # Only 2 dimensions instead of 128\n \"Expected exception for invalid dense vector dimension\",\n ),\n (\n \"Non-existent document ID for by_id query\",\n lambda ref_dense_vector: VectorQuery(\n field_name=\"vector_fp32_field\", id=\"999\"\n ), # Non-existent ID\n \"Expected exception for non-existent document ID\",\n ),\n (\n \"Both vector and id specified (invalid combination)\",\n lambda ref_dense_vector: VectorQuery(\n field_name=\"vector_fp32_field\", vector=ref_dense_vector, id=\"5\"\n ),\n \"Expected exception for specifying both vector and id\",\n ),\n (\n \"Neither vector nor id specified\",\n lambda ref_dense_vector: VectorQuery(\n field_name=\"vector_fp32_field\"\n ), # Neither vector nor id\n \"Expected exception for specifying neither vector nor id\",\n ),\n ],\n )\n def test_query_vector_with_vectors_invalid(\n self,\n full_collection: Collection,\n doc_num,\n test_case_name,\n vector_query,\n expected_error_msg,\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n ref_doc_result = full_collection.fetch(ids=[\"5\"])\n assert \"5\" in ref_doc_result\n ref_doc = ref_doc_result[\"5\"]\n ref_dense_vector = ref_doc.vector(\"vector_fp32_field\")\n\n with pytest.raises(Exception) as exc_info:\n full_collection.query(vectors=[vector_query(ref_dense_vector)])\n assert exc_info.value is not None, expected_error_msg\n\n @pytest.mark.parametrize(\"filter\", [\"int32_field >= 3 and int32_field <= 7\"])\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_query_invalid_param_incompatible_type(\n self, full_collection: Collection, doc_num, filter\n ):\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(doc_num)\n ]\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n\n with pytest.raises(Exception) as exc_info:\n for k, v in DEFAULT_VECTOR_FIELD_NAME.items():\n doc_fields, doc_vectors = generate_vectordict_random(\n full_collection.schema\n )\n query_vector = doc_vectors[v]\n full_collection.query(\n filter=filter,\n vectors=VectorQuery(field_name=v, vector=query_vector),\n param=HnswIndexParam(),\n topk=3,\n )\n\n assert \"query() got an unexpected keyword argument 'param'\" in str(\n exc_info.value\n )\n\n\nclass TestRRFScoreCalculation:\n class MockDoc:\n def __init__(self, id, score=0.0):\n self._id = id\n self._score = score\n\n @property\n def id(self):\n return self._id\n\n @property\n def score(self):\n return self._score\n\n @score.setter\n def score(self, score):\n self._score = score\n\n def test_rrf_score_calculation_formula(self):\n k = 60\n\n assert abs(calculate_rrf_score(0, k) - 1.0 / 61) < 1e-10, (\n \"RRF score for rank 0 should be 1/61\"\n )\n assert abs(calculate_rrf_score(1, k) - 1.0 / 62) < 1e-10, (\n \"RRF score for rank 1 should be 1/62\"\n )\n assert abs(calculate_rrf_score(2, k) - 1.0 / 63) < 1e-10, (\n \"RRF score for rank 2 should be 1/63\"\n )\n assert abs(calculate_rrf_score(10, k) - 1.0 / 71) < 1e-10, (\n \"RRF score for rank 10 should be 1/71\"\n )\n\n k = 10\n assert abs(calculate_rrf_score(0, k) - 1.0 / 11) < 1e-10, (\n \"RRF score for rank 0 with k=10 should be 1/11\"\n )\n assert abs(calculate_rrf_score(1, k) - 1.0 / 12) < 1e-10, (\n \"RRF score for rank 1 with k=10 should be 1/12\"\n )\n\n def test_multi_vector_rrf_scores(self):\n query1_results = [self.MockDoc(\"1\"), self.MockDoc(\"2\"), self.MockDoc(\"3\")]\n query2_results = [self.MockDoc(\"3\"), self.MockDoc(\"1\"), self.MockDoc(\"4\")]\n query3_results = [self.MockDoc(\"2\"), self.MockDoc(\"4\"), self.MockDoc(\"5\")]\n query_results = {\n \"vector1\": query1_results,\n \"vector2\": query2_results,\n \"vector3\": query3_results,\n }\n rrf_scores = calculate_multi_vector_rrf_scores(query_results, k=60)\n\n expected_doc1_score = 1.0 / 61 + 1.0 / 62\n assert abs(rrf_scores[\"1\"] - expected_doc1_score) < 1e-10, (\n f\"RRF score for doc1 mismatch: expected {expected_doc1_score}, got {rrf_scores['1']}\"\n )\n expected_doc2_score = 1.0 / 62 + 1.0 / 61\n assert abs(rrf_scores[\"2\"] - expected_doc2_score) < 1e-10, (\n f\"RRF score for doc2 mismatch: expected {expected_doc2_score}, got {rrf_scores['2']}\"\n )\n expected_doc3_score = 1.0 / 63 + 1.0 / 61\n assert abs(rrf_scores[\"3\"] - expected_doc3_score) < 1e-10, (\n f\"RRF score for doc3 mismatch: expected {expected_doc3_score}, got {rrf_scores['3']}\"\n )\n expected_doc4_score = 1.0 / 63 + 1.0 / 62\n assert abs(rrf_scores[\"4\"] - expected_doc4_score) < 1e-10, (\n f\"RRF score for doc4 mismatch: expected {expected_doc4_score}, got {rrf_scores['4']}\"\n )\n\n expected_doc5_score = 1.0 / 63\n assert abs(rrf_scores[\"5\"] - expected_doc5_score) < 1e-10, (\n f\"RRF score for doc5 mismatch: expected {expected_doc5_score}, got {rrf_scores['5']}\"\n )\n sorted_scores = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)\n expected_order = [\"1\", \"2\", \"3\", \"4\", \"5\"]\n actual_order = [item[0] for item in sorted_scores]\n assert actual_order == expected_order, (\n f\"RRF score ranking mismatch: expected {expected_order}, got {actual_order}\"\n )\n\n\nclass TestCollectionConcurrencyOperations:\n @pytest.mark.parametrize(\"doc_num\", [10])\n def test_concurrent_insert_update_upsert_query(\n self, full_collection: Collection, doc_num\n ):\n import threading\n\n results = []\n errors = []\n multiple_docs = [\n generate_doc(i, full_collection.schema) for i in range(1000, 1010)\n ]\n\n batchdoc_and_check(full_collection, multiple_docs, doc_num, operator=\"insert\")\n\n def insert_operation(thread_id):\n try:\n multiple_docs = [\n generate_doc(i, full_collection.schema)\n for i in range(thread_id, thread_id + 5)\n ]\n result = full_collection.insert(multiple_docs)\n results.append((\"insert\", thread_id, len(result)))\n except Exception as e:\n errors.append((\"insert\", thread_id, str(e)))\n\n def update_operation(thread_id):\n try:\n multiple_docs = [\n generate_doc_random(i, full_collection.schema)\n for i in range(1000, 1001)\n ]\n result = full_collection.update(multiple_docs)\n results.append((\"update\", thread_id, len(result)))\n except Exception as e:\n errors.append((\"update\", thread_id, str(e)))\n\n def upsert_operation(thread_id):\n try:\n multiple_docs = [\n generate_doc(i, full_collection.schema)\n for i in range(thread_id, thread_id + 5)\n ]\n result = full_collection.upsert(multiple_docs)\n results.append((\"upsert\", thread_id, len(result)))\n except Exception as e:\n errors.append((\"upsert\", thread_id, str(e)))\n\n def query_operation(thread_id):\n try:\n if thread_id % 3 == 0:\n result = full_collection.query(filter=\"int32_field > 1\", topk=5)\n elif thread_id % 3 == 1:\n result = full_collection.query(filter=\"bool_field = true\", topk=3)\n else:\n query_vector = [0.1] * 128\n result = full_collection.query(\n VectorQuery(\n field_name=\"vector_fp32_field\", vector=query_vector\n ),\n topk=3,\n )\n\n results.append((\"query\", thread_id, len(result)))\n except Exception as e:\n errors.append((\"query\", thread_id, str(e)))\n\n def delete_operation(thread_id):\n try:\n # Delete some existing documents\n delete_ids = (\n [f\"{thread_id + 1}\", f\"{thread_id + 2}\"]\n if thread_id < 5\n else [f\"{thread_id % 5 + 1}\"]\n )\n result = full_collection.delete(delete_ids)\n results.append((\"delete\", thread_id, len(result)))\n except Exception as e:\n errors.append((\"delete\", thread_id, str(e)))\n\n threads = []\n for i in range(1):\n thread = threading.Thread(target=insert_operation, args=(i,))\n threads.append(thread)\n thread.start()\n for i in range(1):\n thread = threading.Thread(target=update_operation, args=(i,))\n threads.append(thread)\n thread.start()\n for i in range(1):\n thread = threading.Thread(target=upsert_operation, args=(i,))\n threads.append(thread)\n thread.start()\n for i in range(1):\n thread = threading.Thread(target=query_operation, args=(i,))\n threads.append(thread)\n thread.start()\n for i in range(1):\n thread = threading.Thread(target=delete_operation, args=(i,))\n threads.append(thread)\n thread.start()\n\n for thread in threads:\n thread.join()\n\n insert_results = [r for r in results if r[0] == \"insert\"]\n update_results = [r for r in results if r[0] == \"update\"]\n upsert_results = [r for r in results if r[0] == \"upsert\"]\n query_results = [r for r in results if r[0] == \"query\"]\n delete_results = [r for r in results if r[0] == \"delete\"]\n\n assert (\n len(insert_results)\n + len(update_results)\n + len(upsert_results)\n + len(query_results)\n + len(delete_results)\n > 0\n ), f\"No operations succeeded. Errors: {errors}\"\n\n critical_errors = [\n e for e in errors if \"critical\" in e[2].lower() or \"fatal\" in e[2].lower()\n ]\n assert len(critical_errors) == 0, f\"Critical errors occurred: {critical_errors}\"\n" }, { "path": "python/tests/detail/test_collection_exception.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport logging\nimport pytest\nimport numpy as np\nimport zvec\n\nfrom zvec import (\n CollectionOption,\n InvertIndexParam,\n HnswIndexParam,\n DataType,\n Collection,\n Doc,\n FieldSchema,\n VectorSchema,\n VectorQuery,\n)\n\n\nclass TestCollectionExceptionHandling:\n @pytest.fixture(scope=\"function\")\n def test_collection(self, tmp_path_factory):\n \"\"\"Fixture to create a test collection\"\"\"\n collection_schema = zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\",\n DataType.STRING,\n nullable=False,\n index_param=InvertIndexParam(),\n ),\n FieldSchema(\"weight\", DataType.FLOAT, nullable=True),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n VectorSchema(\n \"sparse\", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()\n ),\n ],\n )\n\n collection_option = CollectionOption(read_only=False, enable_mmap=True)\n\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n coll = zvec.create_and_open(\n path=str(collection_path),\n schema=collection_schema,\n option=collection_option,\n )\n\n assert coll is not None, \"Failed to create and open collection\"\n\n yield coll\n\n # Clean up\n if hasattr(coll, \"destroy\") and coll is not None:\n try:\n coll.destroy()\n except Exception as e:\n print(f\"Warning: failed to destroy collection: {e}\")\n\n def test_create_and_open_missing_path(self, tmp_path_factory):\n collection_schema = zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\",\n DataType.STRING,\n nullable=False,\n index_param=InvertIndexParam(),\n ),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n ],\n )\n\n collection_option = CollectionOption(read_only=False, enable_mmap=True)\n\n with pytest.raises(Exception) as exc_info:\n coll = zvec.create_and_open(\n schema=collection_schema, option=collection_option\n )\n assert exc_info.value is not None, (\n \"Expected exception for missing path parameter\"\n )\n\n def test_create_and_open_missing_schema(self, tmp_path_factory):\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n collection_option = CollectionOption(read_only=False, enable_mmap=True)\n\n with pytest.raises(Exception) as exc_info:\n coll = zvec.create_and_open(\n path=str(collection_path), option=collection_option\n )\n assert exc_info.value is not None, (\n \"Expected exception for missing schema parameter\"\n )\n\n def test_open_missing_path(self):\n collection_option = CollectionOption(read_only=False, enable_mmap=True)\n\n with pytest.raises(Exception) as exc_info:\n coll = zvec.open(option=collection_option)\n assert exc_info.value is not None, (\n \"Expected exception for missing path parameter\"\n )\n\n def test_insert_missing_docs(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n result = test_collection.insert()\n assert exc_info.value is not None, (\n \"Expected exception for missing docs parameter\"\n )\n\n def test_update_missing_docs(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n result = test_collection.update()\n assert exc_info.value is not None, (\n \"Expected exception for missing docs parameter\"\n )\n\n def test_upsert_missing_docs(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n result = test_collection.upsert()\n assert exc_info.value is not None, (\n \"Expected exception for missing docs parameter\"\n )\n\n def test_delete_missing_ids(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n result = test_collection.delete()\n assert exc_info.value is not None, (\n \"Expected exception for missing ids parameter\"\n )\n\n def test_fetch_missing_ids(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n result = test_collection.fetch()\n assert exc_info.value is not None, (\n \"Expected exception for missing ids parameter\"\n )\n\n def test_query_missing_vectorquery_field_name(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n result = test_collection.query(vectors=[VectorQuery()])\n assert exc_info.value is not None, (\n \"Expected exception for missing VectorQuery field_name parameter\"\n )\n\n def test_add_column_missing_field_schema(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n test_collection.add_column()\n assert exc_info.value is not None, (\n \"Expected exception for missing field_schema parameter\"\n )\n\n def test_alter_column_missing_old_name(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n test_collection.alter_column(new_name=\"new_name\")\n assert exc_info.value is not None, (\n \"Expected exception for missing old_name parameter\"\n )\n\n def test_alter_column_missing_new_name(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n test_collection.alter_column(old_name=\"old_name\")\n assert exc_info.value is not None, (\n \"Expected exception for missing new_name parameter\"\n )\n\n def test_drop_column_missing_field_name(self, test_collection: Collection):\n with pytest.raises(Exception) as exc_info:\n test_collection.drop_column()\n assert exc_info.value is not None, (\n \"Expected exception for missing field_name parameter\"\n )\n\n def test_invalid_parameter_types(self, test_collection: Collection):\n # This test depends on specific implementation details\n # Generally, we would expect TypeErrors or similar exceptions\n pass\n\n def test_missing_required_parameters(self, test_collection: Collection):\n # This test depends on specific implementation details\n # Generally, we would expect TypeErrors or similar exceptions\n pass\n\n def test_empty_collection_operations(self, tmp_path_factory):\n collection_schema = zvec.CollectionSchema(\n name=\"empty_test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\",\n DataType.STRING,\n nullable=False,\n index_param=InvertIndexParam(),\n ),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n ],\n )\n\n collection_option = CollectionOption(read_only=False, enable_mmap=True)\n\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"empty_test_collection\"\n\n coll = zvec.create_and_open(\n path=str(collection_path),\n schema=collection_schema,\n option=collection_option,\n )\n\n assert coll is not None, \"Failed to create and open collection\"\n\n # Test fetch on empty collection\n result = coll.fetch([\"1\"])\n assert len(result) >= 0 # May be empty or have special handling\n\n # Test query on empty collection\n result = coll.query()\n assert len(result) == 0\n\n # Test update on empty collection\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\"},\n vectors={\"dense\": np.random.random(128).tolist()},\n )\n\n result = coll.update(doc)\n # Should handle gracefully, possibly with NOT_FOUND status\n\n # Clean up\n if hasattr(coll, \"destroy\") and coll is not None:\n try:\n coll.destroy()\n except Exception as e:\n print(f\"Warning: failed to destroy collection: {e}\")\n\n def test_resource_management(self, test_collection: Collection):\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\", \"weight\": 80.5},\n vectors={\n \"dense\": np.random.random(128).tolist(),\n \"sparse\": {1: 1.0, 2: 2.0},\n },\n )\n\n # Insert\n result = test_collection.insert(doc)\n assert result.ok()\n\n # Fetch\n result = test_collection.fetch([\"1\"])\n assert len(result) == 1\n\n # Query\n result = test_collection.query()\n assert len(result) >= 0\n\n # Update\n result = test_collection.update(doc)\n assert result.ok()\n\n # Delete\n result = test_collection.delete(\"1\")\n assert result.ok()\n\n def test_exception_resource_cleanup(self, test_collection: Collection):\n # This test would need to simulate exception conditions\n # which is difficult without specific failure injection points\n pass\n" }, { "path": "python/tests/detail/test_collection_open.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport threading\nimport numpy as np\n\nfrom fixture_helper import *\n\nCOLLECTION_OPTION_TEST_CASES_VALID = [\n # (read_only, enable_mmap, description)\n (False, True, \"Read-write with mmap enabled\"),\n (False, False, \"Read-write with mmap disabled\"),\n (True, True, \"Read-only with mmap enabled\"),\n (True, False, \"Read-only with mmap disabled\"),\n]\n\n# Test data for invalid paths\nINVALID_PATH_LIST = [\n \"/nonexistent/directory/test_collection\",\n \"invalid:path\",\n \"\", # Empty path\n]\n\n\n@pytest.fixture(scope=\"session\")\ndef collection_schema():\n return zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\", DataType.STRING, nullable=False, index_param=InvertIndexParam()\n ),\n FieldSchema(\n \"weight\", DataType.FLOAT, nullable=False, index_param=InvertIndexParam()\n ),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n VectorSchema(\n \"sparse\", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()\n ),\n ],\n )\n\n\n@pytest.fixture\ndef single_doc():\n id = 0\n return Doc(\n id=f\"{id}\",\n fields={\"id\": id, \"name\": \"test\"},\n vectors={\n \"dense\": [id + 0.1] * 128,\n },\n )\n\n\n@pytest.fixture(scope=\"function\")\ndef test_collection(\n tmp_path_factory, collection_schema, collection_option\n) -> Generator[Any, Any, Collection]:\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n coll = zvec.create_and_open(\n path=str(collection_path), schema=collection_schema, option=collection_option\n )\n\n assert coll is not None, \"Failed to create and open collection\"\n assert coll.path == str(collection_path)\n assert coll.schema.name == collection_schema.name\n assert list(coll.schema.fields) == list(collection_schema.fields)\n assert list(coll.schema.vectors) == list(collection_schema.vectors)\n assert coll.option.read_only == collection_option.read_only\n assert coll.option.enable_mmap == collection_option.enable_mmap\n\n try:\n yield coll\n finally:\n if hasattr(coll, \"destroy\") and coll is not None:\n try:\n coll.destroy()\n except Exception as e:\n print(f\"Warning: failed to destroy collection: {e}\")\n\n\nclass TestCollectionOpen:\n def test_open_basic_functionality(\n self, tmp_path_factory, collection_schema, collection_option\n ):\n import sys\n import time\n import os\n\n # Create unique temp directory\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n # Ensure the path exists\n collection_path_str = str(collection_path)\n print(f\"DEBUG: Collection path: {collection_path_str}\")\n print(f\"DEBUG: Temp directory exists: {temp_dir.exists()}\")\n\n # Create and open collection first\n created_coll = zvec.create_and_open(\n path=collection_path_str, schema=collection_schema, option=collection_option\n )\n\n assert created_coll is not None, (\n f\"Failed to create collection, returned None instead of valid Collection object. Path: {collection_path_str}\"\n )\n assert created_coll.path == collection_path_str, (\n f\"Collection path mismatch. Expected: {collection_path_str}, Actual: {created_coll.path}\"\n )\n assert created_coll.schema.name == \"test_collection\", (\n f\"Collection schema name mismatch. Expected: test_collection, Actual: {created_coll.schema.name}\"\n )\n\n # Insert multiple documents to verify persistence\n docs = []\n for i in range(3):\n doc = Doc(\n id=f\"{i}\",\n fields={\"id\": i, \"name\": f\"test_{i}\", \"weight\": float(i * 10)},\n vectors={\n \"dense\": [float(j + i) for j in range(128)],\n \"sparse\": {j: float(j + i) for j in range(5)},\n },\n )\n docs.append(doc)\n\n result = created_coll.insert(docs)\n assert len(result) == 3, f\"Expected 3 insertion results, but got {len(result)}\"\n for i, res in enumerate(result):\n assert res.ok(), (\n f\"Insertion result {i} is not OK. Status code: {res.code()}, Message: {res.message()}\"\n )\n\n # Verify documents were inserted using fetch interface\n fetched_docs_after_insert = created_coll.fetch([\"0\", \"1\", \"2\"])\n assert len(fetched_docs_after_insert) == 3, (\n f\"Expected 3 fetched documents after insertion, but got {len(fetched_docs_after_insert)}\"\n )\n assert \"0\" in fetched_docs_after_insert, (\n \"Document with ID '0' not found in fetched results after insertion\"\n )\n assert \"1\" in fetched_docs_after_insert, (\n \"Document with ID '1' not found in fetched results after insertion\"\n )\n assert \"2\" in fetched_docs_after_insert, (\n \"Document with ID '2' not found in fetched results after insertion\"\n )\n\n # Verify fetched document content after insertion\n for i in range(3):\n doc = fetched_docs_after_insert[f\"{i}\"]\n assert doc is not None, (\n f\"Fetched document with ID '{i}' is None after insertion\"\n )\n assert doc.id == f\"{i}\", (\n f\"Document ID mismatch for document '{i}' after insertion. Expected: {i}, Actual: {doc.id}\"\n )\n assert doc.field(\"id\") == i, (\n f\"Document id field mismatch for document '{i}' after insertion. Expected: {i}, Actual: {doc.field('id')}\"\n )\n assert doc.field(\"name\") == f\"test_{i}\", (\n f\"Document name field mismatch for document '{i}' after insertion. Expected: test_{i}, Actual: {doc.field('name')}\"\n )\n assert doc.field(\"weight\") == float(i * 10), (\n f\"Document weight field mismatch for document '{i}' after insertion. Expected: {float(i * 10)}, Actual: {doc.field('weight')}\"\n )\n\n # Verify vector access after insertion\n assert doc.vector(\"dense\") is not None, (\n f\"Document {i} should have dense vector after insertion\"\n )\n assert doc.vector(\"sparse\") is not None, (\n f\"Document {i} should have sparse vector after insertion\"\n )\n\n # Verify vector types after insertion\n assert isinstance(doc.vector(\"dense\"), list), (\n f\"Document {i} dense vector should be dict after insertion, got {type(doc.vector('dense'))}\"\n )\n assert isinstance(doc.vector(\"sparse\"), dict), (\n f\"Document {i} sparse vector should be dict after insertion, got {type(doc.vector('sparse'))}\"\n )\n\n # Verify documents were inserted using stats\n stats = created_coll.stats\n assert stats is not None, \"Collection stats should not be None\"\n assert stats.doc_count == 3, (\n f\"Document count mismatch after insertion. Expected: 3, Actual: {stats.doc_count}\"\n )\n\n # Store the collection path before cleanup\n collection_path = created_coll.path\n\n # Clean up the created collection reference\n del created_coll\n\n # Wait and verify the path still exists\n print(f\"DEBUG: Collection path after destroy: {collection_path}\")\n print(f\"DEBUG: Path exists after destroy: {os.path.exists(collection_path)}\")\n\n # Now open the existing collection\n try:\n print(f\"DEBUG: Path exists before open: {os.path.exists(collection_path)}\")\n\n # List contents of parent directory for debugging\n parent_dir = os.path.dirname(collection_path)\n if os.path.exists(parent_dir):\n print(f\"DEBUG: Parent directory contents: {os.listdir(parent_dir)}\")\n\n opened_coll = zvec.open(path=collection_path, option=collection_option)\n\n assert opened_coll is not None, (\n f\"Failed to open existing collection at path: {collection_path}. Returned None instead of valid Collection object\"\n )\n assert opened_coll.path == collection_path, (\n f\"Opened collection path mismatch. Expected: {collection_path}, Actual: {opened_coll.path}\"\n )\n assert opened_coll.schema.name == \"test_collection\", (\n f\"Opened collection schema name mismatch. Expected: test_collection, Actual: {opened_coll.schema.name}\"\n )\n\n # Check reference count of opened collection\n opened_ref_count = sys.getrefcount(opened_coll)\n print(f\"DEBUG: Reference count of opened collection: {opened_ref_count}\")\n\n # Verify data persistence\n # Verify data persistence using fetch interface\n fetched_docs = opened_coll.fetch([\"0\", \"1\", \"2\"])\n assert len(fetched_docs) == 3, (\n f\"Expected 3 fetched documents after reopening, but got {len(fetched_docs)}\"\n )\n assert \"0\" in fetched_docs, (\n \"Document with ID '0' not found in fetched results after reopening\"\n )\n assert \"1\" in fetched_docs, (\n \"Document with ID '1' not found in fetched results after reopening\"\n )\n assert \"2\" in fetched_docs, (\n \"Document with ID '2' not found in fetched results after reopening\"\n )\n\n # Verify fetched document content after reopening collection\n for i in range(3):\n doc = fetched_docs[f\"{i}\"]\n assert doc is not None, (\n f\"Fetched document with ID '{i}' is None after reopening collection\"\n )\n assert doc.id == f\"{i}\", (\n f\"Document ID mismatch for document '{i}' after reopening. Expected: {i}, Actual: {doc.id}\"\n )\n assert doc.field(\"id\") == i, (\n f\"Document id field mismatch for document '{i}' after reopening. Expected: {i}, Actual: {doc.field('id')}\"\n )\n assert doc.field(\"name\") == f\"test_{i}\", (\n f\"Document name field mismatch for document '{i}' after reopening. Expected: test_{i}, Actual: {doc.field('name')}\"\n )\n assert doc.field(\"weight\") == float(i * 10), (\n f\"Document weight field mismatch for document '{i}' after reopening. Expected: {float(i * 10)}, Actual: {doc.field('weight')}\"\n )\n\n # Verify vector access after reopening\n assert doc.vector(\"dense\") is not None, (\n f\"Document {i} should have dense vector after reopening\"\n )\n assert doc.vector(\"sparse\") is not None, (\n f\"Document {i} should have sparse vector after reopening\"\n )\n\n # Verify vector types after reopening\n assert isinstance(doc.vector(\"dense\"), list), (\n f\"Document {i} dense vector should be dict after reopening, got {type(doc.vector('dense'))}\"\n )\n assert isinstance(doc.vector(\"sparse\"), dict), (\n f\"Document {i} sparse vector should be dict after reopening, got {type(doc.vector('sparse'))}\"\n )\n\n # Verify score attribute exists\n assert hasattr(doc, \"score\"), (\n f\"Document {i} should have a score attribute after reopening\"\n )\n assert isinstance(doc.score, (int, float)), (\n f\"Document {i} score should be numeric after reopening, got {type(doc.score)}\"\n )\n # For fetch operations, score is typically 0.0\n assert doc.score == 0.0, (\n f\"Document {i} score should be 0.0 for fetch operation after reopening, but got {doc.score}\"\n )\n\n # Test query functionality\n query_result = opened_coll.query(include_vector=True)\n assert len(query_result) == 3, (\n f\"Expected 3 query results, but got {len(query_result)}\"\n )\n\n # Verify query results have proper structure and content with detailed validation\n returned_doc_ids = set()\n for doc in query_result:\n # Verify basic document structure\n assert doc.id is not None, f\"Query result document should have an ID\"\n assert doc.id in [\"0\", \"1\", \"2\"], (\n f\"Query result document ID should be one of ['0', '1', '2'], but got {doc.id}\"\n )\n returned_doc_ids.add(doc.id)\n\n # Verify field access\n assert doc.field(\"id\") is not None, (\n f\"Document {doc.id} should have id field\"\n )\n assert doc.field(\"name\") is not None, (\n f\"Document {doc.id} should have name field\"\n )\n assert doc.field(\"weight\") is not None, (\n f\"Document {doc.id} should have weight field\"\n )\n\n # Verify field values\n expected_id = int(doc.id)\n assert doc.field(\"id\") == expected_id, (\n f\"Document {doc.id} id field mismatch. Expected: {expected_id}, Actual: {doc.field('id')}\"\n )\n assert doc.field(\"name\") == f\"test_{expected_id}\", (\n f\"Document {doc.id} name field mismatch. Expected: test_{expected_id}, Actual: {doc.field('name')}\"\n )\n assert doc.field(\"weight\") == float(expected_id * 10), (\n f\"Document {doc.id} weight field mismatch. Expected: {float(expected_id * 10)}, Actual: {doc.field('weight')}\"\n )\n\n # Verify vector access\n assert doc.vector(\"dense\") is not None, (\n f\"Document {doc.id} should have dense vector\"\n )\n assert doc.vector(\"sparse\") is not None, (\n f\"Document {doc.id} should have sparse vector\"\n )\n\n # Verify vector types\n assert isinstance(doc.vector(\"dense\"), list), (\n f\"Document {doc.id} dense vector should be list, got {type(doc.vector('dense'))}\"\n )\n assert isinstance(doc.vector(\"sparse\"), dict), (\n f\"Document {doc.id} sparse vector should be dict, got {type(doc.vector('sparse'))}\"\n )\n\n # Verify score attribute exists\n assert hasattr(doc, \"score\"), (\n f\"Document {doc.id} should have a score attribute\"\n )\n assert isinstance(doc.score, (int, float)), (\n f\"Document {doc.id} score should be numeric, got {type(doc.score)}\"\n )\n\n # Verify all expected documents are returned\n expected_doc_ids = {\"0\", \"1\", \"2\"}\n assert returned_doc_ids == expected_doc_ids, (\n f\"Query should return all expected documents. Expected: {expected_doc_ids}, Actual: {returned_doc_ids}\"\n )\n\n # === Enhanced validation based on test_collection_dql_operations.py ===\n\n # Verify vector field names accessibility for all documents\n for doc in query_result:\n vector_names = doc.vector_names()\n expected_vector_names = {\"dense\", \"sparse\"}\n assert set(vector_names) == expected_vector_names, (\n f\"Document {doc.id} vector names mismatch. Expected: {expected_vector_names}, Actual: {set(vector_names)}\"\n )\n\n # Verify all vector fields can be accessed\n for vector_name in expected_vector_names:\n vector_data = doc.vector(vector_name)\n assert vector_data is not None, (\n f\"Document {doc.id} should have accessible vector '{vector_name}'\"\n )\n if vector_name == \"dense\":\n assert isinstance(vector_data, list), (\n f\"Document {doc.id} vector '{vector_name}' should be list, got {type(vector_data)}\"\n )\n else:\n assert isinstance(vector_data, dict), (\n f\"Document {doc.id} vector '{vector_name}' should be dict, got {type(vector_data)}\"\n )\n\n # Test query with filter\n filtered_result = opened_coll.query(filter=\"id >= 1\", include_vector=True)\n assert len(filtered_result) == 2, (\n f\"Expected 2 filtered query results (id >= 1), but got {len(filtered_result)}\"\n )\n\n # Verify filtered query results\n filtered_doc_ids = set()\n for doc in filtered_result:\n assert doc.id is not None, (\n f\"Filtered query result document should have an ID\"\n )\n assert doc.id in [\"1\", \"2\"], (\n f\"Filtered query result document ID should be one of ['1', '2'], but got {doc.id}\"\n )\n filtered_doc_ids.add(doc.id)\n\n # Verify filter condition is satisfied\n doc_id = int(doc.id)\n assert doc_id >= 1, (\n f\"Document {doc.id} should satisfy filter condition id >= 1\"\n )\n\n # Verify document structure\n assert doc.field(\"id\") is not None, (\n f\"Document {doc.id} should have id field\"\n )\n assert doc.field(\"name\") is not None, (\n f\"Document {doc.id} should have name field\"\n )\n assert doc.field(\"weight\") is not None, (\n f\"Document {doc.id} should have weight field\"\n )\n\n # Verify field values\n assert doc.field(\"id\") == doc_id, (\n f\"Document {doc.id} id field mismatch. Expected: {doc_id}, Actual: {doc.field('id')}\"\n )\n assert doc.field(\"name\") == f\"test_{doc_id}\", (\n f\"Document {doc.id} name field mismatch. Expected: test_{doc_id}, Actual: {doc.field('name')}\"\n )\n assert doc.field(\"weight\") == float(doc_id * 10), (\n f\"Document {doc.id} weight field mismatch. Expected: {float(doc_id * 10)}, Actual: {doc.field('weight')}\"\n )\n\n # Verify vector access\n assert doc.vector(\"dense\") is not None, (\n f\"Document {doc.id} should have dense vector\"\n )\n assert doc.vector(\"sparse\") is not None, (\n f\"Document {doc.id} should have sparse vector\"\n )\n\n # Verify score attribute exists\n assert hasattr(doc, \"score\"), (\n f\"Document {doc.id} should have a score attribute\"\n )\n assert isinstance(doc.score, (int, float)), (\n f\"Document {doc.id} score should be numeric, got {type(doc.score)}\"\n )\n\n # Verify filtered documents\n expected_filtered_ids = {\"1\", \"2\"}\n assert filtered_doc_ids == expected_filtered_ids, (\n f\"Filtered query should return expected documents. Expected: {expected_filtered_ids}, Actual: {filtered_doc_ids}\"\n )\n\n # Test vector query functionality for dense vectors\n query_vector_dense = [0.1] * 128\n vector_query_result = opened_coll.query(\n VectorQuery(field_name=\"dense\", vector=query_vector_dense)\n )\n assert len(vector_query_result) > 0, (\n f\"Expected at least 1 vector query result, but got {len(vector_query_result)}\"\n )\n\n # Verify vector query results structure\n for doc in vector_query_result[:3]: # Check first 3 results\n assert doc.id is not None, (\n f\"Vector query result document should have an ID\"\n )\n assert doc.id in [\"0\", \"1\", \"2\"], (\n f\"Vector query result document ID should be one of ['0', '1', '2'], but got {doc.id}\"\n )\n\n # Verify document structure\n assert doc.field(\"id\") is not None, (\n f\"Document {doc.id} should have id field\"\n )\n assert doc.field(\"name\") is not None, (\n f\"Document {doc.id} should have name field\"\n )\n assert doc.field(\"weight\") is not None, (\n f\"Document {doc.id} should have weight field\"\n )\n\n # Verify vector access\n assert doc.vector(\"dense\") is not None, (\n f\"Document {doc.id} should have dense vector\"\n )\n assert doc.vector(\"sparse\") is not None, (\n f\"Document {doc.id} should have sparse vector\"\n )\n\n # Verify score attribute exists and is numeric\n assert hasattr(doc, \"score\"), (\n f\"Document {doc.id} should have a score attribute\"\n )\n assert isinstance(doc.score, (int, float)), (\n f\"Document {doc.id} score should be numeric, got {type(doc.score)}\"\n )\n\n # For dense vector queries, score should typically be non-negative (depending on metric)\n # Note: This may vary based on the metric type used\n assert doc.score >= 0 or doc.score < 0, (\n f\"Document {doc.id} score should be a valid number\"\n )\n\n # Test vector query functionality for sparse vectors\n query_vector_sparse = {1: 1.0, 2: 2.0, 3: 3.0}\n sparse_vector_query_result = opened_coll.query(\n VectorQuery(field_name=\"sparse\", vector=query_vector_sparse)\n )\n assert len(sparse_vector_query_result) > 0, (\n f\"Expected at least 1 sparse vector query result, but got {len(sparse_vector_query_result)}\"\n )\n\n # Verify sparse vector query results structure\n for doc in sparse_vector_query_result[:3]: # Check first 3 results\n assert doc.id is not None, (\n f\"Sparse vector query result document should have an ID\"\n )\n assert doc.id in [\"0\", \"1\", \"2\"], (\n f\"Sparse vector query result document ID should be one of ['0', '1', '2'], but got {doc.id}\"\n )\n\n # Verify document structure\n assert doc.field(\"id\") is not None, (\n f\"Document {doc.id} should have id field\"\n )\n assert doc.field(\"name\") is not None, (\n f\"Document {doc.id} should have name field\"\n )\n assert doc.field(\"weight\") is not None, (\n f\"Document {doc.id} should have weight field\"\n )\n\n # Verify vector access\n assert doc.vector(\"dense\") is not None, (\n f\"Document {doc.id} should have dense vector\"\n )\n assert doc.vector(\"sparse\") is not None, (\n f\"Document {doc.id} should have sparse vector\"\n )\n\n # Verify score attribute exists and is numeric\n assert hasattr(doc, \"score\"), (\n f\"Document {doc.id} should have a score attribute\"\n )\n assert isinstance(doc.score, (int, float)), (\n f\"Document {doc.id} score should be numeric, got {type(doc.score)}\"\n )\n\n # Clean up\n if hasattr(opened_coll, \"destroy\") and opened_coll is not None:\n opened_coll.destroy()\n print(\"DEBUG: Opened collection destroyed successfully\")\n\n except Exception as e:\n logging.error(\"Exception occurred: [{}]\".format(e))\n raise e\n\n @pytest.mark.parametrize(\n \"read_only,enable_mmap,description\", COLLECTION_OPTION_TEST_CASES_VALID\n )\n @pytest.mark.parametrize(\"createAndopen_enable_mmap\", [True, False])\n def test_open_with_different_collection_options_valid(\n self,\n tmp_path_factory,\n createAndopen_enable_mmap,\n read_only,\n enable_mmap,\n description,\n collection_schema,\n ):\n # Create collection with initial option\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n initial_option = CollectionOption(\n read_only=False, enable_mmap=createAndopen_enable_mmap\n )\n\n # Create and open collection first\n created_coll = zvec.create_and_open(\n path=str(collection_path), schema=collection_schema, option=initial_option\n )\n\n assert created_coll is not None, \"Failed to create collection\"\n\n # Clean up the created collection reference\n del created_coll\n\n # Now open with different options\n collection_option = CollectionOption(\n read_only=read_only, enable_mmap=enable_mmap\n )\n\n try:\n opened_coll = zvec.open(path=str(collection_path), option=collection_option)\n\n assert opened_coll is not None, (\n f\"Failed to open collection with option: {description}. Returned None instead of valid Collection object. Path: {collection_path}\"\n )\n assert opened_coll.path == str(collection_path), (\n f\"Opened collection path mismatch. Expected: {collection_path}, Actual: {opened_coll.path}\"\n )\n assert opened_coll.schema.name == collection_schema.name, (\n f\"Opened collection schema name mismatch. Expected: {collection_schema.name}, Actual: {opened_coll.schema.name}\"\n )\n assert opened_coll.option.read_only == read_only, (\n f\"Opened collection read_only option mismatch. Expected: {read_only}, Actual: {opened_coll.option.read_only}\"\n )\n assert opened_coll.option.enable_mmap == createAndopen_enable_mmap, (\n f\"Opened collection mmap option mismatch. Expected: {createAndopen_enable_mmap}, Actual: {opened_coll.option.enable_mmap}\"\n )\n\n # Clean up\n if (\n hasattr(opened_coll, \"destroy\")\n and opened_coll is not None\n and read_only == False\n ):\n opened_coll.destroy()\n\n except Exception as e:\n logging.error(\"Exception occurred: [{}]\".format(e))\n pytest.fail(f\"Failed to open collection with different options: {e}\")\n\n def test_open_with_none_option(self, tmp_path_factory, collection_schema):\n # Create collection\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n initial_option = CollectionOption(read_only=False, enable_mmap=True)\n\n # Create and open collection first\n created_coll = zvec.create_and_open(\n path=str(collection_path), schema=collection_schema, option=initial_option\n )\n\n assert created_coll is not None, (\n f\"Failed to create collection. Returned None instead of valid Collection object. Path: {collection_path}\"\n )\n\n # Clean up the created collection reference\n del created_coll\n\n # Now open with None option\n with pytest.raises(Exception) as exc_info:\n zvec.open(path=str(collection_path), option=None)\n\n assert \"incompatible function arguments\" in str(exc_info.value), (\n f\"Expected 'incompatible function arguments' error, but got: {exc_info.value}\"\n )\n\n def test_reopen_collection(self, tmp_path_factory):\n # Prepare schema\n collection_schema = zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\",\n DataType.STRING,\n nullable=False,\n index_param=InvertIndexParam(),\n ),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n ],\n )\n\n collection_option = CollectionOption(read_only=False, enable_mmap=True)\n\n # Create collection\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n # Create and open collection\n coll1 = zvec.create_and_open(\n path=str(collection_path),\n schema=collection_schema,\n option=collection_option,\n )\n\n assert coll1 is not None, \"Failed to create and open collection\"\n\n # Insert some data\n doc = Doc(\n id=\"1\",\n fields={\"id\": 1, \"name\": \"test\"},\n vectors={\"dense\": np.random.random(128).tolist()},\n )\n\n result = coll1.insert(doc)\n assert result.ok()\n\n # Close the first collection (delete reference)\n del coll1\n\n # Reopen the collection\n coll2 = zvec.open(path=str(collection_path), option=collection_option)\n\n assert coll2 is not None, \"Failed to reopen collection\"\n assert coll2.path == str(collection_path)\n assert coll2.schema.name == collection_schema.name\n\n # Verify data is still there\n fetched_docs = coll2.fetch([\"1\"])\n assert \"1\" in fetched_docs\n fetched_doc = fetched_docs[\"1\"]\n assert fetched_doc.id == \"1\"\n assert fetched_doc.field(\"name\") == \"test\"\n\n # Clean up\n if hasattr(coll2, \"destroy\") and coll2 is not None:\n try:\n coll2.destroy()\n except Exception as e:\n print(f\"Warning: failed to destroy collection: {e}\")\n\n def test_open_concurrent_same_path(self, tmp_path_factory):\n # First create a collection\n collection_schema = zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\",\n DataType.STRING,\n nullable=False,\n index_param=InvertIndexParam(),\n ),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n ],\n )\n\n collection_option = CollectionOption(read_only=False, enable_mmap=True)\n\n # Create collection path\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n # First create the collection\n created_coll = zvec.create_and_open(\n path=str(collection_path),\n schema=collection_schema,\n option=collection_option,\n )\n\n assert created_coll is not None, \"Failed to create collection\"\n\n # Close the collection so we can test concurrent opening\n if hasattr(created_coll, \"close\") and created_coll is not None:\n created_coll.close()\n\n # Shared variables to collect results from threads\n results = []\n errors = []\n\n # Lock for thread-safe operations\n lock = threading.Lock()\n # Clean up the created collection reference\n del created_coll\n\n # Function to be executed by each thread\n def open_collection_thread(thread_id):\n try:\n coll = zvec.open(path=str(collection_path), option=collection_option)\n with lock:\n results.append((thread_id, coll))\n # Close the collection if opened successfully\n if hasattr(coll, \"close\") and coll is not None:\n coll.close()\n except Exception as e:\n with lock:\n errors.append((thread_id, str(e)))\n\n # Create 5 threads to call open concurrently\n threads = []\n for i in range(5):\n thread = threading.Thread(target=open_collection_thread, args=(i,))\n threads.append(thread)\n thread.start()\n\n # Wait for all threads to complete\n for thread in threads:\n thread.join()\n\n # Verify concurrency safety: only one should succeed, others should fail\n assert len(results) == 1, (\n f\"Expected exactly one successful open, but got {len(results)}\"\n )\n assert len(errors) == 4, (\n f\"Expected exactly four failures, but got {len(errors)}\"\n )\n\n # Additional verification: check that the successful open has a valid collection\n successful_thread_id, successful_collection = results[0]\n assert successful_collection is not None, (\n \"Successful open should return a valid collection\"\n )\n assert successful_collection.path == str(collection_path), (\n \"Collection path mismatch\"\n )\n\n # Clean up the successfully opened collection\n if (\n hasattr(successful_collection, \"destroy\")\n and successful_collection is not None\n ):\n try:\n successful_collection.destroy()\n except Exception as e:\n print(f\"Warning: failed to destroy collection: {e}\")\n\n def test_open_with_corrupted_files(self, tmp_path_factory):\n # First create a collection\n collection_schema = zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\",\n DataType.STRING,\n nullable=False,\n index_param=InvertIndexParam(),\n ),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n )\n ],\n )\n\n collection_option = CollectionOption(read_only=False, enable_mmap=True)\n\n # Create collection path\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n # First create the collection\n created_coll = zvec.create_and_open(\n path=str(collection_path),\n schema=collection_schema,\n option=collection_option,\n )\n\n assert created_coll is not None, \"Failed to create collection\"\n\n # Close the collection so we can manipulate its files\n if hasattr(created_coll, \"close\") and created_coll is not None:\n created_coll.close()\n\n # Test case 1: Delete some files in the collection directory (simulate partial corruption)\n import os\n import shutil\n import random\n\n # Get the collection directory path\n collection_dir = str(collection_path)\n\n # List all files in the collection directory\n files_in_dir = []\n for root, dirs, files in os.walk(collection_dir):\n for file in files:\n files_in_dir.append(os.path.join(root, file))\n\n # Randomly delete approximately half of the files to simulate partial corruption\n if files_in_dir:\n # Shuffle the list to randomly select files\n random.shuffle(files_in_dir)\n files_to_delete = files_in_dir[: len(files_in_dir) // 2]\n for file_path in files_to_delete:\n try:\n os.remove(file_path)\n except Exception as e:\n pass # Ignore errors during deletion\n\n # Try to open the collection with missing files - should raise an exception\n with pytest.raises(Exception):\n zvec.open(path=str(collection_path), option=collection_option)\n\n # Test case 2: Delete all files in the collection directory (simulate complete corruption)\n # Recreate the collection\n recreated_coll = zvec.create_and_open(\n path=str(collection_path) + \"_all\",\n schema=collection_schema,\n option=collection_option,\n )\n\n assert recreated_coll is not None, \"Failed to recreate collection\"\n\n # Close the collection so we can manipulate its files\n if hasattr(recreated_coll, \"close\") and recreated_coll is not None:\n recreated_coll.close()\n\n # Delete all files in the collection directory\n try:\n shutil.rmtree(collection_dir)\n os.makedirs(collection_dir) # Recreate empty directory\n except Exception as e:\n pass # Ignore errors during deletion\n\n # Try to open the collection with missing files - should raise an exception\n with pytest.raises(Exception):\n zvec.open(path=str(collection_path), option=collection_option)\n" }, { "path": "python/tests/detail/test_collection_recall.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\n\nfrom zvec.typing import DataType, StatusCode, MetricType, QuantizeType\nfrom zvec.model import Collection, Doc, VectorQuery\nfrom zvec.model.param import (\n CollectionOption,\n InvertIndexParam,\n HnswIndexParam,\n FlatIndexParam,\n IVFIndexParam,\n HnswQueryParam,\n IVFQueryParam,\n)\n\nfrom zvec.model.schema import FieldSchema, VectorSchema\nfrom zvec.extension import RrfReRanker, WeightedReRanker, QwenReRanker\nfrom distance_helper import *\n\nfrom zvec import StatusCode\nfrom distance_helper import *\nfrom fixture_helper import *\nfrom doc_helper import *\nfrom params_helper import *\n\nimport time\n\n\n# ==================== helper ====================\ndef batchdoc_and_check(collection: Collection, multiple_docs, operator=\"insert\"):\n if operator == \"insert\":\n result = collection.insert(multiple_docs)\n elif operator == \"upsert\":\n result = collection.upsert(multiple_docs)\n\n elif operator == \"update\":\n result = collection.update(multiple_docs)\n else:\n logging.error(\"operator value is error!\")\n\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok(), (\n f\"result={result},Insert operation failed with code {item.code()}\"\n )\n\n stats = collection.stats\n assert stats is not None, \"Collection stats should not be None\"\n \"\"\"assert stats.doc_count == len(multiple_docs), (\n f\"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}\"\n )\"\"\"\n\n doc_ids = [doc.id for doc in multiple_docs]\n fetched_docs = collection.fetch(doc_ids)\n assert len(fetched_docs) == len(multiple_docs), (\n f\"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}\"\n )\n\n for original_doc in multiple_docs:\n assert original_doc.id in fetched_docs, (\n f\"Expected document ID {original_doc.id} in fetched documents\"\n )\n fetched_doc = fetched_docs[original_doc.id]\n\n assert is_doc_equal(fetched_doc, original_doc, collection.schema)\n\n assert hasattr(fetched_doc, \"score\"), \"Document should have a score attribute\"\n assert fetched_doc.score == 0.0, (\n \"Fetch operation should return default score of 0.0\"\n )\n\n\ndef compute_exact_similarity_scores(\n vectors_a,\n vectors_b,\n metric_type=MetricType.IP,\n DataType=DataType.VECTOR_FP32,\n QuantizeType=QuantizeType.UNDEFINED,\n):\n similarities = []\n for i, vec_a in enumerate(vectors_a):\n for j, vec_b in enumerate(vectors_b):\n similarity = distance_recall(vec_a, vec_b, metric_type, DataType)\n similarities.append((j, similarity))\n\n # For L2,COSINE metric, smaller distances mean higher similarity, so sort in ascending order\n if (\n metric_type in [MetricType.L2]\n and DataType\n in [DataType.VECTOR_FP32, DataType.VECTOR_FP16, DataType.VECTOR_INT8]\n ) or (\n metric_type in [MetricType.COSINE]\n and DataType in [DataType.VECTOR_FP32, DataType.VECTOR_FP16]\n ):\n similarities.sort(key=lambda x: x[1], reverse=False) # Ascending order for L2\n\n else:\n similarities.sort(\n key=lambda x: x[1], reverse=True\n ) # Descending order for others\n\n # Special handling for COSINE in FP16 to address precision issues\n if metric_type == MetricType.COSINE and DataType == DataType.VECTOR_FP16:\n # Clamp values to valid cosine distance range [0, 2] and handle floating point errors\n similarities = [(idx, max(0.0, min(2.0, score))) for idx, score in similarities]\n\n return similarities\n\n\ndef get_ground_truth_for_vector_query(\n collection,\n query_vector,\n field_name,\n all_docs,\n query_idx,\n metric_type,\n k,\n use_exact_computation=False,\n):\n if use_exact_computation:\n all_vectors = [doc.vectors[field_name] for doc in all_docs]\n\n for d, f in DEFAULT_VECTOR_FIELD_NAME.items():\n if field_name == f:\n DataType = d\n break\n similarities = compute_exact_similarity_scores(\n [query_vector],\n all_vectors,\n metric_type,\n DataType=DataType,\n QuantizeType=QuantizeType,\n )\n\n if metric_type == MetricType.COSINE and DataType == DataType.VECTOR_FP16:\n # Filter out tiny non-zero values that may be caused by precision errors\n similarities = [\n (idx, max(0.0, min(2.0, score))) for idx, score in similarities\n ]\n\n ground_truth_ids_scores = similarities[:k]\n print(\"Get the most similar k document IDs k:,ground_truth_ids_scores\")\n print(k, ground_truth_ids_scores)\n return ground_truth_ids_scores\n\n else:\n full_result = collection.query(\n VectorQuery(field_name=field_name, vector=query_vector),\n topk=min(len(all_docs), 1024),\n include_vector=True,\n )\n\n ground_truth_ids_scores = [\n (result.id, result.score) for result in full_result[:k]\n ]\n\n if not ground_truth_ids_scores:\n ground_truth_ids_scores = [(all_docs[query_idx].id, 0)]\n\n return ground_truth_ids_scores\n\n\ndef get_ground_truth_map(collection, test_docs, query_vectors_map, metric_type, k):\n ground_truth_map = {}\n\n for field_name, query_vectors in query_vectors_map.items():\n ground_truth_map[field_name] = {}\n\n for i, query_vector in enumerate(query_vectors):\n # Get the ground truth for this query\n relevant_doc_ids_scores = get_ground_truth_for_vector_query(\n collection, query_vector, field_name, test_docs, i, metric_type, k, True\n )\n ground_truth_map[field_name][i] = relevant_doc_ids_scores\n\n print(\"ground_truth_map:\\n\")\n print(ground_truth_map)\n return ground_truth_map\n\n\ndef calculate_recall_at_k(\n collection: Collection,\n test_docs,\n query_vectors_map,\n schema,\n k=1,\n expected_doc_ids_scores_map=None,\n tolerance=0.01,\n):\n recall_stats = {}\n\n for field_name, query_vectors in query_vectors_map.items():\n recall_stats[field_name] = {\n \"relevant_retrieved_count\": 0,\n \"total_relevant_count\": 0,\n \"retrieved_count\": 0,\n \"recall_at_k\": 0.0,\n }\n\n for i, query_vector in enumerate(query_vectors):\n print(\"Starting %dth query\" % i)\n\n query_result_list = collection.query(\n VectorQuery(field_name=field_name, vector=query_vector),\n topk=1024,\n include_vector=True,\n )\n retrieved_count = len(query_result_list)\n\n query_result_ids_scores = []\n for word in query_result_list:\n query_result_ids_scores.append((word.id, word.score))\n\n recall_stats[field_name][\"retrieved_count\"] += retrieved_count\n\n print(\"expected_doc_ids_scores_map:\\n\")\n print(expected_doc_ids_scores_map)\n if i in (expected_doc_ids_scores_map[field_name]):\n expected_relevant_ids_scores = expected_doc_ids_scores_map[field_name][\n i\n ]\n print(\n \"field_name,i,expected_relevant_ids_scores, query_result_ids_scores:\\n\"\n )\n print(\n field_name,\n i,\n \"\\n\",\n expected_relevant_ids_scores,\n \"\\n\",\n len(query_result_ids_scores),\n query_result_ids_scores,\n )\n\n # Update total relevant documents count\n recall_stats[field_name][\"total_relevant_count\"] += len(\n expected_relevant_ids_scores\n )\n\n relevant_found_count = 0\n for ids_scores_except in expected_relevant_ids_scores:\n for ids_scores_result in query_result_ids_scores[:k]:\n if int(ids_scores_result[0]) == int(ids_scores_except[0]):\n relevant_found_count += 1\n break\n elif (\n int(ids_scores_result[0]) != int(ids_scores_except[0])\n and abs(ids_scores_result[1] - ids_scores_except[1])\n <= tolerance\n ):\n print(\"IDs are not equal, but the error is small, tolerance\")\n print(\n ids_scores_result[0],\n ids_scores_except[0],\n ids_scores_result[1],\n ids_scores_except[1],\n tolerance,\n )\n relevant_found_count += 1\n break\n else:\n continue\n\n recall_stats[field_name][\"relevant_retrieved_count\"] += relevant_found_count\n\n # Calculate Recall@K\n if recall_stats[field_name][\"total_relevant_count\"] > 0:\n recall_stats[field_name][\"recall_at_k\"] = (\n recall_stats[field_name][\"relevant_retrieved_count\"]\n / recall_stats[field_name][\"total_relevant_count\"]\n )\n\n return recall_stats\n\n\nclass TestRecall:\n @pytest.mark.parametrize(\n \"full_schema_new\",\n [\n (True, True, HnswIndexParam()),\n (False, True, IVFIndexParam()),\n (False, True, FlatIndexParam()), # ——ok\n (\n True,\n True,\n HnswIndexParam(\n metric_type=MetricType.IP,\n m=16,\n ef_construction=100,\n ),\n ),\n (\n True,\n True,\n HnswIndexParam(\n metric_type=MetricType.COSINE,\n m=24,\n ef_construction=150,\n ),\n ),\n (\n True,\n True,\n HnswIndexParam(\n metric_type=MetricType.L2,\n m=32,\n ef_construction=200,\n ),\n ),\n (\n False,\n True,\n FlatIndexParam(\n metric_type=MetricType.IP,\n ),\n ),\n (\n True,\n True,\n FlatIndexParam(\n metric_type=MetricType.COSINE,\n ),\n ),\n (\n True,\n True,\n FlatIndexParam(\n metric_type=MetricType.L2,\n ),\n ),\n (\n True,\n True,\n IVFIndexParam(\n metric_type=MetricType.IP,\n n_list=100,\n n_iters=10,\n use_soar=False,\n ),\n ),\n (\n True,\n True,\n IVFIndexParam(\n metric_type=MetricType.L2,\n n_list=200,\n n_iters=20,\n use_soar=True,\n ),\n ),\n (\n True,\n True,\n IVFIndexParam(\n metric_type=MetricType.COSINE,\n n_list=150,\n n_iters=15,\n use_soar=False,\n ),\n ),\n ],\n indirect=True,\n )\n @pytest.mark.parametrize(\"doc_num\", [500])\n @pytest.mark.parametrize(\"query_num\", [10])\n @pytest.mark.parametrize(\"top_k\", [1])\n def test_recall_with_single_vector_valid_500(\n self,\n full_collection_new: Collection,\n doc_num,\n query_num,\n top_k,\n full_schema_new,\n request,\n ):\n full_schema_params = request.getfixturevalue(\"full_schema_new\")\n\n for vector_para in full_schema_params.vectors:\n if vector_para.name == \"vector_fp32_field\":\n metric_type = vector_para.index_param.metric_type\n break\n\n multiple_docs = [\n generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num)\n ]\n print(\"len(multiple_docs):\\n\")\n print(len(multiple_docs))\n # print(multiple_docs)\n\n for i in range(10):\n if i != 0:\n pass\n # print(multiple_docs[i * 1000:1000 * (i + 1)])\n batchdoc_and_check(\n full_collection_new,\n multiple_docs[i * 1000 : 1000 * (i + 1)],\n operator=\"insert\",\n )\n\n stats = full_collection_new.stats\n assert stats.doc_count == len(multiple_docs)\n\n doc_ids = [\"0\", \"1\"]\n fetched_docs = full_collection_new.fetch(doc_ids)\n print(\"fetched_docs,multiple_docs\")\n print(\n fetched_docs[doc_ids[0]].vectors[\"sparse_vector_fp32_field\"],\n fetched_docs[doc_ids[0]].vectors[\"sparse_vector_fp16_field\"],\n fetched_docs[doc_ids[1]].vectors[\"sparse_vector_fp32_field\"],\n fetched_docs[doc_ids[1]].vectors[\"sparse_vector_fp16_field\"],\n \"\\n\",\n multiple_docs[0].vectors[\"sparse_vector_fp32_field\"],\n multiple_docs[0].vectors[\"sparse_vector_fp32_field\"],\n multiple_docs[1].vectors[\"sparse_vector_fp32_field\"],\n multiple_docs[1].vectors[\"sparse_vector_fp16_field\"],\n )\n\n full_collection_new.optimize(option=OptimizeOption())\n\n time.sleep(2)\n\n query_vectors_map = {}\n for field_name in DEFAULT_VECTOR_FIELD_NAME.values():\n query_vectors_map[field_name] = [\n multiple_docs[i].vectors[field_name] for i in range(query_num)\n ]\n\n # Get ground truth mapping\n ground_truth_map = get_ground_truth_map(\n full_collection_new, multiple_docs, query_vectors_map, metric_type, top_k\n )\n\n # Validate ground truth mapping structure\n for field_name in DEFAULT_VECTOR_FIELD_NAME.values():\n assert field_name in ground_truth_map\n field_gt = ground_truth_map[field_name]\n assert len(field_gt) == query_num\n\n for query_idx in range(query_num):\n assert query_idx in field_gt\n relevant_ids = field_gt[query_idx]\n assert isinstance(relevant_ids, list)\n assert len(relevant_ids) <= top_k\n\n # Print ground truth statistics\n print(f\"Ground Truth for Top-{top_k} Retrieval:\")\n for field_name, field_gt in ground_truth_map.items():\n print(f\" {field_name}:\")\n for query_idx, relevant_ids in field_gt.items():\n print(\n f\" Query {query_idx}: {len(relevant_ids)} relevant docs - {relevant_ids[:5]}{'...' if len(relevant_ids) > 5 else ''}\"\n )\n\n # Calculate Recall@K using ground truth\n recall_at_k_stats = calculate_recall_at_k(\n full_collection_new,\n multiple_docs,\n query_vectors_map,\n full_schema_new,\n k=top_k,\n expected_doc_ids_scores_map=ground_truth_map,\n tolerance=0.01,\n )\n print(\"ground_truth_map:\\n\")\n print(ground_truth_map)\n\n print(\"(recall_at_k_stats:\\n\")\n print(recall_at_k_stats)\n print(\"metric_type:\")\n print(metric_type)\n # Print Recall@K statistics\n print(f\"Recall@{top_k} using Ground Truth:\")\n for field_name, stats in recall_at_k_stats.items():\n print(f\" {field_name}:\")\n print(\n f\" Relevant Retrieved: {stats['relevant_retrieved_count']}/{stats['total_relevant_count']}\"\n )\n print(f\" Recall@{top_k}: {stats['recall_at_k']:.4f}\")\n for k, v in recall_at_k_stats.items():\n assert v[\"recall_at_k\"] == 1.0\n\n @pytest.mark.parametrize(\n \"full_schema_new\",\n [\n (True, True, HnswIndexParam()),\n (False, True, IVFIndexParam()),\n (False, True, FlatIndexParam()), # ——ok\n (\n True,\n True,\n HnswIndexParam(\n metric_type=MetricType.IP,\n m=16,\n ef_construction=100,\n ),\n ),\n (\n True,\n True,\n HnswIndexParam(\n metric_type=MetricType.COSINE,\n m=24,\n ef_construction=150,\n ),\n ),\n # (True, True, HnswIndexParam(metric_type=MetricType.L2, m=32, ef_construction=200, )),\n (\n False,\n True,\n FlatIndexParam(\n metric_type=MetricType.IP,\n ),\n ),\n (\n True,\n True,\n FlatIndexParam(\n metric_type=MetricType.COSINE,\n ),\n ),\n # (True, True, FlatIndexParam(metric_type=MetricType.L2, )),\n (\n True,\n True,\n IVFIndexParam(\n metric_type=MetricType.IP,\n n_list=100,\n n_iters=10,\n use_soar=False,\n ),\n ),\n (\n True,\n True,\n IVFIndexParam(\n metric_type=MetricType.L2,\n n_list=200,\n n_iters=20,\n use_soar=True,\n ),\n ),\n # (True, True, IVFIndexParam(metric_type=MetricType.COSINE, n_list=150, n_iters=15, use_soar=False, )),\n ],\n indirect=True,\n )\n @pytest.mark.parametrize(\"doc_num\", [2000])\n @pytest.mark.parametrize(\"query_num\", [2])\n @pytest.mark.parametrize(\"top_k\", [1])\n @pytest.mark.skip(reason=\"known bug\")\n def test_recall_with_single_vector_valid_2000(\n self,\n full_collection_new: Collection,\n doc_num,\n query_num,\n top_k,\n full_schema_new,\n request,\n ):\n full_schema_params = request.getfixturevalue(\"full_schema_new\")\n\n for vector_para in full_schema_params.vectors:\n if vector_para.name == \"vector_fp32_field\":\n metric_type = vector_para.index_param.metric_type\n break\n\n multiple_docs = [\n generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num)\n ]\n print(\"len(multiple_docs):\\n\")\n print(len(multiple_docs))\n # print(multiple_docs)\n\n for i in range(10):\n if i != 0:\n pass\n # print(multiple_docs[i * 1000:1000 * (i + 1)])\n batchdoc_and_check(\n full_collection_new,\n multiple_docs[i * 1000 : 1000 * (i + 1)],\n operator=\"insert\",\n )\n\n stats = full_collection_new.stats\n assert stats.doc_count == len(multiple_docs)\n\n doc_ids = [\"0\", \"1\"]\n fetched_docs = full_collection_new.fetch(doc_ids)\n print(\"fetched_docs,multiple_docs\")\n print(\n fetched_docs[doc_ids[0]].vectors[\"sparse_vector_fp32_field\"],\n fetched_docs[doc_ids[0]].vectors[\"sparse_vector_fp16_field\"],\n fetched_docs[doc_ids[1]].vectors[\"sparse_vector_fp32_field\"],\n fetched_docs[doc_ids[1]].vectors[\"sparse_vector_fp16_field\"],\n \"\\n\",\n multiple_docs[0].vectors[\"sparse_vector_fp32_field\"],\n multiple_docs[0].vectors[\"sparse_vector_fp32_field\"],\n multiple_docs[1].vectors[\"sparse_vector_fp32_field\"],\n multiple_docs[1].vectors[\"sparse_vector_fp16_field\"],\n )\n\n full_collection_new.optimize(option=OptimizeOption())\n\n time.sleep(2)\n\n query_vectors_map = {}\n for field_name in DEFAULT_VECTOR_FIELD_NAME.values():\n query_vectors_map[field_name] = [\n multiple_docs[i].vectors[field_name] for i in range(query_num)\n ]\n\n # Get ground truth mapping\n ground_truth_map = get_ground_truth_map(\n full_collection_new, multiple_docs, query_vectors_map, metric_type, top_k\n )\n\n # Validate ground truth mapping structure\n for field_name in DEFAULT_VECTOR_FIELD_NAME.values():\n assert field_name in ground_truth_map\n field_gt = ground_truth_map[field_name]\n assert len(field_gt) == query_num\n\n for query_idx in range(query_num):\n assert query_idx in field_gt\n relevant_ids = field_gt[query_idx]\n assert isinstance(relevant_ids, list)\n assert len(relevant_ids) <= top_k\n\n # Print ground truth statistics\n print(f\"Ground Truth for Top-{top_k} Retrieval:\")\n for field_name, field_gt in ground_truth_map.items():\n print(f\" {field_name}:\")\n for query_idx, relevant_ids in field_gt.items():\n print(\n f\" Query {query_idx}: {len(relevant_ids)} relevant docs - {relevant_ids[:5]}{'...' if len(relevant_ids) > 5 else ''}\"\n )\n\n # Calculate Recall@K using ground truth\n recall_at_k_stats = calculate_recall_at_k(\n full_collection_new,\n multiple_docs,\n query_vectors_map,\n full_schema_new,\n k=top_k,\n expected_doc_ids_scores_map=ground_truth_map,\n tolerance=0.01,\n )\n print(\"ground_truth_map:\\n\")\n print(ground_truth_map)\n\n print(\"(recall_at_k_stats:\\n\")\n print(recall_at_k_stats)\n print(\"metric_type:\")\n print(metric_type)\n # Print Recall@K statistics\n print(f\"Recall@{top_k} using Ground Truth:\")\n for field_name, stats in recall_at_k_stats.items():\n print(f\" {field_name}:\")\n print(\n f\" Relevant Retrieved: {stats['relevant_retrieved_count']}/{stats['total_relevant_count']}\"\n )\n print(f\" Recall@{top_k}: {stats['recall_at_k']:.4f}\")\n for k, v in recall_at_k_stats.items():\n assert v[\"recall_at_k\"] == 1.0\n" }, { "path": "python/tests/detail/test_db_config.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport pytest\nimport tempfile\nimport os\nimport sys\nimport subprocess\n\nimport zvec\nimport zvec\nfrom zvec import LogType, LogLevel\n\n# Error messages\nINITIALIZATION_ERROR_MSG = \"initialization failed\"\nRUNTIME_ERROR_MSG = \"RuntimeError\"\nVALUE_ERROR_MSG = \"ValueError\"\nTYPE_ERROR_MSG = \"TypeError\"\n\n\n# ==================== helper ====================\ndef run_in_subprocess(func):\n def wrapper(*args, **kwargs):\n if os.getenv(\"RUNNING_IN_SUBPROCESS\"):\n return func(*args, **kwargs)\n\n env = os.environ.copy()\n env[\"RUNNING_IN_SUBPROCESS\"] = \"1\"\n env[\"PYTEST_CURRENT_TEST\"] = func.__name__\n\n import inspect\n\n filepath = inspect.getfile(func)\n qualname = func.__qualname__.replace(\".\", \"::\")\n test_id = f\"{filepath}::{qualname}\"\n\n project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))\n env[\"PYTHONPATH\"] = project_root + \":\" + env.get(\"PYTHONPATH\", \"\")\n\n cmd = [sys.executable, \"-m\", \"pytest\", \"-v\", \"-s\", test_id]\n\n result = subprocess.run(cmd, env=env, capture_output=True, text=True)\n if result.returncode != 0:\n pytest.fail(\n f\"Subprocess test {func.__name__} failed with code {result.returncode}\\n\"\n f\"STDOUT:\\n{result.stdout}\\nSTDERR:\\n{result.stderr}\"\n )\n\n return wrapper\n\n\n# ==================== Fixtures ====================\n@pytest.fixture(scope=\"function\")\ndef temp_log_dir(tmp_path_factory):\n return tmp_path_factory.mktemp(\"logs\")\n\n\n# ==================== Tests ====================\nclass TestDbConfigInitialization:\n @run_in_subprocess\n def test_init_default(self):\n # default config\n # log_type: Optional[LogType] = LogType.CONSOLE,\n # log_level: Optional[LogLevel] = LogLevel.WARN,\n # log_dir: Optional[str] = \"./logs\",\n # log_basename: Optional[str] = \"zvec.log\",\n # log_file_size: Optional[int] = 2048,\n # log_overdue_days: Optional[int] = 7,\n zvec.init()\n\n @run_in_subprocess\n def test_init_file_logger(self):\n from pathlib import Path\n import shutil\n\n zvec.init(\n log_level=LogLevel.DEBUG,\n log_type=LogType.FILE,\n )\n # assert logdir exist\n log_dir = Path(\"./logs\")\n assert log_dir.exists()\n\n # validate write log\n col = zvec.create_and_open(\n \"/tmp/test/1\",\n zvec.CollectionSchema(\n name=\"test\",\n vectors=zvec.VectorSchema(\n dimension=4,\n data_type=zvec.DataType.VECTOR_FP32,\n name=\"image\",\n ),\n ),\n )\n col.insert(docs=[zvec.Doc(id=\"1\", vectors={\"image\": [1.0, 2.0, 3.0, 4.0]})])\n assert any(log_dir.glob(\"zvec.log.*\"))\n\n # clear\n col.destroy()\n shutil.rmtree(log_dir, ignore_errors=True)\n\n @run_in_subprocess\n def test_init_with_mixed_config(self):\n zvec.init(\n memory_limit_mb=128,\n log_type=LogType.FILE,\n query_threads=1,\n log_level=LogLevel.WARN,\n )\n\n @run_in_subprocess\n def test_repeated_initialization(self):\n # Calling init() repeatedly is allowed:\n # it succeeds but becomes a no-op after the first successful init()\n zvec.init()\n\n\nclass TestDbConfigMemoryLimitValidation:\n @run_in_subprocess\n def test_memory_limit_min_valid(self):\n # MIN_MEMORY_LIMIT_BYTES is 100M\n with pytest.raises(RuntimeError):\n zvec.init(memory_limit_mb=99)\n\n @run_in_subprocess\n def test_memory_limit_invalid_value(self):\n # memory_limit_mb must >= 0 and must be int and if None, set default value\n with pytest.raises(ValueError):\n zvec.init(memory_limit_mb=0)\n with pytest.raises(ValueError):\n zvec.init(memory_limit_mb=-1)\n with pytest.raises(TypeError):\n zvec.init(memory_limit_mb=\"512\")\n with pytest.raises(TypeError):\n zvec.init(memory_limit_mb=512.5)\n\n\nclass TestDbConfigThreadValidation:\n @run_in_subprocess\n def test_query_threads(self):\n zvec.init(query_threads=1)\n\n @run_in_subprocess\n def test_query_threads_invalid(self):\n # query_threads must >= 0 and must be int and if None, set default value\n with pytest.raises(ValueError):\n zvec.init(query_threads=0)\n with pytest.raises(ValueError):\n zvec.init(query_threads=-1)\n with pytest.raises(TypeError):\n zvec.init(query_threads=\"value\")\n with pytest.raises(TypeError):\n zvec.init(query_threads=512.5)\n with pytest.raises(TypeError):\n zvec.init(query_threads=\"512\")\n\n @run_in_subprocess\n def test_optimize_threads(self):\n zvec.init(optimize_threads=1)\n\n @run_in_subprocess\n def test_optimize_threads_invalid(self):\n # optimize_threads must >= 0 and must be int and if None, set default value\n with pytest.raises(ValueError):\n zvec.init(optimize_threads=0)\n with pytest.raises(ValueError):\n zvec.init(optimize_threads=-1)\n with pytest.raises(TypeError):\n zvec.init(optimize_threads=\"value\")\n with pytest.raises(TypeError):\n zvec.init(optimize_threads=512.5)\n with pytest.raises(TypeError):\n zvec.init(optimize_threads=\"512\")\n\n\nclass TestDbConfigRatioValidation:\n @run_in_subprocess\n def test_init_invert_to_forward_scan_ratio(self):\n # must be in [0,1]\n zvec.init(invert_to_forward_scan_ratio=0.8)\n\n @run_in_subprocess\n def test_init_invert_to_forward_scan_ratio_invalid(self):\n with pytest.raises(ValueError):\n zvec.init(invert_to_forward_scan_ratio=1.1)\n with pytest.raises(ValueError):\n zvec.init(invert_to_forward_scan_ratio=-0.1)\n with pytest.raises(TypeError):\n zvec.init(invert_to_forward_scan_ratio=\"0.8\")\n\n @run_in_subprocess\n def test_init_brute_force_by_keys_ratio(self):\n zvec.init(brute_force_by_keys_ratio=0.8)\n\n @run_in_subprocess\n def test_init_brute_force_by_keys_ratio_invalid(self):\n with pytest.raises(ValueError):\n zvec.init(brute_force_by_keys_ratio=1.1)\n with pytest.raises(ValueError):\n zvec.init(brute_force_by_keys_ratio=-0.1)\n with pytest.raises(TypeError):\n zvec.init(brute_force_by_keys_ratio=\"0.8\")\n\n\nclass TestDbConfigLogValidation:\n @run_in_subprocess\n def test_log_type_valid(self):\n zvec.init(log_type=LogType.CONSOLE)\n\n @run_in_subprocess\n def test_log_type_invalid(self):\n with pytest.raises(TypeError):\n zvec.init(log_type=\"FILE\")\n with pytest.raises(TypeError):\n zvec.init(log_type=\"\")\n with pytest.raises(TypeError):\n zvec.init(log_type=\"invalid\")\n with pytest.raises(TypeError):\n zvec.init(log_type=123)\n\n @run_in_subprocess\n def test_log_level_valid(self):\n zvec.init(log_level=LogLevel.ERROR)\n\n @run_in_subprocess\n def test_log_level_invalid(self):\n with pytest.raises(TypeError):\n zvec.init(log_level=\"WARN\")\n with pytest.raises(TypeError):\n zvec.init(log_level=\"\")\n with pytest.raises(TypeError):\n zvec.init(log_level=\"invalid\")\n with pytest.raises(TypeError):\n zvec.init(log_level=123)\n\n @run_in_subprocess\n def test_init_file_logger(self):\n from pathlib import Path\n import shutil\n\n temp_dir = tempfile.mkdtemp(prefix=\"log_test_\")\n abs_temp_dir = os.path.abspath(temp_dir)\n\n zvec.init(\n log_level=LogLevel.DEBUG,\n log_type=LogType.FILE,\n log_dir=abs_temp_dir,\n log_basename=\"test\",\n )\n\n # assert logdir exist\n log_dir = Path(abs_temp_dir)\n assert log_dir.exists()\n\n # validate write log\n col = zvec.create_and_open(\n \"/tmp/test/1\",\n zvec.CollectionSchema(\n name=\"test\",\n vectors=zvec.VectorSchema(\n dimension=4,\n data_type=zvec.DataType.VECTOR_FP32,\n name=\"image\",\n ),\n ),\n )\n col.insert(docs=[zvec.Doc(id=\"1\", vectors={\"image\": [1.0, 2.0, 3.0, 4.0]})])\n assert any(log_dir.glob(\"test.*\"))\n\n # clear\n col.destroy()\n shutil.rmtree(log_dir, ignore_errors=True)\n\n @run_in_subprocess\n def test_log_file_size_invalid(self):\n with pytest.raises(TypeError):\n zvec.init(log_type=LogType.FILE, log_file_size=\"df\")\n\n with pytest.raises(ValueError):\n zvec.init(log_type=LogType.FILE, log_file_size=0)\n\n with pytest.raises(ValueError):\n zvec.init(log_type=LogType.FILE, log_file_size=-1)\n\n @run_in_subprocess\n def test_log_overdue_days_invalid(self):\n with pytest.raises(TypeError):\n zvec.init(log_type=LogType.FILE, log_overdue_days=\"df\")\n\n with pytest.raises(ValueError):\n zvec.init(log_type=LogType.FILE, log_overdue_days=0)\n\n with pytest.raises(ValueError):\n zvec.init(log_type=LogType.FILE, log_overdue_days=-1)\n" }, { "path": "python/tests/test_collection.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\n\nimport pytest\nimport zvec\nfrom zvec import (\n Collection,\n CollectionOption,\n DataType,\n Doc,\n FieldSchema,\n HnswIndexParam,\n InvertIndexParam,\n LogLevel,\n LogType,\n VectorSchema,\n StatusCode,\n IndexOption,\n IndexType,\n VectorQuery,\n OptimizeOption,\n)\n\n# ==================== Common ====================\n\n\n@pytest.fixture(scope=\"session\")\ndef collection_schema():\n return zvec.CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\", DataType.STRING, nullable=False, index_param=InvertIndexParam()\n ),\n FieldSchema(\"weight\", DataType.FLOAT, nullable=True),\n FieldSchema(\"height\", DataType.INT32, nullable=True),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n VectorSchema(\n \"sparse\", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()\n ),\n ],\n )\n\n\n@pytest.fixture(scope=\"session\")\ndef collection_option():\n return CollectionOption(read_only=False, enable_mmap=True)\n\n\n@pytest.fixture\ndef single_doc():\n id = 0\n return Doc(\n id=f\"{id}\",\n fields={\"id\": id, \"name\": \"test\", \"weight\": 80.0, \"height\": id + 140},\n vectors={\"dense\": [id + 0.1] * 128, \"sparse\": {1: 1.0, 2: 2.0, 3: 3.0}},\n )\n\n\n@pytest.fixture\ndef multiple_docs():\n return [\n Doc(\n id=f\"{id}\",\n fields={\"id\": id, \"name\": \"test\", \"weight\": 80.0, \"height\": 210},\n vectors={\"dense\": [id + 0.1] * 128, \"sparse\": {1: 1.0, 2: 2.0, 3: 3.0}},\n )\n for id in range(1, 101)\n ]\n\n\n@pytest.fixture(scope=\"function\")\ndef test_collection(\n tmp_path_factory, collection_schema, collection_option\n) -> Collection:\n \"\"\"\n Function-scoped fixture: creates and opens a collection.\n Uses tmp_path_factory to ensure shared temp dir per class.\n \"\"\"\n # Create unique temp directory for this test class\n temp_dir = tmp_path_factory.mktemp(\"zvec\")\n collection_path = temp_dir / \"test_collection\"\n\n coll = zvec.create_and_open(\n path=str(collection_path), schema=collection_schema, option=collection_option\n )\n\n assert coll is not None, \"Failed to create and open collection\"\n assert coll.path == str(collection_path)\n assert coll.schema.name == collection_schema.name\n assert list(coll.schema.fields) == list(collection_schema.fields)\n assert list(coll.schema.vectors) == list(collection_schema.vectors)\n assert coll.option.read_only == collection_option.read_only\n assert coll.option.enable_mmap == collection_option.enable_mmap\n\n try:\n yield coll\n finally:\n if hasattr(coll, \"destroy\") and coll is not None:\n try:\n coll.destroy()\n except Exception as e:\n print(f\"Warning: failed to destroy collection: {e}\")\n\n\n@pytest.fixture\ndef collection_with_single_doc(test_collection: Collection, single_doc) -> Collection:\n # Setup: insert single doc\n assert test_collection.stats.doc_count == 0\n result = test_collection.insert(single_doc)\n assert bool(result)\n assert result.ok()\n assert test_collection.stats.doc_count == 1\n\n yield test_collection\n\n # Teardown: delete single doc\n test_collection.delete(single_doc.id)\n assert test_collection.stats.doc_count == 0\n\n\n@pytest.fixture\ndef collection_with_multiple_docs(\n test_collection: Collection, multiple_docs\n) -> Collection:\n # Setup: insert multiple docs\n assert test_collection.stats.doc_count == 0\n result = test_collection.insert(multiple_docs)\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok()\n assert test_collection.stats.doc_count == len(multiple_docs)\n\n yield test_collection\n\n # Teardown: delete multiple docs\n test_collection.delete([doc.id for doc in multiple_docs])\n\n\n# ==================== Tests ====================\n\n\n# ----------------------------\n# Config Test Case\n# ----------------------------\nclass TestConfig:\n def test_config(self):\n zvec.init(log_type=LogType.CONSOLE, log_level=LogLevel.ERROR, log_dir=\"./log\")\n\n\n# ----------------------------\n# Collection DDL Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionDDL:\n def test_collection_stats(self, test_collection: Collection):\n assert test_collection.stats is not None\n stats = test_collection.stats\n assert stats.doc_count == 0\n assert len(stats.index_completeness) == 2\n assert stats.index_completeness[\"dense\"] == 1\n assert stats.index_completeness[\"sparse\"] == 1\n\n\n# ----------------------------\n# Collection Index DDL Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionIndexDDL:\n def test_create_index(self, test_collection: Collection):\n # before create\n field_schema = test_collection.schema.field(\"weight\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.FLOAT\n assert field_schema.name == \"weight\"\n index_param = field_schema.index_param\n assert index_param is None\n\n # create\n test_collection.create_index(\n field_name=\"weight\", index_param=InvertIndexParam(), option=IndexOption()\n )\n assert test_collection.schema is not None\n field_schema = test_collection.schema.field(\"weight\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.FLOAT\n assert field_schema.name == \"weight\"\n\n index_param = field_schema.index_param\n assert index_param.type == IndexType.INVERT\n assert index_param.enable_range_optimization is False\n assert index_param.enable_extended_wildcard is False\n\n def test_drop_index(self, test_collection: Collection):\n # before drop\n field_schema = test_collection.schema.field(\"name\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.STRING\n assert field_schema.name == \"name\"\n index_param = field_schema.index_param\n assert index_param.type == IndexType.INVERT\n assert index_param.enable_range_optimization is False\n assert index_param.enable_extended_wildcard is False\n\n # drop\n test_collection.drop_index(\"name\")\n field_schema = test_collection.schema.field(\"name\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.STRING\n assert field_schema.name == \"name\"\n\n # without index\n index_param = field_schema.index_param\n assert index_param is None\n\n def test_create_index_field_is_not_exist(self, test_collection: Collection):\n with pytest.raises(Exception) as e:\n test_collection.create_index(\n field_name=\"not_exist\",\n index_param=InvertIndexParam(),\n )\n\n index_param = field_schema.index_param\n assert index_param.type == IndexType.INVERT\n assert index_param.enable_range_optimization is False\n assert index_param.enable_extended_wildcard is False\n\n def test_drop_index(self, test_collection: Collection):\n # before drop\n field_schema = test_collection.schema.field(\"name\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.STRING\n assert field_schema.name == \"name\"\n index_param = field_schema.index_param\n assert index_param.type == IndexType.INVERT\n assert index_param.enable_range_optimization is False\n assert index_param.enable_extended_wildcard is False\n\n # drop\n test_collection.drop_index(\"name\")\n field_schema = test_collection.schema.field(\"name\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.STRING\n assert field_schema.name == \"name\"\n\n # without index\n index_param = field_schema.index_param\n assert index_param is None\n\n def test_create_index_field_is_not_exist(self, test_collection: Collection):\n with pytest.raises(Exception) as e:\n test_collection.create_index(\n field_name=\"not_exist\",\n index_param=InvertIndexParam(),\n )\n\n\n# ----------------------------\n# Collection Column DDL Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionColumnDDL:\n def test_create_column(self, test_collection: Collection):\n # before create column\n field_schema = test_collection.schema.field(\"age\")\n assert field_schema is None\n\n # create\n test_collection.add_column(FieldSchema(\"age\", DataType.INT32, nullable=True))\n\n field_schema = test_collection.schema.field(\"age\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.INT32\n assert field_schema.name == \"age\"\n assert field_schema.index_param is None\n\n def test_create_column_is_nullable(self, test_collection: Collection):\n with pytest.raises(ValueError):\n test_collection.add_column(\n FieldSchema(\"age\", DataType.INT32, nullable=False)\n )\n\n def test_drop_column(self, test_collection: Collection):\n # before drop column\n field_schema = test_collection.schema.field(\"id\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.INT64\n assert field_schema.name == \"id\"\n index_param = field_schema.index_param\n assert index_param is not None\n assert index_param.type == IndexType.INVERT\n\n # drop\n test_collection.drop_column(\"id\")\n field_schema = test_collection.schema.field(\"id\")\n assert field_schema is None\n\n def test_alert_column_to_rename(self, test_collection: Collection):\n # before alert column\n field_schema = test_collection.schema.field(\"id\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.INT64\n assert field_schema.name == \"id\"\n index_param = field_schema.index_param\n assert index_param is not None\n assert index_param.type == IndexType.INVERT\n assert index_param.enable_range_optimization is True\n assert index_param.enable_extended_wildcard is False\n\n # alert rename\n test_collection.alter_column(\"id\", \"doc_id\")\n\n # validate old column\n field_schema = test_collection.schema.field(\"id\")\n assert field_schema is None\n # validate rename column\n field_schema = test_collection.schema.field(\"doc_id\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.INT64\n assert field_schema.name == \"doc_id\"\n assert field_schema.nullable is False\n index_param = field_schema.index_param\n assert index_param is not None\n assert index_param.type == IndexType.INVERT\n assert index_param.enable_range_optimization is True\n assert index_param.enable_extended_wildcard is False\n\n def test_alert_column_to_modify_schema(self, test_collection: Collection):\n # before alert column\n field_schema = test_collection.schema.field(\"id\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.INT64\n assert field_schema.name == \"id\"\n index_param = field_schema.index_param\n assert index_param.type == IndexType.INVERT\n\n test_collection.alter_column(\n old_name=\"id\",\n field_schema=FieldSchema(\"doc_id\", DataType.UINT64, nullable=True),\n )\n field_schema = test_collection.schema.field(\"doc_id\")\n assert field_schema is not None\n assert field_schema.data_type == DataType.UINT64\n assert field_schema.name == \"doc_id\"\n\n def test_column_with_other_dtype(self, test_collection: Collection):\n # only allow number type\n test_collection.add_column(FieldSchema(\"age\", DataType.INT32, nullable=True))\n\n with pytest.raises(ValueError):\n test_collection.add_column(FieldSchema(\"full_name\", DataType.STRING))\n with pytest.raises(ValueError):\n test_collection.drop_column(\"name\")\n with pytest.raises(ValueError):\n test_collection.alter_column(old_name=\"name\", new_name=\"full_name\")\n with pytest.raises(ValueError):\n test_collection.alter_column(\n old_name=\"name\", field_schema=FieldSchema(\"full_name\", DataType.STRING)\n )\n\n\n# ----------------------------\n# Collection Optimize Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionOptimize:\n def test_collection_optimize(self, test_collection: Collection):\n test_collection.optimize(option=OptimizeOption())\n\n\n# ----------------------------\n# Collection Fetch Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionFetch:\n def test_collection_fetch(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n result = collection_with_single_doc.fetch(ids=[single_doc.id])\n assert bool(result)\n assert single_doc.id in result.keys()\n\n doc = result[single_doc.id]\n assert doc is not None\n assert doc.id == single_doc.id\n assert set(doc.field_names()) == set(single_doc.field_names())\n for field_name in doc.field_names():\n if field_name in [\"dense\", \"sparse\"]:\n continue\n assert doc.field(field_name) == single_doc.field(field_name)\n\n def test_collection_fetch_contains_nodata_ids(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n ids = [doc.id for doc in multiple_docs]\n no_data_key = \"x\"\n ids_with_no_data = [no_data_key] + ids\n result = collection_with_multiple_docs.fetch(ids=ids_with_no_data)\n assert bool(result)\n assert len(result) == len(ids)\n assert no_data_key not in result\n\n\n# ----------------------------\n# Collection Insert Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionInsert:\n def test_collection_insert(self, test_collection, single_doc):\n result = test_collection.insert(single_doc)\n assert bool(result)\n assert result.ok()\n stats = test_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_collection_insert_with_nullable_false_field(self, test_collection):\n # id, name's nullable == False\n # weight, height's nullable == True\n\n doc = Doc(\n id=\"0\",\n fields={\n \"id\": 1,\n \"name\": \"test\",\n },\n vectors={\"dense\": [1 + 0.1] * 128, \"sparse\": {1: 1.0, 2: 2.0, 3: 3.0}},\n )\n result = test_collection.insert(doc)\n assert bool(result)\n assert result.ok()\n stats = test_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_collection_insert_without_nullable_false_field(self, test_collection):\n # id, name's nullable == False\n # weight, height's nullable == True\n\n # without id, name\n doc = Doc(\n id=\"0\",\n vectors={\"dense\": [1 + 0.1] * 128, \"sparse\": {1: 1.0, 2: 2.0, 3: 3.0}},\n )\n with pytest.raises(ValueError) as e:\n # ValueError: doc validate failed: field[id] is configured not nullable,\n # but doc does not contain this field\n test_collection.insert(doc)\n assert \"field[id] is configured not nullable\" in str(e.value)\n\n # without name\n doc = Doc(\n id=\"0\",\n fields={\n \"id\": 1,\n },\n vectors={\"dense\": [1 + 0.1] * 128, \"sparse\": {1: 1.0, 2: 2.0, 3: 3.0}},\n )\n with pytest.raises(ValueError) as e:\n test_collection.insert(doc)\n assert \"field[name] is configured not nullable\" in str(e.value)\n\n def test_collection_insert_with_nullable_true_field(self, test_collection):\n # id, name's nullable == False\n # weight, height's nullable == True\n\n doc = Doc(\n id=\"0\",\n fields={\n \"id\": 1,\n \"name\": \"test\",\n },\n vectors={\"dense\": [1 + 0.1] * 128, \"sparse\": {1: 1.0, 2: 2.0, 3: 3.0}},\n )\n result = test_collection.insert(doc)\n assert bool(result)\n assert result.ok()\n stats = test_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n result = test_collection.fetch(ids=[doc.id])\n assert doc.id in result\n ret = result[doc.id]\n assert ret.field(\"id\") == 1\n assert ret.field(\"name\") == \"test\"\n assert ret.field(\"weight\") is None\n assert ret.field(\"height\") is None\n\n def test_collection_insert_batch(self, test_collection, multiple_docs):\n result = test_collection.insert(multiple_docs)\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok()\n\n stats = test_collection.stats\n assert stats is not None\n assert stats.doc_count == len(multiple_docs)\n\n def test_collection_insert_duplicate(\n self, test_collection, single_doc, multiple_docs\n ):\n test_collection.insert(single_doc)\n result = test_collection.insert(single_doc)\n assert bool(result)\n assert result.code() == StatusCode.ALREADY_EXISTS\n\n stats = test_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n\n# ----------------------------\n# Collection Update Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionUpdate:\n def test_empty_collection_update(\n self, test_collection: Collection, single_doc: Doc\n ):\n result = test_collection.update(single_doc)\n assert bool(result)\n assert result.code() == StatusCode.NOT_FOUND\n\n stats = test_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n def test_collection_update_with_nullable_false_field(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n # id, name's nullable == False\n # weight, height's nullable == True\n\n # update doc field id\n doc = Doc(\n id=single_doc.id,\n fields={\"id\": single_doc.field(\"id\") + 1},\n )\n result = collection_with_single_doc.update(doc)\n assert bool(result)\n assert result.ok()\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n # fetch\n result = collection_with_single_doc.fetch(ids=[doc.id])\n assert doc.id in result\n ret = result[doc.id]\n assert ret.field(\"id\") == doc.field(\"id\")\n assert ret.field(\"name\") == single_doc.field(\"name\")\n assert ret.field(\"weight\") == single_doc.field(\"weight\")\n assert ret.field(\"height\") == single_doc.field(\"height\")\n\n def test_collection_update_with_nullable_false_field_is_none(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n # id, name's nullable == False\n # weight, height's nullable == True\n\n # update doc field id\n doc = Doc(\n id=single_doc.id,\n fields={\"id\": None},\n )\n with pytest.raises(ValueError) as e:\n # ValueError: doc validate failed: field[id] is configured not nullable,\n # but doc does not contain this field\n collection_with_single_doc.update(doc)\n\n doc = Doc(\n id=single_doc.id,\n fields={\"id\": single_doc.field(\"id\") + 1, \"weight\": None},\n )\n\n result = collection_with_single_doc.update(doc)\n assert bool(result)\n assert result.ok()\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n ret = collection_with_single_doc.fetch(ids=[doc.id])\n assert doc.id in ret\n ret = ret[doc.id]\n assert ret.field(\"id\") == doc.field(\"id\")\n assert ret.field(\"name\") == single_doc.field(\"name\")\n assert ret.field(\"weight\") is None\n assert ret.field(\"height\") == single_doc.field(\"height\")\n\n def test_collection_update_without_nullable_false_field(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n # id, name's nullable == False\n # weight, height's nullable == True\n\n # update doc field weight\n doc = Doc(\n id=single_doc.id,\n fields={\"weight\": single_doc.field(\"weight\") + 1},\n )\n result = collection_with_single_doc.update(doc)\n assert bool(result)\n assert result.ok()\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n # fetch\n ret = collection_with_single_doc.fetch(ids=[doc.id])\n assert doc.id in ret\n ret = ret[doc.id]\n assert ret.field(\"id\") == single_doc.field(\"id\")\n assert ret.field(\"name\") == single_doc.field(\"name\")\n assert ret.field(\"weight\") == doc.field(\"weight\")\n assert ret.field(\"height\") == single_doc.field(\"height\")\n\n def test_collection_update_without_nullable_false_field_set_null(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n # id, name's nullable == False\n # weight, height's nullable == True\n\n # update doc field weight is None\n doc = Doc(\n id=single_doc.id,\n fields={\"weight\": None},\n )\n result = collection_with_single_doc.update(doc)\n assert bool(result)\n assert result.ok()\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n # fetch\n ret = collection_with_single_doc.fetch(ids=[doc.id])\n assert doc.id in ret\n ret = ret[doc.id]\n assert ret.field(\"id\") == single_doc.field(\"id\")\n assert ret.field(\"name\") == single_doc.field(\"name\")\n assert ret.field(\"weight\") is None\n assert ret.field(\"height\") == single_doc.field(\"height\")\n\n def test_empty_collection_update_batch(\n self, test_collection: Collection, multiple_docs\n ):\n result = test_collection.update(multiple_docs)\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.code() == StatusCode.NOT_FOUND\n\n stats = test_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n def test_collection_update(\n self, collection_with_single_doc: Collection, single_doc\n ):\n result = collection_with_single_doc.update(single_doc)\n assert bool(result) == 1\n assert result.ok()\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_collection_update_batch(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n result = collection_with_multiple_docs.update(multiple_docs)\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok()\n\n stats = collection_with_multiple_docs.stats\n assert stats is not None\n assert stats.doc_count == len(multiple_docs)\n\n\n# ----------------------------\n# Collection Upsert Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionUpsert:\n def test_empty_collection_upsert(self, test_collection: Collection, single_doc):\n result = test_collection.upsert(single_doc)\n assert bool(result)\n assert result.ok()\n\n stats = test_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_empty_collection_upsert_batch(\n self, test_collection: Collection, multiple_docs\n ):\n result = test_collection.upsert(multiple_docs)\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok()\n\n stats = test_collection.stats\n assert stats is not None\n assert stats.doc_count == len(multiple_docs)\n\n def test_collection_upsert(\n self, collection_with_single_doc: Collection, single_doc, multiple_docs\n ):\n # doc is existing\n # upsert => update\n result = collection_with_single_doc.upsert(single_doc)\n assert bool(result)\n assert result.ok()\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_collection_upsert_batch(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n # doc is existing\n # upsert => update\n result = collection_with_multiple_docs.upsert(multiple_docs)\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok()\n\n stats = collection_with_multiple_docs.stats\n assert stats is not None\n assert stats.doc_count == len(multiple_docs)\n\n\n# ----------------------------\n# Collection Upsert Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionDelete:\n def test_empty_collection_delete(self, test_collection: Collection, single_doc):\n result = test_collection.delete(single_doc.id)\n assert bool(result)\n assert result.code() == StatusCode.NOT_FOUND\n\n def test_empty_collection_delete_batch(\n self, test_collection: Collection, multiple_docs\n ):\n result = test_collection.delete([doc.id for doc in multiple_docs])\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.code() == StatusCode.NOT_FOUND\n\n def test_collection_delete(\n self, collection_with_single_doc: Collection, single_doc\n ):\n result = collection_with_single_doc.delete(single_doc.id)\n assert bool(result)\n assert result.ok()\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n result = collection_with_single_doc.insert(single_doc)\n assert bool(result)\n assert result.ok()\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_collection_delete_batch(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n result = collection_with_multiple_docs.delete([doc.id for doc in multiple_docs])\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok()\n stats = collection_with_multiple_docs.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n def test_collection_delete_by_filter(\n self, collection_with_single_doc: Collection, single_doc\n ):\n collection_with_single_doc.delete_by_filter(\n filter=f\"height={single_doc.field('height')}\"\n )\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n def test_collection_delete_by_filter_invert_field(\n self, collection_with_single_doc: Collection, single_doc\n ):\n collection_with_single_doc.delete_by_filter(\n filter=f\"id={single_doc.field('id')}\"\n )\n stats = collection_with_single_doc.stats\n assert stats is not None\n assert stats.doc_count == 0\n\n\n# ----------------------------\n# Collection Upsert Test Case\n# ----------------------------\n@pytest.mark.usefixtures(\"test_collection\")\nclass TestCollectionQuery:\n def test_empty_collection_query(self, test_collection: Collection):\n result = test_collection.query()\n assert len(result) == 0\n\n def test_collection_query(self, collection_with_single_doc: Collection, single_doc):\n result = collection_with_single_doc.query()\n assert len(result) == 1\n doc = result[0]\n assert doc.id == single_doc.id\n assert \"dense\" not in doc.field_names()\n assert \"sparse\" not in doc.field_names()\n field_without_vector = single_doc.field_names()\n assert set(doc.field_names()) == set(field_without_vector)\n for name in field_without_vector:\n assert doc.field(name) == single_doc.field(name)\n\n def test_collection_query_with_include_vector(\n self, collection_with_single_doc: Collection, single_doc\n ):\n result = collection_with_single_doc.query(include_vector=True)\n assert len(result) == 1\n doc = result[0]\n assert doc.vector(\"dense\") is not None\n assert doc.vector(\"sparse\") is not None\n\n def test_collection_query_with_output_fields(\n self, collection_with_single_doc: Collection, single_doc\n ):\n result = collection_with_single_doc.query(output_fields=[\"id\", \"name\"])\n assert len(result) == 1\n doc = result[0]\n assert doc.id == single_doc.id\n assert len(doc.field_names()) == 2\n assert set(doc.field_names()) == {\"id\", \"name\"}\n\n def test_collection_query_with_topk(\n self, collection_with_multiple_docs: Collection\n ):\n result = collection_with_multiple_docs.query()\n assert len(result) == 10\n\n result = collection_with_multiple_docs.query(topk=5)\n assert len(result) == 5\n\n def test_collection_query_with_range_filter_int_field(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n index = 10\n idx = multiple_docs[index].id\n\n result = collection_with_multiple_docs.query(filter=f\"id>{idx}\", topk=100)\n assert len(result) == len(multiple_docs) - index - 1\n\n result = collection_with_multiple_docs.query(filter=f\"id>={idx}\", topk=100)\n assert len(result) == len(multiple_docs) - index\n\n result = collection_with_multiple_docs.query(filter=f\"id<{idx}\", topk=100)\n assert len(result) == index\n\n result = collection_with_multiple_docs.query(filter=f\"id<={idx}\", topk=100)\n assert len(result) == index + 1\n\n result = collection_with_multiple_docs.query(filter=f\"id={idx}\", topk=100)\n assert len(result) == 1\n\n result = collection_with_multiple_docs.query(filter=f\"id!={idx}\", topk=100)\n assert len(result) == len(multiple_docs) - 1\n\n left, right = 10, 90\n l_id, r_id = multiple_docs[left].id, multiple_docs[right].id\n result = collection_with_multiple_docs.query(\n filter=f\"id>{l_id} and id<{r_id}\", topk=100\n )\n assert len(result) == right - left - 1\n\n result = collection_with_multiple_docs.query(\n filter=f\"id>={l_id} and id<{r_id}\", topk=100\n )\n assert len(result) == right - left\n\n result = collection_with_multiple_docs.query(\n filter=f\"id>={l_id} and id<={r_id}\", topk=100\n )\n assert len(result) == right - left + 1\n\n result = collection_with_multiple_docs.query(\n filter=f\"id<{l_id} or id>{r_id}\", topk=100\n )\n assert len(result) == len(multiple_docs) - (right - left) - 1\n\n result = collection_with_multiple_docs.query(\n filter=f\"id<={l_id} or id>{r_id}\", topk=100\n )\n assert len(result) == len(multiple_docs) - (right - left)\n\n result = collection_with_multiple_docs.query(\n filter=f\"id<={l_id} or id>={r_id}\", topk=100\n )\n assert len(result) == len(multiple_docs) - (right - left) + 1\n\n result = collection_with_multiple_docs.query(filter=\"id in (1)\", topk=100)\n assert len(result) == 1\n\n def test_collection_query_with_vector_and_id(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n with pytest.raises(ValueError):\n collection_with_single_doc.query(\n VectorQuery(\n field_name=\"dense\",\n id=single_doc.id,\n vector=single_doc.vector(\"dense\"),\n )\n )\n\n def test_collection_query_with_filter_not_in(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n result = collection_with_multiple_docs.query(filter=\"id not in (1)\", topk=100)\n assert len(result) == len(multiple_docs) - 1\n\n def test_collection_with_error_query_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n query = VectorQuery(\n field_name=\"dense\", vector=multiple_docs[0].vector(\"dense\"), param=[1, 2, 3]\n )\n with pytest.raises(TypeError):\n result = collection_with_multiple_docs.query(\n filter=\"id in (1)\", topk=100, vectors=query\n )\n\n def test_collection_query_by_id(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n result = collection_with_multiple_docs.query(\n VectorQuery(field_name=\"dense\", id=multiple_docs[0].id)\n )\n assert len(result) == 10\n\n def test_collection_query_multi_vector_with_same_field(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n with pytest.raises(ValueError):\n collection_with_multiple_docs.query(\n [\n VectorQuery(\n field_name=\"dense\", vector=multiple_docs[0].vector(\"dense\")\n ),\n VectorQuery(\n field_name=\"dense\", vector=multiple_docs[0].vector(\"dense\")\n ),\n ]\n )\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_by_dense_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_by_sparse_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_by_dense_vector_with_filter(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_by_sparse_vector_with_filter(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_with_rrf_reranker_by_multi_dense_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_with_rrf_reranker_by_multi_sparse_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_with_rrf_reranker_by_hybrid_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_with_weighted_reranker_by_multi_dense_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_with_weighted_reranker_by_multi_sparse_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n\n @pytest.mark.skip(reason=\"TODO: This test case is pending implementation\")\n def test_collection_query_with_weighted_reranker_by_hybrid_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs\n ):\n pass\n" }, { "path": "python/tests/test_collection_hnsw_rabitq.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport platform\nimport sys\n\nimport pytest\nimport math\nimport zvec\n\npytestmark = pytest.mark.skipif(\n not (sys.platform == \"linux\" and platform.machine() in (\"x86_64\", \"AMD64\")),\n reason=\"HNSW RaBitQ only supported on Linux x86_64\",\n)\nfrom zvec import (\n Collection,\n CollectionOption,\n DataType,\n Doc,\n FieldSchema,\n HnswRabitqIndexParam,\n HnswRabitqQueryParam,\n MetricType,\n VectorSchema,\n VectorQuery,\n)\n\n\n# ==================== Fixtures ====================\n\n\n@pytest.fixture(scope=\"session\")\ndef hnsw_rabitq_collection_schema():\n \"\"\"Create a collection schema with HNSW RaBitQ index.\"\"\"\n return zvec.CollectionSchema(\n name=\"test_hnsw_rabitq_collection\",\n fields=[\n FieldSchema(\"id\", DataType.INT64, nullable=False),\n FieldSchema(\"name\", DataType.STRING, nullable=False),\n ],\n vectors=[\n VectorSchema(\n \"embedding\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswRabitqIndexParam(\n metric_type=MetricType.L2,\n m=16,\n ef_construction=200,\n total_bits=7,\n num_clusters=64,\n ),\n ),\n ],\n )\n\n\n@pytest.fixture(scope=\"session\")\ndef collection_option():\n \"\"\"Create collection options.\"\"\"\n return CollectionOption(read_only=False, enable_mmap=True)\n\n\n@pytest.fixture\ndef single_doc():\n \"\"\"Create a single document for testing.\"\"\"\n return Doc(\n id=\"0\",\n fields={\"id\": 0, \"name\": \"test_doc_0\"},\n vectors={\"embedding\": [0.1 + i * 0.01 for i in range(128)]},\n )\n\n\n@pytest.fixture\ndef multiple_docs():\n \"\"\"Create multiple documents for testing.\"\"\"\n return [\n Doc(\n id=f\"{i}\",\n fields={\"id\": i, \"name\": f\"test_doc_{i}\"},\n vectors={\"embedding\": [i * 0.1 + j * 0.01 for j in range(128)]},\n )\n for i in range(1, 101)\n ]\n\n\n@pytest.fixture(scope=\"function\")\ndef hnsw_rabitq_collection(\n tmp_path_factory, hnsw_rabitq_collection_schema, collection_option\n) -> Collection:\n \"\"\"\n Function-scoped fixture: creates and opens a collection with HNSW RaBitQ index.\n \"\"\"\n temp_dir = tmp_path_factory.mktemp(\"zvec_hnsw_rabitq\")\n collection_path = temp_dir / \"test_hnsw_rabitq_collection\"\n\n coll = zvec.create_and_open(\n path=str(collection_path),\n schema=hnsw_rabitq_collection_schema,\n option=collection_option,\n )\n\n assert coll is not None, \"Failed to create and open HNSW RaBitQ collection\"\n assert coll.path == str(collection_path)\n assert coll.schema.name == hnsw_rabitq_collection_schema.name\n\n try:\n yield coll\n finally:\n if hasattr(coll, \"destroy\") and coll is not None:\n try:\n coll.destroy()\n except Exception as e:\n print(f\"Warning: failed to destroy collection: {e}\")\n\n\n@pytest.fixture\ndef collection_with_single_doc(\n hnsw_rabitq_collection: Collection, single_doc: Doc\n) -> Collection:\n \"\"\"Setup: insert single doc into collection.\"\"\"\n assert hnsw_rabitq_collection.stats.doc_count == 0\n result = hnsw_rabitq_collection.insert(single_doc)\n assert bool(result)\n assert result.ok()\n assert hnsw_rabitq_collection.stats.doc_count == 1\n\n yield hnsw_rabitq_collection\n\n # Teardown: delete single doc\n hnsw_rabitq_collection.delete(single_doc.id)\n assert hnsw_rabitq_collection.stats.doc_count == 0\n\n\n@pytest.fixture\ndef collection_with_multiple_docs(\n hnsw_rabitq_collection: Collection, multiple_docs: list[Doc]\n) -> Collection:\n \"\"\"Setup: insert multiple docs into collection.\"\"\"\n assert hnsw_rabitq_collection.stats.doc_count == 0\n result = hnsw_rabitq_collection.insert(multiple_docs)\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok()\n assert hnsw_rabitq_collection.stats.doc_count == len(multiple_docs)\n\n yield hnsw_rabitq_collection\n\n # Teardown: delete multiple docs\n hnsw_rabitq_collection.delete([doc.id for doc in multiple_docs])\n\n\n# ==================== Tests ====================\n\n\n@pytest.mark.usefixtures(\"hnsw_rabitq_collection\")\nclass TestHnswRabitqCollectionCreation:\n \"\"\"Test HNSW RaBitQ collection creation and schema validation.\"\"\"\n\n def test_collection_creation(\n self, hnsw_rabitq_collection: Collection, hnsw_rabitq_collection_schema\n ):\n \"\"\"Test that collection is created with correct schema.\"\"\"\n assert hnsw_rabitq_collection is not None\n assert hnsw_rabitq_collection.schema.name == hnsw_rabitq_collection_schema.name\n assert len(hnsw_rabitq_collection.schema.fields) == len(\n hnsw_rabitq_collection_schema.fields\n )\n assert len(hnsw_rabitq_collection.schema.vectors) == len(\n hnsw_rabitq_collection_schema.vectors\n )\n\n def test_vector_schema_validation(self, hnsw_rabitq_collection: Collection):\n \"\"\"Test that vector schema has correct HNSW RaBitQ configuration.\"\"\"\n vector_schema = hnsw_rabitq_collection.schema.vector(\"embedding\")\n assert vector_schema is not None\n assert vector_schema.name == \"embedding\"\n assert vector_schema.data_type == DataType.VECTOR_FP32\n assert vector_schema.dimension == 128\n\n index_param = vector_schema.index_param\n assert index_param is not None\n assert index_param.metric_type == MetricType.L2\n assert index_param.m == 16\n assert index_param.ef_construction == 200\n assert index_param.total_bits == 7\n assert index_param.num_clusters == 64\n\n def test_collection_stats(self, hnsw_rabitq_collection: Collection):\n \"\"\"Test initial collection statistics.\"\"\"\n stats = hnsw_rabitq_collection.stats\n assert stats is not None\n assert stats.doc_count == 0\n assert len(stats.index_completeness) == 1\n assert stats.index_completeness[\"embedding\"] == 1\n\n\n@pytest.mark.usefixtures(\"hnsw_rabitq_collection\")\nclass TestHnswRabitqCollectionInsert:\n \"\"\"Test document insertion into HNSW RaBitQ collection.\"\"\"\n\n def test_insert_single_doc(\n self, hnsw_rabitq_collection: Collection, single_doc: Doc\n ):\n \"\"\"Test inserting a single document.\"\"\"\n result = hnsw_rabitq_collection.insert(single_doc)\n assert bool(result)\n assert result.ok()\n\n stats = hnsw_rabitq_collection.stats\n assert stats is not None\n assert stats.doc_count == 1\n\n def test_insert_multiple_docs(\n self, hnsw_rabitq_collection: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test inserting multiple documents.\"\"\"\n result = hnsw_rabitq_collection.insert(multiple_docs)\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok()\n\n stats = hnsw_rabitq_collection.stats\n assert stats is not None\n assert stats.doc_count == len(multiple_docs)\n\n\n@pytest.mark.usefixtures(\"hnsw_rabitq_collection\")\nclass TestHnswRabitqCollectionFetch:\n \"\"\"Test document fetching from HNSW RaBitQ collection.\"\"\"\n\n def test_fetch_single_doc(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n \"\"\"Test fetching a single document by ID.\"\"\"\n result = collection_with_single_doc.fetch(ids=[single_doc.id])\n assert bool(result)\n assert single_doc.id in result.keys()\n\n doc = result[single_doc.id]\n assert doc is not None\n assert doc.id == single_doc.id\n assert doc.field(\"id\") == single_doc.field(\"id\")\n assert doc.field(\"name\") == single_doc.field(\"name\")\n\n def test_fetch_multiple_docs(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test fetching multiple documents by IDs.\"\"\"\n ids = [doc.id for doc in multiple_docs[:10]]\n result = collection_with_multiple_docs.fetch(ids=ids)\n assert bool(result)\n assert len(result) == len(ids)\n\n for doc_id in ids:\n assert doc_id in result\n doc = result[doc_id]\n assert doc is not None\n assert doc.id == doc_id\n\n def test_fetch_nonexistent_doc(self, collection_with_single_doc: Collection):\n \"\"\"Test fetching a non-existent document.\"\"\"\n result = collection_with_single_doc.fetch(ids=[\"nonexistent_id\"])\n assert len(result) == 0\n\n\n@pytest.mark.usefixtures(\"hnsw_rabitq_collection\")\nclass TestHnswRabitqCollectionQuery:\n \"\"\"Test vector search queries on HNSW RaBitQ collection.\"\"\"\n\n def test_query_by_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test querying by vector with HNSW RaBitQ index.\"\"\"\n query_vector = multiple_docs[0].vector(\"embedding\")\n query = VectorQuery(\n field_name=\"embedding\",\n vector=query_vector,\n param=HnswRabitqQueryParam(ef=300),\n )\n\n result = collection_with_multiple_docs.query(vectors=query, topk=10)\n assert len(result) > 0\n assert len(result) <= 10\n\n # First result should be the query document itself (or very close)\n first_doc = result[0]\n assert first_doc is not None\n assert first_doc.id is not None\n\n def test_query_by_id(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test querying by document ID with HNSW RaBitQ index.\"\"\"\n query = VectorQuery(\n field_name=\"embedding\",\n id=multiple_docs[0].id,\n param=HnswRabitqQueryParam(ef=300),\n )\n\n result = collection_with_multiple_docs.query(vectors=query, topk=10)\n assert len(result) > 0\n assert len(result) <= 10\n\n def test_query_with_different_ef_values(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test querying with different ef parameter values.\"\"\"\n query_vector = multiple_docs[0].vector(\"embedding\")\n\n # Test with ef=100\n query_100 = VectorQuery(\n field_name=\"embedding\",\n vector=query_vector,\n param=HnswRabitqQueryParam(ef=100),\n )\n result_100 = collection_with_multiple_docs.query(vectors=query_100, topk=10)\n assert len(result_100) > 0\n\n # Test with ef=500\n query_500 = VectorQuery(\n field_name=\"embedding\",\n vector=query_vector,\n param=HnswRabitqQueryParam(ef=500),\n )\n result_500 = collection_with_multiple_docs.query(vectors=query_500, topk=10)\n assert len(result_500) > 0\n\n def test_query_with_topk(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test querying with different topk values.\"\"\"\n query_vector = multiple_docs[0].vector(\"embedding\")\n query = VectorQuery(\n field_name=\"embedding\",\n vector=query_vector,\n param=HnswRabitqQueryParam(ef=300),\n )\n\n # Test topk=5\n result_5 = collection_with_multiple_docs.query(vectors=query, topk=5)\n assert len(result_5) <= 5\n\n # Test topk=20\n result_20 = collection_with_multiple_docs.query(vectors=query, topk=20)\n assert len(result_20) <= 20\n\n def test_query_with_filter(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test querying with filter conditions.\"\"\"\n query_vector = multiple_docs[0].vector(\"embedding\")\n query = VectorQuery(\n field_name=\"embedding\",\n vector=query_vector,\n param=HnswRabitqQueryParam(ef=300),\n )\n\n # Query with id filter\n result = collection_with_multiple_docs.query(\n vectors=query, topk=10, filter=\"id < 50\"\n )\n assert len(result) > 0\n for doc in result:\n assert doc.field(\"id\") < 50\n\n def test_query_with_output_fields(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test querying with specific output fields.\"\"\"\n query_vector = multiple_docs[0].vector(\"embedding\")\n query = VectorQuery(\n field_name=\"embedding\",\n vector=query_vector,\n param=HnswRabitqQueryParam(ef=300),\n )\n\n result = collection_with_multiple_docs.query(\n vectors=query, topk=10, output_fields=[\"id\", \"name\"]\n )\n assert len(result) > 0\n\n first_doc = result[0]\n assert \"id\" in first_doc.field_names()\n assert \"name\" in first_doc.field_names()\n\n def test_query_with_include_vector(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test querying with vector data included in results.\"\"\"\n query_vector = multiple_docs[0].vector(\"embedding\")\n query = VectorQuery(\n field_name=\"embedding\",\n vector=query_vector,\n param=HnswRabitqQueryParam(ef=300),\n )\n\n result = collection_with_multiple_docs.query(\n vectors=query, topk=10, include_vector=True\n )\n assert len(result) > 0\n\n first_doc = result[0]\n assert first_doc.vector(\"embedding\") is not None\n assert len(first_doc.vector(\"embedding\")) == 128\n\n\n@pytest.mark.usefixtures(\"hnsw_rabitq_collection\")\nclass TestHnswRabitqCollectionUpdate:\n \"\"\"Test document update in HNSW RaBitQ collection.\"\"\"\n\n def test_update_doc_fields(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n \"\"\"Test updating document fields.\"\"\"\n updated_doc = Doc(\n id=single_doc.id,\n fields={\"id\": single_doc.field(\"id\"), \"name\": \"updated_name\"},\n )\n\n result = collection_with_single_doc.update(updated_doc)\n assert bool(result)\n assert result.ok()\n\n # Verify update\n fetched = collection_with_single_doc.fetch(ids=[single_doc.id])\n assert single_doc.id in fetched\n doc = fetched[single_doc.id]\n assert doc.field(\"name\") == \"updated_name\"\n\n def test_update_doc_vector(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n \"\"\"Test updating document vector.\"\"\"\n new_vector = [0.5 + i * 0.01 for i in range(128)]\n updated_doc = Doc(\n id=single_doc.id,\n vectors={\"embedding\": new_vector},\n )\n\n result = collection_with_single_doc.update(updated_doc)\n assert bool(result)\n assert result.ok()\n\n # Verify update\n fetched = collection_with_single_doc.fetch(\n ids=[single_doc.id],\n )\n assert single_doc.id in fetched\n doc = fetched[single_doc.id]\n assert doc.vector(\"embedding\") is not None\n embedding = doc.vector(\"embedding\")\n assert len(embedding) == 128\n # Verify vector values are approximately equal (float comparison)\n for i in range(128):\n assert math.isclose(embedding[i], new_vector[i], rel_tol=1e-5)\n\n\n@pytest.mark.usefixtures(\"hnsw_rabitq_collection\")\nclass TestHnswRabitqCollectionDelete:\n \"\"\"Test document deletion from HNSW RaBitQ collection.\"\"\"\n\n def test_delete_single_doc(\n self, collection_with_single_doc: Collection, single_doc: Doc\n ):\n \"\"\"Test deleting a single document.\"\"\"\n result = collection_with_single_doc.delete(single_doc.id)\n assert bool(result)\n assert result.ok()\n\n stats = collection_with_single_doc.stats\n assert stats.doc_count == 0\n\n def test_delete_multiple_docs(\n self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]\n ):\n \"\"\"Test deleting multiple documents.\"\"\"\n ids_to_delete = [doc.id for doc in multiple_docs[:10]]\n result = collection_with_multiple_docs.delete(ids_to_delete)\n assert len(result) == len(ids_to_delete)\n for item in result:\n assert item.ok()\n\n stats = collection_with_multiple_docs.stats\n assert stats.doc_count == len(multiple_docs) - len(ids_to_delete)\n\n\n@pytest.mark.usefixtures(\"hnsw_rabitq_collection\")\nclass TestHnswRabitqCollectionOptimizeAndReopen:\n \"\"\"Test collection optimize and reopen functionality.\"\"\"\n\n def test_optimize_close_reopen_and_query(\n self,\n tmp_path_factory,\n hnsw_rabitq_collection_schema,\n collection_option,\n multiple_docs: list[Doc],\n ):\n \"\"\"Test inserting 100 docs, optimize, close, reopen and query.\"\"\"\n # Create collection and insert 100 documents\n temp_dir = tmp_path_factory.mktemp(\"zvec_hnsw_rabitq_optimize\")\n collection_path = temp_dir / \"test_optimize_collection\"\n\n coll = zvec.create_and_open(\n path=str(collection_path),\n schema=hnsw_rabitq_collection_schema,\n option=collection_option,\n )\n\n assert coll is not None\n assert coll.stats.doc_count == 0\n\n # Insert 100 documents\n result = coll.insert(multiple_docs)\n assert len(result) == len(multiple_docs)\n for item in result:\n assert item.ok()\n assert coll.stats.doc_count == len(multiple_docs)\n\n # Call optimize\n from zvec import OptimizeOption\n\n coll.optimize(option=OptimizeOption())\n\n # Verify data is still accessible after optimize\n query_vector = multiple_docs[0].vector(\"embedding\")\n query = VectorQuery(\n field_name=\"embedding\",\n vector=query_vector,\n param=HnswRabitqQueryParam(ef=300),\n )\n result_before_close = coll.query(vectors=query, topk=10)\n assert len(result_before_close) > 0\n\n # Close collection (destroy will close it)\n collection_path_str = str(collection_path)\n del coll\n\n # Reopen collection\n reopened_coll = zvec.open(path=collection_path_str, option=collection_option)\n assert reopened_coll is not None\n assert reopened_coll.stats.doc_count == len(multiple_docs)\n\n # Execute query on reopened collection\n query_after_reopen = VectorQuery(\n field_name=\"embedding\",\n vector=query_vector,\n param=HnswRabitqQueryParam(ef=300),\n )\n result_after_reopen = reopened_coll.query(vectors=query_after_reopen, topk=10)\n assert len(result_after_reopen) > 0\n assert len(result_after_reopen) <= 10\n\n # Verify query results are valid\n first_doc = result_after_reopen[0]\n assert first_doc is not None\n assert first_doc.id is not None\n assert first_doc.field(\"id\") is not None\n assert first_doc.field(\"name\") is not None\n\n # Cleanup\n reopened_coll.destroy()\n" }, { "path": "python/tests/test_convert.py", "content": "from __future__ import annotations\n\nimport math\n\nimport pytest\nfrom _zvec import _Doc\nfrom zvec.model.convert import convert_to_py_doc, convert_to_cpp_doc\nfrom zvec import Doc, CollectionSchema, DataType, FieldSchema, VectorSchema\n\n\n# ----------------------------\n# Convert Cpp Doc Test Case\n# ----------------------------\nclass TestConvertCppDoc:\n def test_default(self):\n doc = Doc(id=\"1\")\n schema = CollectionSchema(\n name=\"test_collection\",\n fields=FieldSchema(\"name\", DataType.STRING),\n )\n\n cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)\n\n assert cpp_doc is not None\n assert cpp_doc.pk() == doc.id\n\n def test_with_field_notin_schema(self):\n doc = Doc(id=\"1\", fields={\"name\": \"Tom\"})\n schema = CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\"id\", DataType.UINT64),\n FieldSchema(\"salary\", DataType.UINT32),\n FieldSchema(\"age\", DataType.INT32),\n FieldSchema(\"create_at\", DataType.INT64),\n FieldSchema(\"author\", DataType.STRING),\n FieldSchema(\"weight\", DataType.FLOAT),\n ],\n )\n with pytest.raises(ValueError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n def test_with_scalar_fields(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\"id\", DataType.UINT64),\n FieldSchema(\"salary\", DataType.UINT32),\n FieldSchema(\"age\", DataType.INT32),\n FieldSchema(\"create_at\", DataType.INT64),\n FieldSchema(\"author\", DataType.STRING),\n FieldSchema(\"weight\", DataType.FLOAT),\n FieldSchema(\"bmi\", DataType.DOUBLE),\n FieldSchema(\"is_male\", DataType.BOOL),\n ],\n )\n doc = Doc(\n id=\"1\",\n fields={\n \"id\": 1,\n \"salary\": 1000,\n \"age\": 18,\n \"create_at\": 1640995200,\n \"bmi\": 80.0 / 200.0,\n \"author\": \"Tom\",\n \"weight\": 80.0,\n \"is_male\": True,\n },\n )\n cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)\n assert cpp_doc is not None\n assert cpp_doc.pk() == doc.id\n assert cpp_doc.get_any(\"id\", DataType.UINT64) == 1\n assert cpp_doc.get_any(\"salary\", DataType.UINT32) == 1000\n assert cpp_doc.get_any(\"age\", DataType.INT32) == 18\n assert cpp_doc.get_any(\"create_at\", DataType.INT64) == 1640995200\n assert cpp_doc.get_any(\"author\", DataType.STRING) == \"Tom\"\n assert math.isclose(\n cpp_doc.get_any(\"weight\", DataType.FLOAT), 80.0, rel_tol=1e-6\n )\n assert math.isclose(\n cpp_doc.get_any(\"bmi\", DataType.DOUBLE), 80.0 / 200.0, rel_tol=1e-6\n )\n assert cpp_doc.get_any(\"is_male\", DataType.BOOL) == True\n\n def test_with_array_fields(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\"tags\", DataType.ARRAY_STRING),\n FieldSchema(\"ids\", DataType.ARRAY_UINT64),\n FieldSchema(\"marks\", DataType.ARRAY_UINT32),\n FieldSchema(\"x\", DataType.ARRAY_INT32),\n FieldSchema(\"y\", DataType.ARRAY_INT64),\n FieldSchema(\"scores\", DataType.ARRAY_FLOAT),\n FieldSchema(\"ratios\", DataType.ARRAY_DOUBLE),\n FieldSchema(\"results\", DataType.ARRAY_BOOL),\n ],\n )\n\n doc = Doc(\n id=\"1\",\n fields={\n \"tags\": [\"tag1\", \"tag2\", \"tag3\"],\n \"ids\": [111111111111, 222222222222, 333333333333],\n \"marks\": [100, 200, 300],\n \"x\": [1, 2, 3],\n \"y\": [100, 200, 300],\n \"scores\": [1.1, 2.2, 3.3],\n \"ratios\": [0.1, 0.2, 0.3],\n \"results\": [True, False, True],\n },\n )\n cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)\n\n assert cpp_doc is not None\n assert cpp_doc.pk() == doc.id\n assert cpp_doc.get_any(\"tags\", DataType.ARRAY_STRING) == doc.field(\"tags\")\n assert cpp_doc.get_any(\"ids\", DataType.ARRAY_UINT64) == doc.field(\"ids\")\n assert cpp_doc.get_any(\"marks\", DataType.ARRAY_UINT32) == doc.field(\"marks\")\n assert cpp_doc.get_any(\"x\", DataType.ARRAY_INT32) == doc.field(\"x\")\n assert cpp_doc.get_any(\"y\", DataType.ARRAY_INT64) == doc.field(\"y\")\n scores = cpp_doc.get_any(\"scores\", DataType.ARRAY_FLOAT)\n for i in range(len(doc.field(\"scores\"))):\n assert math.isclose(scores[i], doc.field(\"scores\")[i], rel_tol=1e-1)\n ratios = cpp_doc.get_any(\"ratios\", DataType.ARRAY_DOUBLE)\n for i in range(len(doc.field(\"ratios\"))):\n assert math.isclose(ratios[i], doc.field(\"ratios\")[i], rel_tol=1e-1)\n results = cpp_doc.get_any(\"results\", DataType.ARRAY_BOOL)\n for i in range(len(doc.field(\"results\"))):\n assert results[i] == doc.field(\"results\")[i]\n\n def test_with_dense_vector_fields(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n vectors=[\n VectorSchema(\n name=\"embedding\",\n data_type=DataType.VECTOR_FP16,\n dimension=4,\n ),\n VectorSchema(\n name=\"image\",\n data_type=DataType.VECTOR_FP32,\n dimension=8,\n ),\n VectorSchema(\n name=\"text\",\n data_type=DataType.VECTOR_INT8,\n dimension=32,\n ),\n ],\n )\n\n doc = Doc(\n id=\"1\",\n vectors={\n \"embedding\": [1.1] * 4,\n \"image\": [2.2] * 8,\n \"text\": [4] * 32,\n },\n )\n cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)\n assert cpp_doc is not None\n assert cpp_doc.pk() == doc.id\n\n embedding_vector = cpp_doc.get_any(\"embedding\", DataType.VECTOR_FP16)\n assert len(embedding_vector) == 4\n for i in range(4):\n assert math.isclose(\n embedding_vector[i], doc.vector(\"embedding\")[i], rel_tol=1e-1\n )\n\n image_vector = cpp_doc.get_any(\"image\", DataType.VECTOR_FP32)\n assert len(image_vector) == 8\n for i in range(8):\n assert math.isclose(image_vector[i], doc.vector(\"image\")[i], rel_tol=1e-1)\n\n text_vector = cpp_doc.get_any(\"text\", DataType.VECTOR_INT8)\n assert len(text_vector) == 32\n for i in range(32):\n assert text_vector[i] == doc.vectors[\"text\"][i]\n\n def test_with_sparse_vector_fields(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n vectors=[\n VectorSchema(\n name=\"author\",\n data_type=DataType.SPARSE_VECTOR_FP32,\n ),\n VectorSchema(\n name=\"content\",\n data_type=DataType.SPARSE_VECTOR_FP16,\n ),\n ],\n )\n doc = Doc(\n id=\"1\",\n vectors={\n \"author\": {1: 1.1, 2: 2.2, 3: 3.3},\n \"content\": {4: 4.4, 5: 5.5, 6: 6.6},\n },\n )\n\n cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)\n assert cpp_doc is not None\n assert cpp_doc.pk() == doc.id\n\n author_vector = cpp_doc.get_any(\"author\", DataType.SPARSE_VECTOR_FP32)\n assert isinstance(author_vector, dict)\n for key, value in doc.vector(\"author\").items():\n assert math.isclose(author_vector[key], value, rel_tol=1e-1)\n\n content_vector = cpp_doc.get_any(\"content\", DataType.SPARSE_VECTOR_FP16)\n assert isinstance(content_vector, dict)\n for key, value in doc.vector(\"content\").items():\n assert math.isclose(content_vector[key], value, rel_tol=1e-1)\n\n def test_with_scalar_fields_error_datatype(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\"id\", DataType.UINT64),\n FieldSchema(\"salary\", DataType.UINT32),\n FieldSchema(\"age\", DataType.INT32),\n FieldSchema(\"create_at\", DataType.INT64),\n FieldSchema(\"author\", DataType.STRING),\n FieldSchema(\"weight\", DataType.FLOAT),\n FieldSchema(\"bmi\", DataType.DOUBLE),\n FieldSchema(\"is_male\", DataType.BOOL),\n ],\n )\n doc = Doc(\n id=\"1\",\n fields={\n \"id\": \"1\",\n },\n )\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"salary\": \"1000\"})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"age\": \"18\"})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"create_at\": \"2021-01-01\"})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"author\": 1})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"weight\": \"80.5\"})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"bmi\": \"25.0\"})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"is_male\": \"true\"})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n def test_with_array_fields_error_datatype(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\"tags\", DataType.ARRAY_STRING),\n FieldSchema(\"ids\", DataType.ARRAY_UINT64),\n FieldSchema(\"marks\", DataType.ARRAY_UINT32),\n FieldSchema(\"x\", DataType.ARRAY_INT32),\n FieldSchema(\"y\", DataType.ARRAY_INT64),\n FieldSchema(\"scores\", DataType.ARRAY_FLOAT),\n FieldSchema(\"ratios\", DataType.ARRAY_DOUBLE),\n FieldSchema(\"results\", DataType.ARRAY_BOOL),\n ],\n )\n\n doc = Doc(id=\"1\", fields={\"tags\": [1, 2, 3]})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"ids\": [\"1\", \"2\", \"3\"]})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"marks\": [1.1, 2.2, 3.3]})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"x\": [1.1, 2.2, 3.3]})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"y\": [1.1, 2.2, 3.3]})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"scores\": [\"1\", \"2\", \"3\"]})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"ratios\": [\"1\", \"2\", \"3\"]})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", fields={\"results\": [\"1\", \"2\", \"3\"]})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n def test_with_vector_fields_error_datatype(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n vectors=[\n VectorSchema(\n name=\"embedding\",\n data_type=DataType.VECTOR_FP16,\n dimension=4,\n ),\n VectorSchema(\n name=\"image\",\n data_type=DataType.VECTOR_FP32,\n dimension=8,\n ),\n VectorSchema(\n name=\"text\",\n data_type=DataType.VECTOR_INT8,\n dimension=32,\n ),\n ],\n )\n\n doc = Doc(id=\"1\", vectors={\"image\": [\"1.1\"] * 4})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", vectors={\"text\": [\"1\"] * 4})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(id=\"1\", vectors={\"embedding\": [\"1\"] * 4})\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n def test_with_sparse_vector_error_datatype(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n vectors=[\n VectorSchema(\n name=\"author\",\n data_type=DataType.SPARSE_VECTOR_FP32,\n ),\n VectorSchema(\n name=\"content\",\n data_type=DataType.SPARSE_VECTOR_FP16,\n ),\n ],\n )\n doc = Doc(\n id=\"1\",\n vectors={\n \"author\": {\"1\": 1.1, \"2\": 2.2, \"3\": 3.3},\n },\n )\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(\n id=\"1\",\n vectors={\n \"content\": {\"1\": 1.1, \"2\": 2.2, \"3\": 3.3},\n },\n )\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n doc = Doc(\n id=\"1\",\n vectors={\n \"author\": {1: \"1\", 2: \"2\", 3: \"3\"},\n },\n )\n with pytest.raises(TypeError):\n convert_to_cpp_doc(doc, collection_schema=schema)\n\n\n# ----------------------------\n# Convert Py Doc Test Case\n# ----------------------------\nclass TestConvertPyDoc:\n def test_default(self):\n doc = _Doc()\n doc.set_pk(\"1\")\n doc.set_score(1.0)\n\n schema = CollectionSchema(\n name=\"test_collection\",\n fields=FieldSchema(\"name\", DataType.STRING),\n )\n\n py_doc = convert_to_py_doc(doc, schema)\n assert py_doc.id == \"1\"\n assert py_doc.score == 1.0\n\n def test_with_scalar_fields(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\"id\", DataType.UINT64),\n FieldSchema(\"salary\", DataType.UINT32),\n FieldSchema(\"age\", DataType.INT32),\n FieldSchema(\"create_at\", DataType.INT64),\n FieldSchema(\"author\", DataType.STRING),\n FieldSchema(\"weight\", DataType.FLOAT),\n FieldSchema(\"bmi\", DataType.DOUBLE),\n FieldSchema(\"is_male\", DataType.BOOL),\n ],\n )\n doc = _Doc()\n doc.set_pk(\"1\")\n doc.set_any(\"id\", schema.field(\"id\")._get_object(), 1111111111111111)\n doc.set_any(\"salary\", schema.field(\"salary\")._get_object(), 1000)\n doc.set_any(\"age\", schema.field(\"age\")._get_object(), 18)\n doc.set_any(\"create_at\", schema.field(\"create_at\")._get_object(), 1640995200)\n doc.set_any(\"author\", schema.field(\"author\")._get_object(), \"Tom\")\n doc.set_any(\"weight\", schema.field(\"weight\")._get_object(), 80.0)\n doc.set_any(\"bmi\", schema.field(\"bmi\")._get_object(), 80.0 / 200.0)\n doc.set_any(\"is_male\", schema.field(\"is_male\")._get_object(), True)\n\n py_doc = convert_to_py_doc(doc, schema)\n assert py_doc.id == \"1\"\n assert py_doc.field(\"id\") == 1111111111111111\n assert py_doc.field(\"salary\") == 1000\n assert py_doc.field(\"age\") == 18\n assert py_doc.field(\"create_at\") == 1640995200\n assert py_doc.field(\"author\") == \"Tom\"\n assert py_doc.field(\"weight\") == 80.0\n assert py_doc.field(\"bmi\") == 80.0 / 200.0\n assert py_doc.field(\"is_male\") == True\n\n def test_with_array_fields(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\"tags\", DataType.ARRAY_STRING),\n FieldSchema(\"ids\", DataType.ARRAY_UINT64),\n FieldSchema(\"marks\", DataType.ARRAY_UINT32),\n FieldSchema(\"x\", DataType.ARRAY_INT32),\n FieldSchema(\"y\", DataType.ARRAY_INT64),\n FieldSchema(\"scores\", DataType.ARRAY_FLOAT),\n FieldSchema(\"ratios\", DataType.ARRAY_DOUBLE),\n FieldSchema(\"results\", DataType.ARRAY_BOOL),\n ],\n )\n\n doc = _Doc()\n doc.set_pk(\"1\")\n doc.set_any(\n \"tags\", schema.field(\"tags\")._get_object(), [\"tag1\", \"tag2\", \"tag3\"]\n )\n doc.set_any(\n \"ids\",\n schema.field(\"ids\")._get_object(),\n [111111111111, 222222222222, 3333333333333],\n )\n doc.set_any(\"marks\", schema.field(\"marks\")._get_object(), [1000, 2000, 3000])\n doc.set_any(\"x\", schema.field(\"x\")._get_object(), [1, 2, 3])\n doc.set_any(\"y\", schema.field(\"y\")._get_object(), [100, 200, 300])\n doc.set_any(\"scores\", schema.field(\"scores\")._get_object(), [0.1, 0.2, 0.3])\n doc.set_any(\"ratios\", schema.field(\"ratios\")._get_object(), [0.1, 0.2, 0.3])\n doc.set_any(\n \"results\", schema.field(\"results\")._get_object(), [True, False, True]\n )\n\n py_doc = convert_to_py_doc(doc, schema)\n assert py_doc.field(\"tags\") == [\"tag1\", \"tag2\", \"tag3\"]\n assert py_doc.field(\"ids\") == [111111111111, 222222222222, 3333333333333]\n assert py_doc.field(\"marks\") == [1000, 2000, 3000]\n assert py_doc.field(\"x\") == [1, 2, 3]\n assert py_doc.field(\"y\") == [100, 200, 300]\n\n scores = doc.get_any(\"scores\", DataType.ARRAY_FLOAT)\n for i in range(len(scores)):\n assert math.isclose(scores[i], py_doc.field(\"scores\")[i], rel_tol=1e-1)\n ratios = doc.get_any(\"ratios\", DataType.ARRAY_DOUBLE)\n for i in range(len(ratios)):\n assert math.isclose(ratios[i], py_doc.field(\"ratios\")[i], rel_tol=1e-1)\n results = doc.get_any(\"results\", DataType.ARRAY_BOOL)\n for i in range(len(results)):\n assert results[i] == py_doc.field(\"results\")[i]\n\n def test_with_dense_vector_fields(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n vectors=[\n VectorSchema(\n name=\"embedding\",\n data_type=DataType.VECTOR_FP16,\n dimension=4,\n ),\n VectorSchema(\n name=\"image\",\n data_type=DataType.VECTOR_FP32,\n dimension=8,\n ),\n VectorSchema(\n name=\"text\",\n data_type=DataType.VECTOR_INT8,\n dimension=32,\n ),\n ],\n )\n\n doc = _Doc()\n doc.set_pk(\"1\")\n doc.set_any(\"embedding\", schema.vector(\"embedding\")._get_object(), [1.1] * 4)\n doc.set_any(\"image\", schema.vector(\"image\")._get_object(), [2.2] * 8)\n doc.set_any(\"text\", schema.vector(\"text\")._get_object(), [4] * 32)\n\n py_doc = convert_to_py_doc(doc, schema)\n assert py_doc.id == \"1\"\n\n embedding_vector = py_doc.vector(\"embedding\")\n assert len(embedding_vector) == 4\n for i in range(4):\n assert math.isclose(\n py_doc.vector(\"embedding\")[i], embedding_vector[i], rel_tol=1e-1\n )\n\n image_vector = py_doc.vector(\"image\")\n assert len(image_vector) == 8\n for i in range(8):\n assert math.isclose(\n py_doc.vector(\"image\")[i], image_vector[i], rel_tol=1e-1\n )\n\n text_vector = py_doc.vector(\"text\")\n assert len(text_vector) == 32\n for i in range(32):\n assert py_doc.vector(\"text\")[i] == text_vector[i]\n\n def test_with_sparse_vector_fields(self):\n schema = CollectionSchema(\n name=\"test_collection\",\n vectors=[\n VectorSchema(\n name=\"author\",\n data_type=DataType.SPARSE_VECTOR_FP32,\n ),\n VectorSchema(\n name=\"content\",\n data_type=DataType.SPARSE_VECTOR_FP16,\n ),\n ],\n )\n\n doc = _Doc()\n doc.set_pk(\"1\")\n doc.set_any(\n \"author\", schema.vector(\"author\")._get_object(), {1: 1.1, 2: 2.2, 3: 3.3}\n )\n doc.set_any(\n \"content\", schema.vector(\"content\")._get_object(), {4: 4.4, 5: 5.5, 6: 6.6}\n )\n\n py_doc = convert_to_py_doc(doc, schema)\n assert py_doc.id == \"1\"\n\n author_vector = py_doc.vector(\"author\")\n assert isinstance(author_vector, dict)\n for key, value in doc.get_any(\"author\", DataType.SPARSE_VECTOR_FP32).items():\n assert math.isclose(author_vector[key], value, rel_tol=1e-1)\n\n content_vector = py_doc.vector(\"content\")\n assert isinstance(content_vector, dict)\n for key, value in doc.get_any(\"content\", DataType.SPARSE_VECTOR_FP16).items():\n assert math.isclose(content_vector[key], value, rel_tol=1e-1)\n" }, { "path": "python/tests/test_doc.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport math\nimport pytest\n\n\nfrom _zvec import _Doc\nfrom zvec import FieldSchema, VectorSchema, Doc, DataType\n\n\n# ----------------------------\n# PyDoc Test Case\n# ----------------------------\nclass TestPyDoc:\n def test_default(self):\n Doc(id=\"1\")\n\n def test_with_single_vector(self):\n doc = Doc(id=\"1\", vectors={\"dense\": [1, 2, 3]})\n assert doc is not None\n assert doc.id == \"1\"\n assert doc.vector(\"dense\") == [1, 2, 3]\n\n def test_with_hybrid_vectors(self):\n doc = Doc(\n id=\"1\", vectors={\"dense\": [1, 2, 3], \"sparse\": {1: 1.0, 2: 2.0, 3: 3.0}}\n )\n assert doc is not None\n assert doc.id == \"1\"\n assert doc.vector(\"dense\") == [1, 2, 3]\n assert doc.vector(\"sparse\") == {1: 1.0, 2: 2.0, 3: 3.0}\n\n def test_with_multi_vectors(self):\n doc = Doc(\n id=\"1\",\n vectors={\n \"image\": [1, 2, 3],\n \"description\": [4, 5, 6],\n \"keys\": {1: 1.0, 2: 2.0, 3: 3.0},\n },\n fields={\"author\": \"Tom\", \"age\": 19, \"is_male\": True, \"weight\": 60.5},\n )\n assert doc is not None\n assert doc.id == \"1\"\n assert doc.vector(\"image\") == [1, 2, 3]\n assert doc.vector(\"description\") == [4, 5, 6]\n assert doc.vector(\"keys\") == {1: 1.0, 2: 2.0, 3: 3.0}\n assert doc.field(\"author\") == \"Tom\"\n assert doc.field(\"age\") == 19\n assert doc.field(\"is_male\") == True\n assert doc.field(\"weight\") == 60.5\n\n def test_with_numpy_array(self):\n import numpy as np\n\n doc = Doc._from_tuple(\n (\n \"1\",\n 0.0,\n None,\n {\n \"image\": np.array([1, 2, 3]),\n \"description\": np.random.random(512),\n \"keys\": {1: 1.0, 2: 2.0, 3: 3.0},\n },\n )\n )\n assert doc is not None\n assert doc.id == \"1\"\n assert doc.vector(\"image\") == [1, 2, 3]\n assert doc.vector(\"keys\") == {1: 1.0, 2: 2.0, 3: 3.0}\n\n\n# ----------------------------\n# CppDoc Test Case\n# ----------------------------\nclass TestCppDoc:\n def test_default(self):\n doc = _Doc()\n assert doc is not None\n\n def test_doc_set_pk(self):\n doc = _Doc()\n doc.set_pk(\"1\")\n assert doc.pk() == \"1\"\n\n def test_doc_set_score(self):\n doc = _Doc()\n doc.set_score(0.9)\n assert math.isclose(doc.score(), 0.9, rel_tol=1e-6)\n\n def test_doc_get_null_field(self):\n doc = _Doc()\n schema = FieldSchema(\"author\", DataType.STRING, nullable=True)\n doc.set_any(\"author\", schema._get_object(), None)\n assert doc.has_field(\"author\")\n assert doc.get_any(\"author\", schema.data_type) is None\n\n def test_doc_get_set_has_null_field(self):\n doc = _Doc()\n schema = FieldSchema(\"author\", DataType.STRING, nullable=False)\n with pytest.raises(ValueError):\n doc.set_any(\"author\", schema._get_object(), None)\n\n def test_doc_get_set_has_string_field(self):\n doc = _Doc()\n schema = FieldSchema(\"author\", DataType.STRING)\n doc.set_any(\"author\", schema._get_object(), \"Tom\")\n assert doc.has_field(\"author\")\n assert doc.get_any(\"author\", DataType.STRING) == \"Tom\"\n\n def test_doc_get_set_has_bool_field(self):\n doc = _Doc()\n schema = FieldSchema(\"is_male\", DataType.BOOL)\n doc.set_any(\"is_male\", schema._get_object(), True)\n assert doc.has_field(\"is_male\")\n assert doc.get_any(\"is_male\", DataType.BOOL) == True\n\n def test_doc_get_set_has_int32_field(self):\n doc = _Doc()\n schema = FieldSchema(\"age\", DataType.INT32)\n doc.set_any(\"age\", schema._get_object(), 19)\n assert doc.has_field(\"age\")\n assert doc.get_any(\"age\", DataType.INT32) == 19\n\n def test_doc_get_set_has_int64_field(self):\n doc = _Doc()\n schema = FieldSchema(\"id\", DataType.INT64)\n doc.set_any(\"id\", schema._get_object(), 1111111111111111111)\n assert doc.has_field(\"id\")\n assert doc.get_any(\"id\", DataType.INT64) == 1111111111111111111\n\n def test_doc_get_set_has_float_field(self):\n doc = _Doc()\n schema = FieldSchema(\"weight\", DataType.FLOAT)\n doc.set_any(\"weight\", schema._get_object(), 60.5)\n assert doc.has_field(\"weight\")\n assert math.isclose(doc.get_any(\"weight\", DataType.FLOAT), 60.5, rel_tol=1e-6)\n\n def test_doc_get_set_has_double_field(self):\n doc = _Doc()\n schema = FieldSchema(\"height\", DataType.DOUBLE)\n doc.set_any(\"height\", schema._get_object(), 1.77777777777)\n assert doc.has_field(\"height\")\n assert math.isclose(\n doc.get_any(\"height\", DataType.DOUBLE), 1.7777777777, rel_tol=1e-9\n )\n\n def test_doc_get_set_has_uint32_field(self):\n doc = _Doc()\n schema = FieldSchema(\"id\", DataType.UINT32)\n doc.set_any(\"id\", schema._get_object(), 4294967295)\n assert doc.has_field(\"id\")\n assert doc.get_any(\"id\", DataType.UINT32) == 4294967295\n\n def test_doc_get_set_has_uint64_field(self):\n doc = _Doc()\n schema = FieldSchema(\"id\", DataType.UINT64)\n doc.set_any(\"id\", schema._get_object(), 18446744073709551615)\n assert doc.has_field(\"id\")\n assert doc.get_any(\"id\", DataType.UINT64) == 18446744073709551615\n\n def test_doc_get_set_has_array_string_field(self):\n doc = _Doc()\n schema = FieldSchema(\"tags\", DataType.ARRAY_STRING)\n doc.set_any(\"tags\", schema._get_object(), [\"tag1\", \"tag2\", \"tag3\"])\n assert doc.has_field(\"tags\")\n assert doc.get_any(\"tags\", DataType.ARRAY_STRING) == [\"tag1\", \"tag2\", \"tag3\"]\n\n def test_doc_get_set_has_array_int32_field(self):\n doc = _Doc()\n schema = FieldSchema(\"ids\", DataType.ARRAY_INT32)\n doc.set_any(\"ids\", schema._get_object(), [1, 2, 3])\n assert doc.has_field(\"ids\")\n assert doc.get_any(\"ids\", DataType.ARRAY_INT32) == [1, 2, 3]\n\n def test_doc_get_set_has_array_int64_field(self):\n doc = _Doc()\n schema = FieldSchema(\"ids\", DataType.ARRAY_INT64)\n doc.set_any(\"ids\", schema._get_object(), [1, 2, 3])\n assert doc.has_field(\"ids\")\n assert doc.get_any(\"ids\", DataType.ARRAY_INT64) == [1, 2, 3]\n\n def test_doc_get_set_has_array_float_field(self):\n doc = _Doc()\n schema = FieldSchema(\"weights\", DataType.ARRAY_FLOAT)\n doc.set_any(\"weights\", schema._get_object(), [1.0, 2.0, 3.0])\n assert doc.has_field(\"weights\")\n assert doc.get_any(\"weights\", DataType.ARRAY_FLOAT) == [1.0, 2.0, 3.0]\n\n def test_doc_get_set_has_array_double_field(self):\n doc = _Doc()\n schema = FieldSchema(\"heights\", DataType.ARRAY_DOUBLE)\n doc.set_any(\"heights\", schema._get_object(), [1.0, 2.0, 3.0])\n assert doc.has_field(\"heights\")\n assert doc.get_any(\"heights\", DataType.ARRAY_DOUBLE) == [1.0, 2.0, 3.0]\n\n def test_doc_get_set_has_array_bool_field(self):\n doc = _Doc()\n schema = FieldSchema(\"bools\", DataType.ARRAY_BOOL)\n doc.set_any(\"bools\", schema._get_object(), [True, False, True])\n assert doc.has_field(\"bools\")\n assert doc.get_any(\"bools\", DataType.ARRAY_BOOL) == [True, False, True]\n\n def test_doc_get_set_has_vector_fp16(self):\n doc = _Doc()\n schema = VectorSchema(\"image\", DataType.VECTOR_FP16)\n doc.set_any(\"image\", schema._get_object(), [1.0, 2.0, 3.0])\n assert doc.has_field(\"image\")\n image_vector = doc.get_any(\"image\", DataType.VECTOR_FP16)\n assert image_vector is not None\n for i in range(len(image_vector)):\n assert math.isclose(image_vector[i], [1.0, 2.0, 3.0][i], rel_tol=1e-6)\n\n def test_doc_get_set_has_vector_fp32(self):\n doc = _Doc()\n schema = VectorSchema(\"image\", DataType.VECTOR_FP32)\n doc.set_any(\"image\", schema._get_object(), [1.111111, 2.222222, 3.333333])\n assert doc.has_field(\"image\")\n vector = doc.get_any(\"image\", DataType.VECTOR_FP32)\n assert vector is not None\n for i in range(len(vector)):\n assert math.isclose(\n vector[i], [1.111111, 2.222222, 3.333333][i], rel_tol=1e-6\n )\n\n def test_doc_get_set_has_vector_int8(self):\n doc = _Doc()\n schema = VectorSchema(\"image\", DataType.VECTOR_INT8)\n doc.set_any(\"image\", schema._get_object(), [1, 2, 3])\n assert doc.has_field(\"image\")\n assert doc.get_any(\"image\", DataType.VECTOR_INT8) == [1, 2, 3]\n\n def test_doc_get_set_has_sparse_vector_fp32(self):\n doc = _Doc()\n sparse = {1: 1.111111, 2: 2.222222, 3: 3.333333}\n schema = VectorSchema(\"key\", DataType.SPARSE_VECTOR_FP32)\n doc.set_any(\"key\", schema._get_object(), sparse)\n assert doc.has_field(\"key\")\n vector = doc.get_any(\"key\", DataType.SPARSE_VECTOR_FP32)\n assert vector is not None\n assert isinstance(vector, dict)\n for key, value in sparse.items():\n assert math.isclose(vector[key], value, rel_tol=1e-6)\n\n def test_doc_get_set_has_sparse_vector_fp16(self):\n doc = _Doc()\n sparse = {1: 1.1, 2: 2.2, 3: 3.3}\n schema = VectorSchema(\"key\", DataType.SPARSE_VECTOR_FP16)\n doc.set_any(\"key\", schema._get_object(), sparse)\n assert doc.has_field(\"key\")\n vector = doc.get_any(\"key\", DataType.SPARSE_VECTOR_FP16)\n assert vector is not None\n assert isinstance(vector, dict)\n for key, value in sparse.items():\n assert math.isclose(vector[key], value, rel_tol=1e-1)\n" }, { "path": "python/tests/test_embedding.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport os\nfrom http import HTTPStatus\nfrom unittest.mock import MagicMock, patch, Mock\n\nimport numpy as np\nimport pytest\nfrom zvec.extension import (\n BM25EmbeddingFunction,\n DefaultLocalDenseEmbedding,\n DefaultLocalSparseEmbedding,\n OpenAIDenseEmbedding,\n QwenDenseEmbedding,\n QwenSparseEmbedding,\n)\n\n# Environment variable to control integration tests\n# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API/model tests\nRUN_INTEGRATION_TESTS = os.environ.get(\"ZVEC_RUN_INTEGRATION_TESTS\", \"0\") == \"1\"\n\n\n# ----------------------------\n# QwenDenseEmbedding Test Case\n# ----------------------------\nclass TestQwenDenseEmbedding:\n def test_init_with_api_key(self):\n # Test initialization with explicit API key\n embedding_func = QwenDenseEmbedding(dimension=128, api_key=\"test_key\")\n assert embedding_func.dimension == 128\n assert embedding_func.model == \"text-embedding-v4\"\n assert embedding_func._api_key == \"test_key\"\n\n @patch.dict(os.environ, {\"DASHSCOPE_API_KEY\": \"env_key\"})\n def test_init_with_env_api_key(self):\n # Test initialization with API key from environment\n embedding_func = QwenDenseEmbedding(dimension=128)\n assert embedding_func._api_key == \"env_key\"\n\n @patch.dict(os.environ, {\"DASHSCOPE_API_KEY\": \"\"})\n def test_init_with_empty_env_api_key(self):\n # Test initialization with empty API key from environment\n with pytest.raises(ValueError, match=\"DashScope API key is required\"):\n QwenDenseEmbedding(dimension=128)\n\n def test_model_property(self):\n embedding_func = QwenDenseEmbedding(dimension=128, api_key=\"test_key\")\n assert embedding_func.model == \"text-embedding-v4\"\n\n embedding_func = QwenDenseEmbedding(\n dimension=128, model=\"custom-model\", api_key=\"test_key\"\n )\n assert embedding_func.model == \"custom-model\"\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_with_empty_text(self, mock_require_module):\n # Test embed method with empty text raises ValueError\n embedding_func = QwenDenseEmbedding(dimension=128, api_key=\"test_key\")\n\n with pytest.raises(\n ValueError, match=\"Input text cannot be empty or whitespace only\"\n ):\n embedding_func.embed(\"\")\n\n with pytest.raises(TypeError):\n embedding_func.embed(None)\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_success(self, mock_require_module):\n # Test successful embedding\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.OK\n mock_response.output = {\"embeddings\": [{\"embedding\": [0.1, 0.2, 0.3]}]}\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenDenseEmbedding(dimension=3, api_key=\"test_key\")\n # Clear cache to avoid interference\n embedding_func.embed.cache_clear()\n result = embedding_func.embed(\"test text\")\n\n assert result == [0.1, 0.2, 0.3]\n mock_dashscope.TextEmbedding.call.assert_called_once_with(\n model=\"text-embedding-v4\",\n input=\"test text\",\n dimension=3,\n output_type=\"dense\",\n )\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_http_error(self, mock_require_module):\n # Test embedding with HTTP error\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.BAD_REQUEST\n mock_response.message = \"Bad Request\"\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenDenseEmbedding(dimension=128, api_key=\"test_key\")\n embedding_func.embed.cache_clear()\n\n with pytest.raises(ValueError):\n embedding_func.embed(\"test text\")\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_invalid_response(self, mock_require_module):\n # Test embedding with invalid response (wrong number of embeddings)\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.OK\n mock_response.output = {\"embeddings\": []}\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenDenseEmbedding(dimension=128, api_key=\"test_key\")\n embedding_func.embed.cache_clear()\n\n with pytest.raises(ValueError):\n embedding_func.embed(\"test text\")\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n def test_real_embed_success(self):\n \"\"\"Integration test with real DashScope API.\n\n To run this test, set environment variable:\n export ZVEC_RUN_INTEGRATION_TESTS=1\n export DASHSCOPE_API_KEY=your-api-key\n \"\"\"\n embedding_func = QwenDenseEmbedding(dimension=128)\n dense = embedding_func(\"test text\")\n assert len(dense) == 128\n\n\n# ----------------------------\n# QwenSparseEmbedding Test Case\n# ----------------------------\nclass TestQwenSparseEmbedding:\n \"\"\"Test suite for QwenSparseEmbedding (Qwen sparse embedding via DashScope API).\"\"\"\n\n def test_init_with_api_key(self):\n \"\"\"Test initialization with explicit API key.\"\"\"\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n assert embedding_func._dimension == 1024\n assert embedding_func.model == \"text-embedding-v4\"\n assert embedding_func._api_key == \"test_key\"\n # encoding_type defaults to \"query\" via extra_params\n assert embedding_func.extra_params.get(\"encoding_type\", \"query\") == \"query\"\n\n def test_init_with_custom_encoding_type(self):\n \"\"\"Test initialization with custom encoding type.\"\"\"\n embedding_func = QwenSparseEmbedding(\n dimension=1024, encoding_type=\"document\", api_key=\"test_key\"\n )\n assert embedding_func.extra_params.get(\"encoding_type\") == \"document\"\n\n @patch.dict(os.environ, {\"DASHSCOPE_API_KEY\": \"env_key\"})\n def test_init_with_env_api_key(self):\n \"\"\"Test initialization with API key from environment.\"\"\"\n embedding_func = QwenSparseEmbedding(dimension=1024)\n assert embedding_func._api_key == \"env_key\"\n\n @patch.dict(os.environ, {\"DASHSCOPE_API_KEY\": \"\"})\n def test_init_without_api_key(self):\n \"\"\"Test initialization fails without API key.\"\"\"\n with pytest.raises(ValueError, match=\"DashScope API key is required\"):\n QwenSparseEmbedding(dimension=1024)\n\n def test_model_property(self):\n \"\"\"Test model property.\"\"\"\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n assert embedding_func.model == \"text-embedding-v4\"\n\n embedding_func = QwenSparseEmbedding(\n dimension=1024, model=\"text-embedding-v3\", api_key=\"test_key\"\n )\n assert embedding_func.model == \"text-embedding-v3\"\n\n def test_encoding_type_property(self):\n \"\"\"Test encoding_type via extra_params.\"\"\"\n query_emb = QwenSparseEmbedding(\n dimension=1024, encoding_type=\"query\", api_key=\"test_key\"\n )\n assert query_emb.extra_params.get(\"encoding_type\") == \"query\"\n\n doc_emb = QwenSparseEmbedding(\n dimension=1024, encoding_type=\"document\", api_key=\"test_key\"\n )\n assert doc_emb.extra_params.get(\"encoding_type\") == \"document\"\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_with_empty_text(self, mock_require_module):\n \"\"\"Test embed method with empty text raises ValueError.\"\"\"\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n\n with pytest.raises(\n ValueError, match=\"Input text cannot be empty or whitespace only\"\n ):\n embedding_func.embed(\"\")\n\n with pytest.raises(\n ValueError, match=\"Input text cannot be empty or whitespace only\"\n ):\n embedding_func.embed(\" \")\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_with_non_string_input(self, mock_require_module):\n \"\"\"Test embed method with non-string input raises TypeError.\"\"\"\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n embedding_func.embed(123)\n\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n embedding_func.embed(None)\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_success(self, mock_require_module):\n \"\"\"Test successful sparse embedding generation.\"\"\"\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.OK\n # Sparse embedding returns array of {index, value, token} objects\n mock_response.output = {\n \"embeddings\": [\n {\n \"sparse_embedding\": [\n {\"index\": 10, \"value\": 0.5, \"token\": \"机器\"},\n {\"index\": 245, \"value\": 0.8, \"token\": \"学习\"},\n {\"index\": 1023, \"value\": 1.2, \"token\": \"算法\"},\n ]\n }\n ]\n }\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n # Clear cache to avoid interference\n embedding_func.embed.cache_clear()\n result = embedding_func.embed(\"test text\")\n\n # Verify result is a dict\n assert isinstance(result, dict)\n # Verify keys are integers\n assert all(isinstance(k, int) for k in result.keys())\n # Verify values are floats\n assert all(isinstance(v, float) for v in result.values())\n # Verify all values are positive\n assert all(v > 0 for v in result.values())\n # Verify sorted by indices\n keys = list(result.keys())\n assert keys == sorted(keys)\n # Verify specific keys\n assert keys == [10, 245, 1023]\n\n mock_dashscope.TextEmbedding.call.assert_called_once_with(\n model=\"text-embedding-v4\",\n input=\"test text\",\n dimension=1024,\n output_type=\"sparse\",\n text_type=\"query\",\n )\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_with_document_encoding_type(self, mock_require_module):\n \"\"\"Test embedding with document encoding type.\"\"\"\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.OK\n mock_response.output = {\n \"embeddings\": [\n {\n \"sparse_embedding\": [\n {\"index\": 5, \"value\": 0.3, \"token\": \"文档\"},\n {\"index\": 100, \"value\": 0.7, \"token\": \"内容\"},\n {\"index\": 500, \"value\": 0.9, \"token\": \"检索\"},\n ]\n }\n ]\n }\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenSparseEmbedding(\n dimension=1024, encoding_type=\"document\", api_key=\"test_key\"\n )\n embedding_func.embed.cache_clear()\n result = embedding_func.embed(\"test document\")\n\n assert isinstance(result, dict)\n assert list(result.keys()) == [5, 100, 500]\n\n # Verify text_type parameter is \"document\"\n call_args = mock_dashscope.TextEmbedding.call.call_args\n assert call_args[1][\"text_type\"] == \"document\"\n assert call_args[1][\"output_type\"] == \"sparse\"\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_output_sorted_by_indices(self, mock_require_module):\n \"\"\"Test that output is always sorted by indices in ascending order.\"\"\"\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.OK\n # Return unsorted indices\n mock_response.output = {\n \"embeddings\": [\n {\n \"sparse_embedding\": [\n {\"index\": 9999, \"value\": 1.5, \"token\": \"A\"},\n {\"index\": 5, \"value\": 2.0, \"token\": \"B\"},\n {\"index\": 1234, \"value\": 0.8, \"token\": \"C\"},\n {\"index\": 77, \"value\": 3.2, \"token\": \"D\"},\n {\"index\": 500, \"value\": 1.1, \"token\": \"E\"},\n ]\n }\n ]\n }\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n embedding_func.embed.cache_clear()\n result = embedding_func.embed(\"test sorting\")\n\n # Verify keys are sorted\n result_keys = list(result.keys())\n assert result_keys == sorted(result_keys)\n # Verify expected sorted order\n assert result_keys == [5, 77, 500, 1234, 9999]\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_filters_zero_values(self, mock_require_module):\n \"\"\"Test that zero and negative values are filtered out.\"\"\"\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.OK\n # Include zero and negative values\n mock_response.output = {\n \"embeddings\": [\n {\n \"sparse_embedding\": [\n {\"index\": 10, \"value\": 0.5, \"token\": \"正\"},\n {\n \"index\": 20,\n \"value\": 0.0,\n \"token\": \"零\",\n }, # Should be filtered\n {\n \"index\": 30,\n \"value\": -0.3,\n \"token\": \"负\",\n }, # Should be filtered\n {\"index\": 40, \"value\": 0.8, \"token\": \"正\"},\n {\n \"index\": 50,\n \"value\": 0.0,\n \"token\": \"零\",\n }, # Should be filtered\n ]\n }\n ]\n }\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n embedding_func.embed.cache_clear()\n result = embedding_func.embed(\"test filtering\")\n\n # Only positive values should remain\n assert list(result.keys()) == [10, 40]\n assert all(v > 0 for v in result.values())\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_http_error(self, mock_require_module):\n \"\"\"Test embedding with HTTP error.\"\"\"\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.BAD_REQUEST\n mock_response.message = \"Bad Request\"\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n embedding_func.embed.cache_clear()\n\n with pytest.raises(ValueError, match=\"DashScope API error\"):\n embedding_func.embed(\"test text\")\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_invalid_response_no_embeddings(self, mock_require_module):\n \"\"\"Test embedding with invalid response (no embeddings).\"\"\"\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.OK\n mock_response.output = {\"embeddings\": []}\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n embedding_func.embed.cache_clear()\n\n with pytest.raises(ValueError, match=\"Expected exactly 1 embedding\"):\n embedding_func.embed(\"test text\")\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_invalid_response_not_dict(self, mock_require_module):\n \"\"\"Test embedding with invalid response (sparse_embedding not list).\"\"\"\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.OK\n # sparse_embedding should be list, not dict\n mock_response.output = {\n \"embeddings\": [{\"sparse_embedding\": {\"index\": 10, \"value\": 0.5}}]\n }\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n embedding_func.embed.cache_clear()\n\n with pytest.raises(\n ValueError, match=\"'sparse_embedding' field is missing or not a list\"\n ):\n embedding_func.embed(\"test text\")\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_callable_interface(self, mock_require_module):\n \"\"\"Test that embedding function is callable.\"\"\"\n mock_dashscope = MagicMock()\n mock_response = MagicMock()\n mock_response.status_code = HTTPStatus.OK\n mock_response.output = {\n \"embeddings\": [\n {\n \"sparse_embedding\": [\n {\"index\": 100, \"value\": 1.0, \"token\": \"测试\"},\n {\"index\": 200, \"value\": 0.5, \"token\": \"调用\"},\n ]\n }\n ]\n }\n mock_dashscope.TextEmbedding.call.return_value = mock_response\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n embedding_func.embed.cache_clear()\n\n # Test calling the function directly\n result = embedding_func(\"test text\")\n assert isinstance(result, dict)\n assert list(result.keys()) == [100, 200]\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_embed_api_connection_error(self, mock_require_module):\n \"\"\"Test handling of API connection errors.\"\"\"\n mock_dashscope = MagicMock()\n mock_dashscope.TextEmbedding.call.side_effect = Exception(\"Connection timeout\")\n mock_require_module.return_value = mock_dashscope\n\n embedding_func = QwenSparseEmbedding(dimension=1024, api_key=\"test_key\")\n embedding_func.embed.cache_clear()\n\n with pytest.raises(RuntimeError, match=\"Failed to call DashScope API\"):\n embedding_func.embed(\"test text\")\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n def test_real_embed_success(self):\n \"\"\"Integration test with real DashScope API.\n\n To run this test, set environment variable:\n export ZVEC_RUN_INTEGRATION_TESTS=1\n export DASHSCOPE_API_KEY=your-api-key\n \"\"\"\n # Test query embedding\n query_emb = QwenSparseEmbedding(dimension=1024, encoding_type=\"query\")\n query_vec = query_emb.embed(\"machine learning\")\n\n assert isinstance(query_vec, dict)\n assert len(query_vec) > 0\n assert all(isinstance(k, int) for k in query_vec.keys())\n assert all(isinstance(v, float) and v > 0 for v in query_vec.values())\n\n # Verify sorted output\n keys = list(query_vec.keys())\n assert keys == sorted(keys)\n\n # Test document embedding\n doc_emb = QwenSparseEmbedding(dimension=1024, encoding_type=\"document\")\n doc_vec = doc_emb.embed(\"Machine learning is a subset of AI\")\n\n assert isinstance(doc_vec, dict)\n assert len(doc_vec) > 0\n\n # Verify sorted output\n doc_keys = list(doc_vec.keys())\n assert doc_keys == sorted(doc_keys)\n\n\n# ----------------------------\n# OpenAIDenseEmbedding Test Case\n# ----------------------------\nclass TestOpenAIDenseEmbedding:\n def test_init_with_api_key(self):\n \"\"\"Test initialization with explicit API key.\"\"\"\n embedding_func = OpenAIDenseEmbedding(api_key=\"sk-test-key\")\n assert embedding_func.dimension == 1536 # Default for text-embedding-3-small\n assert embedding_func.model == \"text-embedding-3-small\"\n assert embedding_func._api_key == \"sk-test-key\"\n\n @patch.dict(os.environ, {\"OPENAI_API_KEY\": \"sk-env-key\"})\n def test_init_with_env_api_key(self):\n \"\"\"Test initialization with API key from environment.\"\"\"\n embedding_func = OpenAIDenseEmbedding()\n assert embedding_func._api_key == \"sk-env-key\"\n\n @patch.dict(os.environ, {\"OPENAI_API_KEY\": \"\"})\n def test_init_without_api_key(self):\n \"\"\"Test initialization fails without API key.\"\"\"\n with pytest.raises(ValueError, match=\"OpenAI API key is required\"):\n OpenAIDenseEmbedding()\n\n def test_init_with_custom_dimension(self):\n \"\"\"Test initialization with custom dimension.\"\"\"\n embedding_func = OpenAIDenseEmbedding(\n model=\"text-embedding-3-large\", dimension=1024, api_key=\"sk-test\"\n )\n assert embedding_func.dimension == 1024\n assert embedding_func.model == \"text-embedding-3-large\"\n\n def test_init_with_base_url(self):\n \"\"\"Test initialization with custom base URL.\"\"\"\n embedding_func = OpenAIDenseEmbedding(\n api_key=\"sk-test\", base_url=\"https://custom.openai.com/\"\n )\n assert embedding_func._base_url == \"https://custom.openai.com/\"\n\n def test_model_property(self):\n \"\"\"Test model property.\"\"\"\n embedding_func = OpenAIDenseEmbedding(api_key=\"sk-test\")\n assert embedding_func.model == \"text-embedding-3-small\"\n\n embedding_func = OpenAIDenseEmbedding(\n model=\"text-embedding-ada-002\", api_key=\"sk-test\"\n )\n assert embedding_func.model == \"text-embedding-ada-002\"\n\n def test_extra_params(self):\n \"\"\"Test extra_params property.\"\"\"\n # Test without extra params\n embedding_func = OpenAIDenseEmbedding(api_key=\"sk-test\")\n assert embedding_func.extra_params == {}\n\n # Test with extra params\n embedding_func = OpenAIDenseEmbedding(\n api_key=\"sk-test\",\n encoding_format=\"float\",\n user=\"test-user\",\n )\n assert embedding_func.extra_params == {\n \"encoding_format\": \"float\",\n \"user\": \"test-user\",\n }\n\n @patch(\"zvec.extension.openai_function.require_module\")\n def test_embed_with_empty_text(self, mock_require_module):\n \"\"\"Test embed method with empty text raises ValueError.\"\"\"\n embedding_func = OpenAIDenseEmbedding(api_key=\"sk-test\")\n\n with pytest.raises(\n ValueError, match=\"Input text cannot be empty or whitespace only\"\n ):\n embedding_func.embed(\"\")\n\n with pytest.raises(\n ValueError, match=\"Input text cannot be empty or whitespace only\"\n ):\n embedding_func.embed(\" \")\n\n @patch(\"zvec.extension.openai_function.require_module\")\n def test_embed_with_non_string_input(self, mock_require_module):\n \"\"\"Test embed method with non-string input raises TypeError.\"\"\"\n embedding_func = OpenAIDenseEmbedding(api_key=\"sk-test\")\n\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n embedding_func.embed(123)\n\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n embedding_func.embed(None)\n\n @patch(\"zvec.extension.openai_function.require_module\")\n def test_embed_success(self, mock_require_module):\n \"\"\"Test successful embedding generation.\"\"\"\n # Mock OpenAI client\n mock_openai = Mock()\n mock_client = Mock()\n mock_response = Mock()\n\n # Create mock embedding data\n fake_embedding = [0.1, 0.2, 0.3]\n mock_embedding_obj = Mock()\n mock_embedding_obj.embedding = fake_embedding\n mock_response.data = [mock_embedding_obj]\n\n mock_client.embeddings.create.return_value = mock_response\n mock_openai.OpenAI.return_value = mock_client\n mock_require_module.return_value = mock_openai\n\n embedding_func = OpenAIDenseEmbedding(dimension=3, api_key=\"sk-test\")\n embedding_func.embed.cache_clear()\n result = embedding_func.embed(\"test text\")\n\n assert result == [0.1, 0.2, 0.3]\n mock_client.embeddings.create.assert_called_once_with(\n model=\"text-embedding-3-small\", input=\"test text\", dimensions=3\n )\n\n @patch(\"zvec.extension.openai_function.require_module\")\n def test_embed_with_custom_model(self, mock_require_module):\n \"\"\"Test embedding with custom model.\"\"\"\n mock_openai = Mock()\n mock_client = Mock()\n mock_response = Mock()\n\n fake_embedding = [0.1] * 1536\n mock_embedding_obj = Mock()\n mock_embedding_obj.embedding = fake_embedding\n mock_response.data = [mock_embedding_obj]\n\n mock_client.embeddings.create.return_value = mock_response\n mock_openai.OpenAI.return_value = mock_client\n mock_require_module.return_value = mock_openai\n\n embedding_func = OpenAIDenseEmbedding(\n model=\"text-embedding-ada-002\", api_key=\"sk-test\"\n )\n embedding_func.embed.cache_clear()\n result = embedding_func.embed(\"test text\")\n\n assert len(result) == 1536\n mock_client.embeddings.create.assert_called_once_with(\n model=\"text-embedding-ada-002\", input=\"test text\"\n )\n\n @patch(\"zvec.extension.openai_function.require_module\")\n def test_embed_api_error(self, mock_require_module):\n \"\"\"Test handling of API errors.\"\"\"\n mock_openai = Mock()\n mock_client = Mock()\n\n # Simulate API error\n api_error = Mock()\n api_error.__class__.__name__ = \"APIError\"\n mock_openai.APIError = type(\"APIError\", (Exception,), {})\n mock_openai.APIConnectionError = type(\"APIConnectionError\", (Exception,), {})\n\n mock_client.embeddings.create.side_effect = mock_openai.APIError(\n \"Rate limit exceeded\"\n )\n mock_openai.OpenAI.return_value = mock_client\n mock_require_module.return_value = mock_openai\n\n embedding_func = OpenAIDenseEmbedding(api_key=\"sk-test\")\n embedding_func.embed.cache_clear()\n\n with pytest.raises(RuntimeError, match=\"Failed to call OpenAI API\"):\n embedding_func.embed(\"test text\")\n\n @patch(\"zvec.extension.openai_function.require_module\")\n def test_embed_invalid_response(self, mock_require_module):\n \"\"\"Test handling of invalid API response.\"\"\"\n mock_openai = Mock()\n mock_client = Mock()\n mock_response = Mock()\n\n # Empty response data\n mock_response.data = []\n\n mock_client.embeddings.create.return_value = mock_response\n mock_openai.OpenAI.return_value = mock_client\n mock_openai.APIError = type(\"APIError\", (Exception,), {})\n mock_openai.APIConnectionError = type(\"APIConnectionError\", (Exception,), {})\n mock_require_module.return_value = mock_openai\n\n embedding_func = OpenAIDenseEmbedding(api_key=\"sk-test\")\n embedding_func.embed.cache_clear()\n\n with pytest.raises(ValueError, match=\"no embedding data returned\"):\n embedding_func.embed(\"test text\")\n\n @patch(\"zvec.extension.openai_function.require_module\")\n def test_embed_dimension_mismatch(self, mock_require_module):\n \"\"\"Test handling of dimension mismatch.\"\"\"\n mock_openai = Mock()\n mock_client = Mock()\n mock_response = Mock()\n\n # Return embedding with wrong dimension\n fake_embedding = [0.1] * 512\n mock_embedding_obj = Mock()\n mock_embedding_obj.embedding = fake_embedding\n mock_response.data = [mock_embedding_obj]\n\n mock_client.embeddings.create.return_value = mock_response\n mock_openai.OpenAI.return_value = mock_client\n mock_openai.APIError = type(\"APIError\", (Exception,), {})\n mock_openai.APIConnectionError = type(\"APIConnectionError\", (Exception,), {})\n mock_require_module.return_value = mock_openai\n\n embedding_func = OpenAIDenseEmbedding(dimension=1536, api_key=\"sk-test\")\n embedding_func.embed.cache_clear()\n\n with pytest.raises(ValueError, match=\"Dimension mismatch\"):\n embedding_func.embed(\"test text\")\n\n @patch(\"zvec.extension.openai_function.require_module\")\n def test_embed_callable(self, mock_require_module):\n \"\"\"Test that embedding function is callable.\"\"\"\n mock_openai = Mock()\n mock_client = Mock()\n mock_response = Mock()\n\n fake_embedding = [0.1] * 1536\n mock_embedding_obj = Mock()\n mock_embedding_obj.embedding = fake_embedding\n mock_response.data = [mock_embedding_obj]\n\n mock_client.embeddings.create.return_value = mock_response\n mock_openai.OpenAI.return_value = mock_client\n mock_openai.APIError = type(\"APIError\", (Exception,), {})\n mock_openai.APIConnectionError = type(\"APIConnectionError\", (Exception,), {})\n mock_require_module.return_value = mock_openai\n\n embedding_func = OpenAIDenseEmbedding(api_key=\"sk-test\")\n embedding_func.embed.cache_clear()\n\n # Test calling the function directly\n result = embedding_func(\"test text\")\n assert isinstance(result, list)\n assert len(result) == 1536\n\n @patch(\"zvec.extension.openai_function.require_module\")\n def test_embed_with_base_url(self, mock_require_module):\n \"\"\"Test embedding with custom base URL.\"\"\"\n mock_openai = Mock()\n mock_client = Mock()\n mock_response = Mock()\n\n fake_embedding = [0.1] * 1536\n mock_embedding_obj = Mock()\n mock_embedding_obj.embedding = fake_embedding\n mock_response.data = [mock_embedding_obj]\n\n mock_client.embeddings.create.return_value = mock_response\n mock_openai.OpenAI.return_value = mock_client\n mock_openai.APIError = type(\"APIError\", (Exception,), {})\n mock_openai.APIConnectionError = type(\"APIConnectionError\", (Exception,), {})\n mock_require_module.return_value = mock_openai\n\n embedding_func = OpenAIDenseEmbedding(\n api_key=\"sk-test\", base_url=\"https://custom.openai.com/\"\n )\n embedding_func.embed.cache_clear()\n result = embedding_func.embed(\"test text\")\n\n # Verify client was created with custom base URL\n mock_openai.OpenAI.assert_called_once_with(\n api_key=\"sk-test\", base_url=\"https://custom.openai.com/\"\n )\n assert len(result) == 1536\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n def test_real_embed_success(self):\n \"\"\"Integration test with real OpenAI API.\n\n To run this test, set environment variable:\n export ZVEC_RUN_INTEGRATION_TESTS=1\n export OPENAI_API_KEY=sk-...\n \"\"\"\n embedding_func = OpenAIDenseEmbedding(\n model=\"text-embedding-v4\",\n dimension=256,\n base_url=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n )\n vector = embedding_func.embed(\"Hello, world!\")\n assert len(vector) == 256\n assert isinstance(vector, list)\n assert all(isinstance(x, float) for x in vector)\n\n\n# ----------------------------\n# DefaultLocalDenseEmbedding Test Case\n# ----------------------------\nclass TestDefaultLocalDenseEmbedding:\n \"\"\"Test cases for DefaultLocalDenseEmbedding.\"\"\"\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_init_success(self, mock_require_module):\n \"\"\"Test successful initialization with mocked model.\"\"\"\n # Mock sentence_transformers module\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n mock_model.device = \"cpu\"\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n # Initialize embedding function\n emb_func = DefaultLocalDenseEmbedding()\n\n # Assertions\n assert emb_func.dimension == 384\n assert emb_func.model_name == \"all-MiniLM-L6-v2\"\n assert emb_func.model_source == \"huggingface\"\n assert emb_func.device == \"cpu\"\n mock_st.SentenceTransformer.assert_called_once_with(\n \"all-MiniLM-L6-v2\", device=None, trust_remote_code=True\n )\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_init_with_custom_device(self, mock_require_module):\n \"\"\"Test initialization with custom device.\"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n mock_model.device = \"cuda\"\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n emb_func = DefaultLocalDenseEmbedding(device=\"cuda\")\n\n assert emb_func.device == \"cuda\"\n mock_st.SentenceTransformer.assert_called_once_with(\n \"all-MiniLM-L6-v2\", device=\"cuda\", trust_remote_code=True\n )\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_init_with_modelscope(self, mock_require_module):\n \"\"\"Test initialization with ModelScope as model source.\"\"\"\n mock_st = Mock()\n mock_ms = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n mock_model.device = \"cpu\"\n mock_st.SentenceTransformer.return_value = mock_model\n\n def require_module_side_effect(module_name):\n if module_name == \"sentence_transformers\":\n return mock_st\n elif module_name == \"modelscope\":\n return mock_ms\n raise ImportError(f\"No module named '{module_name}'\")\n\n mock_require_module.side_effect = require_module_side_effect\n\n # Mock snapshot_download at the correct import location\n with patch(\n \"modelscope.hub.snapshot_download.snapshot_download\",\n return_value=\"/path/to/cached/model\",\n ):\n emb_func = DefaultLocalDenseEmbedding(model_source=\"modelscope\")\n\n # Assertions\n assert emb_func.dimension == 384\n assert emb_func.model_name == \"iic/nlp_gte_sentence-embedding_chinese-small\"\n assert emb_func.model_source == \"modelscope\"\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_init_with_invalid_model_source(self, mock_require_module):\n \"\"\"Test initialization with invalid model_source raises ValueError.\"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n with pytest.raises(ValueError, match=\"Invalid model_source\"):\n DefaultLocalDenseEmbedding(model_source=\"invalid_source\")\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_embed_success(self, mock_require_module):\n \"\"\"Test successful embedding generation.\"\"\"\n # Mock embedding output\n fake_embedding = np.random.rand(384).astype(np.float32)\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n\n # Configure encode method\n mock_model.encode = Mock(return_value=fake_embedding)\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n emb_func = DefaultLocalDenseEmbedding()\n result = emb_func.embed(\"Hello, world!\")\n\n # Assertions\n assert isinstance(result, list)\n assert len(result) == 384\n assert all(isinstance(x, float) for x in result)\n mock_model.encode.assert_called_once_with(\n \"Hello, world!\",\n convert_to_numpy=True,\n normalize_embeddings=True,\n batch_size=32,\n )\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_embed_with_normalization(self, mock_require_module):\n \"\"\"Test embedding with L2 normalization.\"\"\"\n # Create a normalized vector\n fake_embedding = np.random.rand(384).astype(np.float32)\n fake_embedding = fake_embedding / np.linalg.norm(fake_embedding)\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n\n # Configure encode method\n mock_model.encode = Mock(return_value=fake_embedding)\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n emb_func = DefaultLocalDenseEmbedding(normalize_embeddings=True)\n result = emb_func.embed(\"Test sentence\")\n\n # Check if vector is normalized (L2 norm should be close to 1.0)\n result_array = np.array(result)\n norm = np.linalg.norm(result_array)\n assert abs(norm - 1.0) < 1e-5\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_embed_empty_string(self, mock_require_module):\n \"\"\"Test embedding with empty string raises ValueError.\"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n emb_func = DefaultLocalDenseEmbedding()\n\n with pytest.raises(ValueError, match=\"Input text cannot be empty\"):\n emb_func.embed(\"\")\n\n with pytest.raises(ValueError, match=\"Input text cannot be empty\"):\n emb_func.embed(\" \")\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_embed_non_string_input(self, mock_require_module):\n \"\"\"Test embedding with non-string input raises TypeError.\"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n emb_func = DefaultLocalDenseEmbedding()\n\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n emb_func.embed(123)\n\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n emb_func.embed(None)\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_embed_callable(self, mock_require_module):\n \"\"\"Test that embedding function is callable.\"\"\"\n fake_embedding = np.random.rand(384).astype(np.float32)\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n\n # Configure encode method\n mock_model.encode = Mock(return_value=fake_embedding)\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n emb_func = DefaultLocalDenseEmbedding()\n\n # Test calling the function directly\n result = emb_func(\"Test text\")\n assert isinstance(result, list)\n assert len(result) == 384\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_semantic_similarity(self, mock_require_module):\n \"\"\"Test semantic similarity between similar and different texts.\"\"\"\n # Create mock embeddings for similar and different texts\n similar_emb_1 = np.array([1.0, 0.0, 0.0] + [0.0] * 381, dtype=np.float32)\n similar_emb_2 = np.array([0.9, 0.1, 0.0] + [0.0] * 381, dtype=np.float32)\n different_emb = np.array([0.0, 0.0, 1.0] + [0.0] * 381, dtype=np.float32)\n\n # Normalize\n similar_emb_1 = similar_emb_1 / np.linalg.norm(similar_emb_1)\n similar_emb_2 = similar_emb_2 / np.linalg.norm(similar_emb_2)\n different_emb = different_emb / np.linalg.norm(different_emb)\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n\n # Configure encode method with side_effect for multiple calls\n mock_model.encode = Mock(\n side_effect=[similar_emb_1, similar_emb_2, different_emb]\n )\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n emb_func = DefaultLocalDenseEmbedding()\n\n v1 = emb_func.embed(\"The cat sits on the mat\")\n v2 = emb_func.embed(\"A feline rests on a rug\")\n v3 = emb_func.embed(\"Python programming\")\n\n # Calculate similarities\n similarity_high = np.dot(v1, v2)\n similarity_low = np.dot(v1, v3)\n\n assert similarity_high > similarity_low\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_model_loading_error(self, mock_require_module):\n \"\"\"Test handling of model loading failure.\"\"\"\n # Clear model cache\n from zvec.extension.sentence_transformer_embedding_function import (\n DefaultLocalSparseEmbedding,\n )\n\n DefaultLocalSparseEmbedding.clear_cache()\n mock_st = Mock()\n mock_st.SentenceTransformer.side_effect = Exception(\"Model not found\")\n mock_require_module.return_value = mock_st\n\n with pytest.raises(\n ValueError, match=\"Failed to load Sentence Transformer model\"\n ):\n DefaultLocalDenseEmbedding()\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_modelscope_import_error(self, mock_require_module):\n \"\"\"Test handling of ModelScope import error.\"\"\"\n mock_st = Mock()\n\n def require_module_side_effect(module_name):\n if module_name == \"sentence_transformers\":\n return mock_st\n elif module_name == \"modelscope\":\n raise ImportError(\"No module named 'modelscope'\")\n\n mock_require_module.side_effect = require_module_side_effect\n\n with pytest.raises(\n ImportError, match=\"ModelScope support requires the 'modelscope' package\"\n ):\n DefaultLocalDenseEmbedding(model_source=\"modelscope\")\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_embed_dimension_mismatch(self, mock_require_module):\n \"\"\"Test handling of dimension mismatch in embedding output.\"\"\"\n # Return embedding with wrong dimension\n fake_embedding = np.random.rand(256).astype(np.float32)\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n\n # Configure encode method\n mock_model.encode = Mock(return_value=fake_embedding)\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n emb_func = DefaultLocalDenseEmbedding()\n\n with pytest.raises(ValueError, match=\"Dimension mismatch\"):\n emb_func.embed(\"Test text\")\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n def test_real_embedding_generation(self):\n \"\"\"Integration test with real model (requires sentence-transformers).\n\n To run this test, set environment variable:\n export ZVEC_RUN_INTEGRATION_TESTS=1\n\n Note: First run will download the model (~80MB).\n \"\"\"\n emb_func = DefaultLocalDenseEmbedding()\n\n # Test basic embedding\n vector = emb_func.embed(\"Hello, world!\")\n assert len(vector) == 384\n assert isinstance(vector, list)\n assert all(isinstance(x, float) for x in vector)\n\n # Test normalization\n norm = np.linalg.norm(vector)\n assert abs(norm - 1.0) < 1e-5\n\n # Test semantic similarity\n v1 = emb_func.embed(\"The cat sits on the mat\")\n v2 = emb_func.embed(\"A feline rests on a rug\")\n v3 = emb_func.embed(\"Python programming language\")\n\n similarity_high = np.dot(v1, v2)\n similarity_low = np.dot(v1, v3)\n assert similarity_high > similarity_low\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_model_properties(self, mock_require_module):\n \"\"\"Test model_name and model_source properties.\"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_model.get_sentence_embedding_dimension.return_value = 384\n mock_model.device = \"cpu\"\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n # Test Hugging Face\n emb_func_hf = DefaultLocalDenseEmbedding(model_source=\"huggingface\")\n assert emb_func_hf.model_name == \"all-MiniLM-L6-v2\"\n assert emb_func_hf.model_source == \"huggingface\"\n\n # Test ModelScope\n with patch(\n \"modelscope.hub.snapshot_download.snapshot_download\",\n return_value=\"/path/to/model\",\n ):\n mock_ms = Mock()\n mock_require_module.side_effect = (\n lambda m: mock_st if m == \"sentence_transformers\" else mock_ms\n )\n emb_func_ms = DefaultLocalDenseEmbedding(model_source=\"modelscope\")\n assert (\n emb_func_ms.model_name == \"iic/nlp_gte_sentence-embedding_chinese-small\"\n )\n assert emb_func_ms.model_source == \"modelscope\"\n\n\n# -----------------------------------\n# DefaultLocalSparseEmbedding Test Case\n# -----------------------------------\nclass TestDefaultLocalSparseEmbedding:\n \"\"\"Test suite for DefaultLocalSparseEmbedding (SPLADE sparse embedding).\n\n Note:\n DefaultLocalSparseEmbedding uses naver/splade-cocondenser-ensembledistil\n instead of naver/splade-v3 because:\n\n - splade-v3 is a gated model requiring Hugging Face authentication\n - cocondenser-ensembledistil is publicly accessible\n - Performance difference is minimal (~2%)\n - Avoids \"Access to model is restricted\" errors\n\n This allows all users to run tests without authentication setup.\n \"\"\"\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_init_success(self, mock_require_module):\n \"\"\"Test successful initialization.\n\n Verifies that DefaultLocalSparseEmbedding initializes with the publicly\n accessible naver/splade-cocondenser-ensembledistil model instead of\n the gated naver/splade-v3 model.\n \"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_model.device = \"cpu\"\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding()\n\n assert sparse_emb.model_name == \"naver/splade-cocondenser-ensembledistil\"\n assert sparse_emb.model_source == \"huggingface\"\n assert sparse_emb.device == \"cpu\"\n mock_st.SentenceTransformer.assert_called_once_with(\n \"naver/splade-cocondenser-ensembledistil\",\n device=None,\n trust_remote_code=True,\n )\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_init_with_custom_device(self, mock_require_module):\n \"\"\"Test initialization with custom device.\"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_model.device = \"cuda\"\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding(device=\"cuda\")\n\n assert sparse_emb.device == \"cuda\"\n mock_st.SentenceTransformer.assert_called_once_with(\n \"naver/splade-cocondenser-ensembledistil\",\n device=\"cuda\",\n trust_remote_code=True,\n )\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_embed_success(self, mock_require_module):\n \"\"\"Test successful sparse embedding generation with official API.\"\"\"\n import numpy as np\n\n # Clear model cache to ensure fresh mock\n from zvec.extension.sentence_transformer_embedding_function import (\n DefaultLocalSparseEmbedding,\n )\n\n DefaultLocalSparseEmbedding.clear_cache()\n\n # Create a mock sparse matrix that simulates scipy.sparse behavior\n # The code will call: sparse_matrix[0].toarray().flatten()\n mock_sparse_matrix = Mock()\n\n # Create a dense array representation with vocab_size=30522\n vocab_size = 30522\n dense_array = np.zeros(vocab_size)\n # Set specific non-zero values at indices [10, 245, 1023, 5678]\n dense_array[10] = 0.5\n dense_array[245] = 0.8\n dense_array[1023] = 1.2\n dense_array[5678] = 0.3\n\n # Mock the method chain: sparse_matrix[0].toarray().flatten()\n mock_row = Mock()\n mock_dense = Mock()\n mock_row.toarray.return_value = mock_dense\n mock_dense.flatten.return_value = dense_array\n mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row)\n\n # Also mock hasattr check for 'toarray'\n mock_sparse_matrix.toarray = Mock()\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.device = \"cpu\"\n\n # Configure mock methods to return sparse matrix\n # Must set return_value BEFORE hasattr() check in the code\n mock_model.encode_query = Mock(return_value=mock_sparse_matrix)\n mock_model.encode_document = Mock(return_value=mock_sparse_matrix)\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding()\n result = sparse_emb.embed(\"machine learning\")\n\n # Verify result is a dictionary\n assert isinstance(result, dict)\n # Verify keys are integers and values are floats\n assert all(isinstance(k, int) for k in result.keys())\n assert all(isinstance(v, float) for v in result.values())\n # Verify all values are positive\n assert all(v > 0 for v in result.values())\n # Sparse vectors should have specific dimensions\n assert len(result) == 4\n\n # Verify output is sorted by indices (keys)\n keys = list(result.keys())\n assert keys == sorted(keys), (\n \"Sparse vector keys must be sorted in ascending order\"\n )\n\n # Verify expected keys\n assert keys == [10, 245, 1023, 5678]\n\n # Verify encode_query was called with a list\n mock_model.encode_query.assert_called_once()\n call_args = mock_model.encode_query.call_args[0][0]\n assert isinstance(call_args, list)\n assert call_args == [\"machine learning\"]\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_embed_empty_input(self, mock_require_module):\n \"\"\"Test embedding with empty input.\"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding()\n\n with pytest.raises(ValueError, match=\"Input text cannot be empty\"):\n sparse_emb.embed(\"\")\n\n with pytest.raises(ValueError, match=\"Input text cannot be empty\"):\n sparse_emb.embed(\" \")\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_embed_non_string_input(self, mock_require_module):\n \"\"\"Test embedding with non-string input.\"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding()\n\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n sparse_emb.embed(123)\n\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n sparse_emb.embed([\"text\"])\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_callable_interface(self, mock_require_module):\n \"\"\"Test that DefaultSparseEmbedding is callable.\"\"\"\n import numpy as np\n\n # Clear model cache\n from zvec.extension.sentence_transformer_embedding_function import (\n DefaultLocalSparseEmbedding,\n )\n\n DefaultLocalSparseEmbedding.clear_cache()\n\n # Create a mock sparse matrix\n mock_sparse_matrix = Mock()\n\n # Create a dense array representation with vocab_size=30522\n vocab_size = 30522\n dense_array = np.zeros(vocab_size)\n # Set specific non-zero values at indices [100, 200, 300]\n dense_array[100] = 1.0\n dense_array[200] = 0.5\n dense_array[300] = 0.8\n\n # Mock the method chain: sparse_matrix[0].toarray().flatten()\n mock_row = Mock()\n mock_dense = Mock()\n mock_row.toarray.return_value = mock_dense\n mock_dense.flatten.return_value = dense_array\n mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row)\n\n # Also mock hasattr check for 'toarray'\n mock_sparse_matrix.toarray = Mock()\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.device = \"cpu\"\n\n # Configure mock methods\n mock_model.encode_query = Mock(return_value=mock_sparse_matrix)\n mock_model.encode_document = Mock(return_value=mock_sparse_matrix)\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding()\n\n # Test callable interface\n result = sparse_emb(\"test input\")\n assert isinstance(result, dict)\n assert all(isinstance(k, int) for k in result.keys())\n\n # Verify sorted output\n keys = list(result.keys())\n assert keys == sorted(keys), \"Callable interface must also return sorted keys\"\n assert keys == [100, 200, 300]\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_model_loading_failure(self, mock_require_module):\n \"\"\"Test handling of model loading failure.\"\"\"\n # Clear model cache to ensure the test actually tries to load the model\n from zvec.extension.sentence_transformer_embedding_function import (\n DefaultLocalSparseEmbedding,\n )\n\n DefaultLocalSparseEmbedding.clear_cache()\n\n mock_st = Mock()\n mock_st.SentenceTransformer.side_effect = Exception(\"Model not found\")\n mock_require_module.return_value = mock_st\n\n with pytest.raises(\n ValueError, match=\"Failed to load Sentence Transformer model\"\n ):\n DefaultLocalSparseEmbedding()\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_inference_failure(self, mock_require_module):\n \"\"\"Test handling of inference failure.\"\"\"\n # Clear model cache\n from zvec.extension.sentence_transformer_embedding_function import (\n DefaultLocalSparseEmbedding,\n )\n\n DefaultLocalSparseEmbedding.clear_cache()\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.device = \"cpu\"\n\n # Configure mock methods to raise RuntimeError\n mock_model.encode_query = Mock(side_effect=RuntimeError(\"CUDA out of memory\"))\n mock_model.encode_document = Mock(\n side_effect=RuntimeError(\"CUDA out of memory\")\n )\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding()\n\n with pytest.raises(RuntimeError, match=\"Failed to generate sparse embedding\"):\n sparse_emb.embed(\"test input\")\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_sparse_vector_properties(self, mock_require_module):\n \"\"\"Test properties of sparse vectors (sparsity, non-zero values, sorted order).\"\"\"\n import numpy as np\n\n # Clear model cache\n from zvec.extension.sentence_transformer_embedding_function import (\n DefaultLocalSparseEmbedding,\n )\n\n DefaultLocalSparseEmbedding.clear_cache()\n\n # Create a mock sparse matrix that simulates scipy.sparse behavior\n # The code will call: sparse_matrix[0].toarray().flatten()\n mock_sparse_matrix = Mock()\n\n # Create a dense array representation with vocab_size=30522\n vocab_size = 30522\n dense_array = np.zeros(vocab_size)\n # Set specific non-zero values at indices [50, 100, 200, 400, 500]\n dense_array[50] = 3.0\n dense_array[100] = 2.0\n dense_array[200] = 1.5\n dense_array[400] = 2.5\n dense_array[500] = 1.8\n\n # Mock the method chain: sparse_matrix[0].toarray().flatten()\n mock_row = Mock()\n mock_dense = Mock()\n mock_row.toarray.return_value = mock_dense\n mock_dense.flatten.return_value = dense_array\n mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row)\n\n # Also mock hasattr check for 'toarray'\n mock_sparse_matrix.toarray = Mock()\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.device = \"cpu\"\n\n # Configure mock methods\n mock_model.encode_query = Mock(return_value=mock_sparse_matrix)\n mock_model.encode_document = Mock(return_value=mock_sparse_matrix)\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding()\n result = sparse_emb.embed(\"test\")\n\n # Verify sparsity: result should have much fewer dimensions than vocab_size\n assert len(result) < vocab_size\n # All values should be positive\n assert all(v > 0 for v in result.values())\n\n # Verify keys are sorted in ascending order\n keys = list(result.keys())\n assert keys == sorted(keys), \"Sparse vector keys must be sorted\"\n\n # Verify the specific non-zero indices are present and sorted\n # Expected order: [50, 100, 200, 400, 500] (sorted)\n expected_keys = [50, 100, 200, 400, 500]\n assert keys == expected_keys, f\"Expected {expected_keys}, got {keys}\"\n\n # First key should be smallest\n if len(result) > 0:\n first_key = next(iter(result.keys()))\n assert first_key == min(result.keys()), \"First key must be the smallest\"\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_output_sorted_by_indices(self, mock_require_module):\n \"\"\"Test that output dictionary is always sorted by indices (keys) in ascending order.\"\"\"\n import numpy as np\n\n # Clear model cache\n from zvec.extension.sentence_transformer_embedding_function import (\n DefaultLocalSparseEmbedding,\n )\n\n DefaultLocalSparseEmbedding.clear_cache()\n\n # Create sparse output with deliberately out-of-order indices\n # Non-sequential indices: 9999, 5, 1234, 77, 500\n mock_sparse_matrix = Mock()\n\n # Create a dense array representation with vocab_size=30522\n vocab_size = 30522\n dense_array = np.zeros(vocab_size)\n # Set specific non-zero values at out-of-order indices\n dense_array[9999] = 1.5\n dense_array[5] = 2.0\n dense_array[1234] = 0.8\n dense_array[77] = 3.2\n dense_array[500] = 1.1\n\n # Mock the method chain: sparse_matrix[0].toarray().flatten()\n mock_row = Mock()\n mock_dense = Mock()\n mock_row.toarray.return_value = mock_dense\n mock_dense.flatten.return_value = dense_array\n mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row)\n\n # Also mock hasattr check for 'toarray'\n mock_sparse_matrix.toarray = Mock()\n\n mock_st = Mock()\n mock_model = Mock()\n mock_model.device = \"cpu\"\n\n # Configure mock methods\n mock_model.encode_query = Mock(return_value=mock_sparse_matrix)\n mock_model.encode_document = Mock(return_value=mock_sparse_matrix)\n\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding()\n result = sparse_emb.embed(\"test sorting\")\n\n # Extract keys from result\n result_keys = list(result.keys())\n\n # Verify keys are sorted\n assert result_keys == sorted(result_keys), (\n f\"Keys must be sorted in ascending order. \"\n f\"Got: {result_keys}, Expected: {sorted(result_keys)}\"\n )\n\n # Verify expected keys are present and in correct order\n # Expected sorted order: [5, 77, 500, 1234, 9999]\n expected_sorted_keys = [5, 77, 500, 1234, 9999]\n assert result_keys == expected_sorted_keys, (\n f\"All expected keys should be present in sorted order. \"\n f\"Expected: {expected_sorted_keys}, Got: {result_keys}\"\n )\n\n # Verify first and last keys\n assert result_keys[0] == 5, \"First key must be minimum\"\n assert result_keys[-1] == 9999, \"Last key must be maximum\"\n\n # Verify iteration order matches sorted order\n for i, (key, value) in enumerate(result.items()):\n if i > 0:\n prev_key = list(result.keys())[i - 1]\n assert key > prev_key, (\n f\"Key at position {i} must be greater than previous key\"\n )\n\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_device_property(self, mock_require_module):\n \"\"\"Test device property returns correct device.\"\"\"\n mock_st = Mock()\n mock_model = Mock()\n mock_model.device = \"cuda\"\n mock_st.SentenceTransformer.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n sparse_emb = DefaultLocalSparseEmbedding(device=\"cuda\")\n assert sparse_emb.device == \"cuda\"\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1 and model download\",\n )\n @patch(\"zvec.extension.sentence_transformer_function.require_module\")\n def test_modelscope_source(self, mock_require_module):\n \"\"\"Test initialization with ModelScope source.\"\"\"\n mock_st = Mock()\n mock_ms = Mock()\n mock_model = Mock()\n mock_model.device = \"cpu\"\n mock_st.SentenceTransformer.return_value = mock_model\n\n # Mock ModelScope snapshot_download\n with patch(\n \"modelscope.hub.snapshot_download.snapshot_download\",\n return_value=\"/cache/splade-cocondenser\",\n ):\n mock_require_module.side_effect = (\n lambda m: mock_st if m == \"sentence_transformers\" else mock_ms\n )\n\n sparse_emb = DefaultLocalSparseEmbedding(model_source=\"modelscope\")\n\n assert sparse_emb.model_name == \"naver/splade-cocondenser-ensembledistil\"\n assert sparse_emb.model_source == \"modelscope\"\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1 and model download\",\n )\n def test_integration_real_model(self):\n \"\"\"Integration test with real SPLADE model (requires model download).\n\n This test uses naver/splade-cocondenser-ensembledistil instead of\n naver/splade-v3 because splade-v3 requires Hugging Face authentication.\n The cocondenser-ensembledistil model is publicly accessible and provides\n comparable performance.\n\n To run this test:\n export ZVEC_RUN_INTEGRATION_TESTS=1\n pytest tests/test_embedding.py::TestDefaultSparseEmbedding::test_integration_real_model -v\n\n Note: First run will download ~100MB model from Hugging Face.\n\n Alternative models:\n If you have access to splade-v3, you can create a custom embedding\n class following the example in DefaultSparseEmbedding docstring.\n \"\"\"\n # Clear model cache to ensure fresh load\n from zvec.extension.sentence_transformer_embedding_function import (\n DefaultLocalSparseEmbedding,\n )\n\n DefaultLocalSparseEmbedding.clear_cache()\n\n sparse_emb = DefaultLocalSparseEmbedding()\n\n # Test with real input\n text = \"machine learning and artificial intelligence\"\n result = sparse_emb.embed(text)\n\n # Verify result structure\n assert isinstance(result, dict)\n assert len(result) > 0\n assert all(isinstance(k, int) and k >= 0 for k in result.keys())\n assert all(isinstance(v, float) and v > 0 for v in result.values())\n\n # SPLADE typically produces 100-300 non-zero dimensions\n assert 50 < len(result) < 500\n\n # Verify keys are sorted in ascending order\n keys = list(result.keys())\n assert keys == sorted(keys), \"Real model output must be sorted by indices\"\n\n # Test callable interface\n result2 = sparse_emb(text)\n assert result == result2\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1\",\n )\n def test_integration_multiple_inputs(self):\n \"\"\"Integration test with multiple different inputs.\"\"\"\n # Clear model cache\n from zvec.extension.sentence_transformer_embedding_function import (\n DefaultLocalSparseEmbedding,\n )\n\n DefaultLocalSparseEmbedding.clear_cache()\n\n sparse_emb = DefaultLocalSparseEmbedding()\n\n texts = [\n \"Hello, world!\",\n \"Machine learning is fascinating\",\n \"Python programming language\",\n ]\n\n results = [sparse_emb.embed(text) for text in texts]\n\n # All results should be different\n assert len(results) == 3\n assert all(isinstance(r, dict) for r in results)\n\n # Different inputs should produce different sparse vectors\n assert results[0] != results[1]\n assert results[1] != results[2]\n\n # All results must be sorted by indices\n for i, result in enumerate(results):\n keys = list(result.keys())\n assert keys == sorted(keys), f\"Result {i} must have sorted keys\"\n\n\n# ----------------------------\n# BM25EmbeddingFunction Test Case\n# ----------------------------\nclass TestBM25EmbeddingFunction:\n \"\"\"Test suite for BM25EmbeddingFunction (BM25-based sparse embedding using DashText SDK).\"\"\"\n\n def test_init_with_built_in_encoder(self):\n \"\"\"Test successful initialization with built-in encoder (no corpus).\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n # Test with default language (Chinese)\n bm25 = BM25EmbeddingFunction()\n\n assert bm25.corpus_size == 0\n assert bm25.encoding_type == \"query\"\n assert bm25.language == \"zh\"\n mock_dashtext.SparseVectorEncoder.default.assert_called_once_with(name=\"zh\")\n\n def test_init_with_custom_encoder(self):\n \"\"\"Test successful initialization with custom encoder (with corpus).\"\"\"\n corpus = [\n \"a cat is a feline and likes to purr\",\n \"a dog is the human's best friend\",\n \"a bird is a beautiful animal that can fly\",\n ]\n\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n mock_dashtext.SparseVectorEncoder.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n bm25 = BM25EmbeddingFunction(corpus=corpus, b=0.75, k1=1.2)\n\n assert bm25.corpus_size == 3\n assert bm25.encoding_type == \"query\"\n mock_dashtext.SparseVectorEncoder.assert_called_once_with(b=0.75, k1=1.2)\n mock_encoder.train.assert_called_once_with(corpus)\n\n def test_init_with_empty_corpus(self):\n \"\"\"Test initialization with empty corpus raises ValueError.\"\"\"\n with pytest.raises(ValueError, match=\"Corpus must be a non-empty list\"):\n BM25EmbeddingFunction(corpus=[])\n\n def test_init_with_invalid_corpus(self):\n \"\"\"Test initialization with invalid corpus elements.\"\"\"\n with pytest.raises(ValueError, match=\"All corpus documents must be strings\"):\n BM25EmbeddingFunction(corpus=[\"text\", 123, \"another\"])\n\n with pytest.raises(ValueError, match=\"All corpus documents must be strings\"):\n BM25EmbeddingFunction(corpus=[None, \"text\"])\n\n def test_init_with_language_parameter(self):\n \"\"\"Test initialization with different language settings.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n # Test English language\n bm25_en = BM25EmbeddingFunction(language=\"en\")\n assert bm25_en.language == \"en\"\n mock_dashtext.SparseVectorEncoder.default.assert_called_with(name=\"en\")\n\n def test_init_with_encoding_type(self):\n \"\"\"Test initialization with different encoding types.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n # Test document encoding type\n bm25_doc = BM25EmbeddingFunction(encoding_type=\"document\")\n assert bm25_doc.encoding_type == \"document\"\n\n def test_init_with_missing_dashtext_library(self):\n \"\"\"Test initialization fails when dashtext library is not installed.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_require.side_effect = ImportError(\"dashtext package is required\")\n\n with pytest.raises(ImportError, match=\"dashtext package is required\"):\n BM25EmbeddingFunction()\n\n def test_embed_with_query_encoding(self):\n \"\"\"Test successful sparse embedding generation with query encoding.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n\n # Mock encode_queries to return sparse vector\n mock_encoder.encode_queries.return_value = {\n 5: 0.89,\n 12: 1.45,\n 23: 0.67,\n 45: 1.12,\n }\n\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n bm25 = BM25EmbeddingFunction(encoding_type=\"query\")\n # Clear LRU cache to ensure fresh call\n bm25.embed.cache_clear()\n result = bm25.embed(\"cat purr loud\")\n\n # Verify result structure\n assert isinstance(result, dict)\n assert all(isinstance(k, int) for k in result.keys())\n assert all(isinstance(v, float) for v in result.values())\n\n # Verify all values are positive\n assert all(v > 0 for v in result.values())\n\n # Verify output is sorted by indices\n keys = list(result.keys())\n assert keys == sorted(keys), \"Output must be sorted by indices\"\n\n # Verify expected keys from mock response\n assert result == {5: 0.89, 12: 1.45, 23: 0.67, 45: 1.12}\n\n # Verify encode_queries was called\n mock_encoder.encode_queries.assert_called_once_with(\"cat purr loud\")\n\n def test_embed_with_document_encoding(self):\n \"\"\"Test successful sparse embedding generation with document encoding.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n\n # Mock encode_documents to return sparse vector\n mock_encoder.encode_documents.return_value = {10: 1.5, 20: 2.3}\n\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n bm25 = BM25EmbeddingFunction(encoding_type=\"document\")\n bm25.embed.cache_clear()\n result = bm25.embed(\"document text\")\n\n assert result == {10: 1.5, 20: 2.3}\n mock_encoder.encode_documents.assert_called_once_with(\"document text\")\n\n def test_embed_with_empty_input(self):\n \"\"\"Test embedding with empty input raises ValueError.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n bm25 = BM25EmbeddingFunction()\n\n with pytest.raises(ValueError, match=\"Input text cannot be empty\"):\n bm25.embed(\"\")\n\n with pytest.raises(ValueError, match=\"Input text cannot be empty\"):\n bm25.embed(\" \")\n\n def test_embed_with_non_string_input(self):\n \"\"\"Test embedding with non-string input raises TypeError.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n bm25 = BM25EmbeddingFunction()\n\n # Test with hashable non-string types - should get our custom error message\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n bm25.embed(123)\n\n with pytest.raises(TypeError, match=\"Expected 'input' to be str\"):\n bm25.embed(None)\n\n # Test with unhashable type (list)\n # Note: lru_cache raises TypeError(\"unhashable type: 'list'\") before our type check\n # This is still a valid type error, just caught at a different layer\n with pytest.raises(TypeError, match=\"unhashable type\"):\n bm25.embed([\"text\"])\n\n def test_embed_callable_interface(self):\n \"\"\"Test that BM25EmbeddingFunction is callable.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n mock_encoder.encode_queries.return_value = {10: 1.5}\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n bm25 = BM25EmbeddingFunction()\n bm25.embed.cache_clear()\n\n # Test callable interface\n result = bm25(\"test query\")\n assert isinstance(result, dict)\n assert 10 in result\n\n def test_embed_output_sorted_by_indices(self):\n \"\"\"Test that output is always sorted by indices in ascending order.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n\n # Mock encode_queries with unsorted indices\n mock_encoder.encode_queries.return_value = {\n 9999: 1.5,\n 5: 2.0,\n 1234: 0.8,\n 77: 3.2,\n 500: 1.1,\n }\n\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n bm25 = BM25EmbeddingFunction()\n bm25.embed.cache_clear()\n result = bm25.embed(\"test query\")\n\n # Verify keys are sorted\n result_keys = list(result.keys())\n assert result_keys == sorted(result_keys), (\n f\"Keys must be sorted. Got: {result_keys}, Expected: {sorted(result_keys)}\"\n )\n\n # Verify expected sorted order: [5, 77, 500, 1234, 9999]\n expected_keys = [5, 77, 500, 1234, 9999]\n assert result_keys == expected_keys\n\n def test_embed_filters_zero_values(self):\n \"\"\"Test that zero and negative values are filtered out.\"\"\"\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n\n # Mock encode_queries with zero and negative values\n mock_encoder.encode_queries.return_value = {\n 0: 1.5, # Positive - should be included\n 1: 0.0, # Zero - should be filtered\n 2: -0.5, # Negative - should be filtered\n }\n\n mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n bm25 = BM25EmbeddingFunction()\n bm25.embed.cache_clear()\n result = bm25.embed(\"test\")\n\n # Only positive token should be in result\n assert 0 in result\n assert 1 not in result # Zero value filtered\n assert 2 not in result # Negative value filtered\n assert all(v > 0 for v in result.values())\n\n def test_properties(self):\n \"\"\"Test property accessors.\"\"\"\n corpus = [\"doc1\", \"doc2\", \"doc3\"]\n\n with patch(\n \"zvec.extension.bm25_embedding_function.require_module\"\n ) as mock_require:\n mock_dashtext = Mock()\n mock_encoder = Mock()\n mock_dashtext.SparseVectorEncoder.return_value = mock_encoder\n mock_require.return_value = mock_dashtext\n\n bm25 = BM25EmbeddingFunction(\n corpus=corpus,\n encoding_type=\"document\",\n language=\"en\",\n b=0.8,\n k1=1.5,\n custom_param=\"test\",\n )\n\n assert bm25.corpus_size == 3\n assert bm25.encoding_type == \"document\"\n assert bm25.language == \"en\"\n assert bm25.extra_params == {\"custom_param\": \"test\"}\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n def test_real_dashtext_bm25_embedding(self):\n \"\"\"Integration test with real DashText library.\n\n To run this test:\n export ZVEC_RUN_INTEGRATION_TESTS=1\n pip install dashtext\n\n Note: This test requires the dashtext package to be installed.\n \"\"\"\n # Test built-in encoder (Chinese)\n bm25_zh = BM25EmbeddingFunction(language=\"zh\", encoding_type=\"query\")\n\n query_zh = \"什么是向量检索服务\"\n result_zh = bm25_zh.embed(query_zh)\n\n assert isinstance(result_zh, dict)\n assert len(result_zh) > 0\n assert all(isinstance(k, int) for k in result_zh.keys())\n assert all(isinstance(v, float) and v > 0 for v in result_zh.values())\n\n # Verify sorted output\n keys = list(result_zh.keys())\n assert keys == sorted(keys), \"Real DashText BM25 output must be sorted\"\n\n # Test custom corpus\n corpus = [\n \"The cat sits on the mat\",\n \"The dog plays in the garden\",\n \"Birds fly in the sky\",\n \"Fish swim in the water\",\n ]\n\n bm25_custom = BM25EmbeddingFunction(corpus=corpus, encoding_type=\"query\")\n\n query_en = \"cat on mat\"\n result_en = bm25_custom.embed(query_en)\n\n assert isinstance(result_en, dict)\n assert len(result_en) > 0\n assert all(isinstance(k, int) for k in result_en.keys())\n assert all(isinstance(v, float) and v > 0 for v in result_en.values())\n\n # Test callable interface\n result2 = bm25_custom(query_en)\n assert result_en == result2\n\n # Verify properties\n assert bm25_custom.corpus_size == 4\n" }, { "path": "python/tests/test_params.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport sys\nimport time\n\n\nimport numpy as np\nimport pytest\nfrom zvec import (\n AddColumnOption,\n AlterColumnOption,\n CollectionOption,\n FlatIndexParam,\n HnswIndexParam,\n IndexOption,\n InvertIndexParam,\n IVFIndexParam,\n OptimizeOption,\n HnswQueryParam,\n IVFQueryParam,\n VectorQuery,\n IndexType,\n MetricType,\n QuantizeType,\n DataType,\n VectorSchema,\n)\n\nfrom _zvec.param import _VectorQuery\n\n# ----------------------------\n# Invert Index Param Test Case\n# ----------------------------\n\n\nclass TestInvertIndexParam:\n def test_default(self):\n param = InvertIndexParam()\n assert param.enable_range_optimization is False\n assert param.enable_extended_wildcard is False\n assert param.type == IndexType.INVERT\n\n def test_custom(self):\n param = InvertIndexParam(\n enable_range_optimization=True, enable_extended_wildcard=True\n )\n assert param.enable_range_optimization is True\n assert param.enable_extended_wildcard is True\n\n def test_readonly(self):\n param = InvertIndexParam()\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n param.enable_range_optimization = False\n param.enable_extended_wildcard = False\n\n\n# ----------------------------\n# Hnsw Index Param Test Case\n# ----------------------------\n\n\nclass TestHnswIndexParam:\n def test_default(self):\n param = HnswIndexParam()\n assert param.metric_type == MetricType.IP\n assert param.m == 50\n assert param.ef_construction == 500\n assert param.quantize_type == QuantizeType.UNDEFINED\n assert param.type == IndexType.HNSW\n\n def test_custom(self):\n param = HnswIndexParam(\n metric_type=MetricType.L2,\n m=10,\n ef_construction=1000,\n quantize_type=QuantizeType.FP16,\n )\n assert param.metric_type == MetricType.L2\n assert param.m == 10\n assert param.ef_construction == 1000\n assert param.quantize_type == QuantizeType.FP16\n\n @pytest.mark.parametrize(\n \"attr\", [\"metric_type\", \"m\", \"ef_construction\", \"quantize_type\"]\n )\n def test_readonly_attributes(self, attr):\n param = HnswIndexParam()\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n setattr(param, attr, getattr(param, attr))\n\n\n# ----------------------------\n# Flat Index Param Test Case\n# ----------------------------\nclass TestFlatIndexParam:\n def test_default(self):\n param = FlatIndexParam()\n assert param.type == IndexType.FLAT\n assert param.quantize_type == QuantizeType.UNDEFINED\n assert param.metric_type == MetricType.IP\n\n def test_custom(self):\n param = FlatIndexParam(\n metric_type=MetricType.L2, quantize_type=QuantizeType.INT8\n )\n assert param.metric_type == MetricType.L2\n assert param.quantize_type == QuantizeType.INT8\n\n @pytest.mark.parametrize(\"attr\", [\"metric_type\", \"quantize_type\"])\n def test_readonly_attributes(self, attr):\n param = FlatIndexParam()\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n setattr(param, attr, getattr(param, attr))\n\n\n# ----------------------------\n# Ivf Index Param Test Case\n# ----------------------------\nclass TestIVFIndexParam:\n def test_default(self):\n param = IVFIndexParam()\n assert param.metric_type == MetricType.IP\n assert param.n_list == 0\n assert param.quantize_type == QuantizeType.UNDEFINED\n assert param.type == IndexType.IVF\n\n def test_custom(self):\n param = IVFIndexParam(\n metric_type=MetricType.L2, n_list=1000, quantize_type=QuantizeType.FP16\n )\n assert param.metric_type == MetricType.L2\n assert param.n_list == 1000\n assert param.quantize_type == QuantizeType.FP16\n assert param.type == IndexType.IVF\n\n @pytest.mark.parametrize(\"attr\", [\"metric_type\", \"n_list\", \"quantize_type\"])\n def test_readonly_attributes(self, attr):\n param = IVFIndexParam()\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n setattr(param, attr, getattr(param, attr))\n\n\n# ----------------------------\n# CollectionOption Test Case\n# ----------------------------\nclass TestCollectionOption:\n def test_default(self):\n option = CollectionOption()\n assert option is not None\n assert option.read_only == False\n assert option.enable_mmap == True\n\n def test_custom(self):\n option = CollectionOption(read_only=True, enable_mmap=False)\n assert option.read_only == True\n assert option.enable_mmap == False\n\n option = CollectionOption(read_only=False, enable_mmap=True)\n assert option.read_only == False\n assert option.enable_mmap == True\n\n @pytest.mark.parametrize(\"attr\", [\"read_only\", \"enable_mmap\"])\n def test_readonly_attributes(self, attr):\n param = CollectionOption()\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n setattr(param, attr, getattr(param, attr))\n\n\n# ----------------------------\n# IndexOption Test Case\n# ----------------------------\nclass TestIndexOption:\n def test_default(self):\n option = IndexOption()\n assert option is not None\n assert option.concurrency == 0\n\n def test_custom(self):\n option = IndexOption(concurrency=10)\n assert option.concurrency == 10\n\n @pytest.mark.parametrize(\"attr\", [\"concurrency\"])\n def test_readonly_attributes(self, attr):\n param = IndexOption()\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n setattr(param, attr, getattr(param, attr))\n\n\n# ----------------------------\n# AddColumnOption Test Case\n# ----------------------------\nclass TestAddColumnOption:\n def test_default(self):\n option = AddColumnOption()\n assert option is not None\n assert option.concurrency == 0\n\n def test_custom(self):\n option = AddColumnOption(concurrency=10)\n assert option.concurrency == 10\n\n @pytest.mark.parametrize(\"attr\", [\"concurrency\"])\n def test_readonly_attributes(self, attr):\n param = AddColumnOption()\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n setattr(param, attr, getattr(param, attr))\n\n\n# ----------------------------\n# AlterColumnOption Test Case\n# ----------------------------\nclass TestAlterColumnOption:\n def test_default(self):\n option = AlterColumnOption()\n assert option is not None\n assert option.concurrency == 0\n\n def test_custom(self):\n option = AlterColumnOption(concurrency=10)\n assert option.concurrency == 10\n\n @pytest.mark.parametrize(\"attr\", [\"concurrency\"])\n def test_readonly_attributes(self, attr):\n param = AlterColumnOption()\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n setattr(param, attr, getattr(param, attr))\n\n\n# ----------------------------\n# OptimizeOption Test Case\n# ----------------------------\nclass TestOptimizeOption:\n def test_default(self):\n option = OptimizeOption()\n assert option is not None\n assert option.concurrency == 0\n\n def test_custom(self):\n option = OptimizeOption(concurrency=10)\n assert option.concurrency == 10\n\n @pytest.mark.parametrize(\"attr\", [\"concurrency\"])\n def test_readonly_attributes(self, attr):\n param = OptimizeOption()\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n setattr(param, attr, getattr(param, attr))\n\n\n# ----------------------------\n# HnswQueryParam Test Case\n# ----------------------------\nclass TestHnswQueryParam:\n def test_default(self):\n param = HnswQueryParam()\n assert param is not None\n assert param.ef == 300\n assert param.is_using_refiner == False\n assert param.radius == 0\n assert param.is_linear == False\n\n def test_custom(self):\n param = HnswQueryParam(ef=10, is_using_refiner=True, radius=30, is_linear=True)\n assert param.ef == 10\n assert param.is_using_refiner == True\n assert param.radius == 30\n assert param.is_linear == True\n\n def test_readonly_attributes(self):\n param = HnswQueryParam()\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n param.ef = 10\n param.is_using_refiner = True\n param.radius = 30\n param.is_linear = True\n\n\n# # ----------------------------\n# # IVFQueryParam Test Case\n# # ----------------------------\n# class TestIVFQueryParam:\n# def test_default(self):\n# param = IVFQueryParam()\n# assert param is not None\n# assert param.nprobe == 10\n# assert param.is_using_refiner == False\n# assert param.radius == 0\n# assert param.is_linear == False\n# assert param.scale_factor == 10\n#\n# def test_custom(self):\n# param = IVFQueryParam(\n# nprobe=20,\n# is_using_refiner=True,\n# radius=30,\n# is_linear=True,\n# scale_factor=40\n# )\n# assert param.nprobe == 20\n# assert param.is_using_refiner == True\n# assert param.radius == 30\n# assert param.is_linear == True\n# assert param.scale_factor == 40\n\n\nclass TestVectorQuery:\n def test_init_with_valid_id(self):\n vq = VectorQuery(field_name=\"embedding\", id=\"doc123\")\n assert vq.field_name == \"embedding\"\n assert vq.id == \"doc123\"\n assert vq.vector is None\n assert vq.param is None\n\n def test_init_with_valid_vector(self):\n vec = [0.1, 0.2, 0.3]\n param = HnswQueryParam(ef=300)\n vq = VectorQuery(field_name=\"embedding\", vector=vec, param=param)\n assert vq.field_name == \"embedding\"\n assert vq.vector == vec\n assert vq.param == param\n\n def test_init_both_id_and_vector_raises_error(self):\n with pytest.raises(ValueError):\n VectorQuery(field_name=\"embedding\", id=\"doc123\", vector=[0.1])._validate()\n\n def test_init_without_field_name_raises_error(self):\n with pytest.raises(ValueError):\n VectorQuery(field_name=None)._validate()\n\n def test_has_id_returns_true_when_id_set(self):\n vq = VectorQuery(field_name=\"embedding\", id=\"doc123\")\n assert vq.has_id()\n\n def test_has_id_returns_false_when_no_id(self):\n vq = VectorQuery(field_name=\"embedding\", vector=[0.1])\n assert not vq.has_id()\n\n def test_has_vector_returns_true_with_non_empty_vector(self):\n vq = VectorQuery(field_name=\"embedding\", vector=[0.1])\n assert vq.has_vector()\n\n def test_validate_fails_on_both_id_and_vector(self):\n vq = VectorQuery(field_name=\"test\", id=\"doc123\", vector=[0.1])\n with pytest.raises(ValueError):\n vq._validate()\n" }, { "path": "python/tests/test_query_executor.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom typing import Dict, Union\nfrom unittest.mock import MagicMock\n\nimport numpy as np\nimport math\nfrom _zvec.param import _VectorQuery\n\nimport pytest\nfrom zvec.executor.query_executor import (\n MultiVectorQueryExecutor,\n NoVectorQueryExecutor,\n QueryContext,\n QueryExecutor,\n QueryExecutorFactory,\n SingleVectorQueryExecutor,\n VectorQuery,\n)\nfrom zvec import RrfReRanker, HnswQueryParam, CollectionSchema, VectorSchema, DataType\n\n\n# ----------------------------\n# Mock Vector Schema\n# ----------------------------\nclass MockVectorSchema(VectorSchema):\n def __init__(self, name=\"test_vector\"):\n self._name = name\n\n @property\n def name(self):\n return self._name\n\n def _get_object(self):\n return MagicMock()\n\n\n# ----------------------------\n# Mock Collection Schema\n# ----------------------------\nclass MockCollectionSchema(CollectionSchema):\n def __init__(self, vectors=Union[VectorSchema, Dict[str, VectorSchema]]):\n self._vectors = (\n [vectors] if not isinstance(vectors, Dict) else list(vectors.values())\n )\n\n @property\n def vectors(self):\n return self._vectors\n\n\n# ----------------------------\n# VectorQuery Test Case\n# ----------------------------\nclass TestVectorQuery:\n def test_init(self):\n query = VectorQuery(field_name=\"test_field\")\n assert query.field_name == \"test_field\"\n assert query.id is None\n assert query.vector is None\n assert query.param is None\n\n param = HnswQueryParam()\n query = VectorQuery(\n field_name=\"test_field\", id=\"test_id\", vector=[1, 2, 3], param=param\n )\n assert query.field_name == \"test_field\"\n assert query.id == \"test_id\"\n assert query.vector == [1, 2, 3]\n assert query.param == param\n\n def test_has_id(self):\n query = VectorQuery(field_name=\"test_field\")\n assert not query.has_id()\n\n query = VectorQuery(field_name=\"test_field\", id=\"test_id\")\n assert query.has_id()\n\n def test_has_vector(self):\n query = VectorQuery(field_name=\"test_field\")\n assert not query.has_vector()\n\n query = VectorQuery(field_name=\"test_field\", vector=[])\n assert not query.has_vector()\n\n query = VectorQuery(field_name=\"test_field\", vector=[1, 2, 3])\n assert query.has_vector()\n\n def test_validate_dense_fp16_convert(self):\n v = _VectorQuery()\n schema = VectorSchema(name=\"test\", data_type=DataType.VECTOR_FP16)\n vec = np.array([1.1, 2.1, 3.1], dtype=np.float16)\n v.set_vector(schema._get_object(), vec)\n ret = v.get_vector(schema._get_object())\n assert np.array_equal(vec, ret)\n\n def test_validate_dense_fp32_convert(self):\n v = _VectorQuery()\n schema = VectorSchema(name=\"test\", data_type=DataType.VECTOR_FP32)\n vec = np.array([1.1, 2.1, 3.1], dtype=np.float32)\n v.set_vector(schema._get_object(), vec)\n ret = v.get_vector(schema._get_object())\n assert np.array_equal(vec, ret)\n\n def test_validate_dense_fp64_convert(self):\n v = _VectorQuery()\n schema = VectorSchema(name=\"test\", data_type=DataType.VECTOR_FP64)\n vec = np.array([1.1, 2.1, 3.1], dtype=np.float64)\n v.set_vector(schema._get_object(), vec)\n ret = v.get_vector(schema._get_object())\n assert np.array_equal(vec, ret)\n\n def test_validate_dense_int8_convert(self):\n v = _VectorQuery()\n schema = VectorSchema(name=\"test\", data_type=DataType.VECTOR_INT8)\n vec = np.array([1, 2, 3], dtype=np.int8)\n v.set_vector(schema._get_object(), vec)\n ret = v.get_vector(schema._get_object())\n assert np.array_equal(vec, ret)\n\n def test_validate_sparse_fp32_convert(self):\n v = _VectorQuery()\n schema = VectorSchema(name=\"test\", data_type=DataType.SPARSE_VECTOR_FP32)\n vec = {1: 1.1, 2: 2.2, 3: 3.3}\n v.set_vector(schema._get_object(), vec)\n ret = v.get_vector(schema._get_object())\n for k in vec.keys():\n assert math.isclose(vec[k], ret[k], abs_tol=1e-6)\n\n def test_validate_sparse_fp16_convert(self):\n v = _VectorQuery()\n schema = VectorSchema(name=\"test\", data_type=DataType.SPARSE_VECTOR_FP16)\n vec = {1: 1.1, 2: 2.2, 3: 3.3}\n v.set_vector(schema._get_object(), vec)\n ret = v.get_vector(schema._get_object())\n for k in vec.keys():\n assert math.isclose(np.float16(vec[k]), ret[k], abs_tol=1e-6)\n\n\nclass TestQueryContext:\n def test_init(self):\n ctx = QueryContext(topk=10)\n assert ctx.topk == 10\n assert ctx.queries == []\n assert ctx.filter is None\n assert ctx.reranker is None\n assert ctx.output_fields is None\n assert ctx.include_vector is False\n assert ctx.core_vectors == []\n\n def test_properties(self):\n queries = [VectorQuery(field_name=\"test\")]\n reranker = RrfReRanker()\n output_fields = [\"field1\", \"field2\"]\n\n ctx = QueryContext(\n topk=5,\n filter=\"test_filter\",\n include_vector=True,\n queries=queries,\n output_fields=output_fields,\n reranker=reranker,\n )\n\n assert ctx.topk == 5\n assert ctx.queries == queries\n assert ctx.filter == \"test_filter\"\n assert ctx.reranker == reranker\n assert ctx.output_fields == output_fields\n assert ctx.include_vector is True\n\n def test_core_vectors_setter(self):\n ctx = QueryContext(topk=10)\n core_vectors = [MagicMock()]\n ctx.core_vectors = core_vectors\n assert ctx.core_vectors == core_vectors\n\n\nclass TestNoVectorQueryExecutor:\n def test_init(self):\n schema = MockCollectionSchema()\n executor = NoVectorQueryExecutor(schema)\n assert isinstance(executor, QueryExecutor)\n\n def test_do_validate_with_queries(self):\n schema = MockCollectionSchema()\n executor = NoVectorQueryExecutor(schema)\n ctx = QueryContext(topk=10, queries=[VectorQuery(field_name=\"test\")])\n\n with pytest.raises(\n ValueError, match=\"Collection does not support query with vector or id\"\n ):\n executor._do_validate(ctx)\n\n def test_do_validate_without_queries(self):\n schema = MockCollectionSchema()\n executor = NoVectorQueryExecutor(schema)\n ctx = QueryContext(topk=10)\n\n executor._do_validate(ctx)\n\n def test_do_build(self):\n schema = MockCollectionSchema()\n executor = NoVectorQueryExecutor(schema)\n ctx = QueryContext(topk=5, filter=\"test_filter\")\n\n result = executor._do_build(ctx, MagicMock())\n assert len(result) == 1\n assert result[0].topk == 5\n assert result[0].filter == \"test_filter\"\n\n\nclass TestSingleVectorQueryExecutor:\n def test_init(self):\n schema = MockCollectionSchema()\n executor = SingleVectorQueryExecutor(schema)\n assert isinstance(executor, NoVectorQueryExecutor)\n\n def test_do_validate_multiple_queries(self):\n schema = MockCollectionSchema()\n executor = SingleVectorQueryExecutor(schema)\n queries = [VectorQuery(field_name=\"test1\"), VectorQuery(field_name=\"test2\")]\n ctx = QueryContext(topk=10, queries=queries)\n\n with pytest.raises(\n ValueError,\n match=\"Collection has only one vector field, cannot query with multiple vectors\",\n ):\n executor._do_validate(ctx)\n\n def test_do_build_without_queries(self):\n schema = MockCollectionSchema()\n executor = SingleVectorQueryExecutor(schema)\n ctx = QueryContext(topk=5)\n\n result = executor._do_build(ctx, MagicMock())\n assert len(result) == 1\n assert result[0].topk == 5\n\n\nclass TestMultiVectorQueryExecutor:\n def test_init(self):\n schema = MockCollectionSchema()\n executor = MultiVectorQueryExecutor(schema)\n assert isinstance(executor, SingleVectorQueryExecutor)\n\n def test_do_validate_multiple_queries_without_reranker(self):\n schema = MockCollectionSchema()\n executor = MultiVectorQueryExecutor(schema)\n queries = [VectorQuery(field_name=\"test1\"), VectorQuery(field_name=\"test2\")]\n ctx = QueryContext(topk=10, queries=queries)\n\n with pytest.raises(\n ValueError, match=\"Reranker is required for multi-vector query\"\n ):\n executor._do_validate(ctx)\n\n def test_do_validate_multiple_queries_with_reranker(self):\n schema = MockCollectionSchema()\n executor = MultiVectorQueryExecutor(schema)\n queries = [VectorQuery(field_name=\"test1\"), VectorQuery(field_name=\"test2\")]\n reranker = RrfReRanker()\n ctx = QueryContext(topk=10, queries=queries, reranker=reranker)\n\n executor._do_validate(ctx)\n\n\nclass TestQueryExecutorFactory:\n def test_create_no_vectors(self):\n schema = MockCollectionSchema()\n executor = QueryExecutorFactory.create(schema)\n assert isinstance(executor, NoVectorQueryExecutor)\n\n def test_create_single_vector(self):\n schema = MockCollectionSchema(vectors=MockVectorSchema())\n executor = QueryExecutorFactory.create(schema)\n assert isinstance(executor, SingleVectorQueryExecutor)\n\n def test_create_multiple_vectors(self):\n schema = MockCollectionSchema(\n vectors={\"test1\": MockVectorSchema(), \"test2\": MockVectorSchema()}\n )\n executor = QueryExecutorFactory.create(schema)\n assert isinstance(executor, MultiVectorQueryExecutor)\n" }, { "path": "python/tests/test_reranker.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom unittest.mock import patch, MagicMock\nimport pytest\nimport math\nimport os\n\nfrom zvec import Doc, MetricType\nfrom zvec.extension.multi_vector_reranker import (\n RrfReRanker,\n WeightedReRanker,\n)\nfrom zvec.extension.sentence_transformer_rerank_function import (\n DefaultLocalReRanker,\n)\nfrom zvec.extension.qwen_rerank_function import QwenReRanker\n\n# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API tests\nRUN_INTEGRATION_TESTS = os.environ.get(\"ZVEC_RUN_INTEGRATION_TESTS\", \"0\") == \"1\"\n\n\n# ----------------------------\n# RrfRanker Test Case\n# ----------------------------\nclass TestRrfReRanker:\n def test_init(self):\n reranker = RrfReRanker(topn=5, rerank_field=\"content\", rank_constant=100)\n assert reranker.topn == 5\n assert reranker.rerank_field == \"content\"\n assert reranker.rank_constant == 100\n\n def test_rrf_score(self):\n reranker = RrfReRanker(rank_constant=60)\n # 根据公式 1.0 / (k + rank + 1)，其中k=60\n assert reranker._rrf_score(0) == 1.0 / (60 + 0 + 1)\n assert reranker._rrf_score(1) == 1.0 / (60 + 1 + 1)\n assert reranker._rrf_score(10) == 1.0 / (60 + 10 + 1)\n\n def test_rerank(self):\n reranker = RrfReRanker(topn=3)\n\n doc1 = Doc(id=\"1\", score=0.8)\n doc2 = Doc(id=\"2\", score=0.7)\n doc3 = Doc(id=\"3\", score=0.9)\n doc4 = Doc(id=\"4\", score=0.6)\n\n query_results = {\"vector1\": [doc1, doc2, doc3], \"vector2\": [doc3, doc1, doc4]}\n\n results = reranker.rerank(query_results)\n\n assert len(results) <= reranker.topn\n\n for doc in results:\n assert hasattr(doc, \"score\")\n\n scores = [doc.score for doc in results]\n assert scores == sorted(scores, reverse=True)\n\n\n# ----------------------------\n# WeightedRanker Test Case\n# ----------------------------\nclass TestWeightedReRanker:\n def test_init(self):\n weights = {\"vector1\": 0.7, \"vector2\": 0.3}\n reranker = WeightedReRanker(\n topn=5,\n rerank_field=\"content\",\n metric=MetricType.L2,\n weights=weights,\n )\n assert reranker.topn == 5\n assert reranker.rerank_field == \"content\"\n assert reranker.metric == MetricType.L2\n assert reranker.weights == weights\n\n def test_normalize_score(self):\n reranker = WeightedReRanker()\n\n score = reranker._normalize_score(1.0, MetricType.L2)\n expected = 1.0 - 2 * math.atan(1.0) / math.pi\n assert score == expected\n\n score = reranker._normalize_score(1.0, MetricType.IP)\n expected = 0.5 + math.atan(1.0) / math.pi\n assert score == expected\n\n score = reranker._normalize_score(1.0, MetricType.COSINE)\n expected = 1.0 - 1.0 / 2.0\n assert score == expected\n\n with pytest.raises(ValueError, match=\"Unsupported metric type\"):\n reranker._normalize_score(1.0, \"unsupported_metric\")\n\n def test_rerank(self):\n weights = {\"vector1\": 0.7, \"vector2\": 0.3}\n reranker = WeightedReRanker(topn=3, weights=weights, metric=MetricType.L2)\n\n doc1 = Doc(id=\"1\", score=0.8)\n doc2 = Doc(id=\"2\", score=0.7)\n doc3 = Doc(id=\"3\", score=0.9)\n\n query_results = {\"vector1\": [doc1, doc2], \"vector2\": [doc2, doc3]}\n\n results = reranker.rerank(query_results)\n\n assert len(results) <= reranker.topn\n\n for doc in results:\n assert hasattr(doc, \"score\")\n\n scores = [doc.score for doc in results]\n assert scores == sorted(scores, reverse=True)\n\n\n# ----------------------------\n# QwenReRanker Test Case\n# ----------------------------\nclass TestQwenReRanker:\n def test_init_without_query(self):\n with pytest.raises(ValueError, match=\"Query is required for QwenReRanker\"):\n QwenReRanker(api_key=\"test_key\")\n\n def test_init_without_api_key(self):\n with patch.dict(os.environ, {}, clear=True):\n with pytest.raises(ValueError, match=\"DashScope API key is required\"):\n QwenReRanker(query=\"test\")\n\n @patch.dict(os.environ, {\"DASHSCOPE_API_KEY\": \"test_key\"})\n def test_init_with_env_api_key(self):\n reranker = QwenReRanker(query=\"test\", rerank_field=\"content\")\n assert reranker.query == \"test\"\n assert reranker._api_key == \"test_key\"\n assert reranker.rerank_field == \"content\"\n\n def test_init_with_explicit_api_key(self):\n reranker = QwenReRanker(\n query=\"test\", api_key=\"explicit_key\", rerank_field=\"content\"\n )\n assert reranker.query == \"test\"\n assert reranker._api_key == \"explicit_key\"\n\n def test_model_property(self):\n reranker = QwenReRanker(\n query=\"test\", api_key=\"test_key\", rerank_field=\"content\"\n )\n assert reranker.model == \"gte-rerank-v2\"\n\n reranker = QwenReRanker(\n query=\"test\",\n model=\"custom-model\",\n api_key=\"test_key\",\n rerank_field=\"content\",\n )\n assert reranker.model == \"custom-model\"\n\n def test_query_property(self):\n reranker = QwenReRanker(\n query=\"test query\", api_key=\"test_key\", rerank_field=\"content\"\n )\n assert reranker.query == \"test query\"\n\n def test_topn_property(self):\n reranker = QwenReRanker(\n query=\"test\", topn=5, api_key=\"test_key\", rerank_field=\"content\"\n )\n assert reranker.topn == 5\n\n def test_rerank_field_property(self):\n reranker = QwenReRanker(query=\"test\", api_key=\"test_key\", rerank_field=\"title\")\n assert reranker.rerank_field == \"title\"\n\n def test_rerank_empty_results(self):\n reranker = QwenReRanker(\n query=\"test\", api_key=\"test_key\", rerank_field=\"content\"\n )\n results = reranker.rerank({})\n assert results == []\n\n def test_rerank_no_valid_documents(self):\n reranker = QwenReRanker(\n query=\"test\", api_key=\"test_key\", rerank_field=\"content\"\n )\n # Document without the rerank_field\n query_results = {\"vector1\": [Doc(id=\"1\")]}\n with pytest.raises(ValueError, match=\"No documents to rerank\"):\n reranker.rerank(query_results)\n\n def test_rerank_skip_empty_content(self):\n reranker = QwenReRanker(\n query=\"test\", api_key=\"test_key\", rerank_field=\"content\"\n )\n query_results = {\n \"vector1\": [\n Doc(id=\"1\", fields={\"content\": \"\"}),\n Doc(id=\"2\", fields={\"content\": \" \"}),\n ]\n }\n with pytest.raises(ValueError, match=\"No documents to rerank\"):\n reranker.rerank(query_results)\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_rerank_success(self, mock_require_module):\n # Mock dashscope module\n mock_dashscope = MagicMock()\n mock_require_module.return_value = mock_dashscope\n\n # Mock API response\n mock_response = MagicMock()\n mock_response.status_code = 200\n mock_response.output = {\n \"results\": [\n {\"index\": 0, \"relevance_score\": 0.95},\n {\"index\": 1, \"relevance_score\": 0.85},\n ]\n }\n mock_dashscope.TextReRank.call.return_value = mock_response\n\n reranker = QwenReRanker(\n query=\"test query\", topn=2, api_key=\"test_key\", rerank_field=\"content\"\n )\n\n query_results = {\n \"vector1\": [\n Doc(id=\"1\", fields={\"content\": \"Document 1\"}),\n Doc(id=\"2\", fields={\"content\": \"Document 2\"}),\n ]\n }\n\n results = reranker.rerank(query_results)\n\n assert len(results) == 2\n assert results[0].id == \"1\"\n assert results[0].score == 0.95\n assert results[1].id == \"2\"\n assert results[1].score == 0.85\n\n # Verify API call\n mock_dashscope.TextReRank.call.assert_called_once_with(\n model=\"gte-rerank-v2\",\n query=\"test query\",\n documents=[\"Document 1\", \"Document 2\"],\n top_n=2,\n return_documents=False,\n )\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_rerank_deduplicate_documents(self, mock_require_module):\n # Mock dashscope module\n mock_dashscope = MagicMock()\n mock_require_module.return_value = mock_dashscope\n\n # Mock API response\n mock_response = MagicMock()\n mock_response.status_code = 200\n mock_response.output = {\n \"results\": [\n {\"index\": 0, \"relevance_score\": 0.9},\n ]\n }\n mock_dashscope.TextReRank.call.return_value = mock_response\n\n reranker = QwenReRanker(\n query=\"test\", topn=5, api_key=\"test_key\", rerank_field=\"content\"\n )\n\n # Same document in multiple vector results\n doc1 = Doc(id=\"1\", fields={\"content\": \"Document 1\"})\n query_results = {\"vector1\": [doc1], \"vector2\": [doc1]}\n\n results = reranker.rerank(query_results)\n\n # Should only call API with document once\n call_args = mock_dashscope.TextReRank.call.call_args\n assert len(call_args[1][\"documents\"]) == 1\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_rerank_api_error(self, mock_require_module):\n # Mock dashscope module\n mock_dashscope = MagicMock()\n mock_require_module.return_value = mock_dashscope\n\n # Mock API error response\n mock_response = MagicMock()\n mock_response.status_code = 400\n mock_response.message = \"Invalid request\"\n mock_response.code = \"InvalidParameter\"\n mock_dashscope.TextReRank.call.return_value = mock_response\n\n reranker = QwenReRanker(\n query=\"test\", api_key=\"test_key\", rerank_field=\"content\"\n )\n\n query_results = {\"vector1\": [Doc(id=\"1\", fields={\"content\": \"Document 1\"})]}\n\n with pytest.raises(ValueError, match=\"DashScope API error\"):\n reranker.rerank(query_results)\n\n @patch(\"zvec.extension.qwen_function.require_module\")\n def test_rerank_runtime_error(self, mock_require_module):\n # Mock dashscope module that raises exception\n mock_dashscope = MagicMock()\n mock_require_module.return_value = mock_dashscope\n mock_dashscope.TextReRank.call.side_effect = Exception(\"Network error\")\n\n reranker = QwenReRanker(\n query=\"test\", api_key=\"test_key\", rerank_field=\"content\"\n )\n\n query_results = {\"vector1\": [Doc(id=\"1\", fields={\"content\": \"Document 1\"})]}\n\n with pytest.raises(RuntimeError, match=\"Failed to call DashScope API\"):\n reranker.rerank(query_results)\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n def test_real_qwen_rerank(self):\n \"\"\"Integration test with real DashScope TextReRank API.\n\n To run this test, set environment variables:\n export ZVEC_RUN_INTEGRATION_TESTS=1\n export DASHSCOPE_API_KEY=your-api-key\n \"\"\"\n # Create reranker with real API\n reranker = QwenReRanker(\n query=\"What is machine learning?\",\n topn=3,\n rerank_field=\"content\",\n model=\"gte-rerank-v2\",\n )\n\n # Prepare test documents\n query_results = {\n \"vector1\": [\n Doc(\n id=\"1\",\n score=0.8,\n fields={\n \"content\": \"Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data.\"\n },\n ),\n Doc(\n id=\"2\",\n score=0.7,\n fields={\n \"content\": \"The weather is nice today with clear skies and sunshine.\"\n },\n ),\n Doc(\n id=\"3\",\n score=0.75,\n fields={\n \"content\": \"Deep learning is a specialized branch of machine learning using neural networks with multiple layers.\"\n },\n ),\n ],\n \"vector2\": [\n Doc(\n id=\"4\",\n score=0.6,\n fields={\n \"content\": \"Python is a popular programming language for data science and machine learning applications.\"\n },\n ),\n Doc(\n id=\"5\",\n score=0.65,\n fields={\n \"content\": \"A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder.\"\n },\n ),\n ],\n }\n\n # Call real API\n results = reranker.rerank(query_results)\n\n # Verify results\n assert len(results) <= 3, \"Should return at most topn documents\"\n assert len(results) > 0, \"Should return at least one document\"\n\n # All results should have valid scores\n for doc in results:\n assert hasattr(doc, \"score\"), \"Each document should have a score\"\n assert isinstance(doc.score, (int, float)), \"Score should be numeric\"\n assert doc.score > 0, \"Score should be positive\"\n\n # Verify scores are in descending order\n scores = [doc.score for doc in results]\n assert scores == sorted(scores, reverse=True), (\n \"Results should be sorted by score in descending order\"\n )\n\n # Verify relevant documents are ranked higher\n # Document 1 and 3 are about machine learning, should rank higher than weather/recipe docs\n result_ids = [doc.id for doc in results]\n\n # At least one of the ML-related documents should be in top results\n ml_related_docs = {\"1\", \"3\", \"4\"}\n assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (\n \"ML-related documents should rank higher\"\n )\n\n # Print results for manual verification (useful during development)\n print(\"\\nReranking results:\")\n for i, doc in enumerate(results, 1):\n print(f\"{i}. ID={doc.id}, Score={doc.score:.4f}\")\n if doc.fields:\n content = doc.field(\"content\")\n if content:\n print(f\" Content: {content[:80]}...\")\n\n\n# ----------------------------\n# DefaultLocalReRanker Test Case\n# ----------------------------\nclass TestDefaultLocalReRanker:\n \"\"\"Test cases for DefaultLocalReRanker.\"\"\"\n\n def test_init_without_query(self):\n \"\"\"Test initialization fails without query.\"\"\"\n with pytest.raises(\n ValueError, match=\"Query is required for DefaultLocalReRanker\"\n ):\n DefaultLocalReRanker(rerank_field=\"content\")\n\n def test_init_with_empty_query(self):\n \"\"\"Test initialization fails with empty query.\"\"\"\n with pytest.raises(\n ValueError, match=\"Query is required for DefaultLocalReRanker\"\n ):\n DefaultLocalReRanker(query=\"\", rerank_field=\"content\")\n\n @patch(\"zvec.extension.sentence_transformer_rerank_function.require_module\")\n def test_init_success(self, mock_require_module):\n \"\"\"Test successful initialization with mocked model.\"\"\"\n # Mock sentence_transformers module\n mock_st = MagicMock()\n mock_model = MagicMock()\n mock_model.predict = MagicMock() # Cross-encoder has predict method\n mock_model.device = \"cpu\"\n mock_st.CrossEncoder.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n reranker = DefaultLocalReRanker(\n query=\"test query\",\n topn=5,\n rerank_field=\"content\",\n model_name=\"cross-encoder/ms-marco-MiniLM-L6-v2\",\n )\n\n assert reranker.query == \"test query\"\n assert reranker.topn == 5\n assert reranker.rerank_field == \"content\"\n assert reranker.model_name == \"cross-encoder/ms-marco-MiniLM-L6-v2\"\n assert reranker.model_source == \"huggingface\"\n assert reranker.batch_size == 32\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n @patch(\"zvec.extension.sentence_transformer_rerank_function.require_module\")\n def test_init_with_custom_params(self, mock_require_module):\n \"\"\"Test initialization with custom parameters.\"\"\"\n mock_st = MagicMock()\n mock_model = MagicMock()\n mock_model.predict = MagicMock()\n mock_model.device = \"cuda\"\n mock_st.CrossEncoder.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n reranker = DefaultLocalReRanker(\n query=\"custom query\",\n topn=10,\n rerank_field=\"title\",\n model_name=\"cross-encoder/ms-marco-MiniLM-L12-v2\",\n model_source=\"modelscope\",\n device=\"cuda\",\n batch_size=64,\n )\n\n assert reranker.query == \"custom query\"\n assert reranker.topn == 10\n assert reranker.rerank_field == \"title\"\n assert reranker.model_name == \"cross-encoder/ms-marco-MiniLM-L12-v2\"\n assert reranker.model_source == \"modelscope\"\n assert reranker.batch_size == 64\n\n @patch(\"zvec.extension.sentence_transformer_rerank_function.require_module\")\n def test_init_invalid_model(self, mock_require_module):\n \"\"\"Test initialization fails with non-cross-encoder model.\"\"\"\n # Mock a model without predict method (not a cross-encoder)\n mock_st = MagicMock()\n mock_model = MagicMock(spec=[]) # No predict method\n mock_st.CrossEncoder.return_value = mock_model\n mock_require_module.return_value = mock_st\n\n with pytest.raises(ValueError, match=\"does not appear to be a cross-encoder\"):\n DefaultLocalReRanker(query=\"test\", rerank_field=\"content\")\n\n def test_query_property(self):\n \"\"\"Test query property.\"\"\"\n mock_model = MagicMock()\n mock_model.predict = MagicMock()\n\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(query=\"test query\", rerank_field=\"content\")\n assert reranker.query == \"test query\"\n\n def test_topn_property(self):\n \"\"\"Test topn property.\"\"\"\n mock_model = MagicMock()\n mock_model.predict = MagicMock()\n\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(\n query=\"test\", topn=15, rerank_field=\"content\"\n )\n assert reranker.topn == 15\n\n def test_rerank_field_property(self):\n \"\"\"Test rerank_field property.\"\"\"\n mock_model = MagicMock()\n mock_model.predict = MagicMock()\n\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(query=\"test\", rerank_field=\"title\")\n assert reranker.rerank_field == \"title\"\n\n def test_batch_size_property(self):\n \"\"\"Test batch_size property.\"\"\"\n mock_model = MagicMock()\n mock_model.predict = MagicMock()\n\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(\n query=\"test\", rerank_field=\"content\", batch_size=128\n )\n assert reranker.batch_size == 128\n\n def test_rerank_empty_results(self):\n \"\"\"Test rerank with empty query_results.\"\"\"\n mock_model = MagicMock()\n mock_model.predict = MagicMock()\n\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(query=\"test\", rerank_field=\"content\")\n results = reranker.rerank({})\n assert results == []\n\n def test_rerank_no_valid_documents(self):\n \"\"\"Test rerank with documents missing rerank_field.\"\"\"\n mock_model = MagicMock()\n mock_model.predict = MagicMock()\n\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(query=\"test\", rerank_field=\"content\")\n\n # Document without the rerank_field\n query_results = {\"vector1\": [Doc(id=\"1\")]}\n with pytest.raises(ValueError, match=\"No documents to rerank\"):\n reranker.rerank(query_results)\n\n def test_rerank_skip_empty_content(self):\n \"\"\"Test rerank skips documents with empty content.\"\"\"\n mock_model = MagicMock()\n mock_model.predict = MagicMock()\n\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(query=\"test\", rerank_field=\"content\")\n\n query_results = {\n \"vector1\": [\n Doc(id=\"1\", fields={\"content\": \"\"}),\n Doc(id=\"2\", fields={\"content\": \" \"}),\n ]\n }\n with pytest.raises(ValueError, match=\"No documents to rerank\"):\n reranker.rerank(query_results)\n\n def test_rerank_success(self):\n \"\"\"Test successful rerank with mocked model.\"\"\"\n # Mock standard cross-encoder model\n mock_model = MagicMock()\n\n # Mock predict method to return scores\n import numpy as np\n\n mock_scores = np.array([0.95, 0.85, 0.75])\n mock_model.predict.return_value = mock_scores\n mock_model.device = \"cpu\"\n\n # Mock sentence_transformers module\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(\n query=\"test query\", topn=3, rerank_field=\"content\"\n )\n\n query_results = {\n \"vector1\": [\n Doc(id=\"1\", score=0.8, fields={\"content\": \"Document 1\"}),\n Doc(id=\"2\", score=0.7, fields={\"content\": \"Document 2\"}),\n Doc(id=\"3\", score=0.6, fields={\"content\": \"Document 3\"}),\n ]\n }\n\n results = reranker.rerank(query_results)\n\n # Verify results\n assert len(results) == 3\n assert results[0].id == \"1\"\n assert results[0].score == 0.95\n assert results[1].id == \"2\"\n assert results[1].score == 0.85\n assert results[2].id == \"3\"\n assert results[2].score == 0.75\n\n # Verify model.predict was called correctly\n assert mock_model.predict.called\n call_args = mock_model.predict.call_args\n pairs = call_args[0][0]\n assert len(pairs) == 3\n assert pairs[0] == [\"test query\", \"Document 1\"]\n assert pairs[1] == [\"test query\", \"Document 2\"]\n assert pairs[2] == [\"test query\", \"Document 3\"]\n assert call_args[1][\"batch_size\"] == 32\n assert call_args[1][\"show_progress_bar\"] is False\n\n def test_rerank_with_topn_limit(self):\n \"\"\"Test rerank respects topn limit.\"\"\"\n mock_model = MagicMock()\n\n import numpy as np\n\n mock_scores = np.array([0.9, 0.8, 0.7, 0.6, 0.5])\n mock_model.predict.return_value = mock_scores\n\n # Mock sentence_transformers module\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(\n query=\"test\", topn=2, rerank_field=\"content\"\n )\n\n query_results = {\n \"vector1\": [\n Doc(id=\"1\", fields={\"content\": \"Doc 1\"}),\n Doc(id=\"2\", fields={\"content\": \"Doc 2\"}),\n Doc(id=\"3\", fields={\"content\": \"Doc 3\"}),\n Doc(id=\"4\", fields={\"content\": \"Doc 4\"}),\n Doc(id=\"5\", fields={\"content\": \"Doc 5\"}),\n ]\n }\n\n results = reranker.rerank(query_results)\n\n # Should only return top 2\n assert len(results) == 2\n assert results[0].id == \"1\"\n assert results[0].score == 0.9\n assert results[1].id == \"2\"\n assert results[1].score == 0.8\n\n def test_rerank_deduplicate_documents(self):\n \"\"\"Test rerank deduplicates documents across multiple vectors.\"\"\"\n mock_model = MagicMock()\n\n import numpy as np\n\n mock_scores = np.array([0.95, 0.85])\n mock_model.predict.return_value = mock_scores\n\n # Mock sentence_transformers module\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(\n query=\"test\", topn=5, rerank_field=\"content\"\n )\n\n # Same document in multiple vector results\n doc1 = Doc(id=\"1\", fields={\"content\": \"Document 1\"})\n doc2 = Doc(id=\"2\", fields={\"content\": \"Document 2\"})\n\n query_results = {\n \"vector1\": [doc1, doc2],\n \"vector2\": [doc1], # doc1 appears in both\n }\n\n results = reranker.rerank(query_results)\n\n # Should only process each document once\n assert len(results) == 2\n assert mock_model.predict.call_count == 1\n\n call_args = mock_model.predict.call_args\n pairs = call_args[0][0]\n assert len(pairs) == 2 # Only 2 unique documents\n\n def test_rerank_sorting(self):\n \"\"\"Test rerank sorts documents by score in descending order.\"\"\"\n mock_model = MagicMock()\n\n import numpy as np\n\n # Return scores in non-sorted order\n mock_scores = np.array([0.6, 0.9, 0.7])\n mock_model.predict.return_value = mock_scores\n\n # Mock sentence_transformers module\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(\n query=\"test\", topn=3, rerank_field=\"content\"\n )\n\n query_results = {\n \"vector1\": [\n Doc(id=\"1\", fields={\"content\": \"Doc 1\"}),\n Doc(id=\"2\", fields={\"content\": \"Doc 2\"}),\n Doc(id=\"3\", fields={\"content\": \"Doc 3\"}),\n ]\n }\n\n results = reranker.rerank(query_results)\n\n # Should be sorted by score (descending)\n assert len(results) == 3\n assert results[0].id == \"2\" # score 0.9\n assert results[0].score == 0.9\n assert results[1].id == \"3\" # score 0.7\n assert results[1].score == 0.7\n assert results[2].id == \"1\" # score 0.6\n assert results[2].score == 0.6\n\n def test_rerank_model_error(self):\n \"\"\"Test rerank handles model prediction errors.\"\"\"\n mock_model = MagicMock()\n\n # Mock predict to raise exception\n mock_model.predict.side_effect = Exception(\"Model inference error\")\n\n # Mock sentence_transformers module\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(query=\"test\", rerank_field=\"content\")\n\n query_results = {\"vector1\": [Doc(id=\"1\", fields={\"content\": \"Document 1\"})]}\n\n with pytest.raises(RuntimeError, match=\"Failed to compute rerank scores\"):\n reranker.rerank(query_results)\n\n def test_rerank_with_custom_batch_size(self):\n \"\"\"Test rerank uses custom batch_size.\"\"\"\n mock_model = MagicMock()\n\n import numpy as np\n\n mock_scores = np.array([0.9, 0.8])\n mock_model.predict.return_value = mock_scores\n\n # Mock sentence_transformers module\n mock_st = MagicMock()\n mock_st.CrossEncoder.return_value = mock_model\n\n with patch(\n \"zvec.extension.sentence_transformer_rerank_function.require_module\",\n return_value=mock_st,\n ):\n reranker = DefaultLocalReRanker(\n query=\"test\", rerank_field=\"content\", batch_size=64\n )\n\n query_results = {\n \"vector1\": [\n Doc(id=\"1\", fields={\"content\": \"Doc 1\"}),\n Doc(id=\"2\", fields={\"content\": \"Doc 2\"}),\n ]\n }\n\n reranker.rerank(query_results)\n\n # Verify batch_size is passed to predict\n call_args = mock_model.predict.call_args\n assert call_args[1][\"batch_size\"] == 64\n\n @pytest.mark.skipif(\n not RUN_INTEGRATION_TESTS,\n reason=\"Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.\",\n )\n def test_real_sentence_transformer_rerank(self):\n \"\"\"Integration test with real SentenceTransformer cross-encoder model.\n\n To run this test, set environment variable:\n export ZVEC_RUN_INTEGRATION_TESTS=1\n\n Note: This test requires sentence-transformers package and will\n download the MS MARCO MiniLM model (~80MB) on first run.\n \"\"\"\n # Create reranker with real model (using default lightweight model)\n reranker = DefaultLocalReRanker(\n query=\"What is machine learning?\",\n topn=3,\n rerank_field=\"content\",\n )\n\n # Prepare test documents\n query_results = {\n \"vector1\": [\n Doc(\n id=\"1\",\n score=0.8,\n fields={\n \"content\": \"Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data.\"\n },\n ),\n Doc(\n id=\"2\",\n score=0.7,\n fields={\n \"content\": \"The weather is nice today with clear skies and sunshine.\"\n },\n ),\n Doc(\n id=\"3\",\n score=0.75,\n fields={\n \"content\": \"Deep learning is a specialized branch of machine learning using neural networks with multiple layers.\"\n },\n ),\n ],\n \"vector2\": [\n Doc(\n id=\"4\",\n score=0.6,\n fields={\n \"content\": \"Python is a popular programming language for data science and machine learning applications.\"\n },\n ),\n Doc(\n id=\"5\",\n score=0.65,\n fields={\n \"content\": \"A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder.\"\n },\n ),\n ],\n }\n\n # Call real model\n results = reranker.rerank(query_results)\n\n # Verify results\n assert len(results) <= 3, \"Should return at most topn documents\"\n assert len(results) > 0, \"Should return at least one document\"\n\n # All results should have valid scores\n for doc in results:\n assert hasattr(doc, \"score\"), \"Each document should have a score\"\n assert isinstance(doc.score, (int, float)), \"Score should be numeric\"\n\n # Verify scores are in descending order\n scores = [doc.score for doc in results]\n assert scores == sorted(scores, reverse=True), (\n \"Results should be sorted by score in descending order\"\n )\n\n # Verify relevant documents are ranked higher\n # Documents 1, 3, and 4 are about machine learning, should rank higher\n result_ids = [doc.id for doc in results]\n\n # At least one of the ML-related documents should be in top results\n ml_related_docs = {\"1\", \"3\", \"4\"}\n assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (\n \"ML-related documents should rank higher\"\n )\n\n # Print results for manual verification (useful during development)\n print(\"\\nSentenceTransformer Reranking results:\")\n for i, doc in enumerate(results, 1):\n print(f\"{i}. ID={doc.id}, Score={doc.score:.4f}\")\n if doc.fields:\n content = doc.field(\"content\")\n if content:\n print(f\" Content: {content[:80]}...\")\n" }, { "path": "python/tests/test_schema.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport pytest\nfrom zvec import (\n CollectionSchema,\n CollectionStats,\n FieldSchema,\n VectorSchema,\n HnswIndexParam,\n InvertIndexParam,\n DataType,\n IndexType,\n MetricType,\n)\n\n# ----------------------------\n# FieldSchema Test Case\n# ----------------------------\n\n\nclass TestFieldSchema:\n def test_default(self):\n field = FieldSchema(\"field\", data_type=DataType.FLOAT)\n assert field.name == \"field\"\n assert field.data_type == DataType.FLOAT\n assert field.nullable is False\n assert field.index_param is None\n\n def test_custom(self):\n field_1 = FieldSchema(\n name=\"float\",\n data_type=DataType.FLOAT,\n nullable=True,\n index_param=InvertIndexParam(),\n )\n assert field_1.name == \"float\"\n assert field_1.data_type == DataType.FLOAT\n assert field_1.nullable is True\n assert field_1.index_param.enable_range_optimization is False\n\n field_2 = FieldSchema(\n name=\"str\",\n data_type=DataType.STRING,\n nullable=True,\n index_param=InvertIndexParam(enable_range_optimization=True),\n )\n assert field_2.name == \"str\"\n assert field_2.data_type == DataType.STRING\n assert field_2.nullable is True\n assert field_2.index_param.enable_range_optimization is True\n\n def test_readonly(self):\n field = FieldSchema(\n name=\"float\",\n data_type=DataType.FLOAT,\n nullable=True,\n index_param=InvertIndexParam(),\n )\n\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n field.index_param = InvertIndexParam(enable_range_optimization=True)\n\n\n# ----------------------------\n# VectorSchema Test Case\n# ----------------------------\nclass TestVectorSchema:\n def test_default(self):\n field = VectorSchema(\"vector\", data_type=DataType.VECTOR_FP32, dimension=128)\n assert field.name == \"vector\"\n assert field.data_type == DataType.VECTOR_FP32\n assert field.dimension == 128\n assert field.index_param is not None\n assert field.index_param.type == IndexType.FLAT\n assert field.index_param.metric_type == MetricType.IP\n\n def test_custom(self):\n field = VectorSchema(\n name=\"vector\",\n data_type=DataType.VECTOR_INT8,\n dimension=512,\n index_param=HnswIndexParam(\n metric_type=MetricType.COSINE, m=15, ef_construction=300\n ),\n )\n assert field.name == \"vector\"\n assert field.data_type == DataType.VECTOR_INT8\n assert field.index_param.metric_type == MetricType.COSINE\n assert field.index_param.m == 15\n assert field.index_param.ef_construction == 300\n\n def test_readonly(self):\n field = VectorSchema(\n name=\"vector\",\n dimension=128,\n data_type=DataType.VECTOR_INT8,\n )\n\n import sys\n\n if sys.version_info >= (3, 11):\n match_pattern = r\"(can't set attribute|has no setter|readonly attribute)\"\n else:\n match_pattern = r\"can't set attribute\"\n with pytest.raises(AttributeError, match=match_pattern):\n field.dimension = 4\n\n\n# ----------------------------\n# CollectionSchema Test Case\n# ----------------------------\nclass TestCollectionSchema:\n def test_collection_schema_with_single_field(self):\n collection_schema = CollectionSchema(\n name=\"test_collection\",\n fields=FieldSchema(\n name=\"id\",\n data_type=DataType.INT64,\n index_param=InvertIndexParam(),\n nullable=False,\n ),\n vectors=VectorSchema(\n name=\"vector\",\n data_type=DataType.VECTOR_INT8,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n )\n\n assert collection_schema is not None\n assert collection_schema.name == \"test_collection\"\n assert len(collection_schema.fields) == 1\n assert len(collection_schema.vectors) == 1\n\n field = collection_schema.field(\"id\")\n assert field is not None\n assert field.name == \"id\"\n assert field.data_type == DataType.INT64\n assert not field.nullable\n assert field.index_param.type == IndexType.INVERT\n assert not field.index_param.enable_range_optimization\n\n vector = collection_schema.vector(\"vector\")\n assert vector is not None\n assert vector.name == \"vector\"\n assert vector.data_type == DataType.VECTOR_INT8\n assert vector.dimension == 128\n assert vector.index_param.type == IndexType.HNSW\n assert vector.index_param.m == 50\n assert vector.index_param.ef_construction == 500\n assert vector.index_param.metric_type == MetricType.IP\n\n def test_collection_schema_with_multi_fields(self):\n collection_schema = CollectionSchema(\n name=\"test_collection\",\n fields=[\n FieldSchema(\n \"id\",\n DataType.INT64,\n nullable=False,\n index_param=InvertIndexParam(enable_range_optimization=True),\n ),\n FieldSchema(\n \"name\",\n DataType.STRING,\n nullable=False,\n index_param=InvertIndexParam(),\n ),\n FieldSchema(\n \"weight\",\n DataType.INT32,\n nullable=True,\n ),\n ],\n vectors=[\n VectorSchema(\n \"dense\",\n DataType.VECTOR_FP32,\n dimension=128,\n index_param=HnswIndexParam(),\n ),\n VectorSchema(\n \"sparse\", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()\n ),\n ],\n )\n assert collection_schema is not None\n assert collection_schema.name == \"test_collection\"\n assert len(collection_schema.fields) == 3\n assert len(collection_schema.vectors) == 2\n\n field_id = collection_schema.field(\"id\")\n assert field_id is not None\n assert field_id.name == \"id\"\n assert field_id.data_type == DataType.INT64\n assert not field_id.nullable\n assert field_id.index_param.type == IndexType.INVERT\n\n dense = collection_schema.vector(\"dense\")\n assert dense is not None\n assert dense.name == \"dense\"\n assert dense.data_type == DataType.VECTOR_FP32\n assert dense.dimension == 128\n assert dense.index_param.type == IndexType.HNSW\n\n sparse = collection_schema.vector(\"sparse\")\n assert sparse is not None\n assert sparse.name == \"sparse\"\n assert sparse.data_type == DataType.SPARSE_VECTOR_FP32\n assert sparse.dimension == 0\n assert sparse.index_param.type == IndexType.HNSW\n\n assert str(collection_schema) is not None\n\n\n# ----------------------------\n# CollectionStats Test Case\n# ----------------------------\nclass TestCollectionStats:\n \"\"\"\n The constructor of CollectionStats is not provided.\n It can only be obtained through collection.stats()\n \"\"\"\n\n def test_collection_stats(self):\n stats = CollectionStats()\n assert stats is not None\n" }, { "path": "python/tests/test_typing.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport pytest\nfrom zvec import (\n DataType,\n IndexType,\n MetricType,\n QuantizeType,\n Status,\n StatusCode,\n)\n\n\n# ----------------------------\n# Enum Test Case\n# ----------------------------\n@pytest.mark.parametrize(\n \"member, name\",\n [\n (DataType.FLOAT, \"FLOAT\"),\n (IndexType.HNSW, \"HNSW\"),\n (MetricType.COSINE, \"COSINE\"),\n (QuantizeType.INT8, \"INT8\"),\n (StatusCode.OK, \"OK\"),\n ],\n)\ndef test_enum_names(member, name):\n assert member.name == name\n\n\n@pytest.mark.parametrize(\n \"member, value\",\n [\n (DataType.FLOAT, 8),\n (IndexType.HNSW, 1),\n (MetricType.COSINE, 3),\n (QuantizeType.INT8, 2),\n (StatusCode.OK, 0),\n ],\n)\ndef test_enum_values(member, value):\n assert member.value == value\n\n\n@pytest.mark.parametrize(\"member\", [\"L2\", \"IP\", \"COSINE\"])\ndef test_metric_type_has_member(member):\n assert member in MetricType.__members__\n\n\n@pytest.mark.parametrize(\n \"member\",\n [\n \"STRING\",\n \"BOOL\",\n \"INT32\",\n \"INT64\",\n \"FLOAT\",\n \"DOUBLE\",\n \"UINT32\",\n \"UINT64\",\n \"VECTOR_FP16\",\n \"VECTOR_FP32\",\n \"VECTOR_FP64\",\n \"VECTOR_INT8\",\n \"SPARSE_VECTOR_FP32\",\n \"SPARSE_VECTOR_FP16\",\n \"ARRAY_STRING\",\n \"ARRAY_INT32\",\n \"ARRAY_INT64\",\n \"ARRAY_FLOAT\",\n \"ARRAY_DOUBLE\",\n \"ARRAY_BOOL\",\n \"ARRAY_UINT32\",\n \"ARRAY_UINT64\",\n ],\n)\ndef test_data_type_has_member(member):\n assert member in DataType.__members__\n\n\n@pytest.mark.parametrize(\"member\", [\"HNSW\", \"IVF\", \"FLAT\", \"INVERT\"])\ndef test_index_type_has_member(member):\n assert member in IndexType.__members__\n\n\n@pytest.mark.parametrize(\"member\", [\"FP16\", \"INT8\", \"INT4\", \"UNDEFINED\"])\ndef test_quantize_type_has_member(member):\n assert member in QuantizeType.__members__\n\n\n@pytest.mark.parametrize(\n \"member\",\n [\n \"OK\",\n \"UNKNOWN\",\n \"NOT_FOUND\",\n \"ALREADY_EXISTS\",\n \"INVALID_ARGUMENT\",\n \"PERMISSION_DENIED\",\n \"FAILED_PRECONDITION\",\n \"RESOURCE_EXHAUSTED\",\n \"UNAVAILABLE\",\n \"INTERNAL_ERROR\",\n \"NOT_SUPPORTED\",\n ],\n)\ndef test_status_code_has_member(member):\n assert member in StatusCode.__members__\n\n\n# ----------------------------\n# Status Test Case\n# ----------------------------\nclass TestStatus:\n def test_status_code(self):\n status = Status(StatusCode.OK)\n assert status.code() == StatusCode.OK\n\n def test_status_message(self):\n status = Status(StatusCode.OK, \"OK\")\n assert status.message() == \"OK\"\n\n status = Status(StatusCode.NOT_FOUND, \"Not Found\")\n assert status.message() == \"Not Found\"\n\n def test_status_ok(self):\n status = Status(StatusCode.OK)\n assert status.ok()\n" }, { "path": "python/tests/test_util.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom unittest.mock import MagicMock, patch\n\nimport pytest\nfrom zvec import require_module\n\n\n# ----------------------------\n# require_module func Test Case\n# ----------------------------\ndef test_require_module_success():\n module = require_module(\"os\")\n assert module is not None\n assert hasattr(module, \"path\")\n\n\ndef test_require_module_with_submodule_success():\n module = require_module(\"os.path\")\n assert module is not None\n assert hasattr(module, \"join\")\n\n\ndef test_require_module_import_error():\n with pytest.raises(ImportError) as exc_info:\n require_module(\"nonexistent_module\")\n\n exception_msg = str(exc_info.value)\n assert \"Required package 'nonexistent_module' is not installed.\" in exception_msg\n\n\ndef test_require_module_with_mitigation_import_error():\n with pytest.raises(ImportError) as exc_info:\n require_module(\"nonexistent_module.submodule\", mitigation=\"custom_package\")\n\n exception_msg = str(exc_info.value)\n assert \"Required package 'custom_package' is not installed.\" in exception_msg\n assert (\n \"Module 'nonexistent_module.submodule' is part of 'nonexistent_module'\"\n in exception_msg\n )\n assert \"please pip install 'custom_package'.\" in exception_msg\n\n\ndef test_require_module_submodule_import_error():\n with pytest.raises(ImportError) as exc_info:\n require_module(\"os.nonexistent_submodule\")\n\n exception_msg = str(exc_info.value)\n assert (\n \"Required package 'os.nonexistent_submodule' is not installed.\" in exception_msg\n )\n assert \"Module 'os.nonexistent_submodule' is part of 'os'\" in exception_msg\n assert \"please pip install 'os'.\" in exception_msg\n\n\n@patch(\"importlib.import_module\")\ndef test_require_module_wraps_original_exception(mock_import_module):\n original_exception = ImportError(\"Original error\")\n mock_import_module.side_effect = original_exception\n\n with pytest.raises(ImportError) as exc_info:\n require_module(\"some_module\")\n\n assert exc_info.value.__cause__ is original_exception\n\n\n@patch(\"importlib.import_module\")\ndef test_require_module_calls_importlib(mock_import_module):\n mock_module = MagicMock()\n mock_import_module.return_value = mock_module\n\n result = require_module(\"test_module\")\n\n mock_import_module.assert_called_once_with(\"test_module\")\n assert result is mock_module\n" }, { "path": "python/zvec/__init__.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import annotations\n\nimport sys\nfrom typing import TYPE_CHECKING\n\nif TYPE_CHECKING:\n from importlib.metadata import PackageNotFoundError\n\n\n# ==============================\n# Public API — grouped by category\n# ==============================\n\nfrom . import model as model\n\n# —— Extensions ——\nfrom .extension import (\n BM25EmbeddingFunction,\n DefaultLocalDenseEmbedding,\n DefaultLocalReRanker,\n DefaultLocalSparseEmbedding,\n DenseEmbeddingFunction,\n OpenAIDenseEmbedding,\n OpenAIFunctionBase,\n QwenDenseEmbedding,\n QwenFunctionBase,\n QwenReRanker,\n QwenSparseEmbedding,\n ReRanker,\n RrfReRanker,\n SentenceTransformerFunctionBase,\n SparseEmbeddingFunction,\n WeightedReRanker,\n)\n\n# —— Typing ——\nfrom .model import param as param\nfrom .model import schema as schema\n\n# —— Core data structures ——\nfrom .model.collection import Collection\nfrom .model.doc import Doc\n\n# —— Query & index parameters ——\nfrom .model.param import (\n AddColumnOption,\n AlterColumnOption,\n CollectionOption,\n FlatIndexParam,\n HnswIndexParam,\n HnswQueryParam,\n HnswRabitqIndexParam,\n HnswRabitqQueryParam,\n IndexOption,\n InvertIndexParam,\n IVFIndexParam,\n IVFQueryParam,\n OptimizeOption,\n)\nfrom .model.param.vector_query import VectorQuery\n\n# —— Schema & field definitions ——\nfrom .model.schema import CollectionSchema, CollectionStats, FieldSchema, VectorSchema\n\n# —— tools ——\nfrom .tool import require_module\nfrom .typing import (\n DataType,\n IndexType,\n MetricType,\n QuantizeType,\n Status,\n StatusCode,\n)\nfrom .typing.enum import LogLevel, LogType\n\n# —— lifecycle ——\nfrom .zvec import create_and_open, init, open\n\n# ==============================\n# Public interface declaration\n# ==============================\n__all__ = [\n # Zvec functions\n \"create_and_open\",\n \"init\",\n \"open\",\n # Core classes\n \"Collection\",\n \"Doc\",\n # Schema\n \"CollectionSchema\",\n \"FieldSchema\",\n \"VectorSchema\",\n \"CollectionStats\",\n # Parameters\n \"VectorQuery\",\n \"InvertIndexParam\",\n \"HnswIndexParam\",\n \"HnswRabitqIndexParam\",\n \"HnswRabitqQueryParam\",\n \"FlatIndexParam\",\n \"IVFIndexParam\",\n \"CollectionOption\",\n \"IndexOption\",\n \"OptimizeOption\",\n \"AddColumnOption\",\n \"AlterColumnOption\",\n \"HnswQueryParam\",\n \"IVFQueryParam\",\n # Extensions\n \"DenseEmbeddingFunction\",\n \"SparseEmbeddingFunction\",\n \"QwenFunctionBase\",\n \"OpenAIFunctionBase\",\n \"SentenceTransformerFunctionBase\",\n \"ReRanker\",\n \"DefaultLocalDenseEmbedding\",\n \"DefaultLocalSparseEmbedding\",\n \"BM25EmbeddingFunction\",\n \"OpenAIDenseEmbedding\",\n \"QwenDenseEmbedding\",\n \"QwenSparseEmbedding\",\n \"RrfReRanker\",\n \"WeightedReRanker\",\n \"DefaultLocalReRanker\",\n \"QwenReRanker\",\n # Typing\n \"DataType\",\n \"MetricType\",\n \"QuantizeType\",\n \"IndexType\",\n \"LogLevel\",\n \"LogType\",\n \"Status\",\n \"StatusCode\",\n # Tools\n \"require_module\",\n]\n\n# ==============================\n# Version handling\n# ==============================\n__version__: str\n\ntry:\n from importlib.metadata import version\nexcept ImportError:\n from importlib_metadata import version # Python < 3.8\n\ntry:\n __version__ = version(\"zvec\")\nexcept Exception:\n __version__ = \"unknown\"\n" }, { "path": "python/zvec/__init__.pyi", "content": "\"\"\"\nZvec core module\n\"\"\"\n\nfrom __future__ import annotations\n\nimport collections\n\nfrom . import typing\nfrom .extension import ReRanker, RrfReRanker, WeightedReRanker\nfrom .extension.embedding import DenseEmbeddingFunction\nfrom .model import param, schema\nfrom .model.collection import Collection\nfrom .model.doc import Doc\nfrom .model.param import (\n AddColumnOption,\n AlterColumnOption,\n CollectionOption,\n FlatIndexParam,\n HnswIndexParam,\n HnswQueryParam,\n IndexOption,\n InvertIndexParam,\n IVFIndexParam,\n IVFQueryParam,\n OptimizeOption,\n)\nfrom .model.param.vector_query import VectorQuery\nfrom .model.schema import CollectionSchema, CollectionStats, FieldSchema, VectorSchema\nfrom .tool import require_module\nfrom .typing import (\n DataType,\n IndexType,\n MetricType,\n QuantizeType,\n Status,\n StatusCode,\n)\nfrom .typing.enum import LogLevel, LogType\nfrom .zvec import create_and_open, init, open\n\n__all__: list = [\n \"AddColumnOption\",\n \"AlterColumnOption\",\n \"Collection\",\n \"CollectionOption\",\n \"CollectionSchema\",\n \"CollectionStats\",\n \"DataType\",\n \"DenseEmbeddingFunction\",\n \"DenseEmbeddingFunction\",\n \"Doc\",\n \"FieldSchema\",\n \"FlatIndexParam\",\n \"HnswIndexParam\",\n \"HnswQueryParam\",\n \"IVFIndexParam\",\n \"IVFQueryParam\",\n \"IndexOption\",\n \"IndexType\",\n \"InvertIndexParam\",\n \"LogLevel\",\n \"LogType\",\n \"MetricType\",\n \"OptimizeOption\",\n \"QuantizeType\",\n \"ReRanker\",\n \"ReRanker\",\n \"RrfReRanker\",\n \"Status\",\n \"StatusCode\",\n \"VectorQuery\",\n \"VectorSchema\",\n \"WeightedReRanker\",\n \"create_and_open\",\n \"init\",\n \"open\",\n \"require_module\",\n]\n\nclass _Collection:\n @staticmethod\n def CreateAndOpen(\n arg0: str, arg1: schema._CollectionSchema, arg2: param.CollectionOption\n ) -> _Collection: ...\n @staticmethod\n def Open(arg0: str, arg1: param.CollectionOption) -> _Collection: ...\n def AddColumn(\n self,\n arg0: schema._FieldSchema,\n arg1: str,\n arg2: param.AddColumnOption,\n ) -> None: ...\n def AlterColumn(\n self,\n arg0: str,\n arg1: str,\n arg2: schema._FieldSchema,\n arg3: param.AlterColumnOption,\n ) -> None: ...\n def CreateIndex(\n self, arg0: str, arg1: param.IndexParam, arg2: param.IndexOption\n ) -> None: ...\n def Delete(self, arg0: collections.abc.Sequence[str]) -> list[typing.Status]: ...\n def DeleteByFilter(self, arg0: str) -> None: ...\n def Destroy(self) -> None: ...\n def DropColumn(self, arg0: str) -> None: ...\n def DropIndex(self, arg0: str) -> None: ...\n def Fetch(self, arg0: collections.abc.Sequence[str]) -> dict[str, _Doc]: ...\n def Flush(self) -> None: ...\n def GroupByQuery(self, arg0: ...) -> list[...]: ...\n def Insert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...\n def Optimize(self, arg0: param.OptimizeOption) -> None: ...\n def Options(self) -> param.CollectionOption: ...\n def Path(self) -> str: ...\n def Query(self, arg0: param._VectorQuery) -> list[_Doc]: ...\n def Schema(self) -> schema._CollectionSchema: ...\n def Stats(self) -> schema.CollectionStats: ...\n def Update(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...\n def Upsert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...\n def __getstate__(self) -> tuple: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n\nclass _Doc:\n def __getstate__(self) -> bytes: ...\n def __init__(self) -> None: ...\n def __setstate__(self, arg0: bytes) -> None: ...\n def field_names(self) -> list[str]: ...\n def get_any(self, arg0: str, arg1: typing.DataType) -> typing.Any: ...\n def has_field(self, arg0: str) -> bool: ...\n def pk(self) -> str: ...\n def score(self) -> float: ...\n def set_any(self, arg0: str, arg1: typing.DataType, arg2: typing.Any) -> bool: ...\n def set_pk(self, arg0: str) -> None: ...\n def set_score(self, arg0: typing.SupportsFloat) -> None: ...\n\nclass _DocOp:\n \"\"\"\n Members:\n\n INSERT\n\n UPDATE\n\n DELETE\n\n UPSERT\n \"\"\"\n\n DELETE: typing.ClassVar[_DocOp] # value = <_DocOp.DELETE: 3>\n INSERT: typing.ClassVar[_DocOp] # value = <_DocOp.INSERT: 0>\n UPDATE: typing.ClassVar[_DocOp] # value = <_DocOp.UPDATE: 2>\n UPSERT: typing.ClassVar[_DocOp] # value = <_DocOp.UPSERT: 1>\n __members__: typing.ClassVar[\n dict[str, _DocOp]\n ] # value = {'INSERT': <_DocOp.INSERT: 0>, 'UPDATE': <_DocOp.UPDATE: 2>, 'DELETE': <_DocOp.DELETE: 3>, 'UPSERT': <_DocOp.UPSERT: 1>}\n\n def __eq__(self, other: typing.Any) -> bool: ...\n def __getstate__(self) -> int: ...\n def __hash__(self) -> int: ...\n def __index__(self) -> int: ...\n def __init__(self, value: typing.SupportsInt) -> None: ...\n def __int__(self) -> int: ...\n def __ne__(self, other: typing.Any) -> bool: ...\n def __repr__(self) -> str: ...\n def __setstate__(self, state: typing.SupportsInt) -> None: ...\n def __str__(self) -> str: ...\n @property\n def name(self) -> str: ...\n @property\n def value(self) -> int: ...\n" }, { "path": "python/zvec/common/__init__.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom .constants import DenseVectorType, SparseVectorType, VectorType\n\n__all__ = [\"DenseVectorType\", \"SparseVectorType\", \"VectorType\"]\n" }, { "path": "python/zvec/common/constants.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom typing import Optional, TypeVar, Union\n\nimport numpy as np\n\n# VectorType: DenseVectorType | SparseVectorType\nDenseVectorType = Union[list[float], list[int], np.ndarray]\nSparseVectorType = dict[int, float]\nVectorType = Optional[Union[DenseVectorType, SparseVectorType]]\n\n# Embeddable: Text | Image | Audio\nTEXT = str\nIMAGE = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array\nAUDIO = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array\n\nEmbeddable = Optional[Union[TEXT, IMAGE, AUDIO]]\n\n# Multimodal Embeddable\nMD = TypeVar(\"MD\", bound=Embeddable, contravariant=True)\n" }, { "path": "python/zvec/executor/__init__.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom .query_executor import (\n QueryContext,\n QueryExecutor,\n QueryExecutorFactory,\n)\n\n__all__ = [\n \"QueryContext\",\n \"QueryExecutor\",\n \"QueryExecutorFactory\",\n]\n" }, { "path": "python/zvec/executor/query_executor.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport os\nfrom abc import ABC, abstractmethod\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom typing import Optional, Union, final\n\nimport numpy as np\nfrom _zvec import _Collection\nfrom _zvec.param import _VectorQuery\n\nfrom ..extension import ReRanker, RrfReRanker, WeightedReRanker\nfrom ..model.convert import convert_to_py_doc\nfrom ..model.doc import Doc\nfrom ..model.param.vector_query import VectorQuery\nfrom ..model.schema import CollectionSchema\nfrom ..typing import DataType\n\n__all__ = [\n \"QueryContext\",\n \"QueryExecutor\",\n \"QueryExecutorFactory\",\n]\n\nDTYPE_MAP = {\n DataType.VECTOR_FP16.value: np.float16,\n DataType.VECTOR_FP32.value: np.float32,\n DataType.VECTOR_FP64.value: np.float64,\n DataType.VECTOR_INT8.value: np.int8,\n}\n\n\ndef convert_to_numpy(vec: Union[list, np.ndarray], dtype: np.dtype) -> np.ndarray:\n if isinstance(vec, np.ndarray):\n if vec.dtype == dtype and vec.ndim == 1:\n return vec\n return np.asarray(vec, dtype=dtype).flatten()\n\n try:\n arr = np.asarray(vec, dtype=dtype)\n if arr.ndim != 1:\n arr = arr.flatten()\n return arr\n except (ValueError, TypeError) as e:\n raise TypeError(\n f\"Cannot convert input to 1D numpy array with dtype={dtype}: {type(vec)}\"\n ) from e\n\n\nclass QueryContext:\n def __init__(\n self,\n topk: int,\n filter: Optional[str] = None,\n include_vector: bool = False,\n queries: Optional[list[VectorQuery]] = None,\n output_fields: Optional[list[str]] = None,\n reranker: Optional[ReRanker] = None,\n ):\n # query param\n self._filter = filter\n self._queries = queries or []\n self._topk = topk\n self._include_vector = include_vector\n self._output_fields = output_fields\n\n # reranker\n self._reranker = reranker\n\n # core vectors\n self._core_vectors = []\n\n @property\n def topk(self):\n return self._topk\n\n @property\n def queries(self):\n return self._queries\n\n @property\n def filter(self):\n return self._filter\n\n @property\n def reranker(self):\n return self._reranker\n\n @property\n def output_fields(self):\n return self._output_fields\n\n @property\n def include_vector(self):\n return self._include_vector\n\n @property\n def core_vectors(self):\n return self._core_vectors\n\n @core_vectors.setter\n def core_vectors(self, core_vectors: list[_VectorQuery]):\n self._core_vectors = core_vectors\n\n\nclass QueryExecutor(ABC):\n def __init__(self, schema: CollectionSchema):\n self._schema = schema\n self._concurrency = max(1, int(os.getenv(\"ZVEC_QUERY_CONCURRENCY\", \"1\")))\n\n @abstractmethod\n def _do_validate(self, ctx: QueryContext) -> None:\n pass\n\n @abstractmethod\n def _do_build(\n self, ctx: QueryContext, collection: _Collection\n ) -> list[_VectorQuery]:\n pass\n\n def _do_build_query_wo_vector(self, ctx: QueryContext) -> _VectorQuery:\n core_vector = _VectorQuery()\n core_vector.topk = ctx.topk\n core_vector.include_vector = ctx.include_vector\n if ctx.filter:\n core_vector.filter = ctx.filter\n if ctx.output_fields:\n core_vector.output_fields = ctx.output_fields\n return core_vector\n\n def _do_build_query_with_vector(\n self, ctx: QueryContext, query: VectorQuery, collection: _Collection\n ) -> _VectorQuery:\n core_vector = self._do_build_query_wo_vector(ctx)\n core_vector.field_name = query.field_name\n if query.param:\n core_vector.query_params = query.param\n\n vector_schema = (\n self._schema.vector(query.field_name) if query else self._schema.vectors[0]\n )\n\n if vector_schema is None:\n raise ValueError(\"No vector field found\")\n\n # set output_fields\n core_vector.output_fields = ctx.output_fields\n\n # set vector\n if query.has_vector():\n vec_data = query.vector\n else:\n fetched = collection.Fetch([query.id])\n doc = next(iter(fetched.values()))\n if not doc:\n return core_vector\n vec_data = doc.get_any(vector_schema.name, vector_schema.data_type)\n\n target_dtype = DTYPE_MAP.get(vector_schema.data_type.value)\n core_vector.set_vector(\n vector_schema._get_object(),\n convert_to_numpy(vec_data, target_dtype) if target_dtype else vec_data,\n )\n return core_vector\n\n def _do_execute(\n self, vectors: list[_VectorQuery], collection: _Collection\n ) -> dict[str, list[Doc]]:\n query_cnt = len(vectors)\n if query_cnt == 0:\n raise ValueError(\"No query to execute\")\n\n if len(vectors) == 1 or self._concurrency == 1:\n results = {}\n for query in vectors:\n docs = collection.Query(query)\n results[query.field_name] = [\n convert_to_py_doc(doc, self._schema) for doc in docs\n ]\n return results\n\n results = {}\n with ThreadPoolExecutor(max_workers=self._concurrency) as executor:\n future_to_query = {\n executor.submit(collection.Query, query): query.field_name\n for query in vectors\n }\n\n for future in as_completed(future_to_query):\n field_name = future_to_query[future]\n try:\n docs = future.result()\n results[field_name] = [\n convert_to_py_doc(doc, self._schema) for doc in docs\n ]\n except Exception as e:\n raise e\n return results\n\n def _do_merge_rerank_results(\n self, ctx: QueryContext, docs_map: dict[str, list[Doc]]\n ) -> list[Doc]:\n query_result_cnt = len(docs_map) if docs_map else 0\n if query_result_cnt == 0:\n raise ValueError(\"Query results is none and dost not to rerank\")\n if query_result_cnt == 1:\n if not ctx.reranker or isinstance(\n ctx.reranker, (RrfReRanker, WeightedReRanker)\n ):\n return next(iter(docs_map.values()))\n return ctx.reranker.rerank(docs_map)\n return ctx.reranker.rerank(docs_map)\n\n @final\n def execute(self, ctx: QueryContext, collection: _Collection) -> list[Doc]:\n # 1. validate query\n self._do_validate(ctx)\n # 2. build query vector\n query_vectors = self._do_build(ctx, collection)\n if not query_vectors:\n raise ValueError(\"No query to execute\")\n # 3. execute query\n docs = self._do_execute(query_vectors, collection)\n # 4. merge and rerank result\n return self._do_merge_rerank_results(ctx, docs)\n\n\nclass NoVectorQueryExecutor(QueryExecutor):\n def __init__(self, schema: CollectionSchema):\n super().__init__(schema)\n\n def _do_validate(self, ctx: QueryContext) -> None:\n if len(ctx.queries) > 0:\n raise ValueError(\"Collection does not support query with vector or id\")\n\n def _do_build(\n self, ctx: QueryContext, _collection: _Collection\n ) -> list[_VectorQuery]:\n return [self._do_build_query_wo_vector(ctx)]\n\n\nclass SingleVectorQueryExecutor(NoVectorQueryExecutor):\n def __init__(self, schema: CollectionSchema) -> None:\n super().__init__(schema)\n\n def _do_validate(self, ctx: QueryContext) -> None:\n if len(ctx.queries) > 1:\n raise ValueError(\n \"Collection has only one vector field, cannot query with multiple vectors\"\n )\n for query in ctx.queries:\n query._validate()\n\n def _do_build(\n self, ctx: QueryContext, collection: _Collection\n ) -> list[_VectorQuery]:\n if len(ctx.queries) == 0:\n return [self._do_build_query_wo_vector(ctx)]\n vectors = []\n for query in ctx.queries:\n vectors.append(self._do_build_query_with_vector(ctx, query, collection))\n return vectors\n\n\nclass MultiVectorQueryExecutor(SingleVectorQueryExecutor):\n def __init__(self, schema: CollectionSchema) -> None:\n super().__init__(schema)\n\n def _do_validate(self, ctx: QueryContext) -> None:\n if len(ctx.queries) > 1 and ctx.reranker is None:\n raise ValueError(\"Reranker is required for multi-vector query\")\n seen_fields = set()\n for query in ctx.queries:\n query._validate()\n field = query.field_name\n if field in seen_fields:\n raise ValueError(f\"Query field name '{field}' appears more than once\")\n seen_fields.add(field)\n\n def _do_execute(\n self, vectors: list[_VectorQuery], collection: _Collection\n ) -> dict[str, list[Doc]]:\n return super()._do_execute(vectors, collection)\n\n\nclass QueryExecutorFactory:\n @staticmethod\n def create(schema: CollectionSchema) -> QueryExecutor:\n vectors = schema.vectors\n if len(vectors) == 0:\n return NoVectorQueryExecutor(schema)\n if len(vectors) == 1:\n return SingleVectorQueryExecutor(schema)\n return MultiVectorQueryExecutor(schema)\n" }, { "path": "python/zvec/extension/__init__.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom .bm25_embedding_function import BM25EmbeddingFunction\nfrom .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction\nfrom .http_embedding_function import HTTPDenseEmbedding\nfrom .jina_embedding_function import JinaDenseEmbedding\nfrom .jina_function import JinaFunctionBase\nfrom .multi_vector_reranker import RrfReRanker, WeightedReRanker\nfrom .openai_embedding_function import OpenAIDenseEmbedding\nfrom .openai_function import OpenAIFunctionBase\nfrom .qwen_embedding_function import QwenDenseEmbedding, QwenSparseEmbedding\nfrom .qwen_function import QwenFunctionBase\nfrom .qwen_rerank_function import QwenReRanker\nfrom .rerank_function import RerankFunction as ReRanker\nfrom .sentence_transformer_embedding_function import (\n DefaultLocalDenseEmbedding,\n DefaultLocalSparseEmbedding,\n)\nfrom .sentence_transformer_function import SentenceTransformerFunctionBase\nfrom .sentence_transformer_rerank_function import DefaultLocalReRanker\n\n__all__ = [\n \"BM25EmbeddingFunction\",\n \"DefaultLocalDenseEmbedding\",\n \"DefaultLocalReRanker\",\n \"DefaultLocalSparseEmbedding\",\n \"DenseEmbeddingFunction\",\n \"HTTPDenseEmbedding\",\n \"JinaDenseEmbedding\",\n \"JinaFunctionBase\",\n \"OpenAIDenseEmbedding\",\n \"OpenAIFunctionBase\",\n \"QwenDenseEmbedding\",\n \"QwenFunctionBase\",\n \"QwenReRanker\",\n \"QwenSparseEmbedding\",\n \"ReRanker\",\n \"RrfReRanker\",\n \"SentenceTransformerFunctionBase\",\n \"SparseEmbeddingFunction\",\n \"WeightedReRanker\",\n]\n" }, { "path": "python/zvec/extension/bm25_embedding_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom functools import lru_cache\nfrom typing import Literal, Optional\n\nfrom ..common.constants import TEXT, SparseVectorType\nfrom ..tool import require_module\nfrom .embedding_function import SparseEmbeddingFunction\n\n\nclass BM25EmbeddingFunction(SparseEmbeddingFunction[TEXT]):\n \"\"\"BM25-based sparse embedding function using DashText SDK.\n\n This class provides text-to-sparse-vector embedding capabilities using\n the DashText library with BM25 algorithm. BM25 (Best Matching 25) is a\n probabilistic retrieval function used for lexical search and document\n ranking based on term frequency and inverse document frequency.\n\n BM25 generates sparse vectors where each dimension corresponds to a term in\n the vocabulary, and the value represents the BM25 score for that term. It's\n particularly effective for:\n\n - Lexical search and keyword matching\n - Document ranking and information retrieval\n - Combining with dense embeddings for hybrid search\n - Traditional IR tasks where exact term matching is important\n\n This implementation uses DashText's SparseVectorEncoder, which provides\n efficient BM25 computation for Chinese and English text using either a\n built-in encoder or custom corpus training.\n\n Args:\n corpus (Optional[list[str]], optional): List of documents to train the\n BM25 encoder. If provided, creates a custom encoder trained on this\n corpus for better domain-specific accuracy. If ``None``, uses the\n built-in encoder. Defaults to ``None``.\n encoding_type (Literal[\"query\", \"document\"], optional): Encoding mode\n for text processing. Use ``\"query\"`` for search queries (default) and\n ``\"document\"`` for document indexing. This distinction optimizes the\n BM25 scoring for asymmetric retrieval tasks. Defaults to ``\"query\"``.\n language (Literal[\"zh\", \"en\"], optional): Language for built-in encoder.\n Only used when corpus is None. ``\"zh\"`` for Chinese (trained on Chinese\n Wikipedia), ``\"en\"`` for English. Defaults to ``\"zh\"``.\n b (float, optional): Document length normalization parameter for BM25.\n Range [0, 1]. 0 means no normalization, 1 means full normalization.\n Only used with custom corpus. Defaults to ``0.75``.\n k1 (float, optional): Term frequency saturation parameter for BM25.\n Higher values give more weight to term frequency. Only used with\n custom corpus. Defaults to ``1.2``.\n **kwargs: Additional parameters for DashText encoder customization.\n\n Attributes:\n corpus_size (int): Number of documents in the training corpus (0 if using built-in encoder).\n encoding_type (str): The encoding type being used (\"query\" or \"document\").\n language (str): The language of the built-in encoder (\"zh\" or \"en\").\n\n Raises:\n ValueError: If corpus is provided but empty or contains non-string elements.\n TypeError: If input to ``embed()`` is not a string.\n RuntimeError: If DashText encoder initialization or training fails.\n\n Note:\n - Requires Python 3.10, 3.11, or 3.12\n - Requires the ``dashtext`` package: ``pip install dashtext``\n - Two encoder options available:\n\n 1. **Built-in encoder** (no corpus needed): Pre-trained models for\n Chinese (zh) and English (en), good generalization, works out-of-the-box\n 2. **Custom encoder** (corpus required): Better accuracy for domain-specific\n terminology, requires training on your full corpus with BM25 parameters\n\n - Encoding types:\n\n * ``encoding_type=\"query\"``: Optimized for search queries (shorter text)\n * ``encoding_type=\"document\"``: Optimized for document indexing (longer text)\n\n - BM25 parameters (b, k1) only apply to custom encoder training\n - Output is sorted by indices (vocabulary term IDs) for consistency\n - Results are cached (LRU cache, maxsize=10) to reduce computation\n - No API key or network connectivity required (local computation)\n\n Examples:\n >>> # Option 1: Using built-in encoder for Chinese (no corpus needed)\n >>> from zvec.extension import BM25EmbeddingFunction\n >>>\n >>> # For query encoding (Chinese)\n >>> bm25_query_zh = BM25EmbeddingFunction(language=\"zh\", encoding_type=\"query\")\n >>> query_vec = bm25_query_zh.embed(\"什么是机器学习\")\n >>> isinstance(query_vec, dict)\n True\n >>> # query_vec: {1169440797: 0.29, 2045788977: 0.70, ...}\n\n >>> # For document encoding (Chinese)\n >>> bm25_doc_zh = BM25EmbeddingFunction(language=\"zh\", encoding_type=\"document\")\n >>> doc_vec = bm25_doc_zh.embed(\"机器学习是人工智能的一个重要分支...\")\n >>> isinstance(doc_vec, dict)\n True\n\n >>> # Using built-in encoder for English\n >>> bm25_query_en = BM25EmbeddingFunction(language=\"en\", encoding_type=\"query\")\n >>> query_vec_en = bm25_query_en.embed(\"what is vector search service\")\n >>> isinstance(query_vec_en, dict)\n True\n\n >>> # Option 2: Using custom corpus for domain-specific accuracy\n >>> corpus = [\n ... \"机器学习是人工智能的一个重要分支\",\n ... \"深度学习使用多层神经网络进行特征提取\",\n ... \"自然语言处理技术用于理解和生成人类语言\"\n ... ]\n >>> bm25_custom = BM25EmbeddingFunction(\n ... corpus=corpus,\n ... encoding_type=\"query\",\n ... b=0.75,\n ... k1=1.2\n ... )\n >>> custom_vec = bm25_custom.embed(\"机器学习算法\")\n >>> isinstance(custom_vec, dict)\n True\n\n >>> # Hybrid search: combining with dense embeddings\n >>> from zvec.extension import DefaultLocalDenseEmbedding\n >>> dense_emb = DefaultLocalDenseEmbedding()\n >>> bm25_emb = BM25EmbeddingFunction(language=\"zh\", encoding_type=\"query\")\n >>>\n >>> query = \"machine learning algorithms\"\n >>> dense_vec = dense_emb.embed(query) # Semantic similarity\n >>> sparse_vec = bm25_emb.embed(query) # Lexical matching\n >>> # Combine scores for hybrid retrieval\n\n >>> # Callable interface\n >>> sparse_vec = bm25_query_zh(\"information retrieval\")\n >>> isinstance(sparse_vec, dict)\n True\n\n >>> # Error handling\n >>> try:\n ... bm25_query_zh.embed(\"\") # Empty query\n ... except ValueError as e:\n ... print(f\"Error: {e}\")\n Error: Input text cannot be empty or whitespace only\n\n See Also:\n - ``SparseEmbeddingFunction``: Base class for sparse embeddings\n - ``DefaultLocalSparseEmbedding``: SPLADE-based sparse embedding\n - ``QwenSparseEmbedding``: API-based sparse embedding using Qwen\n - ``DefaultLocalDenseEmbedding``: Dense embedding for semantic search\n\n References:\n - DashText Documentation: https://help.aliyun.com/zh/document_detail/2546039.html\n - DashText PyPI: https://pypi.org/project/dashtext/\n - BM25 Algorithm: Robertson & Zaragoza (2009)\n \"\"\"\n\n def __init__(\n self,\n corpus: Optional[list[str]] = None,\n encoding_type: Literal[\"query\", \"document\"] = \"query\",\n language: Literal[\"zh\", \"en\"] = \"zh\",\n b: float = 0.75,\n k1: float = 1.2,\n **kwargs,\n ):\n \"\"\"Initialize the BM25 embedding function.\n\n Args:\n corpus (Optional[list[str]]): Optional corpus for training custom encoder.\n If None, uses built-in encoder. Defaults to None.\n encoding_type (Literal[\"query\", \"document\"]): Text encoding mode.\n Use \"query\" for search queries, \"document\" for indexing.\n Defaults to \"query\".\n language (Literal[\"zh\", \"en\"]): Language for built-in encoder.\n \"zh\" for Chinese, \"en\" for English. Defaults to \"zh\".\n b (float): Document length normalization for BM25 [0, 1].\n Only used with custom corpus. Defaults to 0.75.\n k1 (float): Term frequency saturation for BM25.\n Only used with custom corpus. Defaults to 1.2.\n **kwargs: Additional DashText encoder parameters.\n\n Raises:\n ValueError: If corpus is provided but empty or invalid.\n ImportError: If dashtext package is not installed.\n RuntimeError: If encoder initialization or training fails.\n \"\"\"\n # Validate corpus if provided\n if corpus is not None:\n if not corpus or not isinstance(corpus, list):\n raise ValueError(\"Corpus must be a non-empty list of strings\")\n\n if not all(isinstance(doc, str) for doc in corpus):\n raise ValueError(\"All corpus documents must be strings\")\n\n # Import dashtext\n self._dashtext = require_module(\"dashtext\")\n\n self._corpus = corpus\n self._encoding_type = encoding_type\n self._language = language\n self._b = b\n self._k1 = k1\n self._extra_params = kwargs\n\n # Initialize the BM25 encoder\n self._build_encoder()\n\n def _build_encoder(self):\n \"\"\"Build the BM25 sparse vector encoder.\n\n Creates either a built-in encoder (pre-trained) or a custom encoder\n trained on the provided corpus.\n\n Raises:\n RuntimeError: If encoder initialization or training fails.\n ImportError: If dashtext package is not installed.\n \"\"\"\n try:\n if self._corpus is None:\n # Use built-in encoder (pre-trained on Wikipedia)\n # language: 'zh' for Chinese, 'en' for English\n self._encoder = self._dashtext.SparseVectorEncoder.default(\n name=self._language\n )\n else:\n # Create custom encoder with BM25 parameters\n self._encoder = self._dashtext.SparseVectorEncoder(\n b=self._b, k1=self._k1, **self._extra_params\n )\n\n # Train encoder with the corpus\n self._encoder.train(self._corpus)\n\n except ImportError as e:\n raise ImportError(\n \"dashtext package is required for BM25EmbeddingFunction. \"\n \"Install it with: pip install dashtext\"\n ) from e\n except Exception as e:\n if isinstance(e, (ValueError, RuntimeError)):\n raise\n raise RuntimeError(f\"Failed to build BM25 encoder: {e!s}\") from e\n\n @property\n def corpus_size(self) -> int:\n \"\"\"int: Number of documents in the training corpus (0 if using built-in encoder).\"\"\"\n return len(self._corpus) if self._corpus is not None else 0\n\n @property\n def encoding_type(self) -> str:\n \"\"\"str: The encoding type being used (\"query\" or \"document\").\"\"\"\n return self._encoding_type\n\n @property\n def language(self) -> str:\n \"\"\"str: The language of the built-in encoder (\"zh\" or \"en\").\"\"\"\n return self._language\n\n @property\n def extra_params(self) -> dict:\n \"\"\"dict: Extra parameters for DashText encoder customization.\"\"\"\n return self._extra_params\n\n def __call__(self, input: TEXT) -> SparseVectorType:\n \"\"\"Make the embedding function callable.\n\n Args:\n input (TEXT): Input text to embed.\n\n Returns:\n SparseVectorType: Sparse vector as dictionary.\n \"\"\"\n return self.embed(input)\n\n @lru_cache(maxsize=10)\n def embed(self, input: TEXT) -> SparseVectorType:\n \"\"\"Generate BM25 sparse embedding for the input text.\n\n This method computes BM25 scores for the input text using DashText's\n SparseVectorEncoder. The encoding behavior depends on the encoding_type:\n\n - ``encoding_type=\"query\"``: Uses ``encode_queries()`` for search queries\n - ``encoding_type=\"document\"``: Uses ``encode_documents()`` for documents\n\n The result is a sparse vector where keys are term indices in the\n vocabulary and values are BM25 scores.\n\n Args:\n input (TEXT): Input text string to embed. Must be non-empty after\n stripping whitespace.\n\n Returns:\n SparseVectorType: A dictionary mapping vocabulary term index to BM25 score.\n Only non-zero scores are included. The dictionary is sorted by indices\n (keys) in ascending order for consistent output.\n Example: ``{1169440797: 0.29, 2045788977: 0.70, ...}``\n\n Raises:\n TypeError: If ``input`` is not a string.\n ValueError: If input is empty or whitespace-only.\n RuntimeError: If BM25 encoding fails.\n\n Examples:\n >>> bm25 = BM25EmbeddingFunction(language=\"zh\", encoding_type=\"query\")\n >>> sparse_vec = bm25.embed(\"query text\")\n >>> isinstance(sparse_vec, dict)\n True\n >>> all(isinstance(k, int) and isinstance(v, float) for k, v in sparse_vec.items())\n True\n\n >>> # Verify sorted output\n >>> keys = list(sparse_vec.keys())\n >>> keys == sorted(keys)\n True\n\n >>> # Error: empty input\n >>> bm25.embed(\" \")\n ValueError: Input text cannot be empty or whitespace only\n\n >>> # Error: non-string input\n >>> bm25.embed(123)\n TypeError: Expected 'input' to be str, got int\n\n Note:\n - BM25 scores are relative to the vocabulary statistics\n - Output dictionary is always sorted by indices for consistency\n - Terms not in the vocabulary will have zero scores (not included)\n - This method is cached (maxsize=10) for performance\n - DashText automatically handles Chinese/English text segmentation\n \"\"\"\n if not isinstance(input, str):\n raise TypeError(f\"Expected 'input' to be str, got {type(input).__name__}\")\n\n input = input.strip()\n if not input:\n raise ValueError(\"Input text cannot be empty or whitespace only\")\n\n try:\n # Encode based on encoding_type\n if self._encoding_type == \"query\":\n sparse_vector = self._encoder.encode_queries(input)\n else: # encoding_type == \"document\"\n sparse_vector = self._encoder.encode_documents(input)\n\n # DashText returns dict with int/long keys and float values\n # Convert to standard format: {int: float}\n sparse_dict: dict[int, float] = {}\n for key, value in sparse_vector.items():\n try:\n idx = int(key)\n val = float(value)\n if val > 0:\n sparse_dict[idx] = val\n except (ValueError, TypeError):\n # Skip invalid entries\n continue\n\n # Sort by indices (keys) to ensure consistent ordering\n return dict(sorted(sparse_dict.items()))\n\n except Exception as e:\n if isinstance(e, (TypeError, ValueError)):\n raise\n raise RuntimeError(f\"Failed to generate BM25 embedding: {e!s}\") from e\n" }, { "path": "python/zvec/extension/embedding_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom abc import abstractmethod\nfrom typing import Protocol, runtime_checkable\n\nfrom ..common.constants import MD, DenseVectorType, SparseVectorType\n\n\n@runtime_checkable\nclass DenseEmbeddingFunction(Protocol[MD]):\n \"\"\"Protocol for dense vector embedding functions.\n\n Dense embedding functions map multimodal input (text, image, or audio) to\n fixed-length real-valued vectors. This is a Protocol class that defines\n the interface - implementations should provide their own initialization\n and properties.\n\n Type Parameters:\n MD: The type of input data (bound to Embeddable: TEXT, IMAGE, or AUDIO).\n\n Note:\n - This is a Protocol class - it only defines the ``embed()`` interface.\n - Implementations are free to define their own ``__init__``, properties,\n and additional methods as needed.\n - The ``embed()`` method is the only required interface.\n\n Examples:\n >>> # Custom text embedding implementation\n >>> class MyTextEmbedding:\n ... def __init__(self, dimension: int, model_name: str):\n ... self.dimension = dimension\n ... self.model = load_model(model_name)\n ...\n ... def embed(self, input: str) -> list[float]:\n ... return self.model.encode(input).tolist()\n\n >>> # Custom image embedding implementation\n >>> class MyImageEmbedding:\n ... def __init__(self, dimension: int = 512):\n ... self.dimension = dimension\n ... self.model = load_image_model()\n ...\n ... def embed(self, input: Union[str, bytes, np.ndarray]) -> list[float]:\n ... if isinstance(input, str):\n ... image = load_image_from_path(input)\n ... else:\n ... image = input\n ... return self.model.extract_features(image).tolist()\n\n >>> # Using built-in implementations\n >>> from zvec.extension import QwenDenseEmbedding\n >>> text_emb = QwenDenseEmbedding(dimension=768, api_key=\"sk-xxx\")\n >>> vector = text_emb.embed(\"Hello world\")\n \"\"\"\n\n @abstractmethod\n def embed(self, input: MD) -> DenseVectorType:\n \"\"\"Generate a dense embedding vector for the input data.\n\n Args:\n input (MD): Multimodal input data to embed. Can be:\n - TEXT (str): Text string\n - IMAGE (str | bytes | np.ndarray): Image file path, raw bytes, or array\n - AUDIO (str | bytes | np.ndarray): Audio file path, raw bytes, or array\n\n Returns:\n DenseVectorType: A dense vector representing the embedding.\n Can be list[float], list[int], or np.ndarray.\n Length should match the implementation's dimension.\n \"\"\"\n ...\n\n\n@runtime_checkable\nclass SparseEmbeddingFunction(Protocol[MD]):\n \"\"\"Abstract base class for sparse vector embedding functions.\n\n Sparse embedding functions map multimodal input (text, image, or audio) to\n a dictionary of {index: weight}, where only non-zero dimensions are stored.\n You can inherit this class to create custom sparse embedding functions.\n\n Type Parameters:\n MD: The type of input data (bound to Embeddable: TEXT, IMAGE, or AUDIO).\n\n Note:\n Subclasses must implement the ``embed()`` method.\n\n Examples:\n >>> # Using built-in text sparse embedding (e.g., BM25, TF-IDF)\n >>> sparse_emb = SomeSparseEmbedding()\n >>> vector = sparse_emb.embed(\"Hello world\")\n >>> # Returns: {0: 0.5, 42: 1.2, 100: 0.8}\n\n >>> # Custom BM25 sparse embedding function\n >>> class MyBM25Embedding(SparseEmbeddingFunction):\n ... def __init__(self, vocab_size: int = 10000):\n ... self.vocab_size = vocab_size\n ... self.tokenizer = MyTokenizer()\n ...\n ... def embed(self, input: str) -> dict[int, float]:\n ... tokens = self.tokenizer.tokenize(input)\n ... sparse_vector = {}\n ... for token_id, weight in self._calculate_bm25(tokens):\n ... if weight > 0:\n ... sparse_vector[token_id] = weight\n ... return sparse_vector\n ...\n ... def _calculate_bm25(self, tokens):\n ... # BM25 calculation logic\n ... pass\n\n >>> # Custom sparse image feature extractor\n >>> class MySparseImageEmbedding(SparseEmbeddingFunction):\n ... def embed(self, input: Union[str, bytes, np.ndarray]) -> dict[int, float]:\n ... image = self._load_image(input)\n ... features = self._extract_sparse_features(image)\n ... return {idx: val for idx, val in enumerate(features) if val != 0}\n \"\"\"\n\n @abstractmethod\n def embed(self, input: MD) -> SparseVectorType:\n \"\"\"Generate a sparse embedding for the input data.\n\n Args:\n input (MD): Multimodal input data to embed. Can be:\n - TEXT (str): Text string\n - IMAGE (str | bytes | np.ndarray): Image file path, raw bytes, or array\n - AUDIO (str | bytes | np.ndarray): Audio file path, raw bytes, or array\n\n Returns:\n SparseVectorType: Mapping from dimension index to non-zero weight.\n Only dimensions with non-zero values are included.\n \"\"\"\n ...\n" }, { "path": "python/zvec/extension/http_embedding_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport json\nimport os\nimport urllib.request\nfrom functools import lru_cache\nfrom typing import Optional\n\nfrom ..common.constants import TEXT, DenseVectorType\nfrom .embedding_function import DenseEmbeddingFunction\n\n\nclass HTTPDenseEmbedding(DenseEmbeddingFunction[TEXT]):\n \"\"\"Dense text embedding function using any OpenAI-compatible HTTP endpoint.\n\n This class calls any server that implements the ``/v1/embeddings`` API\n (LM Studio, Ollama, vLLM, LocalAI, etc.) using only the Python standard\n library — no extra dependencies are required.\n\n The embedding dimension is detected automatically from the first server\n response.\n\n Args:\n base_url (str, optional): Base URL of the embedding server.\n Defaults to ``\"http://localhost:1234\"`` (LM Studio).\n Common values:\n\n - ``\"http://localhost:1234\"`` — LM Studio\n - ``\"http://localhost:11434\"`` — Ollama\n model (str, optional): Model identifier as expected by the server.\n Defaults to ``\"text-embedding-nomic-embed-text-v1.5@f16\"``.\n api_key (Optional[str], optional): Bearer token for authenticated\n endpoints. Falls back to the ``OPENAI_API_KEY`` environment\n variable. Leave as ``None`` for local servers that do not\n require authentication.\n timeout (int, optional): HTTP request timeout in seconds.\n Defaults to 30.\n\n Attributes:\n dimension (int): Embedding vector dimensionality (auto-detected).\n\n Raises:\n TypeError: If ``embed()`` receives a non-string input.\n ValueError: If input is empty/whitespace-only or the server returns\n an unexpected response format.\n RuntimeError: If the HTTP request fails or the server is unreachable.\n\n Examples:\n >>> from zvec.extension import HTTPDenseEmbedding\n >>>\n >>> # LM Studio (default)\n >>> emb = HTTPDenseEmbedding()\n >>> vector = emb.embed(\"Hello, world!\")\n >>> len(vector)\n 768\n >>>\n >>> # Ollama\n >>> emb = HTTPDenseEmbedding(\n ... base_url=\"http://localhost:11434\",\n ... model=\"nomic-embed-text\",\n ... )\n >>> vector = emb.embed(\"Semantic search with local models\")\n\n See Also:\n - ``DenseEmbeddingFunction``: Protocol for dense embeddings.\n - ``OpenAIDenseEmbedding``: Cloud embedding via the OpenAI API.\n \"\"\"\n\n ENDPOINT = \"/v1/embeddings\"\n\n def __init__(\n self,\n base_url: str = \"http://localhost:1234\",\n model: str = \"text-embedding-nomic-embed-text-v1.5@f16\",\n api_key: Optional[str] = None,\n timeout: int = 30,\n ) -> None:\n self._base_url = base_url.rstrip(\"/\")\n self._model = model\n self._api_key = api_key or os.environ.get(\"OPENAI_API_KEY\", \"\")\n self._timeout = timeout\n self._dimension: Optional[int] = None\n\n @property\n def dimension(self) -> int:\n \"\"\"int: Embedding vector dimensionality (auto-detected on first call).\"\"\"\n if self._dimension is None:\n self._dimension = len(self.embed(\"dimension probe\"))\n return self._dimension\n\n def __call__(self, input: TEXT) -> DenseVectorType:\n \"\"\"Make the embedding function callable.\"\"\"\n return self.embed(input)\n\n @lru_cache(maxsize=256)\n def embed(self, input: TEXT) -> DenseVectorType:\n \"\"\"Generate a dense embedding vector for the input text.\n\n Results are cached (LRU, up to 256 entries) so repeated strings\n do not trigger extra HTTP requests.\n\n Args:\n input (TEXT): Input text string to embed. Must be non-empty\n after stripping whitespace.\n\n Returns:\n DenseVectorType: A list of floats representing the embedding.\n\n Raises:\n TypeError: If *input* is not a string.\n ValueError: If *input* is empty/whitespace-only or the server\n returns an unexpected response format.\n RuntimeError: If the HTTP request fails.\n \"\"\"\n if not isinstance(input, TEXT):\n raise TypeError(f\"Expected 'input' to be str, got {type(input).__name__}\")\n\n input = input.strip()\n if not input:\n raise ValueError(\"Input text cannot be empty or whitespace only\")\n\n url = self._base_url + self.ENDPOINT\n payload = json.dumps({\"model\": self._model, \"input\": input}).encode()\n\n headers: dict[str, str] = {\"Content-Type\": \"application/json\"}\n if self._api_key:\n headers[\"Authorization\"] = f\"Bearer {self._api_key}\"\n\n req = urllib.request.Request(url, data=payload, headers=headers, method=\"POST\")\n try:\n with urllib.request.urlopen(req, timeout=self._timeout) as resp:\n body = json.loads(resp.read())\n except urllib.error.HTTPError as exc:\n raise RuntimeError(\n f\"Embedding server returned HTTP {exc.code}: {exc.read().decode()}\"\n ) from exc\n except OSError as exc:\n raise RuntimeError(\n f\"Could not reach embedding server at {url}: {exc}\"\n ) from exc\n\n try:\n vector: list[float] = body[\"data\"][0][\"embedding\"]\n except (KeyError, IndexError) as exc:\n raise ValueError(\n f\"Unexpected response format from embedding server: {body}\"\n ) from exc\n\n return vector\n" }, { "path": "python/zvec/extension/jina_embedding_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom functools import lru_cache\nfrom typing import Optional\n\nfrom ..common.constants import TEXT, DenseVectorType\nfrom .embedding_function import DenseEmbeddingFunction\nfrom .jina_function import JinaFunctionBase\n\n\nclass JinaDenseEmbedding(JinaFunctionBase, DenseEmbeddingFunction[TEXT]):\n \"\"\"Dense text embedding function using Jina AI API.\n\n This class provides text-to-vector embedding capabilities using Jina AI's\n embedding models. It inherits from ``DenseEmbeddingFunction`` and implements\n dense text embedding via the Jina Embeddings API (OpenAI-compatible).\n\n Jina Embeddings v5 models support task-specific embedding through the\n ``task`` parameter, which optimizes the embedding for different use cases\n such as retrieval, text matching, or classification. They also support\n Matryoshka Representation Learning, allowing flexible output dimensions.\n\n Args:\n model (str, optional): Jina embedding model identifier.\n Defaults to ``\"jina-embeddings-v5-text-nano\"``. Available models:\n - ``\"jina-embeddings-v5-text-nano\"``: 768 dims, 239M params, 8K context\n - ``\"jina-embeddings-v5-text-small\"``: 1024 dims, 677M params, 32K context\n dimension (Optional[int], optional): Desired output embedding dimension.\n If ``None``, uses model's default dimension. Supports Matryoshka\n dimensions: 32, 64, 128, 256, 512, 768 (nano) / 1024 (small).\n Defaults to ``None``.\n api_key (Optional[str], optional): Jina API authentication key.\n If ``None``, reads from ``JINA_API_KEY`` environment variable.\n Obtain your key from: https://jina.ai/api-dashboard\n task (Optional[str], optional): Task type to optimize embeddings for.\n Defaults to ``None``. Valid values:\n - ``\"retrieval.query\"``: For search queries\n - ``\"retrieval.passage\"``: For documents/passages to be searched\n - ``\"text-matching\"``: For symmetric text similarity\n - ``\"classification\"``: For text classification\n - ``\"separation\"``: For clustering/separation tasks\n\n Attributes:\n dimension (int): The embedding vector dimension.\n data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation.\n model (str): The Jina model name being used.\n task (Optional[str]): The task type for embedding optimization.\n\n Raises:\n ValueError: If API key is not provided and not found in environment,\n if task is not a valid task type, or if API returns an error response.\n TypeError: If input to ``embed()`` is not a string.\n RuntimeError: If network error or Jina service error occurs.\n\n Note:\n - Requires Python 3.10, 3.11, or 3.12\n - Requires the ``openai`` package: ``pip install openai``\n - Jina API is OpenAI-compatible, so it uses the ``openai`` Python client\n - Embedding results are cached (LRU cache, maxsize=10) to reduce API calls\n - For retrieval tasks, use ``\"retrieval.query\"`` for queries and\n ``\"retrieval.passage\"`` for documents\n - API usage requires a Jina API key from https://jina.ai/api-dashboard\n\n Examples:\n >>> # Basic usage with default model\n >>> from zvec.extension import JinaDenseEmbedding\n >>> import os\n >>> os.environ[\"JINA_API_KEY\"] = \"jina_...\"\n >>>\n >>> emb_func = JinaDenseEmbedding()\n >>> vector = emb_func.embed(\"Hello, world!\")\n >>> len(vector)\n 768\n\n >>> # Retrieval use case: embed queries and documents differently\n >>> query_emb = JinaDenseEmbedding(task=\"retrieval.query\")\n >>> doc_emb = JinaDenseEmbedding(task=\"retrieval.passage\")\n >>>\n >>> query_vector = query_emb.embed(\"What is machine learning?\")\n >>> doc_vector = doc_emb.embed(\"Machine learning is a subset of AI...\")\n\n >>> # Using larger model with custom dimension (Matryoshka)\n >>> emb_func = JinaDenseEmbedding(\n ... model=\"jina-embeddings-v5-text-small\",\n ... dimension=256,\n ... api_key=\"jina_...\",\n ... task=\"text-matching\",\n ... )\n >>> vector = emb_func.embed(\"Semantic similarity comparison\")\n >>> len(vector)\n 256\n\n >>> # Using with zvec collection\n >>> import zvec\n >>> emb_func = JinaDenseEmbedding(task=\"retrieval.passage\")\n >>> schema = zvec.CollectionSchema(\n ... name=\"docs\",\n ... vectors=zvec.VectorSchema(\n ... \"embedding\", zvec.DataType.VECTOR_FP32, emb_func.dimension\n ... ),\n ... )\n >>> collection = zvec.create_and_open(path=\"./my_docs\", schema=schema)\n\n See Also:\n - ``DenseEmbeddingFunction``: Base class for dense embeddings\n - ``OpenAIDenseEmbedding``: Alternative using OpenAI API\n - ``QwenDenseEmbedding``: Alternative using Qwen/DashScope API\n - ``DefaultLocalDenseEmbedding``: Local model without API calls\n \"\"\"\n\n def __init__(\n self,\n model: str = \"jina-embeddings-v5-text-nano\",\n dimension: Optional[int] = None,\n api_key: Optional[str] = None,\n task: Optional[str] = None,\n **kwargs,\n ):\n \"\"\"Initialize the Jina dense embedding function.\n\n Args:\n model (str): Jina model name. Defaults to \"jina-embeddings-v5-text-nano\".\n dimension (Optional[int]): Target embedding dimension or None for default.\n api_key (Optional[str]): API key or None to use environment variable.\n task (Optional[str]): Task type for embedding optimization or None.\n **kwargs: Additional parameters for API calls.\n\n Raises:\n ValueError: If API key is not provided and not in environment,\n or if task is not a valid task type.\n \"\"\"\n # Initialize base class for API connection\n JinaFunctionBase.__init__(self, model=model, api_key=api_key, task=task)\n\n # Store dimension configuration\n self._custom_dimension = dimension\n\n # Determine actual dimension\n if dimension is None:\n self._dimension = self._MODEL_DIMENSIONS.get(model, 768)\n else:\n self._dimension = dimension\n\n # Store extra attributes\n self._extra_params = kwargs\n\n @property\n def dimension(self) -> int:\n \"\"\"int: The expected dimensionality of the embedding vector.\"\"\"\n return self._dimension\n\n @property\n def extra_params(self) -> dict:\n \"\"\"dict: Extra parameters for model-specific customization.\"\"\"\n return self._extra_params\n\n def __call__(self, input: TEXT) -> DenseVectorType:\n \"\"\"Make the embedding function callable.\"\"\"\n return self.embed(input)\n\n @lru_cache(maxsize=10)\n def embed(self, input: TEXT) -> DenseVectorType:\n \"\"\"Generate dense embedding vector for the input text.\n\n This method calls the Jina Embeddings API to convert input text\n into a dense vector representation. Results are cached to improve\n performance for repeated inputs.\n\n Args:\n input (TEXT): Input text string to embed. Must be non-empty after\n stripping whitespace. Maximum length depends on model:\n 8192 tokens for v5-nano, 32768 tokens for v5-small.\n\n Returns:\n DenseVectorType: A list of floats representing the embedding vector.\n Length equals ``self.dimension``. Example:\n ``[0.123, -0.456, 0.789, ...]``\n\n Raises:\n TypeError: If ``input`` is not a string.\n ValueError: If input is empty/whitespace-only, or if the API returns\n an error or malformed response.\n RuntimeError: If network connectivity issues or Jina service\n errors occur.\n\n Examples:\n >>> emb = JinaDenseEmbedding(task=\"retrieval.query\")\n >>> vector = emb.embed(\"What is deep learning?\")\n >>> len(vector)\n 768\n >>> isinstance(vector[0], float)\n True\n\n >>> # Error: empty input\n >>> emb.embed(\" \")\n ValueError: Input text cannot be empty or whitespace only\n\n >>> # Error: non-string input\n >>> emb.embed(123)\n TypeError: Expected 'input' to be str, got int\n\n Note:\n - This method is cached (maxsize=10). Identical inputs return cached results.\n - The cache is based on exact string match (case-sensitive).\n - Task type affects embedding optimization but not caching behavior.\n \"\"\"\n if not isinstance(input, TEXT):\n raise TypeError(f\"Expected 'input' to be str, got {type(input).__name__}\")\n\n input = input.strip()\n if not input:\n raise ValueError(\"Input text cannot be empty or whitespace only\")\n\n # Call API\n embedding_vector = self._call_text_embedding_api(\n input=input,\n dimension=self._custom_dimension,\n )\n\n # Verify dimension\n if len(embedding_vector) != self.dimension:\n raise ValueError(\n f\"Dimension mismatch: expected {self.dimension}, \"\n f\"got {len(embedding_vector)}\"\n )\n\n return embedding_vector\n" }, { "path": "python/zvec/extension/jina_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport os\nfrom typing import ClassVar, Optional\n\nfrom ..common.constants import TEXT\nfrom ..tool import require_module\n\n\nclass JinaFunctionBase:\n \"\"\"Base class for Jina AI functions.\n\n This base class provides common functionality for calling Jina AI APIs\n and handling responses. It supports embeddings (dense) operations via\n the OpenAI-compatible Jina Embeddings API.\n\n This class is not meant to be used directly. Use concrete implementations:\n - ``JinaDenseEmbedding`` for dense embeddings\n\n Args:\n model (str): Jina embedding model identifier.\n api_key (Optional[str]): Jina API authentication key.\n task (Optional[str]): Task type for the embedding model.\n\n Note:\n - This is an internal base class for code reuse across Jina features\n - Subclasses should inherit from appropriate Protocol\n - Provides unified API connection and response handling\n - Jina API is OpenAI-compatible, using the ``openai`` Python client\n \"\"\"\n\n _BASE_URL: ClassVar[str] = \"https://api.jina.ai/v1\"\n\n # Model default dimensions\n _MODEL_DIMENSIONS: ClassVar[dict[str, int]] = {\n \"jina-embeddings-v5-text-nano\": 768,\n \"jina-embeddings-v5-text-small\": 1024,\n }\n\n # Model max tokens\n _MODEL_MAX_TOKENS: ClassVar[dict[str, int]] = {\n \"jina-embeddings-v5-text-nano\": 8192,\n \"jina-embeddings-v5-text-small\": 32768,\n }\n\n # Valid task types\n _VALID_TASKS: ClassVar[tuple[str, ...]] = (\n \"retrieval.query\",\n \"retrieval.passage\",\n \"text-matching\",\n \"classification\",\n \"separation\",\n )\n\n def __init__(\n self,\n model: str,\n api_key: Optional[str] = None,\n task: Optional[str] = None,\n ):\n \"\"\"Initialize the base Jina functionality.\n\n Args:\n model (str): Jina model name.\n api_key (Optional[str]): API key or None to use environment variable.\n task (Optional[str]): Task type for the embedding model.\n Valid values: \"retrieval.query\", \"retrieval.passage\",\n \"text-matching\", \"classification\", \"separation\".\n\n Raises:\n ValueError: If API key is not provided and not in environment,\n or if task is not a valid task type.\n \"\"\"\n self._model = model\n self._api_key = api_key or os.environ.get(\"JINA_API_KEY\")\n self._task = task\n\n if not self._api_key:\n raise ValueError(\n \"Jina API key is required. Please provide 'api_key' parameter \"\n \"or set the 'JINA_API_KEY' environment variable. \"\n \"Get your key from: https://jina.ai/api-dashboard\"\n )\n\n if task is not None and task not in self._VALID_TASKS:\n raise ValueError(\n f\"Invalid task '{task}'. Valid tasks: {', '.join(self._VALID_TASKS)}\"\n )\n\n @property\n def model(self) -> str:\n \"\"\"str: The Jina model name currently in use.\"\"\"\n return self._model\n\n @property\n def task(self) -> Optional[str]:\n \"\"\"Optional[str]: The task type for the embedding model.\"\"\"\n return self._task\n\n def _get_client(self):\n \"\"\"Get OpenAI-compatible client instance configured for Jina API.\n\n Returns:\n OpenAI: Configured OpenAI client pointing to Jina API.\n\n Raises:\n ImportError: If openai package is not installed.\n \"\"\"\n openai = require_module(\"openai\")\n return openai.OpenAI(api_key=self._api_key, base_url=self._BASE_URL)\n\n def _call_text_embedding_api(\n self,\n input: TEXT,\n dimension: Optional[int] = None,\n ) -> list:\n \"\"\"Call Jina Embeddings API.\n\n Args:\n input (TEXT): Input text to embed.\n dimension (Optional[int]): Target dimension for Matryoshka embeddings.\n\n Returns:\n list: Embedding vector as list of floats.\n\n Raises:\n RuntimeError: If API call fails.\n ValueError: If API returns error response.\n \"\"\"\n try:\n client = self._get_client()\n\n # Prepare embedding parameters\n params = {\"model\": self.model, \"input\": input}\n\n # Add dimension parameter for Matryoshka support\n if dimension is not None:\n params[\"dimensions\"] = dimension\n\n # Add task parameter via extra_body\n if self._task is not None:\n params[\"extra_body\"] = {\"task\": self._task}\n\n # Call Jina API (OpenAI-compatible)\n response = client.embeddings.create(**params)\n\n except Exception as e:\n # Check if it's an OpenAI API error\n openai = require_module(\"openai\")\n if isinstance(e, (openai.APIError, openai.APIConnectionError)):\n raise RuntimeError(f\"Failed to call Jina API: {e!s}\") from e\n raise RuntimeError(f\"Unexpected error during API call: {e!s}\") from e\n\n # Extract embedding from response\n try:\n if not response.data:\n raise ValueError(\"Invalid API response: no embedding data returned\")\n\n embedding_vector = response.data[0].embedding\n\n if not isinstance(embedding_vector, list):\n raise ValueError(\n \"Invalid API response: embedding is not a list of numbers\"\n )\n\n return embedding_vector\n\n except (AttributeError, IndexError, TypeError) as e:\n raise ValueError(f\"Failed to parse API response: {e!s}\") from e\n" }, { "path": "python/zvec/extension/multi_vector_reranker.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport heapq\nimport math\nfrom collections import defaultdict\nfrom typing import Optional\n\nfrom ..model.doc import Doc\nfrom ..typing import MetricType\nfrom .rerank_function import RerankFunction\n\n\nclass RrfReRanker(RerankFunction):\n \"\"\"Re-ranker using Reciprocal Rank Fusion (RRF) for multi-vector search.\n\n RRF combines results from multiple vector queries without requiring relevance scores.\n It assigns higher weight to documents that appear early in multiple result lists.\n\n The RRF score for a document at rank ``r`` is: ``1 / (k + r + 1)``,\n where ``k`` is the rank constant.\n\n Note:\n This re-ranker is specifically designed for multi-vector scenarios where\n query results from multiple vector fields need to be combined.\n\n Args:\n topn (int, optional): Number of top documents to return. Defaults to 10.\n rerank_field (Optional[str], optional): Ignored by RRF. Defaults to None.\n rank_constant (int, optional): Smoothing constant ``k`` in RRF formula.\n Larger values reduce the impact of early ranks. Defaults to 60.\n \"\"\"\n\n def __init__(\n self,\n topn: int = 10,\n rerank_field: Optional[str] = None,\n rank_constant: int = 60,\n ):\n super().__init__(topn=topn, rerank_field=rerank_field)\n self._rank_constant = rank_constant\n\n @property\n def rank_constant(self) -> int:\n return self._rank_constant\n\n def _rrf_score(self, rank: int) -> float:\n return 1.0 / (self._rank_constant + rank + 1)\n\n def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:\n \"\"\"Apply Reciprocal Rank Fusion to combine multiple query results.\n\n Args:\n query_results (dict[str, list[Doc]]): Results from one or more vector queries.\n\n Returns:\n list[Doc]: Re-ranked documents with RRF scores in the ``score`` field.\n \"\"\"\n rrf_scores: dict[str, float] = defaultdict(float)\n id_to_doc: dict[str, Doc] = {}\n\n for _, query_result in query_results.items():\n for rank, doc in enumerate(query_result):\n doc_id = doc.id\n rrf_score = self._rrf_score(rank)\n rrf_scores[doc_id] += rrf_score\n if doc_id not in id_to_doc:\n id_to_doc[doc_id] = doc\n\n top_docs = heapq.nlargest(self.topn, rrf_scores.items(), key=lambda x: x[1])\n results: list[Doc] = []\n for doc_id, rrf_score in top_docs:\n doc = id_to_doc[doc_id]\n new_doc = doc._replace(score=rrf_score)\n results.append(new_doc)\n return results\n\n\nclass WeightedReRanker(RerankFunction):\n \"\"\"Re-ranker that combines scores from multiple vector fields using weights.\n\n Each vector field's relevance score is normalized based on its metric type,\n then scaled by a user-provided weight. Final scores are summed across fields.\n\n Note:\n This re-ranker is specifically designed for multi-vector scenarios where\n query results from multiple vector fields need to be combined with\n configurable weights.\n\n Args:\n topn (int, optional): Number of top documents to return. Defaults to 10.\n rerank_field (Optional[str], optional): Ignored. Defaults to None.\n metric (MetricType, optional): Distance metric used for score normalization.\n Defaults to ``MetricType.L2``.\n weights (Optional[dict[str, float]], optional): Weight per vector field.\n Fields not listed use weight 1.0. Defaults to None.\n\n Note:\n Supported metrics: L2, IP, COSINE. Scores are normalized to [0, 1].\n \"\"\"\n\n def __init__(\n self,\n topn: int = 10,\n rerank_field: Optional[str] = None,\n metric: MetricType = MetricType.L2,\n weights: Optional[dict[str, float]] = None,\n ):\n super().__init__(topn=topn, rerank_field=rerank_field)\n self._weights = weights or {}\n self._metric = metric\n\n @property\n def weights(self) -> dict[str, float]:\n \"\"\"dict[str, float]: Weight mapping for vector fields.\"\"\"\n return self._weights\n\n @property\n def metric(self) -> MetricType:\n \"\"\"MetricType: Distance metric used for score normalization.\"\"\"\n return self._metric\n\n def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:\n \"\"\"Combine scores from multiple vector fields using weighted sum.\n\n Args:\n query_results (dict[str, list[Doc]]): Results per vector field.\n\n Returns:\n list[Doc]: Re-ranked documents with combined scores in ``score`` field.\n \"\"\"\n weighted_scores: dict[str, float] = defaultdict(float)\n id_to_doc: dict[str, Doc] = {}\n\n for vector_name, query_result in query_results.items():\n for _, doc in enumerate(query_result):\n doc_id = doc.id\n weighted_score = self._normalize_score(\n doc.score, self.metric\n ) * self.weights.get(vector_name, 1.0)\n weighted_scores[doc_id] += weighted_score\n if doc_id not in id_to_doc:\n id_to_doc[doc_id] = doc\n\n top_docs = heapq.nlargest(\n self.topn, weighted_scores.items(), key=lambda x: x[1]\n )\n results: list[Doc] = []\n for doc_id, weighted_score in top_docs:\n doc = id_to_doc[doc_id]\n new_doc = doc._replace(score=weighted_score)\n results.append(new_doc)\n return results\n\n def _normalize_score(self, score: float, metric: MetricType) -> float:\n if metric == MetricType.L2:\n return 1.0 - 2 * math.atan(score) / math.pi\n if metric == MetricType.IP:\n return 0.5 + math.atan(score) / math.pi\n if metric == MetricType.COSINE:\n return 1.0 - score / 2.0\n raise ValueError(\"Unsupported metric type\")\n" }, { "path": "python/zvec/extension/openai_embedding_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom functools import lru_cache\nfrom typing import Optional\n\nfrom ..common.constants import TEXT, DenseVectorType\nfrom .embedding_function import DenseEmbeddingFunction\nfrom .openai_function import OpenAIFunctionBase\n\n\nclass OpenAIDenseEmbedding(OpenAIFunctionBase, DenseEmbeddingFunction[TEXT]):\n \"\"\"Dense text embedding function using OpenAI API.\n\n This class provides text-to-vector embedding capabilities using OpenAI's\n embedding models. It inherits from ``DenseEmbeddingFunction`` and implements\n dense text embedding via the OpenAI API.\n\n The implementation supports various OpenAI embedding models with different\n dimensions and includes automatic result caching for improved performance.\n\n Args:\n model (str, optional): OpenAI embedding model identifier.\n Defaults to ``\"text-embedding-3-small\"``. Common options:\n - ``\"text-embedding-3-small\"``: 1536 dims, cost-efficient, good performance\n - ``\"text-embedding-3-large\"``: 3072 dims, highest quality\n - ``\"text-embedding-ada-002\"``: 1536 dims, legacy model\n dimension (Optional[int], optional): Desired output embedding dimension.\n If ``None``, uses model's default dimension. For text-embedding-3 models,\n you can specify custom dimensions (e.g., 256, 512, 1024, 1536).\n Defaults to ``None``.\n api_key (Optional[str], optional): OpenAI API authentication key.\n If ``None``, reads from ``OPENAI_API_KEY`` environment variable.\n Obtain your key from: https://platform.openai.com/api-keys\n base_url (Optional[str], optional): Custom API base URL for OpenAI-compatible\n services. Defaults to ``None`` (uses official OpenAI endpoint).\n\n Attributes:\n dimension (int): The embedding vector dimension.\n data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation.\n model (str): The OpenAI model name being used.\n\n Raises:\n ValueError: If API key is not provided and not found in environment,\n or if API returns an error response.\n TypeError: If input to ``embed()`` is not a string.\n RuntimeError: If network error or OpenAI service error occurs.\n\n Note:\n - Requires Python 3.10, 3.11, or 3.12\n - Requires the ``openai`` package: ``pip install openai``\n - Embedding results are cached (LRU cache, maxsize=10) to reduce API calls\n - Network connectivity to OpenAI API endpoints is required\n - API usage incurs costs based on your OpenAI subscription plan\n - Rate limits apply based on your OpenAI account tier\n\n Examples:\n >>> # Basic usage with default model\n >>> from zvec.extension import OpenAIDenseEmbedding\n >>> import os\n >>> os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n >>>\n >>> emb_func = OpenAIDenseEmbedding()\n >>> vector = emb_func.embed(\"Hello, world!\")\n >>> len(vector)\n 1536\n\n >>> # Using specific model with custom dimension\n >>> emb_func = OpenAIDenseEmbedding(\n ... model=\"text-embedding-3-large\",\n ... dimension=1024,\n ... api_key=\"sk-...\"\n ... )\n >>> vector = emb_func.embed(\"Machine learning is fascinating\")\n >>> len(vector)\n 1024\n\n >>> # Using with custom base URL (e.g., Azure OpenAI)\n >>> emb_func = OpenAIDenseEmbedding(\n ... model=\"text-embedding-ada-002\",\n ... api_key=\"your-azure-key\",\n ... base_url=\"https://your-resource.openai.azure.com/\"\n ... )\n >>> vector = emb_func(\"Natural language processing\")\n >>> isinstance(vector, list)\n True\n\n >>> # Batch processing with caching benefit\n >>> texts = [\"First text\", \"Second text\", \"First text\"]\n >>> vectors = [emb_func.embed(text) for text in texts]\n >>> # Third call uses cached result for \"First text\"\n\n >>> # Error handling\n >>> try:\n ... emb_func.embed(\"\") # Empty string\n ... except ValueError as e:\n ... print(f\"Error: {e}\")\n Error: Input text cannot be empty or whitespace only\n\n See Also:\n - ``DenseEmbeddingFunction``: Base class for dense embeddings\n - ``QwenDenseEmbedding``: Alternative using Qwen/DashScope API\n - ``DefaultDenseEmbedding``: Local model without API calls\n - ``SparseEmbeddingFunction``: Base class for sparse embeddings\n \"\"\"\n\n def __init__(\n self,\n model: str = \"text-embedding-3-small\",\n dimension: Optional[int] = None,\n api_key: Optional[str] = None,\n base_url: Optional[str] = None,\n **kwargs,\n ):\n \"\"\"Initialize the OpenAI dense embedding function.\n\n Args:\n model (str): OpenAI model name. Defaults to \"text-embedding-3-small\".\n dimension (Optional[int]): Target embedding dimension or None for default.\n api_key (Optional[str]): API key or None to use environment variable.\n base_url (Optional[str]): Custom API base URL or None for default.\n **kwargs: Additional parameters for API calls. Examples:\n - ``encoding_format`` (str): Format of embeddings, \"float\" or \"base64\".\n - ``user`` (str): User identifier for tracking.\n\n Raises:\n ValueError: If API key is not provided and not in environment.\n \"\"\"\n # Initialize base class for API connection\n OpenAIFunctionBase.__init__(\n self, model=model, api_key=api_key, base_url=base_url\n )\n\n # Store dimension configuration\n self._custom_dimension = dimension\n\n # Determine actual dimension\n if dimension is None:\n # Use model default dimension\n self._dimension = self._MODEL_DIMENSIONS.get(model, 1536)\n else:\n self._dimension = dimension\n\n # Store dense-specific attributes\n self._extra_params = kwargs\n\n @property\n def dimension(self) -> int:\n \"\"\"int: The expected dimensionality of the embedding vector.\"\"\"\n return self._dimension\n\n @property\n def extra_params(self) -> dict:\n \"\"\"dict: Extra parameters for model-specific customization.\"\"\"\n return self._extra_params\n\n def __call__(self, input: TEXT) -> DenseVectorType:\n \"\"\"Make the embedding function callable.\"\"\"\n return self.embed(input)\n\n @lru_cache(maxsize=10)\n def embed(self, input: TEXT) -> DenseVectorType:\n \"\"\"Generate dense embedding vector for the input text.\n\n This method calls the OpenAI Embeddings API to convert input text\n into a dense vector representation. Results are cached to improve\n performance for repeated inputs.\n\n Args:\n input (TEXT): Input text string to embed. Must be non-empty after\n stripping whitespace. Maximum length is 8191 tokens for most models.\n\n Returns:\n DenseVectorType: A list of floats representing the embedding vector.\n Length equals ``self.dimension``. Example:\n ``[0.123, -0.456, 0.789, ...]``\n\n Raises:\n TypeError: If ``input`` is not a string.\n ValueError: If input is empty/whitespace-only, or if the API returns\n an error or malformed response.\n RuntimeError: If network connectivity issues or OpenAI service\n errors occur.\n\n Examples:\n >>> emb = OpenAIDenseEmbedding()\n >>> vector = emb.embed(\"Natural language processing\")\n >>> len(vector)\n 1536\n >>> isinstance(vector[0], float)\n True\n\n >>> # Error: empty input\n >>> emb.embed(\" \")\n ValueError: Input text cannot be empty or whitespace only\n\n >>> # Error: non-string input\n >>> emb.embed(123)\n TypeError: Expected 'input' to be str, got int\n\n Note:\n - This method is cached (maxsize=10). Identical inputs return cached results.\n - The cache is based on exact string match (case-sensitive).\n - Consider pre-processing text (lowercasing, normalization) for better caching.\n \"\"\"\n if not isinstance(input, TEXT):\n raise TypeError(f\"Expected 'input' to be str, got {type(input).__name__}\")\n\n input = input.strip()\n if not input:\n raise ValueError(\"Input text cannot be empty or whitespace only\")\n\n # Call API\n embedding_vector = self._call_text_embedding_api(\n input=input,\n dimension=self._custom_dimension,\n )\n\n # Verify dimension\n if len(embedding_vector) != self.dimension:\n raise ValueError(\n f\"Dimension mismatch: expected {self.dimension}, \"\n f\"got {len(embedding_vector)}\"\n )\n\n return embedding_vector\n" }, { "path": "python/zvec/extension/openai_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport os\nfrom typing import ClassVar, Optional\n\nfrom ..common.constants import TEXT\nfrom ..tool import require_module\n\n\nclass OpenAIFunctionBase:\n \"\"\"Base class for OpenAI functions.\n\n This base class provides common functionality for calling OpenAI APIs\n and handling responses. It supports embeddings (dense) operations.\n\n This class is not meant to be used directly. Use concrete implementations:\n - ``OpenAIDenseEmbedding`` for dense embeddings\n\n Args:\n model (str): OpenAI model identifier.\n api_key (Optional[str]): OpenAI API authentication key.\n base_url (Optional[str]): Custom API base URL.\n\n Note:\n - This is an internal base class for code reuse across OpenAI features\n - Subclasses should inherit from appropriate Protocol\n - Provides unified API connection and response handling\n \"\"\"\n\n # Model default dimensions\n _MODEL_DIMENSIONS: ClassVar[dict[str, int]] = {\n \"text-embedding-3-small\": 1536,\n \"text-embedding-3-large\": 3072,\n \"text-embedding-ada-002\": 1536,\n }\n\n def __init__(\n self,\n model: str,\n api_key: Optional[str] = None,\n base_url: Optional[str] = None,\n ):\n \"\"\"Initialize the base OpenAI functionality.\n\n Args:\n model (str): OpenAI model name.\n api_key (Optional[str]): API key or None to use environment variable.\n base_url (Optional[str]): Custom API base URL or None for default.\n\n Raises:\n ValueError: If API key is not provided and not in environment.\n \"\"\"\n self._model = model\n self._api_key = api_key or os.environ.get(\"OPENAI_API_KEY\")\n self._base_url = base_url\n\n if not self._api_key:\n raise ValueError(\n \"OpenAI API key is required. Please provide 'api_key' parameter \"\n \"or set the 'OPENAI_API_KEY' environment variable.\"\n )\n\n @property\n def model(self) -> str:\n \"\"\"str: The OpenAI model name currently in use.\"\"\"\n return self._model\n\n def _get_client(self):\n \"\"\"Get OpenAI client instance.\n\n Returns:\n OpenAI: Configured OpenAI client.\n\n Raises:\n ImportError: If openai package is not installed.\n \"\"\"\n openai = require_module(\"openai\")\n\n if self._base_url:\n return openai.OpenAI(api_key=self._api_key, base_url=self._base_url)\n return openai.OpenAI(api_key=self._api_key)\n\n def _call_text_embedding_api(\n self,\n input: TEXT,\n dimension: Optional[int] = None,\n ) -> list:\n \"\"\"Call OpenAI Embeddings API.\n\n Args:\n input (TEXT): Input text to embed.\n dimension (Optional[int]): Target dimension (for models that support it).\n\n Returns:\n list: Embedding vector as list of floats.\n\n Raises:\n RuntimeError: If API call fails.\n ValueError: If API returns error response.\n \"\"\"\n try:\n client = self._get_client()\n\n # Prepare embedding parameters\n params = {\"model\": self.model, \"input\": input}\n\n # Add dimension parameter for models that support it\n if dimension is not None:\n params[\"dimensions\"] = dimension\n\n # Call OpenAI API\n response = client.embeddings.create(**params)\n\n except Exception as e:\n # Check if it's an OpenAI API error\n openai = require_module(\"openai\")\n if isinstance(e, (openai.APIError, openai.APIConnectionError)):\n raise RuntimeError(f\"Failed to call OpenAI API: {e!s}\") from e\n raise RuntimeError(f\"Unexpected error during API call: {e!s}\") from e\n\n # Extract embedding from response\n try:\n if not response.data:\n raise ValueError(\"Invalid API response: no embedding data returned\")\n\n embedding_vector = response.data[0].embedding\n\n if not isinstance(embedding_vector, list):\n raise ValueError(\n \"Invalid API response: embedding is not a list of numbers\"\n )\n\n return embedding_vector\n\n except (AttributeError, IndexError, TypeError) as e:\n raise ValueError(f\"Failed to parse API response: {e!s}\") from e\n" }, { "path": "python/zvec/extension/qwen_embedding_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom functools import lru_cache\nfrom typing import Optional\n\nfrom ..common.constants import TEXT, DenseVectorType, SparseVectorType\nfrom .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction\nfrom .qwen_function import QwenFunctionBase\n\n\nclass QwenDenseEmbedding(QwenFunctionBase, DenseEmbeddingFunction[TEXT]):\n \"\"\"Dense text embedding function using Qwen (DashScope) API.\n\n This class provides text-to-vector embedding capabilities using Alibaba Cloud's\n DashScope service and Qwen embedding models. It inherits from\n ``DenseEmbeddingFunction`` and implements dense text embedding.\n\n The implementation supports various Qwen embedding models with configurable\n dimensions and includes automatic result caching for improved performance.\n\n Args:\n dimension (int): Desired output embedding dimension. Common values:\n - 512: Balanced performance and accuracy\n - 1024: Higher accuracy, larger storage\n - 1536: Maximum accuracy for supported models\n model (str, optional): DashScope embedding model identifier.\n Defaults to ``\"text-embedding-v4\"``. Other options include:\n - ``\"text-embedding-v3\"``\n - ``\"text-embedding-v2\"``\n - ``\"text-embedding-v1\"``\n api_key (Optional[str], optional): DashScope API authentication key.\n If ``None``, reads from ``DASHSCOPE_API_KEY`` environment variable.\n Obtain your key from: https://dashscope.console.aliyun.com/\n **kwargs: Additional DashScope API parameters. Supported options:\n - ``text_type`` (str): Specifies the text role in retrieval tasks.\n Options: ``\"query\"`` (search query) or ``\"document\"`` (indexed content).\n This parameter optimizes embeddings for asymmetric search scenarios.\n\n Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api\n\n Attributes:\n dimension (int): The embedding vector dimension.\n data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation.\n model (str): The DashScope model name being used.\n\n Raises:\n ValueError: If API key is not provided and not found in environment,\n or if API returns an error response.\n TypeError: If input to ``embed()`` is not a string.\n RuntimeError: If network error or DashScope service error occurs.\n\n Note:\n - Requires Python 3.10, 3.11, or 3.12\n - Requires the ``dashscope`` package: ``pip install dashscope``\n - Embedding results are cached (LRU cache, maxsize=10) to reduce API calls\n - Network connectivity to DashScope API endpoints is required\n - API usage may incur costs based on your DashScope subscription plan\n\n **Parameter Guidelines:**\n\n - Use ``text_type=\"query\"`` for search queries and ``text_type=\"document\"``\n for indexed content to optimize asymmetric retrieval tasks.\n - For detailed API specifications and parameter usage, refer to:\n https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api\n\n Examples:\n >>> # Basic usage with default model\n >>> from zvec.extension import QwenDenseEmbedding\n >>> import os\n >>> os.environ[\"DASHSCOPE_API_KEY\"] = \"your-api-key\"\n >>>\n >>> emb_func = QwenDenseEmbedding(dimension=1024)\n >>> vector = emb_func.embed(\"Hello, world!\")\n >>> len(vector)\n 1024\n\n >>> # Using specific model with explicit API key\n >>> emb_func = QwenDenseEmbedding(\n ... dimension=512,\n ... model=\"text-embedding-v3\",\n ... api_key=\"sk-xxxxx\"\n ... )\n >>> vector = emb_func(\"Machine learning is fascinating\")\n >>> isinstance(vector, list)\n True\n\n >>> # Using with custom parameters (text_type)\n >>> # For search queries - optimize for query-document matching\n >>> emb_func = QwenDenseEmbedding(\n ... dimension=1024,\n ... text_type=\"query\"\n ... )\n >>> query_vector = emb_func.embed(\"What is machine learning?\")\n >>>\n >>> # For document embeddings - optimize for being matched by queries\n >>> doc_emb_func = QwenDenseEmbedding(\n ... dimension=1024,\n ... text_type=\"document\"\n ... )\n >>> doc_vector = doc_emb_func.embed(\n ... \"Machine learning is a subset of artificial intelligence...\"\n ... )\n\n >>> # Batch processing with caching benefit\n >>> texts = [\"First text\", \"Second text\", \"First text\"]\n >>> vectors = [emb_func.embed(text) for text in texts]\n >>> # Third call uses cached result for \"First text\"\n\n >>> # Error handling\n >>> try:\n ... emb_func.embed(\"\") # Empty string\n ... except ValueError as e:\n ... print(f\"Error: {e}\")\n Error: Input text cannot be empty or whitespace only\n\n See Also:\n - ``DenseEmbeddingFunction``: Base class for dense embeddings\n - ``SparseEmbeddingFunction``: Base class for sparse embeddings\n \"\"\"\n\n def __init__(\n self,\n dimension: int,\n model: str = \"text-embedding-v4\",\n api_key: Optional[str] = None,\n **kwargs,\n ):\n \"\"\"Initialize the Qwen dense embedding function.\n\n Args:\n dimension (int): Target embedding dimension.\n model (str): DashScope model name. Defaults to \"text-embedding-v4\".\n api_key (Optional[str]): API key or None to use environment variable.\n **kwargs: Additional DashScope API parameters. Supported options:\n - ``text_type`` (str): Text role in asymmetric retrieval.\n * ``\"query\"``: Optimize for search queries (short, question-like).\n * ``\"document\"``: Optimize for indexed documents (longer content).\n Using appropriate text_type improves retrieval accuracy by\n optimizing the embedding space for query-document matching.\n\n For detailed API documentation, see:\n https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api\n\n Raises:\n ValueError: If API key is not provided and not in environment.\n \"\"\"\n # Initialize base class for API connection\n QwenFunctionBase.__init__(self, model=model, api_key=api_key)\n\n # Store dense-specific attributes\n self._dimension = dimension\n self._extra_params = kwargs\n\n @property\n def dimension(self) -> int:\n \"\"\"int: The expected dimensionality of the embedding vector.\"\"\"\n return self._dimension\n\n @property\n def extra_params(self) -> dict:\n \"\"\"dict: Extra parameters for model-specific customization.\"\"\"\n return self._extra_params\n\n def __call__(self, input: TEXT) -> DenseVectorType:\n \"\"\"Make the embedding function callable.\"\"\"\n return self.embed(input)\n\n @lru_cache(maxsize=10)\n def embed(self, input: TEXT) -> DenseVectorType:\n \"\"\"Generate dense embedding vector for the input text.\n\n This method calls the DashScope TextEmbedding API to convert input text\n into a dense vector representation. Results are cached to improve\n performance for repeated inputs.\n\n Args:\n input (TEXT): Input text string to embed. Must be non-empty after\n stripping whitespace. Maximum length depends on the model used\n (typically 2048-8192 tokens).\n\n Returns:\n DenseVectorType: A list of floats representing the embedding vector.\n Length equals ``self.dimension``. Example:\n ``[0.123, -0.456, 0.789, ...]``\n\n Raises:\n TypeError: If ``input`` is not a string.\n ValueError: If input is empty/whitespace-only, or if the API returns\n an error or malformed response.\n RuntimeError: If network connectivity issues or DashScope service\n errors occur.\n\n Examples:\n >>> emb = QwenDenseEmbedding(dimension=1024)\n >>> vector = emb.embed(\"Natural language processing\")\n >>> len(vector)\n 1024\n >>> isinstance(vector[0], float)\n True\n\n >>> # Error: empty input\n >>> emb.embed(\" \")\n ValueError: Input text cannot be empty or whitespace only\n\n >>> # Error: non-string input\n >>> emb.embed(123)\n TypeError: Expected 'input' to be str, got int\n\n Note:\n - This method is cached (maxsize=10). Identical inputs return cached results.\n - The cache is based on exact string match (case-sensitive).\n - Consider pre-processing text (lowercasing, normalization) for better caching.\n \"\"\"\n if not isinstance(input, TEXT):\n raise TypeError(f\"Expected 'input' to be str, got {type(input).__name__}\")\n\n input = input.strip()\n if not input:\n raise ValueError(\"Input text cannot be empty or whitespace only\")\n\n # Call API with dense output type\n output = self._call_text_embedding_api(\n input=input,\n dimension=self.dimension,\n output_type=\"dense\",\n text_type=self.extra_params.get(\"text_type\"),\n )\n\n embeddings = output.get(\"embeddings\")\n if not isinstance(embeddings, list):\n raise ValueError(\n \"Invalid API response: 'embeddings' field is missing or not a list\"\n )\n\n if len(embeddings) != 1:\n raise ValueError(\n f\"Expected exactly 1 embedding in response, got {len(embeddings)}\"\n )\n\n first_emb = embeddings[0]\n if not isinstance(first_emb, dict):\n raise ValueError(\"Invalid API response: embedding item is not a dictionary\")\n\n embedding_vector = first_emb.get(\"embedding\")\n if not isinstance(embedding_vector, list):\n raise ValueError(\n \"Invalid API response: 'embedding' field is missing or not a list\"\n )\n\n if len(embedding_vector) != self.dimension:\n raise ValueError(\n f\"Dimension mismatch: expected {self.dimension}, \"\n f\"got {len(embedding_vector)}\"\n )\n\n return list(embedding_vector)\n\n\nclass QwenSparseEmbedding(QwenFunctionBase, SparseEmbeddingFunction[TEXT]):\n \"\"\"Sparse text embedding function using Qwen (DashScope) API.\n\n This class provides text-to-sparse-vector embedding capabilities using\n Alibaba Cloud's DashScope service and Qwen embedding models. It generates\n sparse keyword-weighted vectors suitable for lexical matching and BM25-style\n retrieval scenarios.\n\n Sparse embeddings are particularly useful for:\n - Keyword-based search and exact matching\n - Hybrid retrieval (combining with dense embeddings)\n - Interpretable search results (weights show term importance)\n\n Args:\n dimension (int): Desired output embedding dimension. Common values:\n - 512: Balanced performance and accuracy\n - 1024: Higher accuracy, larger storage\n - 1536: Maximum accuracy for supported models\n model (str, optional): DashScope embedding model identifier.\n Defaults to ``\"text-embedding-v4\"``. Other options include:\n - ``\"text-embedding-v3\"``\n - ``\"text-embedding-v2\"``\n api_key (Optional[str], optional): DashScope API authentication key.\n If ``None``, reads from ``DASHSCOPE_API_KEY`` environment variable.\n Obtain your key from: https://dashscope.console.aliyun.com/\n **kwargs: Additional DashScope API parameters. Supported options:\n - ``encoding_type`` (Literal[\"query\", \"document\"]): Encoding type.\n * ``\"query\"``: Optimize for search queries (default).\n * ``\"document\"``: Optimize for indexed documents.\n This distinction is important for asymmetric retrieval tasks.\n\n Attributes:\n model (str): The DashScope model name being used.\n encoding_type (str): The encoding type (\"query\" or \"document\").\n\n Raises:\n ValueError: If API key is not provided and not found in environment,\n or if API returns an error response.\n TypeError: If input to ``embed()`` is not a string.\n RuntimeError: If network error or DashScope service error occurs.\n\n Note:\n - Requires Python 3.10, 3.11, or 3.12\n - Requires the ``dashscope`` package: ``pip install dashscope``\n - Embedding results are cached (LRU cache, maxsize=10) to reduce API calls\n - Network connectivity to DashScope API endpoints is required\n - API usage may incur costs based on your DashScope subscription plan\n - Sparse vectors have only non-zero dimensions stored as dict\n - Output is sorted by indices (keys) in ascending order\n\n **Parameter Guidelines:**\n\n - Use ``encoding_type=\"query\"`` for search queries and\n ``encoding_type=\"document\"`` for indexed content to optimize\n asymmetric retrieval tasks.\n - For detailed API specifications, refer to:\n https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api\n\n Examples:\n >>> # Basic usage for query embedding\n >>> from zvec.extension import QwenSparseEmbedding\n >>> import os\n >>> os.environ[\"DASHSCOPE_API_KEY\"] = \"your-api-key\"\n >>>\n >>> query_emb = QwenSparseEmbedding(dimension=1024, encoding_type=\"query\")\n >>> query_vec = query_emb.embed(\"machine learning\")\n >>> type(query_vec)\n \n >>> len(query_vec) # Only non-zero dimensions\n 156\n\n >>> # Document embedding\n >>> doc_emb = QwenSparseEmbedding(dimension=1024, encoding_type=\"document\")\n >>> doc_vec = doc_emb.embed(\"Machine learning is a subset of AI\")\n >>> isinstance(doc_vec, dict)\n True\n\n >>> # Asymmetric retrieval example\n >>> query_vec = query_emb.embed(\"what causes aging fast\")\n >>> doc_vec = doc_emb.embed(\n ... \"UV-A light causes tanning, skin aging, and cataracts...\"\n ... )\n >>>\n >>> # Calculate similarity (dot product for sparse vectors)\n >>> similarity = sum(\n ... query_vec.get(k, 0) * doc_vec.get(k, 0)\n ... for k in set(query_vec) | set(doc_vec)\n ... )\n\n >>> # Output is sorted by indices\n >>> list(query_vec.items())[:5] # First 5 dimensions (by index)\n [(10, 0.45), (23, 0.87), (56, 0.32), (89, 1.12), (120, 0.65)]\n\n >>> # Hybrid retrieval (combining dense + sparse)\n >>> from zvec.extension import QwenDenseEmbedding\n >>> dense_emb = QwenDenseEmbedding(dimension=1024)\n >>> sparse_emb = QwenSparseEmbedding(dimension=1024)\n >>>\n >>> query = \"deep learning neural networks\"\n >>> dense_vec = dense_emb.embed(query) # [0.1, -0.3, 0.5, ...]\n >>> sparse_vec = sparse_emb.embed(query) # {12: 0.8, 45: 1.2, ...}\n\n >>> # Error handling\n >>> try:\n ... sparse_emb.embed(\"\") # Empty string\n ... except ValueError as e:\n ... print(f\"Error: {e}\")\n Error: Input text cannot be empty or whitespace only\n\n See Also:\n - ``SparseEmbeddingFunction``: Base class for sparse embeddings\n - ``QwenDenseEmbedding``: Dense embedding using Qwen API\n - ``DefaultSparseEmbedding``: Sparse embedding with SPLADE model\n \"\"\"\n\n def __init__(\n self,\n dimension: int,\n model: str = \"text-embedding-v4\",\n api_key: Optional[str] = None,\n **kwargs,\n ):\n \"\"\"Initialize the Qwen sparse embedding function.\n\n Args:\n dimension (int): Target embedding dimension.\n model (str): DashScope model name. Defaults to \"text-embedding-v4\".\n api_key (Optional[str]): API key or None to use environment variable.\n **kwargs: Additional DashScope API parameters. Supported options:\n - ``encoding_type`` (Literal[\"query\", \"document\"]): Encoding type.\n * ``\"query\"``: Optimize for search queries (default).\n * ``\"document\"``: Optimize for indexed documents.\n This distinction is important for asymmetric retrieval tasks.\n\n Raises:\n ValueError: If API key is not provided and not in environment.\n \"\"\"\n # Initialize base class for API connection\n QwenFunctionBase.__init__(self, model=model, api_key=api_key)\n\n self._dimension = dimension\n self._extra_params = kwargs\n\n @property\n def extra_params(self) -> dict:\n \"\"\"dict: Extra parameters for model-specific customization.\"\"\"\n return self._extra_params\n\n def __call__(self, input: TEXT) -> SparseVectorType:\n \"\"\"Make the embedding function callable.\"\"\"\n return self.embed(input)\n\n @lru_cache(maxsize=10)\n def embed(self, input: TEXT) -> SparseVectorType:\n \"\"\"Generate sparse embedding vector for the input text.\n\n This method calls the DashScope TextEmbedding API with sparse output type\n to convert input text into a sparse vector representation. The result is\n a dictionary where keys are dimension indices and values are importance\n weights (only non-zero values included).\n\n The embedding is optimized based on the ``encoding_type`` specified during\n initialization: \"query\" for search queries or \"document\" for indexed content.\n\n Args:\n input (TEXT): Input text string to embed. Must be non-empty after\n stripping whitespace. Maximum length depends on the model used\n (typically 2048-8192 tokens).\n\n Returns:\n SparseVectorType: A dictionary mapping dimension index to weight.\n Only non-zero dimensions are included. The dictionary is sorted\n by indices (keys) in ascending order for consistent output.\n Example: ``{10: 0.5, 245: 0.8, 1023: 1.2, 5678: 0.5}``\n\n Raises:\n TypeError: If ``input`` is not a string.\n ValueError: If input is empty/whitespace-only, or if the API returns\n an error or malformed response.\n RuntimeError: If network connectivity issues or DashScope service\n errors occur.\n\n Examples:\n >>> emb = QwenSparseEmbedding(dimension=1024, encoding_type=\"query\")\n >>> sparse_vec = emb.embed(\"machine learning\")\n >>> isinstance(sparse_vec, dict)\n True\n >>>\n >>> # Verify sorted output\n >>> keys = list(sparse_vec.keys())\n >>> keys == sorted(keys)\n True\n\n >>> # Error: empty input\n >>> emb.embed(\" \")\n ValueError: Input text cannot be empty or whitespace only\n\n >>> # Error: non-string input\n >>> emb.embed(123)\n TypeError: Expected 'input' to be str, got int\n\n Note:\n - This method is cached (maxsize=10). Identical inputs return cached results.\n - The cache is based on exact string match (case-sensitive).\n - Output dictionary is always sorted by indices for consistency.\n \"\"\"\n if not isinstance(input, TEXT):\n raise TypeError(f\"Expected 'input' to be str, got {type(input).__name__}\")\n\n input = input.strip()\n if not input:\n raise ValueError(\"Input text cannot be empty or whitespace only\")\n\n # Call API with sparse output type\n output = self._call_text_embedding_api(\n input=input,\n dimension=self._dimension,\n output_type=\"sparse\",\n text_type=self.extra_params.get(\"encoding_type\", \"query\"),\n )\n\n embeddings = output.get(\"embeddings\")\n if not isinstance(embeddings, list):\n raise ValueError(\n \"Invalid API response: 'embeddings' field is missing or not a list\"\n )\n\n if len(embeddings) != 1:\n raise ValueError(\n f\"Expected exactly 1 embedding in response, got {len(embeddings)}\"\n )\n\n first_emb = embeddings[0]\n if not isinstance(first_emb, dict):\n raise ValueError(\"Invalid API response: embedding item is not a dictionary\")\n\n sparse_embedding = first_emb.get(\"sparse_embedding\")\n if not isinstance(sparse_embedding, list):\n raise ValueError(\n \"Invalid API response: 'sparse_embedding' field is missing or not a list\"\n )\n\n # Parse sparse embedding: convert array of {index, value, token} to dict\n sparse_dict = {}\n for item in sparse_embedding:\n if not isinstance(item, dict):\n raise ValueError(\n \"Invalid API response: sparse_embedding item is not a dictionary\"\n )\n\n index = item.get(\"index\")\n value = item.get(\"value\")\n\n if index is None or value is None:\n raise ValueError(\n \"Invalid API response: sparse_embedding item missing 'index' or 'value'\"\n )\n\n # Convert to int and float, filter positive values\n idx = int(index)\n val = float(value)\n if val > 0:\n sparse_dict[idx] = val\n\n # Sort by indices (keys) to ensure consistent ordering\n return dict(sorted(sparse_dict.items()))\n" }, { "path": "python/zvec/extension/qwen_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport os\nfrom http import HTTPStatus\nfrom typing import Optional\n\nfrom ..common.constants import TEXT\nfrom ..tool import require_module\n\n\nclass QwenFunctionBase:\n \"\"\"Base class for Qwen (DashScope) functions.\n\n This base class provides common functionality for calling DashScope APIs\n and handling responses. It supports embeddings (dense and sparse) and\n re-ranking operations.\n\n This class is not meant to be used directly. Use concrete implementations:\n - ``QwenDenseEmbedding`` for dense embeddings\n - ``QwenSparseEmbedding`` for sparse embeddings\n - ``QwenReRanker`` for semantic re-ranking\n\n Args:\n model (str): DashScope model identifier.\n api_key (Optional[str]): DashScope API authentication key.\n\n Note:\n - This is an internal base class for code reuse across Qwen features\n - Subclasses should inherit from appropriate Protocol/ABC\n - Provides unified API connection and response handling\n \"\"\"\n\n def __init__(\n self,\n model: str,\n api_key: Optional[str] = None,\n ):\n \"\"\"Initialize the base Qwen embedding functionality.\n\n Args:\n model (str): DashScope model name.\n api_key (Optional[str]): API key or None to use environment variable.\n\n Raises:\n ValueError: If API key is not provided and not in environment.\n \"\"\"\n self._model = model\n self._api_key = api_key or os.environ.get(\"DASHSCOPE_API_KEY\")\n if not self._api_key:\n raise ValueError(\n \"DashScope API key is required. Please provide 'api_key' parameter \"\n \"or set the 'DASHSCOPE_API_KEY' environment variable.\"\n )\n\n @property\n def model(self) -> str:\n \"\"\"str: The DashScope embedding model name currently in use.\"\"\"\n return self._model\n\n def _get_connection(self):\n \"\"\"Establish connection to DashScope API.\n\n Returns:\n module: The dashscope module with API key configured.\n\n Raises:\n ImportError: If dashscope package is not installed.\n \"\"\"\n dashscope = require_module(\"dashscope\")\n dashscope.api_key = self._api_key\n return dashscope\n\n def _call_text_embedding_api(\n self,\n input: TEXT,\n dimension: int,\n output_type: str,\n text_type: Optional[str] = None,\n ) -> dict:\n \"\"\"Call DashScope TextEmbedding API.\n\n Args:\n input (TEXT): Input text to embed.\n dimension (int): Target embedding dimension.\n output_type (str): Output type (\"dense\" or \"sparse\").\n text_type (Optional[str]): Text type (\"query\" or \"document\").\n\n Returns:\n dict: API response output field.\n\n Raises:\n RuntimeError: If API call fails.\n ValueError: If API returns error response.\n \"\"\"\n try:\n # Prepare API call parameters\n call_params = {\n \"model\": self.model,\n \"input\": input,\n \"dimension\": dimension,\n \"output_type\": output_type,\n }\n\n # Add optional text_type parameter if provided\n if text_type is not None:\n call_params[\"text_type\"] = text_type\n\n resp = self._get_connection().TextEmbedding.call(**call_params)\n except Exception as e:\n raise RuntimeError(f\"Failed to call DashScope API: {e!s}\") from e\n\n if resp.status_code != HTTPStatus.OK:\n error_msg = getattr(resp, \"message\", \"Unknown error\")\n error_code = getattr(resp, \"code\", \"N/A\")\n raise ValueError(\n f\"DashScope API error: [Code={error_code}, \"\n f\"Status={resp.status_code}] {error_msg}\"\n )\n\n output = getattr(resp, \"output\", None)\n if not isinstance(output, dict):\n raise ValueError(\n \"Invalid API response: missing or malformed 'output' field\"\n )\n\n return output\n\n def _call_rerank_api(\n self,\n query: str,\n documents: list[str],\n top_n: int,\n ) -> dict:\n \"\"\"Call DashScope TextReRank API.\n\n Args:\n query (str): Query text for semantic matching.\n documents (list[str]): List of document texts to re-rank.\n top_n (int): Maximum number of documents to return.\n\n Returns:\n dict: API response output field containing re-ranked results.\n\n Raises:\n RuntimeError: If API call fails.\n ValueError: If API returns error response.\n \"\"\"\n try:\n resp = self._get_connection().TextReRank.call(\n model=self.model,\n query=query,\n documents=documents,\n top_n=top_n,\n return_documents=False,\n )\n except Exception as e:\n raise RuntimeError(f\"Failed to call DashScope API: {e!s}\") from e\n\n if resp.status_code != HTTPStatus.OK:\n error_msg = getattr(resp, \"message\", \"Unknown error\")\n error_code = getattr(resp, \"code\", \"N/A\")\n raise ValueError(\n f\"DashScope API error: [Code={error_code}, \"\n f\"Status={resp.status_code}] {error_msg}\"\n )\n\n output = getattr(resp, \"output\", None)\n if not isinstance(output, dict):\n raise ValueError(\n \"Invalid API response: missing or malformed 'output' field\"\n )\n\n return output\n" }, { "path": "python/zvec/extension/qwen_rerank_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom typing import Optional\n\nfrom ..model.doc import Doc\nfrom .qwen_function import QwenFunctionBase\nfrom .rerank_function import RerankFunction\n\n\nclass QwenReRanker(QwenFunctionBase, RerankFunction):\n \"\"\"Re-ranker using Qwen (DashScope) cross-encoder API for semantic re-ranking.\n\n This re-ranker leverages DashScope's TextReRank service to perform\n cross-encoder style re-ranking. It sends query and document pairs to the\n API and receives relevance scores based on deep semantic understanding.\n\n The re-ranker is suitable for single-vector or multi-vector search scenarios\n where semantic relevance to a specific query is required.\n\n Args:\n query (str): Query text for semantic re-ranking. **Required**.\n topn (int, optional): Maximum number of documents to return after re-ranking.\n Defaults to 10.\n rerank_field (str): Document field name to use as re-ranking input text.\n **Required** (e.g., \"content\", \"title\", \"body\").\n model (str, optional): DashScope re-ranking model identifier.\n Defaults to ``\"gte-rerank-v2\"``.\n api_key (Optional[str], optional): DashScope API authentication key.\n If not provided, reads from ``DASHSCOPE_API_KEY`` environment variable.\n\n Raises:\n ValueError: If ``query`` is empty/None, ``rerank_field`` is None,\n or API key is not available.\n\n Note:\n - Requires ``dashscope`` Python package installed\n - Documents without valid content in ``rerank_field`` are skipped\n - API rate limits and quotas apply per DashScope subscription\n\n Example:\n >>> reranker = QwenReRanker(\n ... query=\"machine learning algorithms\",\n ... topn=5,\n ... rerank_field=\"content\",\n ... model=\"gte-rerank-v2\",\n ... api_key=\"your-api-key\"\n ... )\n >>> # Use in collection.query(reranker=reranker)\n \"\"\"\n\n def __init__(\n self,\n query: Optional[str] = None,\n topn: int = 10,\n rerank_field: Optional[str] = None,\n model: str = \"gte-rerank-v2\",\n api_key: Optional[str] = None,\n ):\n \"\"\"Initialize QwenReRanker with query and configuration.\n\n Args:\n query (Optional[str]): Query text for semantic matching. Required.\n topn (int): Number of top results to return.\n rerank_field (Optional[str]): Document field for re-ranking input.\n model (str): DashScope model name.\n api_key (Optional[str]): API key or None to use environment variable.\n\n Raises:\n ValueError: If query is empty or API key is unavailable.\n \"\"\"\n QwenFunctionBase.__init__(self, model=model, api_key=api_key)\n RerankFunction.__init__(self, topn=topn, rerank_field=rerank_field)\n\n if not query:\n raise ValueError(\"Query is required for QwenReRanker\")\n self._query = query\n\n @property\n def query(self) -> str:\n \"\"\"str: Query text used for semantic re-ranking.\"\"\"\n return self._query\n\n def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:\n \"\"\"Re-rank documents using Qwen's TextReRank API.\n\n Sends document texts to DashScope TextReRank service along with the query.\n Returns documents sorted by relevance scores from the cross-encoder model.\n\n Args:\n query_results (dict[str, list[Doc]]): Mapping from vector field names\n to lists of retrieved documents. Documents from all fields are\n deduplicated and re-ranked together.\n\n Returns:\n list[Doc]: Re-ranked documents (up to ``topn``) with updated ``score``\n fields containing relevance scores from the API.\n\n Raises:\n ValueError: If no valid documents are found or API call fails.\n\n Note:\n - Duplicate documents (same ID) across fields are processed once\n - Documents with empty/missing ``rerank_field`` content are skipped\n - Returned scores are relevance scores from the cross-encoder model\n \"\"\"\n if not query_results:\n return []\n\n # Collect and deduplicate documents\n id_to_doc: dict[str, Doc] = {}\n doc_ids: list[str] = []\n contents: list[str] = []\n\n for _, query_result in query_results.items():\n for doc in query_result:\n doc_id = doc.id\n if doc_id in id_to_doc:\n continue\n\n # Extract text content from specified field\n field_value = doc.field(self.rerank_field)\n rank_content = str(field_value).strip() if field_value else \"\"\n if not rank_content:\n continue\n\n id_to_doc[doc_id] = doc\n doc_ids.append(doc_id)\n contents.append(rank_content)\n\n if not contents:\n raise ValueError(\"No documents to rerank\")\n\n # Call DashScope TextReRank API\n output = self._call_rerank_api(\n query=self.query,\n documents=contents,\n top_n=self.topn,\n )\n\n # Build result list with updated scores\n results: list[Doc] = []\n for item in output[\"results\"]:\n idx = item[\"index\"]\n doc_id = doc_ids[idx]\n doc = id_to_doc[doc_id]\n new_doc = doc._replace(score=item[\"relevance_score\"])\n results.append(new_doc)\n\n return results\n" }, { "path": "python/zvec/extension/rerank_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom abc import ABC, abstractmethod\nfrom typing import Optional\n\nfrom ..model.doc import Doc\n\n\nclass RerankFunction(ABC):\n \"\"\"Abstract base class for re-ranking search results.\n\n Re-rankers refine the output of one or more vector queries by applying\n a secondary scoring strategy. They are used in the ``query()`` method of\n ``Collection`` via the ``reranker`` parameter.\n\n Args:\n topn (int, optional): Number of top documents to return after re-ranking.\n Defaults to 10.\n rerank_field (Optional[str], optional): Field name used as input for\n re-ranking (e.g., document title or body). Defaults to None.\n\n Note:\n Subclasses must implement the ``rerank()`` method.\n \"\"\"\n\n def __init__(\n self,\n topn: int = 10,\n rerank_field: Optional[str] = None,\n ):\n self._topn = topn\n self._rerank_field = rerank_field\n\n @property\n def topn(self) -> int:\n \"\"\"int: Number of top documents to return after re-ranking.\"\"\"\n return self._topn\n\n @property\n def rerank_field(self) -> Optional[str]:\n \"\"\"Optional[str]: Field name used as re-ranking input.\"\"\"\n return self._rerank_field\n\n @abstractmethod\n def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:\n \"\"\"Re-rank documents from one or more vector queries.\n\n Args:\n query_results (dict[str, list[Doc]]): Mapping from vector field name\n to list of retrieved documents (sorted by relevance).\n\n Returns:\n list[Doc]: Re-ranked list of documents (length ≤ ``topn``),\n with updated ``score`` fields.\n \"\"\"\n ...\n" }, { "path": "python/zvec/extension/sentence_transformer_embedding_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom typing import ClassVar, Literal, Optional\n\nimport numpy as np\n\nfrom ..common.constants import TEXT, DenseVectorType, SparseVectorType\nfrom .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction\nfrom .sentence_transformer_function import SentenceTransformerFunctionBase\n\n\nclass DefaultLocalDenseEmbedding(\n SentenceTransformerFunctionBase, DenseEmbeddingFunction[TEXT]\n):\n \"\"\"Default local dense embedding using all-MiniLM-L6-v2 model.\n\n This is the default implementation for dense text embedding that uses the\n ``all-MiniLM-L6-v2`` model from Hugging Face by default. This model provides\n a good balance between speed and quality for general-purpose text embedding.\n\n The class provides text-to-vector dense embedding capabilities using the\n sentence-transformers library. It supports models from Hugging Face Hub and\n ModelScope, runs locally without API calls, and supports CPU/GPU acceleration.\n\n The model produces 384-dimensional embeddings and is optimized for semantic\n similarity tasks. It runs locally without requiring API keys.\n\n Args:\n model_source (Literal[\"huggingface\", \"modelscope\"], optional): Model source.\n - ``\"huggingface\"``: Use Hugging Face Hub (default, for international users)\n - ``\"modelscope\"``: Use ModelScope (recommended for users in China)\n Defaults to ``\"huggingface\"``.\n device (Optional[str], optional): Device to run the model on.\n Options: ``\"cpu\"``, ``\"cuda\"``, ``\"mps\"`` (for Apple Silicon), or ``None``\n for automatic detection. Defaults to ``None``.\n normalize_embeddings (bool, optional): Whether to normalize embeddings to\n unit length (L2 normalization). Useful for cosine similarity.\n Defaults to ``True``.\n batch_size (int, optional): Batch size for encoding. Defaults to ``32``.\n **kwargs: Additional parameters for future extension.\n\n Attributes:\n dimension (int): Always 384 for both models.\n model_name (str): \"all-MiniLM-L6-v2\" (HF) or \"iic/nlp_gte_sentence-embedding_chinese-small\" (MS).\n model_source (str): The model source being used.\n device (str): The device the model is running on.\n\n Raises:\n ValueError: If the model cannot be loaded or input is invalid.\n TypeError: If input to ``embed()`` is not a string.\n RuntimeError: If model inference fails.\n\n Note:\n - Requires Python 3.10, 3.11, or 3.12\n - Requires the ``sentence-transformers`` package:\n ``pip install sentence-transformers``\n - For ModelScope, also requires: ``pip install modelscope``\n - First run downloads the model (~50-80MB) from chosen source\n - Hugging Face cache: ``~/.cache/torch/sentence_transformers/``\n - ModelScope cache: ``~/.cache/modelscope/hub/``\n - No API keys or network required after initial download\n - Inference speed: ~1000 sentences/sec on CPU, ~10000 on GPU\n\n **For users in China:**\n\n If you encounter Hugging Face access issues, use ModelScope instead:\n\n .. code-block:: python\n\n # Recommended for users in China\n emb = DefaultLocalDenseEmbedding(model_source=\"modelscope\")\n\n Alternatively, use Hugging Face mirror:\n\n .. code-block:: bash\n\n export HF_ENDPOINT=https://hf-mirror.com\n # Then use default Hugging Face mode\n\n Examples:\n >>> # Basic usage with Hugging Face (default)\n >>> from zvec.extension import DefaultLocalDenseEmbedding\n >>>\n >>> emb_func = DefaultLocalDenseEmbedding()\n >>> vector = emb_func.embed(\"Hello, world!\")\n >>> len(vector)\n 384\n >>> isinstance(vector, list)\n True\n\n >>> # Recommended for users in China (uses ModelScope)\n >>> emb_func = DefaultLocalDenseEmbedding(model_source=\"modelscope\")\n >>> vector = emb_func.embed(\"你好，世界！\") # Works well with Chinese text\n >>> len(vector)\n 384\n\n >>> # Alternative for China users: Use Hugging Face mirror\n >>> import os\n >>> os.environ[\"HF_ENDPOINT\"] = \"https://hf-mirror.com\"\n >>> emb_func = DefaultLocalDenseEmbedding() # Uses HF mirror\n >>> vector = emb_func.embed(\"Hello, world!\")\n\n >>> # Using GPU for faster inference\n >>> emb_func = DefaultLocalDenseEmbedding(device=\"cuda\")\n >>> vector = emb_func(\"Machine learning is fascinating\")\n >>> # Normalized vector has unit length\n >>> import numpy as np\n >>> np.linalg.norm(vector)\n 1.0\n\n >>> # Batch processing\n >>> texts = [\"First text\", \"Second text\", \"Third text\"]\n >>> vectors = [emb_func.embed(text) for text in texts]\n >>> len(vectors)\n 3\n >>> all(len(v) == 384 for v in vectors)\n True\n\n >>> # Semantic similarity\n >>> v1 = emb_func.embed(\"The cat sits on the mat\")\n >>> v2 = emb_func.embed(\"A feline rests on a rug\")\n >>> v3 = emb_func.embed(\"Python programming\")\n >>> similarity_high = np.dot(v1, v2) # Similar sentences\n >>> similarity_low = np.dot(v1, v3) # Different topics\n >>> similarity_high > similarity_low\n True\n\n >>> # Error handling\n >>> try:\n ... emb_func.embed(\"\") # Empty string\n ... except ValueError as e:\n ... print(f\"Error: {e}\")\n Error: Input text cannot be empty or whitespace only\n\n See Also:\n - ``DenseEmbeddingFunction``: Base class for dense embeddings\n - ``DefaultLocalSparseEmbedding``: Sparse embedding with SPLADE\n - ``QwenDenseEmbedding``: Alternative using Qwen API\n \"\"\"\n\n def __init__(\n self,\n model_source: Literal[\"huggingface\", \"modelscope\"] = \"huggingface\",\n device: Optional[str] = None,\n normalize_embeddings: bool = True,\n batch_size: int = 32,\n **kwargs,\n ):\n \"\"\"Initialize with all-MiniLM-L6-v2 model.\n\n Args:\n model_source (Literal[\"huggingface\", \"modelscope\"]): Model source.\n Defaults to \"huggingface\".\n device (Optional[str]): Target device (\"cpu\", \"cuda\", \"mps\", or None).\n Defaults to None (automatic detection).\n normalize_embeddings (bool): Whether to L2-normalize output vectors.\n Defaults to True.\n batch_size (int): Batch size for encoding. Defaults to 32.\n **kwargs: Additional parameters for future extension.\n\n Raises:\n ImportError: If sentence-transformers or modelscope is not installed.\n ValueError: If model cannot be loaded.\n \"\"\"\n # Use different models based on source\n if model_source == \"modelscope\":\n # Use Chinese-optimized model for ModelScope (better for Chinese text)\n model_name = \"iic/nlp_gte_sentence-embedding_chinese-small\"\n else:\n model_name = \"all-MiniLM-L6-v2\"\n\n # Initialize base class for model loading\n SentenceTransformerFunctionBase.__init__(\n self, model_name=model_name, model_source=model_source, device=device\n )\n\n self._normalize_embeddings = normalize_embeddings\n self._batch_size = batch_size\n\n # Load model and get dimension\n model = self._get_model()\n self._dimension = model.get_sentence_embedding_dimension()\n\n # Store extra parameters\n self._extra_params = kwargs\n\n @property\n def dimension(self) -> int:\n \"\"\"int: The expected dimensionality of the embedding vector.\"\"\"\n return self._dimension\n\n @property\n def extra_params(self) -> dict:\n \"\"\"dict: Extra parameters for model-specific customization.\"\"\"\n return self._extra_params\n\n def __call__(self, input: str) -> DenseVectorType:\n \"\"\"Make the embedding function callable.\"\"\"\n return self.embed(input)\n\n def embed(self, input: str) -> DenseVectorType:\n \"\"\"Generate dense embedding vector for the input text.\n\n This method uses the Sentence Transformer model to convert input text\n into a dense vector representation. The model runs locally without\n requiring API calls.\n\n Args:\n input (str): Input text string to embed. Must be non-empty after\n stripping whitespace. Maximum length depends on the model used\n (typically 128-512 tokens for most models).\n\n Returns:\n DenseVectorType: A list of floats representing the embedding vector.\n Length equals ``self.dimension``. If ``normalize_embeddings=True``,\n the vector has unit length. Example:\n ``[0.123, -0.456, 0.789, ...]``\n\n Raises:\n TypeError: If ``input`` is not a string.\n ValueError: If input is empty or whitespace-only.\n RuntimeError: If model inference fails.\n\n Examples:\n >>> emb = DefaultLocalDenseEmbedding()\n >>> vector = emb.embed(\"Natural language processing\")\n >>> len(vector)\n 384\n >>> isinstance(vector[0], float)\n True\n\n >>> # Normalized vectors have unit length\n >>> import numpy as np\n >>> emb = DefaultLocalDenseEmbedding(normalize_embeddings=True)\n >>> vector = emb.embed(\"Test sentence\")\n >>> np.linalg.norm(vector)\n 1.0\n\n >>> # Error: empty input\n >>> emb.embed(\" \")\n ValueError: Input text cannot be empty or whitespace only\n\n >>> # Error: non-string input\n >>> emb.embed(123)\n TypeError: Expected 'input' to be str, got int\n\n >>> # Semantic similarity example\n >>> v1 = emb.embed(\"The cat sits on the mat\")\n >>> v2 = emb.embed(\"A feline rests on a rug\")\n >>> similarity = np.dot(v1, v2) # High similarity due to semantic meaning\n >>> similarity > 0.7\n True\n\n Note:\n - First call may be slower due to model loading\n - Subsequent calls are much faster as the model stays in memory\n - For batch processing, consider encoding multiple texts together\n (though this method handles single texts only)\n - GPU acceleration provides 5-10x speedup over CPU\n \"\"\"\n if not isinstance(input, str):\n raise TypeError(f\"Expected 'input' to be str, got {type(input).__name__}\")\n\n input = input.strip()\n if not input:\n raise ValueError(\"Input text cannot be empty or whitespace only\")\n\n try:\n model = self._get_model()\n embedding = model.encode(\n input,\n convert_to_numpy=True,\n normalize_embeddings=self._normalize_embeddings,\n batch_size=self._batch_size,\n )\n\n # Convert numpy array to list\n if isinstance(embedding, np.ndarray):\n embedding_list = embedding.tolist()\n else:\n embedding_list = list(embedding)\n\n # Validate dimension\n if len(embedding_list) != self.dimension:\n raise ValueError(\n f\"Dimension mismatch: expected {self.dimension}, \"\n f\"got {len(embedding_list)}\"\n )\n\n return embedding_list\n\n except Exception as e:\n if isinstance(e, (TypeError, ValueError)):\n raise\n raise RuntimeError(f\"Failed to generate embedding: {e!s}\") from e\n\n\nclass DefaultLocalSparseEmbedding(\n SentenceTransformerFunctionBase, SparseEmbeddingFunction[TEXT]\n):\n \"\"\"Default local sparse embedding using SPLADE model.\n\n This class provides sparse vector embedding using the SPLADE (SParse Lexical\n AnD Expansion) model. SPLADE generates sparse, interpretable representations\n where each dimension corresponds to a vocabulary term with learned importance\n weights. It's ideal for lexical matching, BM25-style retrieval, and hybrid\n search scenarios.\n\n The default model is ``naver/splade-cocondenser-ensembledistil``, which is\n publicly available without authentication. It produces sparse vectors with\n thousands of dimensions but only hundreds of non-zero values, making them\n efficient for storage and retrieval while maintaining strong lexical matching.\n\n **Model Caching:**\n\n This class uses class-level caching to share the SPLADE model across all instances\n with the same configuration (model_source, device). This significantly reduces\n memory usage when creating multiple instances for different encoding types\n (query vs document).\n\n **Cache Management:**\n\n The class provides methods to manage the model cache:\n\n - ``clear_cache()``: Clear all cached models to free memory\n - ``get_cache_info()``: Get information about cached models\n - ``remove_from_cache(model_source, device)``: Remove a specific model from cache\n\n .. note::\n **Why not use splade-v3?**\n\n The newer ``naver/splade-v3`` model is gated (requires access approval).\n We use ``naver/splade-cocondenser-ensembledistil`` instead.\n\n **To use splade-v3 (if you have access):**\n\n 1. Request access at https://huggingface.co/naver/splade-v3\n 2. Get your Hugging Face token from https://huggingface.co/settings/tokens\n 3. Set environment variable:\n\n .. code-block:: bash\n\n export HF_TOKEN=\"your_huggingface_token\"\n\n 4. Or login programmatically:\n\n .. code-block:: python\n\n from huggingface_hub import login\n login(token=\"your_huggingface_token\")\n\n 5. To use a custom SPLADE model, you can subclass this class and override\n the model_name in ``__init__``, or create your own implementation\n inheriting from ``SentenceTransformerFunctionBase`` and\n ``SparseEmbeddingFunction``.\n\n Args:\n model_source (Literal[\"huggingface\", \"modelscope\"], optional): Model source.\n Defaults to ``\"huggingface\"``. ModelScope support may vary for SPLADE models.\n device (Optional[str], optional): Device to run the model on.\n Options: ``\"cpu\"``, ``\"cuda\"``, ``\"mps\"`` (for Apple Silicon), or ``None``\n for automatic detection. Defaults to ``None``.\n encoding_type (Literal[\"query\", \"document\"], optional): Encoding type.\n - ``\"query\"``: Optimize for search queries (default)\n - ``\"document\"``: Optimize for indexed documents\n **kwargs: Additional parameters (currently unused, for future extension).\n\n Attributes:\n model_name (str): Model identifier.\n model_source (str): The model source being used.\n device (str): The device the model is running on.\n\n Raises:\n ValueError: If the model cannot be loaded or input is invalid.\n TypeError: If input to ``embed()`` is not a string.\n RuntimeError: If model inference fails.\n\n Note:\n - Requires Python 3.10, 3.11, or 3.12\n - Requires the ``sentence-transformers`` package:\n ``pip install sentence-transformers``\n - First run downloads the model (~100MB) from Hugging Face\n - Cache location: ``~/.cache/torch/sentence_transformers/``\n - No API keys or authentication required\n - Sparse vectors have ~30k dimensions but only ~100-200 non-zero values\n - Best combined with dense embeddings for hybrid retrieval\n\n **SPLADE vs Dense Embeddings:**\n\n - **Dense**: Continuous semantic vectors, good for semantic similarity\n - **Sparse**: Lexical keyword-based, interpretable, good for exact matching\n - **Hybrid**: Combine both for best retrieval performance\n\n Examples:\n >>> # Memory-efficient: both instances share the same model (~200MB)\n >>> from zvec.extension import DefaultLocalSparseEmbedding\n >>>\n >>> # Query embedding\n >>> query_emb = DefaultLocalSparseEmbedding(encoding_type=\"query\")\n >>> query_vec = query_emb.embed(\"machine learning algorithms\")\n >>> type(query_vec)\n \n >>> len(query_vec) # Only non-zero dimensions\n 156\n\n >>> # Document embedding (shares model with query_emb)\n >>> doc_emb = DefaultLocalSparseEmbedding(encoding_type=\"document\")\n >>> doc_vec = doc_emb.embed(\"Machine learning is a subset of AI\")\n >>> # Total memory: ~200MB (not 400MB) thanks to model caching\n\n >>> # Asymmetric retrieval example\n >>> query_vec = query_emb.embed(\"what causes aging fast\")\n >>> doc_vec = doc_emb.embed(\n ... \"UV-A light causes tanning, skin aging, and cataracts...\"\n ... )\n >>>\n >>> # Calculate similarity (dot product for sparse vectors)\n >>> similarity = sum(\n ... query_vec.get(k, 0) * doc_vec.get(k, 0)\n ... for k in set(query_vec) | set(doc_vec)\n ... )\n\n >>> # Batch processing\n >>> queries = [\"query 1\", \"query 2\", \"query 3\"]\n >>> query_vecs = [query_emb.embed(q) for q in queries]\n >>>\n >>> documents = [\"doc 1\", \"doc 2\", \"doc 3\"]\n >>> doc_vecs = [doc_emb.embed(d) for d in documents]\n\n >>> # Inspecting sparse dimensions (output is sorted by indices)\n >>> query_vec = query_emb.embed(\"machine learning\")\n >>> list(query_vec.items())[:5] # First 5 dimensions (by index)\n [(10, 0.45), (23, 0.87), (56, 0.32), (89, 1.12), (120, 0.65)]\n >>>\n >>> # Sort by weight to find most important terms\n >>> sorted_by_weight = sorted(query_vec.items(), key=lambda x: x[1], reverse=True)\n >>> top_5 = sorted_by_weight[:5] # Top 5 most important terms\n >>> top_5\n [(1023, 1.45), (245, 1.23), (8901, 0.98), (5678, 0.87), (12034, 0.76)]\n\n >>> # Using GPU for faster inference\n >>> sparse_emb = DefaultLocalSparseEmbedding(device=\"cuda\")\n >>> vector = sparse_emb.embed(\"natural language processing\")\n\n >>> # Hybrid retrieval example (combining dense + sparse)\n >>> from zvec.extension import DefaultDenseEmbedding\n >>> dense_emb = DefaultDenseEmbedding()\n >>> sparse_emb = DefaultLocalSparseEmbedding()\n >>>\n >>> query = \"deep learning neural networks\"\n >>> dense_vec = dense_emb.embed(query) # [0.1, -0.3, 0.5, ...]\n >>> sparse_vec = sparse_emb.embed(query) # {12: 0.8, 45: 1.2, ...}\n\n >>> # Error handling\n >>> try:\n ... sparse_emb.embed(\"\") # Empty string\n ... except ValueError as e:\n ... print(f\"Error: {e}\")\n Error: Input text cannot be empty or whitespace only\n\n >>> # Cache management\n >>> # Check cache status\n >>> info = DefaultLocalSparseEmbedding.get_cache_info()\n >>> print(f\"Cached models: {info['cached_models']}\")\n Cached models: 1\n >>>\n >>> # Clear cache to free memory\n >>> DefaultLocalSparseEmbedding.clear_cache()\n >>> info = DefaultLocalSparseEmbedding.get_cache_info()\n >>> print(f\"Cached models: {info['cached_models']}\")\n Cached models: 0\n >>>\n >>> # Remove specific model from cache\n >>> query_emb = DefaultLocalSparseEmbedding() # Creates CPU model\n >>> cuda_emb = DefaultLocalSparseEmbedding(device=\"cuda\") # Creates CUDA model\n >>> info = DefaultLocalSparseEmbedding.get_cache_info()\n >>> print(f\"Cached models: {info['cached_models']}\")\n Cached models: 2\n >>>\n >>> # Remove only CPU model\n >>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device=None)\n >>> print(f\"Removed: {removed}\")\n True\n >>> info = DefaultLocalSparseEmbedding.get_cache_info()\n >>> print(f\"Cached models: {info['cached_models']}\")\n Cached models: 1\n\n See Also:\n - ``SparseEmbeddingFunction``: Base class for sparse embeddings\n - ``DefaultDenseEmbedding``: Dense embedding with all-MiniLM-L6-v2\n - ``QwenDenseEmbedding``: Alternative using Qwen API\n\n References:\n - SPLADE Paper: https://arxiv.org/abs/2109.10086\n - Model: https://huggingface.co/naver/splade-cocondenser-ensembledistil\n \"\"\"\n\n # Class-level model cache: {(model_name, model_source, device): model}\n # Shared across all DefaultLocalSparseEmbedding instances to save memory\n _model_cache: ClassVar[dict] = {}\n\n @classmethod\n def clear_cache(cls) -> None:\n \"\"\"Clear all cached SPLADE models from memory.\n\n This is useful for:\n - Freeing memory when models are no longer needed\n - Forcing a fresh model reload\n - Testing and debugging\n Examples:\n >>> # Clear cache to free memory\n >>> DefaultLocalSparseEmbedding.clear_cache()\n\n >>> # Or in tests to ensure fresh model loading\n >>> def test_something():\n ... DefaultLocalSparseEmbedding.clear_cache()\n ... emb = DefaultLocalSparseEmbedding()\n ... # Test with fresh model\n \"\"\"\n cls._model_cache.clear()\n\n @classmethod\n def get_cache_info(cls) -> dict:\n \"\"\"Get information about currently cached models.\n\n Returns:\n dict: Dictionary with cache statistics:\n - cached_models (int): Number of cached model instances\n - cache_keys (list): List of cache keys (model_name, model_source, device)\n\n Examples:\n >>> info = DefaultLocalSparseEmbedding.get_cache_info()\n >>> print(f\"Cached models: {info['cached_models']}\")\n Cached models: 2\n >>> print(f\"Cache keys: {info['cache_keys']}\")\n Cache keys: [('naver/splade-cocondenser-ensembledistil', 'huggingface', None),\n ('naver/splade-cocondenser-ensembledistil', 'huggingface', 'cuda')]\n \"\"\"\n return {\n \"cached_models\": len(cls._model_cache),\n \"cache_keys\": list(cls._model_cache.keys()),\n }\n\n @classmethod\n def remove_from_cache(\n cls, model_source: str = \"huggingface\", device: Optional[str] = None\n ) -> bool:\n \"\"\"Remove a specific model from cache.\n\n Args:\n model_source (str): Model source (\"huggingface\" or \"modelscope\").\n Defaults to \"huggingface\".\n device (Optional[str]): Device identifier. Defaults to None.\n\n Returns:\n bool: True if model was found and removed, False otherwise.\n\n Examples:\n >>> # Remove CPU model from cache\n >>> removed = DefaultLocalSparseEmbedding.remove_from_cache()\n >>> print(f\"Removed: {removed}\")\n True\n\n >>> # Remove CUDA model from cache\n >>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device=\"cuda\")\n >>> print(f\"Removed: {removed}\")\n True\n \"\"\"\n model_name = \"naver/splade-cocondenser-ensembledistil\"\n cache_key = (model_name, model_source, device)\n\n if cache_key in cls._model_cache:\n del cls._model_cache[cache_key]\n return True\n return False\n\n def __init__(\n self,\n model_source: Literal[\"huggingface\", \"modelscope\"] = \"huggingface\",\n device: Optional[str] = None,\n encoding_type: Literal[\"query\", \"document\"] = \"query\",\n **kwargs,\n ):\n \"\"\"Initialize with SPLADE model.\n\n Args:\n model_source (Literal[\"huggingface\", \"modelscope\"]): Model source.\n Defaults to \"huggingface\".\n device (Optional[str]): Target device (\"cpu\", \"cuda\", \"mps\", or None).\n Defaults to None (automatic detection).\n encoding_type (Literal[\"query\", \"document\"]): Encoding type for embeddings.\n - \"query\": Optimize for search queries (default)\n - \"document\": Optimize for indexed documents\n This distinction is important for asymmetric retrieval tasks.\n **kwargs: Additional parameters (reserved for future use).\n\n Raises:\n ImportError: If sentence-transformers is not installed.\n ValueError: If model cannot be loaded.\n\n Note:\n Multiple instances with the same (model_source, device) configuration\n will share the same underlying model to save memory. Different\n instances can use different encoding_type settings while sharing\n the model.\n\n **Model Selection:**\n\n Uses ``naver/splade-cocondenser-ensembledistil`` instead of the newer\n ``naver/splade-v3`` because splade-v3 is a gated model requiring\n Hugging Face authentication. The cocondenser-ensembledistil variant:\n\n - Does not require authentication or API tokens\n - Is immediately available for all users\n - Provides comparable retrieval performance (~2% difference)\n - Avoids \"Access to model is restricted\" errors\n\n If you need splade-v3 and have obtained access, you can subclass\n this class and override the model_name parameter.\n\n Examples:\n >>> # Both instances share the same model (saves memory)\n >>> query_emb = DefaultLocalSparseEmbedding(encoding_type=\"query\")\n >>> doc_emb = DefaultLocalSparseEmbedding(encoding_type=\"document\")\n >>> # Only one model is loaded in memory\n \"\"\"\n # Use publicly available SPLADE model (no gated access required)\n # Note: naver/splade-v3 requires authentication, so we use the\n # cocondenser-ensembledistil variant which is publicly accessible\n model_name = \"naver/splade-cocondenser-ensembledistil\"\n\n # Initialize base class for model loading\n SentenceTransformerFunctionBase.__init__(\n self, model_name=model_name, model_source=model_source, device=device\n )\n\n self._encoding_type = encoding_type\n self._extra_params = kwargs\n\n # Create cache key for this model configuration\n self._cache_key = (model_name, model_source, device)\n\n # Load model to ensure it's available (will use cache if exists)\n self._get_model()\n\n @property\n def extra_params(self) -> dict:\n \"\"\"dict: Extra parameters for model-specific customization.\"\"\"\n return self._extra_params\n\n def __call__(self, input: str) -> SparseVectorType:\n \"\"\"Make the embedding function callable.\"\"\"\n return self.embed(input)\n\n def embed(self, input: str) -> SparseVectorType:\n \"\"\"Generate sparse embedding vector for the input text.\n\n This method uses the SPLADE model to convert input text into a sparse\n vector representation. The result is a dictionary where keys are dimension\n indices and values are importance weights (only non-zero values included).\n\n The embedding is optimized based on the ``encoding_type`` specified during\n initialization: \"query\" for search queries or \"document\" for indexed content.\n\n Args:\n input (str): Input text string to embed. Must be non-empty after\n stripping whitespace.\n\n Returns:\n SparseVectorType: A dictionary mapping dimension index to weight.\n Only non-zero dimensions are included. The dictionary is sorted\n by indices (keys) in ascending order for consistent output.\n Example: ``{10: 0.5, 245: 0.8, 1023: 1.2, 5678: 0.5}``\n\n Raises:\n TypeError: If ``input`` is not a string.\n ValueError: If input is empty or whitespace-only.\n RuntimeError: If model inference fails.\n\n Examples:\n >>> # Query embedding\n >>> query_emb = DefaultLocalSparseEmbedding(encoding_type=\"query\")\n >>> query_vec = query_emb.embed(\"machine learning\")\n >>> isinstance(query_vec, dict)\n True\n\n Note:\n - First call may be slower due to model loading\n - Subsequent calls are much faster as the model stays in memory\n - GPU acceleration provides significant speedup\n - Sparse vectors are memory-efficient (only store non-zero values)\n \"\"\"\n if not isinstance(input, str):\n raise TypeError(f\"Expected 'input' to be str, got {type(input).__name__}\")\n\n input = input.strip()\n if not input:\n raise ValueError(\"Input text cannot be empty or whitespace only\")\n\n try:\n model = self._get_model()\n\n # Use appropriate encoding method based on type\n if self._encoding_type == \"document\" and hasattr(model, \"encode_document\"):\n # Use document encoding\n sparse_matrix = model.encode_document([input])\n elif hasattr(model, \"encode_query\"):\n # Use query encoding (default)\n sparse_matrix = model.encode_query([input])\n else:\n # Fallback: manual implementation for older sentence-transformers\n return self._manual_sparse_encode(input)\n\n # Convert sparse matrix to dictionary\n # SPLADE returns shape [1, vocab_size] for single input\n\n # Check if it's a sparse matrix (duck typing - has toarray method)\n if hasattr(sparse_matrix, \"toarray\"):\n # Sparse matrix (CSR/CSC/etc.) - convert to dense array\n sparse_array = sparse_matrix[0].toarray().flatten()\n sparse_dict = {\n int(idx): float(val)\n for idx, val in enumerate(sparse_array)\n if val > 0\n }\n else:\n # Dense array format (numpy array or similar)\n if isinstance(sparse_matrix, np.ndarray):\n sparse_array = sparse_matrix[0]\n else:\n sparse_array = sparse_matrix\n\n sparse_dict = {\n int(idx): float(val)\n for idx, val in enumerate(sparse_array)\n if val > 0\n }\n\n # Sort by indices (keys) to ensure consistent ordering\n return dict(sorted(sparse_dict.items()))\n\n except Exception as e:\n if isinstance(e, (TypeError, ValueError)):\n raise\n raise RuntimeError(f\"Failed to generate sparse embedding: {e!s}\") from e\n\n def _manual_sparse_encode(self, input: str) -> SparseVectorType:\n \"\"\"Fallback manual SPLADE encoding for older sentence-transformers.\n\n Args:\n input (str): Input text to encode.\n\n Returns:\n SparseVectorType: Sparse vector as dictionary.\n \"\"\"\n import torch\n\n model = self._get_model()\n\n # Tokenize input\n features = model.tokenize([input])\n\n # Move to correct device\n features = {k: v.to(model.device) for k, v in features.items()}\n\n # Forward pass with no gradient\n with torch.no_grad():\n embeddings = model.forward(features)\n\n # Get logits from model output\n # SPLADE models typically output 'token_embeddings'\n if isinstance(embeddings, dict) and \"token_embeddings\" in embeddings:\n logits = embeddings[\"token_embeddings\"][0] # First batch item\n elif hasattr(embeddings, \"token_embeddings\"):\n logits = embeddings.token_embeddings[0]\n # Fallback: try to get first value\n elif isinstance(embeddings, dict):\n logits = next(iter(embeddings.values()))[0]\n else:\n logits = embeddings[0]\n\n # Apply SPLADE activation: log(1 + relu(x))\n relu_log = torch.log(1 + torch.relu(logits))\n\n # Max pooling over token dimension (reduce to vocab size)\n if relu_log.dim() > 1:\n sparse_vec, _ = torch.max(relu_log, dim=0)\n else:\n sparse_vec = relu_log\n\n # Convert to sparse dictionary (only non-zero values)\n sparse_vec_np = sparse_vec.cpu().numpy()\n sparse_dict = {\n int(idx): float(val) for idx, val in enumerate(sparse_vec_np) if val > 0\n }\n\n # Sort by indices (keys) to ensure consistent ordering\n return dict(sorted(sparse_dict.items()))\n\n def _get_model(self):\n \"\"\"Load or retrieve the SPLADE model from class-level cache.\n\n Returns:\n SentenceTransformer: The loaded SPLADE model instance.\n\n Raises:\n ImportError: If required packages are not installed.\n ValueError: If model cannot be loaded.\n\n Note:\n Models are cached at class level and shared across all instances\n with the same (model_name, model_source, device) configuration.\n This allows memory-efficient usage when creating multiple instances\n with different encoding_type settings.\n \"\"\"\n # Check class-level cache first\n if self._cache_key in self._model_cache:\n return self._model_cache[self._cache_key]\n\n # Use parent class method to load model\n model = super()._get_model()\n\n # Cache the model at class level\n self._model_cache[self._cache_key] = model\n\n return model\n" }, { "path": "python/zvec/extension/sentence_transformer_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom typing import Literal, Optional\n\nfrom ..tool import require_module\n\n\nclass SentenceTransformerFunctionBase:\n \"\"\"Base class for Sentence Transformer functions (both dense and sparse).\n\n This base class provides common functionality for loading and managing\n sentence-transformers models from Hugging Face or ModelScope. It supports\n both dense models (e.g., all-MiniLM-L6-v2) and sparse models (e.g., SPLADE).\n\n This class is not meant to be used directly. Use concrete implementations:\n - ``SentenceTransformerEmbeddingFunction`` for dense embeddings\n - ``SentenceTransformerSparseEmbeddingFunction`` for sparse embeddings\n - ``DefaultDenseEmbedding`` for default dense embeddings\n - ``DefaultSparseEmbedding`` for default sparse embeddings\n\n Args:\n model_name (str): Model identifier or local path.\n model_source (Literal[\"huggingface\", \"modelscope\"]): Model source.\n device (Optional[str]): Device to run the model on.\n\n Note:\n - This is an internal base class for code reuse\n - Subclasses should inherit from appropriate Protocol (Dense/Sparse)\n - Provides model loading and management functionality\n \"\"\"\n\n def __init__(\n self,\n model_name: str,\n model_source: Literal[\"huggingface\", \"modelscope\"] = \"huggingface\",\n device: Optional[str] = None,\n ):\n \"\"\"Initialize the base Sentence Transformer functionality.\n\n Args:\n model_name (str): Model identifier or local path.\n model_source (Literal[\"huggingface\", \"modelscope\"]): Model source.\n device (Optional[str]): Device to run the model on.\n\n Raises:\n ValueError: If model_source is invalid.\n \"\"\"\n # Validate model_source\n if model_source not in (\"huggingface\", \"modelscope\"):\n raise ValueError(\n f\"Invalid model_source: '{model_source}'. \"\n \"Must be 'huggingface' or 'modelscope'.\"\n )\n\n self._model_name = model_name\n self._model_source = model_source\n self._device = device\n self._model = None\n\n @property\n def model_name(self) -> str:\n \"\"\"str: The Sentence Transformer model name currently in use.\"\"\"\n return self._model_name\n\n @property\n def model_source(self) -> str:\n \"\"\"str: The model source being used (\"huggingface\" or \"modelscope\").\"\"\"\n return self._model_source\n\n @property\n def device(self) -> str:\n \"\"\"str: The device the model is running on.\"\"\"\n model = self._get_model()\n if model is not None:\n return str(model.device)\n return self._device or \"cpu\"\n\n def _get_model(self):\n \"\"\"Load or retrieve the Sentence Transformer model.\n\n Returns:\n SentenceTransformer or SparseEncoder: The loaded model instance.\n\n Raises:\n ImportError: If required packages are not installed.\n ValueError: If model cannot be loaded.\n \"\"\"\n # Return cached model if exists\n if self._model is not None:\n return self._model\n\n # Load model\n try:\n sentence_transformers = require_module(\"sentence_transformers\")\n\n if self._model_source == \"modelscope\":\n # Load from ModelScope\n require_module(\"modelscope\")\n from modelscope.hub.snapshot_download import snapshot_download\n\n # Download model to cache\n model_dir = snapshot_download(self._model_name)\n\n # Load from local path\n self._model = sentence_transformers.SentenceTransformer(\n model_dir, device=self._device, trust_remote_code=True\n )\n else:\n # Load from Hugging Face (default)\n self._model = sentence_transformers.SentenceTransformer(\n self._model_name, device=self._device, trust_remote_code=True\n )\n\n return self._model\n\n except ImportError as e:\n if \"modelscope\" in str(e) and self._model_source == \"modelscope\":\n raise ImportError(\n \"ModelScope support requires the 'modelscope' package. \"\n \"Please install it with: pip install modelscope\"\n ) from e\n raise\n except Exception as e:\n raise ValueError(\n f\"Failed to load Sentence Transformer model '{self._model_name}' \"\n f\"from {self._model_source}: {e!s}\"\n ) from e\n\n def _is_sparse_model(self) -> bool:\n \"\"\"Check if the loaded model is a sparse encoder (e.g., SPLADE).\n\n Returns:\n bool: True if model supports sparse encoding.\n \"\"\"\n model = self._get_model()\n # Check if model has sparse encoding methods\n return hasattr(model, \"encode_query\") or hasattr(model, \"encode_document\")\n" }, { "path": "python/zvec/extension/sentence_transformer_rerank_function.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom typing import Literal, Optional\n\nfrom ..model.doc import Doc\nfrom ..tool import require_module\nfrom .rerank_function import RerankFunction\nfrom .sentence_transformer_function import SentenceTransformerFunctionBase\n\n\nclass DefaultLocalReRanker(SentenceTransformerFunctionBase, RerankFunction):\n \"\"\"Re-ranker using Sentence Transformer cross-encoder models for semantic re-ranking.\n\n This re-ranker leverages pre-trained cross-encoder models to perform deep semantic\n re-ranking of search results. It runs locally without API calls, supports GPU\n acceleration, and works with models from Hugging Face or ModelScope.\n\n Cross-encoder models evaluate query-document pairs jointly, providing more\n accurate relevance scores than bi-encoder (embedding-based) similarity.\n\n Args:\n query (str): Query text for semantic re-ranking. **Required**.\n topn (int, optional): Maximum number of documents to return after re-ranking.\n Defaults to 10.\n rerank_field (Optional[str], optional): Document field name to use as\n re-ranking input text. **Required** (e.g., \"content\", \"title\", \"body\").\n model_name (str, optional): Cross-encoder model identifier or local path.\n Defaults to ``\"cross-encoder/ms-marco-MiniLM-L6-v2\"`` (MS MARCO MiniLM).\n Common options:\n - ``\"cross-encoder/ms-marco-MiniLM-L6-v2\"``: Lightweight, fast (~80MB, recommended)\n - ``\"cross-encoder/ms-marco-MiniLM-L12-v2\"``: Better accuracy (~120MB)\n - ``\"BAAI/bge-reranker-base\"``: BGE Reranker Base (~280MB)\n - ``\"BAAI/bge-reranker-large\"``: BGE Reranker Large (highest quality, ~560MB)\n model_source (Literal[\"huggingface\", \"modelscope\"], optional): Model source.\n Defaults to ``\"huggingface\"``.\n - ``\"huggingface\"``: Load from Hugging Face Hub\n - ``\"modelscope\"``: Load from ModelScope (recommended for users in China)\n device (Optional[str], optional): Device to run the model on.\n Options: ``\"cpu\"``, ``\"cuda\"``, ``\"mps\"`` (for Apple Silicon), or ``None``\n for automatic detection. Defaults to ``None``.\n batch_size (int, optional): Batch size for processing query-document pairs.\n Larger values speed up processing but use more memory. Defaults to ``32``.\n\n Attributes:\n query (str): The query text used for re-ranking.\n topn (int): Maximum number of documents to return.\n rerank_field (Optional[str]): Field name used for re-ranking input.\n model_name (str): The cross-encoder model being used.\n model_source (str): The model source (\"huggingface\" or \"modelscope\").\n device (str): The device the model is running on.\n\n Raises:\n ValueError: If ``query`` is empty/None, ``rerank_field`` is None,\n or model cannot be loaded.\n TypeError: If input types are invalid.\n RuntimeError: If model inference fails.\n\n Note:\n - Requires Python 3.10, 3.11, or 3.12\n - Requires ``sentence-transformers`` package: ``pip install sentence-transformers``\n - For ModelScope support, also requires: ``pip install modelscope``\n - First run downloads the model (~80-560MB depending on model) from chosen source\n - No API keys or network required after initial download\n - Cross-encoders are slower than bi-encoders but more accurate\n - GPU acceleration provides significant speedup (5-10x)\n\n **MS MARCO MiniLM-L6-v2 Model (Default):**\n\n The default model ``cross-encoder/ms-marco-MiniLM-L6-v2`` is a lightweight and\n efficient cross-encoder trained on MS MARCO dataset. It provides:\n\n - Fast inference speed (suitable for real-time applications)\n - Small model size (~80MB, quick to download)\n - Good balance between speed and accuracy\n - Trained on 500K+ query-document pairs\n - Public availability without authentication\n\n **For users in China:**\n\n If you encounter Hugging Face access issues, use ModelScope instead:\n\n .. code-block:: python\n\n # Recommended for users in China\n reranker = SentenceTransformerReRanker(\n query=\"机器学习算法\",\n rerank_field=\"content\",\n model_source=\"modelscope\"\n )\n\n Alternatively, use Hugging Face mirror:\n\n .. code-block:: bash\n\n export HF_ENDPOINT=https://hf-mirror.com\n\n Examples:\n >>> # Basic usage with default MS MARCO MiniLM model\n >>> from zvec.extension import SentenceTransformerReRanker\n >>>\n >>> reranker = SentenceTransformerReRanker(\n ... query=\"machine learning algorithms\",\n ... topn=5,\n ... rerank_field=\"content\"\n ... )\n >>>\n >>> # Use in collection.query()\n >>> results = collection.query(\n ... data={\"vector_field\": query_vector},\n ... reranker=reranker,\n ... topk=20\n ... )\n\n >>> # Using ModelScope for users in China\n >>> reranker = SentenceTransformerReRanker(\n ... query=\"深度学习\",\n ... topn=10,\n ... rerank_field=\"content\",\n ... model_source=\"modelscope\"\n ... )\n\n >>> # Using larger model for better quality\n >>> reranker = SentenceTransformerReRanker(\n ... query=\"neural networks\",\n ... topn=5,\n ... rerank_field=\"content\",\n ... model_name=\"BAAI/bge-reranker-large\",\n ... device=\"cuda\",\n ... batch_size=64\n ... )\n\n >>> # Direct rerank call (for testing)\n >>> query_results = {\n ... \"vector1\": [\n ... Doc(id=\"1\", score=0.9, fields={\"content\": \"Machine learning is...\"}),\n ... Doc(id=\"2\", score=0.8, fields={\"content\": \"Deep learning is...\"}),\n ... ]\n ... }\n >>> reranked = reranker.rerank(query_results)\n >>> for doc in reranked:\n ... print(f\"ID: {doc.id}, Score: {doc.score:.4f}\")\n ID: 2, Score: 0.9234\n ID: 1, Score: 0.8567\n\n See Also:\n - ``RerankFunction``: Abstract base class for re-rankers\n - ``QwenReRanker``: Re-ranker using Qwen API\n - ``RrfReRanker``: Multi-vector re-ranker using RRF\n - ``WeightedReRanker``: Multi-vector re-ranker using weighted scores\n\n References:\n - MS MARCO Cross-Encoder: https://huggingface.co/cross-encoder/ms-marco-MiniLM-L6-v2\n - BGE Reranker: https://huggingface.co/BAAI/bge-reranker-base\n - Cross-Encoder vs Bi-Encoder: https://www.sbert.net/examples/applications/cross-encoder/README.html\n \"\"\"\n\n def __init__(\n self,\n query: Optional[str] = None,\n topn: int = 10,\n rerank_field: Optional[str] = None,\n model_name: str = \"cross-encoder/ms-marco-MiniLM-L6-v2\",\n model_source: Literal[\"huggingface\", \"modelscope\"] = \"huggingface\",\n device: Optional[str] = None,\n batch_size: int = 32,\n ):\n \"\"\"Initialize SentenceTransformerReRanker with query and configuration.\n\n Args:\n query (Optional[str]): Query text for semantic matching. Required.\n topn (int): Number of top results to return.\n rerank_field (Optional[str]): Document field for re-ranking input.\n model_name (str): Cross-encoder model identifier.\n model_source (Literal[\"huggingface\", \"modelscope\"]): Model source.\n device (Optional[str]): Target device (\"cpu\", \"cuda\", \"mps\", or None).\n batch_size (int): Batch size for processing query-document pairs.\n\n Raises:\n ValueError: If query is empty or model cannot be loaded.\n \"\"\"\n # Initialize base class for model loading\n SentenceTransformerFunctionBase.__init__(\n self, model_name=model_name, model_source=model_source, device=device\n )\n\n # Initialize rerank function\n RerankFunction.__init__(self, topn=topn, rerank_field=rerank_field)\n\n # Validate query\n if not query:\n raise ValueError(\"Query is required for DefaultLocalReRanker\")\n self._query = query\n self._batch_size = batch_size\n\n # Load and validate cross-encoder model\n model = self._get_model()\n if not hasattr(model, \"predict\"):\n raise ValueError(\n f\"Model '{model_name}' does not appear to be a cross-encoder model. \"\n \"Cross-encoder models should have a 'predict' method.\"\n )\n self._model = model\n\n def _get_model(self):\n \"\"\"Load or retrieve the CrossEncoder model.\n\n This overrides the base class method to load CrossEncoder instead of\n SentenceTransformer, as reranking requires cross-encoder models.\n\n Returns:\n CrossEncoder: The loaded cross-encoder model instance.\n\n Raises:\n ImportError: If required packages are not installed.\n ValueError: If model cannot be loaded.\n \"\"\"\n # Return cached model if exists\n if self._model is not None:\n return self._model\n\n # Load cross-encoder model\n try:\n sentence_transformers = require_module(\"sentence_transformers\")\n\n if self._model_source == \"modelscope\":\n # Load from ModelScope\n require_module(\"modelscope\")\n from modelscope.hub.snapshot_download import snapshot_download\n\n # Download model to cache\n model_dir = snapshot_download(self._model_name)\n\n # Load CrossEncoder from local path\n model = sentence_transformers.CrossEncoder(\n model_dir, device=self._device\n )\n else:\n # Load CrossEncoder from Hugging Face (default)\n model = sentence_transformers.CrossEncoder(\n self._model_name, device=self._device\n )\n\n return model\n\n except ImportError as e:\n if \"modelscope\" in str(e) and self._model_source == \"modelscope\":\n raise ImportError(\n \"ModelScope support requires the 'modelscope' package. \"\n \"Please install it with: pip install modelscope\"\n ) from e\n raise\n except Exception as e:\n raise ValueError(\n f\"Failed to load CrossEncoder model '{self._model_name}' \"\n f\"from {self._model_source}: {e!s}\"\n ) from e\n\n @property\n def query(self) -> str:\n \"\"\"str: Query text used for semantic re-ranking.\"\"\"\n return self._query\n\n @property\n def batch_size(self) -> int:\n \"\"\"int: Batch size for processing query-document pairs.\"\"\"\n return self._batch_size\n\n def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:\n \"\"\"Re-rank documents using Sentence Transformer cross-encoder model.\n\n Evaluates each query-document pair using the cross-encoder model to compute\n relevance scores. Documents are then sorted by these scores and the top-k\n results are returned.\n\n Args:\n query_results (dict[str, list[Doc]]): Mapping from vector field names\n to lists of retrieved documents. Documents from all fields are\n deduplicated and re-ranked together.\n\n Returns:\n list[Doc]: Re-ranked documents (up to ``topn``) with updated ``score``\n fields containing relevance scores from the cross-encoder model.\n\n Raises:\n ValueError: If no valid documents are found or model inference fails.\n\n Note:\n - Duplicate documents (same ID) across fields are processed once\n - Documents with empty/missing ``rerank_field`` content are skipped\n - Returned scores are logits from the cross-encoder model\n - Higher scores indicate higher relevance\n - Processing time is O(n) where n is the number of documents\n\n Examples:\n >>> reranker = SentenceTransformerReRanker(\n ... query=\"machine learning\",\n ... topn=3,\n ... rerank_field=\"content\"\n ... )\n >>> query_results = {\n ... \"vector1\": [\n ... Doc(id=\"1\", score=0.9, fields={\"content\": \"ML basics\"}),\n ... Doc(id=\"2\", score=0.8, fields={\"content\": \"DL tutorial\"}),\n ... ]\n ... }\n >>> reranked = reranker.rerank(query_results)\n >>> len(reranked) <= 3\n True\n \"\"\"\n if not query_results:\n return []\n\n # Collect and deduplicate documents\n id_to_doc: dict[str, Doc] = {}\n doc_ids: list[str] = []\n contents: list[str] = []\n\n for _, query_result in query_results.items():\n for doc in query_result:\n doc_id = doc.id\n if doc_id in id_to_doc:\n continue\n\n # Extract text content from specified field\n field_value = doc.field(self.rerank_field)\n rank_content = str(field_value).strip() if field_value else \"\"\n if not rank_content:\n continue\n\n id_to_doc[doc_id] = doc\n doc_ids.append(doc_id)\n contents.append(rank_content)\n\n if not contents:\n raise ValueError(\"No documents to rerank\")\n\n try:\n # Use standard cross-encoder predict method\n pairs = [[self.query, content] for content in contents]\n scores = self._model.predict(\n pairs,\n batch_size=self.batch_size,\n show_progress_bar=False,\n convert_to_numpy=True,\n )\n\n # Convert to float list if needed\n if hasattr(scores, \"tolist\"):\n scores = scores.tolist()\n else:\n scores = [float(s) for s in scores]\n\n except Exception as e:\n raise RuntimeError(f\"Failed to compute rerank scores: {e!s}\") from e\n\n # Create scored documents\n scored_docs = [\n (doc_ids[i], id_to_doc[doc_ids[i]], scores[i]) for i in range(len(doc_ids))\n ]\n\n # Sort by score (descending) and take top-k\n scored_docs.sort(key=lambda x: x[2], reverse=True)\n top_scored_docs = scored_docs[: self.topn]\n\n # Build result list with updated scores\n results: list[Doc] = []\n for _, doc, score in top_scored_docs:\n new_doc = doc._replace(score=score)\n results.append(new_doc)\n\n return results\n" }, { "path": "python/zvec/model/__init__.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom .collection import Collection\nfrom .doc import Doc\nfrom .param.vector_query import VectorQuery\nfrom .schema.collection_schema import CollectionSchema\nfrom .schema.field_schema import FieldSchema\n\n__all__ = [\"Collection\", \"CollectionSchema\", \"Doc\", \"FieldSchema\", \"VectorQuery\"]\n" }, { "path": "python/zvec/model/collection.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom typing import Optional, Union, overload\n\nfrom _zvec import _Collection\n\nfrom ..executor import QueryContext, QueryExecutorFactory\nfrom ..extension import ReRanker\nfrom ..typing import Status\nfrom .convert import convert_to_cpp_doc, convert_to_py_doc\nfrom .doc import Doc\nfrom .param import (\n AddColumnOption,\n AlterColumnOption,\n CollectionOption,\n FlatIndexParam,\n HnswIndexParam,\n IndexOption,\n InvertIndexParam,\n IVFIndexParam,\n OptimizeOption,\n)\nfrom .param.vector_query import VectorQuery\nfrom .schema import CollectionSchema, CollectionStats, FieldSchema\n\n__all__ = [\"Collection\"]\n\n_VECTOR_INDEX_TYPES = (HnswIndexParam, IVFIndexParam, FlatIndexParam)\n\n\nclass Collection:\n \"\"\"Represents an opened collection in Zvec.\n\n A `Collection` provides methods for data definition (DDL), data manipulation (DML),\n and querying (DQL). It is obtained via `create_and_open()` or `open()`.\n\n This class is not meant to be instantiated directly; use factory functions instead.\n \"\"\"\n\n def __init__(self, obj: _Collection):\n self._obj = obj\n self._schema = None\n self._querier = None\n\n @classmethod\n def _from_core(cls, core_collection: _Collection) -> Collection:\n if not core_collection:\n raise ValueError(\"Collection is None\")\n inst = cls.__new__(cls)\n inst._obj = core_collection\n schema = CollectionSchema._from_core(core_collection.Schema())\n inst._schema = schema\n inst._querier = QueryExecutorFactory.create(schema)\n return inst\n\n @property\n def path(self) -> str:\n \"\"\"str: The filesystem path of the collection.\"\"\"\n return self._obj.Path()\n\n @property\n def option(self) -> CollectionOption:\n \"\"\"CollectionOption: The options used to open the collection.\"\"\"\n return self._obj.Options()\n\n @property\n def schema(self) -> CollectionSchema:\n \"\"\"CollectionSchema: The schema defining the structure of the collection.\"\"\"\n return self._schema\n\n @property\n def stats(self) -> CollectionStats:\n \"\"\"CollectionStats: Runtime statistics about the collection (e.g., doc count, size).\"\"\"\n return self._obj.Stats()\n\n # ========== Collection DDL Methods ==========\n def destroy(self) -> None:\n \"\"\"Permanently delete the collection from disk.\n\n Warning:\n This operation is irreversible. All data will be lost.\n \"\"\"\n self._obj.Destroy()\n\n def flush(self) -> None:\n \"\"\"Force all pending writes to disk.\n\n Ensures durability of recent inserts/updates.\n \"\"\"\n self._obj.Flush()\n\n # ========== Index DDL Methods ==========\n def create_index(\n self,\n field_name: str,\n index_param: Union[\n HnswIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam\n ],\n option: IndexOption = IndexOption(),\n ) -> None:\n \"\"\"Create an index on a field.\n\n Vector index types (HNSW, IVF, FLAT) can only be applied to vector fields.\n Inverted index (`InvertIndexParam`) is for scalar fields.\n\n Args:\n field_name (str): Name of the field to index.\n index_param (Union[HnswIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam]):\n Index configuration.\n option (Optional[IndexOption], optional): Index creation options.\n Defaults to ``IndexOption()``.\n\n Raises:\n ValueError: If a vector index is applied to a non-vector field.\n \"\"\"\n if index_param in _VECTOR_INDEX_TYPES and not self.schema.vector(field_name):\n supported_types = \", \".join(cls.__name__ for cls in _VECTOR_INDEX_TYPES)\n raise ValueError(\n f\"Cannot apply vector index to non-vector field '{field_name}'. \"\n f\"The field must be of vector type to use index types like {supported_types}.\"\n )\n self._obj.CreateIndex(field_name, index_param, option)\n self._schema = CollectionSchema._from_core(self._obj.Schema())\n\n def drop_index(self, field_name: str) -> None:\n \"\"\"Remove the index from a field.\n\n Args:\n field_name (str): Name of the indexed field.\n \"\"\"\n self._obj.DropIndex(field_name)\n self._schema = CollectionSchema._from_core(self._obj.Schema())\n\n def optimize(self, option: OptimizeOption = OptimizeOption()) -> None:\n \"\"\"Optimize the collection (e.g., merge segments, rebuild index).\n\n Args:\n option (Optional[OptimizeOption], optional): Optimization options.\n Defaults to ``OptimizeOption()``.\n \"\"\"\n self._obj.Optimize(option)\n\n # ========== COLUMN DDL Methods ==========\n def add_column(\n self,\n field_schema: FieldSchema,\n expression: str = \"\",\n option: AddColumnOption = AddColumnOption(),\n ) -> None:\n \"\"\"Add a new column to the collection.\n\n The column is populated using the provided expression (e.g., SQL-like formula).\n\n Args:\n field_schema (FieldSchema): Schema definition for the new column.\n expression (str): Expression to compute values for existing documents.\n option (Optional[AddColumnOption], optional): Options for the operation.\n Defaults to ``AddColumnOption()``.\n \"\"\"\n self._obj.AddColumn(field_schema._get_object(), expression, option)\n self._schema = CollectionSchema._from_core(self._obj.Schema())\n\n def drop_column(self, field_name: str) -> None:\n \"\"\"Remove a column from the collection.\n\n Args:\n field_name (str): Name of the column to drop.\n \"\"\"\n self._obj.DropColumn(field_name)\n self._schema = CollectionSchema._from_core(self._obj.Schema())\n\n def alter_column(\n self,\n old_name: str,\n new_name: Optional[str] = None,\n field_schema: Optional[FieldSchema] = None,\n option: AlterColumnOption = AlterColumnOption(),\n ) -> None:\n \"\"\"Rename a column, update its schema.\n\n This method supports three atomic operations:\n 1. Rename only (when `field_schema` is None).\n 2. Modify schema only (when `new_name` is None or empty string).\n\n Args:\n old_name (str): The current name of the column to be altered.\n new_name (Optional[str]): The new name for the column.\n - If provided and non-empty, the column will be renamed.\n - If `None` or empty string, no rename occurs.\n field_schema (Optional[FieldSchema]): The new schema definition.\n - If provided, the column's type, dimension, or other properties will be updated.\n - If `None`, only renaming (if requested) is performed.\n option (AlterColumnOption, optional): Options controlling the alteration behavior.\n Defaults to ``AlterColumnOption()``.\n\n **Limitation**: This operation **only supports scalar numeric columns**. such as:\n - `DOUBLE`, `FLOAT`,\n - `INT32`, `INT64`, `UINT32`, `UINT64`\n\n Note:\n - Schema modification may trigger data migration or index rebuild.\n\n Examples:\n >>> # Rename column only\n >>> results = collection.alter_column(old_name=\"id\", new_name=\"doc_id\")\n\n >>> # Modify schema only\n >>> new_schema = FieldSchema(name=\"doc_id\", dtype=DataType.INT64)\n >>> collection.alter_column(\"id\", field_schema=new_schema)\n \"\"\"\n self._obj.AlterColumn(\n old_name,\n new_name or \"\",\n field_schema._get_object() if field_schema else None,\n option,\n )\n self._schema = CollectionSchema._from_core(self._obj.Schema())\n\n # ========== Collection DDL Methods ==========\n @overload\n def insert(self, docs: Doc) -> Status:\n pass\n\n @overload\n def insert(self, docs: list[Doc]) -> list[Status]:\n pass\n\n def insert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:\n \"\"\"Insert new documents into the collection.\n\n Documents must have unique IDs and conform to the schema.\n\n Args:\n docs (Union[Doc, list[Doc]]): One or more documents to insert.\n\n Returns:\n Union[Status, list[Status]]: If a single Doc was given, returns its Status;\n if a list was given, returns a list of Status objects.\n \"\"\"\n is_single = isinstance(docs, Doc)\n doc_list = [docs] if is_single else docs\n results = self._obj.Insert(\n [convert_to_cpp_doc(doc, self.schema) for doc in doc_list]\n )\n return results[0] if is_single else results\n\n @overload\n def upsert(self, docs: Doc) -> Status:\n pass\n\n @overload\n def upsert(self, docs: list[Doc]) -> list[Status]:\n pass\n\n def upsert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:\n \"\"\"Insert new documents or update existing ones by ID.\n\n Args:\n docs (Union[Doc, list[Doc]]): Documents to upsert.\n\n Returns:\n Union[Status, list[Status]]: If a single Doc was given, returns its Status;\n if a list was given, returns a list of Status objects.\n \"\"\"\n is_single = isinstance(docs, Doc)\n doc_list = [docs] if is_single else docs\n results = self._obj.Upsert(\n [convert_to_cpp_doc(doc, self.schema) for doc in doc_list]\n )\n return results[0] if is_single else results\n\n @overload\n def update(self, docs: Doc) -> Status:\n pass\n\n @overload\n def update(self, docs: list[Doc]) -> list[Status]:\n pass\n\n def update(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:\n \"\"\"Update existing documents by ID.\n\n Only specified fields are updated; others remain unchanged.\n\n Args:\n docs (Union[Doc, list[Doc]]): Documents containing updated fields.\n\n Returns:\n Union[Status, list[Status]]: If a single Doc was given, returns its Status;\n if a list was given, returns a list of Status objects.\n \"\"\"\n is_single = isinstance(docs, Doc)\n doc_list = [docs] if is_single else docs\n results = self._obj.Update(\n [convert_to_cpp_doc(doc, self.schema) for doc in doc_list]\n )\n return results[0] if is_single else results\n\n @overload\n def delete(self, ids: str) -> Status:\n pass\n\n @overload\n def delete(self, ids: list[str]) -> list[Status]:\n pass\n\n def delete(self, ids: Union[str, list[str]]) -> Union[Status, list[Status]]:\n \"\"\"Delete documents by ID.\n\n Args:\n ids (Union[str, list[str]]): One or more document IDs to delete.\n\n Returns:\n Union[Status, list[Status]]: If a single id was given, returns its Status;\n if a list was given, returns a list of Status objects.\n \"\"\"\n is_single = isinstance(ids, str)\n id_list = [ids] if isinstance(ids, str) else ids\n results = self._obj.Delete(id_list)\n return results[0] if is_single else results\n\n def delete_by_filter(self, filter: str) -> None:\n \"\"\"Delete documents matching a filter expression.\n\n Args:\n filter (str): Boolean expression (e.g., ``\"age > 30\"``).\n \"\"\"\n self._obj.DeleteByFilter(filter)\n\n # ========== Collection DQL-fetch Methods ==========\n def fetch(self, ids: Union[str, list[str]]) -> dict[str, Doc]:\n \"\"\"Retrieve documents by ID.\n\n Args:\n ids (Union[str, list[str]]): Document IDs to fetch.\n\n Returns:\n dict[str, Doc]: Mapping from ID to document. Missing IDs are omitted.\n \"\"\"\n ids = [ids] if isinstance(ids, str) else ids\n docs = self._obj.Fetch(ids)\n return {\n doc_id: py_doc\n for doc_id, core_doc in docs.items()\n if (py_doc := convert_to_py_doc(core_doc, self.schema)) is not None\n }\n\n # ========== Collection DQL-Query Methods ==========\n\n def query(\n self,\n vectors: Optional[Union[VectorQuery, list[VectorQuery]]] = None,\n *,\n topk: int = 10,\n filter: Optional[str] = None,\n include_vector: bool = False,\n output_fields: Optional[list[str]] = None,\n reranker: Optional[ReRanker] = None,\n ) -> list[Doc]:\n \"\"\"Perform vector similarity search with optional filtering and re-ranking.\n\n At least one `VectorQuery` must be provided.\n\n Args:\n vectors (Optional[Union[VectorQuery, list[VectorQuery]]], optional):\n One or more vector queries. Defaults to None.\n topk (int, optional): Number of nearest neighbors to return.\n Defaults to 10.\n filter (Optional[str], optional): Boolean expression to pre-filter candidates.\n Defaults to None.\n include_vector (bool, optional): Whether to include vector data in results.\n Defaults to False.\n output_fields (Optional[list[str]], optional): Scalar fields to include.\n If None, all fields are returned. Defaults to None.\n reranker (Optional[ReRanker], optional): Re-ranker to refine results.\n Defaults to None.\n\n Returns:\n list[Doc]: Top-k matching documents, sorted by relevance score.\n\n Examples:\n >>> from zvec import VectorQuery\n >>> results = collection.query(\n ... vectors=VectorQuery(\"embedding\", vector=[0.1, 0.2]),\n ... topk=5,\n ... filter=\"category == 'tech'\",\n ... output_fields=[\"title\", \"url\"]\n ... )\n \"\"\"\n ctx = QueryContext(\n topk=topk,\n filter=filter,\n queries=[vectors] if isinstance(vectors, VectorQuery) else vectors,\n include_vector=include_vector,\n output_fields=output_fields,\n reranker=reranker,\n )\n return self._querier.execute(ctx, self._obj)\n" }, { "path": "python/zvec/model/convert.py", "content": "# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom _zvec import _Doc\n\nfrom .doc import Doc\nfrom .schema import CollectionSchema\n\n\ndef convert_to_cpp_doc(doc: Doc, collection_schema: CollectionSchema) -> _Doc:\n if not doc or not collection_schema:\n return None\n\n _doc = _Doc()\n\n # set pk\n _doc.set_pk(doc.id)\n\n # set scalar fields\n for k, v in doc.fields.items():\n field_schema = collection_schema.field(k)\n if not field_schema:\n raise ValueError(\n f\"schema validate failed: {k} not found in collection schema\"\n )\n _doc.set_any(k, field_schema._get_object(), v)\n\n # set vector fields\n for k, v in doc.vectors.items():\n vector_schema = collection_schema.vector(k)\n if not vector_schema:\n raise ValueError(\n f\"schema validate failed: {k} not found in collection schema\"\n )\n _doc.set_any(k, vector_schema._get_object(), v)\n return _doc\n\n\ndef convert_to_py_doc(doc: _Doc, collection_schema: CollectionSchema) -> Doc:\n if not doc or not collection_schema:\n return None\n\n data_tuple = doc.get_all(collection_schema._get_object())\n return Doc._from_tuple(data_tuple)\n" }, { "path": "python/zvec/model/doc.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport json\nfrom typing import Any, Optional\n\nfrom ..common import VectorType\n\n__all__ = [\n \"Doc\",\n]\n\n\nclass Doc:\n \"\"\"Represents a retrieved document with optional metadata, fields, and vectors.\n\n This immutable data class encapsulates the result of a search or retrieval\n operation. It includes the document ID, relevance score (if applicable),\n scalar fields, and vector embeddings.\n\n During initialization, any `numpy.ndarray` in `vectors` is automatically\n converted to a plain Python list for JSON serialization and immutability.\n\n Attributes:\n id (str): Unique identifier of the document.\n score (Optional[float], optional): Relevance score from search.\n Defaults to None.\n vectors (Optional[dict[str, VectorType]], optional): Named vector\n embeddings associated with the document. Values are converted to\n lists if originally `np.ndarray`. Defaults to None.\n fields (Optional[dict[str, Any]], optional): Scalar metadata fields\n (e.g., title, timestamp). Defaults to None.\n\n Examples:\n >>> import numpy as np\n >>> import zvec\n >>> doc = zvec.Doc(\n ... id=\"doc1\",\n ... score=0.95,\n ... vectors={\"emb\": np.array([0.1, 0.2, 0.3])},\n ... fields={\"title\": \"Hello World\"}\n ... )\n >>> print(doc.vector(\"emb\"))\n [0.1, 0.2, 0.3]\n >>> print(doc.has_field(\"title\"))\n True\n \"\"\"\n\n __slots__ = (\"id\", \"score\", \"vectors\", \"fields\")\n\n def __init__(\n self,\n id: str,\n score: Optional[float] = None,\n vectors: Optional[dict[str, VectorType]] = None,\n fields: Optional[dict[str, Any]] = None,\n ):\n self.id = id\n self.score = score\n self.vectors = vectors or {}\n self.fields = fields or {}\n\n def has_field(self, name: str) -> bool:\n \"\"\"Check if the document contains a scalar field with the given name.\n\n Args:\n name (str): Name of the field to check.\n\n Returns:\n bool: True if the field exists, False otherwise.\n \"\"\"\n return name in self.fields\n\n def has_vector(self, name: str) -> bool:\n \"\"\"Check if the document contains a vector with the given name.\n\n Args:\n name (str): Name of the vector to check.\n\n Returns:\n bool: True if the vector exists, False otherwise.\n \"\"\"\n return name in self.vectors\n\n def vector(self, name: str):\n \"\"\"Get a vector by name.\n\n Args:\n name (str): Name of the vector.\n\n Returns:\n Any: The vector (as a list) if it exists, otherwise None.\n \"\"\"\n return self.vectors and self.vectors.get(name)\n\n def field(self, name: str):\n \"\"\"Get a scalar field by name.\n\n Args:\n name (str): Name of the field.\n\n Returns:\n Any: The field value if it exists, otherwise None.\n \"\"\"\n return self.fields and self.fields.get(name)\n\n def vector_names(self) -> list[str]:\n \"\"\"Get the list of all vector names in this document.\n\n Returns:\n list[str]: A list of vector field names. Empty if no vectors.\n \"\"\"\n return [] if not self.vectors else list(self.vectors.keys())\n\n def field_names(self) -> list[str]:\n \"\"\"Get the list of all scalar field names in this document.\n\n Returns:\n list[str]: A list of field names. Empty if no fields.\n \"\"\"\n return [] if not self.fields else list(self.fields.keys())\n\n def __repr__(self) -> str:\n try:\n schema = {\n \"id\": self.id,\n \"score\": self.score,\n \"fields\": self.fields,\n \"vectors\": self.vectors,\n }\n return json.dumps(schema, indent=2, ensure_ascii=False)\n except Exception as e:\n return f\"\"\n\n def _replace(self, **changes):\n new_tuple = (\n changes.get(\"id\", self.id),\n changes.get(\"score\", self.score),\n changes.get(\"fields\", self.fields.copy() if self.fields else None),\n changes.get(\"vectors\", self.vectors.copy() if self.vectors else None),\n )\n return type(self)._from_tuple(new_tuple)\n\n @classmethod\n def _from_tuple(\n cls, data_tuple: tuple[str, float, dict[str, Any], dict[str, VectorType]]\n ):\n obj = object.__new__(cls)\n obj.id = data_tuple[0]\n obj.score = data_tuple[1]\n obj.fields = data_tuple[2] or {}\n\n vectors = data_tuple[3]\n if vectors is not None:\n obj.vectors = {\n name: (vec.tolist() if hasattr(vec, \"tolist\") else vec)\n for name, vec in vectors.items()\n }\n else:\n obj.vectors = {}\n return obj\n" }, { "path": "python/zvec/model/param/__init__.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom _zvec.param import (\n AddColumnOption,\n AlterColumnOption,\n CollectionOption,\n FlatIndexParam,\n HnswIndexParam,\n HnswQueryParam,\n HnswRabitqIndexParam,\n HnswRabitqQueryParam,\n IndexOption,\n InvertIndexParam,\n IVFIndexParam,\n IVFQueryParam,\n OptimizeOption,\n)\n\n__all__ = [\n \"AddColumnOption\",\n \"AlterColumnOption\",\n \"CollectionOption\",\n \"FlatIndexParam\",\n \"HnswIndexParam\",\n \"HnswQueryParam\",\n \"HnswRabitqIndexParam\",\n \"HnswRabitqQueryParam\",\n \"IVFIndexParam\",\n \"IVFQueryParam\",\n \"IndexOption\",\n \"InvertIndexParam\",\n \"OptimizeOption\",\n]\n" }, { "path": "python/zvec/model/param/__init__.pyi", "content": "\"\"\"\nThis module contains the params of Zvec\n\"\"\"\n\nfrom __future__ import annotations\n\nimport collections\nimport typing\n\nimport _zvec.typing\n\n__all__: list[str] = [\n \"AddColumnOption\",\n \"AlterColumnOption\",\n \"CollectionOption\",\n \"FlatIndexParam\",\n \"HnswIndexParam\",\n \"HnswQueryParam\",\n \"HnswRabitqIndexParam\",\n \"HnswRabitqQueryParam\",\n \"IVFIndexParam\",\n \"IVFQueryParam\",\n \"IndexOption\",\n \"IndexParam\",\n \"InvertIndexParam\",\n \"OptimizeOption\",\n \"QueryParam\",\n \"SegmentOption\",\n \"VectorIndexParam\",\n]\n\nclass AddColumnOption:\n \"\"\"\n\n Options for adding a new column to a collection.\n\n Attributes:\n concurrency (int): Number of threads to use when backfilling data\n for the new column. If 0, auto-detect is used. Default is 0.\n\n Examples:\n >>> opt = AddColumnOption(concurrency=1)\n >>> print(opt.concurrency)\n 1\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(self, concurrency: typing.SupportsInt = 0) -> None:\n \"\"\"\n Constructs an AddColumnOption instance.\n\n Args:\n concurrency (int, optional): Number of threads for data backfill.\n 0 means auto-detect. Defaults to 0.\n \"\"\"\n\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def concurrency(self) -> int:\n \"\"\"\n int: Number of threads used when adding a column (0 = auto).\n \"\"\"\n\nclass AlterColumnOption:\n \"\"\"\n\n Options for altering an existing column (e.g., changing index settings).\n\n Attributes:\n concurrency (int): Number of threads to use during the alteration process.\n If 0, the system will choose an optimal value automatically.\n Default is 0.\n\n Examples:\n >>> opt = AlterColumnOption(concurrency=1)\n >>> print(opt.concurrency)\n 1\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(self, concurrency: typing.SupportsInt = 0) -> None:\n \"\"\"\n Constructs an AlterColumnOption instance.\n\n Args:\n concurrency (int, optional): Number of threads for column alteration.\n 0 means auto-detect. Defaults to 0.\n \"\"\"\n\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def concurrency(self) -> int:\n \"\"\"\n int: Number of threads used when altering a column (0 = auto).\n \"\"\"\n\nclass CollectionOption:\n \"\"\"\n\n Options for opening or creating a collection.\n\n Attributes:\n read_only (bool): Whether the collection is opened in read-only mode.\n Default is False.\n enable_mmap (bool): Whether to use memory-mapped I/O for data files.\n Default is True.\n\n Examples:\n >>> opt = CollectionOption(read_only=True, enable_mmap=False)\n >>> print(opt.read_only)\n True\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(self, read_only: bool = False, enable_mmap: bool = True) -> None:\n \"\"\"\n Constructs a CollectionOption instance.\n\n Args:\n read_only (bool, optional): Open collection in read-only mode.\n Defaults to False.\n enable_mmap (bool, optional): Enable memory-mapped I/O.\n Defaults to True.\n \"\"\"\n\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def enable_mmap(self) -> bool: ...\n @property\n def read_only(self) -> bool: ...\n\nclass FlatIndexParam(VectorIndexParam):\n \"\"\"\n\n Parameters for configuring a flat (brute-force) index.\n\n A flat index performs exact nearest neighbor search by comparing the query\n vector against all vectors in the collection. It is simple, accurate, and\n suitable for small to medium datasets or as a baseline.\n\n Attributes:\n metric_type (MetricType): Distance metric used for similarity computation.\n Default is ``MetricType.IP`` (inner product).\n quantize_type (QuantizeType): Optional quantization type for vector\n compression (e.g., FP16, INT8). Use ``QuantizeType.UNDEFINED`` to\n disable quantization. Default is ``QuantizeType.UNDEFINED``.\n\n Examples:\n >>> from zvec.typing import MetricType, QuantizeType\n >>> params = FlatIndexParam(\n ... metric_type=MetricType.L2,\n ... quantize_type=QuantizeType.FP16\n ... )\n >>> print(params)\n {'metric_type': 'L2', 'quantize_type': 'FP16'}\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(\n self,\n metric_type: _zvec.typing.MetricType = ...,\n quantize_type: _zvec.typing.QuantizeType = ...,\n ) -> None:\n \"\"\"\n Constructs a FlatIndexParam instance.\n\n Args:\n metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.\n quantize_type (QuantizeType, optional): Vector quantization type.\n Defaults to QuantizeType.UNDEFINED (no quantization).\n \"\"\"\n\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n def to_dict(self) -> dict:\n \"\"\"\n Convert to dictionary with all fields\n \"\"\"\n\nclass HnswIndexParam(VectorIndexParam):\n \"\"\"\n\n Parameters for configuring an HNSW (Hierarchical Navigable Small World) index.\n\n HNSW is a graph-based approximate nearest neighbor search index. This class\n encapsulates its construction hyperparameters.\n\n Attributes:\n metric_type (MetricType): Distance metric used for similarity computation.\n Default is ``MetricType.IP`` (inner product).\n m (int): Number of bi-directional links created for every new element\n during construction. Higher values improve accuracy but increase\n memory usage and construction time. Default is 50.\n ef_construction (int): Size of the dynamic candidate list for nearest\n neighbors during index construction. Larger values yield better\n graph quality at the cost of slower build time. Default is 500.\n quantize_type (QuantizeType): Optional quantization type for vector\n compression (e.g., FP16, INT8). Default is `QuantizeType.UNDEFINED` to\n disable quantization.\n\n Examples:\n >>> from zvec.typing import MetricType, QuantizeType\n >>> params = HnswIndexParam(\n ... metric_type=MetricType.COSINE,\n ... m=16,\n ... ef_construction=200,\n ... quantize_type=QuantizeType.INT8\n ... )\n >>> print(params)\n {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8'}\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(\n self,\n metric_type: _zvec.typing.MetricType = ...,\n m: typing.SupportsInt = 50,\n ef_construction: typing.SupportsInt = 500,\n quantize_type: _zvec.typing.QuantizeType = ...,\n ) -> None: ...\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n def to_dict(self) -> dict:\n \"\"\"\n Convert to dictionary with all fields\n \"\"\"\n\n @property\n def ef_construction(self) -> int:\n \"\"\"\n int: Candidate list size during index construction.\n \"\"\"\n\n @property\n def m(self) -> int:\n \"\"\"\n int: Maximum number of neighbors per node in upper layers.\n \"\"\"\n\nclass HnswQueryParam(QueryParam):\n \"\"\"\n\n Query parameters for HNSW (Hierarchical Navigable Small World) index.\n\n Controls the trade-off between search speed and accuracy via the `ef` parameter.\n\n Attributes:\n type (IndexType): Always ``IndexType.HNSW``.\n ef (int): Size of the dynamic candidate list during search.\n Larger values improve recall but slow down search.\n Default is 300.\n radius (float): Search radius for range queries. Default is 0.0.\n is_linear (bool): Force linear search. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n\n Examples:\n >>> params = HnswQueryParam(ef=300)\n >>> print(params.ef)\n 300\n >>> print(params.to_dict() if hasattr(params, 'to_dict') else params)\n {\"type\":\"HNSW\", \"ef\":300}\n \"\"\"\n def __getstate__(self) -> tuple: ...\n def __init__(\n self,\n ef: typing.SupportsInt = 300,\n radius: typing.SupportsFloat = 0.0,\n is_linear: bool = False,\n is_using_refiner: bool = False,\n ) -> None:\n \"\"\"\n Constructs an HnswQueryParam instance.\n\n Args:\n ef (int, optional): Search-time candidate list size.\n Higher values improve accuracy. Defaults to 300.\n radius (float, optional): Search radius for range queries. Default is 0.0.\n is_linear (bool, optional): Force linear search. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n \"\"\"\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def ef(self) -> int:\n \"\"\"\n int: Size of the dynamic candidate list during HNSW search.\n \"\"\"\n\nclass HnswRabitqIndexParam(VectorIndexParam):\n \"\"\"\n\n Parameters for configuring an HNSW (Hierarchical Navigable Small World) index with RabitQ quantization.\n\n HNSW is a graph-based approximate nearest neighbor search index. RabitQ is a\n quantization method that provides high compression with minimal accuracy loss.\n\n Attributes:\n metric_type (MetricType): Distance metric used for similarity computation.\n Default is ``MetricType.IP`` (inner product).\n total_bits (int): Total bits for RabitQ quantization. Default is 7.\n num_clusters (int): Number of clusters for RabitQ. Default is 16.\n m (int): Number of bi-directional links created for every new element\n during construction. Higher values improve accuracy but increase\n memory usage and construction time. Default is 50.\n ef_construction (int): Size of the dynamic candidate list for nearest\n neighbors during index construction. Larger values yield better\n graph quality at the cost of slower build time. Default is 500.\n sample_count (int): Sample count for RabitQ training. Default is 0.\n\n Examples:\n >>> from zvec.typing import MetricType\n >>> params = HnswRabitqIndexParam(\n ... metric_type=MetricType.COSINE,\n ... total_bits=8,\n ... num_clusters=256,\n ... m=16,\n ... ef_construction=200,\n ... sample_count=10000\n ... )\n >>> print(params)\n {'metric_type': 'COSINE', 'total_bits': 8, 'num_clusters': 256, 'm': 16, 'ef_construction': 200, 'sample_count': 10000}\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(\n self,\n metric_type: _zvec.typing.MetricType = ...,\n total_bits: typing.SupportsInt = 7,\n num_clusters: typing.SupportsInt = 16,\n m: typing.SupportsInt = 50,\n ef_construction: typing.SupportsInt = 500,\n sample_count: typing.SupportsInt = 0,\n ) -> None: ...\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n def to_dict(self) -> dict:\n \"\"\"\n Convert to dictionary with all fields\n \"\"\"\n\n @property\n def ef_construction(self) -> int:\n \"\"\"\n int: Candidate list size during index construction.\n \"\"\"\n\n @property\n def m(self) -> int:\n \"\"\"\n int: Maximum number of neighbors per node.\n \"\"\"\n\n @property\n def total_bits(self) -> int:\n \"\"\"\n int: Total bits for RabitQ quantization.\n \"\"\"\n\n @property\n def num_clusters(self) -> int:\n \"\"\"\n int: Number of clusters for RabitQ.\n \"\"\"\n\n @property\n def sample_count(self) -> int:\n \"\"\"\n int: Sample count for RabitQ training.\n \"\"\"\n\nclass HnswRabitqQueryParam(QueryParam):\n \"\"\"\n\n Query parameters for HNSW index with RabitQ quantization.\n\n Controls the trade-off between search speed and accuracy via the `ef` parameter.\n\n Attributes:\n type (IndexType): Always ``IndexType.HNSW_RABITQ``.\n ef (int): Size of the dynamic candidate list during search.\n Larger values improve recall but slow down search.\n Default is 300.\n radius (float): Search radius for range queries. Default is 0.0.\n is_linear (bool): Force linear search. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n\n Examples:\n >>> params = HnswRabitqQueryParam(ef=300)\n >>> print(params.ef)\n 300\n \"\"\"\n def __getstate__(self) -> tuple: ...\n def __init__(\n self,\n ef: typing.SupportsInt = 300,\n radius: typing.SupportsFloat = 0.0,\n is_linear: bool = False,\n is_using_refiner: bool = False,\n ) -> None:\n \"\"\"\n Constructs an HnswRabitqQueryParam instance.\n\n Args:\n ef (int, optional): Search-time candidate list size.\n Higher values improve accuracy. Defaults to 300.\n radius (float, optional): Search radius for range queries. Default is 0.0.\n is_linear (bool, optional): Force linear search. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n \"\"\"\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def ef(self) -> int:\n \"\"\"\n int: Size of the dynamic candidate list during HNSW search.\n \"\"\"\n\nclass IVFIndexParam(VectorIndexParam):\n \"\"\"\n\n Parameters for configuring an IVF (Inverted File Index) index.\n\n IVF partitions the vector space into clusters (inverted lists). At query time,\n only a subset of clusters is searched, providing a trade-off between speed\n and accuracy.\n\n Attributes:\n metric_type (MetricType): Distance metric used for similarity computation.\n Default is ``MetricType.IP`` (inner product).\n n_list (int): Number of clusters (inverted lists) to partition the dataset into.\n If set to 0, the system will auto-select a reasonable value based on data size.\n Default is 0 (auto).\n n_iters (int): Number of iterations for k-means clustering during index training.\n Higher values yield more stable centroids. Default is 10.\n use_soar (bool): Whether to enable SOAR (Scalable Optimized Adaptive Routing)\n for improved IVF search performance. Default is False.\n quantize_type (QuantizeType): Optional quantization type for vector\n compression (e.g., FP16, INT8). Default is ``QuantizeType.UNDEFINED``.\n\n Examples:\n >>> from zvec.typing import MetricType, QuantizeType\n >>> params = IVFIndexParam(\n ... metric_type=MetricType.COSINE,\n ... n_list=100,\n ... n_iters=15,\n ... use_soar=True,\n ... quantize_type=QuantizeType.INT8\n ... )\n >>> print(params.n_list)\n 100\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(\n self,\n metric_type: _zvec.typing.MetricType = ...,\n n_list: typing.SupportsInt = 0,\n n_iters: typing.SupportsInt = 10,\n use_soar: bool = False,\n quantize_type: _zvec.typing.QuantizeType = ...,\n ) -> None:\n \"\"\"\n Constructs an IVFIndexParam instance.\n\n Args:\n metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.\n n_list (int, optional): Number of inverted lists (clusters). Set to 0 for auto.\n Defaults to 0.\n n_iters (int, optional): Number of k-means iterations during training.\n Defaults to 10.\n use_soar (bool, optional): Enable SOAR optimization. Defaults to False.\n quantize_type (QuantizeType, optional): Vector quantization type.\n Defaults to QuantizeType.UNDEFINED.\n \"\"\"\n\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n def to_dict(self) -> dict:\n \"\"\"\n Convert to dictionary with all fields\n \"\"\"\n\n @property\n def n_iters(self) -> int:\n \"\"\"\n int: Number of k-means iterations during training.\n \"\"\"\n\n @property\n def n_list(self) -> int:\n \"\"\"\n int: Number of inverted lists (0 = auto).\n \"\"\"\n\n @property\n def use_soar(self) -> bool:\n \"\"\"\n bool: Whether SOAR optimization is enabled.\n \"\"\"\n\nclass IVFQueryParam(QueryParam):\n \"\"\"\n\n Query parameters for IVF (Inverted File Index) index.\n\n Controls how many inverted lists (`nprobe`) to visit during search.\n\n Attributes:\n type (IndexType): Always ``IndexType.IVF``.\n nprobe (int): Number of closest clusters (inverted lists) to search.\n Higher values improve recall but increase latency.\n Default is 10.\n radius (float): Search radius for range queries. Default is 0.0.\n is_linear (bool): Force linear search. Default is False.\n\n Examples:\n >>> params = IVFQueryParam(nprobe=20)\n >>> print(params.nprobe)\n 20\n \"\"\"\n def __getstate__(self) -> tuple: ...\n def __init__(self, nprobe: typing.SupportsInt = 10) -> None:\n \"\"\"\n Constructs an IVFQueryParam instance.\n\n Args:\n nprobe (int, optional): Number of inverted lists to probe during search.\n Higher values improve accuracy. Defaults to 10.\n \"\"\"\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def nprobe(self) -> int:\n \"\"\"\n int: Number of inverted lists to search during IVF query.\n \"\"\"\n\nclass IndexOption:\n \"\"\"\n\n Options for creating an index.\n\n Attributes:\n concurrency (int): Number of threads to use during index creation.\n If 0, the system will choose an optimal value automatically.\n Default is 0.\n\n Examples:\n >>> opt = IndexOption(concurrency=4)\n >>> print(opt.concurrency)\n 4\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(self, concurrency: typing.SupportsInt = 0) -> None:\n \"\"\"\n Constructs an IndexOption instance.\n\n Args:\n concurrency (int, optional): Number of concurrent threads.\n 0 means auto-detect. Defaults to 0.\n \"\"\"\n\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def concurrency(self) -> int:\n \"\"\"\n int: Number of threads used for index creation (0 = auto).\n \"\"\"\n\nclass IndexParam:\n \"\"\"\n\n Base class for all index parameter configurations.\n\n This abstract base class defines the common interface for index types.\n It should not be instantiated directly; use derived classes instead.\n\n Attributes:\n type (IndexType): The type of the index (e.g., HNSW, FLAT, INVERT).\n \"\"\"\n\n __hash__: typing.ClassVar[None] = None\n\n def __eq__(self, arg0: typing.Any) -> bool: ...\n def __getstate__(self) -> tuple: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n def clone(self) -> IndexParam: ...\n def to_dict(self) -> dict:\n \"\"\"\n Convert to dictionary with all fields\n \"\"\"\n\n @property\n def type(self) -> _zvec.typing.IndexType:\n \"\"\"\n IndexType: The type of the index.\n \"\"\"\n\nclass InvertIndexParam(IndexParam):\n \"\"\"\n\n Parameters for configuring an invert index.\n\n This class controls whether range query\n optimization is enabled for invert index structures.\n\n Attributes:\n type (IndexType): Always `IndexType.INVERTED`.\n enable_range_optimization (bool): Whether range optimization is enabled.\n enable_extended_wildcard (bool): Whether extended wildcard (suffix and infix) search is enabled.\n\n Examples:\n >>> params = InvertIndexParam(enable_range_optimization=True, enable_extended_wildcard=False)\n >>> print(params.enable_range_optimization)\n True\n >>> print(params.enable_extended_wildcard)\n False\n >>> config = params.to_dict()\n >>> print(config)\n {'enable_range_optimization': True, 'enable_extended_wildcard': False}\n \"\"\"\n def __getstate__(self) -> tuple: ...\n def __init__(\n self,\n enable_range_optimization: bool = False,\n enable_extended_wildcard: bool = False,\n ) -> None:\n \"\"\"\n Constructs an InvertIndexParam instance.\n\n Args:\n enable_range_optimization (bool, optional): If True, enables range query\n optimization for the invert index. Defaults to False.\n enable_extended_wildcard (bool, optional): If True, enables extended wildcard\n search including suffix and infix patterns. Defaults to False.\n \"\"\"\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n def to_dict(self) -> dict:\n \"\"\"\n Convert to dictionary with all fields\n \"\"\"\n @property\n def enable_extended_wildcard(self) -> bool:\n \"\"\"\n bool: Whether extended wildcard (suffix and infix) search is enabled.\n Note: Prefix search is always enabled regardless of this setting.\n \"\"\"\n @property\n def enable_range_optimization(self) -> bool:\n \"\"\"\n bool: Whether range optimization is enabled for this inverted index.\n \"\"\"\n\nclass OptimizeOption:\n \"\"\"\n\n Options for optimizing a collection (e.g., merging segments).\n\n Attributes:\n concurrency (int): Number of threads to use during optimization.\n If 0, the system will choose an optimal value automatically.\n Default is 0.\n\n Examples:\n >>> opt = OptimizeOption(concurrency=2)\n >>> print(opt.concurrency)\n 2\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(self, concurrency: typing.SupportsInt = 0) -> None:\n \"\"\"\n Constructs an OptimizeOption instance.\n\n Args:\n concurrency (int, optional): Number of concurrent threads.\n 0 means auto-detect. Defaults to 0.\n \"\"\"\n\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def concurrency(self) -> int:\n \"\"\"\n int: Number of threads used for optimization (0 = auto).\n \"\"\"\n\nclass QueryParam:\n \"\"\"\n\n Base class for all query parameter configurations.\n\n This abstract base class defines common query settings such as search radius\n and whether to force linear (brute-force) search. It should not be instantiated\n directly; use derived classes like `HnswQueryParam` or `IVFQueryParam`.\n\n Attributes:\n type (IndexType): The index type this query is configured for.\n radius (float): Search radius for range queries. Used in combination with\n top-k to filter results. Default is 0.0 (disabled).\n is_linear (bool): If True, forces brute-force linear search instead of\n using the index. Useful for debugging or small datasets. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n \"\"\"\n def __getstate__(self) -> tuple: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def is_linear(self) -> bool:\n \"\"\"\n bool: Whether to bypass the index and use brute-force linear search.\n \"\"\"\n @property\n def is_using_refiner(self) -> bool:\n \"\"\"\n bool: Whether to use refiner for the query.\n \"\"\"\n @property\n def radius(self) -> float:\n \"\"\"\n IndexType: The type of index this query targets.\n \"\"\"\n @property\n def type(self) -> _zvec.typing.IndexType:\n \"\"\"\n IndexType: The type of index this query targets.\n \"\"\"\n\nclass SegmentOption:\n \"\"\"\n\n Options for segment-level operations.\n\n Currently, this class mirrors CollectionOption and is used internally.\n It supports read-only mode, memory mapping, and buffer configuration.\n\n Note:\n This class is primarily for internal use. Most users should use\n CollectionOption instead.\n\n Examples:\n >>> opt = SegmentOption()\n >>> print(opt.enable_mmap)\n True\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __init__(self) -> None:\n \"\"\"\n Constructs a SegmentOption with default settings.\n \"\"\"\n\n def __repr__(self) -> str: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n @property\n def enable_mmap(self) -> bool:\n \"\"\"\n bool: Whether memory-mapped I/O is enabled.\n \"\"\"\n\n @property\n def max_buffer_size(self) -> int:\n \"\"\"\n int: Maximum buffer size in bytes (internal use).\n \"\"\"\n\n @property\n def read_only(self) -> bool:\n \"\"\"\n bool: Whether the segment is read-only.\n \"\"\"\n\nclass VectorIndexParam(IndexParam):\n \"\"\"\n\n Base class for vector index parameter configurations.\n\n Encapsulates common settings for all vector index types.\n\n Attributes:\n type (IndexType): The specific vector index type (e.g., HNSW, FLAT).\n metric_type (MetricType): Distance metric used for similarity search.\n quantize_type (QuantizeType): Optional vector quantization type.\n \"\"\"\n\n def __getstate__(self) -> tuple: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n def to_dict(self) -> dict:\n \"\"\"\n Convert to dictionary with all fields\n \"\"\"\n\n @property\n def metric_type(self) -> _zvec.typing.MetricType:\n \"\"\"\n MetricType: Distance metric (e.g., IP, COSINE, L2).\n \"\"\"\n\n @property\n def quantize_type(self) -> _zvec.typing.QuantizeType:\n \"\"\"\n QuantizeType: Vector quantization type (e.g., FP16, INT8).\n \"\"\"\n\nclass _VectorQuery:\n field_name: str\n filter: str\n include_vector: bool\n query_params: QueryParam\n\n def __getstate__(self) -> tuple: ...\n def __init__(self) -> None: ...\n def __setstate__(self, arg0: tuple) -> None: ...\n def set_vector(self, arg0: ..., arg1: typing.Any) -> None: ...\n @property\n def output_fields(self) -> list[str] | None: ...\n @output_fields.setter\n def output_fields(self, arg0: collections.abc.Sequence[str] | None) -> None: ...\n @property\n def topk(self) -> int: ...\n @topk.setter\n def topk(self, arg0: typing.SupportsInt) -> None: ...\n" }, { "path": "python/zvec/model/param/vector_query.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass\nfrom typing import Optional, Union\n\nfrom ...common import VectorType\nfrom . import HnswQueryParam, IVFQueryParam\n\n__all__ = [\"VectorQuery\"]\n\n\n@dataclass(frozen=True)\nclass VectorQuery:\n \"\"\"Represents a vector search query for a specific field in a collection.\n\n A `VectorQuery` can be constructed using either a document ID (to look up\n its vector) or an explicit vector. It may optionally include index-specific\n query parameters to control search behavior (e.g., `ef` for HNSW, `nprobe` for IVF).\n\n Exactly one of `id` or `vector` should be provided. If both are given,\n behavior is implementation-defined (typically `id` takes precedence).\n\n Attributes:\n field_name (str): Name of the vector field to query.\n id (Optional[str], optional): Document ID to fetch vector from. Default is None.\n vector (VectorType, optional): Explicit query vector. Default is None.\n param (Optional[Union[HnswQueryParam, IVFQueryParam]], optional):\n Index-specific query parameters. Default is None.\n\n Examples:\n >>> import zvec\n >>> # Query by ID\n >>> q1 = zvec.VectorQuery(field_name=\"embedding\", id=\"doc123\")\n >>> # Query by vector\n >>> q2 = zvec.VectorQuery(\n ... field_name=\"embedding\",\n ... vector=[0.1, 0.2, 0.3],\n ... param=HnswQueryParam(ef=300)\n ... )\n \"\"\"\n\n field_name: str\n id: Optional[str] = None\n vector: VectorType = None\n param: Optional[Union[HnswQueryParam, IVFQueryParam]] = None\n\n def has_id(self) -> bool:\n \"\"\"Check if the query is based on a document ID.\n\n Returns:\n bool: True if `id` is set, False otherwise.\n \"\"\"\n return self.id is not None\n\n def has_vector(self) -> bool:\n \"\"\"Check if the query contains an explicit vector.\n\n Returns:\n bool: True if `vector` is non-empty, False otherwise.\n \"\"\"\n return self.vector is not None and len(self.vector) > 0\n\n def _validate(self) -> None:\n if self.field_name is None:\n raise ValueError(\"Field name cannot be empty\")\n if self.id and self.vector:\n raise ValueError(\"Cannot provide both id and vector\")\n" }, { "path": "python/zvec/model/schema/__init__.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom _zvec.schema import CollectionStats\n\nfrom .collection_schema import CollectionSchema\nfrom .field_schema import FieldSchema, VectorSchema\n\n__all__ = [\"CollectionSchema\", \"CollectionStats\", \"FieldSchema\", \"VectorSchema\"]\n" }, { "path": "python/zvec/model/schema/__init__.pyi", "content": "\"\"\"\nThis module contains the schema of Zvec\n\"\"\"\n\nfrom __future__ import annotations\n\nimport collections.abc\nimport typing\n\nimport _zvec.param\nimport _zvec.typing\n\nfrom .collection_schema import CollectionSchema\nfrom .field_schema import FieldSchema, VectorSchema\n\n__all__: list[str] = [\n \"CollectionSchema\",\n \"CollectionStats\",\n \"FieldSchema\",\n \"VectorSchema\",\n]\n\nclass CollectionStats:\n def __init__(self) -> None: ...\n def __repr__(self) -> str: ...\n @property\n def doc_count(self) -> int: ...\n @property\n def index_completeness(self) -> dict[str, float]: ...\n\nclass _CollectionSchema:\n __hash__: typing.ClassVar[None] = None\n\n def __eq__(self, arg0: _CollectionSchema) -> bool: ...\n def __init__(\n self, name: str, fields: collections.abc.Sequence[_FieldSchema]\n ) -> None:\n \"\"\"\n Construct with name and list of fields\n \"\"\"\n\n def __ne__(self, arg0: _CollectionSchema) -> bool: ...\n def fields(self) -> list[_FieldSchema]:\n \"\"\"\n Return list of all field schemas.\n \"\"\"\n\n def forward_fields(self) -> list[_FieldSchema]:\n \"\"\"\n Return list of forward-indexed fields.\n \"\"\"\n\n def get_field(self, field_name: str) -> _FieldSchema:\n \"\"\"\n Get field by name (const pointer), returns None if not found.\n \"\"\"\n\n def get_forward_field(self, field_name: str) -> _FieldSchema:\n \"\"\"\n Get forward field (used for filtering).\n \"\"\"\n\n def get_vector_field(self, field_name: str) -> _FieldSchema:\n \"\"\"\n Get vector field by name.\n \"\"\"\n\n def has_field(self, field_name: str) -> bool:\n \"\"\"\n Check if a field exists.\n \"\"\"\n\n def vector_fields(self) -> list[_FieldSchema]:\n \"\"\"\n Return list of vector fields.\n \"\"\"\n\n @property\n def name(self) -> str: ...\n\nclass _FieldSchema:\n __hash__: typing.ClassVar[None] = None\n\n def __eq__(self, arg0: _FieldSchema) -> bool: ...\n def __init__(\n self,\n name: str,\n data_type: _zvec.typing.DataType,\n nullable: bool = False,\n dimension: typing.SupportsInt = 0,\n index_param: _zvec.param.IndexParam = None,\n ) -> None: ...\n def __ne__(self, arg0: _FieldSchema) -> bool: ...\n @property\n def data_type(self) -> _zvec.typing.DataType: ...\n @property\n def dimension(self) -> int: ...\n @property\n def index_param(self) -> typing.Any: ...\n @property\n def index_type(self) -> _zvec.typing.IndexType: ...\n @property\n def is_dense_vector(self) -> bool: ...\n @property\n def is_sparse_vector(self) -> bool: ...\n @property\n def name(self) -> str: ...\n @property\n def nullable(self) -> bool: ...\n" }, { "path": "python/zvec/model/schema/collection_schema.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport json\nfrom typing import Optional, Union\n\nfrom _zvec.schema import _CollectionSchema, _FieldSchema\n\nfrom .field_schema import FieldSchema, VectorSchema\n\n__all__ = [\n \"CollectionSchema\",\n]\n\n\nclass CollectionSchema:\n \"\"\"Defines the structure of a collection in Zvec.\n\n A collection schema specifies the name of the collection and its fields,\n including both scalar fields (e.g., int, string) and vector fields.\n Field names must be unique across both scalar and vector fields.\n\n Args:\n name (str): Name of the collection.\n fields (Optional[Union[FieldSchema, list[FieldSchema]]], optional):\n One or more scalar field definitions. Defaults to None.\n vectors (Optional[Union[VectorSchema, list[VectorSchema]]], optional):\n One or more vector field definitions. Defaults to None.\n\n Raises:\n TypeError: If `fields` or `vectors` are of unsupported types.\n ValueError: If any field or vector name is duplicated.\n\n Examples:\n >>> from zvec import FieldSchema, VectorSchema, DataType, IndexType\n >>> id_field = FieldSchema(\"id\", DataType.INT64, is_primary=True)\n >>> emb_field = VectorSchema(\"embedding\", dim=128, data_type=DataType.VECTOR_FP32)\n >>> schema = CollectionSchema(\n ... name=\"my_collection\",\n ... fields=id_field,\n ... vectors=emb_field\n ... )\n >>> print(schema.name)\n my_collection\n \"\"\"\n\n def __init__(\n self,\n name: str,\n fields: Optional[Union[FieldSchema, list[FieldSchema]]] = None,\n vectors: Optional[Union[VectorSchema, list[VectorSchema]]] = None,\n ):\n if name is None or not isinstance(name, str):\n raise ValueError(\n f\"schema validate failed: collection name must be str, got {type(name).__name__}\"\n )\n\n # handle fields\n _fields_name: list[str] = []\n _fields_list: list[_FieldSchema] = []\n\n self._check_fields(fields, _fields_name, _fields_list)\n self._check_vectors(vectors, _fields_name, _fields_list)\n\n # init\n self._cpp_obj = _CollectionSchema(\n name=name,\n fields=_fields_list,\n )\n\n def _check_fields(\n self,\n fields: Optional[Union[FieldSchema, list[FieldSchema]]],\n _fields_name: list[str],\n _fields_list: list[_FieldSchema],\n ) -> None:\n field_items = []\n\n if isinstance(fields, FieldSchema):\n field_items = [fields]\n elif isinstance(fields, list):\n field_items = fields\n elif fields is None:\n field_items = []\n else:\n raise TypeError(\n f\"schema validate failed: invalid 'fields' type, expected FieldSchema or list[FieldSchema], \"\n f\"got {type(fields).__name__}\"\n )\n\n for idx, field in enumerate(field_items):\n if not isinstance(field, FieldSchema):\n raise TypeError(\n f\"schema validate failed: invalid field type in 'fields' list, expected FieldSchema, \"\n f\"got {type(field).__name__} at index {idx}\"\n )\n\n if field.name in _fields_name:\n raise ValueError(\n f\"schema validate failed: duplicate field name '{field.name}': field names must be unique\"\n )\n _fields_name.append(field.name)\n _fields_list.append(field._get_object())\n\n def _check_vectors(\n self,\n vectors: Optional[Union[VectorSchema, list[VectorSchema]]],\n _fields_name: list[str],\n _fields_list: list[_FieldSchema],\n ) -> None:\n # handle vector\n if isinstance(vectors, VectorSchema):\n vectors_items = [vectors]\n elif isinstance(vectors, list):\n vectors_items = vectors\n elif vectors is None:\n vectors_items = []\n else:\n raise TypeError(\n f\"schema validate failed: invalid 'vectors' type, expected VectorSchema or list[VectorSchema], \"\n f\"got {type(vectors).__name__}\"\n )\n\n for idx, vector in enumerate(vectors_items):\n if not isinstance(vector, VectorSchema):\n raise TypeError(\n f\"schema validate failed: invalid vector type in 'vectors' list, expected VectorSchema, \"\n f\"got {type(vector).__name__} at index {idx}\"\n )\n\n if vector.name in _fields_name:\n raise ValueError(\n f\"schema validate failed: duplicate vector name '{vector.name}', vector names must be unique \"\n f\"(conflicts with existing field or vector)\"\n )\n _fields_name.append(vector.name)\n _fields_list.append(vector._get_object())\n\n @classmethod\n def _from_core(cls, core_collection_schema: _CollectionSchema):\n inst = cls.__new__(cls)\n if not core_collection_schema:\n raise ValueError(\"schema validate failed: schema is null\")\n inst._cpp_obj = core_collection_schema\n return inst\n\n @property\n def name(self) -> str:\n \"\"\"str: The name of the collection.\"\"\"\n return self._cpp_obj.name\n\n def field(self, name: str) -> Optional[FieldSchema]:\n \"\"\"Retrieve a scalar field by name.\n\n Args:\n name (str): Name of the field.\n\n Returns:\n Optional[FieldSchema]: The field if found, otherwise None.\n \"\"\"\n _field = self._cpp_obj.get_forward_field(name)\n return FieldSchema._from_core(_field) if _field else None\n\n def vector(self, name: str) -> Optional[VectorSchema]:\n \"\"\"Retrieve a vector field by name.\n\n Args:\n name (str): Name of the vector field.\n\n Returns:\n Optional[VectorSchema]: The vector field if found, otherwise None.\n \"\"\"\n _field = self._cpp_obj.get_vector_field(name)\n return VectorSchema._from_core(_field) if _field else None\n\n @property\n def fields(self) -> list[FieldSchema]:\n \"\"\"list[FieldSchema]: All scalar (non-vector) fields in the schema.\"\"\"\n _fields = self._cpp_obj.forward_fields()\n return [FieldSchema._from_core(_field) for _field in _fields]\n\n @property\n def vectors(self) -> list[VectorSchema]:\n \"\"\"list[VectorSchema]: All vector fields in the schema.\"\"\"\n _vectors = self._cpp_obj.vector_fields()\n return [VectorSchema._from_core(_vector) for _vector in _vectors]\n\n def _get_object(self) -> _CollectionSchema:\n return self._cpp_obj\n\n def __repr__(self) -> str:\n try:\n schema = {\n \"name\": self.name,\n \"fields\": {field.name: field.__dict__() for field in self.fields},\n \"vectors\": {vector.name: vector.__dict__() for vector in self.vectors},\n }\n return json.dumps(schema, indent=2, ensure_ascii=False)\n except Exception as e:\n return f\"\"\n\n def __str__(self) -> str:\n return self.__repr__()\n" }, { "path": "python/zvec/model/schema/field_schema.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport json\nfrom typing import Any, Optional, Union\n\nfrom _zvec.schema import _FieldSchema\n\nfrom zvec.model.param import (\n FlatIndexParam,\n HnswIndexParam,\n InvertIndexParam,\n IVFIndexParam,\n)\nfrom zvec.typing import DataType\n\n__all__ = [\n \"FieldSchema\",\n \"VectorSchema\",\n]\n\nSUPPORT_VECTOR_DATA_TYPE = [\n DataType.VECTOR_FP16,\n DataType.VECTOR_FP32,\n DataType.VECTOR_FP64,\n DataType.VECTOR_INT8,\n DataType.SPARSE_VECTOR_FP16,\n DataType.SPARSE_VECTOR_FP32,\n]\n\nSUPPORT_SCALAR_DATA_TYPE = [\n DataType.INT32,\n DataType.INT64,\n DataType.UINT32,\n DataType.UINT64,\n DataType.FLOAT,\n DataType.DOUBLE,\n DataType.STRING,\n DataType.BOOL,\n DataType.ARRAY_INT32,\n DataType.ARRAY_INT64,\n DataType.ARRAY_UINT32,\n DataType.ARRAY_UINT64,\n DataType.ARRAY_FLOAT,\n DataType.ARRAY_DOUBLE,\n DataType.ARRAY_STRING,\n DataType.ARRAY_BOOL,\n]\n\n\nclass FieldSchema:\n \"\"\"Represents a scalar (non-vector) field in a collection schema.\n\n A `FieldSchema` defines the name, data type, nullability, and optional\n inverted index configuration for a regular field (e.g., ID, timestamp, category).\n\n Args:\n name (str): Name of the field. Must be unique within the collection.\n data_type (DataType): Data type of the field (e.g., INT64, STRING).\n nullable (bool, optional): Whether the field can contain null values.\n Defaults to False.\n index_param (Optional[InvertIndexParam], optional): Inverted index\n parameters for this field. Only applicable to fields that support\n indexing (e.g., scalar fields used in filtering). Defaults to None.\n\n Examples:\n >>> from zvec.typing import DataType\n >>> from zvec.model.param import InvertIndexParam\n >>> id_field = FieldSchema(\n ... name=\"id\",\n ... data_type=DataType.INT64,\n ... nullable=False,\n ... index_param=InvertIndexParam(enable_range_optimization=True)\n ... )\n \"\"\"\n\n def __init__(\n self,\n name: str,\n data_type: DataType,\n nullable: bool = False,\n index_param: Optional[InvertIndexParam] = None,\n ):\n if name is None or not isinstance(name, str):\n raise ValueError(\n f\"schema validate failed: field name must be str, got {type(name).__name__}\"\n )\n\n if data_type not in SUPPORT_SCALAR_DATA_TYPE:\n raise ValueError(\n f\"schema validate failed: scalar_field's data_type must be one of \"\n f\"{', '.join(str(dt) for dt in SUPPORT_SCALAR_DATA_TYPE)}, \"\n f\"but field[{name}]'s data_type is {data_type}\"\n )\n\n self._cpp_obj = _FieldSchema(\n name=name,\n data_type=data_type,\n dimension=0,\n nullable=nullable,\n index_param=index_param,\n )\n\n @classmethod\n def _from_core(cls, core_field_schema: _FieldSchema):\n if core_field_schema is None:\n raise ValueError(\"schema validate failed: field schema is None\")\n inst = cls.__new__(cls)\n inst._cpp_obj = core_field_schema\n return inst\n\n def _get_object(self) -> _FieldSchema:\n return self._cpp_obj\n\n @property\n def name(self) -> str:\n \"\"\"str: The name of the field.\"\"\"\n return self._cpp_obj.name\n\n @property\n def data_type(self) -> DataType:\n \"\"\"DataType: The data type of the field (e.g., INT64, STRING).\"\"\"\n return self._cpp_obj.data_type\n\n @property\n def nullable(self) -> bool:\n \"\"\"bool: Whether the field allows null values.\"\"\"\n return self._cpp_obj.nullable\n\n @property\n def index_param(self) -> Optional[InvertIndexParam]:\n \"\"\"Optional[InvertIndexParam]: Inverted index configuration, if any.\"\"\"\n return self._cpp_obj.index_param\n\n def __dict__(self) -> dict[str, Any]:\n return {\n \"name\": self.name,\n \"data_type\": (\n self.data_type.name\n if hasattr(self.data_type, \"name\")\n else str(self.data_type)\n ),\n \"nullable\": self.nullable,\n \"index_param\": (\n self.index_param.to_dict() if self.index_param is not None else None\n ),\n }\n\n def __repr__(self) -> str:\n try:\n schema = self.__dict__()\n return json.dumps(schema, indent=2, ensure_ascii=False)\n except Exception as e:\n return f\"\"\n\n def __str__(self) -> str:\n return self.__repr__()\n\n def __eq__(self, other: object) -> bool:\n if not isinstance(other, FieldSchema):\n return False\n return self._cpp_obj == other._cpp_obj\n\n def __hash__(self) -> int:\n return hash((self.name, self.data_type, self.nullable))\n\n\nclass VectorSchema:\n \"\"\"Represents a vector field in a collection schema.\n\n A `VectorSchema` defines the name, data type, dimensionality, and index\n configuration for a vector field used in similarity search.\n\n Args:\n name (str): Name of the vector field. Must be unique within the collection.\n data_type (DataType): Vector data type (e.g., VECTOR_FP32, VECTOR_INT8).\n dimension (int, optional): Dimensionality of the vector. Must be > 0 for dense vectors;\n may be `None` for sparse vectors.\n index_param (Union[HnswIndexParam, IVFIndexParam, FlatIndexParam], optional):\n Index configuration for this vector field. Defaults to\n ``HnswIndexParam()``.\n\n Examples:\n >>> from zvec.typing import DataType\n >>> from zvec.model.param import HnswIndexParam\n >>> emb_field = VectorSchema(\n ... name=\"embedding\",\n ... data_type=DataType.VECTOR_FP32,\n ... dimension=128,\n ... index_param=HnswIndexParam(ef_construction=200, m=16)\n ... )\n \"\"\"\n\n def __init__(\n self,\n name: str,\n data_type: DataType,\n dimension: Optional[int] = 0,\n index_param: Optional[\n Union[HnswIndexParam, FlatIndexParam, IVFIndexParam]\n ] = None,\n ):\n if name is None or not isinstance(name, str):\n raise ValueError(\n f\"schema validate failed: field name must be str, got {type(name).__name__}\"\n )\n\n if not isinstance(dimension, int) or dimension < 0:\n raise ValueError(\"schema validate failed: vector's dimension must be >= 0\")\n\n if data_type not in SUPPORT_VECTOR_DATA_TYPE:\n raise ValueError(\n f\"schema validate failed: vector's data_type must be one of \"\n f\"{', '.join(str(dt) for dt in SUPPORT_VECTOR_DATA_TYPE)}, \"\n f\"but field[{name}]'s data_type is {data_type}\"\n )\n\n if index_param is None:\n index_param = FlatIndexParam()\n\n self._cpp_obj = _FieldSchema(\n name=name,\n data_type=data_type,\n dimension=dimension,\n nullable=False,\n index_param=index_param,\n )\n\n @classmethod\n def _from_core(cls, core_field_schema: _FieldSchema):\n inst = cls.__new__(cls)\n inst._cpp_obj = core_field_schema\n return inst\n\n def _get_object(self) -> _FieldSchema:\n return self._cpp_obj\n\n @property\n def name(self) -> str:\n \"\"\"str: The name of the vector field.\"\"\"\n return self._cpp_obj.name\n\n @property\n def data_type(self) -> DataType:\n \"\"\"DataType: The vector data type (e.g., VECTOR_FP32).\"\"\"\n return self._cpp_obj.data_type\n\n @property\n def dimension(self) -> int:\n \"\"\"int: The dimensionality of the vector.\"\"\"\n return self._cpp_obj.dimension\n\n @property\n def index_param(self) -> Union[HnswIndexParam, IVFIndexParam, FlatIndexParam]:\n \"\"\"Union[HnswIndexParam, IVFIndexParam, FlatIndexParam]: Index configuration for the vector.\"\"\"\n return self._cpp_obj.index_param\n\n def __dict__(self) -> dict[str, Any]:\n return {\n \"name\": self.name,\n \"data_type\": (\n self.data_type.name\n if hasattr(self.data_type, \"name\")\n else str(self.data_type)\n ),\n \"dimension\": self.dimension,\n \"index_param\": (\n self.index_param.to_dict() if self.index_param is not None else None\n ),\n }\n\n def __repr__(self) -> str:\n try:\n schema = self.__dict__()\n return json.dumps(schema, indent=2, ensure_ascii=False)\n except Exception as e:\n return f\"\"\n\n def __str__(self) -> str:\n return self.__repr__()\n\n def __eq__(self, other: object) -> bool:\n if not isinstance(other, VectorSchema):\n return False\n return self._cpp_obj == other._cpp_obj\n\n def __hash__(self) -> int:\n return hash((self.name, self.data_type, self.dimension))\n" }, { "path": "python/zvec/py.typed", "content": "" }, { "path": "python/zvec/tool/__init__.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom .util import require_module\n\n__all__ = [\"require_module\"]\n" }, { "path": "python/zvec/tool/util.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport importlib\nfrom typing import Any, Optional\n\n\ndef require_module(module: str, mitigation: Optional[str] = None) -> Any:\n \"\"\"Import a Python module and raise a user-friendly error if it is not available.\n\n This utility helps provide actionable error messages when optional dependencies\n are missing. It attempts to import the given module and, on failure, suggests\n a `pip install` command based on either the module name or an optional\n mitigation package name.\n\n Args:\n module (str): The full module name to import (e.g., ``\"numpy\"``, ``\"pandas.io.parquet\"``).\n mitigation (Optional[str], optional): The package name to suggest for installation\n if the import fails. If not provided, the top-level package of `module`\n will be used (e.g., ``\"pandas\"`` for ``\"pandas.io.parquet\"``).\n\n Returns:\n Any: The imported module object.\n\n Raises:\n ImportError: If the module cannot be imported, with a clear installation hint.\n\n Examples:\n >>> import zvec\n >>> np = zvec.require_module(\"numpy\")\n >>> pq = zvec.require_module(\"pyarrow.parquet\", mitigation=\"pyarrow\")\n\n Note:\n This function is intended for lazy-loading optional dependencies\n with helpful error messages, not for core dependencies.\n \"\"\"\n try:\n return importlib.import_module(module)\n except ImportError as e:\n package = mitigation or module\n msg = f\"Required package '{package}' is not installed. \"\n if \".\" in module:\n top_level = module.split(\".\", maxsplit=1)[0]\n msg += f\"Module '{module}' is part of '{top_level}', \"\n if mitigation:\n msg += f\"please pip install '{mitigation}'.\"\n else:\n msg += f\"please pip install '{top_level}'.\"\n else:\n msg += f\"Please pip install '{package}'.\"\n raise ImportError(msg) from e\n" }, { "path": "python/zvec/typing/__init__.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom _zvec.typing import (\n DataType,\n IndexType,\n MetricType,\n QuantizeType,\n Status,\n StatusCode,\n)\n\n__all__ = [\n \"DataType\",\n \"IndexType\",\n \"MetricType\",\n \"QuantizeType\",\n \"Status\",\n \"StatusCode\",\n]\n" }, { "path": "python/zvec/typing/__init__.pyi", "content": "\"\"\"\nThis module contains the basic data types of Zvec\n\"\"\"\n\nfrom __future__ import annotations\n\nimport typing\n\n__all__: list[str] = [\n \"DataType\",\n \"IndexType\",\n \"MetricType\",\n \"QuantizeType\",\n \"Status\",\n \"StatusCode\",\n]\n\nclass DataType:\n \"\"\"\n\n Enumeration of supported data types in Zvec.\n\n Includes scalar types, dense/sparse vector types, and array types.\n\n Examples:\n >>> import zvec\n >>> print(zvec.DataType.FLOAT)\n DataType.FLOAT\n >>> print(zvec.DataType.VECTOR_FP32)\n DataType.VECTOR_FP32\n\n\n Members:\n\n STRING\n\n BOOL\n\n INT32\n\n INT64\n\n FLOAT\n\n DOUBLE\n\n UINT32\n\n UINT64\n\n VECTOR_FP16\n\n VECTOR_FP32\n\n VECTOR_FP64\n\n VECTOR_INT8\n\n SPARSE_VECTOR_FP32\n\n SPARSE_VECTOR_FP16\n\n ARRAY_STRING\n\n ARRAY_INT32\n\n ARRAY_INT64\n\n ARRAY_FLOAT\n\n ARRAY_DOUBLE\n\n ARRAY_BOOL\n\n ARRAY_UINT32\n\n ARRAY_UINT64\n \"\"\"\n\n ARRAY_BOOL: typing.ClassVar[DataType] # value = \n ARRAY_DOUBLE: typing.ClassVar[DataType] # value = \n ARRAY_FLOAT: typing.ClassVar[DataType] # value = \n ARRAY_INT32: typing.ClassVar[DataType] # value = \n ARRAY_INT64: typing.ClassVar[DataType] # value = \n ARRAY_STRING: typing.ClassVar[DataType] # value = \n ARRAY_UINT32: typing.ClassVar[DataType] # value = \n ARRAY_UINT64: typing.ClassVar[DataType] # value = \n BOOL: typing.ClassVar[DataType] # value = \n DOUBLE: typing.ClassVar[DataType] # value = \n FLOAT: typing.ClassVar[DataType] # value = \n INT32: typing.ClassVar[DataType] # value = \n INT64: typing.ClassVar[DataType] # value = \n SPARSE_VECTOR_FP16: typing.ClassVar[\n DataType\n ] # value = \n SPARSE_VECTOR_FP32: typing.ClassVar[\n DataType\n ] # value = \n STRING: typing.ClassVar[DataType] # value = \n UINT32: typing.ClassVar[DataType] # value = \n UINT64: typing.ClassVar[DataType] # value = \n VECTOR_FP16: typing.ClassVar[DataType] # value = \n VECTOR_FP32: typing.ClassVar[DataType] # value = \n VECTOR_FP64: typing.ClassVar[DataType] # value = \n VECTOR_INT8: typing.ClassVar[DataType] # value = \n __members__: typing.ClassVar[\n dict[str, DataType]\n ] # 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': }\n\n def __eq__(self, other: typing.Any) -> bool: ...\n def __getstate__(self) -> int: ...\n def __hash__(self) -> int: ...\n def __index__(self) -> int: ...\n def __init__(self, value: typing.SupportsInt) -> None: ...\n def __int__(self) -> int: ...\n def __ne__(self, other: typing.Any) -> bool: ...\n def __repr__(self) -> str: ...\n def __setstate__(self, state: typing.SupportsInt) -> None: ...\n def __str__(self) -> str: ...\n @property\n def name(self) -> str: ...\n @property\n def value(self) -> int: ...\n\nclass IndexType:\n \"\"\"\n\n Enumeration of supported index types in Zvec.\n\n Examples:\n >>> import zvec\n >>> print(zvec.IndexType.HNSW)\n IndexType.HNSW\n\n\n Members:\n\n UNDEFINED\n\n HNSW\n\n IVF\n\n FLAT\n\n INVERT\n \"\"\"\n\n FLAT: typing.ClassVar[IndexType] # value = \n HNSW: typing.ClassVar[IndexType] # value = \n INVERT: typing.ClassVar[IndexType] # value = \n IVF: typing.ClassVar[IndexType] # value = \n UNDEFINED: typing.ClassVar[IndexType] # value = \n __members__: typing.ClassVar[\n dict[str, IndexType]\n ] # value = {'UNDEFINED': , 'HNSW': , 'IVF': , 'FLAT': , 'INVERT': }\n\n def __eq__(self, other: typing.Any) -> bool: ...\n def __getstate__(self) -> int: ...\n def __hash__(self) -> int: ...\n def __index__(self) -> int: ...\n def __init__(self, value: typing.SupportsInt) -> None: ...\n def __int__(self) -> int: ...\n def __ne__(self, other: typing.Any) -> bool: ...\n def __repr__(self) -> str: ...\n def __setstate__(self, state: typing.SupportsInt) -> None: ...\n def __str__(self) -> str: ...\n @property\n def name(self) -> str: ...\n @property\n def value(self) -> int: ...\n\nclass MetricType:\n \"\"\"\n\n Enumeration of supported distance/similarity metrics.\n\n - COSINE: Cosine similarity.\n - IP: Inner product (dot product).\n - L2: Euclidean distance (L2 norm).\n\n Examples:\n >>> import zvec\n >>> print(zvec.MetricType.COSINE)\n MetricType.COSINE\n\n\n Members:\n\n COSINE\n\n IP\n\n L2\n \"\"\"\n\n COSINE: typing.ClassVar[MetricType] # value = \n IP: typing.ClassVar[MetricType] # value = \n L2: typing.ClassVar[MetricType] # value = \n __members__: typing.ClassVar[\n dict[str, MetricType]\n ] # value = {'COSINE': , 'IP': , 'L2': }\n\n def __eq__(self, other: typing.Any) -> bool: ...\n def __getstate__(self) -> int: ...\n def __hash__(self) -> int: ...\n def __index__(self) -> int: ...\n def __init__(self, value: typing.SupportsInt) -> None: ...\n def __int__(self) -> int: ...\n def __ne__(self, other: typing.Any) -> bool: ...\n def __repr__(self) -> str: ...\n def __setstate__(self, state: typing.SupportsInt) -> None: ...\n def __str__(self) -> str: ...\n @property\n def name(self) -> str: ...\n @property\n def value(self) -> int: ...\n\nclass QuantizeType:\n \"\"\"\n\n Enumeration of supported quantization types for vector compression.\n\n Examples:\n >>> import zvec\n >>> print(zvec.QuantizeType.INT8)\n QuantizeType.INT8\n\n\n Members:\n\n UNDEFINED\n\n FP16\n\n INT8\n\n INT4\n \"\"\"\n\n FP16: typing.ClassVar[QuantizeType] # value = \n INT4: typing.ClassVar[QuantizeType] # value = \n INT8: typing.ClassVar[QuantizeType] # value = \n UNDEFINED: typing.ClassVar[QuantizeType] # value = \n __members__: typing.ClassVar[\n dict[str, QuantizeType]\n ] # value = {'UNDEFINED': , 'FP16': , 'INT8': , 'INT4': }\n\n def __eq__(self, other: typing.Any) -> bool: ...\n def __getstate__(self) -> int: ...\n def __hash__(self) -> int: ...\n def __index__(self) -> int: ...\n def __init__(self, value: typing.SupportsInt) -> None: ...\n def __int__(self) -> int: ...\n def __ne__(self, other: typing.Any) -> bool: ...\n def __repr__(self) -> str: ...\n def __setstate__(self, state: typing.SupportsInt) -> None: ...\n def __str__(self) -> str: ...\n @property\n def name(self) -> str: ...\n @property\n def value(self) -> int: ...\n\nclass Status:\n \"\"\"\n\n Represents the outcome of a Zvec operation.\n\n A `Status` object is either OK (success) or carries an error code and message.\n\n Examples:\n >>> from zvec.typing import Status, StatusCode\n >>> s = Status()\n >>> print(s.ok())\n True\n >>> s = Status(StatusCode.INVALID_ARGUMENT, \"Field not found\")\n >>> print(s.code() == StatusCode.INVALID_ARGUMENT)\n True\n >>> print(s.message())\n Field not found\n \"\"\"\n\n __hash__: typing.ClassVar[None] = None\n\n @staticmethod\n def AlreadyExists(message: str) -> Status: ...\n @staticmethod\n def InternalError(message: str) -> Status: ...\n @staticmethod\n def InvalidArgument(message: str) -> Status: ...\n @staticmethod\n def NotFound(message: str) -> Status: ...\n @staticmethod\n def OK() -> Status:\n \"\"\"\n Create an OK status.\n \"\"\"\n\n @staticmethod\n def PermissionDenied(message: str) -> Status: ...\n def __eq__(self, arg0: Status) -> bool: ...\n @typing.overload\n def __init__(self) -> None: ...\n @typing.overload\n def __init__(self, code: StatusCode, message: str = \"\") -> None:\n \"\"\"\n Construct a status with the given code and optional message.\n\n Args:\n code (StatusCode): The status code.\n message (str, optional): Error message. Defaults to empty string.\n \"\"\"\n\n def __ne__(self, arg0: Status) -> bool: ...\n def __repr__(self) -> str: ...\n def code(self) -> StatusCode:\n \"\"\"\n StatusCode: Returns the status code.\n \"\"\"\n\n def message(self) -> str:\n \"\"\"\n str: Returns the error message (may be empty).\n \"\"\"\n\n def ok(self) -> bool:\n \"\"\"\n bool: Returns True if the status is OK.\n \"\"\"\n\nclass StatusCode:\n \"\"\"\n\n Enumeration of possible status codes for Zvec operations.\n\n Used by the `Status` class to indicate success or failure reason.\n\n\n Members:\n\n OK\n\n NOT_FOUND\n\n ALREADY_EXISTS\n\n INVALID_ARGUMENT\n\n PERMISSION_DENIED\n\n FAILED_PRECONDITION\n\n RESOURCE_EXHAUSTED\n\n UNAVAILABLE\n\n INTERNAL_ERROR\n\n NOT_SUPPORTED\n\n UNKNOWN\n \"\"\"\n\n ALREADY_EXISTS: typing.ClassVar[\n StatusCode\n ] # value = \n FAILED_PRECONDITION: typing.ClassVar[\n StatusCode\n ] # value = \n INTERNAL_ERROR: typing.ClassVar[\n StatusCode\n ] # value = \n INVALID_ARGUMENT: typing.ClassVar[\n StatusCode\n ] # value = \n NOT_FOUND: typing.ClassVar[StatusCode] # value = \n NOT_SUPPORTED: typing.ClassVar[StatusCode] # value = \n OK: typing.ClassVar[StatusCode] # value = \n PERMISSION_DENIED: typing.ClassVar[\n StatusCode\n ] # value = \n RESOURCE_EXHAUSTED: typing.ClassVar[\n StatusCode\n ] # value = \n UNAVAILABLE: typing.ClassVar[StatusCode] # value = \n UNKNOWN: typing.ClassVar[StatusCode] # value = \n __members__: typing.ClassVar[\n dict[str, StatusCode]\n ] # value = {'OK': , 'NOT_FOUND': , 'ALREADY_EXISTS': , 'INVALID_ARGUMENT': , 'PERMISSION_DENIED': , 'FAILED_PRECONDITION': , 'RESOURCE_EXHAUSTED': , 'UNAVAILABLE': , 'INTERNAL_ERROR': , 'NOT_SUPPORTED': , 'UNKNOWN': }\n\n def __eq__(self, other: typing.Any) -> bool: ...\n def __getstate__(self) -> int: ...\n def __hash__(self) -> int: ...\n def __index__(self) -> int: ...\n def __init__(self, value: typing.SupportsInt) -> None: ...\n def __int__(self) -> int: ...\n def __ne__(self, other: typing.Any) -> bool: ...\n def __repr__(self) -> str: ...\n def __setstate__(self, state: typing.SupportsInt) -> None: ...\n def __str__(self) -> str: ...\n @property\n def name(self) -> str: ...\n @property\n def value(self) -> int: ...\n" }, { "path": "python/zvec/typing/enum.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom enum import IntEnum\n\n__all__ = [\"LogLevel\", \"LogType\"]\n\n\nclass LogLevel(IntEnum):\n \"\"\"Enumeration of logging severity levels, ordered from lowest to highest priority.\n\n Used to control verbosity and filtering of log messages. Higher numeric values\n indicate more severe conditions.\n\n Note:\n ``WARNING`` is an alias for ``WARN`` to match Python's built-in :mod:`logging`\n module convention.\n\n Attributes:\n DEBUG (int): Detailed information, typically of interest only when diagnosing problems.\n INFO (int): Confirmation that things are working as expected.\n WARN (int): An indication that something unexpected happened, or indicative of\n potential future problems. (Alias: ``WARNING``)\n WARNING (int): Same as ``WARN``.\n ERROR (int): Due to a more serious problem, the software has not been able\n to perform some function.\n FATAL (int): A serious error, indicating that the program itself may be unable\n to continue running.\n \"\"\"\n\n DEBUG = 0\n INFO = 1\n WARN = 2\n WARNING = 2\n ERROR = 3\n FATAL = 4\n\n\nclass LogType(IntEnum):\n \"\"\"Enumeration of supported log output destinations.\n\n Specifies where log messages should be written.\n\n Attributes:\n CONSOLE (int): Output logs to standard output/error (e.g., terminal or IDE console).\n FILE (int): Write logs to a persistent file on disk.\n \"\"\"\n\n CONSOLE = 0\n FILE = 1\n" }, { "path": "python/zvec/zvec.py", "content": "# Copyright 2025-present the zvec project\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom typing import Optional\n\nfrom _zvec import Initialize, _Collection\n\nfrom .model import Collection\nfrom .model.param import CollectionOption\nfrom .model.schema import CollectionSchema\n\n__all__ = [\"create_and_open\", \"init\", \"open\"]\n\nfrom .typing.enum import LogLevel, LogType\n\n\ndef init(\n *,\n log_type: Optional[LogType] = LogType.CONSOLE,\n log_level: Optional[LogLevel] = LogLevel.WARN,\n log_dir: Optional[str] = \"./logs\",\n log_basename: Optional[str] = \"zvec.log\",\n log_file_size: Optional[int] = 2048,\n log_overdue_days: Optional[int] = 7,\n query_threads: Optional[int] = None,\n optimize_threads: Optional[int] = None,\n invert_to_forward_scan_ratio: Optional[float] = None,\n brute_force_by_keys_ratio: Optional[float] = None,\n memory_limit_mb: Optional[int] = None,\n) -> None:\n \"\"\"Initialize Zvec with configuration options.\n\n This function must be called before any other operation.\n It can only be called once — subsequent calls raise a ``RuntimeError``.\n\n Parameters set to ``None`` are **omitted** from the configuration and\n fall back to Zvec's internal defaults, which may be derived from the runtime\n environment (e.g., cgroup CPU/memory limits). Explicitly provided values\n always override defaults.\n\n Args:\n log_type (Optional[LogType], optional): Logger destination.\n - ``LogType.CONSOLE`` (default if omitted or set to this)\n - ``LogType.FILE``\n - If ``None``, uses internal default (currently ``CONSOLE``).\n log_level (Optional[LogLevel], optional): Minimum log severity.\n Default: ``LogLevel.WARN``.\n Accepted values: ``DEBUG``, ``INFO``, ``WARN``, ``ERROR``, ``FATAL``.\n If ``None``, uses internal default (``WARN``).\n log_dir (Optional[str], optional):\n Directory for log files (only used when ``log_type=FILE``).\n Parent directories are **not** created automatically.\n Default: ``\"./logs\"``.\n If ``None``, internal default is used.\n log_basename (Optional[str], optional):\n Base name for rotated log files (e.g., ``zvec.log.1``, ``zvec.log.2``).\n Default: ``\"zvec.log\"``.\n log_file_size (Optional[int], optional):\n Max size per log file in **MB** before rotation.\n Default: ``2048`` MB (2 GB).\n log_overdue_days (Optional[int], optional):\n Days to retain rotated log files before deletion.\n Default: ``7`` days.\n query_threads (Optional[int], optional):\n Number of threads for query execution.\n If ``None`` (default), inferred from available CPU cores (via cgroup).\n Must be ≥ 1 if provided.\n optimize_threads (Optional[int], optional):\n Threads for background tasks (e.g., compaction, indexing).\n If ``None``, defaults to same as ``query_threads`` or CPU count.\n invert_to_forward_scan_ratio (Optional[float], optional):\n Threshold to switch from inverted index to full forward scan.\n Range: [0.0, 1.0]. Higher → more aggressive index skipping.\n Default: ``0.9`` (if omitted).\n brute_force_by_keys_ratio (Optional[float], optional):\n Threshold to use brute-force key lookup over index.\n Lower → prefer index; higher → prefer brute-force.\n Range: [0.0, 1.0]. Default: ``0.1``.\n memory_limit_mb (Optional[int], optional):\n Soft memory cap in MB. Zvec may throttle or fail operations\n approaching this limit.\n If ``None``, inferred from cgroup memory limit * 0.8 (e.g., in Docker).\n Must be > 0 if provided.\n\n Raises:\n RuntimeError: If Zvec is already initialized.\n ValueError: On invalid values (e.g., negative thread count, log level out of range).\n TypeError: If a value has incorrect type (e.g., string for ``query_threads``).\n\n Note:\n - All ``None`` arguments are **excluded** from the configuration payload,\n allowing the core library to apply environment-aware defaults.\n - This design ensures container-friendliness: in Kubernetes/Docker,\n omitting ``memory_limit_mb`` and thread counts lets Zvec auto-adapt.\n\n Examples:\n Initialize with defaults (log to console, auto-detect resources):\n >>> import zvec\n >>> zvec.init()\n\n Customize logging to file with rotation:\n >>> zvec.init(\n ... log_type=LogType.FILE,\n ... log_dir=\"/var/log/zvec\",\n ... log_file_size=1024,\n ... log_overdue_days=30\n ... )\n\n Limit resources explicitly:\n >>> zvec.init(\n ... memory_limit_mb=2048,\n ... query_threads=4,\n ... optimize_threads=2\n ... )\n\n Fine-tune query heuristics:\n >>> zvec.init(\n ... invert_to_forward_scan_ratio=0.95,\n ... brute_force_by_keys_ratio=0.05\n ... )\n \"\"\"\n # Build config dict, skipping None values\n config_dict = {}\n if log_type is not None:\n if not isinstance(log_type, LogType):\n raise TypeError(\"log_type must be LogType\")\n config_dict[\"log_type\"] = log_type.name\n if log_level is not None:\n if not isinstance(log_level, LogLevel):\n raise TypeError(\"log_level must be LogLevel\")\n config_dict[\"log_level\"] = log_level.name\n if log_dir is not None:\n config_dict[\"log_dir\"] = log_dir\n if log_basename is not None:\n config_dict[\"log_basename\"] = log_basename\n if log_file_size is not None:\n config_dict[\"log_file_size\"] = log_file_size\n if log_overdue_days is not None:\n config_dict[\"log_overdue_days\"] = log_overdue_days\n if query_threads is not None:\n config_dict[\"query_threads\"] = query_threads\n if optimize_threads is not None:\n config_dict[\"optimize_threads\"] = optimize_threads\n if invert_to_forward_scan_ratio is not None:\n config_dict[\"invert_to_forward_scan_ratio\"] = invert_to_forward_scan_ratio\n if brute_force_by_keys_ratio is not None:\n config_dict[\"brute_force_by_keys_ratio\"] = brute_force_by_keys_ratio\n if memory_limit_mb is not None:\n config_dict[\"memory_limit_mb\"] = memory_limit_mb\n\n Initialize(config_dict)\n\n\ndef create_and_open(\n path: str,\n schema: CollectionSchema,\n option: Optional[CollectionOption] = None,\n) -> Collection:\n \"\"\"Create a new collection and open it for use.\n\n If a collection already exists at the given path, it may raise an error\n depending on the underlying implementation.\n\n Args:\n path (str): Path or name of the collection to create.\n schema (CollectionSchema): Schema defining the structure of the collection.\n option (CollectionOption): Configuration options\n for opening the collection. Defaults to a default-constructed\n ``CollectionOption()`` if not provided.\n\n Returns:\n Collection: An opened collection instance ready for operations.\n\n Examples:\n >>> import zvec\n >>> schema = zvec.CollectionSchema(\n ... name=\"my_collection\",\n ... fields=[zvec.FieldSchema(\"id\", zvec.DataType.INT64, nullable=True)]\n ... )\n >>> coll = create_and_open(\"./my_collection\", schema)\n \"\"\"\n if not isinstance(path, str):\n raise TypeError(\"path must be a string\")\n if not isinstance(schema, CollectionSchema):\n raise TypeError(\"schema must be a CollectionSchema\")\n\n option = option or CollectionOption()\n if not isinstance(option, CollectionOption):\n raise TypeError(\"option must be a CollectionOption\")\n\n _collection = _Collection.CreateAndOpen(path, schema._get_object(), option)\n return Collection._from_core(_collection)\n\n\ndef open(path: str, option: CollectionOption = CollectionOption()) -> Collection:\n \"\"\"Open an existing collection from disk.\n\n The collection must have been previously created with ``create_and_open``.\n\n Args:\n path (str): Path or name of the existing collection.\n option (CollectionOption): Configuration options\n for opening the collection. Defaults to a default-constructed\n ``CollectionOption()`` if not provided.\n\n Returns:\n Collection: An opened collection instance.\n\n Examples:\n >>> import zvec\n >>> coll = zvec.open(\"./my_collection\")\n \"\"\"\n _collection = _Collection.Open(path, option)\n return Collection._from_core(_collection)\n" }, { "path": "scripts/README.md", "content": "" }, { "path": "scripts/build_android.sh", "content": "#!/bin/bash\nset -e\nCURRENT_DIR=$(pwd)\n\nABI=${1:-\"arm64-v8a\"}\nAPI_LEVEL=${2:-21}\nBUILD_TYPE=${3:-\"Release\"}\n\n# step1: use host env to compile protoc\necho \"step1: building protoc for host...\"\nHOST_BUILD_DIR=\"build_host\"\nmkdir -p $HOST_BUILD_DIR\ncd $HOST_BUILD_DIR\n\ncmake -DCMAKE_BUILD_TYPE=\"$BUILD_TYPE\" ..\nmake -j protoc\nPROTOC_EXECUTABLE=$CURRENT_DIR/$HOST_BUILD_DIR/bin/protoc\ncd $CURRENT_DIR\n\necho \"step1: Done!!!\"\n\n# step2: cross build zvec based on android ndk\necho \"step2: building zvec for android...\"\n\n# reset thirdparty directory\ngit submodule foreach --recursive 'git stash --include-untracked'\n\nexport ANDROID_SDK_ROOT=$HOME/Library/Android/sdk\nexport ANDROID_HOME=$ANDROID_SDK_ROOT\nexport ANDROID_NDK_HOME=$ANDROID_SDK_ROOT/ndk/28.2.13676358\nexport CMAKE_TOOLCHAIN_FILE=$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake\n\nexport PATH=$PATH:$ANDROID_SDK_ROOT/cmdline-tools/latest/bin\nexport PATH=$PATH:$ANDROID_SDK_ROOT/platform-tools\nexport PATH=$PATH:$ANDROID_NDK_HOME\n\nif [ -z \"$ANDROID_NDK_HOME\" ]; then\n echo \"error: ANDROID_NDK_HOME env not set\"\n echo \"please install NDK and set env variable ANDROID_NDK_HOME\"\n exit 1\nfi\n\nBUILD_DIR=\"build_android_${ABI}\"\nmkdir -p $BUILD_DIR\ncd $BUILD_DIR\n\necho \"configure CMake...\"\ncmake \\\n -DANDROID_NDK=\"$ANDROID_NDK_HOME\" \\\n -DCMAKE_TOOLCHAIN_FILE=\"$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake\" \\\n -DANDROID_ABI=\"$ABI\" \\\n -DANDROID_NATIVE_API_LEVEL=\"$API_LEVEL\" \\\n -DANDROID_STL=\"c++_static\" \\\n -DCMAKE_BUILD_TYPE=\"$BUILD_TYPE\" \\\n -DBUILD_PYTHON_BINDINGS=OFF \\\n -DBUILD_TOOLS=OFF \\\n -DCMAKE_INSTALL_PREFIX=\"./install\" \\\n -DGLOBAL_CC_PROTOBUF_PROTOC=$PROTOC_EXECUTABLE \\\n ../\n\necho \"building...\"\nCORE_COUNT=$(sysctl -n hw.ncpu)\nmake -j$CORE_COUNT\n\necho \"step2: Done!!!\"" }, { "path": "scripts/gcov.sh", "content": "#!/bin/bash\n\nproject_name=proxima-zvec\ngcov_tool=gcov\nzip_html=false\noutput_name=html\nkeep_info=false\n\nscript_dir=$(cd \"$(dirname \"$0\")\"; pwd)\nsource_base=$(dirname \"$script_dir\")\nfilter_list=\"'*/tests/*' '*/thirdparty/*' '*/deps/*' '*/proto/*' '*/external/*' '*/sqlengine/antlr/gen/*'\"\n\nwhile getopts t:p:o:zk option; do\n case \"$option\" in\n t)\n gcov_tool=$OPTARG;;\n p)\n project_name=$OPTARG;;\n o)\n output_name=$OPTARG;;\n z)\n zip_html=true;;\n k)\n keep_info=true;;\n esac\ndone\n\n# Process sources\nlcov -c -b \"$source_base\" -d . -o $project_name.lcov.info --gcov-tool=$gcov_tool --no-external || exit 1\neval $(echo lcov -r $project_name.lcov.info -o $project_name-filtered.lcov.info $filter_list) || exit 1\n\n# Gather HTML files\ngenhtml -t \"$project_name\" -o $output_name $project_name-filtered.lcov.info || exit 1\nif [ \"$keep_info\" = false ]; then\n rm -rf *.lcov.info\nfi\n\n# Zip HTML files\nif $zip_html ; then\n zip -r $output_name.zip $output_name/\nfi\n" }, { "path": "src/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\n# Retrieve version from git repository\ngit_version(ZVEC_VERSION ${CMAKE_CURRENT_SOURCE_DIR})\n\n# Add repository\ncc_directory(ailego)\ncc_directory(turbo)\ncc_directory(core)\ncc_directory(db)\nif(BUILD_PYTHON_BINDINGS)\n cc_directory(binding)\nendif()\n" }, { "path": "src/ailego/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\nfind_package(Threads REQUIRED)\n\nif(UNIX AND NOT APPLE AND NOT ANDROID)\n find_library(LIB_RT NAMES rt)\nelse()\n set(LIB_RT \"\")\nendif()\n\ngit_version(GIT_SRCS_VER ${CMAKE_CURRENT_SOURCE_DIR})\nfile(GLOB_RECURSE ALL_SRCS *.cc *.c *.h)\n\nset(EXTRA_LIBS ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS})\n\nif(UNIX AND NOT APPLE)\n list(APPEND EXTRA_LIBS ${LIB_RT})\nendif()\n\nif(NOT ANDROID AND AUTO_DETECT_ARCH)\n if(CMAKE_SYSTEM_PROCESSOR MATCHES \"x86_64|i686|i386|x64\")\n setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512 MATH_MARCH_FLAG_AVX512FP16)\n message(STATUS \"best compiler march, sse: \" ${MATH_MARCH_FLAG_SSE} \", avx2: \" ${MATH_MARCH_FLAG_AVX2} \", avx512: \" ${MATH_MARCH_FLAG_AVX512} \", avx512fp16: \" ${MATH_MARCH_FLAG_AVX512FP16})\n\n file(GLOB_RECURSE MATH_FILES_SSE\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_sse.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_sse.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_sse.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_sse.c\n )\n\n file(GLOB_RECURSE MATH_FILES_AVX2\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx2.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx2.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx2.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx2.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx.c\n )\n\n file(GLOB_RECURSE MATH_FILES_AVX512\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512.c\n )\n\n file(GLOB_RECURSE MATH_FILES_AVX512FP16\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512fp16.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_avx512fp16.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512fp16.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_avx512fp16.c\n )\n\n foreach(MATH_FILE ${MATH_FILES_SSE})\n set_source_files_properties(\n ${MATH_FILE}\n PROPERTIES\n COMPILE_FLAGS \"${MATH_MARCH_FLAG_SSE}\"\n )\n endforeach()\n\n foreach(MATH_FILE ${MATH_FILES_AVX2})\n set_source_files_properties(\n ${MATH_FILE}\n PROPERTIES\n COMPILE_FLAGS \"${MATH_MARCH_FLAG_AVX2}\"\n )\n endforeach()\n\n foreach(MATH_FILE ${MATH_FILES_AVX512})\n set_source_files_properties(\n ${MATH_FILE}\n PROPERTIES\n COMPILE_FLAGS \"${MATH_MARCH_FLAG_AVX512}\"\n )\n endforeach()\n\n foreach(MATH_FILE ${MATH_FILES_AVX512FP16})\n set_source_files_properties(\n ${MATH_FILE}\n PROPERTIES\n COMPILE_FLAGS \"${MATH_MARCH_FLAG_AVX512FP16}\"\n )\n endforeach()\n elseif(CMAKE_SYSTEM_PROCESSOR MATCHES \"aarch64|arm64|ARM64\")\n # set(CMAKE_CXX_FLAGS \"-march=armv8-a\")\n # set(CMAKE_C_FLAGS \"-march=armv8-a\")\n set(MATH_MARCH_FLAG_NEON \"-march=armv8-a\")\n\n file(GLOB_RECURSE MATH_FILES_NEON\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_dispatch.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_dispatch.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_neon.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math/*_neon.c\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_neon.cc\n ${CMAKE_CURRENT_SOURCE_DIR}/math_batch/*_neon.c\n )\n\n foreach(MATH_FILE ${MATH_FILES_NEON})\n set_source_files_properties(\n ${MATH_FILE}\n PROPERTIES\n COMPILE_FLAGS \"${MATH_MARCH_FLAG_NEON}\"\n )\n endforeach()\n endif()\nendif()\n\ncc_library(\n NAME zvec_ailego STATIC STRICT PACKED\n SRCS ${ALL_SRCS}\n LIBS ${EXTRA_LIBS}\n Arrow::arrow_static\n Arrow::parquet_static\n VERSION \"${GIT_SRCS_VER}\"\n)\n" }, { "path": "src/ailego/algorithm/binary_quantizer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"binary_quantizer.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n//! Feed the training data\nbool BinaryQuantizer::feed(const float *vec, size_t dim) {\n for (size_t i = 0; i < dim; ++i) {\n data_.emplace_back(vec[i]);\n }\n return true;\n}\n\n//! Train the quantizer\nbool BinaryQuantizer::train(void) {\n return true;\n}\n\n//! Quantize data: encode the float input to uint32_t output\nvoid BinaryQuantizer::encode(const float *in, size_t dim, uint32_t *out) const {\n for (size_t i = 0; i < dim; i += 32) {\n size_t remain = i + 32 <= dim ? 32 : dim - i;\n uint32_t data = 0;\n uint32_t mask = 1;\n\n for (size_t j = 0; j < remain; j++) {\n if (in[i + j] >= threshold_) {\n data |= mask;\n }\n\n mask <<= 1;\n }\n\n *out = data;\n out++;\n }\n}\n\n//! De-quantize data: decode the input uint32_t to float output\n//! bit value 1 will be mapped to 1.0\n//! bit value 0 will be mapped to -1.0\nvoid BinaryQuantizer::decode(const uint32_t *in, size_t dim, float *out) const {\n for (size_t i = 0; i < dim; ++i) {\n uint8_t bit = (in[i >> 5] >> (i & 31)) & 0x01;\n\n if (bit == 1) {\n out[i] = 1.0f;\n } else {\n out[i] = -1.0f;\n }\n\n // std::cout << \"dim: \" << i << \", value: \" << (size_t)bit << std::endl;\n }\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/algorithm/binary_quantizer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Binary Quantization Algorithm\n */\nclass BinaryQuantizer {\n public:\n //! Constructor\n BinaryQuantizer(void) {}\n\n //! Feed the training data\n bool feed(const float *vec, size_t dim);\n\n //! Train the quantizer\n bool train(void);\n\n //! Quantize data: encode the float input to uint32_t output\n void encode(const float *in, size_t dim, uint32_t *out) const;\n\n //! De-quantize data: decode the input uint32_t to float output\n void decode(const uint32_t *in, size_t dim, float *out) const;\n\n //! Get encoded elements in type of uint32_t\n static size_t EncodedSizeInBinary32(size_t dim) {\n return (dim + 31) / 32;\n }\n\n //! Set quantization threshold\n void set_threshold(float threshold) {\n threshold_ = threshold;\n }\n\n //! Get quantization threshold\n float threshold(void) const {\n return threshold_;\n }\n\n private:\n //! Disable them\n BinaryQuantizer(const BinaryQuantizer &) = delete;\n BinaryQuantizer &operator=(const BinaryQuantizer &) = delete;\n\n private:\n //! Members\n std::vector data_{};\n float threshold_{0.0f};\n};\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/algorithm/integer_quantizer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"integer_quantizer.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n//! Make smooth the distribution to eliminate zero in hist\nstatic inline void MakeSmooth(std::vector &dist) {\n constexpr float epsilon = std::numeric_limits::epsilon();\n\n // L1 Normalize first\n float norm = 1.0f;\n Normalizer::L1(dist.data(), dist.size(), &norm);\n\n size_t zero_count = std::count_if(dist.begin(), dist.end(), [](float val) {\n return (std::abs(val) < std::numeric_limits::epsilon());\n });\n size_t nonzero_count = dist.size() - zero_count;\n\n // Double check\n if (nonzero_count == 0 || zero_count == 0) {\n return;\n }\n\n float y = epsilon * zero_count / static_cast(nonzero_count);\n for (auto &it : dist) {\n if (std::abs(it) < epsilon) {\n it += epsilon;\n } else {\n it -= y;\n }\n } // end of for\n}\n\n//! Compute the Entropy of distribution p/q by Kullback-Leibler Divergence\nstatic inline double ComputeKlDivergence(const std::vector &p,\n const std::vector &q) {\n if (p.size() != q.size() || p.size() == 0) {\n return std::numeric_limits::max();\n }\n\n double v = 0.0f;\n for (size_t i = 0; i != p.size(); ++i) {\n if (p[i] == 0 || q[i] == 0) {\n return std::numeric_limits::max();\n }\n v += p[i] * std::log(static_cast(p[i]) / static_cast(q[i]));\n }\n return v;\n}\n\n//! Expand the quantization distribution to origin distribution in\n//! [-threshold, threshold]\nstatic inline void ExpandCandidateDistribution(\n const std::vector &distribution,\n const std::vector &quantized_distribution, size_t threshold,\n std::vector *expand_distribution) {\n expand_distribution->resize(threshold * 2, 0);\n float merged_cnt = static_cast(expand_distribution->size()) /\n quantized_distribution.size();\n size_t left_boundary = distribution.size() / 2 - threshold;\n\n for (size_t i = 0; i < quantized_distribution.size(); ++i) {\n float start = i * merged_cnt;\n float end = start + merged_cnt;\n const size_t start_ceil = static_cast(std::ceil(start));\n const size_t end_floor = static_cast(std::floor(end));\n float left_ratio = static_cast(start_ceil) - start;\n float right_ratio = end - static_cast(end_floor);\n float nonzero_count = 0;\n\n //! Count the non-zeros bins, if the histogram bin is partially included,\n //! non-zero bins is also partially counted\n if (left_ratio > 0 && left_boundary + start_ceil > 0) {\n if (distribution[left_boundary + start_ceil - 1] != 0) {\n nonzero_count += left_ratio;\n }\n }\n if (right_ratio > 0 && left_boundary + end_floor < distribution.size()) {\n if (distribution[left_boundary + end_floor] != 0) {\n nonzero_count += right_ratio;\n }\n }\n for (size_t j = start_ceil; j < end_floor; j++) {\n nonzero_count += distribution[left_boundary + j] != 0;\n }\n if (nonzero_count == 0) {\n continue;\n }\n\n //! expand the quantized value\n float value = quantized_distribution[i] / nonzero_count;\n if (left_ratio > 0 && start_ceil > 0) {\n (*expand_distribution)[start_ceil - 1] += value * left_ratio;\n }\n if (right_ratio > 0 && end_floor < expand_distribution->size()) {\n (*expand_distribution)[end_floor] += value * right_ratio;\n }\n for (size_t j = start_ceil; j < end_floor; j++) {\n if (distribution[left_boundary + j] != 0) {\n (*expand_distribution)[j] = value;\n }\n } // end of for\n } // end of for\n}\n\n/*! Compute quantization threshold bins\n * Implement Int8 Quantization Algorithm ref:\n * http://on-demand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf\n */\nstatic inline size_t ComputeThreshold(const std::vector &hist,\n const size_t target_bins) {\n std::vector P_distribution(hist.size());\n size_t zero_point_index = hist.size() / 2;\n\n size_t start_bin = target_bins / 2;\n size_t end_bin = hist.size() / 2;\n size_t negative_outliers_count = 0;\n size_t positive_outliers_count = 0;\n double min_divergence = std::numeric_limits::max();\n size_t target_threshold = end_bin;\n\n for (size_t threshold = start_bin; threshold <= end_bin; ++threshold) {\n negative_outliers_count += hist[zero_point_index - threshold];\n positive_outliers_count += hist[zero_point_index + threshold - 1];\n }\n\n //! for each zero-axised quantization range: [-threshold, threshold], search\n //! the best solution\n for (size_t threshold = start_bin; threshold <= end_bin; ++threshold) {\n P_distribution.resize(threshold * 2);\n auto p_hist = &hist[zero_point_index - threshold];\n for (size_t i = 0; i != P_distribution.size(); ++i) {\n P_distribution[i] = static_cast(p_hist[i]);\n }\n\n negative_outliers_count -= hist[zero_point_index - threshold];\n positive_outliers_count -= hist[zero_point_index + threshold - 1];\n P_distribution[0] += negative_outliers_count;\n P_distribution[P_distribution.size() - 1] += positive_outliers_count;\n\n //! Quantize the bins in range [-threshold, threshold] to target_bins\n std::vector Q_distribution(target_bins, 0);\n float merged_cnt = static_cast(threshold * 2) / target_bins;\n size_t left_boundary = zero_point_index - threshold;\n for (size_t i = 0; i < target_bins; ++i) {\n float start = i * merged_cnt;\n float end = start + merged_cnt;\n const size_t start_ceil = static_cast(std::ceil(start));\n const size_t end_floor = static_cast(std::floor(end));\n if (left_boundary + start_ceil > 0) {\n Q_distribution[i] +=\n ((float)start_ceil - start) * hist[left_boundary + start_ceil - 1];\n }\n if (left_boundary + end_floor < hist.size()) {\n Q_distribution[i] +=\n (end - (float)end_floor) * hist[left_boundary + end_floor];\n }\n\n for (size_t j = start_ceil; j < end_floor; j++) {\n Q_distribution[i] += hist[left_boundary + j];\n }\n }\n std::vector Q_expand_distribution;\n ExpandCandidateDistribution(hist, Q_distribution, threshold,\n &Q_expand_distribution);\n\n //! Compute Kullback-Leibler Divergence, normalize the smooth the data\n //! first. Ref: http://hanj.cs.illinois.edu/cs412/bk3/KL-divergence.pdf\n MakeSmooth(P_distribution);\n MakeSmooth(Q_expand_distribution);\n double divergence =\n ComputeKlDivergence(P_distribution, Q_expand_distribution);\n\n if (divergence < min_divergence) {\n min_divergence = divergence;\n target_threshold = threshold;\n }\n }\n return target_threshold;\n}\n\n// Quantize the value in range\ntemplate \nstatic inline float QuantizeValue(float val, float scale, float bias) {\n val = (val + bias) * scale;\n\n if (val > RANGE_MAX) {\n val = RANGE_MAX;\n } else if (val < RANGE_MIN) {\n val = RANGE_MIN;\n }\n return val;\n}\n\n// Init the historgram params\n#define INIT_HISTOGRAM() \\\n { \\\n if (histogram_bins_ == 0) { \\\n size_t range = non_bias_ \\\n ? std::max(std::abs(MIN_VALUE), std::abs(MAX_VALUE)) \\\n : (MAX_VALUE - MIN_VALUE); \\\n histogram_bins_ = std::max(4096u, range * 8); \\\n } \\\n histogram_.resize((histogram_bins_ + 1) >> 1 << 1); \\\n if (non_bias_) { \\\n bias_ = 0.0f; \\\n auto val = std::max(std::abs(max_), std::abs(min_)); \\\n left_boundary_ = -val; \\\n hist_interval_ = (val * 2) / static_cast(histogram_.size()); \\\n } else { \\\n bias_ = -static_cast(min_ + (max_ - min_) * 0.5); \\\n left_boundary_ = min_; \\\n hist_interval_ = (max_ - min_) / static_cast(histogram_.size()); \\\n } \\\n }\n\n// Feed vector and update the historgram\n#define UPDATE_HISTOGRAM(vec, dim) \\\n { \\\n if (max_ < min_) { \\\n return false; \\\n } \\\n if (histogram_.size() == 0) { \\\n INIT_HISTOGRAM() \\\n } \\\n for (size_t i = 0; i < dim; ++i) { \\\n ssize_t index = 0; \\\n if (hist_interval_ > 0.0) { \\\n index = \\\n static_cast((vec[i] - left_boundary_) / hist_interval_); \\\n } \\\n if (index < 0) { \\\n index = 0; \\\n } else if ((size_t)index >= histogram_.size()) { \\\n index = histogram_.size() - 1; \\\n } \\\n ailego_assert_with((size_t)index < histogram_.size(), \"Invalid index\"); \\\n histogram_[index] += 1; \\\n } \\\n return true; \\\n }\n\n// Train the quantizer\n#define TRAIN_QUANTIZER() \\\n { \\\n auto sum = std::accumulate(histogram_.begin(), histogram_.end(), 0); \\\n if (sum == 0) { \\\n return false; \\\n } \\\n size_t target_bins = \\\n ailego_align(static_cast(MAX_VALUE - MIN_VALUE), 2); \\\n auto threshold_bins = ComputeThreshold(histogram_, target_bins); \\\n auto threshold = \\\n (static_cast(threshold_bins) + 0.5f) * hist_interval_; \\\n scale_ = target_bins / 2 / threshold; \\\n if (!non_bias_) { \\\n bias_ += (MAX_VALUE + MIN_VALUE) * 0.5f / scale_; \\\n } \\\n scale_reciprocal_ = 1 / scale_; \\\n return true; \\\n }\n\n// Feed the INT16 quantizer\nbool EntropyInt16Quantizer::feed(const float *vec, size_t dim) {\n UPDATE_HISTOGRAM(vec, dim)\n}\n\n// Train the INT16 quantizer\nbool EntropyInt16Quantizer::train(void) {\n TRAIN_QUANTIZER()\n}\n\n// Encode to INT16\nvoid EntropyInt16Quantizer::encode(const float *in, size_t dim,\n int16_t *out) const {\n for (size_t i = 0; i < dim; ++i) {\n out[i] = static_cast(\n std::round(QuantizeValue(in[i], scale_, bias_)));\n }\n}\n\n// Decode from INT16\nvoid EntropyInt16Quantizer::decode(const int16_t *in, size_t dim,\n float *out) const {\n for (size_t i = 0; i < dim; ++i) {\n out[i] = in[i] * this->scale_reciprocal() - this->bias();\n }\n}\n\n// Feed the UINT16 quantizer\nbool EntropyUInt16Quantizer::feed(const float *vec, size_t dim) {\n UPDATE_HISTOGRAM(vec, dim)\n}\n\n// Train the UINT16 quantizer\nbool EntropyUInt16Quantizer::train(void) {\n TRAIN_QUANTIZER()\n}\n\n// Encode to UINT16\nvoid EntropyUInt16Quantizer::encode(const float *in, size_t dim,\n uint16_t *out) const {\n for (size_t i = 0; i < dim; ++i) {\n out[i] = static_cast_from_float_to_uint16(\n std::round(QuantizeValue(in[i], scale_, bias_)));\n }\n}\n\n// Decode from INT16\nvoid EntropyUInt16Quantizer::decode(const uint16_t *in, size_t dim,\n float *out) const {\n for (size_t i = 0; i < dim; ++i) {\n out[i] = in[i] * this->scale_reciprocal() - this->bias();\n }\n}\n\n// Feed the INT8 quantizer\nbool EntropyInt8Quantizer::feed(const float *vec, size_t dim) {\n UPDATE_HISTOGRAM(vec, dim)\n}\n\n// Train the INT8 quantizer\nbool EntropyInt8Quantizer::train(void) {\n TRAIN_QUANTIZER()\n}\n\n// Encode to INT8\nvoid EntropyInt8Quantizer::encode(const float *in, size_t dim,\n int8_t *out) const {\n for (size_t i = 0; i < dim; ++i) {\n out[i] = static_cast(\n std::round(QuantizeValue(in[i], scale_, bias_)));\n }\n}\n\n// Decode from INT8\nvoid EntropyInt8Quantizer::decode(const int8_t *in, size_t dim,\n float *out) const {\n for (size_t i = 0; i < dim; ++i) {\n out[i] = in[i] * this->scale_reciprocal() - this->bias();\n }\n}\n\n// Feed the UINT8 quantizer\nbool EntropyUInt8Quantizer::feed(const float *vec, size_t dim) {\n UPDATE_HISTOGRAM(vec, dim)\n}\n\n// Train the UINT8 quantizer\nbool EntropyUInt8Quantizer::train(void) {\n TRAIN_QUANTIZER()\n}\n\n// Encode to INT8\nvoid EntropyUInt8Quantizer::encode(const float *in, size_t dim,\n uint8_t *out) const {\n for (size_t i = 0; i < dim; ++i) {\n out[i] = static_cast_from_float_to_uint8(\n std::round(QuantizeValue(in[i], scale_, bias_)));\n }\n}\n\n// Decode from UINT8\nvoid EntropyUInt8Quantizer::decode(const uint8_t *in, size_t dim,\n float *out) const {\n for (size_t i = 0; i < dim; ++i) {\n out[i] = in[i] * this->scale_reciprocal() - this->bias();\n }\n}\n\n// Feed the INT4 quantizer\nbool EntropyInt4Quantizer::feed(const float *vec, size_t dim) {\n UPDATE_HISTOGRAM(vec, dim)\n}\n\n// Train the INT4 quantizer\nbool EntropyInt4Quantizer::train(void) {\n TRAIN_QUANTIZER()\n}\n\n// Encode to INT4\nvoid EntropyInt4Quantizer::encode(const float *in, size_t dim,\n uint8_t *out) const {\n ailego_assert_with(dim % 2 == 0, \"Dimension must be aligned with 2\");\n\n for (size_t i = 0; i < dim; i += 2) {\n float lo = QuantizeValue(in[i], scale_, bias_);\n float hi = QuantizeValue(in[i + 1], scale_, bias_);\n out[i / 2] = (static_cast_from_float_to_uint8(std::round(hi)) << 4) |\n (static_cast_from_float_to_uint8(std::round(lo)) & 0xF);\n }\n}\n\n// Decode from INT4\nvoid EntropyInt4Quantizer::decode(const uint8_t *in, size_t dim,\n float *out) const {\n ailego_assert_with(dim % 2 == 0, \"Dimension must be aligned with 2\");\n\n size_t size = dim / 2;\n for (size_t i = 0; i < size; i += 1) {\n uint8_t v = in[i];\n int8_t lo = (static_cast(v << 4) >> 4);\n int8_t hi = (static_cast(v & 0xf0) >> 4);\n out[2 * i] = lo * this->scale_reciprocal() - this->bias();\n out[2 * i + 1] = hi * this->scale_reciprocal() - this->bias();\n }\n}\n\n// Feed the UINT4 quantizer\nbool EntropyUInt4Quantizer::feed(const float *vec, size_t dim) {\n UPDATE_HISTOGRAM(vec, dim)\n}\n\n// Train the UINT4 quantizer\nbool EntropyUInt4Quantizer::train(void) {\n TRAIN_QUANTIZER()\n}\n\n// Encode to INT4\nvoid EntropyUInt4Quantizer::encode(const float *in, size_t dim,\n uint8_t *out) const {\n ailego_assert_with(dim % 2 == 0, \"Dimension must be aligned with 2\");\n\n for (size_t i = 0; i < dim; i += 2) {\n float lo = QuantizeValue(in[i], scale_, bias_);\n float hi = QuantizeValue(in[i + 1], scale_, bias_);\n out[i / 2] = (static_cast_from_float_to_uint8(std::round(hi)) << 4) |\n (static_cast_from_float_to_uint8(std::round(lo)) & 0xF);\n }\n}\n\n// Decode from INT4\nvoid EntropyUInt4Quantizer::decode(const uint8_t *in, size_t dim,\n float *out) const {\n ailego_assert_with(dim % 2 == 0, \"Dimension must be aligned with 2\");\n\n size_t size = dim / 2;\n for (size_t i = 0; i < size; i += 1) {\n uint8_t v = in[i];\n out[2 * i] = (v & 0xf) * this->scale_reciprocal() - this->bias();\n out[2 * i + 1] = (v >> 4) * this->scale_reciprocal() - this->bias();\n }\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/algorithm/integer_quantizer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Entropy-based Integer Quantization Algorithm\n */\ntemplate \nclass EntropyIntegerQuantizer {\n public:\n //! Primitive Built-in Types to store the quantized data\n using ValueType = typename std::remove_cv::type;\n\n //! Constants\n constexpr static int MIN_VALUE = RANGE_MIN;\n constexpr static int MAX_VALUE = RANGE_MAX;\n\n // Check supporting type\n static_assert(std::is_integral::value, \"ValueType must be integral\");\n\n // Check template values\n static_assert(RANGE_MIN < RANGE_MAX, \"Invalid value range\");\n\n //! Constructor\n EntropyIntegerQuantizer(void) {}\n\n //! Set histogram bins in train\n void set_histogram_bins(size_t bins) {\n if (bins > (RANGE_MAX - RANGE_MIN)) {\n histogram_bins_ = bins;\n }\n }\n\n //! Set quantization params scale\n void set_scale(float val) {\n if (val > 0.0f) {\n scale_ = val;\n scale_reciprocal_ = 1 / scale_;\n }\n }\n\n //! Set quantization params bias\n void set_bias(float val) {\n bias_ = val;\n }\n\n //! Set quantization params max\n void set_max(float val) {\n max_ = val;\n }\n\n //! Set quantization params min\n void set_min(float val) {\n min_ = val;\n }\n\n //! Set quantization params non bias\n void set_non_bias(bool val) {\n non_bias_ = val;\n }\n\n //! Get histogram bins in train\n size_t histogram_bins(void) const {\n return histogram_bins_;\n }\n\n //! Get quantization params scale\n float scale(void) const {\n return scale_;\n }\n\n //! Get quantization params bias\n float bias(void) const {\n return bias_;\n }\n\n //! Get quantization params max\n float max(void) const {\n return max_;\n }\n\n //! Get quantization params min\n float min(void) const {\n return min_;\n }\n\n //! Get quantization params non bias\n bool non_bias(void) const {\n return non_bias_;\n }\n\n //! Retrieve the scale reciprocal for decoding\n float scale_reciprocal(void) const {\n return scale_reciprocal_;\n }\n\n protected:\n //! Disable them\n EntropyIntegerQuantizer(const EntropyIntegerQuantizer &) = delete;\n EntropyIntegerQuantizer &operator=(const EntropyIntegerQuantizer &) = delete;\n\n //! Members\n size_t histogram_bins_{0};\n float hist_interval_{1.0f};\n float max_{std::numeric_limits::min()};\n float min_{std::numeric_limits::max()};\n float bias_{0.0f};\n float scale_{0.0f};\n float scale_reciprocal_{0.0f};\n float left_boundary_{0.0f};\n bool non_bias_{false};\n std::vector histogram_{};\n};\n\n/*! INT16 Quantizer\n */\nclass EntropyInt16Quantizer\n : public EntropyIntegerQuantizer {\n public:\n //! Feed the training data\n bool feed(const float *vec, size_t dim);\n\n //! Train the quantizer\n bool train(void);\n\n //! Encode float vector to int16\n void encode(const float *in, size_t dim, ValueType *out) const;\n\n //! Decode to float vector from int16\n void decode(const ValueType *in, size_t dim, float *out) const;\n};\n\n/*! UINT16 Quantizer\n */\nclass EntropyUInt16Quantizer\n : public EntropyIntegerQuantizer {\n public:\n //! Feed the training data\n bool feed(const float *vec, size_t dim);\n\n //! Train the quantizer\n bool train(void);\n\n //! Encode float vector to uint16\n void encode(const float *in, size_t dim, ValueType *out) const;\n\n //! Decode to float vector from uint16\n void decode(const ValueType *in, size_t dim, float *out) const;\n};\n\n/*! INT8 Quantizer\n */\nclass EntropyInt8Quantizer : public EntropyIntegerQuantizer {\n public:\n //! Feed the training data\n bool feed(const float *vec, size_t dim);\n\n //! Train the quantizer\n bool train(void);\n\n //! Encode float vector to int8\n void encode(const float *in, size_t dim, ValueType *out) const;\n\n //! Decode to float vector from int8\n void decode(const ValueType *in, size_t dim, float *out) const;\n};\n\n/*! UINT8 Quantizer\n */\nclass EntropyUInt8Quantizer : public EntropyIntegerQuantizer {\n public:\n //! Feed the training data\n bool feed(const float *vec, size_t dim);\n\n //! Train the quantizer\n bool train(void);\n\n //! Encode float vector to uint8\n void encode(const float *in, size_t dim, ValueType *out) const;\n\n //! Decode to float vector from uint8\n void decode(const ValueType *in, size_t dim, float *out) const;\n};\n\n/*! INT4 Quantizer\n */\nclass EntropyInt4Quantizer : public EntropyIntegerQuantizer {\n public:\n //! Feed the training data\n bool feed(const float *vec, size_t dim);\n\n //! Train the quantizer\n bool train(void);\n\n //! Encode float vector to int4\n void encode(const float *in, size_t dim, ValueType *out) const;\n\n //! Decode to float vector from int4\n void decode(const ValueType *in, size_t dim, float *out) const;\n};\n\n/*! UINT4 Quantizer\n */\nclass EntropyUInt4Quantizer : public EntropyIntegerQuantizer {\n public:\n //! Feed the training data\n bool feed(const float *vec, size_t dim);\n\n //! Train the quantizer\n bool train(void);\n\n //! Encode float vector to uint4\n void encode(const float *in, size_t dim, ValueType *out) const;\n\n //! Decode to float vector from uint4\n void decode(const ValueType *in, size_t dim, float *out) const;\n};\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/algorithm/kmeans.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"lloyd_cluster.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n/*! K-MC2 Centroids Generator\n */\ntemplate \nclass Kmc2CentroidsGenerator {\n public:\n //! Type of values\n using OwnerType = typename std::decay::type;\n using ContainerType = typename OwnerType::ContainerType;\n using ContextType = typename OwnerType::ContextType;\n using ValueType = typename OwnerType::ValueType;\n using StoreType = typename OwnerType::StoreType;\n using ThreadPoolType = TPool;\n\n //! constexpr variables\n constexpr static size_t BatchCount = OwnerType::BatchCount;\n\n //! Generate centroids\n void operator()(OwnerType *owner, ThreadPoolType &pool) const {\n if (chain_length_ == 0) {\n this->init_centroids_random(owner);\n } else if (!assumption_free_) {\n this->init_centroids_kmc2(owner, pool);\n } else {\n this->init_centroids_afkmc2(owner, pool);\n }\n }\n\n //! Retrieve the markov chain length\n size_t chain_length(void) const {\n return chain_length_;\n }\n\n //! Set the mutable markov chain length\n void set_chain_length(size_t len) {\n chain_length_ = len;\n }\n\n //! Retrieve assumption free option\n bool assumption_free(void) const {\n return assumption_free_;\n }\n\n //! Set the assumption free option\n void set_assumption_free(bool val) {\n assumption_free_ = val;\n }\n\n protected:\n //! Initialize centroids randomly\n void init_centroids_random(OwnerType *owner) const {\n RandomSelectBenches(owner->feature_cache(), owner->feature_matrix(),\n owner->k_value(), owner->mutable_centroids());\n }\n\n //! Initialize centroids with K-MC2\n void init_centroids_kmc2(OwnerType *owner, ThreadPoolType &pool) const {\n const auto &matrix = owner->feature_matrix();\n const auto &cache = owner->feature_cache();\n auto *centroids = owner->mutable_centroids();\n\n std::mt19937 mt((std::random_device())());\n std::uniform_real_distribution dist(0.0, 1.0);\n\n ContainerType benches(cache.dimension());\n std::vector scores;\n\n // Sample first center uniformly\n RandomSelectBenches(cache, matrix, 1, centroids);\n\n // Make a thread group\n auto group = pool.make_group();\n\n for (size_t i = 1, k = owner->k_value(); i < k; ++i) {\n RandomSelectBenches(cache, matrix, chain_length_, &benches);\n\n // Update bench scores\n scores.resize(benches.count());\n for (size_t j = 0; j != scores.size(); ++j) {\n group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateBenchScores,\n centroids, benches[j], &scores[j]));\n }\n group->wait_finish();\n\n //! Select the better centroid randomly\n float x = scores[0];\n size_t xj = 0;\n for (size_t j = 1; j != scores.size(); ++j) {\n float y = scores[j];\n\n if (x == 0.0f || x * dist(mt) < y) {\n x = y;\n xj = j;\n }\n }\n centroids->append(benches[xj], benches.dimension());\n } // end of for\n }\n\n //! Initialize centroids with K-MC2\n void init_centroids_afkmc2(OwnerType *owner, ThreadPoolType &pool) const {\n const auto &matrix = owner->feature_matrix();\n const auto &cache = owner->feature_cache();\n\n // Probability\n std::vector probs(matrix.count() + cache.count());\n\n // Sample first center uniformly\n RandomSelectBenches(cache, matrix, 1, owner->mutable_centroids());\n\n // Make a thread group\n auto group = pool.make_group();\n if (!matrix.empty()) {\n size_t n = matrix.count() / BatchCount;\n size_t c = std::max(n / pool.count() / 2u, 1u);\n size_t m = n / c * c;\n\n for (size_t i = 0; i != m; i += c) {\n group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateMatrixScores,\n owner, i, i + c, &probs[0]));\n }\n for (size_t i = m; i != n; i += 1) {\n group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateMatrixScores,\n owner, i, i + 1, &probs[0]));\n }\n }\n if (!cache.empty()) {\n group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateCacheScores,\n owner, &probs[matrix.count()]));\n }\n group->wait_finish();\n\n // Update probabilities\n double p_sum = std::accumulate(probs.begin(), probs.end(), 0.0);\n for (auto it = probs.begin(); it != probs.end(); ++it) {\n *it = static_cast((*it / p_sum + 1.0 / probs.size()) * 0.5);\n }\n\n std::mt19937 mt((std::random_device())());\n std::uniform_real_distribution dist(0.0, 1.0);\n ContainerType benches(cache.dimension());\n std::vector scores;\n std::vector bench_probs;\n\n for (size_t i = 1; i < owner->k_value(); ++i) {\n RandomSelectBenches(cache, matrix, chain_length_, probs, &benches,\n &bench_probs);\n\n // Update bench scores\n scores.resize(benches.count());\n for (size_t j = 0; j != scores.size(); ++j) {\n group->submit(Closure::New(&Kmc2CentroidsGenerator::UpdateBenchScores,\n owner->mutable_centroids(), benches[j],\n &scores[j]));\n }\n group->wait_finish();\n\n // Update scores with probabilities\n for (size_t j = 0; j != scores.size(); ++j) {\n scores[j] /= bench_probs[j];\n }\n\n //! Select the better centroid randomly\n float x = scores[0];\n size_t xj = 0;\n for (size_t j = 1; j != scores.size(); ++j) {\n float y = scores[j];\n\n if (x == 0.0f || x * dist(mt) < y) {\n x = y;\n xj = j;\n }\n }\n owner->mutable_centroids()->append(benches[xj], benches.dimension());\n } // end of for\n }\n\n //! Update matrix score\n static void UpdateMatrixScores(const OwnerType *owner, size_t first,\n size_t last, float *out) {\n const auto &matrix = owner->feature_matrix();\n const auto *bench = owner->centroids().data();\n\n for (size_t i = first * BatchCount; i != last * BatchCount;\n i += BatchCount) {\n ContextType::template BatchDistance<1>(matrix[i], bench,\n matrix.dimension(), &out[i]);\n }\n }\n\n //! Update cache score\n static void UpdateCacheScores(const OwnerType *owner, float *out) {\n const auto &cache = owner->feature_cache();\n const auto *bench = owner->centroids().data();\n\n for (size_t i = 0, n = cache.count(); i != n; ++i) {\n ContextType::Distance(bench, cache[i], cache.dimension(), &out[i]);\n }\n }\n\n //! Update bench score\n static void UpdateBenchScores(const ContainerType *benches,\n const StoreType *feat, float *out) {\n float min_score = std::numeric_limits::max();\n\n for (size_t i = 0, c = benches->count(); i != c; ++i) {\n float new_score;\n ContextType::Distance(benches->at(i), feat, benches->dimension(),\n &new_score);\n\n if (new_score < min_score) {\n min_score = new_score;\n }\n }\n *out = min_score;\n }\n\n //! Select k benches randomly\n static void RandomSelectBenches(const ContainerType &cache,\n const ContainerType &matrix, size_t k,\n ContainerType *benches) {\n ContainerType rows(cache.dimension());\n size_t m = matrix.count();\n size_t n = m + cache.count();\n std::mt19937 mt((std::random_device())());\n\n rows.resize(BatchCount);\n benches->reset(cache.dimension());\n benches->reserve(k);\n\n for (size_t i = 0; k > 0 && i < n; ++i) {\n if (mt() % (n - i) >= k) {\n continue;\n }\n // Selected a feature\n if (i < m) {\n ContextType::MatrixReverseTranspose(matrix[i / BatchCount * BatchCount],\n matrix.dimension(), rows.data());\n benches->append(rows[i & (BatchCount - 1u)], matrix.dimension());\n } else {\n benches->append(cache[i - m], cache.dimension());\n }\n --k;\n } // end of for\n }\n\n //! Select k benches randomly\n static void RandomSelectBenches(const ContainerType &cache,\n const ContainerType &matrix, size_t k,\n const std::vector &probs,\n ContainerType *benches,\n std::vector *bench_probs) {\n std::mt19937 mt((std::random_device())());\n std::uniform_real_distribution dist(0.0, 1.0);\n\n // Sample features\n KeyValueHeap> samples(k);\n for (size_t i = 0; i < probs.size(); ++i) {\n samples.emplace(i, std::pow(dist(mt), 1.0 / probs[i]));\n }\n\n ContainerType rows(cache.dimension());\n size_t matrix_count = matrix.count();\n\n rows.resize(BatchCount);\n benches->reset(cache.dimension());\n benches->reserve(k);\n bench_probs->clear();\n bench_probs->reserve(k);\n\n for (const auto &it : samples) {\n // Selected a feature\n if (it.first < matrix_count) {\n ContextType::MatrixReverseTranspose(\n matrix[it.first / BatchCount * BatchCount], matrix.dimension(),\n rows.data());\n benches->append(rows[it.first & (BatchCount - 1u)], matrix.dimension());\n } else {\n benches->append(cache[it.first - matrix_count], cache.dimension());\n }\n bench_probs->push_back(probs[it.first]);\n }\n }\n\n private:\n size_t chain_length_{32};\n bool assumption_free_{false};\n};\n\n/*! Numerical K-Means Context\n */\ntemplate \nclass NumericalKmeansContext {\n public:\n //! constexpr variables\n constexpr static size_t BatchCount = BATCH_COUNT;\n\n //! Type of values\n using ValueType = typename std::remove_cv::type;\n using StoreType = typename std::remove_cv::type;\n\n // Check supporting type\n static_assert(IsSignedArithmetic::value,\n \"ValueType must be signed arithmetic\");\n\n /*! K-Means Context Cluster\n */\n class Cluster {\n public:\n //! Constructor\n Cluster(size_t dim) : accum_(dim, 0.0) {}\n\n //! Constructor\n Cluster(const Cluster &rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(rhs.accum_) {}\n\n //! Constructor\n Cluster(Cluster &&rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(std::move(rhs.accum_)) {}\n\n //! Assignment\n Cluster &operator=(const Cluster &rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = rhs.accum_;\n return *this;\n }\n\n //! Assignment\n Cluster &operator=(Cluster &&rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = std::move(rhs.accum_);\n return *this;\n }\n\n //! Append a vector\n void append(const ValueType *vec, size_t dim, float dist) {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n mutex_.lock();\n cost_ += dist;\n count_ += 1;\n\n for (size_t i = 0; i != dim; ++i) {\n accum_[i] += vec[i];\n }\n mutex_.unlock();\n }\n\n //! Retrieve the centroid of vectors\n void centroid(ValueType *out, size_t dim) const {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n for (size_t i = 0; i != dim; ++i) {\n out[i] = count_ == 0 ? FloatCast(NAN)\n : FloatCast(accum_[i] / count_);\n }\n }\n\n //! Retrieve squared error\n double cost(void) const {\n return cost_;\n }\n\n //! Retrieve feature count\n size_t count(void) const {\n return count_;\n }\n\n protected:\n //! Convert float type to another type\n template \n static auto FloatCast(const double val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(val);\n }\n\n //! Convert float type to another type\n template \n static auto FloatCast(const double val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(std::round(val));\n }\n\n private:\n SpinMutex mutex_{};\n double cost_{0.0};\n size_t count_{0u};\n std::vector accum_{};\n };\n\n //! operator []\n const Cluster &operator[](size_t i) const {\n return clusters_[i];\n }\n\n //! operator []\n Cluster &operator[](size_t i) {\n return clusters_[i];\n }\n\n //! Clear the context\n void clear(void) {\n clusters_.clear();\n }\n\n //! Reset the context\n void reset(size_t k_value, size_t dim) {\n clusters_.clear();\n clusters_.resize(k_value, dim);\n }\n\n //! Retrieve context of clusters\n const std::vector &clusters(void) const {\n return clusters_;\n }\n\n //! Compute the distance between matrix and query (batch)\n template \n static void BatchDistance(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n SquaredEuclideanDistanceMatrix::Compute(m, q, dim,\n out);\n }\n\n //! Compute the distance between matrix and query (single)\n static void Distance(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out);\n }\n\n //! Transpose a matrix\n template \n static auto MatrixTranspose(const U *src, size_t dim, T *dst) ->\n typename std::enable_if= 2>::type {\n MatrixHelper::Transpose(src, dim, dst);\n }\n\n //! Transpose a matrix\n template \n static auto MatrixTranspose(const U *src, size_t dim, U *dst) ->\n typename std::enable_if::type {\n MatrixHelper::Transpose(src, dim >> 2, dst);\n }\n\n //! Reverse transpose a matrix\n template \n static auto MatrixReverseTranspose(const U *src, size_t dim, U *dst) ->\n typename std::enable_if= 2>::type {\n MatrixHelper::ReverseTranspose(src, dim, dst);\n }\n\n //! Reverse transpose a matrix\n template \n static auto MatrixReverseTranspose(const U *src, size_t dim, U *dst) ->\n typename std::enable_if::type {\n MatrixHelper::ReverseTranspose(src, dim >> 2, dst);\n }\n\n //! Compute Norm2\n template ::value>::type>\n static void Norm2(ValueType *data, size_t dim, float *norm) {\n Normalizer::L2(data, dim, norm);\n }\n\n //! Compute Norm2, for non-float do nothing\n static void Norm2(ValueType * /*data*/, size_t /*dim*/, float *norm) {\n *norm = 0.0f;\n }\n\n private:\n //! Members\n std::vector clusters_{};\n};\n\n/*! Nibble K-Means Context (INT4)\n */\ntemplate \nclass NibbleKmeansContext {\n public:\n //! constexpr variables\n constexpr static size_t BatchCount = BATCH_COUNT;\n\n //! Type of values\n using ValueType = typename std::remove_cv::type;\n using StoreType = typename std::make_unsigned::type;\n\n // Check supporting type\n static_assert(std::is_same::value ||\n std::is_same::value,\n \"ValueType must be int32_t or int64_t\");\n\n /*! K-Means Context Cluster\n */\n class Cluster {\n public:\n //! Constructor\n Cluster(size_t dim) : accum_(dim, 0.0) {}\n\n //! Constructor\n Cluster(const Cluster &rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(rhs.accum_) {}\n\n //! Constructor\n Cluster(Cluster &&rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(std::move(rhs.accum_)) {}\n\n //! Assignment\n Cluster &operator=(const Cluster &rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = rhs.accum_;\n return *this;\n }\n\n //! Assignment\n Cluster &operator=(Cluster &&rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = std::move(rhs.accum_);\n return *this;\n }\n\n //! Append a vector\n void append(const StoreType *vec, size_t dim, float dist) {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n mutex_.lock();\n cost_ += dist;\n count_ += 1;\n\n const uint8_t *arr = reinterpret_cast(vec);\n dim = (dim >> 1) << 1;\n for (size_t i = 0; i != dim; i += 2) {\n uint8_t val = arr[i >> 1];\n accum_[i] += ((int8_t)(val << 4) >> 4);\n accum_[i + 1] += ((int8_t)(val) >> 4);\n }\n mutex_.unlock();\n }\n\n //! Retrieve the centroid of vectors\n void centroid(StoreType *out, size_t dim) const {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n uint8_t *arr = reinterpret_cast(out);\n dim = (dim >> 1) << 1;\n for (size_t i = 0; i != dim; i += 2) {\n int lo =\n count_ == 0 ? 0 : static_cast(std::round(accum_[i] / count_));\n int hi = count_ == 0\n ? 0\n : static_cast(std::round(accum_[i + 1] / count_));\n arr[i >> 1] = (uint8_t)((hi << 4) & 0xf0) | (uint8_t)(lo & 0xf);\n }\n }\n\n //! Retrieve squared error\n double cost(void) const {\n return cost_;\n }\n\n //! Retrieve feature count\n size_t count(void) const {\n return count_;\n }\n\n private:\n SpinMutex mutex_{};\n double cost_{0.0};\n size_t count_{0u};\n std::vector accum_{};\n };\n\n //! operator []\n const Cluster &operator[](size_t i) const {\n return clusters_[i];\n }\n\n //! operator []\n Cluster &operator[](size_t i) {\n return clusters_[i];\n }\n\n //! Clear the context\n void clear(void) {\n clusters_.clear();\n }\n\n //! Reset the context\n void reset(size_t k_value, size_t dim) {\n clusters_.clear();\n clusters_.resize(k_value, dim);\n }\n\n //! Retrieve context of clusters\n const std::vector &clusters(void) const {\n return clusters_;\n }\n\n //! Compute the distance between matrix and query (batch)\n template \n static void BatchDistance(const StoreType *m, const StoreType *q, size_t dim,\n float *out) {\n SquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, out);\n }\n\n //! Compute the distance between matrix and query (single)\n static void Distance(const StoreType *m, const StoreType *q, size_t dim,\n float *out) {\n SquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, out);\n }\n\n //! Transpose a matrix\n static void MatrixTranspose(const StoreType *src, size_t dim,\n StoreType *dst) {\n MatrixHelper::Transpose(src, dim >> 3, dst);\n }\n\n //! Reverse transpose a matrix\n static void MatrixReverseTranspose(const StoreType *src, size_t dim,\n StoreType *dst) {\n MatrixHelper::ReverseTranspose(src, dim >> 3, dst);\n }\n\n //! Compute and do norm2\n static void Norm2(StoreType * /*data*/, size_t /*dim*/, float *norm) {\n *norm = 0;\n }\n\n private:\n //! Members\n std::vector clusters_{};\n};\n\n/*! Binary K-Means Context\n */\ntemplate \nclass BinaryKmeansContext {\n public:\n //! constexpr variables\n constexpr static size_t BatchCount = BATCH_COUNT;\n\n //! Type of values\n using ValueType = typename std::remove_cv::type;\n using StoreType = typename std::remove_cv::type;\n\n // Check supporting type\n static_assert(std::is_same::value ||\n std::is_same::value,\n \"ValueType must be uint32_t or uint64_t\");\n\n /*! K-Means Context Cluster\n */\n class Cluster {\n public:\n //! Constructor\n Cluster(size_t dim) : accum_(dim, 0) {}\n\n //! Constructor\n Cluster(const Cluster &rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(rhs.accum_) {}\n\n //! Constructor\n Cluster(Cluster &&rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(std::move(rhs.accum_)) {}\n\n //! Assignment\n Cluster &operator=(const Cluster &rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = rhs.accum_;\n return *this;\n }\n\n //! Assignment\n Cluster &operator=(Cluster &&rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = std::move(rhs.accum_);\n return *this;\n }\n\n //! Append a vector\n void append(const ValueType *vec, size_t dim, float dist) {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n mutex_.lock();\n cost_ += dist;\n count_ += 1;\n\n const uint8_t *arr = reinterpret_cast(vec);\n for (size_t i = 0; i != dim; ++i) {\n if (arr[i >> 3] & (1u << (i & 7))) {\n accum_[i] += 1;\n }\n }\n mutex_.unlock();\n }\n\n //! Retrieve the centroid of vectors\n void centroid(ValueType *out, size_t dim) const {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n uint8_t *arr = reinterpret_cast(out);\n size_t half = count_ >> 1;\n for (size_t i = 0; i != dim; ++i) {\n if (accum_[i] > half) {\n arr[i >> 3] |= static_cast(1 << (i & 0x7));\n } else {\n arr[i >> 3] &= ~static_cast(1 << (i & 0x7));\n }\n }\n }\n\n //! Retrieve squared error\n double cost(void) const {\n return cost_;\n }\n\n //! Retrieve feature count\n size_t count(void) const {\n return count_;\n }\n\n private:\n SpinMutex mutex_{};\n double cost_{0.0};\n size_t count_{0u};\n std::vector accum_{};\n };\n\n //! operator []\n const Cluster &operator[](size_t i) const {\n return clusters_[i];\n }\n\n //! operator []\n Cluster &operator[](size_t i) {\n return clusters_[i];\n }\n\n //! Clear the context\n void clear(void) {\n clusters_.clear();\n }\n\n //! Reset the context\n void reset(size_t k_value, size_t dim) {\n clusters_.clear();\n clusters_.resize(k_value, dim);\n }\n\n //! Retrieve context of clusters\n const std::vector &clusters(void) const {\n return clusters_;\n }\n\n //! Compute the distance between matrix and query (batch)\n template \n static void BatchDistance(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n HammingDistanceMatrix::Compute(m, q, dim, out);\n }\n\n //! Compute the distance between matrix and query (single)\n static void Distance(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n HammingDistanceMatrix::Compute(m, q, dim, out);\n }\n\n //! Transpose a matrix\n static void MatrixTranspose(const ValueType *src, size_t dim, T *dst) {\n MatrixHelper::Transpose(\n src, (dim >> 3) / sizeof(ValueType), dst);\n }\n\n //! Reverse transpose a matrix\n static void MatrixReverseTranspose(const ValueType *src, size_t dim, T *dst) {\n MatrixHelper::ReverseTranspose(\n src, (dim >> 3) / sizeof(ValueType), dst);\n }\n\n //! Compute Norm2\n static void Norm2(ValueType * /*data*/, size_t /*dim*/, float *norm) {\n *norm = 0;\n }\n\n private:\n //! Members\n std::vector clusters_{};\n};\n\n/*! Numerical InnerProduct K-Means Context\n */\ntemplate \nclass NumericalInnerProductKmeansContext {\n public:\n //! constexpr variables\n constexpr static size_t BatchCount = BATCH_COUNT;\n\n //! Type of values\n using ValueType = typename std::remove_cv::type;\n using StoreType = typename std::remove_cv::type;\n\n // Check supporting type\n static_assert(IsSignedArithmetic::value,\n \"ValueType must be signed arithmetic\");\n\n /*! K-Means Context Cluster\n */\n class Cluster {\n public:\n //! Constructor\n Cluster(size_t dim) : accum_(dim, 0.0) {}\n\n //! Constructor\n Cluster(const Cluster &rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(rhs.accum_) {}\n\n //! Constructor\n Cluster(Cluster &&rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(std::move(rhs.accum_)) {}\n\n //! Assignment\n Cluster &operator=(const Cluster &rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = rhs.accum_;\n return *this;\n }\n\n //! Assignment\n Cluster &operator=(Cluster &&rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = std::move(rhs.accum_);\n return *this;\n }\n\n //! Append a vector\n void append(const ValueType *vec, size_t dim, float dist) {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n mutex_.lock();\n cost_ += dist;\n count_ += 1;\n\n for (size_t i = 0; i != dim; ++i) {\n accum_[i] += vec[i];\n }\n mutex_.unlock();\n }\n\n //! Retrieve the centroid of vectors\n void centroid(ValueType *out, size_t dim) const {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n for (size_t i = 0; i != dim; ++i) {\n out[i] = count_ == 0 ? FloatCast(NAN)\n : FloatCast(accum_[i] / count_);\n }\n }\n\n //! Retrieve squared error\n double cost(void) const {\n return cost_;\n }\n\n //! Retrieve feature count\n size_t count(void) const {\n return count_;\n }\n\n protected:\n //! Convert float type to another type\n template \n static auto FloatCast(const double val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(val);\n }\n\n //! Convert float type to another type\n template \n static auto FloatCast(const double val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(std::round(val));\n }\n\n private:\n SpinMutex mutex_{};\n double cost_{0.0};\n size_t count_{0u};\n std::vector accum_{};\n };\n\n //! operator []\n const Cluster &operator[](size_t i) const {\n return clusters_[i];\n }\n\n //! operator []\n Cluster &operator[](size_t i) {\n return clusters_[i];\n }\n\n //! Clear the context\n void clear(void) {\n clusters_.clear();\n }\n\n //! Reset the context\n void reset(size_t k_value, size_t dim) {\n clusters_.clear();\n clusters_.resize(k_value, dim);\n }\n\n //! Retrieve context of clusters\n const std::vector &clusters(void) const {\n return clusters_;\n }\n\n //! Compute the distance between matrix and query (batch)\n template \n static void BatchDistance(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n MinusInnerProductMatrix::Compute(m, q, dim, out);\n }\n\n //! Compute the distance between matrix and query (single)\n static void Distance(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n MinusInnerProductMatrix::Compute(m, q, dim, out);\n }\n\n //! Transpose a matrix\n template \n static auto MatrixTranspose(const U *src, size_t dim, T *dst) ->\n typename std::enable_if= 2>::type {\n MatrixHelper::Transpose(src, dim, dst);\n }\n\n //! Transpose a matrix\n template \n static auto MatrixTranspose(const U *src, size_t dim, U *dst) ->\n typename std::enable_if::type {\n MatrixHelper::Transpose(src, dim >> 2, dst);\n }\n\n //! Reverse transpose a matrix\n template \n static auto MatrixReverseTranspose(const U *src, size_t dim, U *dst) ->\n typename std::enable_if= 2>::type {\n MatrixHelper::ReverseTranspose(src, dim, dst);\n }\n\n //! Reverse transpose a matrix\n template \n static auto MatrixReverseTranspose(const U *src, size_t dim, U *dst) ->\n typename std::enable_if::type {\n MatrixHelper::ReverseTranspose(src, dim >> 2, dst);\n }\n\n //! Compute Norm2\n template ::value>::type>\n static void Norm2(ValueType *data, size_t dim, float *norm) {\n Normalizer::L2(data, dim, norm);\n }\n\n //! Compute Norm2, for non-float do nothing\n static void Norm2(ValueType * /*data*/, size_t /*dim*/, float *norm) {\n *norm = 0.0f;\n }\n\n private:\n //! Members\n std::vector clusters_{};\n};\n\n/*! Nibble InnerProduct K-Means Context (INT4)\n */\ntemplate \nclass NibbleInnerProductKmeansContext {\n public:\n //! constexpr variables\n constexpr static size_t BatchCount = BATCH_COUNT;\n\n //! Type of values\n using ValueType = typename std::remove_cv::type;\n using StoreType = typename std::make_unsigned::type;\n\n // Check supporting type\n static_assert(std::is_same::value ||\n std::is_same::value,\n \"ValueType must be int32_t or int64_t\");\n\n /*! K-Means Context Cluster\n */\n class Cluster {\n public:\n //! Constructor\n Cluster(size_t dim) : accum_(dim, 0.0) {}\n\n //! Constructor\n Cluster(const Cluster &rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(rhs.accum_) {}\n\n //! Constructor\n Cluster(Cluster &&rhs)\n : cost_(rhs.cost_), count_(rhs.count_), accum_(std::move(rhs.accum_)) {}\n\n //! Assignment\n Cluster &operator=(const Cluster &rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = rhs.accum_;\n return *this;\n }\n\n //! Assignment\n Cluster &operator=(Cluster &&rhs) {\n cost_ = rhs.cost_;\n count_ = rhs.count_;\n accum_ = std::move(rhs.accum_);\n return *this;\n }\n\n //! Append a vector\n void append(const StoreType *vec, size_t dim, float dist) {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n mutex_.lock();\n cost_ += dist;\n count_ += 1;\n\n const uint8_t *arr = reinterpret_cast(vec);\n dim = (dim >> 1) << 1;\n for (size_t i = 0; i != dim; i += 2) {\n uint8_t val = arr[i >> 1];\n accum_[i] += ((int8_t)(val << 4) >> 4);\n accum_[i + 1] += ((int8_t)(val) >> 4);\n }\n mutex_.unlock();\n }\n\n //! Retrieve the centroid of vectors\n void centroid(StoreType *out, size_t dim) const {\n ailego_check_with(dim == accum_.size(), \"Unmatched dimension\");\n\n uint8_t *arr = reinterpret_cast(out);\n dim = (dim >> 1) << 1;\n for (size_t i = 0; i != dim; i += 2) {\n int lo =\n count_ == 0 ? 0 : static_cast(std::round(accum_[i] / count_));\n int hi = count_ == 0\n ? 0\n : static_cast(std::round(accum_[i + 1] / count_));\n arr[i >> 1] = (uint8_t)((hi << 4) & 0xf0) | (uint8_t)(lo & 0xf);\n }\n }\n\n //! Retrieve squared error\n double cost(void) const {\n return cost_;\n }\n\n //! Retrieve feature count\n size_t count(void) const {\n return count_;\n }\n\n private:\n SpinMutex mutex_{};\n double cost_{0.0};\n size_t count_{0u};\n std::vector accum_{};\n };\n\n //! operator []\n const Cluster &operator[](size_t i) const {\n return clusters_[i];\n }\n\n //! operator []\n Cluster &operator[](size_t i) {\n return clusters_[i];\n }\n\n //! Clear the context\n void clear(void) {\n clusters_.clear();\n }\n\n //! Reset the context\n void reset(size_t k_value, size_t dim) {\n clusters_.clear();\n clusters_.resize(k_value, dim);\n }\n\n //! Retrieve context of clusters\n const std::vector &clusters(void) const {\n return clusters_;\n }\n\n //! Compute the distance between matrix and query (batch)\n template \n static void BatchDistance(const StoreType *m, const StoreType *q, size_t dim,\n float *out) {\n MinusInnerProductMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, out);\n }\n\n //! Compute the distance between matrix and query (single)\n static void Distance(const StoreType *m, const StoreType *q, size_t dim,\n float *out) {\n MinusInnerProductMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, out);\n }\n\n //! Transpose a matrix\n static void MatrixTranspose(const StoreType *src, size_t dim,\n StoreType *dst) {\n MatrixHelper::Transpose(src, dim >> 3, dst);\n }\n\n //! Reverse transpose a matrix\n static void MatrixReverseTranspose(const StoreType *src, size_t dim,\n StoreType *dst) {\n MatrixHelper::ReverseTranspose(src, dim >> 3, dst);\n }\n\n //! Compute Norm2\n static void Norm2(StoreType * /*data*/, size_t /*dim*/, float *norm) {\n *norm = 0;\n }\n\n private:\n //! Members\n std::vector clusters_{};\n};\n\n/*! Numerical K-Means cluster algorithm\n */\ntemplate >\nusing NumericalKmeans =\n LloydCluster>;\n\n/*! Nibble K-Means cluster algorithm\n */\ntemplate >\nusing NibbleKmeans = LloydCluster>;\n\n/*! Binary K-Means cluster algorithm\n */\ntemplate >\nusing BinaryKmeans = LloydCluster>;\n\n/*! Numerical K-Means cluster algorithm\n */\ntemplate >\nusing NumericalInnerProductKmeans =\n LloydCluster>;\n\n/*! Nibble K-Means cluster algorithm\n */\ntemplate >\nusing NibbleInnerProductKmeans =\n LloydCluster>;\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/algorithm/lloyd_cluster.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Random Centroids Generator\n */\ntemplate \nstruct RandomCentroidsGenerator {\n //! Type of values\n using OwnerType = typename std::decay::type;\n using ContainerType = typename OwnerType::ContainerType;\n using ContextType = typename OwnerType::ContextType;\n using ThreadPoolType = TPool;\n\n //! constexpr variables\n constexpr static size_t BatchCount = OwnerType::BatchCount;\n\n //! Generate centroids\n void operator()(OwnerType *owner, ThreadPoolType &) const {\n const auto &matrix = owner->feature_matrix();\n const auto &cache = owner->feature_cache();\n auto *centroids = owner->mutable_centroids();\n\n ContainerType rows(cache.dimension());\n size_t m = matrix.count();\n size_t n = m + cache.count();\n size_t k = owner->k_value();\n std::mt19937 mt((std::random_device())());\n\n rows.resize(BatchCount);\n centroids->reset(cache.dimension());\n centroids->reserve(k);\n\n for (size_t i = 0; k > 0 && i < n; ++i) {\n if (mt() % (n - i) >= k) {\n continue;\n }\n // Selected a feature\n if (i < m) {\n ContextType::MatrixReverseTranspose(matrix[i / BatchCount * BatchCount],\n matrix.dimension(), rows.data());\n centroids->append(rows[i & (BatchCount - 1u)], matrix.dimension());\n } else {\n centroids->append(cache[i - m], cache.dimension());\n }\n --k;\n }\n }\n};\n\n/*! Lloyd's algorithm cluster\n */\ntemplate \nclass LloydCluster {\n public:\n //! constexpr variables\n constexpr static size_t BatchCount = TContext::BatchCount;\n\n //! Type of values\n using ThreadPoolType = TPool;\n using ContainerType = TContainer;\n using ContextType = TContext;\n using ValueType = typename TContext::ValueType;\n using StoreType = typename TContext::StoreType;\n\n //! Constructor\n LloydCluster(size_t k, size_t dim)\n : k_value_(k),\n feature_cache_(dim),\n feature_matrix_(dim),\n centroids_matrix_(dim),\n centroids_(dim) {}\n\n //! Constructor\n LloydCluster(size_t k, size_t dim, bool spherical)\n : k_value_(k),\n feature_cache_(dim),\n feature_matrix_(dim),\n centroids_matrix_(dim),\n centroids_(dim),\n spherical_{spherical} {}\n\n //! Constructor\n LloydCluster(void) {}\n\n //! Destructor\n ~LloydCluster(void) {}\n\n //! Append a feature\n void append(const StoreType *arr, size_t dim) {\n feature_cache_.append(arr, dim);\n\n if (feature_cache_.count() == BatchCount) {\n size_t pos = feature_matrix_.count();\n feature_matrix_.resize(pos + BatchCount);\n ContextType::MatrixTranspose(feature_cache_.data(), dim,\n feature_matrix_[pos]);\n feature_cache_.clear();\n }\n }\n\n //! Reset cluster\n void reset(size_t k, size_t dim) {\n k_value_ = k;\n feature_cache_.reset(dim);\n feature_matrix_.reset(dim);\n centroids_.reset(dim);\n centroids_matrix_.reset(dim);\n context_.clear();\n }\n\n //! Reset cluster\n void reset(size_t k, size_t dim, bool spherical) {\n k_value_ = k;\n feature_cache_.reset(dim);\n feature_matrix_.reset(dim);\n centroids_.reset(dim);\n centroids_matrix_.reset(dim);\n context_.clear();\n spherical_ = spherical;\n }\n\n //! Initialize centroids\n template >\n void init_centroids(ThreadPoolType &pool, const G &g = G()) {\n g(this, pool);\n }\n\n //! Cluster one time\n template \n bool cluster_once(ThreadPoolType &pool, double *cost) {\n if (centroids_.empty()) {\n RandomCentroidsGenerator g;\n this->init_centroids(pool, g);\n }\n if (centroids_.count() != k_value_) {\n return false;\n }\n context_.reset(centroids_.count(), centroids_.dimension());\n\n size_t count = centroids_.count() / BatchCount * BatchCount;\n centroids_matrix_.resize(count);\n for (size_t i = 0; i != count; i += BatchCount) {\n ContextType::MatrixTranspose(centroids_[i], centroids_.dimension(),\n centroids_matrix_[i]);\n }\n size_t remain = static_cast(centroids_.count() - count);\n if (remain > 0) {\n centroids_matrix_.append(centroids_[count], centroids_.dimension(),\n remain);\n }\n\n // Using thread pool\n auto group = pool.make_group();\n if (!feature_matrix_.empty()) {\n size_t n = feature_matrix_.count() / BatchCount;\n size_t c = std::max(n / pool.count() / 2u, 1u);\n size_t m = n / c * c;\n\n for (size_t i = 0; i != m; i += c) {\n group->submit(Closure::New(this, &LloydCluster::cluster_matrix_features,\n i, i + c));\n }\n for (size_t i = m; i != n; i += 1) {\n group->submit(Closure::New(this, &LloydCluster::cluster_matrix_features,\n i, i + 1));\n }\n }\n if (!feature_cache_.empty()) {\n group->submit(Closure::New(this, &LloydCluster::cluster_cache_features));\n }\n group->wait_finish();\n\n *cost = 0.0;\n for (size_t i = 0, n = centroids_.count(); i != n; ++i) {\n const auto &item = context_[i];\n item.centroid(centroids_[i], centroids_.dimension());\n *cost += item.cost();\n }\n\n if (spherical_) {\n for (size_t i = 0, n = centroids_.count(); i != n; ++i) {\n float norm;\n ContextType::Norm2(centroids_[i], centroids_.dimension(), &norm);\n }\n }\n\n return true;\n }\n\n //! Retrieve the controids\n ContainerType *mutable_centroids(void) {\n return ¢roids_;\n }\n\n //! Retrieve the controids\n const ContainerType ¢roids(void) const {\n return centroids_;\n }\n\n //! Retrieve the K value\n size_t k_value(void) const {\n return k_value_;\n }\n\n //! Retrieve context\n const ContextType &context(void) const {\n return context_;\n }\n\n //! Retrieve the feature cache\n const ContainerType &feature_cache(void) const {\n return feature_cache_;\n }\n\n //! Retrieve the feature matrix\n const ContainerType &feature_matrix(void) const {\n return feature_matrix_;\n }\n\n //! Reserve the feature matrix\n void feature_matrix_reserve(size_t count) {\n feature_matrix_.reserve(count);\n }\n\n protected:\n //! Cluster the cache features\n void cluster_cache_features(void) {\n std::array scores;\n\n for (size_t i = 0, n = feature_cache_.count(); i != n; ++i) {\n size_t count = centroids_matrix_.count() / BatchCount * BatchCount;\n const StoreType *feature = feature_cache_[i];\n float nearest_score = std::numeric_limits::max();\n size_t nearest_index = 0;\n\n for (size_t j = 0; j != count; j += BatchCount) {\n ContextType::template BatchDistance<1>(centroids_matrix_[j], feature,\n centroids_matrix_.dimension(),\n scores.data());\n\n for (size_t k = 0; k < BatchCount; ++k) {\n if (scores[k] < nearest_score) {\n nearest_score = scores[k];\n nearest_index = j + k;\n }\n }\n } // end of for\n\n for (size_t j = count, total = centroids_matrix_.count(); j != total;\n ++j) {\n ContextType::Distance(centroids_matrix_[j], feature,\n centroids_matrix_.dimension(), scores.data());\n\n if (scores[0] < nearest_score) {\n nearest_score = scores[0];\n nearest_index = j;\n }\n }\n context_[nearest_index].append(feature, feature_cache_.dimension(),\n nearest_score);\n } // end of for\n }\n\n //! Cluster the matrix features\n void cluster_matrix_features(size_t first, size_t last) {\n std::array scores;\n ContainerType rows(centroids_matrix_.dimension());\n\n auto comp = [](float i, float j) {\n if (std::isnan(i)) return false;\n if (std::isnan(j)) return true;\n\n return i < j;\n };\n\n std::array nearest_scores;\n std::array nearest_indexes;\n\n rows.resize(BatchCount);\n for (size_t i = first * BatchCount; i != last * BatchCount;\n i += BatchCount) {\n size_t count = centroids_matrix_.count() / BatchCount * BatchCount;\n const StoreType *block = feature_matrix_[i];\n\n std::fill(nearest_indexes.data(), nearest_indexes.data() + BatchCount, 0);\n std::fill(nearest_scores.data(), nearest_scores.data() + BatchCount,\n std::numeric_limits::max());\n\n for (size_t j = 0; j != count; j += BatchCount) {\n ContextType::template BatchDistance(\n centroids_matrix_[j], block, centroids_matrix_.dimension(),\n scores.data());\n\n for (size_t k = 0; k < BatchCount; ++k) {\n const float *start = &scores[k * BatchCount];\n const float *result =\n std::min_element(start, start + BatchCount, comp);\n if (*result < nearest_scores[k]) {\n nearest_scores[k] = *result;\n nearest_indexes[k] = j + (result - start);\n }\n }\n } // end of for\n\n for (size_t j = count, total = centroids_matrix_.count(); j != total;\n ++j) {\n ContextType::template BatchDistance<1>(block, centroids_matrix_[j],\n centroids_matrix_.dimension(),\n scores.data());\n\n for (size_t k = 0; k < BatchCount; ++k) {\n float score = scores[k];\n if (score < nearest_scores[k]) {\n nearest_scores[k] = score;\n nearest_indexes[k] = j;\n }\n }\n } // end of for\n\n ContextType::MatrixReverseTranspose(block, feature_matrix_.dimension(),\n rows.data());\n for (size_t k = 0; k < BatchCount; ++k) {\n context_[nearest_indexes[k]].append(\n rows[k], feature_matrix_.dimension(), nearest_scores[k]);\n }\n } // end of for\n }\n\n private:\n //! Members\n size_t k_value_{0u};\n ContainerType feature_cache_{};\n ContainerType feature_matrix_{};\n ContainerType centroids_matrix_{};\n ContainerType centroids_{};\n ContextType context_{};\n bool spherical_{false};\n};\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/buffer/buffer_manager.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#ifdef __clang__\n#pragma clang diagnostic push\n#pragma clang diagnostic ignored \"-Wunused-parameter\"\n#pragma clang diagnostic ignored \"-Wshadow\"\n#elif defined(__GNUC__) || defined(__GNUG__)\n#pragma GCC diagnostic push\n#pragma GCC diagnostic ignored \"-Wunused-parameter\"\n#pragma GCC diagnostic ignored \"-Wshadow\"\n#endif\n\n#include \n\n#ifdef __clang__\n#pragma clang diagnostic pop\n#elif defined(__GNUC__) || defined(__GNUG__)\n#pragma GCC diagnostic pop\n#endif\n\n\nnamespace zvec {\n\n\nnamespace ailego {\n\n\nnamespace {\n\n\nstruct IDHash {\n size_t operator()(const BufferID &buffer_id) const {\n size_t hash = std::hash{}(static_cast(buffer_id.type));\n hash = hash ^ (std::hash{}(buffer_id.file_id));\n if (buffer_id.type == BufferID::TYPE::PARQUET) {\n hash = hash * 31 + std::hash{}(buffer_id.parquet().column);\n hash = hash * 31 + std::hash{}(buffer_id.parquet().row_group);\n } else if (buffer_id.type == BufferID::TYPE::VECTOR) {\n hash = hash * 31 + std::hash{}(buffer_id.vector().offset);\n }\n return hash;\n }\n};\n\n\nstruct IDEqual {\n bool operator()(const BufferID &a, const BufferID &b) const {\n if (a.type != b.type) {\n return false;\n }\n if (a.file_name != b.file_name) {\n return false;\n }\n if (a.file_id != b.file_id) {\n return false;\n }\n if (a.mtime != b.mtime) {\n return false;\n }\n if (a.type == BufferID::TYPE::PARQUET) {\n return a.parquet().column == b.parquet().column &&\n a.parquet().row_group == b.parquet().row_group;\n } else if (a.type == BufferID::TYPE::VECTOR) {\n return a.vector().offset == b.vector().offset;\n } else {\n return false;\n }\n }\n};\n\n\n} // namespace\n\n\nstruct BufferManager::BufferContext {\n BufferContext(const BufferID &id, BufferPool *p) : id(id), pool(p) {};\n BufferContext(const BufferContext &) = delete;\n BufferContext(BufferContext &&) = delete;\n BufferContext &operator=(const BufferContext &) = delete;\n BufferContext &operator=(BufferContext &&) = delete;\n\n\n ~BufferContext() {\n if (vector) {\n free(vector);\n }\n }\n\n\n typedef std::unique_ptr Pointer;\n\n\n enum State : uint32_t {\n IDLE = 0, // Empty and not held by any users, not in LRU\n RESERVED = 1, // Pinned by a user but no data yet, not in LRU\n IN_USE = 2, // Pinned by a user and data is present, not in LRU\n CACHED = 3, // Data is present but not held by any users, in LRU\n ERROR = 4 // Something went wrong, not in LRU\n };\n\n\n // Identifier for the buffer\n BufferID id;\n\n // Current state\n State state{IDLE};\n\n // The size of the buffer\n uint32_t size{0};\n\n // Handle of the file backing this buffer\n File file;\n\n // The number of external references to this buffer (via pin/unpin)\n std::atomic refs_buf{0};\n\n // The number of external references to this context (via BufferHandle)\n std::atomic refs_context{0};\n\n BufferPool *pool{nullptr};\n\n // A shared pointer to the buffers allocated for arrow parquet data\n std::shared_ptr arrow{nullptr};\n\n // Guard original arrow buffers to prevent premature deletion\n std::vector> arrow_refs{};\n\n // A pointer to the buffer allocated for vector data\n void *vector{nullptr};\n\n // Doubly linked LRU list\n BufferContext *next{nullptr};\n BufferContext *prev{nullptr};\n\n\n // Return a string representation of the status\n const std::string status_string() const;\n\n // Populate the buffer with parquet data\n arrow::Status read_arrow_parquet();\n\n // Populate the buffer with vector data\n bool read_vector();\n};\n\n\nconst std::string BufferManager::BufferContext::status_string() const {\n std::string msg{id.to_string() + \": \"};\n switch (state) {\n case State::IDLE: {\n msg += \"Idle\";\n break;\n }\n case State::RESERVED: {\n msg += \"Reserved\";\n break;\n }\n case State::IN_USE: {\n msg += \"In use\";\n break;\n }\n case State::CACHED: {\n msg += \"Cached\";\n break;\n }\n case State::ERROR: {\n msg += \"Error\";\n break;\n }\n }\n return msg;\n}\n\n\narrow::Status BufferManager::BufferContext::read_arrow_parquet() {\n // TODO: file handler and memory pool can be optimized\n arrow::MemoryPool *mem_pool = arrow::default_memory_pool();\n\n // Open file\n std::shared_ptr input;\n const auto &file_name = id.file_name;\n ARROW_ASSIGN_OR_RAISE(input, arrow::io::ReadableFile::Open(file_name));\n\n // Open reader\n std::unique_ptr reader;\n ARROW_ASSIGN_OR_RAISE(reader, parquet::arrow::OpenFile(input, mem_pool));\n\n // Perform read\n int row_group = id.parquet().row_group;\n int column = id.parquet().column;\n auto s = reader->RowGroup(row_group)->Column(column)->Read(&arrow);\n if (!s.ok()) {\n LOG_ERROR(\"Failed to read parquet file[%s]\", file_name.c_str());\n arrow = nullptr;\n return s;\n }\n\n // Compute the memory usage and hijack Arrow's buffers with our implementation\n for (auto &array : arrow->chunks()) {\n auto &buffers = array->data()->buffers;\n for (size_t buf_idx = 0; buf_idx < buffers.size(); ++buf_idx) {\n if (buffers[buf_idx] == nullptr) {\n continue;\n }\n // Keep references to original buffers to prevent premature deletion\n arrow_refs.emplace_back(buffers[buf_idx]);\n size += buffers[buf_idx]->capacity();\n // Create hijacked buffer with custom deleter that notifies us when Arrow\n // is finished with the buffer\n std::shared_ptr hijacked_buffer(\n buffers[buf_idx].get(), BufferManager::ArrowBufferDeleter(this));\n buffers[buf_idx] = hijacked_buffer;\n }\n }\n\n return arrow::Status::OK();\n}\n\n\nbool BufferManager::BufferContext::read_vector() {\n const auto &file_name = id.file_name;\n if (!file.is_valid()) {\n if (!File::IsExist(file_name)) {\n LOG_ERROR(\"File[%s] does not exist\", file_name.c_str());\n return false;\n }\n if (!File::IsRegular(file_name)) {\n LOG_ERROR(\"[%s] is not a regular file\", file_name.c_str());\n return false;\n }\n if (!file.open(file_name.c_str(), true, false)) {\n LOG_ERROR(\"Failed to open file[%s]\", file_name.c_str());\n return false;\n }\n }\n AILEGO_DEFER([this] { file.close(); });\n uint32_t len = id.vector().length;\n auto ret = posix_memalign((void **)&vector, 64, len); // 64-byte alignment\n if (ret != 0 || vector == nullptr) {\n LOG_ERROR(\"Failed to allocate buffer for file[%s]\", file_name.c_str());\n return false;\n }\n uint32_t offset = id.vector().offset;\n if (file.read(offset, vector, len) != len) {\n LOG_ERROR(\"Failed to read file[%s]\", file_name.c_str());\n free(vector);\n vector = nullptr;\n return false;\n }\n size = len;\n return true;\n}\n\n\n// Thread-safe buffer pool implementation.\n//\n// BufferContext states:\n// 1. Must exist in the lookup (hash) table.\n// 2. LRU list presence:\n// - In LRU: holds memory but not pinned by any users\n// - Not in LRU: either holds memory pinned by users, or doesn't hold memory\n// 3. External references: when an external user acquires a context and pins the\n// memory, that context is removed from LRU list; when they unpins the\n// memory, that context is moved to LRU list if it was the last reference.\n//\n// Any operation on the hash table is protected by mutex_table_.\n// Any change to context state and LRU list is protected by mutex_context_.\n//\nclass BufferManager::BufferPool {\n public:\n explicit BufferPool(uint64_t limit) : limit_(limit) {\n sentinel_.next = &sentinel_;\n sentinel_.prev = &sentinel_;\n }\n\n\n BufferContext *acquire_locked(BufferID &id) {\n std::lock_guard lock(mutex_context_);\n if (auto iter = table_.find(id); iter != table_.end()) {\n return iter->second.get();\n }\n auto [iter, _] =\n table_.emplace(id, std::make_unique(id, this));\n return iter->second.get();\n }\n\n\n void try_release_context_locked(BufferContext *context) {\n if (context->refs_context.load() != 0) {\n return;\n }\n std::lock_guard lock(mutex_table_);\n if (context->refs_context.load() != 0) {\n return;\n }\n if (context->state == BufferContext::State::IDLE) {\n table_.erase(context->id);\n }\n }\n\n\n void pin_locked(BufferContext *ctx) {\n std::lock_guard lock(mutex_context_);\n if (ctx->state == BufferContext::State::IDLE) {\n return pin_at_IDLE(ctx);\n }\n if (ctx->state == BufferContext::State::IN_USE) {\n return pin_at_IN_USE(ctx);\n }\n if (ctx->state == BufferContext::State::CACHED) {\n return pin_at_CACHED(ctx);\n }\n if (ctx->state == BufferContext::State::ERROR) {\n return;\n }\n }\n\n\n bool unpin_locked(BufferContext *ctx) {\n uint32_t prev_refs = ctx->refs_buf.fetch_sub(1);\n if (prev_refs > 1) {\n return false;\n }\n std::lock_guard lock(mutex_context_);\n if (ctx->refs_buf.load() == 0 &&\n ctx->state != BufferContext::State::CACHED) {\n ctx->state = BufferContext::State::CACHED;\n LRU_insert(ctx);\n return true;\n } else {\n return false;\n }\n }\n\n\n void LRU_insert_locked(BufferContext *context) {\n std::lock_guard lock(mutex_context_);\n LRU_insert(context);\n }\n\n\n void LRU_remove_locked(BufferContext *context) {\n std::lock_guard lock(mutex_context_);\n LRU_remove(context);\n }\n\n\n uint64_t usage() const {\n return usage_;\n }\n\n\n private:\n void pin_at_IDLE(BufferContext *ctx) {\n ctx->state = BufferContext::State::RESERVED;\n\n while (usage_ >= limit_) {\n // The tail of LRU list is the least recently used context\n BufferContext *victim = sentinel_.prev;\n if (victim == &sentinel_) { // No victim could be found\n ctx->state = BufferContext::State::ERROR;\n return;\n }\n if (victim->state == BufferContext::State::ERROR) {\n LRU_remove(victim);\n try_release_context_locked(ctx);\n continue;\n }\n if (victim->id.type == BufferID::TYPE::PARQUET) {\n victim->arrow_refs.clear();\n } else {\n free(victim->vector);\n victim->vector = nullptr;\n }\n victim->state = BufferContext::State::IDLE;\n LRU_remove(victim);\n try_release_context_locked(ctx);\n usage_ -= victim->size;\n }\n\n if (ctx->id.type == BufferID::TYPE::PARQUET) {\n if (ctx->read_arrow_parquet().ok()) {\n ctx->state = BufferContext::State::IN_USE;\n ctx->refs_buf.fetch_add(ctx->arrow_refs.size());\n usage_ += ctx->size;\n } else {\n LOG_ERROR(\"Failed to read to %s\", ctx->id.to_string().c_str());\n ctx->state = BufferContext::State::ERROR;\n }\n } else {\n if (ctx->read_vector()) {\n ctx->state = BufferContext::State::IN_USE;\n ctx->refs_buf.fetch_add(1);\n usage_ += ctx->size;\n } else {\n LOG_ERROR(\"Failed to read to %s\", ctx->id.to_string().c_str());\n ctx->state = BufferContext::State::ERROR;\n }\n }\n }\n\n\n void pin_at_IN_USE(BufferContext *ctx) {\n if (ctx->id.type == BufferID::TYPE::PARQUET) {\n ctx->refs_buf.fetch_add(ctx->arrow_refs.size());\n } else {\n ctx->refs_buf.fetch_add(1);\n }\n }\n\n\n void pin_at_CACHED(BufferContext *ctx) {\n if (ctx->id.type == BufferID::TYPE::PARQUET) {\n ctx->refs_buf.fetch_add(ctx->arrow_refs.size());\n } else {\n ctx->refs_buf.fetch_add(1);\n }\n LRU_remove(ctx);\n ctx->state = BufferContext::State::IN_USE;\n }\n\n\n void LRU_insert(BufferContext *context) {\n if (context->refs_buf > 0) {\n return; // Already pinned, should not be evicted\n }\n if (context->next != nullptr || context->prev != nullptr) {\n return;\n }\n // Insert the context to the head of LRU list\n context->next = sentinel_.next;\n context->prev = &sentinel_;\n sentinel_.next = context;\n context->next->prev = context;\n inactive_ += context->size;\n }\n\n\n void LRU_remove(BufferContext *context) {\n if (context->next == nullptr) {\n return; // Not in LRU list\n }\n context->next->prev = context->prev;\n context->prev->next = context->next;\n context->next = nullptr;\n context->prev = nullptr;\n inactive_ -= context->size;\n }\n\n private:\n using Table =\n std::unordered_map;\n\n uint64_t limit_;\n std::atomic usage_{0};\n std::atomic inactive_{0};\n\n Table table_{};\n std::mutex mutex_table_{};\n BufferContext sentinel_{BufferID{}, this}; // LRU list sentinel\n std::mutex mutex_context_{};\n};\n\n\nBufferManager::ArrowBufferDeleter::ArrowBufferDeleter(BufferContext *c)\n : context(c) {}\n\n\nvoid BufferManager::ArrowBufferDeleter::operator()(arrow::Buffer *) {\n context->pool->unpin_locked(context);\n}\n\n\nBufferHandle::BufferHandle(BufferContext *context) : context_(context) {\n if (context_ != nullptr) {\n pool_ = context_->pool;\n context_->refs_context.fetch_add(1);\n }\n}\n\n\nBufferHandle::~BufferHandle() {\n if (context_ != nullptr) {\n uint32_t prev_refs = context_->refs_context.fetch_sub(1);\n if (prev_refs > 1) {\n return;\n }\n if (context_->state == BufferContext::State::IDLE) {\n pool_->try_release_context_locked(context_);\n }\n }\n}\n\n\nstd::shared_ptr BufferHandle::pin_parquet_data() {\n pool_->pin_locked(context_);\n return context_->arrow;\n}\n\n\nvoid *BufferHandle::pin_vector_data() {\n if (!context_) {\n return nullptr;\n }\n pool_->pin_locked(context_);\n return context_->vector;\n}\n\n\nbool BufferHandle::unpin_vector_data() {\n if (!context_) {\n return true;\n }\n return pool_->unpin_locked(context_);\n}\n\n\nuint32_t BufferHandle::references() const {\n return context_->refs_buf.load();\n}\n\n\nuint32_t BufferHandle::size() const {\n return context_->size;\n}\n\n\nvoid BufferManager::init(uint64_t limit, uint32_t num_shards) {\n pools_.clear();\n uint64_t limit_per_shard = ailego_align(limit / num_shards, 4096);\n for (uint32_t i = 0; i < num_shards; ++i) {\n auto pool = new BufferPool(limit_per_shard);\n pools_.push_back(pool);\n }\n LOG_INFO(\n \"BufferManager initialized with [%u] buffer pools, [%zu] bytes memory \"\n \"limit per pool, total memory limit [%zu] bytes\",\n num_shards, (size_t)limit_per_shard, (size_t)limit);\n}\n\n\nBufferHandle BufferManager::acquire(BufferID &buffer_id) {\n static IDHash id_hash{};\n auto hash_val = id_hash(buffer_id);\n auto ctx = pools_[hash_val % pools_.size()]->acquire_locked(buffer_id);\n return BufferHandle(ctx);\n}\n\n\nstd::unique_ptr BufferManager::acquire_ptr(BufferID &buffer_id) {\n static IDHash id_hash{};\n auto hash_val = id_hash(buffer_id);\n auto ctx = pools_[hash_val % pools_.size()]->acquire_locked(buffer_id);\n return std::make_unique(ctx);\n}\n\n\nuint64_t BufferManager::total_size_in_bytes() const {\n uint64_t total_usage = 0;\n for (auto pool : pools_) {\n total_usage += pool->usage();\n }\n return total_usage;\n}\n\n\nBufferManager::~BufferManager() {\n for (auto pool : pools_) {\n delete pool;\n }\n}\n\n\n} // namespace ailego\n\n\n} // namespace zvec" }, { "path": "src/ailego/buffer/buffer_pool.cc", "content": "#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\nint LRUCache::init(size_t block_size) {\n block_size_ = block_size;\n for (size_t i = 0; i < CATCH_QUEUE_NUM; i++) {\n queues_.push_back(ConcurrentQueue(block_size));\n }\n return 0;\n}\n\nbool LRUCache::evict_single_block(BlockType &item) {\n bool found = false;\n for (size_t i = 0; i < CATCH_QUEUE_NUM; i++) {\n found = queues_[i].try_dequeue(item);\n if (found) {\n break;\n }\n }\n return found;\n}\n\nbool LRUCache::add_single_block(const LPMap *lp_map, const BlockType &block,\n int block_type) {\n bool ok = queues_[block_type].enqueue(block);\n if (!ok) {\n LOG_ERROR(\"enqueue failed.\");\n return false;\n }\n evict_queue_insertions_.fetch_add(1, std::memory_order_relaxed);\n if (evict_queue_insertions_ % block_size_ == 0) {\n this->clear_dead_node(lp_map);\n }\n return true;\n}\n\nvoid LRUCache::clear_dead_node(const LPMap *lp_map) {\n for (size_t i = 0; i < CATCH_QUEUE_NUM; i++) {\n size_t clear_size = block_size_ * 2;\n if (queues_[i].size_approx() < clear_size * 4) {\n continue;\n }\n size_t clear_count = 0;\n ConcurrentQueue tmp(block_size_);\n BlockType item;\n while (queues_[i].try_dequeue(item) && (clear_count++ < clear_size)) {\n if (!lp_map->isDeadBlock(item)) {\n if (!tmp.enqueue(item)) {\n LOG_ERROR(\"enqueue failed.\");\n }\n }\n }\n while (tmp.try_dequeue(item)) {\n if (!lp_map->isDeadBlock(item)) {\n if (!queues_[i].enqueue(item)) {\n LOG_ERROR(\"enqueue failed.\");\n }\n }\n }\n }\n}\n\nvoid LPMap::init(size_t entry_num) {\n if (entries_) {\n delete[] entries_;\n }\n entry_num_ = entry_num;\n entries_ = new Entry[entry_num_];\n for (size_t i = 0; i < entry_num_; i++) {\n entries_[i].ref_count.store(std::numeric_limits::min());\n entries_[i].load_count.store(0);\n entries_[i].buffer = nullptr;\n }\n cache_.init(entry_num * 4);\n}\n\nchar *LPMap::acquire_block(block_id_t block_id, bool lru_mode) {\n assert(block_id < entry_num_);\n Entry &entry = entries_[block_id];\n if (!lru_mode) {\n return entry.buffer;\n }\n while (true) {\n int current_count = entry.ref_count.load(std::memory_order_acquire);\n if (current_count < 0) {\n return nullptr;\n }\n if (entry.ref_count.compare_exchange_weak(current_count, current_count + 1,\n std::memory_order_acq_rel,\n std::memory_order_acquire)) {\n if (current_count == 0) {\n entry.load_count.fetch_add(1, std::memory_order_relaxed);\n }\n return entry.buffer;\n }\n }\n}\n\nvoid LPMap::release_block(block_id_t block_id) {\n assert(block_id < entry_num_);\n Entry &entry = entries_[block_id];\n\n if (entry.ref_count.fetch_sub(1, std::memory_order_release) == 1) {\n std::atomic_thread_fence(std::memory_order_acquire);\n LRUCache::BlockType block;\n block.first = block_id;\n block.second = entry.load_count.load();\n cache_.add_single_block(this, block, 0);\n }\n}\n\nchar *LPMap::evict_block(block_id_t block_id) {\n assert(block_id < entry_num_);\n Entry &entry = entries_[block_id];\n int expected = 0;\n if (entry.ref_count.compare_exchange_strong(\n expected, std::numeric_limits::min())) {\n char *buffer = entry.buffer;\n entry.buffer = nullptr;\n return buffer;\n } else {\n return nullptr;\n }\n}\n\nchar *LPMap::set_block_acquired(block_id_t block_id, char *buffer) {\n assert(block_id < entry_num_);\n Entry &entry = entries_[block_id];\n while (true) {\n int current_count = entry.ref_count.load(std::memory_order_relaxed);\n if (current_count >= 0) {\n if (entry.ref_count.compare_exchange_weak(\n current_count, current_count + 1, std::memory_order_acq_rel,\n std::memory_order_acquire)) {\n return entry.buffer;\n }\n } else {\n if (entry.ref_count.compare_exchange_weak(current_count, 1,\n std::memory_order_acq_rel,\n std::memory_order_acquire)) {\n entry.buffer = buffer;\n entry.load_count.fetch_add(1, std::memory_order_relaxed);\n return entry.buffer;\n }\n }\n }\n}\n\nvoid LPMap::recycle(moodycamel::ConcurrentQueue &free_buffers) {\n LRUCache::BlockType block;\n do {\n bool ok = cache_.evict_single_block(block);\n if (!ok) {\n return;\n }\n } while (isDeadBlock(block));\n char *buffer = evict_block(block.first);\n if (buffer) {\n if (!free_buffers.enqueue(buffer)) {\n LOG_ERROR(\"recycle buffer enqueue failed.\");\n ailego_free(buffer);\n }\n }\n}\n\nVecBufferPool::VecBufferPool(const std::string &filename) {\n fd_ = open(filename.c_str(), O_RDONLY);\n if (fd_ < 0) {\n throw std::runtime_error(\"Failed to open file: \" + filename);\n }\n struct stat st;\n if (fstat(fd_, &st) < 0) {\n ::close(fd_);\n throw std::runtime_error(\"Failed to stat file: \" + filename);\n }\n file_size_ = st.st_size;\n}\n\nint VecBufferPool::init(size_t pool_capacity, size_t block_size,\n size_t segment_count) {\n if (block_size == 0) {\n LOG_ERROR(\"block_size must not be 0\");\n return -1;\n }\n pool_capacity_ = pool_capacity;\n size_t buffer_num = pool_capacity_ / block_size + 10;\n size_t block_num = segment_count + 10;\n lp_map_.init(block_num);\n mutex_vec_.reserve(block_num);\n for (int i = 0; i < block_num; i++) {\n mutex_vec_.emplace_back(std::make_unique());\n }\n for (size_t i = 0; i < buffer_num; i++) {\n char *buffer = (char *)ailego_malloc(block_size);\n if (buffer != nullptr) {\n if (!free_buffers_.enqueue(buffer)) {\n LOG_ERROR(\"recycle buffer enqueue failed.\");\n ailego_free(buffer);\n return -1;\n }\n } else {\n LOG_ERROR(\"aligned_alloc %zu(size: %zu) failed\", i, block_size);\n return -1;\n }\n }\n LOG_DEBUG(\"Buffer pool num: %zu, entry num: %zu\", buffer_num,\n lp_map_.entry_num());\n no_lru_mode_ = false;\n if (lp_map_.entry_num() <= buffer_num) {\n no_lru_mode_ = true;\n }\n return 0;\n}\n\nVecBufferPoolHandle VecBufferPool::get_handle() {\n return VecBufferPoolHandle(*this);\n}\n\nchar *VecBufferPool::acquire_buffer(block_id_t block_id, size_t offset,\n size_t size, int retry) {\n char *buffer = lp_map_.acquire_block(block_id, !no_lru_mode());\n if (buffer) {\n return buffer;\n }\n std::lock_guard lock(*mutex_vec_[block_id]);\n buffer = lp_map_.acquire_block(block_id, !no_lru_mode());\n if (buffer) {\n return buffer;\n }\n {\n bool found = free_buffers_.try_dequeue(buffer);\n if (!found && !no_lru_mode_) {\n for (int i = 0; i < retry; i++) {\n lp_map_.recycle(free_buffers_);\n found = free_buffers_.try_dequeue(buffer);\n if (found) {\n break;\n }\n }\n }\n if (!found) {\n LOG_ERROR(\"Buffer pool failed to get free buffer\");\n return nullptr;\n }\n }\n\n ssize_t read_bytes = pread(fd_, buffer, size, offset);\n if (read_bytes != static_cast(size)) {\n LOG_ERROR(\"Buffer pool failed to read file at offset: %zu\", offset);\n free_buffers_.enqueue(buffer);\n return nullptr;\n }\n return lp_map_.set_block_acquired(block_id, buffer);\n}\n\nint VecBufferPool::get_meta(size_t offset, size_t length, char *buffer) {\n ssize_t read_bytes = pread(fd_, buffer, length, offset);\n if (read_bytes != static_cast(length)) {\n LOG_ERROR(\"Buffer pool failed to read file at offset: %zu\", offset);\n return -1;\n }\n return 0;\n}\n\nchar *VecBufferPoolHandle::get_block(size_t offset, size_t size,\n size_t block_id) {\n char *buffer = pool_.acquire_buffer(block_id, offset, size, 5);\n return buffer;\n}\n\nint VecBufferPoolHandle::get_meta(size_t offset, size_t length, char *buffer) {\n return pool_.get_meta(offset, length, buffer);\n}\n\nvoid VecBufferPoolHandle::release_one(block_id_t block_id) {\n if (!pool_.no_lru_mode()) {\n pool_.lp_map_.release_block(block_id);\n }\n}\n\nvoid VecBufferPoolHandle::acquire_one(block_id_t block_id) {\n if (!pool_.no_lru_mode()) {\n pool_.lp_map_.acquire_block(block_id, true);\n }\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/container/bitmap.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"bitmap.h\"\n\nnamespace zvec {\nnamespace ailego {\n\nsize_t Bitset::BitwiseAndCardinality(const Bitset &lhs, const Bitset &rhs) {\n return BitsetHelper::BitwiseAndCardinality(\n lhs.array_.data(), rhs.array_.data(),\n std::min(lhs.array_.size(), rhs.array_.size()));\n}\n\nsize_t Bitset::BitwiseAndnotCardinality(const Bitset &lhs, const Bitset &rhs) {\n size_t lsize = lhs.array_.size();\n size_t rsize = rhs.array_.size();\n\n if (lsize > rsize) {\n return (\n BitsetHelper::BitwiseAndnotCardinality(lhs.array_.data(),\n rhs.array_.data(), rsize) +\n BitsetHelper::Cardinality(lhs.array_.data() + rsize, lsize - rsize));\n }\n return BitsetHelper::BitwiseAndnotCardinality(lhs.array_.data(),\n rhs.array_.data(), lsize);\n}\n\nsize_t Bitset::BitwiseXorCardinality(const Bitset &lhs, const Bitset &rhs) {\n size_t lsize = lhs.array_.size();\n size_t rsize = rhs.array_.size();\n\n if (lsize < rsize) {\n return (\n BitsetHelper::BitwiseXorCardinality(lhs.array_.data(),\n rhs.array_.data(), lsize) +\n BitsetHelper::Cardinality(rhs.array_.data() + lsize, rsize - lsize));\n } else if (lsize > rsize) {\n return (\n BitsetHelper::BitwiseXorCardinality(lhs.array_.data(),\n rhs.array_.data(), rsize) +\n BitsetHelper::Cardinality(lhs.array_.data() + rsize, lsize - rsize));\n }\n return BitsetHelper::BitwiseXorCardinality(lhs.array_.data(),\n rhs.array_.data(), lsize);\n}\n\nsize_t Bitset::BitwiseOrCardinality(const Bitset &lhs, const Bitset &rhs) {\n size_t lsize = lhs.array_.size();\n size_t rsize = rhs.array_.size();\n\n if (lsize < rsize) {\n return (\n BitsetHelper::BitwiseOrCardinality(lhs.array_.data(), rhs.array_.data(),\n lsize) +\n BitsetHelper::Cardinality(rhs.array_.data() + lsize, rsize - lsize));\n } else if (lsize > rsize) {\n return (\n BitsetHelper::BitwiseOrCardinality(lhs.array_.data(), rhs.array_.data(),\n rsize) +\n BitsetHelper::Cardinality(lhs.array_.data() + rsize, lsize - rsize));\n }\n return BitsetHelper::BitwiseOrCardinality(lhs.array_.data(),\n rhs.array_.data(), lsize);\n}\n\nvoid Bitmap::clear(void) {\n for (std::vector::iterator iter = array_.begin();\n iter != array_.end(); ++iter) {\n delete (*iter);\n }\n array_.clear();\n}\n\nvoid Bitmap::copy(const Bitmap &rhs) {\n this->clear();\n\n for (std::vector::const_iterator iter = rhs.array_.begin();\n iter != rhs.array_.end(); ++iter) {\n Bucket *bucket = NULL;\n if (*iter) {\n bucket = new Bucket(*(*iter));\n }\n array_.push_back(bucket);\n }\n}\n\nvoid Bitmap::shrink_to_fit(void) {\n size_t shrink_count = 0;\n std::vector::reverse_iterator iter;\n\n for (iter = array_.rbegin(); iter != array_.rend(); ++iter) {\n if (*iter) {\n if (!(*iter)->test_none()) {\n break;\n }\n delete (*iter);\n *iter = NULL;\n }\n ++shrink_count;\n }\n for (; iter != array_.rend(); ++iter) {\n if ((*iter) && (*iter)->test_none()) {\n delete (*iter);\n *iter = NULL;\n }\n }\n if (shrink_count != 0) {\n array_.resize(array_.size() - shrink_count);\n }\n}\n\nbool Bitmap::test(size_t num) const {\n // High 16 bits\n size_t offset = num >> 16;\n\n if (offset < array_.size()) {\n const Bucket *bucket = array_[offset];\n if (bucket) {\n // Low 16 bits\n return bucket->test(static_cast(num));\n }\n }\n return false;\n}\n\nvoid Bitmap::set(size_t num) {\n // High 16 bits\n size_t offset = num >> 16;\n if (offset >= array_.size()) {\n array_.resize(offset + 1, NULL);\n }\n\n Bucket *&bucket = array_[offset];\n if (!bucket) {\n bucket = new Bucket;\n }\n // Low 16 bits\n bucket->set(static_cast(num));\n}\n\nvoid Bitmap::reset(size_t num) {\n // High 16 bits\n size_t offset = num >> 16;\n if (offset >= array_.size()) {\n array_.resize(offset + 1, NULL);\n }\n\n if (offset < array_.size()) {\n Bucket *bucket = array_[offset];\n if (bucket) {\n // Low 16 bits\n bucket->reset(static_cast(num));\n }\n }\n}\n\nvoid Bitmap::flip(size_t num) {\n // High 16 bits\n uint16_t offset = (uint16_t)(num >> 16);\n if (offset >= array_.size()) {\n array_.resize(offset + 1, NULL);\n }\n\n Bucket *&bucket = array_[offset];\n if (!bucket) {\n bucket = new Bucket;\n }\n // Low 16 bits\n bucket->flip(static_cast(num));\n}\n\nvoid Bitmap::bitwise_and(const Bitmap &rhs) {\n size_t overlap = std::min(array_.size(), rhs.array_.size());\n\n for (size_t i = 0; i < overlap; ++i) {\n Bucket *&dst = array_[i];\n\n if (dst) {\n const Bucket *src = rhs.array_[i];\n if (src) {\n dst->bitwise_and(*src);\n } else {\n delete dst;\n dst = NULL;\n }\n }\n }\n for (size_t i = overlap; i < array_.size(); ++i) {\n Bucket *&dst = array_[i];\n delete dst;\n dst = NULL;\n }\n}\n\nvoid Bitmap::bitwise_andnot(const Bitmap &rhs) {\n size_t overlap = std::min(array_.size(), rhs.array_.size());\n\n for (size_t i = 0; i < overlap; ++i) {\n Bucket *&dst = array_[i];\n\n if (dst) {\n const Bucket *src = rhs.array_[i];\n if (src) {\n dst->bitwise_andnot(*src);\n }\n }\n }\n}\n\nvoid Bitmap::bitwise_or(const Bitmap &rhs) {\n size_t overlap = std::min(array_.size(), rhs.array_.size());\n\n for (size_t i = 0; i < overlap; ++i) {\n const Bucket *src = rhs.array_[i];\n\n if (src) {\n Bucket *&dst = array_[i];\n\n if (dst) {\n dst->bitwise_or(*src);\n } else {\n dst = new Bucket(*src);\n }\n }\n }\n for (size_t i = overlap; i < rhs.array_.size(); ++i) {\n const Bucket *src = rhs.array_[i];\n Bucket *bucket = NULL;\n\n if (src) {\n bucket = new Bucket(*src);\n }\n array_.push_back(bucket);\n }\n}\n\nvoid Bitmap::bitwise_xor(const Bitmap &rhs) {\n size_t overlap = std::min(array_.size(), rhs.array_.size());\n\n for (size_t i = 0; i < overlap; ++i) {\n const Bucket *src = rhs.array_[i];\n\n if (src) {\n Bucket *&dst = array_[i];\n\n if (dst) {\n dst->bitwise_xor(*src);\n } else {\n dst = new Bucket(*src);\n }\n }\n }\n for (size_t i = overlap; i < rhs.array_.size(); ++i) {\n const Bucket *src = rhs.array_[i];\n Bucket *bucket = NULL;\n\n if (src) {\n bucket = new Bucket(*src);\n }\n array_.push_back(bucket);\n }\n}\n\nvoid Bitmap::bitwise_not(void) {\n for (std::vector::iterator iter = array_.begin();\n iter != array_.end(); ++iter) {\n Bucket *&bucket = *iter;\n if (!bucket) {\n bucket = new Bucket;\n }\n bucket->bitwise_not();\n }\n}\n\nbool Bitmap::test_all(void) const {\n if (array_.empty()) {\n return false;\n }\n for (std::vector::const_iterator iter = array_.begin();\n iter != array_.end(); ++iter) {\n if (!(*iter) || !(*iter)->test_all()) {\n return false;\n }\n }\n return true;\n}\n\nbool Bitmap::test_any(void) const {\n for (std::vector::const_iterator iter = array_.begin();\n iter != array_.end(); ++iter) {\n if (*iter && (*iter)->test_any()) {\n return true;\n }\n }\n return false;\n}\n\nbool Bitmap::test_none(void) const {\n for (std::vector::const_iterator iter = array_.begin();\n iter != array_.end(); ++iter) {\n if (*iter && !(*iter)->test_none()) {\n return false;\n }\n }\n return true;\n}\n\nsize_t Bitmap::cardinality(void) const {\n size_t result = 0;\n for (std::vector::const_iterator iter = array_.begin();\n iter != array_.end(); ++iter) {\n if (*iter) {\n result += (*iter)->cardinality();\n }\n }\n return result;\n}\n\nvoid Bitmap::extract(size_t base, std::vector *out) const {\n for (std::vector::const_iterator iter = array_.begin();\n iter != array_.end(); ++iter) {\n if (*iter) {\n (*iter)->extract(base, out);\n }\n base += Bucket::MAX_SIZE;\n }\n}\n\nsize_t Bitmap::BitwiseAndCardinality(const Bitmap &lhs, const Bitmap &rhs) {\n size_t overlap = std::min(lhs.array_.size(), rhs.array_.size());\n size_t dist = 0;\n\n for (size_t i = 0; i < overlap; ++i) {\n const Bucket *l = lhs.array_[i];\n const Bucket *r = rhs.array_[i];\n\n if (l && r) {\n dist += Bucket::BitwiseAndCardinality(*l, *r);\n }\n }\n return dist;\n}\n\nsize_t Bitmap::BitwiseAndnotCardinality(const Bitmap &lhs, const Bitmap &rhs) {\n size_t overlap = std::min(lhs.array_.size(), rhs.array_.size());\n size_t dist = 0;\n\n for (size_t i = 0; i < overlap; ++i) {\n const Bucket *l = lhs.array_[i];\n if (l) {\n const Bucket *r = rhs.array_[i];\n if (r) {\n dist += Bucket::BitwiseAndnotCardinality(*l, *r);\n } else {\n dist += l->cardinality();\n }\n }\n }\n for (size_t i = overlap; i < lhs.array_.size(); ++i) {\n const Bucket *l = lhs.array_[i];\n if (l) {\n dist += l->cardinality();\n }\n }\n return dist;\n}\n\nsize_t Bitmap::BitwiseXorCardinality(const Bitmap &lhs, const Bitmap &rhs) {\n size_t overlap = std::min(lhs.array_.size(), rhs.array_.size());\n size_t dist = 0;\n\n for (size_t i = 0; i < overlap; ++i) {\n const Bucket *l = lhs.array_[i];\n const Bucket *r = rhs.array_[i];\n\n if (l && r) {\n dist += Bucket::BitwiseXorCardinality(*l, *r);\n } else if (l) {\n dist += l->cardinality();\n } else if (r) {\n dist += r->cardinality();\n }\n }\n for (size_t i = overlap; i < lhs.array_.size(); ++i) {\n const Bucket *l = lhs.array_[i];\n if (l) {\n dist += l->cardinality();\n }\n }\n for (size_t i = overlap; i < rhs.array_.size(); ++i) {\n const Bucket *r = rhs.array_[i];\n if (r) {\n dist += r->cardinality();\n }\n }\n return dist;\n}\n\nsize_t Bitmap::BitwiseOrCardinality(const Bitmap &lhs, const Bitmap &rhs) {\n size_t overlap = std::min(lhs.array_.size(), rhs.array_.size());\n size_t dist = 0;\n\n for (size_t i = 0; i < overlap; ++i) {\n const Bucket *l = lhs.array_[i];\n const Bucket *r = rhs.array_[i];\n\n if (l && r) {\n dist += Bucket::BitwiseOrCardinality(*l, *r);\n } else if (l) {\n dist += l->cardinality();\n } else if (r) {\n dist += r->cardinality();\n }\n }\n for (size_t i = overlap; i < lhs.array_.size(); ++i) {\n const Bucket *l = lhs.array_[i];\n if (l) {\n dist += l->cardinality();\n }\n }\n for (size_t i = overlap; i < rhs.array_.size(); ++i) {\n const Bucket *r = rhs.array_[i];\n if (r) {\n dist += r->cardinality();\n }\n }\n return dist;\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/container/bitmap.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Fixed Bitset\n */\ntemplate ::type>\nclass FixedBitset {\n public:\n enum { MAX_SIZE = N };\n\n //! Constructor\n FixedBitset(void) {\n memset(array_, 0, sizeof(array_));\n }\n\n //! Constructor\n FixedBitset(const FixedBitset &rhs) {\n memcpy(array_, rhs.array_, sizeof(array_));\n }\n\n //! Destructor\n ~FixedBitset(void) {}\n\n //! Assignment\n FixedBitset &operator=(const FixedBitset &rhs) {\n memcpy(array_, rhs.array_, sizeof(array_));\n return *this;\n }\n\n //! Retrieve data pointer\n uint32_t *data(void) {\n return reinterpret_cast(array_);\n }\n\n //! Retrieve data pointer\n const uint32_t *data(void) const {\n return reinterpret_cast(array_);\n }\n\n //! Retrieve count of bits in set\n constexpr size_t size(void) const {\n return MAX_SIZE;\n }\n\n // ！Clear the bitset\n void clear(void) {\n memset(array_, 0, sizeof(array_));\n }\n\n //! Test a bit in bitset\n bool test(size_t num) const {\n ailego_assert_with(N > num, \"overflow argument\");\n return ((array_[num >> 5] & (1u << (num & 0x1f))) != 0);\n }\n\n //! Set a bit in bitset\n void set(size_t num) {\n ailego_assert_with(N > num, \"overflow argument\");\n uint32_t mask = (1u << (num & 0x1f));\n array_[num >> 5] |= mask;\n }\n\n //! Clear a bit in bitset\n void reset(size_t num) {\n ailego_assert_with(N > num, \"overflow argument\");\n uint32_t mask = (1u << (num & 0x1f));\n array_[num >> 5] &= ~mask;\n }\n\n //! Toggle a bit in bitset\n void flip(size_t num) {\n ailego_assert_with(N > num, \"overflow argument\");\n uint32_t mask = (1u << (num & 0x1f));\n array_[num >> 5] ^= mask;\n }\n\n //! Perform binary AND\n void bitwise_and(const FixedBitset &rhs) {\n BitsetHelper::BitwiseAnd(array_, rhs.array_, ((N + 0x1f) >> 5));\n }\n\n //! Perform binary AND NOT\n void bitwise_andnot(const FixedBitset &rhs) {\n BitsetHelper::BitwiseAndnot(array_, rhs.array_, ((N + 0x1f) >> 5));\n }\n\n //! Perform binary OR\n void bitwise_or(const FixedBitset &rhs) {\n BitsetHelper::BitwiseOr(array_, rhs.array_, ((N + 0x1f) >> 5));\n }\n\n //! Perform binary XOR\n void bitwise_xor(const FixedBitset &rhs) {\n BitsetHelper::BitwiseXor(array_, rhs.array_, ((N + 0x1f) >> 5));\n }\n\n //! Perform binary NOT\n void bitwise_not(void) {\n BitsetHelper::BitwiseNot(array_, ((N + 0x1f) >> 5));\n }\n\n //! Check if all bits are set to true\n bool test_all(void) const {\n return BitsetHelper::TestAll(array_, ((N + 0x1f) >> 5));\n }\n\n //! Check if any bits are set to true\n bool test_any(void) const {\n return BitsetHelper::TestAny(array_, ((N + 0x1f) >> 5));\n }\n\n //! Check if none of the bits are set to true\n bool test_none(void) const {\n return BitsetHelper::TestNone(array_, ((N + 0x1f) >> 5));\n }\n\n //! Compute the cardinality of a bitset\n size_t cardinality(void) const {\n return BitsetHelper::Cardinality(array_, ((N + 0x1f) >> 5));\n }\n\n //! Extract the bitset to an array\n void extract(size_t base, std::vector *out) const {\n const uint32_t *iter = array_;\n const uint32_t *last = array_ + ((N + 0x1f) >> 5);\n\n for (; iter != last; ++iter) {\n uint32_t w = *iter;\n\n while (w != 0) {\n uint32_t c = ailego_ctz32(w);\n w &= ~(1u << c);\n out->push_back(base + c);\n }\n base += 32u;\n }\n }\n\n //! Extract the bitset to an array\n void extract(std::vector *out) const {\n this->extract(0, out);\n }\n\n //! Compute the AND cardinality between two bitsets\n static size_t BitwiseAndCardinality(const FixedBitset &lhs,\n const FixedBitset &rhs) {\n return BitsetHelper::BitwiseAndCardinality(lhs.array_, rhs.array_,\n ((N + 0x1f) >> 5));\n }\n\n //! Compute the ANDNOT cardinality between two bitsets\n static size_t BitwiseAndnotCardinality(const FixedBitset &lhs,\n const FixedBitset &rhs) {\n return BitsetHelper::BitwiseAndnotCardinality(lhs.array_, rhs.array_,\n ((N + 0x1f) >> 5));\n }\n\n //! Compute the XOR cardinality between two bitsets\n static size_t BitwiseXorCardinality(const FixedBitset &lhs,\n const FixedBitset &rhs) {\n return BitsetHelper::BitwiseXorCardinality(lhs.array_, rhs.array_,\n ((N + 0x1f) >> 5));\n }\n\n //! Compute the OR cardinality between two bitsets\n static size_t BitwiseOrCardinality(const FixedBitset &lhs,\n const FixedBitset &rhs) {\n return BitsetHelper::BitwiseOrCardinality(lhs.array_, rhs.array_,\n ((N + 0x1f) >> 5));\n }\n\n //! Convert a array pointer to bitset pointer\n static FixedBitset *Cast(uint32_t *arr) {\n return reinterpret_cast *>(arr);\n }\n\n //! Convert a array pointer to bitset pointer\n static const FixedBitset *Cast(const uint32_t *arr) {\n return reinterpret_cast *>(arr);\n }\n\n //! Convert a array pointer to bitset pointer\n static FixedBitset *Cast(uint64_t *arr) {\n return reinterpret_cast *>(arr);\n }\n\n //! Convert a array pointer to bitset pointer\n static const FixedBitset *Cast(const uint64_t *arr) {\n return reinterpret_cast *>(arr);\n }\n\n private:\n uint32_t array_[(N + 0x1f) >> 5];\n};\n\n/*! Fixed Bitset (Special)\n */\ntemplate <>\nclass FixedBitset<0> {\n public:\n enum { MAX_SIZE = 0 };\n\n //! Retrieve max size of bitset\n constexpr size_t size(void) const {\n return MAX_SIZE;\n }\n};\n\n/*! Bitset\n */\nclass Bitset {\n public:\n //! Constructor\n Bitset(void) : array_() {}\n\n //! Constructor\n Bitset(size_t bits) : array_((bits + 0x1f) >> 5) {}\n\n //! Constructor\n Bitset(const Bitset &rhs) : array_(rhs.array_) {}\n\n //! Constructor\n Bitset(Bitset &&rhs) : array_(std::move(rhs.array_)) {}\n\n //! Destructor\n ~Bitset(void) {}\n\n //! Assignment\n Bitset &operator=(const Bitset &rhs) {\n array_ = rhs.array_;\n return *this;\n }\n\n //! Assignment\n Bitset &operator=(Bitset &&rhs) {\n array_ = std::move(rhs.array_);\n return *this;\n }\n\n //! Retrieve data pointer\n uint32_t *data(void) {\n return array_.data();\n }\n\n //! Retrieve data pointer\n const uint32_t *data(void) const {\n return array_.data();\n }\n\n //! Retrieve count of bits in set\n size_t size(void) const {\n return (array_.size() << 5);\n }\n\n //! Resize the bitset\n void resize(size_t bits) {\n array_.resize((bits + 0x1f) >> 5);\n }\n\n // ！Clear the bitset\n void clear(void) {\n array_.clear();\n }\n\n //! Test a bit in bitset\n bool test(size_t num) const {\n ailego_assert_with(this->size() > num, \"overflow argument\");\n return ((array_[num >> 5] & (1u << (num & 0x1f))) != 0);\n }\n\n //! Set a bit in bitset\n void set(size_t num) {\n ailego_assert_with(this->size() > num, \"overflow argument\");\n uint32_t mask = (1u << (num & 0x1f));\n array_[num >> 5] |= mask;\n }\n\n //! Clear a bit in bitset\n void reset(size_t num) {\n ailego_assert_with(this->size() > num, \"overflow argument\");\n uint32_t mask = (1u << (num & 0x1f));\n array_[num >> 5] &= ~mask;\n }\n\n //! Toggle a bit in bitset\n void flip(size_t num) {\n ailego_assert_with(this->size() > num, \"overflow argument\");\n uint32_t mask = (1u << (num & 0x1f));\n array_[num >> 5] ^= mask;\n }\n\n //! Perform binary AND\n void bitwise_and(const Bitset &rhs) {\n BitsetHelper::BitwiseAnd(array_.data(), rhs.array_.data(),\n std::min(array_.size(), rhs.array_.size()));\n }\n\n //! Perform binary AND NOT\n void bitwise_andnot(const Bitset &rhs) {\n BitsetHelper::BitwiseAndnot(array_.data(), rhs.array_.data(),\n std::min(array_.size(), rhs.array_.size()));\n }\n\n //! Perform binary OR\n void bitwise_or(const Bitset &rhs) {\n BitsetHelper::BitwiseOr(array_.data(), rhs.array_.data(),\n std::min(array_.size(), rhs.array_.size()));\n }\n\n //! Perform binary XOR\n void bitwise_xor(const Bitset &rhs) {\n BitsetHelper::BitwiseXor(array_.data(), rhs.array_.data(),\n std::min(array_.size(), rhs.array_.size()));\n }\n\n //! Perform binary NOT\n void bitwise_not(void) {\n BitsetHelper::BitwiseNot(array_.data(), array_.size());\n }\n\n //! Check if all bits are set to true\n bool test_all(void) const {\n return BitsetHelper::TestAll(array_.data(), array_.size());\n }\n\n //! Check if any bits are set to true\n bool test_any(void) const {\n return BitsetHelper::TestAny(array_.data(), array_.size());\n }\n\n //! Check if none of the bits are set to true\n bool test_none(void) const {\n return BitsetHelper::TestNone(array_.data(), array_.size());\n }\n\n //! Compute the cardinality of a bitset\n size_t cardinality(void) const {\n return BitsetHelper::Cardinality(array_.data(), array_.size());\n }\n\n //! Extract the bitset to an array\n void extract(size_t base, std::vector *out) const {\n const uint32_t *iter = array_.data();\n const uint32_t *last = array_.data() + array_.size();\n\n for (; iter != last; ++iter) {\n uint32_t w = *iter;\n\n while (w != 0) {\n uint32_t c = ailego_ctz32(w);\n w &= ~(1u << c);\n out->push_back(base + c);\n }\n base += 32u;\n }\n }\n\n //! Extract the bitset to an array\n void extract(std::vector *out) const {\n this->extract(0, out);\n }\n\n //! Compute the AND cardinality between two bitsets\n static size_t BitwiseAndCardinality(const Bitset &lhs, const Bitset &rhs);\n\n //! Compute the ANDNOT cardinality between two bitsets\n static size_t BitwiseAndnotCardinality(const Bitset &lhs, const Bitset &rhs);\n\n //! Compute the XOR cardinality between two bitsets\n static size_t BitwiseXorCardinality(const Bitset &lhs, const Bitset &rhs);\n\n //! Compute the OR cardinality between two bitsets\n static size_t BitwiseOrCardinality(const Bitset &lhs, const Bitset &rhs);\n\n private:\n std::vector array_;\n};\n\n/*! Bitmap\n */\nclass Bitmap {\n public:\n typedef FixedBitset<65536u> Bucket;\n\n //! Constructor\n Bitmap(void) : array_() {}\n\n //! Constructor\n Bitmap(const Bitmap &rhs) {\n this->copy(rhs);\n }\n\n //! Destructor\n ~Bitmap(void) {\n this->clear();\n }\n\n //! Assignment\n Bitmap &operator=(const Bitmap &rhs) {\n this->copy(rhs);\n return *this;\n }\n\n //! Retrieve bucket size of bitmap\n size_t bucket_size(void) const {\n return array_.size();\n }\n\n // ！Clear the bitmap\n void clear(void);\n\n //! Remove the none buckets\n void shrink_to_fit(void);\n\n //! Test a bit in bitmap\n bool test(size_t num) const;\n\n //! Set a bit in bitmap\n void set(size_t num);\n\n //! Reset a bit in bitmap\n void reset(size_t num);\n\n //! Toggle a bit in bitmap\n void flip(size_t num);\n\n //! Perform binary AND\n void bitwise_and(const Bitmap &rhs);\n\n //! Perform binary AND NOT\n void bitwise_andnot(const Bitmap &rhs);\n\n //! Perform binary OR\n void bitwise_or(const Bitmap &rhs);\n\n //! Perform binary XOR\n void bitwise_xor(const Bitmap &rhs);\n\n //! Perform binary NOT (It will expand the whole map)\n void bitwise_not(void);\n\n //! Check if all bits are set to true\n bool test_all(void) const;\n\n //! Check if any bits are set to true\n bool test_any(void) const;\n\n //! Check if none of the bits are set to true\n bool test_none(void) const;\n\n //! Compute the cardinality of a bitmap\n size_t cardinality(void) const;\n\n //! Extract the bitmap to an array\n void extract(size_t base, std::vector *out) const;\n\n //! Extract the bitmap to an array\n void extract(std::vector *out) const {\n this->extract(0, out);\n }\n\n //! Compute the AND cardinality between two bitmaps\n static size_t BitwiseAndCardinality(const Bitmap &lhs, const Bitmap &rhs);\n\n //! Compute the ANDNOT cardinality between two bitmaps\n static size_t BitwiseAndnotCardinality(const Bitmap &lhs, const Bitmap &rhs);\n\n //! Compute the XOR cardinality between two bitmaps\n static size_t BitwiseXorCardinality(const Bitmap &lhs, const Bitmap &rhs);\n\n //! Compute the OR cardinality between two bitmaps\n static size_t BitwiseOrCardinality(const Bitmap &lhs, const Bitmap &rhs);\n\n protected:\n //! Copy the content from another bitmap\n void copy(const Bitmap &rhs);\n\n private:\n std::vector array_;\n};\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/container/bloom_filter.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Bloom Filter Calculator\n */\nstruct BloomFilterCalculator {\n /**\n * \\brief Calculate probability of false positives\n * \\param n Number of items in the filter\n * \\param m Number of bits in the filter\n * \\param k Number of hash functions\n * \\return Probability of false positives\n */\n static double Probability(size_t n, size_t m, size_t k) {\n return std::pow(1.0 - std::exp(-((double)k / (double)m * (double)n)), k);\n }\n\n /**\n * \\brief Calculate number of items in the filter\n * \\param m Number of bits in the filter\n * \\param k Number of hash functions\n * \\param p Probability of false positives\n * \\return Number of items in the filter\n */\n static size_t NumberOfItems(size_t m, size_t k, double p) {\n return (size_t)std::ceil(\n -((double)m / (double)k *\n std::log(1.0 - std::exp(std::log(p) / (double)k))));\n }\n\n /**\n * \\brief Calculate number of bits in the filter\n * \\param n Number of items in the filter\n * \\param p Probability of false positives\n * \\return Number of bits in the filter\n */\n static size_t NumberOfBits(size_t n, double p) {\n return (size_t)std::ceil((double)n * std::log(p) /\n std::log(1.0 / std::pow(2.0, std::log(2.0))));\n }\n\n /**\n * \\brief Calculate number of bits in the filter\n * \\param n Number of items in the filter\n * \\param k Number of hash functions\n * \\param p Probability of false positives\n * \\return Number of bits in the filter\n */\n static size_t NumberOfBits(size_t n, size_t k, double p) {\n return (size_t)std::ceil(-((double)k * (double)n /\n std::log(1.0 - std::pow(p, 1.0 / (double)k))));\n }\n\n /**\n * \\brief Calculate number of bytes in the filter\n * \\param n Number of items in the filter\n * \\param p Probability of false positives\n * \\return Number of bytes in the filter\n */\n static size_t NumberOfBytes(size_t n, double p) {\n return ((NumberOfBits(n, p) + 7) >> 3);\n }\n\n /**\n * \\brief Calculate number of bits in the filter\n * \\param n Number of items in the filter\n * \\param k Number of hash functions\n * \\param p Probability of false positives\n * \\return Number of bits in the filter\n */\n static size_t NumberOfBytes(size_t n, size_t k, double p) {\n return ((NumberOfBits(n, k, p) + 7) >> 3);\n }\n\n /**\n * \\brief Calculate number of hash functions\n * \\param n Number of items in the filter\n * \\param m Number of bits in the filter\n * \\return Number of hash functions\n */\n static size_t NumberOfHash(size_t n, size_t m) {\n return (size_t)std::round((double)m / (double)n * std::log(2.0));\n }\n};\n\n/*! Bloom Filter\n */\ntemplate \nclass BloomFilter {\n public:\n //! Constructor\n BloomFilter(void) {}\n\n //! Constructor\n BloomFilter(size_t n, double p) {\n if (n > 0 && p > 0.0 && p < 1.0) {\n capacity_ = n;\n bits_count_ = BloomFilterCalculator::NumberOfBits(n, K, p);\n bits_count_ = ((bits_count_ + 31) >> 5) << 5;\n probability_ = BloomFilterCalculator::Probability(n, bits_count_, K);\n bitset_ = new uint32_t[bits_count_ >> 5];\n memset(bitset_, 0, (bits_count_ >> 3));\n }\n }\n\n //! Constructor\n BloomFilter(BloomFilter &&rhs)\n : bitset_(rhs.bitset_),\n bits_count_(rhs.bits_count_),\n capacity_(rhs.capacity_),\n count_(rhs.count_),\n probability_(rhs.probability_) {\n rhs.bitset_ = nullptr;\n rhs.bits_count_ = 0u;\n rhs.capacity_ = 0u;\n rhs.count_ = 0u;\n rhs.probability_ = 0u;\n }\n\n //! Destructor\n ~BloomFilter(void) {\n delete[] bitset_;\n }\n\n //! Test if the filter is valid\n bool is_valid(void) const {\n return (bitset_ != nullptr);\n }\n\n //! Reset the bloom filter\n bool reset(size_t n, double p) {\n if (n <= 0 || p <= 0.0 || p >= 1.0) {\n return false;\n }\n delete[] bitset_;\n capacity_ = n;\n count_ = 0u;\n bits_count_ = BloomFilterCalculator::NumberOfBits(n, K, p);\n bits_count_ = ((bits_count_ + 31) >> 5) << 5;\n probability_ = BloomFilterCalculator::Probability(n, bits_count_, K);\n bitset_ = new (std::nothrow) uint32_t[bits_count_ >> 5];\n if (!bitset_) {\n return false;\n }\n memset(bitset_, 0, (bits_count_ >> 3));\n return true;\n }\n\n //! Clear the bloom filter\n void clear(void) {\n if (bitset_) {\n memset(bitset_, 0, (bits_count_ >> 3));\n count_ = 0u;\n }\n }\n\n //! Insert a item into bloom filter\n template ...>::value &&\n sizeof...(TArgs) == K>::type>\n bool insert(TArgs... vals) {\n if (count_ >= capacity_) {\n return false;\n }\n this->set_bits(vals...);\n ++count_;\n return true;\n }\n\n //! Force insert a item into bloom filter\n template ...>::value &&\n sizeof...(TArgs) == K>::type>\n void force_insert(TArgs... vals) {\n this->set_bits(vals...);\n ++count_;\n }\n\n //! Insert a item into bloom filter\n template ...>::value &&\n sizeof...(TArgs) == K>::type>\n bool has(TArgs... vals) const {\n return this->test_bits(vals...);\n }\n\n //! Retrieve count of bits in bloom filter\n size_t bits_count(void) const {\n return bits_count_;\n }\n\n //! Retrieve capacity of bloom filter\n size_t capacity(void) const {\n return capacity_;\n }\n\n //! Retrieve count of items in bloom filter\n size_t count(void) const {\n return count_;\n }\n\n //! Retrieve probability of false positives\n double probability(void) const {\n return probability_;\n }\n\n protected:\n //! Disable them\n BloomFilter(const BloomFilter &) = delete;\n BloomFilter &operator=(const BloomFilter &) = delete;\n\n //! Set bits in bloom filter\n template \n void set_bits(TArg val) {\n size_t num = static_cast(val) % bits_count_;\n bitset_[num >> 5] |= (1u << (num & 0x1f));\n }\n\n //! Set bits in bloom filter\n template \n void set_bits(TArg val, TArgs... vals) {\n this->set_bits(val);\n this->set_bits(vals...);\n }\n\n //! Test bits in bloom filter\n template \n bool test_bits(TArg val) const {\n size_t num = static_cast(val) % bits_count_;\n return ((bitset_[num >> 5] & (1u << (num & 0x1f))) != 0);\n }\n\n //! Test bits in bloom filter\n template \n bool test_bits(TArg val, TArgs... vals) const {\n if (!this->test_bits(val)) {\n return false;\n }\n return this->test_bits(vals...);\n }\n\n private:\n uint32_t *bitset_{nullptr};\n size_t bits_count_{0u};\n size_t capacity_{0u};\n size_t count_{0u};\n double probability_{0.0};\n};\n\n/*! Bloom Filter (Special)\n */\ntemplate <>\nstruct BloomFilter<0> {};\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/container/params.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n\n//! Global environ variable\nextern char **environ;\n\nnamespace zvec {\nnamespace ailego {\n\nstatic void ParseFromJsonObject(const ailego::JsonObject &obj, Params *params) {\n for (ailego::JsonObject::const_iterator it = obj.begin(); it != obj.end();\n ++it) {\n const ailego::JsonValue &val = it->value();\n\n if (val.is_boolean()) {\n params->set(it->key().as_stl_string(), val.as_bool());\n } else if (val.is_integer()) {\n params->set(it->key().as_stl_string(),\n static_cast(val.as_integer()));\n } else if (val.is_float()) {\n params->set(it->key().as_stl_string(), val.as_float());\n } else if (val.is_string()) {\n params->set(it->key().as_stl_string(),\n val.as_string().decode().as_stl_string());\n } else if (val.is_object()) {\n Params subparams;\n ParseFromJsonObject(val.as_object(), &subparams);\n params->set(it->key().as_stl_string(), std::move(subparams));\n }\n }\n}\n\nbool Params::ParseFromBuffer(const std::string &buf, Params *params) {\n ailego::JsonValue val;\n ailego::JsonParser parser;\n\n parser.set_comment(true);\n parser.set_simple(true);\n parser.set_squote(true);\n parser.set_unstrict(false);\n if (!parser.parse(buf.c_str(), &val)) {\n return false;\n }\n\n if (!val.is_object()) {\n return false;\n }\n ParseFromJsonObject(val.as_object(), params);\n return true;\n}\n\nvoid Params::ParseFromEnvironment(Params *params) {\n // Dump all environ string\n for (size_t i = 0; environ[i]; ++i) {\n const char *env = environ[i];\n const char *p = std::strchr(env, '=');\n if (p) {\n params->set(std::string(env, p - env), std::string(p + 1));\n }\n }\n}\n\nstatic void SerializeToJsonObject(const Params ¶ms,\n ailego::JsonObject *obj) {\n for (const auto &it : params.hypercube().cubes()) {\n const ailego::Cube &cube = it.second;\n const char *key = it.first.c_str();\n\n if (cube.compatible()) {\n const auto &val = cube.unsafe_cast();\n ailego::JsonString str(val.data(), val.size());\n obj->set(key, ailego::JsonValue(str.encode()));\n } else if (cube.compatible()) {\n obj->set(key,\n ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n obj->set(key, ailego::JsonValue(cube.unsafe_cast()));\n } else if (cube.compatible()) {\n ailego::JsonObject subobj;\n SerializeToJsonObject(cube.unsafe_cast(), &subobj);\n obj->set(key, ailego::JsonValue(subobj));\n } else {\n LOG_WARN(\"Unsupported serializing \\'%s\\' <%s>.\", key, cube.type().name());\n }\n }\n}\n\nvoid Params::SerializeToBuffer(const Params ¶ms, std::string *buf) {\n if (buf != nullptr) {\n ailego::JsonObject obj;\n SerializeToJsonObject(params, &obj);\n buf->assign(ailego::JsonValue(obj).as_json_string().as_stl_string());\n }\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/container/reservoir.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Sampling Reservoir\n */\ntemplate >\nclass Reservoir {\n public:\n //! Constructor\n Reservoir(size_t cnt)\n : samples_(cnt), total_(0), mt_(std::random_device()()), pool_() {\n pool_.reserve(samples_);\n }\n\n //! Constructor\n Reservoir(const Reservoir &rhs)\n : samples_(rhs.samples_),\n total_(rhs.total_),\n mt_(std::random_device()()),\n pool_(rhs.pool_) {}\n\n //! Constructor\n Reservoir(Reservoir &&rhs)\n : samples_(rhs.samples_),\n total_(rhs.total_),\n mt_(std::random_device()()),\n pool_(std::move(rhs.pool_)) {}\n\n //! Destructor\n ~Reservoir(void) {}\n\n //! Assignment\n Reservoir &operator=(const Reservoir &rhs) {\n samples_ = rhs.samples_;\n total_ = rhs.total_;\n pool_ = rhs.pool_;\n return *this;\n }\n\n //! Assignment\n Reservoir &operator=(Reservoir &&rhs) {\n samples_ = rhs.samples_;\n total_ = rhs.total_;\n pool_ = std::move(rhs.pool_);\n return *this;\n }\n\n //! Retrieve pool of reservoir\n std::vector *mutable_pool(void) {\n return &pool_;\n }\n\n //! Retrieve pool of reservoir\n const std::vector &pool(void) const {\n return pool_;\n }\n\n //! Retrieve count of samples\n size_t samples(void) const {\n return samples_;\n }\n\n //! Retrieve total count of filling\n size_t total(void) const {\n return total_;\n }\n\n //! Reset the reservoir\n void reset(void) {\n total_ = 0;\n pool_.clear();\n pool_.reserve(samples_);\n }\n\n //! Fill the reservoir\n void fill(const T &item) {\n if (samples_ > 0) {\n if (pool_.size() >= samples_) {\n std::uniform_int_distribution dt(0, total_);\n size_t i = dt(mt_);\n\n if (i < samples_) {\n pool_[i] = item;\n }\n } else {\n pool_.push_back(item);\n }\n }\n ++total_;\n }\n\n //! Fill the reservoir\n void fill(T &&item) {\n if (samples_ > 0) {\n if (pool_.size() >= samples_) {\n std::uniform_int_distribution dt(0, total_);\n size_t i = dt(mt_);\n\n if (i < samples_) {\n pool_[i] = std::move(item);\n }\n } else {\n pool_.push_back(std::move(item));\n }\n }\n ++total_;\n }\n\n private:\n //! Disable them\n Reservoir(void) = delete;\n\n //! Members\n size_t samples_;\n size_t total_;\n std::mt19937 mt_;\n std::vector pool_;\n};\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/container/vector_array.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Numerical Vector Array\n */\ntemplate ::value>::type>\nclass NumericalVectorArray {\n public:\n //! Type of value\n using ValueType = typename NumericalVector::ValueType;\n\n //! Constructor\n NumericalVectorArray(void) {}\n\n //! Constructor\n explicit NumericalVectorArray(size_t dim) : dimension_(dim) {}\n\n //! Constructor\n NumericalVectorArray(const NumericalVectorArray &rhs)\n : dimension_(rhs.dimension_), buffer_(rhs.buffer_) {}\n\n //! Constructor\n NumericalVectorArray(NumericalVectorArray &&rhs)\n : dimension_(rhs.dimension_), buffer_(std::move(rhs.buffer_)) {}\n\n //! Assignment\n NumericalVectorArray &operator=(const NumericalVectorArray &rhs) {\n dimension_ = rhs.dimension_;\n buffer_ = rhs.buffer_;\n return *this;\n }\n\n //! Assignment\n NumericalVectorArray &operator=(NumericalVectorArray &&rhs) {\n dimension_ = rhs.dimension_;\n buffer_ = std::move(rhs.buffer_);\n return *this;\n }\n\n //! Overloaded operator []\n ValueType *operator[](size_t i) {\n return (reinterpret_cast(&buffer_[0]) + i * dimension_);\n }\n\n //! Overloaded operator []\n const ValueType *operator[](size_t i) const {\n return (reinterpret_cast(buffer_.data()) +\n i * dimension_);\n }\n\n //! Append a vector\n void append(const ValueType *vec, size_t dim) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n buffer_.append(reinterpret_cast(vec),\n dim * sizeof(ValueType));\n }\n\n //! Append vectors\n void append(const ValueType *vec, size_t dim, size_t cnt) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n buffer_.append(reinterpret_cast(vec),\n cnt * dim * sizeof(ValueType));\n }\n\n //! Append a vector\n void append(const NumericalVector &vec) {\n this->append(vec.data(), vec.dimension());\n }\n\n //! Replace a vector\n void replace(size_t index, const ValueType *vec, size_t dim) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n size_t element_size = dim * sizeof(ValueType);\n buffer_.replace(index * element_size, element_size,\n reinterpret_cast(vec), element_size);\n }\n\n //! Replace a vector\n void replace(size_t index, const ValueType *vec, size_t dim, size_t cnt) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n size_t element_size = dim * sizeof(ValueType);\n size_t total = element_size * cnt;\n buffer_.replace(index * element_size, total,\n reinterpret_cast(vec), total);\n }\n\n //! Replace a vector\n void replace(size_t index, const NumericalVector &vec) {\n this->replace(index, vec.data(), vec.dimension());\n }\n\n //! Request a change in capacity\n void reserve(size_t n) {\n buffer_.reserve(n * dimension_ * sizeof(ValueType));\n }\n\n //! Resize the array to a length of n elements\n void resize(size_t n) {\n buffer_.resize(n * dimension_ * sizeof(ValueType));\n }\n\n //! Clear the vector array\n void clear(void) {\n buffer_.clear();\n }\n\n //! Reset the vector array\n void reset(size_t dim) {\n dimension_ = dim;\n buffer_.clear();\n }\n\n //! Requests the removal of unused capacity.\n void shrink_to_fit(void) {\n buffer_.shrink_to_fit();\n }\n\n //! Retrieve pointer of data\n ValueType *data(void) {\n return reinterpret_cast(&buffer_[0]);\n }\n\n //! Retrieve pointer of data\n const ValueType *data(void) const {\n return reinterpret_cast(buffer_.data());\n }\n\n //! Retrieve pointer of data\n ValueType *at(size_t i) {\n if (ailego_unlikely(i >= this->count())) {\n throw std::out_of_range(\"Index overflow\");\n }\n return (reinterpret_cast(&buffer_[0]) + i * dimension_);\n }\n\n //! Retrieve pointer of data\n const ValueType *at(size_t i) const {\n if (ailego_unlikely(i >= this->count())) {\n throw std::out_of_range(\"Index overflow\");\n }\n return (reinterpret_cast(buffer_.data()) +\n i * dimension_);\n }\n\n //! Test if the array is empty\n bool empty(void) const {\n return buffer_.empty();\n }\n\n //! Retrieve count of vectors\n size_t count(void) const {\n return (dimension_ > 0 ? buffer_.size() / (dimension_ * sizeof(ValueType))\n : 0u);\n }\n\n //! Retrieve dimension of vector\n size_t dimension(void) const {\n return dimension_;\n }\n\n //! Retrieve size of array in bytes\n size_t bytes(void) const {\n return buffer_.size();\n }\n\n private:\n size_t dimension_{0u};\n std::string buffer_{};\n};\n\n/*! Nibble Vector Array\n */\ntemplate ::value>::type>\nclass NibbleVectorArray {\n public:\n //! Type of value\n using ValueType = typename NibbleVector::ValueType;\n using StoreType = typename NibbleVector::StoreType;\n\n //! Constructor\n NibbleVectorArray(void) {}\n\n //! Constructor\n explicit NibbleVectorArray(size_t dim)\n : dimension_((dim + (sizeof(ValueType) << 1) - 1) /\n (sizeof(ValueType) << 1) * sizeof(ValueType)\n << 1) {}\n\n //! Constructor\n NibbleVectorArray(const NibbleVectorArray &rhs)\n : dimension_(rhs.dimension_), buffer_(rhs.buffer_) {}\n\n //! Constructor\n NibbleVectorArray(NibbleVectorArray &&rhs)\n : dimension_(rhs.dimension_), buffer_(std::move(rhs.buffer_)) {}\n\n //! Assignment\n NibbleVectorArray &operator=(const NibbleVectorArray &rhs) {\n dimension_ = rhs.dimension_;\n buffer_ = rhs.buffer_;\n return *this;\n }\n\n //! Assignment\n NibbleVectorArray &operator=(NibbleVectorArray &&rhs) {\n dimension_ = rhs.dimension_;\n buffer_ = std::move(rhs.buffer_);\n return *this;\n }\n\n //! Overloaded operator []\n StoreType *operator[](size_t i) {\n return reinterpret_cast(&buffer_[0] + i * (dimension_ >> 1));\n }\n\n //! Overloaded operator []\n const StoreType *operator[](size_t i) const {\n return reinterpret_cast(&buffer_[0] +\n i * (dimension_ >> 1));\n }\n\n //! Append a vector\n void append(const StoreType *vec, size_t dim) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n buffer_.append(reinterpret_cast(vec), dim >> 1);\n }\n\n //! Append vectors\n void append(const StoreType *vec, size_t dim, size_t cnt) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n buffer_.append(reinterpret_cast(vec), cnt * (dim >> 1));\n }\n\n //! Append a vector\n void append(const NibbleVector &vec) {\n this->append(vec.data(), vec.dimension());\n }\n\n //! Replace a vector\n void replace(size_t index, const StoreType *vec, size_t dim) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n size_t element_size = (dim >> 1);\n buffer_.replace(index * element_size, element_size,\n reinterpret_cast(vec), element_size);\n }\n\n //! Replace a vector\n void replace(size_t index, const StoreType *vec, size_t dim, size_t cnt) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n size_t element_size = (dim >> 1);\n size_t total = element_size * cnt;\n buffer_.replace(index * element_size, total,\n reinterpret_cast(vec), total);\n }\n\n //! Replace a vector\n void replace(size_t index, const NibbleVector &vec) {\n this->replace(index, vec.data(), vec.dimension());\n }\n\n //! Request a change in capacity\n void reserve(size_t n) {\n buffer_.reserve(n * (dimension_ >> 1));\n }\n\n //! Resize the array to a length of n elements\n void resize(size_t n) {\n buffer_.resize(n * (dimension_ >> 1));\n }\n\n //! Clear the vector array\n void clear(void) {\n buffer_.clear();\n }\n\n //! Reset the vector array\n void reset(size_t dim) {\n dimension_ = (dim + (sizeof(ValueType) << 1) - 1) /\n (sizeof(ValueType) << 1) * sizeof(ValueType)\n << 1;\n buffer_.clear();\n }\n\n //! Requests the removal of unused capacity.\n void shrink_to_fit(void) {\n buffer_.shrink_to_fit();\n }\n\n //! Retrieve pointer of data\n StoreType *data(void) {\n return reinterpret_cast(&buffer_[0]);\n }\n\n //! Retrieve pointer of data\n const StoreType *data(void) const {\n return reinterpret_cast(buffer_.data());\n }\n\n //! Retrieve pointer of data\n StoreType *at(size_t i) {\n if (ailego_unlikely(i >= this->count())) {\n throw std::out_of_range(\"Index overflow\");\n }\n return reinterpret_cast(&buffer_[0] + i * (dimension_ >> 1));\n }\n\n //! Retrieve pointer of data\n const StoreType *at(size_t i) const {\n if (ailego_unlikely(i >= this->count())) {\n throw std::out_of_range(\"Index overflow\");\n }\n return reinterpret_cast(buffer_.data() +\n i * (dimension_ >> 1));\n }\n\n //! Test if the array is empty\n bool empty(void) const {\n return buffer_.empty();\n }\n\n //! Retrieve count of vectors\n size_t count(void) const {\n return (dimension_ > 1 ? buffer_.size() / (dimension_ >> 1) : 0u);\n }\n\n //! Retrieve dimension of vector\n size_t dimension(void) const {\n return dimension_;\n }\n\n //! Retrieve size of array in bytes\n size_t bytes(void) const {\n return buffer_.size();\n }\n\n private:\n size_t dimension_{0u};\n std::string buffer_{};\n};\n\n/*! Binary Vector Array\n */\ntemplate ::value>::type>\nclass BinaryVectorArray {\n public:\n //! Type of value\n using ValueType = typename BinaryVector::ValueType;\n\n //! Constructor\n BinaryVectorArray(void) {}\n\n //! Constructor\n explicit BinaryVectorArray(size_t dim)\n : dimension_((dim + (sizeof(ValueType) << 3) - 1) /\n (sizeof(ValueType) << 3) * (sizeof(ValueType) << 3)) {}\n\n //! Constructor\n BinaryVectorArray(const BinaryVectorArray &rhs)\n : dimension_(rhs.dimension_), buffer_(rhs.buffer_) {}\n\n //! Constructor\n BinaryVectorArray(BinaryVectorArray &&rhs)\n : dimension_(rhs.dimension_), buffer_(std::move(rhs.buffer_)) {}\n\n //! Assignment\n BinaryVectorArray &operator=(const BinaryVectorArray &rhs) {\n dimension_ = rhs.dimension_;\n buffer_ = rhs.buffer_;\n return *this;\n }\n\n //! Assignment\n BinaryVectorArray &operator=(BinaryVectorArray &&rhs) {\n dimension_ = rhs.dimension_;\n buffer_ = std::move(rhs.buffer_);\n return *this;\n }\n\n //! Overloaded operator []\n ValueType *operator[](size_t i) {\n return reinterpret_cast(&buffer_[0] + i * (dimension_ >> 3));\n }\n\n //! Overloaded operator []\n const ValueType *operator[](size_t i) const {\n return reinterpret_cast(buffer_.data() +\n i * (dimension_ >> 3));\n }\n\n //! Append a vector\n void append(const ValueType *vec, size_t dim) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n buffer_.append(reinterpret_cast(vec), (dim >> 3));\n }\n\n //! Append vectors\n void append(const ValueType *vec, size_t dim, size_t cnt) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n buffer_.append(reinterpret_cast(vec), cnt * (dim >> 3));\n }\n\n //! Append a vector\n void append(const BinaryVector &vec) {\n this->append(vec.data(), vec.dimension());\n }\n\n //! Replace a vector\n void replace(size_t index, const ValueType *vec, size_t dim) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n size_t element_size = (dim >> 3);\n buffer_.replace(index * element_size, element_size,\n reinterpret_cast(vec), element_size);\n }\n\n //! Replace a vector\n void replace(size_t index, const ValueType *vec, size_t dim, size_t cnt) {\n if (ailego_unlikely(dim != dimension_)) {\n throw std::length_error(\"Unmatched dimension\");\n }\n size_t element_size = (dim >> 3);\n size_t total = element_size * cnt;\n buffer_.replace(index * element_size, total,\n reinterpret_cast(vec), total);\n }\n\n //! Replace a vector\n void replace(size_t index, const BinaryVector &vec) {\n this->replace(index, vec.data(), vec.dimension());\n }\n\n //! Request a change in capacity\n void reserve(size_t n) {\n buffer_.reserve(n * (dimension_ >> 3));\n }\n\n //! Resize the array to a length of n elements\n void resize(size_t n) {\n buffer_.resize(n * (dimension_ >> 3));\n }\n\n //! Clear the vector array\n void clear(void) {\n buffer_.clear();\n }\n\n //! Reset the vector array\n void reset(size_t dim) {\n dimension_ = (dim + (sizeof(ValueType) << 3) - 1) /\n (sizeof(ValueType) << 3) * (sizeof(ValueType) << 3);\n buffer_.clear();\n }\n\n //! Requests the removal of unused capacity.\n void shrink_to_fit(void) {\n buffer_.shrink_to_fit();\n }\n\n //! Retrieve pointer of data\n ValueType *data(void) {\n return reinterpret_cast(&buffer_[0]);\n }\n\n //! Retrieve pointer of data\n const ValueType *data(void) const {\n return reinterpret_cast(buffer_.data());\n }\n\n //! Retrieve pointer of data\n ValueType *at(size_t i) {\n if (ailego_unlikely(i >= this->count())) {\n throw std::out_of_range(\"Index overflow\");\n }\n return reinterpret_cast(&buffer_[0] + i * (dimension_ >> 3));\n }\n\n //! Retrieve pointer of data\n const ValueType *at(size_t i) const {\n if (ailego_unlikely(i >= this->count())) {\n throw std::out_of_range(\"Index overflow\");\n }\n return reinterpret_cast(buffer_.data() +\n i * (dimension_ >> 3));\n }\n\n //! Test if the array is empty\n bool empty(void) const {\n return buffer_.empty();\n }\n\n //! Retrieve count of vectors\n size_t count(void) const {\n return (dimension_ > 0 ? buffer_.size() / (dimension_ >> 3) : 0u);\n }\n\n //! Retrieve dimension of vector\n size_t dimension(void) const {\n return dimension_;\n }\n\n //! Retrieve size of array in bytes\n size_t bytes(void) const {\n return buffer_.size();\n }\n\n private:\n size_t dimension_{0u};\n std::string buffer_{};\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/encoding/json/mod_json.c", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n\n#ifndef MOD_JSON_TOKEN_DEFOPTS\n#define MOD_JSON_TOKEN_DEFOPTS 0 /* default options of token */\n#endif\n#ifndef MOD_JSON_TOKEN_DEFOBJDEP\n#define MOD_JSON_TOKEN_DEFOBJDEP 64 /* default objects depth of token */\n#endif\n#ifndef MOD_JSON_TOKEN_DEFARRDEP\n#define MOD_JSON_TOKEN_DEFARRDEP 64 /* default arrays depth of token */\n#endif\n#ifndef MOD_JSON_STRING_DEFSIZE\n#define MOD_JSON_STRING_DEFSIZE 32 /* default started size of string */\n#endif\n#ifndef MOD_JSON_ARRAY_DEFSIZE\n#define MOD_JSON_ARRAY_DEFSIZE 32 /* default started size of array */\n#endif\n#ifndef MOD_JSON_OBJECT_DEFSIZE\n#define MOD_JSON_OBJECT_DEFSIZE 32 /* default started size of object */\n#endif\n\n#ifndef mod_json_malloc\n#define mod_json_malloc malloc\n#endif\n#ifndef mod_json_free\n#define mod_json_free free\n#endif\n\n#ifdef __GNUC__\n#define mod_json_likely(x) __builtin_expect(!!(x), 1)\n#define mod_json_unlikely(x) __builtin_expect(!!(x), 0)\n#else\n#define mod_json_likely(x) (x)\n#define mod_json_unlikely(x) (x)\n#endif\n\n#define mod_json_minus_if_ne_zero(COND) \\\n if (mod_json_unlikely((COND) != 0)) return (-1)\n\n#define mod_json_minus_if_false(COND) \\\n if (mod_json_unlikely(!(COND))) return (-1)\n\n#define mod_json_null_if_ne_zero(COND) \\\n if (mod_json_unlikely((COND) != 0)) return (NULL)\n\n#define mod_json_null_if_false(COND) \\\n if (mod_json_unlikely(!(COND))) return (NULL)\n\n#if defined(_MSC_VER)\n#pragma warning(disable : 4200)\n#define strtoull _strtoui64\n#define snprintf(buf, size, format, ...) \\\n _snprintf_s(buf, size, _TRUNCATE, format, ##__VA_ARGS__)\n#endif\n#define mod_json_utils_snprintf snprintf\n#define mod_json_utils_strtoi strtoull\n#define mod_json_utils_strtof strtod\n#define mod_json_utils_strlen strlen\n\n/*! JSON Token\n */\nstruct mod_json_token {\n mod_json_state_t state;\n mod_json_error_t error;\n mod_json_cchar_t *context;\n mod_json_size_t options;\n mod_json_size_t object_max_depth;\n mod_json_size_t array_max_depth;\n mod_json_size_t object_depth;\n mod_json_size_t array_depth;\n mod_json_event_t event_code;\n mod_json_event_proc event_proc;\n mod_json_void_t *param;\n mod_json_char_t tags[0];\n};\n\ntypedef struct mod_json_parser mod_json_parser_t;\n\n/*! JSON Parser\n */\nstruct mod_json_parser {\n mod_json_string_t *key;\n mod_json_value_t *val_null;\n mod_json_value_t *val_true;\n mod_json_value_t *val_false;\n mod_json_value_t *val_zero;\n mod_json_value_t *val_zerof;\n mod_json_value_t *val_empty;\n mod_json_value_t *vals[0];\n};\n\nstatic inline mod_json_size_t mod_json_utils_clp2(mod_json_size_t n) {\n n = n - 1;\n n = n | (n >> 1);\n n = n | (n >> 2);\n n = n | (n >> 4);\n n = n | (n >> 8);\n n = n | (n >> 16);\n return (n + 1);\n}\n\nstatic inline mod_json_size_t mod_json_utils_itostr(mod_json_char_t *buf,\n mod_json_integer_t val) {\n mod_json_char_t *pos, *first, *last;\n\n pos = buf;\n if (val < 0) {\n *pos++ = '-';\n val = -val;\n }\n\n /* save pointer to first digit */\n first = pos;\n\n do {\n /* convert to ASCII and store */\n *pos++ = (mod_json_char_t)(val % 10 + '0');\n\n /* next digit */\n val /= 10;\n\n } while (val > 0);\n\n *pos = '\\0';\n\n /* save pointer to last digit */\n last = pos - 1;\n\n /* reverse digit string */\n while (first < last) {\n mod_json_char_t temp = *first;\n *first++ = *last;\n *last-- = temp;\n }\n return (mod_json_size_t)(pos - buf);\n}\n\nstatic inline mod_json_float_t mod_json_utils_pow10(int n) {\n /* 1e-308...1e308: 617 * 8 bytes = 4936 bytes */\n static const mod_json_float_t etab[] = {\n 1e-308, 1e-307, 1e-306, 1e-305, 1e-304, 1e-303, 1e-302, 1e-301, 1e-300,\n 1e-299, 1e-298, 1e-297, 1e-296, 1e-295, 1e-294, 1e-293, 1e-292, 1e-291,\n 1e-290, 1e-289, 1e-288, 1e-287, 1e-286, 1e-285, 1e-284, 1e-283, 1e-282,\n 1e-281, 1e-280, 1e-279, 1e-278, 1e-277, 1e-276, 1e-275, 1e-274, 1e-273,\n 1e-272, 1e-271, 1e-270, 1e-269, 1e-268, 1e-267, 1e-266, 1e-265, 1e-264,\n 1e-263, 1e-262, 1e-261, 1e-260, 1e-259, 1e-258, 1e-257, 1e-256, 1e-255,\n 1e-254, 1e-253, 1e-252, 1e-251, 1e-250, 1e-249, 1e-248, 1e-247, 1e-246,\n 1e-245, 1e-244, 1e-243, 1e-242, 1e-241, 1e-240, 1e-239, 1e-238, 1e-237,\n 1e-236, 1e-235, 1e-234, 1e-233, 1e-232, 1e-231, 1e-230, 1e-229, 1e-228,\n 1e-227, 1e-226, 1e-225, 1e-224, 1e-223, 1e-222, 1e-221, 1e-220, 1e-219,\n 1e-218, 1e-217, 1e-216, 1e-215, 1e-214, 1e-213, 1e-212, 1e-211, 1e-210,\n 1e-209, 1e-208, 1e-207, 1e-206, 1e-205, 1e-204, 1e-203, 1e-202, 1e-201,\n 1e-200, 1e-199, 1e-198, 1e-197, 1e-196, 1e-195, 1e-194, 1e-193, 1e-192,\n 1e-191, 1e-190, 1e-189, 1e-188, 1e-187, 1e-186, 1e-185, 1e-184, 1e-183,\n 1e-182, 1e-181, 1e-180, 1e-179, 1e-178, 1e-177, 1e-176, 1e-175, 1e-174,\n 1e-173, 1e-172, 1e-171, 1e-170, 1e-169, 1e-168, 1e-167, 1e-166, 1e-165,\n 1e-164, 1e-163, 1e-162, 1e-161, 1e-160, 1e-159, 1e-158, 1e-157, 1e-156,\n 1e-155, 1e-154, 1e-153, 1e-152, 1e-151, 1e-150, 1e-149, 1e-148, 1e-147,\n 1e-146, 1e-145, 1e-144, 1e-143, 1e-142, 1e-141, 1e-140, 1e-139, 1e-138,\n 1e-137, 1e-136, 1e-135, 1e-134, 1e-133, 1e-132, 1e-131, 1e-130, 1e-129,\n 1e-128, 1e-127, 1e-126, 1e-125, 1e-124, 1e-123, 1e-122, 1e-121, 1e-120,\n 1e-119, 1e-118, 1e-117, 1e-116, 1e-115, 1e-114, 1e-113, 1e-112, 1e-111,\n 1e-110, 1e-109, 1e-108, 1e-107, 1e-106, 1e-105, 1e-104, 1e-103, 1e-102,\n 1e-101, 1e-100, 1e-99, 1e-98, 1e-97, 1e-96, 1e-95, 1e-94, 1e-93,\n 1e-92, 1e-91, 1e-90, 1e-89, 1e-88, 1e-87, 1e-86, 1e-85, 1e-84,\n 1e-83, 1e-82, 1e-81, 1e-80, 1e-79, 1e-78, 1e-77, 1e-76, 1e-75,\n 1e-74, 1e-73, 1e-72, 1e-71, 1e-70, 1e-69, 1e-68, 1e-67, 1e-66,\n 1e-65, 1e-64, 1e-63, 1e-62, 1e-61, 1e-60, 1e-59, 1e-58, 1e-57,\n 1e-56, 1e-55, 1e-54, 1e-53, 1e-52, 1e-51, 1e-50, 1e-49, 1e-48,\n 1e-47, 1e-46, 1e-45, 1e-44, 1e-43, 1e-42, 1e-41, 1e-40, 1e-39,\n 1e-38, 1e-37, 1e-36, 1e-35, 1e-34, 1e-33, 1e-32, 1e-31, 1e-30,\n 1e-29, 1e-28, 1e-27, 1e-26, 1e-25, 1e-24, 1e-23, 1e-22, 1e-21,\n 1e-20, 1e-19, 1e-18, 1e-17, 1e-16, 1e-15, 1e-14, 1e-13, 1e-12,\n 1e-11, 1e-10, 1e-9, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3,\n 1e-2, 1e-1, 1e+0, 1e+1, 1e+2, 1e+3, 1e+4, 1e+5, 1e+6,\n 1e+7, 1e+8, 1e+9, 1e+10, 1e+11, 1e+12, 1e+13, 1e+14, 1e+15,\n 1e+16, 1e+17, 1e+18, 1e+19, 1e+20, 1e+21, 1e+22, 1e+23, 1e+24,\n 1e+25, 1e+26, 1e+27, 1e+28, 1e+29, 1e+30, 1e+31, 1e+32, 1e+33,\n 1e+34, 1e+35, 1e+36, 1e+37, 1e+38, 1e+39, 1e+40, 1e+41, 1e+42,\n 1e+43, 1e+44, 1e+45, 1e+46, 1e+47, 1e+48, 1e+49, 1e+50, 1e+51,\n 1e+52, 1e+53, 1e+54, 1e+55, 1e+56, 1e+57, 1e+58, 1e+59, 1e+60,\n 1e+61, 1e+62, 1e+63, 1e+64, 1e+65, 1e+66, 1e+67, 1e+68, 1e+69,\n 1e+70, 1e+71, 1e+72, 1e+73, 1e+74, 1e+75, 1e+76, 1e+77, 1e+78,\n 1e+79, 1e+80, 1e+81, 1e+82, 1e+83, 1e+84, 1e+85, 1e+86, 1e+87,\n 1e+88, 1e+89, 1e+90, 1e+91, 1e+92, 1e+93, 1e+94, 1e+95, 1e+96,\n 1e+97, 1e+98, 1e+99, 1e+100, 1e+101, 1e+102, 1e+103, 1e+104, 1e+105,\n 1e+106, 1e+107, 1e+108, 1e+109, 1e+110, 1e+111, 1e+112, 1e+113, 1e+114,\n 1e+115, 1e+116, 1e+117, 1e+118, 1e+119, 1e+120, 1e+121, 1e+122, 1e+123,\n 1e+124, 1e+125, 1e+126, 1e+127, 1e+128, 1e+129, 1e+130, 1e+131, 1e+132,\n 1e+133, 1e+134, 1e+135, 1e+136, 1e+137, 1e+138, 1e+139, 1e+140, 1e+141,\n 1e+142, 1e+143, 1e+144, 1e+145, 1e+146, 1e+147, 1e+148, 1e+149, 1e+150,\n 1e+151, 1e+152, 1e+153, 1e+154, 1e+155, 1e+156, 1e+157, 1e+158, 1e+159,\n 1e+160, 1e+161, 1e+162, 1e+163, 1e+164, 1e+165, 1e+166, 1e+167, 1e+168,\n 1e+169, 1e+170, 1e+171, 1e+172, 1e+173, 1e+174, 1e+175, 1e+176, 1e+177,\n 1e+178, 1e+179, 1e+180, 1e+181, 1e+182, 1e+183, 1e+184, 1e+185, 1e+186,\n 1e+187, 1e+188, 1e+189, 1e+190, 1e+191, 1e+192, 1e+193, 1e+194, 1e+195,\n 1e+196, 1e+197, 1e+198, 1e+199, 1e+200, 1e+201, 1e+202, 1e+203, 1e+204,\n 1e+205, 1e+206, 1e+207, 1e+208, 1e+209, 1e+210, 1e+211, 1e+212, 1e+213,\n 1e+214, 1e+215, 1e+216, 1e+217, 1e+218, 1e+219, 1e+220, 1e+221, 1e+222,\n 1e+223, 1e+224, 1e+225, 1e+226, 1e+227, 1e+228, 1e+229, 1e+230, 1e+231,\n 1e+232, 1e+233, 1e+234, 1e+235, 1e+236, 1e+237, 1e+238, 1e+239, 1e+240,\n 1e+241, 1e+242, 1e+243, 1e+244, 1e+245, 1e+246, 1e+247, 1e+248, 1e+249,\n 1e+250, 1e+251, 1e+252, 1e+253, 1e+254, 1e+255, 1e+256, 1e+257, 1e+258,\n 1e+259, 1e+260, 1e+261, 1e+262, 1e+263, 1e+264, 1e+265, 1e+266, 1e+267,\n 1e+268, 1e+269, 1e+270, 1e+271, 1e+272, 1e+273, 1e+274, 1e+275, 1e+276,\n 1e+277, 1e+278, 1e+279, 1e+280, 1e+281, 1e+282, 1e+283, 1e+284, 1e+285,\n 1e+286, 1e+287, 1e+288, 1e+289, 1e+290, 1e+291, 1e+292, 1e+293, 1e+294,\n 1e+295, 1e+296, 1e+297, 1e+298, 1e+299, 1e+300, 1e+301, 1e+302, 1e+303,\n 1e+304, 1e+305, 1e+306, 1e+307, 1e+308};\n return (n < -308 ? 0.0 : etab[n + 308]);\n}\n\nstatic inline mod_json_cchar_t *mod_json_utils_strskpb(mod_json_cchar_t *cstr) {\n static const mod_json_char_t blanks[256] = {\n 0, 0, 0, 0, 0, 0, 0, 0, 0, '\\t', '\\n', '\\v', '\\f', '\\r', 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ' ', 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};\n\n while (*(blanks + *cstr)) {\n ++cstr;\n }\n return cstr;\n}\n\nstatic inline mod_json_cchar_t *mod_json_utils_strskpc1(\n mod_json_cchar_t *cstr) {\n mod_json_char_t c;\n\n while ((c = *cstr++) != '\\0') {\n if (c == '\\r' || c == '\\n') {\n return mod_json_utils_strskpb(cstr);\n }\n }\n return (cstr - 1);\n}\n\nstatic inline mod_json_cchar_t *mod_json_utils_strskpc2(\n mod_json_cchar_t *cstr) {\n mod_json_char_t c;\n\n while ((c = *cstr++) != '\\0') {\n /* asterisk, slash */\n if (c == '*' && *cstr == '/') {\n return mod_json_utils_strskpb(cstr + 1);\n }\n }\n return (cstr - 1);\n}\n\nstatic inline mod_json_cchar_t *mod_json_utils_strskp(mod_json_cchar_t *cstr) {\n cstr = mod_json_utils_strskpb(cstr);\n\n /* treat it as comments? */\n while (*cstr == '/') {\n mod_json_char_t c = *(cstr + 1); /* second char */\n\n if (c == '/') {\n /* two slashes */\n cstr = mod_json_utils_strskpc1(cstr + 2);\n } else if (c == '*') {\n /* slash, asterisk */\n cstr = mod_json_utils_strskpc2(cstr + 2);\n } else {\n /* invalid format */\n break;\n }\n }\n return cstr;\n}\n\nstatic inline int mod_json_utils_char2hex(mod_json_char_t ch) {\n static const mod_json_char_t char2hex[256] = {\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 0, 1, 2, 3, 4, 5,\n 6, 7, 8, 9, 16, 16, 16, 16, 16, 16, 16, 10, 11, 12, 13, 14, 15, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 10, 11, 12, 13, 14, 15, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16,\n 16, 16, 16, 16};\n return *(char2hex + ch);\n}\n\nstatic inline mod_json_cchar_t *mod_json_utils_strfquo(mod_json_cchar_t *cstr,\n mod_json_char_t quo) {\n mod_json_char_t c;\n\n for (c = *cstr; c != quo; c = *(++cstr)) {\n if ((mod_json_uchar_t)c <= 0x1f) {\n return NULL;\n }\n\n if (c != '\\\\') {\n continue;\n }\n\n /* next char */\n switch (*(++cstr)) {\n case '\\\"':\n case '/':\n case 'b':\n case 'f':\n case '\\\\':\n case 'n':\n case 'r':\n case 't':\n /* ignore next char */\n break;\n\n case 'u':\n if (mod_json_utils_char2hex(*(cstr + 1)) > 15) {\n return NULL;\n }\n if (mod_json_utils_char2hex(*(cstr + 2)) > 15) {\n return NULL;\n }\n if (mod_json_utils_char2hex(*(cstr + 3)) > 15) {\n return NULL;\n }\n if (mod_json_utils_char2hex(*(cstr + 4)) > 15) {\n return NULL;\n }\n cstr += 4;\n break;\n\n default:\n /* invalid */\n return NULL;\n }\n }\n /* found it */\n return cstr;\n}\n\nstatic inline mod_json_cchar_t *mod_json_utils_strfquo2(mod_json_cchar_t *cstr,\n mod_json_char_t quo) {\n mod_json_char_t c;\n\n for (c = *cstr; c; c = *(++cstr)) {\n if (c == quo) {\n /* found it */\n return cstr;\n }\n\n if (c == '\\\\') {\n /* ignore next char */\n if (*(++cstr) == '\\0') {\n break;\n }\n }\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_utils_strfsep(mod_json_cchar_t *cstr) {\n mod_json_char_t c;\n\n while ((c = *cstr++) != '\\0') {\n switch (c) {\n case ':':\n case ' ':\n case '\\t':\n case '\\r':\n case '\\n':\n case '\\f':\n case '\\v':\n return (cstr - 1);\n }\n }\n return (cstr - 1);\n}\n\nstatic inline mod_json_cchar_t *mod_json_utils_strfsep2(\n mod_json_cchar_t *cstr) {\n mod_json_char_t c;\n\n while ((c = *cstr++) != '\\0') {\n switch (c) {\n case ':':\n case ' ':\n case '\\t':\n case '\\r':\n case '\\n':\n case '\\f':\n case '\\v':\n return (cstr - 1);\n\n case '/':\n if (*cstr == '/' || *cstr == '*') {\n return (cstr - 1);\n }\n }\n }\n return (cstr - 1);\n}\n\nstatic inline mod_json_char_t *mod_json_utils_uni2utf8(mod_json_char_t *buf,\n mod_json_size_t size,\n mod_json_uchar_t high,\n mod_json_uchar_t low) {\n /* convert to UTF-8 */\n if (high >= 0x8) {\n /* 0800 - FFFF | 1110xxxx 10xxxxxx 10xxxxxx */\n if (size >= 3) {\n *buf++ = (mod_json_char_t)(0xE0 | (high >> 4));\n *buf++ = (mod_json_char_t)(0x80 | ((high & 0xF) << 2) | (low >> 6));\n *buf++ = (mod_json_char_t)(0x80 | (low & 0x3F));\n return buf;\n }\n } else if (high > 0 || low >= 0x80) {\n /* 0080 - 07FF | 110xxxxx 10xxxxxx */\n if (size >= 2) {\n *buf++ = (mod_json_char_t)(0xC0 | (high << 2) | (low >> 6));\n *buf++ = (mod_json_char_t)(0x80 | (low & 0x3F));\n return buf;\n }\n } else {\n /* 0000 - 007F | 0xxxxxxx */\n if (size >= 1) {\n *buf++ = (mod_json_char_t)(low);\n return buf;\n }\n }\n return (mod_json_char_t *)0;\n}\n\nmod_json_value_t *mod_json_value_set_null(void) {\n mod_json_value_t *val;\n\n /* create a value */\n val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t));\n mod_json_null_if_false(val);\n\n val->refer = 1;\n val->type = mod_json_type_null;\n val->data.c_int = 0;\n return val;\n}\n\nmod_json_value_t *mod_json_value_set_object(mod_json_object_t *obj) {\n mod_json_value_t *val;\n\n /* create a value */\n val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t));\n mod_json_null_if_false(val);\n\n val->refer = 1;\n val->type = mod_json_type_object;\n val->data.c_obj = obj ? mod_json_object_grab(obj) : NULL;\n return val;\n}\n\nmod_json_value_t *mod_json_value_set_array(mod_json_array_t *arr) {\n mod_json_value_t *val;\n\n /* create a value */\n val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t));\n mod_json_null_if_false(val);\n\n val->refer = 1;\n val->type = mod_json_type_array;\n val->data.c_arr = arr ? mod_json_array_grab(arr) : NULL;\n return val;\n}\n\nmod_json_value_t *mod_json_value_set_string(mod_json_string_t *str) {\n mod_json_value_t *val;\n\n /* create a value */\n val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t));\n mod_json_null_if_false(val);\n\n val->refer = 1;\n val->type = mod_json_type_string;\n val->data.c_str = str ? mod_json_string_grab(str) : NULL;\n return val;\n}\n\nmod_json_value_t *mod_json_value_set_buffer(mod_json_cchar_t *buf,\n mod_json_size_t len) {\n mod_json_value_t *val;\n mod_json_string_t *str;\n\n /* create a value */\n val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t));\n mod_json_null_if_false(val);\n\n /* create a string */\n str = mod_json_string_set(buf, len);\n if (mod_json_unlikely(!str)) {\n mod_json_free(val);\n return NULL;\n }\n\n val->refer = 1;\n val->type = mod_json_type_string;\n val->data.c_str = str;\n return val;\n}\n\nmod_json_value_t *mod_json_value_set_integer(mod_json_integer_t num) {\n mod_json_value_t *val;\n\n /* create a value */\n val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t));\n mod_json_null_if_false(val);\n\n val->refer = 1;\n val->type = mod_json_type_integer;\n val->data.c_int = num;\n return val;\n}\n\nmod_json_value_t *mod_json_value_set_float(mod_json_float_t dbl) {\n mod_json_value_t *val;\n\n /* create a value */\n val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t));\n mod_json_null_if_false(val);\n\n val->refer = 1;\n val->type = mod_json_type_float;\n val->data.c_float = dbl;\n return val;\n}\n\nmod_json_value_t *mod_json_value_set_boolean(mod_json_boolean_t bol) {\n mod_json_value_t *val;\n\n /* create a value */\n val = (mod_json_value_t *)mod_json_malloc(sizeof(mod_json_value_t));\n mod_json_null_if_false(val);\n\n val->refer = 1;\n val->type = mod_json_type_boolean;\n val->data.c_bool = bol ? MOD_JSON_TRUE : MOD_JSON_FALSE;\n return val;\n}\n\nstatic inline void mod_json_value_clear(mod_json_value_t *val) {\n switch (val->type) {\n case mod_json_type_object:\n mod_json_object_unset(val->data.c_obj);\n break;\n\n case mod_json_type_array:\n mod_json_array_unset(val->data.c_arr);\n break;\n\n case mod_json_type_string:\n mod_json_string_unset(val->data.c_str);\n break;\n\n default:\n break;\n }\n}\n\nvoid mod_json_value_assign_null(mod_json_value_t *val) {\n if (val) {\n mod_json_value_clear(val);\n val->type = mod_json_type_null;\n val->data.c_int = 0;\n }\n}\n\nvoid mod_json_value_assign_object(mod_json_value_t *val,\n mod_json_object_t *obj) {\n if (val) {\n mod_json_value_clear(val);\n val->type = mod_json_type_object;\n val->data.c_obj = obj ? mod_json_object_grab(obj) : NULL;\n }\n}\n\nvoid mod_json_value_assign_array(mod_json_value_t *val, mod_json_array_t *arr) {\n if (val) {\n mod_json_value_clear(val);\n val->type = mod_json_type_array;\n val->data.c_arr = arr ? mod_json_array_grab(arr) : NULL;\n }\n}\n\nvoid mod_json_value_assign_string(mod_json_value_t *val,\n mod_json_string_t *str) {\n if (val) {\n mod_json_value_clear(val);\n val->type = mod_json_type_string;\n val->data.c_str = str ? mod_json_string_grab(str) : NULL;\n }\n}\n\nvoid mod_json_value_assign_integer(mod_json_value_t *val,\n mod_json_integer_t num) {\n if (val) {\n mod_json_value_clear(val);\n val->type = mod_json_type_integer;\n val->data.c_int = num;\n }\n}\n\nvoid mod_json_value_assign_float(mod_json_value_t *val, mod_json_float_t dbl) {\n if (val) {\n mod_json_value_clear(val);\n val->type = mod_json_type_float;\n val->data.c_float = dbl;\n }\n}\n\nvoid mod_json_value_assign_boolean(mod_json_value_t *val,\n mod_json_boolean_t bol) {\n if (val) {\n mod_json_value_clear(val);\n val->type = mod_json_type_boolean;\n val->data.c_bool = bol ? MOD_JSON_TRUE : MOD_JSON_FALSE;\n }\n}\n\nvoid mod_json_value_assign(mod_json_value_t *dst, mod_json_value_t *src) {\n if (!dst || dst == src) {\n return;\n }\n\n if (!src) {\n /* treat as JSON null */\n mod_json_value_assign_null(dst);\n return;\n }\n\n switch (src->type) {\n case mod_json_type_boolean:\n mod_json_value_assign_boolean(dst, src->data.c_bool);\n break;\n\n case mod_json_type_integer:\n mod_json_value_assign_integer(dst, src->data.c_int);\n break;\n\n case mod_json_type_float:\n mod_json_value_assign_float(dst, src->data.c_float);\n break;\n\n case mod_json_type_string:\n mod_json_value_assign_string(dst, src->data.c_str);\n break;\n\n case mod_json_type_array:\n mod_json_value_assign_array(dst, src->data.c_arr);\n break;\n\n case mod_json_type_object:\n mod_json_value_assign_object(dst, src->data.c_obj);\n break;\n\n default:\n mod_json_value_assign_null(dst);\n break;\n }\n}\n\nstatic inline int mod_json_value_merge_array(mod_json_value_t *val,\n mod_json_array_t *arr) {\n if (val->type != mod_json_type_array || !val->data.c_arr) {\n mod_json_value_assign_array(val, arr);\n return 0;\n }\n\n if (arr) {\n if (mod_json_array_is_shared(val->data.c_arr)) {\n mod_json_array_put(val->data.c_arr);\n val->data.c_arr = mod_json_array_clone(val->data.c_arr);\n }\n return mod_json_array_merge(val->data.c_arr, arr);\n }\n return 0;\n}\n\nstatic inline int mod_json_value_merge_object(mod_json_value_t *val,\n mod_json_object_t *obj) {\n if (val->type != mod_json_type_object || !val->data.c_obj) {\n mod_json_value_assign_object(val, obj);\n return 0;\n }\n\n if (obj) {\n if (mod_json_object_is_shared(val->data.c_obj)) {\n mod_json_object_put(val->data.c_obj);\n val->data.c_obj = mod_json_object_clone(val->data.c_obj);\n }\n return mod_json_object_merge(val->data.c_obj, obj);\n }\n return 0;\n}\n\nint mod_json_value_merge(mod_json_value_t *dst, mod_json_value_t *src) {\n mod_json_minus_if_false(dst && dst != src);\n\n if (!src) {\n mod_json_value_assign_null(dst);\n return 0;\n }\n\n switch (src->type) {\n case mod_json_type_boolean:\n mod_json_value_assign_boolean(dst, src->data.c_bool);\n break;\n\n case mod_json_type_integer:\n mod_json_value_assign_integer(dst, src->data.c_int);\n break;\n\n case mod_json_type_float:\n mod_json_value_assign_float(dst, src->data.c_float);\n break;\n\n case mod_json_type_string:\n mod_json_value_assign_string(dst, src->data.c_str);\n break;\n\n case mod_json_type_array:\n return mod_json_value_merge_array(dst, src->data.c_arr);\n\n case mod_json_type_object:\n return mod_json_value_merge_object(dst, src->data.c_obj);\n\n default:\n mod_json_value_assign_null(dst);\n break;\n }\n return 0;\n}\n\nmod_json_object_t *mod_json_value_object(mod_json_value_t *val) {\n if (val && val->type == mod_json_type_object) {\n return (val->data.c_obj);\n }\n return NULL;\n}\n\nmod_json_array_t *mod_json_value_array(mod_json_value_t *val) {\n if (val && val->type == mod_json_type_array) {\n return (val->data.c_arr);\n }\n return NULL;\n}\n\nmod_json_string_t *mod_json_value_string(mod_json_value_t *val) {\n if (val && val->type == mod_json_type_string) {\n return (val->data.c_str);\n }\n return NULL;\n}\n\nmod_json_cchar_t *mod_json_value_cstring(mod_json_value_t *val) {\n if (val && val->type == mod_json_type_string) {\n return mod_json_string_cstr(val->data.c_str);\n }\n return NULL;\n}\n\nmod_json_float_t mod_json_value_float(mod_json_value_t *val) {\n if (val) {\n switch (val->type) {\n case mod_json_type_boolean:\n return (val->data.c_bool ? 1.0 : 0.0);\n\n case mod_json_type_integer:\n return (mod_json_float_t)(val->data.c_int);\n\n case mod_json_type_float:\n return (val->data.c_float);\n\n case mod_json_type_string:\n return mod_json_string_float(val->data.c_str);\n\n default:\n break;\n }\n }\n return (0.0);\n}\n\nmod_json_boolean_t mod_json_value_boolean(mod_json_value_t *val) {\n if (val) {\n switch (val->type) {\n case mod_json_type_null:\n return MOD_JSON_FALSE;\n\n case mod_json_type_object:\n return (mod_json_object_count(val->data.c_obj) != 0);\n\n case mod_json_type_array:\n return (mod_json_array_count(val->data.c_arr) != 0);\n\n case mod_json_type_string:\n return (mod_json_string_length(val->data.c_str) != 0);\n\n case mod_json_type_integer:\n return (val->data.c_int != 0);\n\n case mod_json_type_float:\n return (val->data.c_float != 0);\n\n case mod_json_type_boolean:\n return (val->data.c_bool);\n\n default:\n break;\n }\n }\n return MOD_JSON_FALSE;\n}\n\nmod_json_integer_t mod_json_value_integer(mod_json_value_t *val) {\n if (val) {\n switch (val->type) {\n case mod_json_type_boolean:\n return (val->data.c_bool ? 1 : 0);\n\n case mod_json_type_integer:\n return (val->data.c_int);\n\n case mod_json_type_float:\n return (mod_json_integer_t)(val->data.c_float);\n\n case mod_json_type_string:\n return mod_json_string_integer(val->data.c_str);\n\n default:\n break;\n }\n }\n return (0);\n}\n\nmod_json_value_t *mod_json_value_clone(mod_json_value_t *val) {\n if (val) {\n switch (val->type) {\n case mod_json_type_null:\n return mod_json_value_set_null();\n\n case mod_json_type_object:\n return mod_json_value_set_object(val->data.c_obj);\n\n case mod_json_type_array:\n return mod_json_value_set_array(val->data.c_arr);\n\n case mod_json_type_string:\n return mod_json_value_set_string(val->data.c_str);\n\n case mod_json_type_integer:\n return mod_json_value_set_integer(val->data.c_int);\n\n case mod_json_type_float:\n return mod_json_value_set_float(val->data.c_float);\n\n case mod_json_type_boolean:\n return mod_json_value_set_boolean(val->data.c_bool);\n\n default:\n break;\n }\n }\n return NULL;\n}\n\nstatic inline mod_json_boolean_t mod_json_value_is_equal_float(\n mod_json_float_t lhs, mod_json_float_t rhs) {\n mod_json_float_t diff = lhs - rhs;\n return ((diff < DBL_EPSILON) && (diff > -DBL_EPSILON));\n}\n\nmod_json_boolean_t mod_json_value_is_equal(mod_json_value_t *lhs,\n mod_json_value_t *rhs) {\n if (lhs == rhs) {\n /* The same pointer */\n return MOD_JSON_TRUE;\n }\n\n if (lhs && rhs && lhs->type == rhs->type) {\n switch (lhs->type) {\n case mod_json_type_null:\n return MOD_JSON_TRUE;\n\n case mod_json_type_object:\n return mod_json_object_is_equal(lhs->data.c_obj, rhs->data.c_obj);\n\n case mod_json_type_array:\n return mod_json_array_is_equal(lhs->data.c_arr, rhs->data.c_arr);\n\n case mod_json_type_string:\n return (mod_json_string_compare(lhs->data.c_str, rhs->data.c_str) == 0);\n\n case mod_json_type_integer:\n return (lhs->data.c_int == rhs->data.c_int);\n\n case mod_json_type_float:\n return mod_json_value_is_equal_float(lhs->data.c_float,\n rhs->data.c_float);\n\n case mod_json_type_boolean:\n return ((!lhs->data.c_bool) == (!rhs->data.c_bool));\n\n default:\n break;\n }\n }\n return MOD_JSON_FALSE;\n}\n\nvoid mod_json_value_unset(mod_json_value_t *val) {\n if (val && mod_json_value_put(val) <= 0) {\n mod_json_value_clear(val);\n mod_json_free(val);\n }\n}\n\nstatic inline int mod_json_string_expand(mod_json_string_t *str,\n mod_json_size_t size) {\n mod_json_char_t *cstr;\n mod_json_size_t len;\n\n size = mod_json_utils_clp2(size);\n if (size < MOD_JSON_STRING_DEFSIZE) {\n size = MOD_JSON_STRING_DEFSIZE;\n }\n mod_json_minus_if_false(size > str->size);\n\n cstr = (mod_json_char_t *)mod_json_malloc(size * sizeof(mod_json_char_t));\n mod_json_minus_if_false(cstr);\n\n len = (mod_json_size_t)(str->last - str->first);\n if (len != 0) {\n memcpy(cstr, str->first, len + 1);\n } else {\n *cstr = '\\0'; /* terminal character */\n }\n mod_json_free(str->first);\n str->first = cstr;\n str->last = cstr + len;\n str->size = size;\n\n /* success */\n return 0;\n}\n\nint mod_json_string_reserve(mod_json_string_t *str, mod_json_size_t n) {\n mod_json_minus_if_false(str);\n\n if (str->size >= n + 1) {\n /* needn't grow */\n return 0;\n }\n return mod_json_string_expand(str, n + 1);\n}\n\nstatic inline mod_json_string_t *mod_json_string_malloc(mod_json_size_t size) {\n mod_json_string_t *str;\n mod_json_char_t *buf;\n\n buf = (mod_json_char_t *)mod_json_malloc(size * sizeof(mod_json_char_t));\n mod_json_null_if_false(buf);\n\n str = (mod_json_string_t *)mod_json_malloc(sizeof(mod_json_string_t));\n if (mod_json_unlikely(!str)) {\n mod_json_free(buf);\n return NULL;\n }\n\n str->refer = 1;\n str->size = size;\n str->first = buf;\n str->last = buf;\n *buf = '\\0';\n return str;\n}\n\nint mod_json_string_assign(mod_json_string_t *str, mod_json_cchar_t *cstr,\n mod_json_size_t len) {\n mod_json_string_reset(str);\n mod_json_minus_if_ne_zero(mod_json_string_reserve(str, len));\n\n if (cstr && len) {\n memcpy(str->first, cstr, len);\n }\n str->last = str->first + len;\n *(str->last) = '\\0';\n\n /* success */\n return 0;\n}\n\nstatic inline mod_json_string_t *mod_json_string_set_empty(void) {\n return mod_json_string_malloc(MOD_JSON_STRING_DEFSIZE);\n}\n\nstatic inline mod_json_string_t *mod_json_string_set_cstr(\n mod_json_cchar_t *cstr, mod_json_size_t len) {\n mod_json_string_t *str;\n\n str = mod_json_string_malloc(mod_json_utils_clp2(len + 1));\n mod_json_null_if_false(str);\n\n str->last = str->first + len;\n memcpy(str->first, cstr, len);\n *(str->last) = '\\0';\n return str;\n}\n\nmod_json_string_t *mod_json_string_set(mod_json_cchar_t *cstr,\n mod_json_size_t len) {\n return ((cstr && len) ? mod_json_string_set_cstr(cstr, len)\n : mod_json_string_set_empty());\n}\n\nvoid mod_json_string_unset(mod_json_string_t *str) {\n if (str && mod_json_string_put(str) <= 0) {\n mod_json_free(str->first);\n mod_json_free(str);\n }\n}\n\nvoid mod_json_string_reset(mod_json_string_t *str) {\n if (str) {\n str->last = str->first;\n *(str->first) = '\\0';\n }\n}\n\nstatic inline int mod_json_string_add_char(mod_json_string_t *str,\n mod_json_char_t ch) {\n mod_json_size_t need;\n\n need = (mod_json_size_t)(str->last - str->first) + 2;\n if (need > str->size) {\n mod_json_minus_if_ne_zero(mod_json_string_expand(str, need));\n }\n\n /* append to string */\n *(str->last++) = ch;\n *(str->last) = '\\0';\n\n /* success */\n return 0;\n}\n\nstatic inline int mod_json_string_add_cstr(mod_json_string_t *str,\n mod_json_cchar_t *cstr,\n mod_json_size_t len) {\n if (cstr && len) {\n mod_json_size_t need;\n\n need = len + (mod_json_size_t)(str->last - str->first) + 1;\n if (need > str->size) {\n mod_json_minus_if_ne_zero(mod_json_string_expand(str, need));\n }\n\n /* append to string */\n memcpy(str->last, cstr, len);\n str->last += len;\n *(str->last) = '\\0';\n }\n\n /* success */\n return 0;\n}\n\nstatic inline int mod_json_string_add_jstr(mod_json_string_t *str,\n mod_json_string_t *val) {\n return mod_json_string_add_cstr(str, val->first,\n (mod_json_size_t)(val->last - val->first));\n}\n\nint mod_json_string_add(mod_json_string_t *str, mod_json_string_t *val) {\n return mod_json_string_add_jstr(str, val);\n}\n\nint mod_json_string_append(mod_json_string_t *str, mod_json_cchar_t *cstr,\n mod_json_size_t len) {\n return mod_json_string_add_cstr(str, cstr, len);\n}\n\nmod_json_size_t mod_json_string_hash(mod_json_string_t *str) {\n mod_json_size_t hash = 1;\n\n if (str) {\n mod_json_cchar_t *iter = str->first;\n mod_json_cchar_t *last = str->last;\n\n for (; iter != last; ++iter) {\n mod_json_size_t c = (mod_json_size_t)(*iter);\n hash = hash * 131 + c;\n }\n }\n return hash;\n}\n\nint mod_json_string_compare(mod_json_string_t *str1, mod_json_string_t *str2) {\n mod_json_size_t len1 = 0, len2 = 0;\n\n if (str1 == str2) {\n /* The same pointer */\n return 0;\n }\n\n if (str1) {\n len1 = (mod_json_size_t)(str1->last - str1->first);\n if (str2) {\n len2 = (mod_json_size_t)(str2->last - str2->first);\n if (len1 == len2) {\n return memcmp(str1->first, str2->first, len1);\n }\n }\n } else {\n /* The first string is null, and the second string it not null. */\n len2 = (mod_json_size_t)(str2->last - str2->first);\n }\n return (int)(len1 - len2);\n}\n\nmod_json_integer_t mod_json_string_integer(mod_json_string_t *str) {\n return (str ? (mod_json_integer_t)mod_json_utils_strtoi(str->first, NULL, 0)\n : 0);\n}\n\nmod_json_float_t mod_json_string_float(mod_json_string_t *str) {\n return (str ? mod_json_utils_strtof(str->first, NULL) : 0.0);\n}\n\nstatic inline int mod_json_string_flat(mod_json_string_t *dst,\n mod_json_string_t *src) {\n static mod_json_cchar_t *flattab[32] = {\n \"\\\\u0000\", \"\\\\u0001\", \"\\\\u0002\", \"\\\\u0003\", \"\\\\u0004\", \"\\\\u0005\",\n \"\\\\u0006\", \"\\\\u0007\", \"\\\\b\", \"\\\\t\", \"\\\\n\", \"\\\\u000b\",\n \"\\\\f\", \"\\\\r\", \"\\\\u000e\", \"\\\\u000f\", \"\\\\u0010\", \"\\\\u0011\",\n \"\\\\u0012\", \"\\\\u0013\", \"\\\\u0014\", \"\\\\u0015\", \"\\\\u0016\", \"\\\\u0017\",\n \"\\\\u0018\", \"\\\\u0019\", \"\\\\u001a\", \"\\\\u001b\", \"\\\\u001c\", \"\\\\u001d\",\n \"\\\\u001e\", \"\\\\u001f\"};\n\n /* length of items in flat table */\n static const mod_json_uchar_t flatlen[32] = {6, 6, 6, 6, 6, 6, 6, 6, 2, 2, 2,\n 6, 2, 2, 6, 6, 6, 6, 6, 6, 6, 6,\n 6, 6, 6, 6, 6, 6, 6, 6, 6, 6};\n\n mod_json_cchar_t *first = src->first;\n mod_json_cchar_t *iter = src->first;\n mod_json_cchar_t *last = src->last;\n\n /* the whole string */\n for (; iter != last; ++iter) {\n int c = *iter;\n\n if ((mod_json_uchar_t)c <= 0x1f) {\n if (iter > first) {\n mod_json_minus_if_ne_zero(mod_json_string_add_cstr(\n dst, first, (mod_json_size_t)(iter - first)));\n }\n mod_json_minus_if_ne_zero(\n mod_json_string_add_cstr(dst, flattab[c], flatlen[c]));\n\n /* skip current character */\n first = iter + 1;\n } else if (c == '\\\"' || c == '\\\\') {\n if (iter > first) {\n mod_json_minus_if_ne_zero(mod_json_string_add_cstr(\n dst, first, (mod_json_size_t)(iter - first)));\n }\n mod_json_minus_if_ne_zero(mod_json_string_add_char(dst, '\\\\'));\n\n /* don't skip current character */\n first = iter;\n }\n }\n\n if (iter > first) {\n mod_json_minus_if_ne_zero(\n mod_json_string_add_cstr(dst, first, (mod_json_size_t)(iter - first)));\n }\n\n /* success */\n return 0;\n}\n\nstatic inline int mod_json_string_unflat(mod_json_string_t *dst,\n mod_json_string_t *src) {\n enum {\n state_normal,\n state_rev_slash,\n state_digit_1,\n state_digit_2,\n state_digit_3,\n state_digit_4\n } state;\n\n mod_json_char_t *pbuf = dst->first;\n mod_json_char_t *pend = dst->first + dst->size;\n mod_json_cchar_t *iter = src->first;\n mod_json_cchar_t *last = src->last;\n mod_json_uchar_t high = 0;\n mod_json_uchar_t low = 0;\n\n /* the whole string */\n for (state = state_normal; iter != last; ++iter) {\n int c = *iter;\n\n switch (state) {\n case state_normal:\n if (c != '\\\\') {\n mod_json_minus_if_false(pbuf < pend);\n *pbuf++ = (mod_json_char_t)c;\n } else {\n /* '\\\\' in process */\n state = state_rev_slash;\n }\n break;\n\n case state_rev_slash:\n mod_json_minus_if_false(pbuf < pend);\n\n switch (c) {\n case '\\\"':\n state = state_normal;\n *pbuf++ = '\\\"';\n break;\n case '/':\n state = state_normal;\n *pbuf++ = '/';\n break;\n case 'b':\n state = state_normal;\n *pbuf++ = '\\b';\n break;\n case 'f':\n state = state_normal;\n *pbuf++ = '\\f';\n break;\n case '\\\\':\n state = state_normal;\n *pbuf++ = '\\\\';\n break;\n case 'n':\n state = state_normal;\n *pbuf++ = '\\n';\n break;\n case 'r':\n state = state_normal;\n *pbuf++ = '\\r';\n break;\n case 't':\n state = state_normal;\n *pbuf++ = '\\t';\n break;\n case 'u':\n state = state_digit_1;\n break;\n default:\n return -1;\n }\n break;\n\n case state_digit_1:\n if ((c = mod_json_utils_char2hex((mod_json_char_t)c)) > 15) {\n /* invalid character */\n return -1;\n }\n high = (mod_json_uchar_t)(c << 4);\n state = state_digit_2;\n break;\n\n case state_digit_2:\n if ((c = mod_json_utils_char2hex((mod_json_char_t)c)) > 15) {\n /* invalid character */\n return -1;\n }\n high |= (mod_json_uchar_t)c;\n state = state_digit_3;\n break;\n\n case state_digit_3:\n if ((c = mod_json_utils_char2hex((mod_json_char_t)c)) > 15) {\n /* invalid character */\n return -1;\n }\n low = (mod_json_uchar_t)(c << 4);\n state = state_digit_4;\n break;\n\n case state_digit_4:\n if ((c = mod_json_utils_char2hex((mod_json_char_t)c)) > 15) {\n /* invalid character */\n return -1;\n }\n low |= (mod_json_uchar_t)c;\n\n /* decode as a UTF-8 string */\n pbuf = mod_json_utils_uni2utf8(pbuf, (mod_json_size_t)(pend - pbuf),\n high, low);\n if (!pbuf) {\n /* lack of buffer */\n return -1;\n }\n state = state_normal;\n break;\n }\n }\n\n if (state != state_normal) {\n /* uncompleted state */\n return -1;\n }\n mod_json_minus_if_false(pbuf < pend);\n\n /* update the last pointer */\n *(dst->last = pbuf) = '\\0';\n\n /* success */\n return 0;\n}\n\nmod_json_string_t *mod_json_string_encode(mod_json_string_t *src) {\n mod_json_string_t *dst;\n mod_json_null_if_false(src);\n\n dst = mod_json_string_malloc(\n mod_json_utils_clp2((mod_json_size_t)(src->last - src->first) + 1));\n mod_json_null_if_false(dst);\n\n if (mod_json_unlikely(mod_json_string_flat(dst, src) != 0)) {\n mod_json_string_unset(dst);\n return NULL;\n }\n return dst;\n}\n\nmod_json_string_t *mod_json_string_decode(mod_json_string_t *src) {\n mod_json_string_t *dst;\n mod_json_null_if_false(src);\n\n dst = mod_json_string_malloc(\n mod_json_utils_clp2((mod_json_size_t)(src->last - src->first) + 1));\n mod_json_null_if_false(dst);\n\n if (mod_json_unlikely(mod_json_string_unflat(dst, src) != 0)) {\n mod_json_string_unset(dst);\n return NULL;\n }\n return dst;\n}\n\nmod_json_array_t *mod_json_array_set(mod_json_size_t size) {\n mod_json_array_t *arr;\n mod_json_value_t **buf;\n\n size = (size ? mod_json_utils_clp2(size) : MOD_JSON_ARRAY_DEFSIZE);\n buf = (mod_json_value_t **)mod_json_malloc(size * sizeof(mod_json_value_t *));\n mod_json_null_if_false(buf);\n\n /* create an array */\n arr = (mod_json_array_t *)mod_json_malloc(sizeof(mod_json_array_t));\n if (mod_json_unlikely(!arr)) {\n mod_json_free(buf);\n return NULL;\n }\n\n arr->refer = 1;\n arr->size = size;\n arr->first = buf;\n arr->last = buf;\n return arr;\n}\n\nmod_json_array_t *mod_json_array_clone(mod_json_array_t *arr) {\n mod_json_array_t *arr2 = NULL;\n\n if (arr) {\n arr2 = mod_json_array_set((mod_json_size_t)(arr->last - arr->first));\n if (arr2) {\n mod_json_value_t **iter = arr->first;\n\n /* clone items */\n for (; iter != arr->last; ++iter) {\n *arr2->last++ = *iter ? mod_json_value_grab(*iter) : NULL;\n }\n }\n }\n return arr2;\n}\n\nmod_json_boolean_t mod_json_array_is_equal(mod_json_array_t *lhs,\n mod_json_array_t *rhs) {\n mod_json_value_t **itl, **itr;\n\n if (lhs == rhs) {\n return MOD_JSON_TRUE;\n }\n\n if (!lhs || !rhs || ((lhs->last - lhs->first) != (rhs->last - rhs->first))) {\n return MOD_JSON_FALSE;\n }\n\n /* compare items */\n for (itl = lhs->first, itr = rhs->first; itl != lhs->last; ++itl, ++itr) {\n if (!mod_json_value_is_equal(*itl, *itr)) {\n return MOD_JSON_FALSE;\n }\n }\n return MOD_JSON_TRUE;\n}\n\nvoid mod_json_array_unset(mod_json_array_t *arr) {\n if (arr && mod_json_array_put(arr) <= 0) {\n mod_json_value_t **iter = arr->first;\n\n for (; iter != arr->last; ++iter) {\n mod_json_value_unset(*iter);\n }\n mod_json_free(arr->first);\n mod_json_free(arr);\n }\n}\n\nvoid mod_json_array_reset(mod_json_array_t *arr) {\n if (arr) {\n mod_json_value_t **iter = arr->first;\n\n for (; iter != arr->last; ++iter) {\n mod_json_value_unset(*iter);\n }\n arr->last = arr->first;\n }\n}\n\nstatic inline void mod_json_array_migrate(mod_json_array_t *arr,\n mod_json_value_t **buf,\n mod_json_size_t size) {\n mod_json_size_t count = (mod_json_size_t)(arr->last - arr->first);\n if (count > 0) {\n memcpy(buf, arr->first, count * sizeof(mod_json_value_t *));\n }\n mod_json_free(arr->first);\n\n arr->first = buf;\n arr->last = buf + count;\n arr->size = size;\n}\n\nstatic inline int mod_json_array_expand(mod_json_array_t *arr,\n mod_json_size_t n) {\n mod_json_size_t size;\n mod_json_value_t **vals;\n\n size = mod_json_utils_clp2(n);\n if (size < MOD_JSON_ARRAY_DEFSIZE) {\n size = MOD_JSON_ARRAY_DEFSIZE;\n }\n mod_json_minus_if_false(size > arr->size);\n\n vals =\n (mod_json_value_t **)mod_json_malloc(size * sizeof(mod_json_value_t *));\n mod_json_minus_if_false(vals);\n\n /* use new buffer */\n mod_json_array_migrate(arr, vals, size);\n\n /* success */\n return 0;\n}\n\nint mod_json_array_reserve(mod_json_array_t *arr, mod_json_size_t n) {\n mod_json_minus_if_false(arr);\n\n if (arr->size >= n) {\n /* needn't grow */\n return 0;\n }\n return mod_json_array_expand(arr, n);\n}\n\nvoid mod_json_array_reverse(mod_json_array_t *arr) {\n if (arr) {\n mod_json_value_t **first = arr->first;\n mod_json_value_t **last = arr->last - 1;\n\n while (first < last) {\n mod_json_value_t *temp = *first;\n *first++ = *last;\n *last-- = temp;\n }\n }\n}\n\nint mod_json_array_push(mod_json_array_t *arr, mod_json_value_t *val) {\n mod_json_size_t count;\n mod_json_minus_if_false(arr);\n\n count = (mod_json_size_t)(arr->last - arr->first);\n if (count >= arr->size) {\n mod_json_minus_if_ne_zero(mod_json_array_expand(arr, count + 1));\n }\n\n *arr->last++ = val ? mod_json_value_grab(val) : NULL;\n return 0;\n}\n\nvoid mod_json_array_pop(mod_json_array_t *arr) {\n if (arr && arr->first != arr->last) {\n mod_json_value_unset(*(--arr->last));\n }\n}\n\nvoid mod_json_array_shift(mod_json_array_t *arr) {\n if (arr && arr->first != arr->last) {\n mod_json_value_t **it = arr->first;\n mod_json_value_t **last = --arr->last;\n\n mod_json_value_unset(*it++);\n for (; it <= last; ++it) {\n *(it - 1) = *it;\n }\n }\n}\n\nmod_json_value_t *mod_json_array_at(mod_json_array_t *arr, mod_json_size_t id) {\n if (arr && ((arr->first + id) < arr->last)) {\n return (arr->first[id]);\n }\n return NULL;\n}\n\nint mod_json_array_merge(mod_json_array_t *dst, mod_json_array_t *src) {\n long count, len1, len2;\n\n mod_json_minus_if_false(dst && src && dst != src);\n\n /* update length of array */\n len1 = (mod_json_size_t)(src->last - src->first);\n len2 = (mod_json_size_t)(dst->last - dst->first);\n mod_json_minus_if_false(len1 >= 0 && len2 >= 0);\n\n /* append empty values */\n count = len1 - len2;\n for (; count > 0; --count) {\n mod_json_array_push(dst, NULL);\n }\n\n /* It must be assigned again. */\n len2 = (mod_json_size_t)(dst->last - dst->first);\n count = (len1 < len2 ? len1 : len2);\n\n while ((count--) > 0) {\n mod_json_value_t **iter1 = src->first + count;\n mod_json_value_t **iter2 = dst->first + count;\n\n if (!(*iter2)) {\n *iter2 = *iter1 ? mod_json_value_grab(*iter1) : NULL;\n continue;\n }\n\n if (mod_json_value_is_shared(*iter2)) {\n mod_json_value_put(*iter2);\n *iter2 = mod_json_value_clone(*iter2);\n }\n mod_json_value_merge(*iter2, *iter1);\n }\n\n /* success */\n return 0;\n}\n\nint mod_json_array_resize(mod_json_array_t *arr, mod_json_size_t n,\n mod_json_value_t *val) {\n mod_json_size_t orig;\n\n /* check input */\n mod_json_minus_if_false(arr);\n\n /* original count of array */\n orig = (mod_json_size_t)(arr->last - arr->first);\n\n if (orig < n) {\n mod_json_value_t **iter;\n\n if (arr->size < n) {\n mod_json_minus_if_ne_zero(mod_json_array_expand(arr, n));\n }\n\n iter = arr->last;\n arr->last = arr->first + n;\n\n /* grab the first one, but get the others */\n *iter++ = val = val ? mod_json_value_grab(val) : NULL;\n for (; iter != arr->last; ++iter) {\n *iter = val ? mod_json_value_get(val) : NULL;\n }\n } else if (orig > n) {\n mod_json_value_t **iter = arr->first + n;\n\n for (; iter != arr->last; ++iter) {\n mod_json_value_unset(*iter);\n *iter = NULL;\n }\n arr->last = arr->first + n;\n }\n\n /* success */\n return 0;\n}\n\nstatic inline void mod_json_pair_init(mod_json_pair_t *pair,\n mod_json_string_t *key,\n mod_json_value_t *val) {\n pair->key = mod_json_string_grab(key);\n pair->val = val ? mod_json_value_grab(val) : NULL;\n}\n\nstatic inline void mod_json_pair_cleanup(mod_json_pair_t *pair) {\n mod_json_string_unset(pair->key);\n mod_json_value_unset(pair->val);\n pair->key = NULL;\n pair->val = NULL;\n}\n\nmod_json_object_t *mod_json_object_set(mod_json_size_t size) {\n mod_json_object_t *obj;\n mod_json_pair_t *buf;\n\n size = (size ? mod_json_utils_clp2(size) : MOD_JSON_OBJECT_DEFSIZE);\n buf = (mod_json_pair_t *)mod_json_malloc(size * sizeof(mod_json_pair_t));\n mod_json_null_if_false(buf);\n\n /* create a object */\n obj = (mod_json_object_t *)mod_json_malloc(sizeof(mod_json_object_t));\n if (mod_json_unlikely(!obj)) {\n mod_json_free(buf);\n return NULL;\n }\n\n obj->refer = 1;\n obj->size = size;\n obj->first = buf;\n obj->last = buf;\n return obj;\n}\n\nvoid mod_json_object_unset(mod_json_object_t *obj) {\n if (obj && mod_json_object_put(obj) <= 0) {\n mod_json_pair_t *iter = obj->first;\n\n for (; iter != obj->last; ++iter) {\n mod_json_pair_cleanup(iter);\n }\n mod_json_free(obj->first);\n mod_json_free(obj);\n }\n}\n\nvoid mod_json_object_reset(mod_json_object_t *obj) {\n if (obj) {\n mod_json_pair_t *iter = obj->first;\n\n for (; iter != obj->last; ++iter) {\n mod_json_pair_cleanup(iter);\n }\n obj->last = obj->first;\n }\n}\n\nstatic inline void mod_json_object_migrate(mod_json_object_t *obj,\n mod_json_pair_t *buf,\n mod_json_size_t size) {\n mod_json_size_t count = (mod_json_size_t)(obj->last - obj->first);\n if (count > 0) {\n memcpy(buf, obj->first, count * sizeof(mod_json_pair_t));\n }\n mod_json_free(obj->first);\n\n obj->first = buf;\n obj->last = buf + count;\n obj->size = size;\n}\n\nstatic inline int mod_json_object_expand(mod_json_object_t *obj,\n mod_json_size_t n) {\n mod_json_size_t size;\n mod_json_pair_t *buf;\n\n size = mod_json_utils_clp2(n);\n if (size < MOD_JSON_OBJECT_DEFSIZE) {\n size = MOD_JSON_OBJECT_DEFSIZE;\n }\n mod_json_minus_if_false(size > obj->size);\n\n buf = (mod_json_pair_t *)mod_json_malloc(size * sizeof(mod_json_pair_t));\n mod_json_minus_if_false(buf);\n\n /* use new buffer */\n mod_json_object_migrate(obj, buf, size);\n\n /* success */\n return 0;\n}\n\nstatic inline mod_json_pair_t *mod_json_object_find_pair(mod_json_object_t *obj,\n mod_json_string_t *key,\n mod_json_size_t *out) {\n mod_json_pair_t *first = obj->first;\n mod_json_pair_t *last = obj->last;\n\n while (first < last) {\n mod_json_pair_t *middle = first + ((last - first) >> 2);\n int diff = mod_json_string_compare(middle->key, key);\n\n if (diff < 0) {\n first = middle + 1;\n } else if (diff > 0) {\n last = middle;\n } else /*if (diff == 0)*/\n {\n *out = (mod_json_size_t)(middle - obj->first);\n return middle;\n }\n }\n *out = (mod_json_size_t)(first - obj->first);\n return NULL;\n}\n\nmod_json_pair_t *mod_json_object_insert_force(mod_json_object_t *obj,\n mod_json_size_t npos,\n mod_json_string_t *key,\n mod_json_value_t *val) {\n mod_json_pair_t *iter, *pos;\n mod_json_size_t count;\n\n count = (mod_json_size_t)(obj->last - obj->first);\n if (count >= obj->size) {\n mod_json_null_if_ne_zero(mod_json_object_expand(obj, count + 1));\n }\n\n pos = obj->first + npos;\n iter = obj->last++;\n for (; iter != pos; --iter) {\n mod_json_pair_t *prev = iter - 1;\n iter->key = prev->key;\n iter->val = prev->val;\n }\n mod_json_pair_init(pos, key, val);\n return pos;\n}\n\nmod_json_pair_t *mod_json_object_insert(mod_json_object_t *obj,\n mod_json_string_t *key,\n mod_json_value_t *val) {\n mod_json_size_t npos;\n mod_json_null_if_false(obj && key);\n\n if (mod_json_object_find_pair(obj, key, &npos)) {\n /* One in object */\n return NULL;\n }\n return mod_json_object_insert_force(obj, npos, key, val);\n}\n\nmod_json_pair_t *mod_json_object_assign(mod_json_object_t *obj,\n mod_json_string_t *key,\n mod_json_value_t *val) {\n mod_json_pair_t *elem = NULL;\n\n if (obj && key) {\n mod_json_size_t npos;\n\n elem = mod_json_object_find_pair(obj, key, &npos);\n if (elem) {\n if (!elem->val) {\n elem->val = val ? mod_json_value_grab(val) : NULL;\n } else {\n /* overwrite the old value */\n mod_json_value_assign(elem->val, val);\n }\n } else {\n /* insert a new one */\n elem = mod_json_object_insert_force(obj, npos, key, val);\n }\n }\n return elem;\n}\n\nmod_json_pair_t *mod_json_object_touch(mod_json_object_t *obj,\n mod_json_cchar_t *key) {\n mod_json_pair_t *elem = NULL;\n\n if (obj && key) {\n mod_json_string_t str;\n mod_json_size_t npos;\n\n str.first = (mod_json_char_t *)key;\n str.last = str.first + mod_json_utils_strlen(key);\n\n elem = mod_json_object_find_pair(obj, &str, &npos);\n if (!elem) {\n mod_json_string_t *jkey;\n\n /* insert a new one */\n jkey =\n mod_json_string_set(key, (mod_json_size_t)mod_json_utils_strlen(key));\n elem = mod_json_object_insert_force(obj, npos, jkey, NULL);\n mod_json_string_unset(jkey);\n }\n }\n return elem;\n}\n\nmod_json_object_t *mod_json_object_clone(mod_json_object_t *obj) {\n mod_json_object_t *obj2 = NULL;\n\n if (obj) {\n obj2 = mod_json_object_set((mod_json_size_t)(obj->last - obj->first));\n if (obj2) {\n mod_json_pair_t *iter = obj->first;\n\n /* clone items */\n for (; iter != obj->last; ++iter) {\n mod_json_pair_init(obj2->last++, iter->key, iter->val);\n }\n }\n }\n return obj2;\n}\n\nmod_json_boolean_t mod_json_object_is_equal(mod_json_object_t *lhs,\n mod_json_object_t *rhs) {\n mod_json_pair_t *itl, *itr;\n\n if (lhs == rhs) {\n /* The same pointer */\n return MOD_JSON_TRUE;\n }\n\n if (!lhs || !rhs || ((lhs->last - lhs->first) != (rhs->last - rhs->first))) {\n return MOD_JSON_FALSE;\n }\n\n /* compare items */\n for (itl = lhs->first, itr = rhs->first; itl != lhs->last; ++itl, ++itr) {\n if ((mod_json_string_compare(itl->key, itr->key) != 0) ||\n (!mod_json_value_is_equal(itl->val, itr->val))) {\n return MOD_JSON_FALSE;\n }\n }\n return MOD_JSON_TRUE;\n}\n\nvoid mod_json_object_erase(mod_json_object_t *obj, mod_json_cchar_t *key) {\n if (obj && key) {\n mod_json_string_t str;\n mod_json_pair_t *iter;\n mod_json_size_t npos;\n\n str.first = (mod_json_char_t *)key;\n str.last = str.first + mod_json_utils_strlen(key);\n\n iter = mod_json_object_find_pair(obj, &str, &npos);\n if (iter) {\n mod_json_pair_cleanup(iter++);\n\n for (; iter != obj->last; ++iter) {\n mod_json_pair_t *prev = iter - 1;\n prev->key = iter->key;\n prev->val = iter->val;\n }\n --obj->last;\n }\n }\n}\n\nmod_json_value_t *mod_json_object_at(mod_json_object_t *obj,\n mod_json_cchar_t *key) {\n if (obj && key) {\n mod_json_string_t str;\n mod_json_pair_t *elem;\n mod_json_size_t npos;\n\n str.first = (mod_json_char_t *)key;\n str.last = str.first + mod_json_utils_strlen(key);\n\n elem = mod_json_object_find_pair(obj, &str, &npos);\n if (elem) {\n return (elem->val);\n }\n }\n return NULL;\n}\n\nmod_json_pair_t *mod_json_object_find(mod_json_object_t *obj,\n mod_json_cchar_t *key) {\n if (obj && key) {\n mod_json_string_t str;\n mod_json_size_t npos;\n\n str.first = (mod_json_char_t *)key;\n str.last = str.first + mod_json_utils_strlen(key);\n\n return mod_json_object_find_pair(obj, &str, &npos);\n }\n return NULL;\n}\n\nint mod_json_object_merge(mod_json_object_t *dst, mod_json_object_t *src) {\n mod_json_pair_t *iter;\n\n mod_json_minus_if_false(dst && src && dst != src);\n\n for (iter = src->first; iter != src->last; ++iter) {\n mod_json_pair_t *elem;\n mod_json_size_t npos;\n\n elem = mod_json_object_find_pair(dst, iter->key, &npos);\n if (!elem) {\n /* insert a new one */\n mod_json_object_insert_force(dst, npos, iter->key, iter->val);\n continue;\n }\n\n if (!elem->val) {\n elem->val = iter->val ? mod_json_value_grab(iter->val) : NULL;\n continue;\n }\n\n if (mod_json_value_is_shared(elem->val)) {\n mod_json_value_put(elem->val);\n elem->val = mod_json_value_clone(elem->val);\n }\n mod_json_value_merge(elem->val, iter->val);\n }\n return 0;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_strskp(mod_json_token_t *tok,\n mod_json_cchar_t *cstr) {\n if ((tok->options & MOD_JSON_COMMENT) == 0) {\n return mod_json_utils_strskpb(cstr);\n }\n return mod_json_utils_strskp(cstr);\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_strfquo(mod_json_token_t *tok,\n mod_json_cchar_t *cstr,\n mod_json_char_t quo) {\n if ((tok->options & MOD_JSON_UNSTRICT) == 0) {\n return mod_json_utils_strfquo(cstr, quo);\n }\n return mod_json_utils_strfquo2(cstr, quo);\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_strfsep(mod_json_token_t *tok,\n mod_json_cchar_t *cstr) {\n if ((tok->options & MOD_JSON_COMMENT) == 0) {\n return mod_json_utils_strfsep(cstr);\n }\n return mod_json_utils_strfsep2(cstr);\n}\n\nmod_json_token_t *mod_json_token_create(mod_json_option_t *opt) {\n mod_json_token_t *tok;\n mod_json_size_t opts = MOD_JSON_TOKEN_DEFOPTS;\n mod_json_size_t mobj = MOD_JSON_TOKEN_DEFOBJDEP;\n mod_json_size_t marr = MOD_JSON_TOKEN_DEFARRDEP;\n\n if (opt) {\n opts = opt->options;\n\n if (opt->object_depth > 0) {\n mobj = opt->object_depth;\n }\n if (opt->array_depth > 0) {\n marr = opt->array_depth;\n }\n }\n\n tok = (mod_json_token_t *)mod_json_malloc(\n (mobj + marr) * sizeof(mod_json_char_t) + sizeof(mod_json_token_t));\n mod_json_null_if_false(tok);\n\n memset(tok, 0, sizeof(mod_json_token_t));\n tok->state = mod_json_state_null;\n tok->error = mod_json_error_null;\n tok->options = opts;\n tok->object_max_depth = mobj;\n tok->array_max_depth = marr;\n return tok;\n}\n\nvoid mod_json_token_destroy(mod_json_token_t *tok) {\n mod_json_free(tok);\n}\n\nstatic inline void mod_json_token_set_tag(mod_json_token_t *tok,\n mod_json_char_t tag) {\n mod_json_size_t depth = tok->object_depth + tok->array_depth;\n if (depth != 0) {\n tok->tags[depth - 1] = tag;\n }\n}\n\nstatic inline mod_json_char_t mod_json_token_tag(mod_json_token_t *tok) {\n mod_json_size_t depth = tok->object_depth + tok->array_depth;\n\n /* type of current depth */\n return (depth ? tok->tags[depth - 1] : (mod_json_char_t)-1);\n}\n\nmod_json_error_t mod_json_token_error(mod_json_token_t *tok) {\n return (tok->error);\n}\n\nmod_json_cchar_t *mod_json_token_context(mod_json_token_t *tok) {\n return (tok->context);\n}\n\nmod_json_state_t mod_json_token_state(mod_json_token_t *tok) {\n return (tok->state);\n}\n\nmod_json_size_t mod_json_token_object_depth(mod_json_token_t *tok) {\n return (tok->object_depth);\n}\n\nmod_json_size_t mod_json_token_array_depth(mod_json_token_t *tok) {\n return (tok->array_depth);\n}\n\nmod_json_size_t mod_json_token_depth(mod_json_token_t *tok) {\n return (tok->object_depth + tok->array_depth);\n}\n\nmod_json_size_t mod_json_token_max_object_depth(mod_json_token_t *tok) {\n return (tok->object_max_depth);\n}\n\nmod_json_size_t mod_json_token_max_array_depth(mod_json_token_t *tok) {\n return (tok->array_max_depth);\n}\n\nmod_json_size_t mod_json_token_max_depth(mod_json_token_t *tok) {\n return (tok->object_max_depth + tok->array_max_depth);\n}\n\nmod_json_void_t *mod_json_token_param(mod_json_token_t *tok) {\n return (tok->param);\n}\n\nvoid mod_json_token_set_param(mod_json_token_t *tok, mod_json_void_t *param) {\n tok->param = param;\n}\n\nvoid mod_json_token_set_event(mod_json_token_t *tok, mod_json_event_proc proc) {\n tok->event_proc = proc;\n}\n\nmod_json_event_t mod_json_token_event(mod_json_token_t *tok) {\n return (tok->event_code);\n}\n\nstatic inline int mod_json_token_invoke_field(mod_json_token_t *tok,\n mod_json_cchar_t *val,\n mod_json_size_t len) {\n mod_json_event_proc invoke = tok->event_proc;\n if (invoke) {\n tok->event_code = mod_json_event_field;\n return invoke(tok, (mod_json_void_t *)val, len);\n }\n return 0;\n}\n\nstatic inline int mod_json_token_invoke_object(mod_json_token_t *tok) {\n mod_json_event_proc invoke = tok->event_proc;\n if (invoke) {\n tok->event_code = mod_json_event_object;\n return invoke(tok, NULL, 0);\n }\n return 0;\n}\n\nstatic inline int mod_json_token_invoke_array(mod_json_token_t *tok) {\n mod_json_event_proc invoke = tok->event_proc;\n if (invoke) {\n tok->event_code = mod_json_event_array;\n return invoke(tok, NULL, 0);\n }\n return 0;\n}\n\nstatic inline int mod_json_token_invoke_null(mod_json_token_t *tok) {\n mod_json_event_proc invoke = tok->event_proc;\n if (invoke) {\n tok->event_code = mod_json_event_null;\n return invoke(tok, NULL, 0);\n }\n return 0;\n}\n\nstatic inline int mod_json_token_invoke_boolean(mod_json_token_t *tok,\n mod_json_boolean_t val) {\n mod_json_event_proc invoke = tok->event_proc;\n if (invoke) {\n tok->event_code = mod_json_event_boolean;\n return invoke(tok, &val, sizeof(val));\n }\n return 0;\n}\n\nstatic inline int mod_json_token_invoke_integer(mod_json_token_t *tok,\n mod_json_integer_t val) {\n mod_json_event_proc invoke = tok->event_proc;\n if (invoke) {\n tok->event_code = mod_json_event_integer;\n return invoke(tok, &val, sizeof(val));\n }\n return 0;\n}\n\nstatic inline int mod_json_token_invoke_float(mod_json_token_t *tok,\n mod_json_float_t val) {\n mod_json_event_proc invoke = tok->event_proc;\n if (invoke) {\n tok->event_code = mod_json_event_float;\n return invoke(tok, &val, sizeof(val));\n }\n return 0;\n}\n\nstatic inline int mod_json_token_invoke_string(mod_json_token_t *tok,\n mod_json_cchar_t *val,\n mod_json_size_t len) {\n mod_json_event_proc invoke = tok->event_proc;\n if (invoke) {\n tok->event_code = mod_json_event_string;\n return invoke(tok, (mod_json_void_t *)val, len);\n }\n return 0;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_start(mod_json_token_t *tok,\n mod_json_cchar_t *cstr) {\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case '{':\n tok->state = mod_json_state_object_start;\n return (cstr + 1);\n\n case '[':\n tok->state = mod_json_state_array_start;\n return (cstr + 1);\n\n case '\\0':\n tok->error = mod_json_error_empty;\n tok->context = cstr;\n break;\n\n default:\n tok->error = mod_json_error_start;\n tok->context = cstr;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_value_null(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n mod_json_char_t c1 = *(cstr + 1);\n mod_json_char_t c2 = *(cstr + 2);\n mod_json_char_t c3 = *(cstr + 3);\n\n if ((c1 != 'u' && c1 != 'U') || (c2 != 'l' && c2 != 'L') ||\n (c3 != 'l' && c3 != 'L')) {\n tok->error = mod_json_error_value;\n tok->context = cstr;\n return NULL;\n }\n\n if (mod_json_token_invoke_null(tok) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n return (cstr + 4);\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_value_true(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n mod_json_char_t c1 = *(cstr + 1);\n mod_json_char_t c2 = *(cstr + 2);\n mod_json_char_t c3 = *(cstr + 3);\n\n if ((c1 != 'r' && c1 != 'R') || (c2 != 'u' && c2 != 'U') ||\n (c3 != 'e' && c3 != 'E')) {\n tok->error = mod_json_error_value;\n tok->context = cstr;\n return NULL;\n }\n\n if (mod_json_token_invoke_boolean(tok, MOD_JSON_TRUE) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n return (cstr + 4);\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_value_false(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n mod_json_char_t c1 = *(cstr + 1);\n mod_json_char_t c2 = *(cstr + 2);\n mod_json_char_t c3 = *(cstr + 3);\n mod_json_char_t c4 = *(cstr + 4);\n\n if ((c1 != 'a' && c1 != 'A') || (c2 != 'l' && c2 != 'L') ||\n (c3 != 's' && c3 != 'S') || (c4 != 'e' && c4 != 'E')) {\n tok->error = mod_json_error_value;\n tok->context = cstr;\n return NULL;\n }\n\n if (mod_json_token_invoke_boolean(tok, MOD_JSON_FALSE) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n return (cstr + 5);\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_value_infinity(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n mod_json_char_t c1 = *(cstr + 1);\n mod_json_char_t c2 = *(cstr + 2);\n\n if ((c1 != 'n' && c1 != 'N') || (c2 != 'f' && c2 != 'F')) {\n tok->error = mod_json_error_value;\n tok->context = cstr;\n return NULL;\n }\n\n if (mod_json_token_invoke_float(tok, MOD_JSON_INFINITY) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n return (cstr + 3);\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_value_string(\n mod_json_token_t *tok, mod_json_cchar_t *cstr, mod_json_char_t quo) {\n mod_json_cchar_t *cstr2 = mod_json_token_strfquo(tok, ++cstr, quo);\n if (!cstr2) {\n tok->error = mod_json_error_quote;\n tok->context = cstr;\n return NULL;\n }\n\n if (mod_json_token_invoke_string(tok, cstr,\n (mod_json_size_t)(cstr2 - cstr)) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n return (cstr2 + 1);\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_value_number(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n enum { number_integer, number_float } num_type = number_integer;\n\n mod_json_float_t dbl = 0.0;\n uint32_t dig = 0;\n uint64_t u64 = 0;\n int32_t minus = 0;\n int32_t exp_frac = 0, exp = 0;\n\n /* Parse minus */\n minus = *cstr;\n if (minus == '-' || minus == '+') {\n ++cstr;\n }\n\n /* The first digit */\n if ((dig = (uint32_t)(*cstr - '0')) > 9) {\n return NULL;\n }\n\n /* Save the first digit */\n u64 = dig;\n\n /* Parse as 64bit integer */\n if (minus != '-') {\n while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) {\n if (u64 >= 1844674407370955161uLL) {\n /* 2^64 - 1 = 18446744073709551615 */\n if (u64 != 1844674407370955161uLL || dig > 5) {\n dbl = (mod_json_float_t)u64 * 10 + dig;\n num_type = number_float;\n break;\n }\n }\n u64 = u64 * 10 + dig;\n }\n } else {\n while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) {\n /* 2^63 = 9223372036854775808 */\n if (u64 >= 922337203685477580uLL) {\n if (u64 != 922337203685477580uLL || dig > 8) {\n dbl = (mod_json_float_t)u64 * 10 + dig;\n num_type = number_float;\n break;\n }\n }\n u64 = u64 * 10 + dig;\n }\n }\n\n /* Force double for big integer */\n if (num_type == number_float) {\n while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) {\n if (dbl >= 1E307) {\n /* Number too big to store in double */\n return NULL;\n }\n dbl = dbl * 10 + dig;\n }\n }\n\n /* Parse frac = decimal-point 1*DIGIT */\n if (*cstr == '.') {\n if (num_type != number_float) {\n dbl = (mod_json_float_t)u64;\n num_type = number_float;\n }\n\n if ((dig = (uint32_t)(*(++cstr) - '0')) > 9) {\n /* At least one digit in fraction part */\n return NULL;\n }\n\n dbl = dbl * 10 + dig;\n --exp_frac;\n\n while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) {\n if (exp_frac > -16) {\n dbl = dbl * 10 + dig;\n --exp_frac;\n }\n }\n }\n\n /* Parse exp = e [ minus / plus ] 1*DIGIT */\n if (*cstr == 'e' || *cstr == 'E') {\n int32_t exp_minus = 0;\n\n if (num_type != number_float) {\n dbl = (mod_json_float_t)u64;\n num_type = number_float;\n }\n\n exp_minus = *(++cstr);\n if (exp_minus == '-' || exp_minus == '+') {\n ++cstr;\n }\n\n /* The first number char after 'e/E' */\n if ((dig = (uint32_t)(*cstr - '0')) > 9) {\n return NULL;\n }\n exp = (int32_t)dig;\n\n while ((dig = (uint32_t)(*(++cstr) - '0')) <= 9) {\n exp = exp * 10 + (int32_t)dig;\n if (exp > 308) {\n /* Number too big to store in double */\n return NULL;\n }\n }\n\n if (exp_minus == '-') {\n exp = -exp;\n }\n }\n\n /* Finish parsing, call event according to the type of number. */\n if (num_type == number_float) {\n dbl *= mod_json_utils_pow10(exp + exp_frac);\n if (minus == '-') {\n dbl = -dbl;\n }\n if (mod_json_token_invoke_float(tok, dbl) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n } else {\n if (minus == '-') {\n u64 = (uint64_t)(-(int64_t)u64);\n }\n if (mod_json_token_invoke_integer(tok, (mod_json_integer_t)u64) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n }\n return cstr;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_array_start(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n if (tok->array_depth < tok->array_max_depth) {\n /* callback */\n if (mod_json_token_invoke_array(tok) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n\n /* increase depth */\n ++tok->array_depth;\n\n /* push current tag */\n mod_json_token_set_tag(tok, '[');\n\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case '[':\n tok->state = mod_json_state_array_start;\n return (cstr + 1);\n\n case ']':\n tok->state = mod_json_state_array_finish;\n return (cstr + 1);\n\n case '\\0':\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n break;\n\n default:\n tok->state = mod_json_state_array_half;\n return (cstr);\n }\n } else {\n tok->error = mod_json_error_depth;\n tok->context = cstr;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_array_half(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case ',':\n tok->state = mod_json_state_array_half;\n return (cstr + 1);\n\n case '[':\n tok->state = mod_json_state_array_start;\n return (cstr + 1);\n\n case ']':\n tok->state = mod_json_state_array_finish;\n return (cstr + 1);\n\n case '{':\n tok->state = mod_json_state_object_start;\n return (cstr + 1);\n\n case '\\0':\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n return NULL;\n\n /* value in array */\n case 't':\n case 'T':\n cstr = mod_json_token_value_true(tok, cstr);\n if (!cstr) {\n return NULL;\n }\n break;\n\n case 'f':\n case 'F':\n cstr = mod_json_token_value_false(tok, cstr);\n if (!cstr) {\n return NULL;\n }\n break;\n\n case 'n':\n case 'N':\n cstr = mod_json_token_value_null(tok, cstr);\n if (!cstr) {\n return NULL;\n }\n break;\n\n case '0':\n case '1':\n case '2':\n case '3':\n case '4':\n case '5':\n case '6':\n case '7':\n case '8':\n case '9':\n case '+':\n case '-':\n cstr = mod_json_token_value_number(tok, cstr);\n if (!cstr) {\n return NULL;\n }\n break;\n\n case '\\\"':\n cstr = mod_json_token_value_string(tok, cstr, '\\\"');\n if (!cstr) {\n return NULL;\n }\n break;\n\n case '\\'':\n if (tok->options & MOD_JSON_SQUOTE) {\n cstr = mod_json_token_value_string(tok, cstr, '\\'');\n if (!cstr) {\n return NULL;\n }\n break;\n }\n /* FALLTHRU */\n\n default:\n tok->error = mod_json_error_value;\n tok->context = cstr;\n return NULL;\n }\n\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case ',':\n tok->state = mod_json_state_array_half;\n return (cstr + 1);\n\n case ']':\n tok->state = mod_json_state_array_finish;\n return (cstr + 1);\n\n case '\\0':\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n break;\n\n default:\n tok->error = mod_json_error_value;\n tok->context = cstr;\n break;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_array_finish(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n if (tok->array_depth) {\n /* decrease depth */\n --tok->array_depth;\n\n /* callback */\n if (mod_json_token_invoke_array(tok) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case ']':\n tok->state = mod_json_state_array_finish;\n return (cstr + 1);\n\n case '}':\n tok->state = mod_json_state_object_finish;\n return (cstr + 1);\n\n case '\\0':\n if (tok->object_depth || tok->array_depth) {\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n } else {\n tok->state = mod_json_state_finish;\n }\n break;\n\n case ',':\n if (tok->object_depth || tok->array_depth) {\n mod_json_char_t tag = mod_json_token_tag(tok);\n\n if (tag == '{') {\n tok->state = mod_json_state_object_half1;\n return (cstr + 1);\n } else if (tag == '[') {\n tok->state = mod_json_state_array_half;\n return (cstr + 1);\n }\n }\n /* FALLTHRU */\n\n default:\n tok->error = mod_json_error_array;\n tok->context = cstr;\n }\n } else {\n tok->error = mod_json_error_depth;\n tok->context = cstr;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_object_start(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n if (tok->object_depth < tok->object_max_depth) {\n /* callback */\n if (mod_json_token_invoke_object(tok) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n\n /* increase depth */\n ++tok->object_depth;\n\n /* push current tag */\n mod_json_token_set_tag(tok, '{');\n\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case '}':\n tok->state = mod_json_state_object_finish;\n return (cstr + 1);\n\n case '\\0':\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n break;\n\n default:\n tok->state = mod_json_state_object_half1;\n return (cstr);\n }\n } else {\n tok->error = mod_json_error_depth;\n tok->context = cstr;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_object_quotekey(\n mod_json_token_t *tok, mod_json_cchar_t *cstr, mod_json_char_t quo) {\n mod_json_cchar_t *cstr2 = mod_json_token_strfquo(tok, ++cstr, quo);\n if (cstr2) {\n /* callback */\n if (mod_json_token_invoke_field(tok, cstr,\n (mod_json_size_t)(cstr2 - cstr)) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n\n cstr2 = mod_json_token_strskp(tok, ++cstr2);\n switch (*cstr2) {\n case ':':\n tok->state = mod_json_state_object_half2;\n return (cstr2 + 1);\n\n case '\\0':\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n break;\n\n default:\n tok->error = mod_json_error_key;\n tok->context = cstr2;\n break;\n }\n } else {\n tok->error = mod_json_error_quote;\n tok->context = cstr;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_object_simplekey(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n mod_json_cchar_t *cstr2 = mod_json_token_strfsep(tok, cstr);\n if (cstr2 != cstr) {\n /* callback */\n if (mod_json_token_invoke_field(tok, cstr,\n (mod_json_size_t)(cstr2 - cstr)) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n\n cstr2 = mod_json_token_strskp(tok, cstr2);\n switch (*cstr2) {\n case ':':\n tok->state = mod_json_state_object_half2;\n return (cstr2 + 1);\n\n case '\\0':\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n break;\n\n default:\n tok->error = mod_json_error_key;\n tok->context = cstr2;\n break;\n }\n } else {\n tok->error = mod_json_error_key;\n tok->context = cstr;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_object_half1(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case ',':\n tok->state = mod_json_state_object_half1;\n return (cstr + 1);\n\n case '}':\n tok->state = mod_json_state_object_finish;\n return (cstr + 1);\n\n case '\\0':\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n break;\n\n case '\\\"':\n /* The key with double quotes */\n return mod_json_token_object_quotekey(tok, cstr, '\\\"');\n\n case '\\'':\n if (tok->options & MOD_JSON_SQUOTE) {\n /* The key with single quotes */\n return mod_json_token_object_quotekey(tok, cstr, '\\'');\n }\n /* FALLTHRU */\n\n default:\n /* support simple format? */\n if (tok->options & MOD_JSON_SIMPLE) {\n return mod_json_token_object_simplekey(tok, cstr);\n } else {\n tok->error = mod_json_error_quote;\n tok->context = cstr;\n }\n break;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_object_half2(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case '{':\n tok->state = mod_json_state_object_start;\n return (cstr + 1);\n\n case '[':\n tok->state = mod_json_state_array_start;\n return (cstr + 1);\n\n case ',':\n tok->state = mod_json_state_object_half1;\n return (cstr + 1);\n\n case '}':\n tok->state = mod_json_state_object_finish;\n return (cstr + 1);\n\n case '\\0':\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n return NULL;\n\n case 't':\n case 'T':\n cstr = mod_json_token_value_true(tok, cstr);\n if (!cstr) {\n return NULL;\n }\n break;\n\n case 'f':\n case 'F':\n cstr = mod_json_token_value_false(tok, cstr);\n if (!cstr) {\n return NULL;\n }\n break;\n\n case 'i':\n case 'I':\n cstr = mod_json_token_value_infinity(tok, cstr);\n if (!cstr) {\n return NULL;\n }\n break;\n\n case 'n':\n case 'N':\n cstr = mod_json_token_value_null(tok, cstr);\n if (!cstr) {\n return NULL;\n }\n break;\n\n case '0':\n case '1':\n case '2':\n case '3':\n case '4':\n case '5':\n case '6':\n case '7':\n case '8':\n case '9':\n case '+':\n case '-':\n cstr = mod_json_token_value_number(tok, cstr);\n if (!cstr) {\n return NULL;\n }\n break;\n\n case '\\\"':\n cstr = mod_json_token_value_string(tok, cstr, '\\\"');\n if (!cstr) {\n return NULL;\n }\n break;\n\n case '\\'':\n if (tok->options & MOD_JSON_SQUOTE) {\n cstr = mod_json_token_value_string(tok, cstr, '\\'');\n if (!cstr) {\n return NULL;\n }\n break;\n }\n /* FALLTHRU */\n\n default:\n tok->error = mod_json_error_value;\n tok->context = cstr;\n return NULL;\n }\n\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case ',':\n tok->state = mod_json_state_object_half1;\n return (cstr + 1);\n\n case '}':\n tok->state = mod_json_state_object_finish;\n return (cstr + 1);\n\n case '\\0':\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n break;\n\n default:\n tok->error = mod_json_error_value;\n tok->context = cstr;\n break;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_object_finish(\n mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n if (tok->object_depth) {\n /* decrease depth */\n --tok->object_depth;\n\n /* callback */\n if (mod_json_token_invoke_object(tok) != 0) {\n tok->error = mod_json_error_break;\n tok->context = cstr;\n return NULL;\n }\n\n cstr = mod_json_token_strskp(tok, cstr);\n switch (*cstr) {\n case '}':\n tok->state = mod_json_state_object_finish;\n return (cstr + 1);\n\n case ']':\n tok->state = mod_json_state_array_finish;\n return (cstr + 1);\n\n case '\\0':\n if (tok->object_depth || tok->array_depth) {\n tok->error = mod_json_error_trunc;\n tok->context = cstr;\n } else {\n tok->state = mod_json_state_finish;\n }\n break;\n\n case ',':\n if (tok->object_depth || tok->array_depth) {\n mod_json_char_t tag = mod_json_token_tag(tok);\n\n if (tag == '{') {\n tok->state = mod_json_state_object_half1;\n return (cstr + 1);\n } else if (tag == '[') {\n tok->state = mod_json_state_array_half;\n return (cstr + 1);\n }\n }\n /* FALLTHRU */\n\n default:\n tok->error = mod_json_error_object;\n tok->context = cstr;\n }\n } else {\n tok->error = mod_json_error_depth;\n tok->context = cstr;\n }\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_null(mod_json_token_t *tok,\n mod_json_cchar_t *cstr) {\n if (!cstr || *cstr == '\\0') {\n tok->error = mod_json_error_invalid;\n tok->context = cstr;\n return NULL;\n }\n\n tok->state = mod_json_state_start;\n return cstr;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_finish(mod_json_token_t *tok,\n mod_json_cchar_t *cstr) {\n tok->error = mod_json_error_null;\n (void)cstr;\n return NULL;\n}\n\nstatic inline mod_json_cchar_t *mod_json_token_default(mod_json_token_t *tok,\n mod_json_cchar_t *cstr) {\n tok->error = mod_json_error_state;\n tok->context = cstr;\n return NULL;\n}\n\nint mod_json_token_parse(mod_json_token_t *tok, mod_json_cchar_t *cstr) {\n while (cstr) {\n switch (tok->state) {\n case mod_json_state_start:\n cstr = mod_json_token_start(tok, cstr);\n break;\n\n case mod_json_state_array_start:\n cstr = mod_json_token_array_start(tok, cstr);\n break;\n\n case mod_json_state_array_half:\n cstr = mod_json_token_array_half(tok, cstr);\n break;\n\n case mod_json_state_array_finish:\n cstr = mod_json_token_array_finish(tok, cstr);\n break;\n\n case mod_json_state_object_start:\n cstr = mod_json_token_object_start(tok, cstr);\n break;\n\n case mod_json_state_object_half1:\n cstr = mod_json_token_object_half1(tok, cstr);\n break;\n\n case mod_json_state_object_half2:\n cstr = mod_json_token_object_half2(tok, cstr);\n break;\n\n case mod_json_state_object_finish:\n cstr = mod_json_token_object_finish(tok, cstr);\n break;\n\n case mod_json_state_null:\n cstr = mod_json_token_null(tok, cstr);\n break;\n\n case mod_json_state_finish:\n cstr = mod_json_token_finish(tok, cstr);\n break;\n\n default:\n cstr = mod_json_token_default(tok, cstr);\n break;\n }\n }\n return (tok->error == mod_json_error_null ? 0 : -1);\n}\n\nstatic inline int mod_json_parser_insert(mod_json_parser_t *par,\n mod_json_size_t depth,\n mod_json_value_t *val) {\n if (depth > 0) {\n mod_json_value_t *cur = par->vals[depth - 1];\n\n switch (cur->type) {\n case mod_json_type_object:\n return (mod_json_object_insert(cur->data.c_obj, par->key, val) ? 0\n : -1);\n\n case mod_json_type_array:\n return mod_json_array_push(cur->data.c_arr, val);\n\n default:\n break;\n }\n }\n return -1;\n}\n\nstatic inline int mod_json_parser_insert_object(mod_json_parser_t *par,\n mod_json_size_t depth) {\n mod_json_object_t *obj;\n mod_json_value_t *jval;\n\n obj = mod_json_object_set_default();\n mod_json_minus_if_false(obj);\n\n jval = mod_json_value_set_object(obj);\n mod_json_object_unset(obj);\n mod_json_minus_if_false(jval);\n\n if (depth > 0) {\n int ret = mod_json_parser_insert(par, depth, jval);\n if (ret == 0) {\n par->vals[depth] = jval;\n }\n mod_json_value_unset(jval);\n return ret;\n } else {\n /* It's the root, save the pointer. Don't unset it. */\n par->vals[0] = jval;\n }\n return 0;\n}\n\nstatic inline int mod_json_parser_insert_array(mod_json_parser_t *par,\n mod_json_size_t depth) {\n mod_json_array_t *arr;\n mod_json_value_t *jval;\n\n arr = mod_json_array_set_default();\n mod_json_minus_if_false(arr);\n\n jval = mod_json_value_set_array(arr);\n mod_json_array_unset(arr);\n mod_json_minus_if_false(jval);\n\n if (depth > 0) {\n int ret = mod_json_parser_insert(par, depth, jval);\n if (ret == 0) {\n par->vals[depth] = jval;\n }\n mod_json_value_unset(jval);\n return ret;\n } else {\n /* It's the root, save the pointer. Don't unset it. */\n par->vals[0] = jval;\n }\n return 0;\n}\n\nstatic inline void mod_json_token_set_parser(mod_json_token_t *tok,\n mod_json_parser_t *par) {\n mod_json_token_set_param(tok, par);\n}\n\nstatic inline mod_json_parser_t *mod_json_token_parser(mod_json_token_t *tok) {\n return (mod_json_parser_t *)mod_json_token_param(tok);\n}\n\nstatic inline int mod_json_parser_event_field(mod_json_token_t *tok,\n mod_json_cchar_t *val,\n mod_json_size_t len) {\n mod_json_parser_t *parser;\n\n /* get information */\n parser = mod_json_token_parser(tok);\n\n /* unset previous one */\n mod_json_string_unset(parser->key);\n\n parser->key = mod_json_string_set(val, len);\n return (parser->key ? 0 : -1);\n}\n\nstatic inline int mod_json_parser_event_array(mod_json_token_t *tok) {\n switch (mod_json_token_state(tok)) {\n case mod_json_state_array_finish:\n /* continue */\n return 0;\n\n case mod_json_state_array_start:\n return mod_json_parser_insert_array(mod_json_token_parser(tok),\n mod_json_token_depth(tok));\n\n default:\n break;\n }\n return -1;\n}\n\nstatic inline int mod_json_parser_event_object(mod_json_token_t *tok) {\n switch (mod_json_token_state(tok)) {\n case mod_json_state_object_finish:\n /* continue */\n return 0;\n\n case mod_json_state_object_start:\n return mod_json_parser_insert_object(mod_json_token_parser(tok),\n mod_json_token_depth(tok));\n\n default:\n break;\n }\n return -1;\n}\n\nstatic inline int mod_json_parser_event_null(mod_json_token_t *tok) {\n mod_json_parser_t *parser;\n\n /* get information */\n parser = mod_json_token_parser(tok);\n\n if (!parser->val_null) {\n parser->val_null = mod_json_value_set_null();\n mod_json_minus_if_false(parser->val_null);\n }\n return mod_json_parser_insert(parser, mod_json_token_depth(tok),\n parser->val_null);\n}\n\nstatic inline int mod_json_parser_event_true(mod_json_token_t *tok) {\n mod_json_parser_t *parser;\n\n /* get information */\n parser = mod_json_token_parser(tok);\n\n if (!parser->val_true) {\n parser->val_true = mod_json_value_set_boolean(MOD_JSON_TRUE);\n mod_json_minus_if_false(parser->val_true);\n }\n return mod_json_parser_insert(parser, mod_json_token_depth(tok),\n parser->val_true);\n}\n\nstatic inline int mod_json_parser_event_false(mod_json_token_t *tok) {\n mod_json_parser_t *parser;\n\n /* get information */\n parser = mod_json_token_parser(tok);\n\n if (!parser->val_false) {\n parser->val_false = mod_json_value_set_boolean(MOD_JSON_FALSE);\n mod_json_minus_if_false(parser->val_false);\n }\n return mod_json_parser_insert(parser, mod_json_token_depth(tok),\n parser->val_false);\n}\n\nstatic inline int mod_json_parser_event_boolean(mod_json_token_t *tok,\n mod_json_boolean_t val) {\n if (!val) {\n return mod_json_parser_event_false(tok);\n }\n return mod_json_parser_event_true(tok);\n}\n\nstatic inline int mod_json_parser_event_zero(mod_json_token_t *tok) {\n mod_json_parser_t *parser;\n\n /* get information */\n parser = mod_json_token_parser(tok);\n\n if (!parser->val_zero) {\n parser->val_zero = mod_json_value_set_integer(0);\n mod_json_minus_if_false(parser->val_zero);\n }\n return mod_json_parser_insert(parser, mod_json_token_depth(tok),\n parser->val_zero);\n}\n\nstatic inline int mod_json_parser_event_integer(mod_json_token_t *tok,\n mod_json_integer_t val) {\n int ret = -1;\n\n if (val != 0) {\n mod_json_value_t *jval;\n\n jval = mod_json_value_set_integer(val);\n if (jval) {\n ret = mod_json_parser_insert(mod_json_token_parser(tok),\n mod_json_token_depth(tok), jval);\n mod_json_value_unset(jval);\n }\n } else {\n /* zero event */\n ret = mod_json_parser_event_zero(tok);\n }\n return ret;\n}\n\nstatic inline int mod_json_parser_event_zerof(mod_json_token_t *tok) {\n mod_json_parser_t *parser;\n\n /* get information */\n parser = mod_json_token_parser(tok);\n\n if (!parser->val_zerof) {\n parser->val_zerof = mod_json_value_set_float(0.0);\n mod_json_minus_if_false(parser->val_zerof);\n }\n return mod_json_parser_insert(parser, mod_json_token_depth(tok),\n parser->val_zerof);\n}\n\nstatic inline int mod_json_parser_event_float(mod_json_token_t *tok,\n mod_json_float_t val) {\n int ret = -1;\n\n if (val != 0.0) {\n mod_json_value_t *jval;\n\n jval = mod_json_value_set_float(val);\n if (jval) {\n ret = mod_json_parser_insert(mod_json_token_parser(tok),\n mod_json_token_depth(tok), jval);\n mod_json_value_unset(jval);\n }\n } else {\n /* zero event */\n ret = mod_json_parser_event_zerof(tok);\n }\n return ret;\n}\n\nstatic inline int mod_json_parser_event_empty(mod_json_token_t *tok) {\n mod_json_parser_t *parser;\n\n /* get information */\n parser = mod_json_token_parser(tok);\n\n if (!parser->val_empty) {\n mod_json_string_t *str;\n\n str = mod_json_string_set(\"\", 0);\n mod_json_minus_if_false(str);\n\n parser->val_empty = mod_json_value_set_string(str);\n mod_json_string_unset(str);\n mod_json_minus_if_false(parser->val_empty);\n }\n return mod_json_parser_insert(parser, mod_json_token_depth(tok),\n parser->val_empty);\n}\n\nstatic inline int mod_json_parser_event_string(mod_json_token_t *tok,\n mod_json_cchar_t *val,\n mod_json_size_t len) {\n int ret = -1;\n\n if (len > 0) {\n mod_json_string_t *str;\n mod_json_value_t *jval;\n\n str = mod_json_string_set(val, len);\n if (str) {\n jval = mod_json_value_set_string(str);\n } else {\n jval = NULL;\n }\n mod_json_string_unset(str);\n\n if (jval) {\n ret = mod_json_parser_insert(mod_json_token_parser(tok),\n mod_json_token_depth(tok), jval);\n mod_json_value_unset(jval);\n }\n } else {\n /* empty event */\n ret = mod_json_parser_event_empty(tok);\n }\n return ret;\n}\n\nstatic int mod_json_parser_event(mod_json_token_t *tok, mod_json_void_t *val,\n mod_json_size_t len) {\n switch (tok->event_code) {\n case mod_json_event_field:\n return mod_json_parser_event_field(tok, (mod_json_cchar_t *)val, len);\n\n case mod_json_event_object:\n return mod_json_parser_event_object(tok);\n\n case mod_json_event_array:\n return mod_json_parser_event_array(tok);\n\n case mod_json_event_null:\n return mod_json_parser_event_null(tok);\n\n case mod_json_event_boolean:\n return mod_json_parser_event_boolean(tok, *(mod_json_boolean_t *)val);\n\n case mod_json_event_integer:\n return mod_json_parser_event_integer(tok, *(mod_json_integer_t *)val);\n\n case mod_json_event_float:\n return mod_json_parser_event_float(tok, *(mod_json_float_t *)val);\n\n case mod_json_event_string:\n return mod_json_parser_event_string(tok, (mod_json_cchar_t *)val, len);\n\n default:\n break;\n }\n return -1;\n}\n\nstatic inline mod_json_parser_t *mod_json_parser_create(mod_json_size_t depth) {\n mod_json_parser_t *parser;\n mod_json_null_if_false(depth > 0);\n\n parser = (mod_json_parser_t *)mod_json_malloc(\n depth * sizeof(mod_json_value_t *) + sizeof(mod_json_parser_t));\n mod_json_null_if_false(parser);\n\n memset(parser, 0, sizeof(mod_json_parser_t));\n parser->vals[0] = NULL;\n return parser;\n}\n\nstatic inline void mod_json_parser_destroy(mod_json_parser_t *par) {\n mod_json_value_unset(par->val_null);\n mod_json_value_unset(par->val_true);\n mod_json_value_unset(par->val_false);\n mod_json_value_unset(par->val_zero);\n mod_json_value_unset(par->val_zerof);\n mod_json_value_unset(par->val_empty);\n mod_json_string_unset(par->key);\n mod_json_free(par);\n}\n\nmod_json_value_t *mod_json_parse(mod_json_token_t *tok,\n mod_json_cchar_t *cstr) {\n mod_json_parser_t *parser;\n mod_json_value_t *root;\n mod_json_null_if_false(tok && cstr && *cstr);\n\n parser = mod_json_parser_create(mod_json_token_max_depth(tok));\n mod_json_null_if_false(parser);\n\n mod_json_token_set_parser(tok, parser);\n mod_json_token_set_event(tok, mod_json_parser_event);\n\n if (mod_json_token_parse(tok, cstr) == 0) {\n root = parser->vals[0];\n } else {\n /* error occur */\n root = NULL;\n mod_json_value_unset(parser->vals[0]);\n }\n\n /* clean up */\n mod_json_parser_destroy(parser);\n\n /* success? */\n return root;\n}\n\nmod_json_value_t *mod_json_parse_simply(mod_json_cchar_t *cstr,\n mod_json_size_t opts) {\n mod_json_value_t *val;\n mod_json_token_t *tok;\n mod_json_option_t opt;\n\n opt.options = opts;\n opt.object_depth = 0; /* Use default object depth */\n opt.array_depth = 0; /* Use default array depth */\n\n tok = mod_json_token_create(&opt);\n mod_json_null_if_false(tok);\n\n val = mod_json_parse(tok, cstr);\n mod_json_token_destroy(tok);\n\n /* value of root */\n return val;\n}\n\nstatic inline int mod_json_dump_null(mod_json_string_t *str) {\n return mod_json_string_add_cstr(str, \"null\", 4);\n}\n\nstatic inline int mod_json_dump_boolean(mod_json_string_t *str,\n mod_json_boolean_t bol) {\n if (!bol) {\n return mod_json_string_add_cstr(str, \"false\", 5);\n }\n return mod_json_string_add_cstr(str, \"true\", 4);\n}\n\nstatic inline int mod_json_dump_integer(mod_json_string_t *str,\n mod_json_integer_t num) {\n mod_json_char_t buf[32];\n\n return mod_json_string_add_cstr(str, buf, mod_json_utils_itostr(buf, num));\n}\n\nstatic inline int mod_json_dump_float(mod_json_string_t *str,\n mod_json_float_t dbl) {\n mod_json_char_t buf[32];\n\n return mod_json_string_add_cstr(\n str, buf,\n (mod_json_size_t)mod_json_utils_snprintf(buf, sizeof(buf), \"%g\", dbl));\n}\n\nstatic inline int mod_json_dump_string(mod_json_string_t *str,\n mod_json_string_t *val) {\n mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '\\\"'));\n\n if (val) {\n mod_json_minus_if_ne_zero(mod_json_string_add_jstr(str, val));\n }\n mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '\\\"'));\n return 0;\n}\n\nstatic inline int mod_json_dump_value(mod_json_string_t *str,\n mod_json_value_t *val);\n\nstatic inline int mod_json_dump_array(mod_json_string_t *str,\n mod_json_array_t *arr) {\n mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '['));\n\n if (arr) {\n mod_json_value_t **iter = arr->first;\n\n for (; iter != arr->last; ++iter) {\n mod_json_minus_if_ne_zero(mod_json_dump_value(str, *iter));\n if (iter + 1 != arr->last) {\n mod_json_minus_if_ne_zero(mod_json_string_add_char(str, ','));\n }\n }\n }\n mod_json_minus_if_ne_zero(mod_json_string_add_char(str, ']'));\n return 0;\n}\n\nstatic inline int mod_json_dump_key(mod_json_string_t *str,\n mod_json_string_t *key) {\n mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '\\\"'));\n mod_json_minus_if_ne_zero(mod_json_string_add_jstr(str, key));\n mod_json_minus_if_ne_zero(mod_json_string_add_cstr(str, \"\\\":\", 2));\n return 0;\n}\n\nstatic inline int mod_json_dump_object(mod_json_string_t *str,\n mod_json_object_t *obj) {\n mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '{'));\n\n if (obj) {\n mod_json_pair_t *iter = obj->first;\n\n for (; iter != obj->last; ++iter) {\n mod_json_minus_if_ne_zero(mod_json_dump_key(str, iter->key));\n mod_json_minus_if_ne_zero(mod_json_dump_value(str, iter->val));\n\n if (iter + 1 != obj->last) {\n mod_json_minus_if_ne_zero(mod_json_string_add_char(str, ','));\n }\n }\n }\n mod_json_minus_if_ne_zero(mod_json_string_add_char(str, '}'));\n return 0;\n}\n\nstatic inline int mod_json_dump_value(mod_json_string_t *str,\n mod_json_value_t *val) {\n if (val) {\n switch (val->type) {\n case mod_json_type_null:\n return mod_json_dump_null(str);\n\n case mod_json_type_boolean:\n return mod_json_dump_boolean(str, val->data.c_bool);\n\n case mod_json_type_integer:\n return mod_json_dump_integer(str, val->data.c_int);\n\n case mod_json_type_float:\n return mod_json_dump_float(str, val->data.c_float);\n\n case mod_json_type_string:\n return mod_json_dump_string(str, val->data.c_str);\n\n case mod_json_type_array:\n return mod_json_dump_array(str, val->data.c_arr);\n\n case mod_json_type_object:\n return mod_json_dump_object(str, val->data.c_obj);\n\n default:\n return -1;\n }\n }\n return mod_json_dump_null(str);\n}\n\nmod_json_string_t *mod_json_dump(mod_json_value_t *val) {\n mod_json_string_t *str = mod_json_string_set(\"\", 0);\n mod_json_null_if_false(str);\n\n if (mod_json_unlikely(mod_json_dump_value(str, val) != 0)) {\n /* error occur */\n mod_json_string_unset(str);\n return NULL;\n }\n return str;\n}\n" }, { "path": "src/ailego/hash/crc32c.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n\n#if !defined(__SSE4_2__) && !defined(__ARM_FEATURE_CRC32)\n/**\n * The following CRC lookup table was generated automagically\n * using the following model parameters:\n *\n * Generator Polynomial = ................. 0x1EDC6F41\n * Generator Polynomial Length = .......... 32 bits\n * Reflected Bits = ....................... TRUE\n * Table Generation Offset = .............. 32 bits\n * Number of Slices = ..................... 8 slices\n * Slice Lengths = ........................ 8 8 8 8 8 8 8 8\n */\nstatic uint32_t crc_tableil8_o32[256] = {\n 0x00000000, 0xF26B8303, 0xE13B70F7, 0x1350F3F4, 0xC79A971F, 0x35F1141C,\n 0x26A1E7E8, 0xD4CA64EB, 0x8AD958CF, 0x78B2DBCC, 0x6BE22838, 0x9989AB3B,\n 0x4D43CFD0, 0xBF284CD3, 0xAC78BF27, 0x5E133C24, 0x105EC76F, 0xE235446C,\n 0xF165B798, 0x030E349B, 0xD7C45070, 0x25AFD373, 0x36FF2087, 0xC494A384,\n 0x9A879FA0, 0x68EC1CA3, 0x7BBCEF57, 0x89D76C54, 0x5D1D08BF, 0xAF768BBC,\n 0xBC267848, 0x4E4DFB4B, 0x20BD8EDE, 0xD2D60DDD, 0xC186FE29, 0x33ED7D2A,\n 0xE72719C1, 0x154C9AC2, 0x061C6936, 0xF477EA35, 0xAA64D611, 0x580F5512,\n 0x4B5FA6E6, 0xB93425E5, 0x6DFE410E, 0x9F95C20D, 0x8CC531F9, 0x7EAEB2FA,\n 0x30E349B1, 0xC288CAB2, 0xD1D83946, 0x23B3BA45, 0xF779DEAE, 0x05125DAD,\n 0x1642AE59, 0xE4292D5A, 0xBA3A117E, 0x4851927D, 0x5B016189, 0xA96AE28A,\n 0x7DA08661, 0x8FCB0562, 0x9C9BF696, 0x6EF07595, 0x417B1DBC, 0xB3109EBF,\n 0xA0406D4B, 0x522BEE48, 0x86E18AA3, 0x748A09A0, 0x67DAFA54, 0x95B17957,\n 0xCBA24573, 0x39C9C670, 0x2A993584, 0xD8F2B687, 0x0C38D26C, 0xFE53516F,\n 0xED03A29B, 0x1F682198, 0x5125DAD3, 0xA34E59D0, 0xB01EAA24, 0x42752927,\n 0x96BF4DCC, 0x64D4CECF, 0x77843D3B, 0x85EFBE38, 0xDBFC821C, 0x2997011F,\n 0x3AC7F2EB, 0xC8AC71E8, 0x1C661503, 0xEE0D9600, 0xFD5D65F4, 0x0F36E6F7,\n 0x61C69362, 0x93AD1061, 0x80FDE395, 0x72966096, 0xA65C047D, 0x5437877E,\n 0x4767748A, 0xB50CF789, 0xEB1FCBAD, 0x197448AE, 0x0A24BB5A, 0xF84F3859,\n 0x2C855CB2, 0xDEEEDFB1, 0xCDBE2C45, 0x3FD5AF46, 0x7198540D, 0x83F3D70E,\n 0x90A324FA, 0x62C8A7F9, 0xB602C312, 0x44694011, 0x5739B3E5, 0xA55230E6,\n 0xFB410CC2, 0x092A8FC1, 0x1A7A7C35, 0xE811FF36, 0x3CDB9BDD, 0xCEB018DE,\n 0xDDE0EB2A, 0x2F8B6829, 0x82F63B78, 0x709DB87B, 0x63CD4B8F, 0x91A6C88C,\n 0x456CAC67, 0xB7072F64, 0xA457DC90, 0x563C5F93, 0x082F63B7, 0xFA44E0B4,\n 0xE9141340, 0x1B7F9043, 0xCFB5F4A8, 0x3DDE77AB, 0x2E8E845F, 0xDCE5075C,\n 0x92A8FC17, 0x60C37F14, 0x73938CE0, 0x81F80FE3, 0x55326B08, 0xA759E80B,\n 0xB4091BFF, 0x466298FC, 0x1871A4D8, 0xEA1A27DB, 0xF94AD42F, 0x0B21572C,\n 0xDFEB33C7, 0x2D80B0C4, 0x3ED04330, 0xCCBBC033, 0xA24BB5A6, 0x502036A5,\n 0x4370C551, 0xB11B4652, 0x65D122B9, 0x97BAA1BA, 0x84EA524E, 0x7681D14D,\n 0x2892ED69, 0xDAF96E6A, 0xC9A99D9E, 0x3BC21E9D, 0xEF087A76, 0x1D63F975,\n 0x0E330A81, 0xFC588982, 0xB21572C9, 0x407EF1CA, 0x532E023E, 0xA145813D,\n 0x758FE5D6, 0x87E466D5, 0x94B49521, 0x66DF1622, 0x38CC2A06, 0xCAA7A905,\n 0xD9F75AF1, 0x2B9CD9F2, 0xFF56BD19, 0x0D3D3E1A, 0x1E6DCDEE, 0xEC064EED,\n 0xC38D26C4, 0x31E6A5C7, 0x22B65633, 0xD0DDD530, 0x0417B1DB, 0xF67C32D8,\n 0xE52CC12C, 0x1747422F, 0x49547E0B, 0xBB3FFD08, 0xA86F0EFC, 0x5A048DFF,\n 0x8ECEE914, 0x7CA56A17, 0x6FF599E3, 0x9D9E1AE0, 0xD3D3E1AB, 0x21B862A8,\n 0x32E8915C, 0xC083125F, 0x144976B4, 0xE622F5B7, 0xF5720643, 0x07198540,\n 0x590AB964, 0xAB613A67, 0xB831C993, 0x4A5A4A90, 0x9E902E7B, 0x6CFBAD78,\n 0x7FAB5E8C, 0x8DC0DD8F, 0xE330A81A, 0x115B2B19, 0x020BD8ED, 0xF0605BEE,\n 0x24AA3F05, 0xD6C1BC06, 0xC5914FF2, 0x37FACCF1, 0x69E9F0D5, 0x9B8273D6,\n 0x88D28022, 0x7AB90321, 0xAE7367CA, 0x5C18E4C9, 0x4F48173D, 0xBD23943E,\n 0xF36E6F75, 0x0105EC76, 0x12551F82, 0xE03E9C81, 0x34F4F86A, 0xC69F7B69,\n 0xD5CF889D, 0x27A40B9E, 0x79B737BA, 0x8BDCB4B9, 0x988C474D, 0x6AE7C44E,\n 0xBE2DA0A5, 0x4C4623A6, 0x5F16D052, 0xAD7D5351};\n\n/**\n * The following CRC lookup table was generated automagically\n * using the following model parameters:\n *\n * Generator Polynomial = ................. 0x1EDC6F41\n * Generator Polynomial Length = .......... 32 bits\n * Reflected Bits = ....................... TRUE\n * Table Generation Offset = .............. 32 bits\n * Number of Slices = ..................... 8 slices\n * Slice Lengths = ........................ 8 8 8 8 8 8 8 8\n */\nstatic uint32_t crc_tableil8_o40[256] = {\n 0x00000000, 0x13A29877, 0x274530EE, 0x34E7A899, 0x4E8A61DC, 0x5D28F9AB,\n 0x69CF5132, 0x7A6DC945, 0x9D14C3B8, 0x8EB65BCF, 0xBA51F356, 0xA9F36B21,\n 0xD39EA264, 0xC03C3A13, 0xF4DB928A, 0xE7790AFD, 0x3FC5F181, 0x2C6769F6,\n 0x1880C16F, 0x0B225918, 0x714F905D, 0x62ED082A, 0x560AA0B3, 0x45A838C4,\n 0xA2D13239, 0xB173AA4E, 0x859402D7, 0x96369AA0, 0xEC5B53E5, 0xFFF9CB92,\n 0xCB1E630B, 0xD8BCFB7C, 0x7F8BE302, 0x6C297B75, 0x58CED3EC, 0x4B6C4B9B,\n 0x310182DE, 0x22A31AA9, 0x1644B230, 0x05E62A47, 0xE29F20BA, 0xF13DB8CD,\n 0xC5DA1054, 0xD6788823, 0xAC154166, 0xBFB7D911, 0x8B507188, 0x98F2E9FF,\n 0x404E1283, 0x53EC8AF4, 0x670B226D, 0x74A9BA1A, 0x0EC4735F, 0x1D66EB28,\n 0x298143B1, 0x3A23DBC6, 0xDD5AD13B, 0xCEF8494C, 0xFA1FE1D5, 0xE9BD79A2,\n 0x93D0B0E7, 0x80722890, 0xB4958009, 0xA737187E, 0xFF17C604, 0xECB55E73,\n 0xD852F6EA, 0xCBF06E9D, 0xB19DA7D8, 0xA23F3FAF, 0x96D89736, 0x857A0F41,\n 0x620305BC, 0x71A19DCB, 0x45463552, 0x56E4AD25, 0x2C896460, 0x3F2BFC17,\n 0x0BCC548E, 0x186ECCF9, 0xC0D23785, 0xD370AFF2, 0xE797076B, 0xF4359F1C,\n 0x8E585659, 0x9DFACE2E, 0xA91D66B7, 0xBABFFEC0, 0x5DC6F43D, 0x4E646C4A,\n 0x7A83C4D3, 0x69215CA4, 0x134C95E1, 0x00EE0D96, 0x3409A50F, 0x27AB3D78,\n 0x809C2506, 0x933EBD71, 0xA7D915E8, 0xB47B8D9F, 0xCE1644DA, 0xDDB4DCAD,\n 0xE9537434, 0xFAF1EC43, 0x1D88E6BE, 0x0E2A7EC9, 0x3ACDD650, 0x296F4E27,\n 0x53028762, 0x40A01F15, 0x7447B78C, 0x67E52FFB, 0xBF59D487, 0xACFB4CF0,\n 0x981CE469, 0x8BBE7C1E, 0xF1D3B55B, 0xE2712D2C, 0xD69685B5, 0xC5341DC2,\n 0x224D173F, 0x31EF8F48, 0x050827D1, 0x16AABFA6, 0x6CC776E3, 0x7F65EE94,\n 0x4B82460D, 0x5820DE7A, 0xFBC3FAF9, 0xE861628E, 0xDC86CA17, 0xCF245260,\n 0xB5499B25, 0xA6EB0352, 0x920CABCB, 0x81AE33BC, 0x66D73941, 0x7575A136,\n 0x419209AF, 0x523091D8, 0x285D589D, 0x3BFFC0EA, 0x0F186873, 0x1CBAF004,\n 0xC4060B78, 0xD7A4930F, 0xE3433B96, 0xF0E1A3E1, 0x8A8C6AA4, 0x992EF2D3,\n 0xADC95A4A, 0xBE6BC23D, 0x5912C8C0, 0x4AB050B7, 0x7E57F82E, 0x6DF56059,\n 0x1798A91C, 0x043A316B, 0x30DD99F2, 0x237F0185, 0x844819FB, 0x97EA818C,\n 0xA30D2915, 0xB0AFB162, 0xCAC27827, 0xD960E050, 0xED8748C9, 0xFE25D0BE,\n 0x195CDA43, 0x0AFE4234, 0x3E19EAAD, 0x2DBB72DA, 0x57D6BB9F, 0x447423E8,\n 0x70938B71, 0x63311306, 0xBB8DE87A, 0xA82F700D, 0x9CC8D894, 0x8F6A40E3,\n 0xF50789A6, 0xE6A511D1, 0xD242B948, 0xC1E0213F, 0x26992BC2, 0x353BB3B5,\n 0x01DC1B2C, 0x127E835B, 0x68134A1E, 0x7BB1D269, 0x4F567AF0, 0x5CF4E287,\n 0x04D43CFD, 0x1776A48A, 0x23910C13, 0x30339464, 0x4A5E5D21, 0x59FCC556,\n 0x6D1B6DCF, 0x7EB9F5B8, 0x99C0FF45, 0x8A626732, 0xBE85CFAB, 0xAD2757DC,\n 0xD74A9E99, 0xC4E806EE, 0xF00FAE77, 0xE3AD3600, 0x3B11CD7C, 0x28B3550B,\n 0x1C54FD92, 0x0FF665E5, 0x759BACA0, 0x663934D7, 0x52DE9C4E, 0x417C0439,\n 0xA6050EC4, 0xB5A796B3, 0x81403E2A, 0x92E2A65D, 0xE88F6F18, 0xFB2DF76F,\n 0xCFCA5FF6, 0xDC68C781, 0x7B5FDFFF, 0x68FD4788, 0x5C1AEF11, 0x4FB87766,\n 0x35D5BE23, 0x26772654, 0x12908ECD, 0x013216BA, 0xE64B1C47, 0xF5E98430,\n 0xC10E2CA9, 0xD2ACB4DE, 0xA8C17D9B, 0xBB63E5EC, 0x8F844D75, 0x9C26D502,\n 0x449A2E7E, 0x5738B609, 0x63DF1E90, 0x707D86E7, 0x0A104FA2, 0x19B2D7D5,\n 0x2D557F4C, 0x3EF7E73B, 0xD98EEDC6, 0xCA2C75B1, 0xFECBDD28, 0xED69455F,\n 0x97048C1A, 0x84A6146D, 0xB041BCF4, 0xA3E32483};\n\n/**\n * The following CRC lookup table was generated automagically\n * using the following model parameters:\n *\n * Generator Polynomial = ................. 0x1EDC6F41\n * Generator Polynomial Length = .......... 32 bits\n * Reflected Bits = ....................... TRUE\n * Table Generation Offset = .............. 32 bits\n * Number of Slices = ..................... 8 slices\n * Slice Lengths = ........................ 8 8 8 8 8 8 8 8\n */\nstatic uint32_t crc_tableil8_o48[256] = {\n 0x00000000, 0xA541927E, 0x4F6F520D, 0xEA2EC073, 0x9EDEA41A, 0x3B9F3664,\n 0xD1B1F617, 0x74F06469, 0x38513EC5, 0x9D10ACBB, 0x773E6CC8, 0xD27FFEB6,\n 0xA68F9ADF, 0x03CE08A1, 0xE9E0C8D2, 0x4CA15AAC, 0x70A27D8A, 0xD5E3EFF4,\n 0x3FCD2F87, 0x9A8CBDF9, 0xEE7CD990, 0x4B3D4BEE, 0xA1138B9D, 0x045219E3,\n 0x48F3434F, 0xEDB2D131, 0x079C1142, 0xA2DD833C, 0xD62DE755, 0x736C752B,\n 0x9942B558, 0x3C032726, 0xE144FB14, 0x4405696A, 0xAE2BA919, 0x0B6A3B67,\n 0x7F9A5F0E, 0xDADBCD70, 0x30F50D03, 0x95B49F7D, 0xD915C5D1, 0x7C5457AF,\n 0x967A97DC, 0x333B05A2, 0x47CB61CB, 0xE28AF3B5, 0x08A433C6, 0xADE5A1B8,\n 0x91E6869E, 0x34A714E0, 0xDE89D493, 0x7BC846ED, 0x0F382284, 0xAA79B0FA,\n 0x40577089, 0xE516E2F7, 0xA9B7B85B, 0x0CF62A25, 0xE6D8EA56, 0x43997828,\n 0x37691C41, 0x92288E3F, 0x78064E4C, 0xDD47DC32, 0xC76580D9, 0x622412A7,\n 0x880AD2D4, 0x2D4B40AA, 0x59BB24C3, 0xFCFAB6BD, 0x16D476CE, 0xB395E4B0,\n 0xFF34BE1C, 0x5A752C62, 0xB05BEC11, 0x151A7E6F, 0x61EA1A06, 0xC4AB8878,\n 0x2E85480B, 0x8BC4DA75, 0xB7C7FD53, 0x12866F2D, 0xF8A8AF5E, 0x5DE93D20,\n 0x29195949, 0x8C58CB37, 0x66760B44, 0xC337993A, 0x8F96C396, 0x2AD751E8,\n 0xC0F9919B, 0x65B803E5, 0x1148678C, 0xB409F5F2, 0x5E273581, 0xFB66A7FF,\n 0x26217BCD, 0x8360E9B3, 0x694E29C0, 0xCC0FBBBE, 0xB8FFDFD7, 0x1DBE4DA9,\n 0xF7908DDA, 0x52D11FA4, 0x1E704508, 0xBB31D776, 0x511F1705, 0xF45E857B,\n 0x80AEE112, 0x25EF736C, 0xCFC1B31F, 0x6A802161, 0x56830647, 0xF3C29439,\n 0x19EC544A, 0xBCADC634, 0xC85DA25D, 0x6D1C3023, 0x8732F050, 0x2273622E,\n 0x6ED23882, 0xCB93AAFC, 0x21BD6A8F, 0x84FCF8F1, 0xF00C9C98, 0x554D0EE6,\n 0xBF63CE95, 0x1A225CEB, 0x8B277743, 0x2E66E53D, 0xC448254E, 0x6109B730,\n 0x15F9D359, 0xB0B84127, 0x5A968154, 0xFFD7132A, 0xB3764986, 0x1637DBF8,\n 0xFC191B8B, 0x595889F5, 0x2DA8ED9C, 0x88E97FE2, 0x62C7BF91, 0xC7862DEF,\n 0xFB850AC9, 0x5EC498B7, 0xB4EA58C4, 0x11ABCABA, 0x655BAED3, 0xC01A3CAD,\n 0x2A34FCDE, 0x8F756EA0, 0xC3D4340C, 0x6695A672, 0x8CBB6601, 0x29FAF47F,\n 0x5D0A9016, 0xF84B0268, 0x1265C21B, 0xB7245065, 0x6A638C57, 0xCF221E29,\n 0x250CDE5A, 0x804D4C24, 0xF4BD284D, 0x51FCBA33, 0xBBD27A40, 0x1E93E83E,\n 0x5232B292, 0xF77320EC, 0x1D5DE09F, 0xB81C72E1, 0xCCEC1688, 0x69AD84F6,\n 0x83834485, 0x26C2D6FB, 0x1AC1F1DD, 0xBF8063A3, 0x55AEA3D0, 0xF0EF31AE,\n 0x841F55C7, 0x215EC7B9, 0xCB7007CA, 0x6E3195B4, 0x2290CF18, 0x87D15D66,\n 0x6DFF9D15, 0xC8BE0F6B, 0xBC4E6B02, 0x190FF97C, 0xF321390F, 0x5660AB71,\n 0x4C42F79A, 0xE90365E4, 0x032DA597, 0xA66C37E9, 0xD29C5380, 0x77DDC1FE,\n 0x9DF3018D, 0x38B293F3, 0x7413C95F, 0xD1525B21, 0x3B7C9B52, 0x9E3D092C,\n 0xEACD6D45, 0x4F8CFF3B, 0xA5A23F48, 0x00E3AD36, 0x3CE08A10, 0x99A1186E,\n 0x738FD81D, 0xD6CE4A63, 0xA23E2E0A, 0x077FBC74, 0xED517C07, 0x4810EE79,\n 0x04B1B4D5, 0xA1F026AB, 0x4BDEE6D8, 0xEE9F74A6, 0x9A6F10CF, 0x3F2E82B1,\n 0xD50042C2, 0x7041D0BC, 0xAD060C8E, 0x08479EF0, 0xE2695E83, 0x4728CCFD,\n 0x33D8A894, 0x96993AEA, 0x7CB7FA99, 0xD9F668E7, 0x9557324B, 0x3016A035,\n 0xDA386046, 0x7F79F238, 0x0B899651, 0xAEC8042F, 0x44E6C45C, 0xE1A75622,\n 0xDDA47104, 0x78E5E37A, 0x92CB2309, 0x378AB177, 0x437AD51E, 0xE63B4760,\n 0x0C158713, 0xA954156D, 0xE5F54FC1, 0x40B4DDBF, 0xAA9A1DCC, 0x0FDB8FB2,\n 0x7B2BEBDB, 0xDE6A79A5, 0x3444B9D6, 0x91052BA8};\n\n/**\n * The following CRC lookup table was generated automagically\n * using the following model parameters:\n *\n * Generator Polynomial = ................. 0x1EDC6F41\n * Generator Polynomial Length = .......... 32 bits\n * Reflected Bits = ....................... TRUE\n * Table Generation Offset = .............. 32 bits\n * Number of Slices = ..................... 8 slices\n * Slice Lengths = ........................ 8 8 8 8 8 8 8 8\n */\nstatic uint32_t crc_tableil8_o56[256] = {\n 0x00000000, 0xDD45AAB8, 0xBF672381, 0x62228939, 0x7B2231F3, 0xA6679B4B,\n 0xC4451272, 0x1900B8CA, 0xF64463E6, 0x2B01C95E, 0x49234067, 0x9466EADF,\n 0x8D665215, 0x5023F8AD, 0x32017194, 0xEF44DB2C, 0xE964B13D, 0x34211B85,\n 0x560392BC, 0x8B463804, 0x924680CE, 0x4F032A76, 0x2D21A34F, 0xF06409F7,\n 0x1F20D2DB, 0xC2657863, 0xA047F15A, 0x7D025BE2, 0x6402E328, 0xB9474990,\n 0xDB65C0A9, 0x06206A11, 0xD725148B, 0x0A60BE33, 0x6842370A, 0xB5079DB2,\n 0xAC072578, 0x71428FC0, 0x136006F9, 0xCE25AC41, 0x2161776D, 0xFC24DDD5,\n 0x9E0654EC, 0x4343FE54, 0x5A43469E, 0x8706EC26, 0xE524651F, 0x3861CFA7,\n 0x3E41A5B6, 0xE3040F0E, 0x81268637, 0x5C632C8F, 0x45639445, 0x98263EFD,\n 0xFA04B7C4, 0x27411D7C, 0xC805C650, 0x15406CE8, 0x7762E5D1, 0xAA274F69,\n 0xB327F7A3, 0x6E625D1B, 0x0C40D422, 0xD1057E9A, 0xABA65FE7, 0x76E3F55F,\n 0x14C17C66, 0xC984D6DE, 0xD0846E14, 0x0DC1C4AC, 0x6FE34D95, 0xB2A6E72D,\n 0x5DE23C01, 0x80A796B9, 0xE2851F80, 0x3FC0B538, 0x26C00DF2, 0xFB85A74A,\n 0x99A72E73, 0x44E284CB, 0x42C2EEDA, 0x9F874462, 0xFDA5CD5B, 0x20E067E3,\n 0x39E0DF29, 0xE4A57591, 0x8687FCA8, 0x5BC25610, 0xB4868D3C, 0x69C32784,\n 0x0BE1AEBD, 0xD6A40405, 0xCFA4BCCF, 0x12E11677, 0x70C39F4E, 0xAD8635F6,\n 0x7C834B6C, 0xA1C6E1D4, 0xC3E468ED, 0x1EA1C255, 0x07A17A9F, 0xDAE4D027,\n 0xB8C6591E, 0x6583F3A6, 0x8AC7288A, 0x57828232, 0x35A00B0B, 0xE8E5A1B3,\n 0xF1E51979, 0x2CA0B3C1, 0x4E823AF8, 0x93C79040, 0x95E7FA51, 0x48A250E9,\n 0x2A80D9D0, 0xF7C57368, 0xEEC5CBA2, 0x3380611A, 0x51A2E823, 0x8CE7429B,\n 0x63A399B7, 0xBEE6330F, 0xDCC4BA36, 0x0181108E, 0x1881A844, 0xC5C402FC,\n 0xA7E68BC5, 0x7AA3217D, 0x52A0C93F, 0x8FE56387, 0xEDC7EABE, 0x30824006,\n 0x2982F8CC, 0xF4C75274, 0x96E5DB4D, 0x4BA071F5, 0xA4E4AAD9, 0x79A10061,\n 0x1B838958, 0xC6C623E0, 0xDFC69B2A, 0x02833192, 0x60A1B8AB, 0xBDE41213,\n 0xBBC47802, 0x6681D2BA, 0x04A35B83, 0xD9E6F13B, 0xC0E649F1, 0x1DA3E349,\n 0x7F816A70, 0xA2C4C0C8, 0x4D801BE4, 0x90C5B15C, 0xF2E73865, 0x2FA292DD,\n 0x36A22A17, 0xEBE780AF, 0x89C50996, 0x5480A32E, 0x8585DDB4, 0x58C0770C,\n 0x3AE2FE35, 0xE7A7548D, 0xFEA7EC47, 0x23E246FF, 0x41C0CFC6, 0x9C85657E,\n 0x73C1BE52, 0xAE8414EA, 0xCCA69DD3, 0x11E3376B, 0x08E38FA1, 0xD5A62519,\n 0xB784AC20, 0x6AC10698, 0x6CE16C89, 0xB1A4C631, 0xD3864F08, 0x0EC3E5B0,\n 0x17C35D7A, 0xCA86F7C2, 0xA8A47EFB, 0x75E1D443, 0x9AA50F6F, 0x47E0A5D7,\n 0x25C22CEE, 0xF8878656, 0xE1873E9C, 0x3CC29424, 0x5EE01D1D, 0x83A5B7A5,\n 0xF90696D8, 0x24433C60, 0x4661B559, 0x9B241FE1, 0x8224A72B, 0x5F610D93,\n 0x3D4384AA, 0xE0062E12, 0x0F42F53E, 0xD2075F86, 0xB025D6BF, 0x6D607C07,\n 0x7460C4CD, 0xA9256E75, 0xCB07E74C, 0x16424DF4, 0x106227E5, 0xCD278D5D,\n 0xAF050464, 0x7240AEDC, 0x6B401616, 0xB605BCAE, 0xD4273597, 0x09629F2F,\n 0xE6264403, 0x3B63EEBB, 0x59416782, 0x8404CD3A, 0x9D0475F0, 0x4041DF48,\n 0x22635671, 0xFF26FCC9, 0x2E238253, 0xF36628EB, 0x9144A1D2, 0x4C010B6A,\n 0x5501B3A0, 0x88441918, 0xEA669021, 0x37233A99, 0xD867E1B5, 0x05224B0D,\n 0x6700C234, 0xBA45688C, 0xA345D046, 0x7E007AFE, 0x1C22F3C7, 0xC167597F,\n 0xC747336E, 0x1A0299D6, 0x782010EF, 0xA565BA57, 0xBC65029D, 0x6120A825,\n 0x0302211C, 0xDE478BA4, 0x31035088, 0xEC46FA30, 0x8E647309, 0x5321D9B1,\n 0x4A21617B, 0x9764CBC3, 0xF54642FA, 0x2803E842};\n\n/**\n * The following CRC lookup table was generated automagically\n * using the following model parameters:\n *\n * Generator Polynomial = ................. 0x1EDC6F41\n * Generator Polynomial Length = .......... 32 bits\n * Reflected Bits = ....................... TRUE\n * Table Generation Offset = .............. 32 bits\n * Number of Slices = ..................... 8 slices\n * Slice Lengths = ........................ 8 8 8 8 8 8 8 8\n */\nstatic uint32_t crc_tableil8_o64[256] = {\n 0x00000000, 0x38116FAC, 0x7022DF58, 0x4833B0F4, 0xE045BEB0, 0xD854D11C,\n 0x906761E8, 0xA8760E44, 0xC5670B91, 0xFD76643D, 0xB545D4C9, 0x8D54BB65,\n 0x2522B521, 0x1D33DA8D, 0x55006A79, 0x6D1105D5, 0x8F2261D3, 0xB7330E7F,\n 0xFF00BE8B, 0xC711D127, 0x6F67DF63, 0x5776B0CF, 0x1F45003B, 0x27546F97,\n 0x4A456A42, 0x725405EE, 0x3A67B51A, 0x0276DAB6, 0xAA00D4F2, 0x9211BB5E,\n 0xDA220BAA, 0xE2336406, 0x1BA8B557, 0x23B9DAFB, 0x6B8A6A0F, 0x539B05A3,\n 0xFBED0BE7, 0xC3FC644B, 0x8BCFD4BF, 0xB3DEBB13, 0xDECFBEC6, 0xE6DED16A,\n 0xAEED619E, 0x96FC0E32, 0x3E8A0076, 0x069B6FDA, 0x4EA8DF2E, 0x76B9B082,\n 0x948AD484, 0xAC9BBB28, 0xE4A80BDC, 0xDCB96470, 0x74CF6A34, 0x4CDE0598,\n 0x04EDB56C, 0x3CFCDAC0, 0x51EDDF15, 0x69FCB0B9, 0x21CF004D, 0x19DE6FE1,\n 0xB1A861A5, 0x89B90E09, 0xC18ABEFD, 0xF99BD151, 0x37516AAE, 0x0F400502,\n 0x4773B5F6, 0x7F62DA5A, 0xD714D41E, 0xEF05BBB2, 0xA7360B46, 0x9F2764EA,\n 0xF236613F, 0xCA270E93, 0x8214BE67, 0xBA05D1CB, 0x1273DF8F, 0x2A62B023,\n 0x625100D7, 0x5A406F7B, 0xB8730B7D, 0x806264D1, 0xC851D425, 0xF040BB89,\n 0x5836B5CD, 0x6027DA61, 0x28146A95, 0x10050539, 0x7D1400EC, 0x45056F40,\n 0x0D36DFB4, 0x3527B018, 0x9D51BE5C, 0xA540D1F0, 0xED736104, 0xD5620EA8,\n 0x2CF9DFF9, 0x14E8B055, 0x5CDB00A1, 0x64CA6F0D, 0xCCBC6149, 0xF4AD0EE5,\n 0xBC9EBE11, 0x848FD1BD, 0xE99ED468, 0xD18FBBC4, 0x99BC0B30, 0xA1AD649C,\n 0x09DB6AD8, 0x31CA0574, 0x79F9B580, 0x41E8DA2C, 0xA3DBBE2A, 0x9BCAD186,\n 0xD3F96172, 0xEBE80EDE, 0x439E009A, 0x7B8F6F36, 0x33BCDFC2, 0x0BADB06E,\n 0x66BCB5BB, 0x5EADDA17, 0x169E6AE3, 0x2E8F054F, 0x86F90B0B, 0xBEE864A7,\n 0xF6DBD453, 0xCECABBFF, 0x6EA2D55C, 0x56B3BAF0, 0x1E800A04, 0x269165A8,\n 0x8EE76BEC, 0xB6F60440, 0xFEC5B4B4, 0xC6D4DB18, 0xABC5DECD, 0x93D4B161,\n 0xDBE70195, 0xE3F66E39, 0x4B80607D, 0x73910FD1, 0x3BA2BF25, 0x03B3D089,\n 0xE180B48F, 0xD991DB23, 0x91A26BD7, 0xA9B3047B, 0x01C50A3F, 0x39D46593,\n 0x71E7D567, 0x49F6BACB, 0x24E7BF1E, 0x1CF6D0B2, 0x54C56046, 0x6CD40FEA,\n 0xC4A201AE, 0xFCB36E02, 0xB480DEF6, 0x8C91B15A, 0x750A600B, 0x4D1B0FA7,\n 0x0528BF53, 0x3D39D0FF, 0x954FDEBB, 0xAD5EB117, 0xE56D01E3, 0xDD7C6E4F,\n 0xB06D6B9A, 0x887C0436, 0xC04FB4C2, 0xF85EDB6E, 0x5028D52A, 0x6839BA86,\n 0x200A0A72, 0x181B65DE, 0xFA2801D8, 0xC2396E74, 0x8A0ADE80, 0xB21BB12C,\n 0x1A6DBF68, 0x227CD0C4, 0x6A4F6030, 0x525E0F9C, 0x3F4F0A49, 0x075E65E5,\n 0x4F6DD511, 0x777CBABD, 0xDF0AB4F9, 0xE71BDB55, 0xAF286BA1, 0x9739040D,\n 0x59F3BFF2, 0x61E2D05E, 0x29D160AA, 0x11C00F06, 0xB9B60142, 0x81A76EEE,\n 0xC994DE1A, 0xF185B1B6, 0x9C94B463, 0xA485DBCF, 0xECB66B3B, 0xD4A70497,\n 0x7CD10AD3, 0x44C0657F, 0x0CF3D58B, 0x34E2BA27, 0xD6D1DE21, 0xEEC0B18D,\n 0xA6F30179, 0x9EE26ED5, 0x36946091, 0x0E850F3D, 0x46B6BFC9, 0x7EA7D065,\n 0x13B6D5B0, 0x2BA7BA1C, 0x63940AE8, 0x5B856544, 0xF3F36B00, 0xCBE204AC,\n 0x83D1B458, 0xBBC0DBF4, 0x425B0AA5, 0x7A4A6509, 0x3279D5FD, 0x0A68BA51,\n 0xA21EB415, 0x9A0FDBB9, 0xD23C6B4D, 0xEA2D04E1, 0x873C0134, 0xBF2D6E98,\n 0xF71EDE6C, 0xCF0FB1C0, 0x6779BF84, 0x5F68D028, 0x175B60DC, 0x2F4A0F70,\n 0xCD796B76, 0xF56804DA, 0xBD5BB42E, 0x854ADB82, 0x2D3CD5C6, 0x152DBA6A,\n 0x5D1E0A9E, 0x650F6532, 0x081E60E7, 0x300F0F4B, 0x783CBFBF, 0x402DD013,\n 0xE85BDE57, 0xD04AB1FB, 0x9879010F, 0xA0686EA3};\n\n/**\n * The following CRC lookup table was generated automagically\n * using the following model parameters:\n *\n * Generator Polynomial = ................. 0x1EDC6F41\n * Generator Polynomial Length = .......... 32 bits\n * Reflected Bits = ....................... TRUE\n * Table Generation Offset = .............. 32 bits\n * Number of Slices = ..................... 8 slices\n * Slice Lengths = ........................ 8 8 8 8 8 8 8 8\n */\nstatic uint32_t crc_tableil8_o72[256] = {\n 0x00000000, 0xEF306B19, 0xDB8CA0C3, 0x34BCCBDA, 0xB2F53777, 0x5DC55C6E,\n 0x697997B4, 0x8649FCAD, 0x6006181F, 0x8F367306, 0xBB8AB8DC, 0x54BAD3C5,\n 0xD2F32F68, 0x3DC34471, 0x097F8FAB, 0xE64FE4B2, 0xC00C303E, 0x2F3C5B27,\n 0x1B8090FD, 0xF4B0FBE4, 0x72F90749, 0x9DC96C50, 0xA975A78A, 0x4645CC93,\n 0xA00A2821, 0x4F3A4338, 0x7B8688E2, 0x94B6E3FB, 0x12FF1F56, 0xFDCF744F,\n 0xC973BF95, 0x2643D48C, 0x85F4168D, 0x6AC47D94, 0x5E78B64E, 0xB148DD57,\n 0x370121FA, 0xD8314AE3, 0xEC8D8139, 0x03BDEA20, 0xE5F20E92, 0x0AC2658B,\n 0x3E7EAE51, 0xD14EC548, 0x570739E5, 0xB83752FC, 0x8C8B9926, 0x63BBF23F,\n 0x45F826B3, 0xAAC84DAA, 0x9E748670, 0x7144ED69, 0xF70D11C4, 0x183D7ADD,\n 0x2C81B107, 0xC3B1DA1E, 0x25FE3EAC, 0xCACE55B5, 0xFE729E6F, 0x1142F576,\n 0x970B09DB, 0x783B62C2, 0x4C87A918, 0xA3B7C201, 0x0E045BEB, 0xE13430F2,\n 0xD588FB28, 0x3AB89031, 0xBCF16C9C, 0x53C10785, 0x677DCC5F, 0x884DA746,\n 0x6E0243F4, 0x813228ED, 0xB58EE337, 0x5ABE882E, 0xDCF77483, 0x33C71F9A,\n 0x077BD440, 0xE84BBF59, 0xCE086BD5, 0x213800CC, 0x1584CB16, 0xFAB4A00F,\n 0x7CFD5CA2, 0x93CD37BB, 0xA771FC61, 0x48419778, 0xAE0E73CA, 0x413E18D3,\n 0x7582D309, 0x9AB2B810, 0x1CFB44BD, 0xF3CB2FA4, 0xC777E47E, 0x28478F67,\n 0x8BF04D66, 0x64C0267F, 0x507CEDA5, 0xBF4C86BC, 0x39057A11, 0xD6351108,\n 0xE289DAD2, 0x0DB9B1CB, 0xEBF65579, 0x04C63E60, 0x307AF5BA, 0xDF4A9EA3,\n 0x5903620E, 0xB6330917, 0x828FC2CD, 0x6DBFA9D4, 0x4BFC7D58, 0xA4CC1641,\n 0x9070DD9B, 0x7F40B682, 0xF9094A2F, 0x16392136, 0x2285EAEC, 0xCDB581F5,\n 0x2BFA6547, 0xC4CA0E5E, 0xF076C584, 0x1F46AE9D, 0x990F5230, 0x763F3929,\n 0x4283F2F3, 0xADB399EA, 0x1C08B7D6, 0xF338DCCF, 0xC7841715, 0x28B47C0C,\n 0xAEFD80A1, 0x41CDEBB8, 0x75712062, 0x9A414B7B, 0x7C0EAFC9, 0x933EC4D0,\n 0xA7820F0A, 0x48B26413, 0xCEFB98BE, 0x21CBF3A7, 0x1577387D, 0xFA475364,\n 0xDC0487E8, 0x3334ECF1, 0x0788272B, 0xE8B84C32, 0x6EF1B09F, 0x81C1DB86,\n 0xB57D105C, 0x5A4D7B45, 0xBC029FF7, 0x5332F4EE, 0x678E3F34, 0x88BE542D,\n 0x0EF7A880, 0xE1C7C399, 0xD57B0843, 0x3A4B635A, 0x99FCA15B, 0x76CCCA42,\n 0x42700198, 0xAD406A81, 0x2B09962C, 0xC439FD35, 0xF08536EF, 0x1FB55DF6,\n 0xF9FAB944, 0x16CAD25D, 0x22761987, 0xCD46729E, 0x4B0F8E33, 0xA43FE52A,\n 0x90832EF0, 0x7FB345E9, 0x59F09165, 0xB6C0FA7C, 0x827C31A6, 0x6D4C5ABF,\n 0xEB05A612, 0x0435CD0B, 0x308906D1, 0xDFB96DC8, 0x39F6897A, 0xD6C6E263,\n 0xE27A29B9, 0x0D4A42A0, 0x8B03BE0D, 0x6433D514, 0x508F1ECE, 0xBFBF75D7,\n 0x120CEC3D, 0xFD3C8724, 0xC9804CFE, 0x26B027E7, 0xA0F9DB4A, 0x4FC9B053,\n 0x7B757B89, 0x94451090, 0x720AF422, 0x9D3A9F3B, 0xA98654E1, 0x46B63FF8,\n 0xC0FFC355, 0x2FCFA84C, 0x1B736396, 0xF443088F, 0xD200DC03, 0x3D30B71A,\n 0x098C7CC0, 0xE6BC17D9, 0x60F5EB74, 0x8FC5806D, 0xBB794BB7, 0x544920AE,\n 0xB206C41C, 0x5D36AF05, 0x698A64DF, 0x86BA0FC6, 0x00F3F36B, 0xEFC39872,\n 0xDB7F53A8, 0x344F38B1, 0x97F8FAB0, 0x78C891A9, 0x4C745A73, 0xA344316A,\n 0x250DCDC7, 0xCA3DA6DE, 0xFE816D04, 0x11B1061D, 0xF7FEE2AF, 0x18CE89B6,\n 0x2C72426C, 0xC3422975, 0x450BD5D8, 0xAA3BBEC1, 0x9E87751B, 0x71B71E02,\n 0x57F4CA8E, 0xB8C4A197, 0x8C786A4D, 0x63480154, 0xE501FDF9, 0x0A3196E0,\n 0x3E8D5D3A, 0xD1BD3623, 0x37F2D291, 0xD8C2B988, 0xEC7E7252, 0x034E194B,\n 0x8507E5E6, 0x6A378EFF, 0x5E8B4525, 0xB1BB2E3C};\n\n/**\n * The following CRC lookup table was generated automagically\n * using the following model parameters:\n *\n * Generator Polynomial = ................. 0x1EDC6F41\n * Generator Polynomial Length = .......... 32 bits\n * Reflected Bits = ....................... TRUE\n * Table Generation Offset = .............. 32 bits\n * Number of Slices = ..................... 8 slices\n * Slice Lengths = ........................ 8 8 8 8 8 8 8 8\n */\nstatic uint32_t crc_tableil8_o80[256] = {\n 0x00000000, 0x68032CC8, 0xD0065990, 0xB8057558, 0xA5E0C5D1, 0xCDE3E919,\n 0x75E69C41, 0x1DE5B089, 0x4E2DFD53, 0x262ED19B, 0x9E2BA4C3, 0xF628880B,\n 0xEBCD3882, 0x83CE144A, 0x3BCB6112, 0x53C84DDA, 0x9C5BFAA6, 0xF458D66E,\n 0x4C5DA336, 0x245E8FFE, 0x39BB3F77, 0x51B813BF, 0xE9BD66E7, 0x81BE4A2F,\n 0xD27607F5, 0xBA752B3D, 0x02705E65, 0x6A7372AD, 0x7796C224, 0x1F95EEEC,\n 0xA7909BB4, 0xCF93B77C, 0x3D5B83BD, 0x5558AF75, 0xED5DDA2D, 0x855EF6E5,\n 0x98BB466C, 0xF0B86AA4, 0x48BD1FFC, 0x20BE3334, 0x73767EEE, 0x1B755226,\n 0xA370277E, 0xCB730BB6, 0xD696BB3F, 0xBE9597F7, 0x0690E2AF, 0x6E93CE67,\n 0xA100791B, 0xC90355D3, 0x7106208B, 0x19050C43, 0x04E0BCCA, 0x6CE39002,\n 0xD4E6E55A, 0xBCE5C992, 0xEF2D8448, 0x872EA880, 0x3F2BDDD8, 0x5728F110,\n 0x4ACD4199, 0x22CE6D51, 0x9ACB1809, 0xF2C834C1, 0x7AB7077A, 0x12B42BB2,\n 0xAAB15EEA, 0xC2B27222, 0xDF57C2AB, 0xB754EE63, 0x0F519B3B, 0x6752B7F3,\n 0x349AFA29, 0x5C99D6E1, 0xE49CA3B9, 0x8C9F8F71, 0x917A3FF8, 0xF9791330,\n 0x417C6668, 0x297F4AA0, 0xE6ECFDDC, 0x8EEFD114, 0x36EAA44C, 0x5EE98884,\n 0x430C380D, 0x2B0F14C5, 0x930A619D, 0xFB094D55, 0xA8C1008F, 0xC0C22C47,\n 0x78C7591F, 0x10C475D7, 0x0D21C55E, 0x6522E996, 0xDD279CCE, 0xB524B006,\n 0x47EC84C7, 0x2FEFA80F, 0x97EADD57, 0xFFE9F19F, 0xE20C4116, 0x8A0F6DDE,\n 0x320A1886, 0x5A09344E, 0x09C17994, 0x61C2555C, 0xD9C72004, 0xB1C40CCC,\n 0xAC21BC45, 0xC422908D, 0x7C27E5D5, 0x1424C91D, 0xDBB77E61, 0xB3B452A9,\n 0x0BB127F1, 0x63B20B39, 0x7E57BBB0, 0x16549778, 0xAE51E220, 0xC652CEE8,\n 0x959A8332, 0xFD99AFFA, 0x459CDAA2, 0x2D9FF66A, 0x307A46E3, 0x58796A2B,\n 0xE07C1F73, 0x887F33BB, 0xF56E0EF4, 0x9D6D223C, 0x25685764, 0x4D6B7BAC,\n 0x508ECB25, 0x388DE7ED, 0x808892B5, 0xE88BBE7D, 0xBB43F3A7, 0xD340DF6F,\n 0x6B45AA37, 0x034686FF, 0x1EA33676, 0x76A01ABE, 0xCEA56FE6, 0xA6A6432E,\n 0x6935F452, 0x0136D89A, 0xB933ADC2, 0xD130810A, 0xCCD53183, 0xA4D61D4B,\n 0x1CD36813, 0x74D044DB, 0x27180901, 0x4F1B25C9, 0xF71E5091, 0x9F1D7C59,\n 0x82F8CCD0, 0xEAFBE018, 0x52FE9540, 0x3AFDB988, 0xC8358D49, 0xA036A181,\n 0x1833D4D9, 0x7030F811, 0x6DD54898, 0x05D66450, 0xBDD31108, 0xD5D03DC0,\n 0x8618701A, 0xEE1B5CD2, 0x561E298A, 0x3E1D0542, 0x23F8B5CB, 0x4BFB9903,\n 0xF3FEEC5B, 0x9BFDC093, 0x546E77EF, 0x3C6D5B27, 0x84682E7F, 0xEC6B02B7,\n 0xF18EB23E, 0x998D9EF6, 0x2188EBAE, 0x498BC766, 0x1A438ABC, 0x7240A674,\n 0xCA45D32C, 0xA246FFE4, 0xBFA34F6D, 0xD7A063A5, 0x6FA516FD, 0x07A63A35,\n 0x8FD9098E, 0xE7DA2546, 0x5FDF501E, 0x37DC7CD6, 0x2A39CC5F, 0x423AE097,\n 0xFA3F95CF, 0x923CB907, 0xC1F4F4DD, 0xA9F7D815, 0x11F2AD4D, 0x79F18185,\n 0x6414310C, 0x0C171DC4, 0xB412689C, 0xDC114454, 0x1382F328, 0x7B81DFE0,\n 0xC384AAB8, 0xAB878670, 0xB66236F9, 0xDE611A31, 0x66646F69, 0x0E6743A1,\n 0x5DAF0E7B, 0x35AC22B3, 0x8DA957EB, 0xE5AA7B23, 0xF84FCBAA, 0x904CE762,\n 0x2849923A, 0x404ABEF2, 0xB2828A33, 0xDA81A6FB, 0x6284D3A3, 0x0A87FF6B,\n 0x17624FE2, 0x7F61632A, 0xC7641672, 0xAF673ABA, 0xFCAF7760, 0x94AC5BA8,\n 0x2CA92EF0, 0x44AA0238, 0x594FB2B1, 0x314C9E79, 0x8949EB21, 0xE14AC7E9,\n 0x2ED97095, 0x46DA5C5D, 0xFEDF2905, 0x96DC05CD, 0x8B39B544, 0xE33A998C,\n 0x5B3FECD4, 0x333CC01C, 0x60F48DC6, 0x08F7A10E, 0xB0F2D456, 0xD8F1F89E,\n 0xC5144817, 0xAD1764DF, 0x15121187, 0x7D113D4F};\n\n/**\n * The following CRC lookup table was generated automagically\n * using the following model parameters:\n *\n * Generator Polynomial = ................. 0x1EDC6F41\n * Generator Polynomial Length = .......... 32 bits\n * Reflected Bits = ....................... TRUE\n * Table Generation Offset = .............. 32 bits\n * Number of Slices = ..................... 8 slices\n * Slice Lengths = ........................ 8 8 8 8 8 8 8 8\n */\nstatic uint32_t crc_tableil8_o88[256] = {\n 0x00000000, 0x493C7D27, 0x9278FA4E, 0xDB448769, 0x211D826D, 0x6821FF4A,\n 0xB3657823, 0xFA590504, 0x423B04DA, 0x0B0779FD, 0xD043FE94, 0x997F83B3,\n 0x632686B7, 0x2A1AFB90, 0xF15E7CF9, 0xB86201DE, 0x847609B4, 0xCD4A7493,\n 0x160EF3FA, 0x5F328EDD, 0xA56B8BD9, 0xEC57F6FE, 0x37137197, 0x7E2F0CB0,\n 0xC64D0D6E, 0x8F717049, 0x5435F720, 0x1D098A07, 0xE7508F03, 0xAE6CF224,\n 0x7528754D, 0x3C14086A, 0x0D006599, 0x443C18BE, 0x9F789FD7, 0xD644E2F0,\n 0x2C1DE7F4, 0x65219AD3, 0xBE651DBA, 0xF759609D, 0x4F3B6143, 0x06071C64,\n 0xDD439B0D, 0x947FE62A, 0x6E26E32E, 0x271A9E09, 0xFC5E1960, 0xB5626447,\n 0x89766C2D, 0xC04A110A, 0x1B0E9663, 0x5232EB44, 0xA86BEE40, 0xE1579367,\n 0x3A13140E, 0x732F6929, 0xCB4D68F7, 0x827115D0, 0x593592B9, 0x1009EF9E,\n 0xEA50EA9A, 0xA36C97BD, 0x782810D4, 0x31146DF3, 0x1A00CB32, 0x533CB615,\n 0x8878317C, 0xC1444C5B, 0x3B1D495F, 0x72213478, 0xA965B311, 0xE059CE36,\n 0x583BCFE8, 0x1107B2CF, 0xCA4335A6, 0x837F4881, 0x79264D85, 0x301A30A2,\n 0xEB5EB7CB, 0xA262CAEC, 0x9E76C286, 0xD74ABFA1, 0x0C0E38C8, 0x453245EF,\n 0xBF6B40EB, 0xF6573DCC, 0x2D13BAA5, 0x642FC782, 0xDC4DC65C, 0x9571BB7B,\n 0x4E353C12, 0x07094135, 0xFD504431, 0xB46C3916, 0x6F28BE7F, 0x2614C358,\n 0x1700AEAB, 0x5E3CD38C, 0x857854E5, 0xCC4429C2, 0x361D2CC6, 0x7F2151E1,\n 0xA465D688, 0xED59ABAF, 0x553BAA71, 0x1C07D756, 0xC743503F, 0x8E7F2D18,\n 0x7426281C, 0x3D1A553B, 0xE65ED252, 0xAF62AF75, 0x9376A71F, 0xDA4ADA38,\n 0x010E5D51, 0x48322076, 0xB26B2572, 0xFB575855, 0x2013DF3C, 0x692FA21B,\n 0xD14DA3C5, 0x9871DEE2, 0x4335598B, 0x0A0924AC, 0xF05021A8, 0xB96C5C8F,\n 0x6228DBE6, 0x2B14A6C1, 0x34019664, 0x7D3DEB43, 0xA6796C2A, 0xEF45110D,\n 0x151C1409, 0x5C20692E, 0x8764EE47, 0xCE589360, 0x763A92BE, 0x3F06EF99,\n 0xE44268F0, 0xAD7E15D7, 0x572710D3, 0x1E1B6DF4, 0xC55FEA9D, 0x8C6397BA,\n 0xB0779FD0, 0xF94BE2F7, 0x220F659E, 0x6B3318B9, 0x916A1DBD, 0xD856609A,\n 0x0312E7F3, 0x4A2E9AD4, 0xF24C9B0A, 0xBB70E62D, 0x60346144, 0x29081C63,\n 0xD3511967, 0x9A6D6440, 0x4129E329, 0x08159E0E, 0x3901F3FD, 0x703D8EDA,\n 0xAB7909B3, 0xE2457494, 0x181C7190, 0x51200CB7, 0x8A648BDE, 0xC358F6F9,\n 0x7B3AF727, 0x32068A00, 0xE9420D69, 0xA07E704E, 0x5A27754A, 0x131B086D,\n 0xC85F8F04, 0x8163F223, 0xBD77FA49, 0xF44B876E, 0x2F0F0007, 0x66337D20,\n 0x9C6A7824, 0xD5560503, 0x0E12826A, 0x472EFF4D, 0xFF4CFE93, 0xB67083B4,\n 0x6D3404DD, 0x240879FA, 0xDE517CFE, 0x976D01D9, 0x4C2986B0, 0x0515FB97,\n 0x2E015D56, 0x673D2071, 0xBC79A718, 0xF545DA3F, 0x0F1CDF3B, 0x4620A21C,\n 0x9D642575, 0xD4585852, 0x6C3A598C, 0x250624AB, 0xFE42A3C2, 0xB77EDEE5,\n 0x4D27DBE1, 0x041BA6C6, 0xDF5F21AF, 0x96635C88, 0xAA7754E2, 0xE34B29C5,\n 0x380FAEAC, 0x7133D38B, 0x8B6AD68F, 0xC256ABA8, 0x19122CC1, 0x502E51E6,\n 0xE84C5038, 0xA1702D1F, 0x7A34AA76, 0x3308D751, 0xC951D255, 0x806DAF72,\n 0x5B29281B, 0x1215553C, 0x230138CF, 0x6A3D45E8, 0xB179C281, 0xF845BFA6,\n 0x021CBAA2, 0x4B20C785, 0x906440EC, 0xD9583DCB, 0x613A3C15, 0x28064132,\n 0xF342C65B, 0xBA7EBB7C, 0x4027BE78, 0x091BC35F, 0xD25F4436, 0x9B633911,\n 0xA777317B, 0xEE4B4C5C, 0x350FCB35, 0x7C33B612, 0x866AB316, 0xCF56CE31,\n 0x14124958, 0x5D2E347F, 0xE54C35A1, 0xAC704886, 0x7734CFEF, 0x3E08B2C8,\n 0xC451B7CC, 0x8D6DCAEB, 0x56294D82, 0x1F1530A5};\n\n/**\n * Implementations adapted from Intel's Slicing By 8 Sourceforge Project\n * http://sourceforge.net/projects/slicing-by-8/\n * http://www.evanjones.ca/crc32c.html\n */\nstatic inline uint32_t crc32c_slicing8(const void *data, size_t len,\n uint32_t crc) {\n const uint8_t *p_buf = (const uint8_t *)data;\n\n /* Handle leading misaligned bytes */\n size_t init_bytes =\n (sizeof(int32_t) - (intptr_t)p_buf) & (sizeof(int32_t) - 1);\n if (len < init_bytes) {\n init_bytes = len;\n }\n for (size_t li = 0; li < init_bytes; li++) {\n crc = crc_tableil8_o32[(crc ^ *p_buf++) & 0x000000FF] ^ (crc >> 8);\n }\n\n len -= init_bytes;\n size_t running_length = len & ~(sizeof(uint64_t) - 1);\n size_t end_bytes = len - running_length;\n\n for (size_t li = 0; li < running_length / 8; li++) {\n uint32_t term1, term2;\n\n crc ^= *(uint32_t *)p_buf;\n p_buf += 4;\n term1 = crc_tableil8_o88[crc & 0x000000FF] ^\n crc_tableil8_o80[(crc >> 8) & 0x000000FF];\n term2 = crc >> 16;\n crc = term1 ^ crc_tableil8_o72[term2 & 0x000000FF] ^\n crc_tableil8_o64[(term2 >> 8) & 0x000000FF];\n term1 = crc_tableil8_o56[(*(uint32_t *)p_buf) & 0x000000FF] ^\n crc_tableil8_o48[((*(uint32_t *)p_buf) >> 8) & 0x000000FF];\n\n term2 = (*(uint32_t *)p_buf) >> 16;\n crc = crc ^ term1 ^ crc_tableil8_o40[term2 & 0x000000FF] ^\n crc_tableil8_o32[(term2 >> 8) & 0x000000FF];\n p_buf += 4;\n }\n\n for (size_t li = 0; li < end_bytes; li++) {\n crc = crc_tableil8_o32[(crc ^ *p_buf++) & 0x000000FF] ^ (crc >> 8);\n }\n return crc;\n}\n#endif // !__SSE4_2__\n\n#if defined(__SSE4_2__)\n#if defined(AILEGO_M64)\nstatic inline uint32_t crc32c_sse42(const void *data, size_t len,\n uint32_t crc) {\n const uint8_t *first = (const uint8_t *)data;\n const uint8_t *last = first + ((len >> 3) << 3);\n\n for (; first != last; first += 8) {\n crc = (uint32_t)_mm_crc32_u64(crc, *(uint64_t *)first);\n }\n switch (((uint8_t *)data + len) - last) {\n case 1:\n crc = _mm_crc32_u8(crc, *last);\n break;\n case 2:\n crc = _mm_crc32_u16(crc, *(uint16_t *)last);\n break;\n case 3:\n crc = _mm_crc32_u16(crc, *(uint16_t *)last);\n crc = _mm_crc32_u8(crc, *(last + 2));\n break;\n case 4:\n crc = _mm_crc32_u32(crc, *(uint32_t *)last);\n break;\n case 5:\n crc = _mm_crc32_u32(crc, *(uint32_t *)last);\n crc = _mm_crc32_u8(crc, *(last + 4));\n break;\n case 6:\n crc = _mm_crc32_u32(crc, *(uint32_t *)last);\n crc = _mm_crc32_u16(crc, *(uint16_t *)(last + 4));\n break;\n case 7:\n crc = _mm_crc32_u32(crc, *(uint32_t *)last);\n crc = _mm_crc32_u16(crc, *(uint16_t *)(last + 4));\n crc = _mm_crc32_u8(crc, *(last + 6));\n break;\n }\n return crc;\n}\n#else\nstatic inline uint32_t crc32c_sse42(const void *data, size_t len,\n uint32_t crc) {\n const uint8_t *first = (const uint8_t *)data;\n const uint8_t *last = first + ((len >> 2) << 2);\n\n for (; first != last; first += 4) {\n crc = _mm_crc32_u32(crc, *(uint32_t *)first);\n }\n switch (((uint8_t *)data + len) - last) {\n case 1:\n crc = _mm_crc32_u8(crc, *last);\n break;\n case 2:\n crc = _mm_crc32_u16(crc, *(uint16_t *)last);\n break;\n case 3:\n crc = _mm_crc32_u16(crc, *(uint16_t *)last);\n crc = _mm_crc32_u8(crc, *(last + 2));\n break;\n }\n return crc;\n}\n#endif // AILEGO_M64\n#endif // __SSE4_2__\n\n#if defined(__ARM_FEATURE_CRC32)\nstatic inline uint32_t crc32c_neon(const void *data, size_t len, uint32_t crc) {\n const uint8_t *first = (const uint8_t *)data;\n const uint8_t *last = first + ((len >> 3) << 3);\n\n for (; first != last; first += 8) {\n crc = __crc32cd(crc, *(uint64_t *)first);\n }\n switch (((uint8_t *)data + len) - last) {\n case 1:\n crc = __crc32cb(crc, *last);\n break;\n case 2:\n crc = __crc32ch(crc, *(uint16_t *)last);\n break;\n case 3:\n crc = __crc32ch(crc, *(uint16_t *)last);\n crc = __crc32cb(crc, *(last + 2));\n break;\n case 4:\n crc = __crc32cw(crc, *(uint32_t *)last);\n break;\n case 5:\n crc = __crc32cw(crc, *(uint32_t *)last);\n crc = __crc32cb(crc, *(last + 4));\n break;\n case 6:\n crc = __crc32cw(crc, *(uint32_t *)last);\n crc = __crc32ch(crc, *(uint16_t *)(last + 4));\n break;\n case 7:\n crc = __crc32cw(crc, *(uint32_t *)last);\n crc = __crc32ch(crc, *(uint16_t *)(last + 4));\n crc = __crc32cb(crc, *(last + 6));\n break;\n }\n return crc;\n}\n#endif // __ARM_FEATURE_CRC32\n\nnamespace zvec {\nnamespace ailego {\n\nuint32_t Crc32c::Hash(const void *data, size_t len, uint32_t crc) {\n#if defined(__SSE4_2__)\n return crc32c_sse42(data, len, crc);\n#elif defined(__ARM_FEATURE_CRC32)\n return crc32c_neon(data, len, crc);\n#else\n return crc32c_slicing8(data, len, crc);\n#endif\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/internal/cpu_features.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"cpu_features.h\"\n#include \n\n#if !defined(_MSC_VER) && !defined(__ARM_ARCH)\n#include \n#endif\n\nnamespace zvec {\nnamespace ailego {\nnamespace internal {\n\n//\n// REFER: https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/\n// tree/arch/x86/include/asm/cpufeatures.h\n// https://software.intel.com/sites/default/files/managed/c5/15/\n// architecture-instruction-set-extensions-programming-reference.pdf\n//\n\nCpuFeatures::CpuFlags CpuFeatures::flags_;\n\n#if defined(_MSC_VER)\nCpuFeatures::CpuFlags::CpuFlags(void)\n : L1_ECX(0), L1_EDX(0), L7_EBX(0), L7_ECX(0), L7_EDX(0) {\n int l1[4] = {0, 0, 0, 0};\n int l7[4] = {0, 0, 0, 0};\n\n __cpuidex(l1, 1, 0);\n __cpuidex(l7, 7, 0);\n L1_ECX = l1[2];\n L1_EDX = l1[3];\n L7_EBX = l7[1];\n L7_ECX = l7[2];\n L7_EDX = l7[3];\n}\n#elif !defined(__ARM_ARCH)\nCpuFeatures::CpuFlags::CpuFlags(void)\n : L1_ECX(0), L1_EDX(0), L7_EBX(0), L7_ECX(0), L7_EDX(0) {\n uint32_t eax, ebx, ecx, edx;\n\n if (__get_cpuid(1, &eax, &ebx, &ecx, &edx)) {\n L1_ECX = ecx;\n L1_EDX = edx;\n }\n if (__get_cpuid_max(0, NULL) >= 7) {\n __cpuid_count(7, 0, eax, ebx, ecx, edx);\n L7_EBX = ebx;\n L7_ECX = ecx;\n L7_EDX = edx;\n }\n}\n#else\nCpuFeatures::CpuFlags::CpuFlags(void)\n : L1_ECX(0), L1_EDX(0), L7_EBX(0), L7_ECX(0), L7_EDX(0) {}\n#endif\n\n//! 16-bit FP conversions\nbool CpuFeatures::F16C(void) {\n return !!(flags_.L1_ECX & (1u << 29));\n}\n\n//! Multimedia Extensions\nbool CpuFeatures::MMX(void) {\n return !!(flags_.L1_EDX & (1u << 23));\n}\n\n//! Streaming SIMD Extensions\nbool CpuFeatures::SSE(void) {\n return !!(flags_.L1_EDX & (1u << 25));\n}\n\n//! Streaming SIMD Extensions 2\nbool CpuFeatures::SSE2(void) {\n return !!(flags_.L1_EDX & (1u << 26));\n}\n\n//! Streaming SIMD Extensions 3\nbool CpuFeatures::SSE3(void) {\n return !!(flags_.L1_ECX & (1u << 0));\n}\n\n//! Supplemental Streaming SIMD Extensions 3\nbool CpuFeatures::SSSE3(void) {\n return !!(flags_.L1_ECX & (1u << 9));\n}\n\n//! Streaming SIMD Extensions 4.1\nbool CpuFeatures::SSE4_1(void) {\n return !!(flags_.L1_ECX & (1u << 19));\n}\n\n//! Streaming SIMD Extensions 4.2\nbool CpuFeatures::SSE4_2(void) {\n return !!(flags_.L1_ECX & (1u << 20));\n}\n\n//! Advanced Vector Extensions\nbool CpuFeatures::AVX(void) {\n return !!(flags_.L1_ECX & (1u << 28));\n}\n\n//! Advanced Vector Extensions 2\nbool CpuFeatures::AVX2(void) {\n return !!(flags_.L7_EBX & (1u << 5));\n}\n\n//! AVX-512 Foundation\nbool CpuFeatures::AVX512F(void) {\n return !!(flags_.L7_EBX & (1u << 16));\n}\n\n//! AVX-512 DQ (Double/Quad granular) Instructions\nbool CpuFeatures::AVX512DQ(void) {\n return !!(flags_.L7_EBX & (1u << 17));\n}\n\n//! AVX-512 Prefetch\nbool CpuFeatures::AVX512PF(void) {\n return !!(flags_.L7_EBX & (1u << 26));\n}\n\n//! AVX-512 Exponential and Reciprocal\nbool CpuFeatures::AVX512ER(void) {\n return !!(flags_.L7_EBX & (1u << 27));\n}\n\n//! AVX-512 Conflict Detection\nbool CpuFeatures::AVX512CD(void) {\n return !!(flags_.L7_EBX & (1u << 28));\n}\n\n//! AVX-512 BW (Byte/Word granular) Instructions\nbool CpuFeatures::AVX512BW(void) {\n return !!(flags_.L7_EBX & (1u << 30));\n}\n\n//! AVX-512 VL (128/256 Vector Length) Extensions\nbool CpuFeatures::AVX512VL(void) {\n return !!(flags_.L7_EBX & (1u << 31));\n}\n\n//! AVX-512 Integer Fused Multiply-Add instructions\nbool CpuFeatures::AVX512_IFMA(void) {\n return !!(flags_.L7_EBX & (1u << 21));\n}\n\n//! AVX512 Vector Bit Manipulation instructions\nbool CpuFeatures::AVX512_VBMI(void) {\n return !!(flags_.L7_ECX & (1u << 1));\n}\n\n//! Additional AVX512 Vector Bit Manipulation Instructions\nbool CpuFeatures::AVX512_VBMI2(void) {\n return !!(flags_.L7_ECX & (1u << 6));\n}\n\n//! Vector Neural Network Instructions\nbool CpuFeatures::AVX512_VNNI(void) {\n return !!(flags_.L7_ECX & (1u << 11));\n}\n\n//! Support for VPOPCNT[B,W] and VPSHUF-BITQMB instructions\nbool CpuFeatures::AVX512_BITALG(void) {\n return !!(flags_.L7_ECX & (1u << 12));\n}\n\n//! POPCNT for vectors of DW/QW\nbool CpuFeatures::AVX512_VPOPCNTDQ(void) {\n return !!(flags_.L7_ECX & (1u << 14));\n}\n\n//! AVX-512 Neural Network Instructions\nbool CpuFeatures::AVX512_4VNNIW(void) {\n return !!(flags_.L7_EDX & (1u << 2));\n}\n\n//! AVX-512 Multiply Accumulation Single precision\nbool CpuFeatures::AVX512_4FMAPS(void) {\n return !!(flags_.L7_EDX & (1u << 3));\n}\n\n//! AVX-512 FP16 instructions\nbool CpuFeatures::AVX512_FP16(void) {\n return !!(flags_.L7_EDX & (1u << 23));\n}\n\n//! CMPXCHG8 instruction\nbool CpuFeatures::CX8(void) {\n return !!(flags_.L1_EDX & (1u << 8));\n}\n\n//! CMPXCHG16B instruction\nbool CpuFeatures::CX16(void) {\n return !!(flags_.L1_ECX & (1u << 13));\n}\n\n//! PCLMULQDQ instruction\nbool CpuFeatures::PCLMULQDQ(void) {\n return !!(flags_.L1_ECX & (1u << 1));\n}\n\n//! Carry-Less Multiplication Double Quadword\nbool CpuFeatures::VPCLMULQDQ(void) {\n return !!(flags_.L7_ECX & (1u << 10));\n}\n\n//! CMOV instructions (plus FCMOVcc, FCOMI with FPU)\nbool CpuFeatures::CMOV(void) {\n return !!(flags_.L1_EDX & (1u << 15));\n}\n\n//! MOVBE instruction\nbool CpuFeatures::MOVBE(void) {\n return !!(flags_.L1_ECX & (1u << 22));\n}\n\n//! Enhanced REP MOVSB/STOSB instructions\nbool CpuFeatures::ERMS(void) {\n return !!(flags_.L7_EBX & (1u << 9));\n}\n\n//! POPCNT instruction\nbool CpuFeatures::POPCNT(void) {\n return !!(flags_.L1_ECX & (1u << 23));\n}\n\n//! XSAVE/XRSTOR/XSETBV/XGETBV instructions\nbool CpuFeatures::XSAVE(void) {\n return !!(flags_.L1_ECX & (1u << 26));\n}\n\n//! Fused multiply-add\nbool CpuFeatures::FMA(void) {\n return !!(flags_.L1_ECX & (1u << 12));\n}\n\n//! ADCX and ADOX instructions\nbool CpuFeatures::ADX(void) {\n return !!(flags_.L7_EBX & (1u << 19));\n}\n\n//! Galois Field New Instructions\nbool CpuFeatures::GFNI(void) {\n return !!(flags_.L7_ECX & (1u << 8));\n}\n\n//! AES instructions\nbool CpuFeatures::AES(void) {\n return !!(flags_.L1_ECX & (1u << 25));\n}\n\n//! Vector AES\nbool CpuFeatures::VAES(void) {\n return !!(flags_.L7_ECX & (1u << 9));\n}\n\n//! RDSEED instruction\nbool CpuFeatures::RDSEED(void) {\n return !!(flags_.L7_EBX & (1u << 18));\n}\n\n//! RDRAND instruction\nbool CpuFeatures::RDRAND(void) {\n return !!(flags_.L1_ECX & (1u << 30));\n}\n\n//! SHA1/SHA256 Instruction Extensions\nbool CpuFeatures::SHA(void) {\n return !!(flags_.L7_EBX & (1u << 29));\n}\n\n//! 1st group bit manipulation extensions\nbool CpuFeatures::BMI1(void) {\n return !!(flags_.L7_EBX & (1u << 3));\n}\n\n//! 2nd group bit manipulation extensions\nbool CpuFeatures::BMI2(void) {\n return !!(flags_.L7_EBX & (1u << 8));\n}\n\n//! CLFLUSH instruction\nbool CpuFeatures::CLFLUSH(void) {\n return !!(flags_.L1_EDX & (1u << 19));\n}\n\n//! CLFLUSHOPT instruction\nbool CpuFeatures::CLFLUSHOPT(void) {\n return !!(flags_.L7_EBX & (1u << 23));\n}\n\n//! CLWB instruction\nbool CpuFeatures::CLWB(void) {\n return !!(flags_.L7_EBX & (1u << 24));\n}\n\n//! RDPID instruction\nbool CpuFeatures::RDPID(void) {\n return !!(flags_.L7_ECX & (1u << 22));\n}\n\n//! Onboard FPU\nbool CpuFeatures::FPU(void) {\n return !!(flags_.L1_EDX & (1u << 0));\n}\n\n//! Hyper-Threading\nbool CpuFeatures::HT(void) {\n return !!(flags_.L1_EDX & (1u << 28));\n}\n\n//! Hardware virtualization\nbool CpuFeatures::VMX(void) {\n return !!(flags_.L1_ECX & (1u << 5));\n}\n\n// ！Running on a hypervisor\nbool CpuFeatures::HYPERVISOR(void) {\n return !!(flags_.L1_ECX & (1u << 31));\n}\n\nconst char *CpuFeatures::Intrinsics(void) {\n return \"\"\n#if defined(__ARM_NEON)\n \"Neon\"\n#if defined(__ARM_FEATURE_CRC32)\n \"+CRC\"\n#endif\n#if defined(__ARM_FEATURE_FP16_SCALAR_ARITHMETIC) || \\\n defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)\n \"+FP16\"\n#endif\n#elif defined(__AVX512F__)\n \"AVX512F\"\n#if defined(__AVX512VL__)\n \"+AVX512VL\"\n#endif\n#if defined(__AVX512BW__)\n \"+AVX512BW\"\n#endif\n#if defined(__AVX512DQ__)\n \"+AVX512DQ\"\n#endif\n#if defined(__AVX512CD__)\n \"+AVX512CD\"\n#endif\n#if defined(__AVX512ER__)\n \"+AVX512ER\"\n#endif\n#if defined(__AVX512PF__)\n \"+AVX512PF\"\n#endif\n#if defined(__AVX512IFMA__)\n \"+AVX512IFMA\"\n#endif\n#if defined(__AVX512VBMI__)\n \"+AVX512VBMI\"\n#endif\n#if defined(__AVX512VBMI2__)\n \"+AVX512VBMI2\"\n#endif\n#if defined(__AVX512VNNI__)\n \"+AVX512VNNI\"\n#endif\n#if defined(__AVX512BITALG__)\n \"+AVX512BITALG\"\n#endif\n#if defined(__AVX512VPOPCNTDQ__)\n \"+AVX512VPOPCNTDQ\"\n#endif\n#if defined(__AVX512FP16__)\n \"+AVX512FP16\"\n#endif\n#elif defined(__AVX2__)\n \"AVX2\"\n#elif defined(__AVX__)\n \"AVX\"\n#elif defined(__SSE4_2__)\n \"SSE4.2\"\n#elif defined(__SSE4_1__)\n \"SSE4.1\"\n#elif defined(__SSSE3__)\n \"SSSE3\"\n#elif defined(__SSE3__)\n \"SSE3\"\n#elif defined(__SSE2__)\n \"SSE2\"\n#elif defined(__SSE__)\n \"SSE\"\n#elif defined(__MMX__)\n \"MMX\"\n#endif\n#if defined(__FMA__)\n \"+FMA\"\n#endif\n#if defined(__BMI2__)\n \"+BMI2\"\n#elif defined(__BMI__)\n \"+BMI\"\n#endif\n#if defined(__F16C__)\n \"+F16C\"\n#endif\n ;\n}\n\nCpuFeatures::StaticFlags CpuFeatures::static_flags_;\n} // namespace internal\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/internal/cpu_features.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \nnamespace zvec {\nnamespace ailego {\nnamespace internal {\n\n/*! Cpu Features\n */\nclass CpuFeatures {\n public:\n //! 16-bit FP conversions\n static bool F16C(void);\n\n //! Multimedia Extensions\n static bool MMX(void);\n\n //! Streaming SIMD Extensions\n static bool SSE(void);\n\n //! Streaming SIMD Extensions 2\n static bool SSE2(void);\n\n //! Streaming SIMD Extensions 3\n static bool SSE3(void);\n\n //! Supplemental Streaming SIMD Extensions 3\n static bool SSSE3(void);\n\n //! Streaming SIMD Extensions 4.1\n static bool SSE4_1(void);\n\n //! Streaming SIMD Extensions 4.2\n static bool SSE4_2(void);\n\n //! Advanced Vector Extensions\n static bool AVX(void);\n\n //! Advanced Vector Extensions 2\n static bool AVX2(void);\n\n //! AVX-512 Foundation\n static bool AVX512F(void);\n\n //! AVX-512 DQ (Double/Quad granular) Instructions\n static bool AVX512DQ(void);\n\n //! AVX-512 Prefetch\n static bool AVX512PF(void);\n\n //! AVX-512 Exponential and Reciprocal\n static bool AVX512ER(void);\n\n //! AVX-512 Conflict Detection\n static bool AVX512CD(void);\n\n //! AVX-512 BW (Byte/Word granular) Instructions\n static bool AVX512BW(void);\n\n //! AVX-512 VL (128/256 Vector Length) Extensions\n static bool AVX512VL(void);\n\n //! AVX-512 Integer Fused Multiply-Add instructions\n static bool AVX512_IFMA(void);\n\n //! AVX512 Vector Bit Manipulation instructions\n static bool AVX512_VBMI(void);\n\n //! Additional AVX512 Vector Bit Manipulation Instructions\n static bool AVX512_VBMI2(void);\n\n //! Vector Neural Network Instructions\n static bool AVX512_VNNI(void);\n\n //! Support for VPOPCNT[B,W] and VPSHUF-BITQMB instructions\n static bool AVX512_BITALG(void);\n\n //! POPCNT for vectors of DW/QW\n static bool AVX512_VPOPCNTDQ(void);\n\n //! AVX-512 Neural Network Instructions\n static bool AVX512_4VNNIW(void);\n\n //! AVX-512 Multiply Accumulation Single precision\n static bool AVX512_4FMAPS(void);\n\n //! AVX-512 FP16 instructions\n static bool AVX512_FP16(void);\n\n //! CMPXCHG8 instruction\n static bool CX8(void);\n\n //! CMPXCHG16B instruction\n static bool CX16(void);\n\n //! PCLMULQDQ instruction\n static bool PCLMULQDQ(void);\n\n //! Carry-Less Multiplication Double Quadword\n static bool VPCLMULQDQ(void);\n\n //! CMOV instructions (plus FCMOVcc, FCOMI with FPU)\n static bool CMOV(void);\n\n //! MOVBE instruction\n static bool MOVBE(void);\n\n //! Enhanced REP MOVSB/STOSB instructions\n static bool ERMS(void);\n\n //! POPCNT instruction\n static bool POPCNT(void);\n\n //! XSAVE/XRSTOR/XSETBV/XGETBV instructions\n static bool XSAVE(void);\n\n //! Fused multiply-add\n static bool FMA(void);\n\n //! ADCX and ADOX instructions\n static bool ADX(void);\n\n //! Galois Field New Instructions\n static bool GFNI(void);\n\n //! AES instructions\n static bool AES(void);\n\n //! Vector AES\n static bool VAES(void);\n\n //! RDSEED instruction\n static bool RDSEED(void);\n\n //! RDRAND instruction\n static bool RDRAND(void);\n\n //! SHA1/SHA256 Instruction Extensions\n static bool SHA(void);\n\n //! 1st group bit manipulation extensions\n static bool BMI1(void);\n\n //! 2nd group bit manipulation extensions\n static bool BMI2(void);\n\n //! CLFLUSH instruction\n static bool CLFLUSH(void);\n\n //! CLFLUSHOPT instruction\n static bool CLFLUSHOPT(void);\n\n //! CLWB instruction\n static bool CLWB(void);\n\n //! RDPID instruction\n static bool RDPID(void);\n\n //! Onboard FPU\n static bool FPU(void);\n\n //! Hyper-Threading\n static bool HT(void);\n\n //! Hardware virtualization\n static bool VMX(void);\n\n // ！Running on a hypervisor\n static bool HYPERVISOR(void);\n\n //! Intrinsics of compiling\n static const char *Intrinsics(void);\n\n private:\n struct CpuFlags {\n //! Constructor\n CpuFlags(void);\n\n //! Members\n uint32_t L1_ECX;\n uint32_t L1_EDX;\n uint32_t L7_EBX;\n uint32_t L7_ECX;\n uint32_t L7_EDX;\n };\n\n //! Static Members\n static CpuFlags flags_;\n\n public:\n struct StaticFlags {\n //! 16-bit FP conversions\n bool F16C = CpuFeatures::F16C();\n\n //! Multimedia Extensions\n bool MMX = CpuFeatures::MMX();\n\n //! Streaming SIMD Extensions\n bool SSE = CpuFeatures::SSE();\n\n //! Streaming SIMD Extensions 2\n bool SSE2 = CpuFeatures::SSE2();\n\n //! Streaming SIMD Extensions 3\n bool SSE3 = CpuFeatures::SSE3();\n\n //! Supplemental Streaming SIMD Extensions 3\n bool SSSE3 = CpuFeatures::SSSE3();\n\n //! Streaming SIMD Extensions 4.1\n bool SSE4_1 = CpuFeatures::SSE4_1();\n\n //! Streaming SIMD Extensions 4.2\n bool SSE4_2 = CpuFeatures::SSE4_2();\n\n //! Advanced Vector Extensions\n bool AVX = CpuFeatures::AVX();\n\n //! Advanced Vector Extensions 2\n bool AVX2 = CpuFeatures::AVX2();\n\n //! AVX-512 Foundation\n bool AVX512F = CpuFeatures::AVX512F();\n\n //! AVX-512 DQ (Double/Quad granular) Instructions\n bool AVX512DQ = CpuFeatures::AVX512DQ();\n\n //! AVX-512 Prefetch\n bool AVX512PF = CpuFeatures::AVX512PF();\n\n //! AVX-512 Exponential and Reciprocal\n bool AVX512ER = CpuFeatures::AVX512ER();\n\n //! AVX-512 Conflict Detection\n bool AVX512CD = CpuFeatures::AVX512CD();\n\n //! AVX-512 BW (Byte/Word granular) Instructions\n bool AVX512BW = CpuFeatures::AVX512BW();\n\n //! AVX-512 VL (128/256 Vector Length) Extensions\n bool AVX512VL = CpuFeatures::AVX512VL();\n\n //! AVX-512 Integer Fused Multiply-Add instructions\n bool AVX512_IFMA = CpuFeatures::AVX512_IFMA();\n\n //! AVX512 Vector Bit Manipulation instructions\n bool AVX512_VBMI = CpuFeatures::AVX512_VBMI();\n\n //! Additional AVX512 Vector Bit Manipulation Instructions\n bool AVX512_VBMI2 = CpuFeatures::AVX512_VBMI2();\n\n //! Vector Neural Network Instructions\n bool AVX512_VNNI = CpuFeatures::AVX512_VNNI();\n\n //! Support for VPOPCNT[B,W] and VPSHUF-BITQMB instructions\n bool AVX512_BITALG = CpuFeatures::AVX512_BITALG();\n\n //! POPCNT for vectors of DW/QW\n bool AVX512_VPOPCNTDQ = CpuFeatures::AVX512_VPOPCNTDQ();\n\n //! AVX-512 Neural Network Instructions\n bool AVX512_4VNNIW = CpuFeatures::AVX512_4VNNIW();\n\n //! AVX-512 Multiply Accumulation Single precision\n bool AVX512_4FMAPS = CpuFeatures::AVX512_4FMAPS();\n\n //! AVX-512 FP16 instructions\n bool AVX512_FP16 = CpuFeatures::AVX512_FP16();\n\n //! CMPXCHG8 instruction\n bool CX8 = CpuFeatures::CX8();\n\n //! CMPXCHG16B instruction\n bool CX16 = CpuFeatures::CX16();\n\n //! PCLMULQDQ instruction\n bool PCLMULQDQ = CpuFeatures::PCLMULQDQ();\n\n //! Carry-Less Multiplication Double Quadword\n bool VPCLMULQDQ = CpuFeatures::VPCLMULQDQ();\n\n //! CMOV instructions (plus FCMOVcc, FCOMI with FPU)\n bool CMOV = CpuFeatures::CMOV();\n\n //! MOVBE instruction\n bool MOVBE = CpuFeatures::MOVBE();\n\n //! Enhanced REP MOVSB/STOSB instructions\n bool ERMS = CpuFeatures::ERMS();\n\n //! POPCNT instruction\n bool POPCNT = CpuFeatures::POPCNT();\n\n //! XSAVE/XRSTOR/XSETBV/XGETBV instructions\n bool XSAVE = CpuFeatures::XSAVE();\n\n //! Fused multiply-add\n bool FMA = CpuFeatures::FMA();\n\n //! ADCX and ADOX instructions\n bool ADX = CpuFeatures::ADX();\n\n //! Galois Field New Instructions\n bool GFNI = CpuFeatures::GFNI();\n\n //! AES instructions\n bool AES = CpuFeatures::AES();\n\n //! Vector AES\n bool VAES = CpuFeatures::VAES();\n\n //! RDSEED instruction\n bool RDSEED = CpuFeatures::RDSEED();\n\n //! RDRAND instruction\n bool RDRAND = CpuFeatures::RDRAND();\n\n //! SHA1/SHA256 Instruction Extensions\n bool SHA = CpuFeatures::SHA();\n\n //! 1st group bit manipulation extensions\n bool BMI1 = CpuFeatures::BMI1();\n\n //! 2nd group bit manipulation extensions\n bool BMI2 = CpuFeatures::BMI2();\n\n //! CLFLUSH instruction\n bool CLFLUSH = CpuFeatures::CLFLUSH();\n\n //! CLFLUSHOPT instruction\n bool CLFLUSHOPT = CpuFeatures::CLFLUSHOPT();\n\n //! CLWB instruction\n bool CLWB = CpuFeatures::CLWB();\n\n //! RDPID instruction\n bool RDPID = CpuFeatures::RDPID();\n\n //! Onboard FPU\n bool FPU = CpuFeatures::FPU();\n\n //! Hyper-Threading\n bool HT = CpuFeatures::HT();\n\n //! Hardware virtualization\n bool VMX = CpuFeatures::VMX();\n\n // ！Running on a hypervisor\n bool HYPERVISOR = CpuFeatures::HYPERVISOR();\n };\n static StaticFlags static_flags_;\n};\n\n} // namespace internal\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/io/file.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#if !defined(_WIN64) && !defined(_WIN32)\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#else\n#include \n#endif\n\nnamespace zvec {\nnamespace ailego {\n\n#if !defined(_WIN64) && !defined(_WIN32)\n\nstatic inline int OpenSafely(const char *path, int flags) {\n int fd = open(path, flags, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH);\n while (fd == -1 && errno == EINTR) {\n fd = open(path, flags, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH);\n }\n return fd;\n}\n\nstatic inline void CloseSafely(int fd) {\n int ret = close(fd);\n while (ret == -1 && errno == EINTR) {\n ret = close(fd);\n }\n}\n\nstatic inline ssize_t ReadSafely(int fd, void *buf, size_t count) {\n ssize_t ret = read(fd, buf, count);\n while (ret == -1 && errno == EINTR) {\n ret = read(fd, buf, count);\n }\n return ret;\n}\n\nstatic inline ssize_t PreadSafely(int fd, void *buf, size_t count,\n ssize_t offset) {\n ssize_t ret = pread(fd, buf, count, offset);\n while (ret == -1 && errno == EINTR) {\n ret = pread(fd, buf, count, offset);\n }\n return ret;\n}\n\nstatic inline ssize_t WriteSafely(int fd, const void *buf, size_t count) {\n ssize_t ret = write(fd, buf, count);\n while (ret == -1 && errno == EINTR) {\n ret = write(fd, buf, count);\n }\n return ret;\n}\n\nstatic inline ssize_t PwriteSafely(int fd, const void *buf, size_t count,\n ssize_t offset) {\n ssize_t ret = pwrite(fd, buf, count, offset);\n while (ret == -1 && errno == EINTR) {\n ret = pwrite(fd, buf, count, offset);\n }\n return ret;\n}\n\nstatic inline size_t ReadAll(int fd, void *buf, size_t count) {\n size_t rdlen = 0;\n while (rdlen < count) {\n ssize_t ret = ReadSafely(fd, (char *)buf + rdlen, count - rdlen);\n if (ret <= 0) {\n break;\n }\n rdlen += ret;\n }\n return rdlen;\n}\n\nstatic inline size_t PreadAll(int fd, void *buf, size_t count, ssize_t offset) {\n size_t rdlen = 0;\n while (rdlen < count) {\n ssize_t ret =\n PreadSafely(fd, (char *)buf + rdlen, count - rdlen, offset + rdlen);\n if (ret <= 0) {\n break;\n }\n rdlen += ret;\n }\n return rdlen;\n}\n\nstatic inline size_t WriteAll(int fd, const void *buf, size_t count) {\n size_t wrlen = 0;\n while (wrlen < count) {\n ssize_t ret = WriteSafely(fd, (const char *)buf + wrlen, count - wrlen);\n if (ret <= 0) {\n break;\n }\n wrlen += ret;\n }\n return wrlen;\n}\n\nstatic inline size_t PwriteAll(int fd, const void *buf, size_t count,\n ssize_t offset) {\n size_t wrlen = 0;\n while (wrlen < count) {\n ssize_t ret = PwriteSafely(fd, (const char *)buf + wrlen, count - wrlen,\n offset + wrlen);\n if (ret <= 0) {\n break;\n }\n wrlen += ret;\n }\n return wrlen;\n}\n\nbool File::create(const char *path, size_t len, bool direct) {\n ailego_false_if_false(native_handle_ == File::InvalidHandle && path);\n\n // Try opening or creating a file\n int flags = O_RDWR | O_CREAT;\n#ifdef O_DIRECT\n if (direct) {\n flags |= O_DIRECT;\n }\n#else\n (void)direct;\n#endif\n\n int fd = OpenSafely(path, flags);\n ailego_false_if_lt_zero(fd);\n\n#ifdef F_NOCACHE\n // Direct IO canonical solution for Mac OSX\n if (direct) {\n ailego_false_if_ne_zero(fcntl(fd, F_NOCACHE, 1));\n }\n#endif\n\n // Truncate the file to the specified size\n ailego_do_if_ne_zero(ftruncate(fd, len)) {\n CloseSafely(fd);\n return false;\n }\n\n read_only_ = false;\n native_handle_ = fd;\n return true;\n}\n\nbool File::open(const char *path, bool rdonly, bool direct) {\n ailego_false_if_false(native_handle_ == File::InvalidHandle && path);\n\n // Try opening the file\n int flags = rdonly ? O_RDONLY : O_RDWR;\n#ifdef O_DIRECT\n if (direct) {\n flags |= O_DIRECT;\n }\n#else\n (void)direct;\n#endif\n\n int fd = OpenSafely(path, flags);\n ailego_false_if_lt_zero(fd);\n\n#ifdef F_NOCACHE\n // Direct IO canonical solution for Mac OSX\n if (direct) {\n ailego_false_if_ne_zero(fcntl(fd, F_NOCACHE, 1));\n }\n#endif\n\n read_only_ = rdonly;\n native_handle_ = fd;\n return true;\n}\n\nvoid File::close(void) {\n ailego_return_if_false(native_handle_ != File::InvalidHandle);\n CloseSafely(native_handle_);\n native_handle_ = File::InvalidHandle;\n}\n\nvoid File::reset(void) {\n ailego_return_if_false(native_handle_ != File::InvalidHandle);\n lseek(native_handle_, 0, SEEK_SET);\n}\n\nsize_t File::write(const void *data, size_t len) {\n const size_t block_size = 0x40000000u;\n size_t total = 0u;\n\n for (; len >= block_size; len -= block_size) {\n size_t wrlen =\n WriteAll(native_handle_, (const uint8_t *)data + total, block_size);\n if (wrlen != block_size) {\n return (total + wrlen);\n }\n total += block_size;\n }\n if (len > 0) {\n total += WriteAll(native_handle_, (const uint8_t *)data + total, len);\n }\n return total;\n}\n\nsize_t File::write(ssize_t off, const void *data, size_t len) {\n const size_t block_size = 0x40000000u;\n size_t total = 0u;\n\n for (; len >= block_size; len -= block_size) {\n size_t wrlen = PwriteAll(native_handle_, (const uint8_t *)data + total,\n block_size, off + total);\n if (wrlen != block_size) {\n return (total + wrlen);\n }\n total += block_size;\n }\n if (len > 0) {\n total += PwriteAll(native_handle_, (const uint8_t *)data + total, len,\n off + total);\n }\n return total;\n}\n\nsize_t File::read(void *buf, size_t len) {\n const size_t block_size = 0x40000000u;\n size_t total = 0u;\n\n for (; len >= block_size; len -= block_size) {\n size_t rdlen = ReadAll(native_handle_, (uint8_t *)buf + total, block_size);\n if (rdlen != block_size) {\n return (total + rdlen);\n }\n total += block_size;\n }\n if (len > 0) {\n total += ReadAll(native_handle_, (uint8_t *)buf + total, len);\n }\n return total;\n}\n\nsize_t File::read(ssize_t off, void *buf, size_t len) {\n const size_t block_size = 0x40000000u;\n size_t total = 0u;\n\n for (; len >= block_size; len -= block_size) {\n size_t rdlen = PreadAll(native_handle_, (uint8_t *)buf + total, block_size,\n off + total);\n if (rdlen != block_size) {\n return (total + rdlen);\n }\n total += block_size;\n }\n if (len > 0) {\n total += PreadAll(native_handle_, (uint8_t *)buf + total, len, off + total);\n }\n return total;\n}\n\nbool File::flush(void) {\n ailego_false_if_false(native_handle_ != File::InvalidHandle);\n return (fsync(native_handle_) == 0);\n}\n\nbool File::seek(ssize_t off, Origin origin) {\n ailego_false_if_false(native_handle_ != File::InvalidHandle);\n ailego_false_if_false(lseek(native_handle_, off, (int)origin) != (off_t)-1);\n return true;\n}\n\nbool File::truncate(size_t len) {\n ailego_false_if_false(native_handle_ != File::InvalidHandle);\n ailego_false_if_ne_zero(ftruncate(native_handle_, (off_t)len));\n return true;\n}\n\nsize_t File::size(void) const {\n struct stat fs;\n ailego_zero_if_false(native_handle_ != File::InvalidHandle &&\n fstat(native_handle_, &fs) == 0);\n return (fs.st_size);\n}\n\nssize_t File::offset(void) const {\n off_t off;\n ailego_zero_if_false(native_handle_ != File::InvalidHandle &&\n (off = lseek(native_handle_, 0, SEEK_CUR)) != -1);\n return off;\n}\n\nvoid *File::MemoryMap(NativeHandle handle, ssize_t off, size_t len, int opts) {\n int prot =\n ((opts & File::MMAP_READONLY) ? PROT_READ : PROT_READ | PROT_WRITE);\n\n#if defined(MAP_POPULATE)\n if (opts & File::MMAP_POPULATE) {\n prot |= MAP_POPULATE;\n }\n#endif\n int flags = (opts & File::MMAP_SHARED) ? MAP_SHARED : MAP_PRIVATE;\n#if defined(MAP_HUGETLB)\n if (opts & File::MMAP_HUGE_PAGE) {\n flags |= MAP_HUGETLB;\n }\n#endif\n void *addr = mmap(nullptr, len, prot, flags, handle, off);\n ailego_null_if_false(addr != MAP_FAILED);\n\n if (opts & File::MMAP_LOCKED) {\n mlock(addr, len);\n }\n if (opts & File::MMAP_WARMUP) {\n File::MemoryWarmup(addr, len);\n }\n return addr;\n}\n\n#if !defined(MAP_ANONYMOUS) && defined(MAP_ANON)\n#define MAP_ANONYMOUS MAP_ANON\n#endif\n\nvoid *File::MemoryMap(size_t len, int opts) {\n#if defined(MAP_ANONYMOUS)\n int prot =\n ((opts & File::MMAP_READONLY) ? PROT_READ : PROT_READ | PROT_WRITE);\n\n#if defined(MAP_POPULATE)\n if (opts & File::MMAP_POPULATE) {\n prot |= MAP_POPULATE;\n }\n#endif\n int flags = (opts & File::MMAP_SHARED) ? MAP_SHARED | MAP_ANONYMOUS\n : MAP_PRIVATE | MAP_ANONYMOUS;\n#if defined(MAP_HUGETLB)\n if (opts & File::MMAP_HUGE_PAGE) {\n flags |= MAP_HUGETLB;\n }\n#endif\n void *addr = mmap(nullptr, len, prot, flags, -1, 0);\n ailego_null_if_false(addr != MAP_FAILED);\n return addr;\n#else\n (void)len;\n (void)opts;\n return nullptr;\n#endif // MAP_ANONYMOUS\n}\n\nvoid *File::MemoryRemap(void *oldptr, size_t oldsize, void *newptr,\n size_t newsize) {\n#if defined(__linux) || defined(__linux__)\n return newptr ? mremap(oldptr, oldsize, newsize, MREMAP_FIXED, newptr)\n : mremap(oldptr, oldsize, newsize, MREMAP_MAYMOVE);\n#elif defined(__NetBSD__)\n return newptr ? mremap(oldptr, oldsize, newptr, newsize, MAP_FIXED)\n : mremap(oldptr, oldsize, nullptr, newsize, 0);\n#else\n (void)oldptr;\n (void)oldsize;\n (void)newptr;\n (void)newsize;\n errno = ENOTSUP;\n return nullptr;\n#endif\n}\n\nvoid File::MemoryUnmap(void *addr, size_t len) {\n ailego_return_if_false(addr);\n munmap(addr, len);\n}\n\nbool File::MemoryFlush(void *addr, size_t len) {\n ailego_false_if_false(addr);\n return (msync(addr, len, MS_ASYNC) == 0);\n}\n\nbool File::MemoryLock(void *addr, size_t len) {\n ailego_false_if_false(addr && len);\n return (mlock(addr, len) == 0);\n}\n\nbool File::MemoryUnlock(void *addr, size_t len) {\n ailego_false_if_false(addr && len);\n return (munlock(addr, len) == 0);\n}\n\n#else\n\n//! Create a local file\nbool File::create(const char *path, size_t len, bool direct) {\n ailego_false_if_false(native_handle_ == File::InvalidHandle && path);\n\n // Try opening or creating the file\n HANDLE file_handle =\n CreateFileA(path, GENERIC_WRITE | GENERIC_READ, FILE_SHARE_READ, nullptr,\n CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr);\n ailego_false_if_false(file_handle != INVALID_HANDLE_VALUE);\n\n // Truncate the file to the specified size\n LARGE_INTEGER file_size;\n file_size.QuadPart = len;\n ailego_do_if_false(\n SetFilePointerEx(file_handle, file_size, nullptr, FILE_BEGIN) &&\n SetEndOfFile(file_handle)) {\n CloseHandle(file_handle);\n return false;\n }\n\n if (!direct) {\n // Reset the file pointer\n SetFilePointer(file_handle, 0, nullptr, FILE_BEGIN);\n } else {\n // Close and reopen file\n CloseHandle(file_handle);\n file_handle = CreateFileA(\n path, GENERIC_WRITE | GENERIC_READ, FILE_SHARE_READ, nullptr,\n OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING, nullptr);\n ailego_false_if_false(file_handle != INVALID_HANDLE_VALUE);\n }\n\n read_only_ = false;\n native_handle_ = file_handle;\n return true;\n}\n\n//! Open a local file\nbool File::open(const char *path, bool rdonly, bool direct) {\n ailego_false_if_false(native_handle_ == File::InvalidHandle && path);\n\n // Try opening the file\n DWORD flags = FILE_ATTRIBUTE_NORMAL;\n if (direct) {\n flags |= FILE_FLAG_NO_BUFFERING;\n }\n HANDLE file_handle =\n CreateFileA(path, (rdonly ? GENERIC_READ : GENERIC_READ | GENERIC_WRITE),\n FILE_SHARE_READ, nullptr, OPEN_EXISTING, flags, nullptr);\n ailego_false_if_false(file_handle != INVALID_HANDLE_VALUE);\n\n read_only_ = rdonly;\n native_handle_ = file_handle;\n return true;\n}\n\nvoid File::close(void) {\n ailego_return_if_false(native_handle_ != File::InvalidHandle);\n CloseHandle(native_handle_);\n native_handle_ = File::InvalidHandle;\n}\n\nvoid File::reset(void) {\n ailego_return_if_false(native_handle_ != File::InvalidHandle);\n SetFilePointer(native_handle_, 0, nullptr, FILE_BEGIN);\n}\n\nsize_t File::write(const void *data, size_t len) {\n const DWORD block_size = 0x40000000u;\n DWORD wrlen = 0u;\n size_t total = 0u;\n\n for (; len >= block_size; len -= block_size) {\n if (!WriteFile(native_handle_, (const uint8_t *)data + total, block_size,\n &wrlen, nullptr)) {\n return total;\n }\n if (wrlen != block_size) {\n return (total + wrlen);\n }\n total += block_size;\n }\n if (len > 0 && WriteFile(native_handle_, (const uint8_t *)data + total,\n (DWORD)len, &wrlen, nullptr)) {\n total += wrlen;\n }\n return total;\n}\n\nsize_t File::write(ssize_t off, const void *data, size_t len) {\n const DWORD block_size = 0x40000000u;\n DWORD wrlen = 0u;\n size_t total = 0u;\n\n OVERLAPPED overlapped;\n memset(&overlapped, 0, sizeof(OVERLAPPED));\n\n for (; len >= block_size; len -= block_size) {\n uint64_t current = off + total;\n overlapped.OffsetHigh = (DWORD)(current >> 32);\n overlapped.Offset = (DWORD)(current & 0xffffffffu);\n\n if (!WriteFile(native_handle_, (const uint8_t *)data + total, block_size,\n &wrlen, &overlapped)) {\n return total;\n }\n if (wrlen != block_size) {\n return (total + wrlen);\n }\n total += block_size;\n }\n if (len > 0) {\n uint64_t current = off + total;\n overlapped.OffsetHigh = (DWORD)(current >> 32);\n overlapped.Offset = (DWORD)(current & 0xffffffffu);\n\n if (WriteFile(native_handle_, (const uint8_t *)data + total, (DWORD)len,\n &wrlen, &overlapped)) {\n total += wrlen;\n }\n }\n return total;\n}\n\nsize_t File::read(void *buf, size_t len) {\n const DWORD block_size = 0x40000000u;\n DWORD rdlen = 0u;\n size_t total = 0u;\n\n for (; len >= block_size; len -= block_size) {\n if (!ReadFile(native_handle_, (uint8_t *)buf + total, block_size, &rdlen,\n nullptr)) {\n return total;\n }\n if (rdlen != block_size) {\n return (total + rdlen);\n }\n total += block_size;\n }\n if (len > 0 && ReadFile(native_handle_, (uint8_t *)buf + total, (DWORD)len,\n &rdlen, nullptr)) {\n total += rdlen;\n }\n return total;\n}\n\nsize_t File::read(ssize_t off, void *buf, size_t len) {\n const DWORD block_size = 0x40000000u;\n DWORD rdlen = 0u;\n size_t total = 0u;\n\n OVERLAPPED overlapped;\n memset(&overlapped, 0, sizeof(OVERLAPPED));\n\n for (; len >= block_size; len -= block_size) {\n uint64_t current = off + total;\n overlapped.OffsetHigh = (DWORD)(current >> 32);\n overlapped.Offset = (DWORD)(current & 0xffffffffu);\n\n if (!ReadFile(native_handle_, (uint8_t *)buf + total, block_size, &rdlen,\n &overlapped)) {\n return total;\n }\n if (rdlen != block_size) {\n return (total + rdlen);\n }\n total += block_size;\n }\n if (len > 0) {\n uint64_t current = off + total;\n overlapped.OffsetHigh = (DWORD)(current >> 32);\n overlapped.Offset = (DWORD)(current & 0xffffffffu);\n\n if (ReadFile(native_handle_, (uint8_t *)buf + total, (DWORD)len, &rdlen,\n &overlapped)) {\n total += rdlen;\n }\n }\n return total;\n}\n\nbool File::flush(void) {\n ailego_false_if_false(native_handle_ != File::InvalidHandle);\n return (!!FlushFileBuffers(native_handle_));\n}\n\nbool File::seek(ssize_t off, Origin origin) {\n ailego_false_if_false(native_handle_ != File::InvalidHandle);\n\n LARGE_INTEGER file_offset;\n file_offset.QuadPart = off;\n ailego_false_if_false(SetFilePointerEx(native_handle_, file_offset, nullptr,\n (DWORD)origin) != 0);\n return true;\n}\n\nbool File::truncate(size_t len) {\n ailego_false_if_false(native_handle_ != File::InvalidHandle);\n\n LARGE_INTEGER file_size, orig_file_size;\n file_size.QuadPart = 0;\n orig_file_size.QuadPart = 0;\n ailego_false_if_false(SetFilePointerEx(native_handle_, file_size,\n &orig_file_size, FILE_CURRENT));\n\n // Truncate the file to the specified size\n file_size.QuadPart = len;\n ailego_false_if_false(\n SetFilePointerEx(native_handle_, file_size, nullptr, FILE_BEGIN) &&\n SetEndOfFile(native_handle_));\n\n // Reset the file pointer\n SetFilePointerEx(native_handle_, orig_file_size, nullptr, FILE_BEGIN);\n return true;\n}\n\nsize_t File::size(void) const {\n LARGE_INTEGER file_size;\n ailego_zero_if_false(native_handle_ != File::InvalidHandle &&\n GetFileSizeEx(native_handle_, &file_size));\n return (size_t)file_size.QuadPart;\n}\n\nssize_t File::offset(void) const {\n LARGE_INTEGER file_size;\n LARGE_INTEGER file_size_new;\n file_size.QuadPart = 0;\n ailego_zero_if_false(native_handle_ != File::InvalidHandle &&\n SetFilePointerEx(native_handle_, file_size,\n &file_size_new, FILE_CURRENT));\n return (size_t)file_size_new.QuadPart;\n}\n\nvoid *File::MemoryMap(NativeHandle handle, ssize_t off, size_t len, int opts) {\n LARGE_INTEGER file_size;\n file_size.QuadPart = len;\n\n // Create map object\n HANDLE file_mapping = CreateFileMapping(\n handle, nullptr,\n ((opts & File::MMAP_READONLY) ? PAGE_READONLY : PAGE_READWRITE),\n file_size.HighPart, file_size.LowPart, nullptr);\n ailego_null_if_false(file_mapping != nullptr);\n\n DWORD desired_access = FILE_MAP_READ;\n if (!(opts & File::MMAP_READONLY)) {\n desired_access |= FILE_MAP_WRITE;\n }\n if (!(opts & File::MMAP_SHARED)) {\n desired_access |= FILE_MAP_COPY;\n }\n file_size.QuadPart = off;\n\n // Map the whole file to memory and close handle\n void *addr = MapViewOfFile(file_mapping, desired_access, file_size.HighPart,\n file_size.LowPart, 0);\n CloseHandle(file_mapping);\n\n ailego_null_if_false(addr);\n if (opts & File::MMAP_LOCKED) {\n VirtualLock(addr, len);\n }\n if (opts & File::MMAP_WARMUP) {\n File::MemoryWarmup(addr, len);\n }\n return addr;\n}\n\nvoid *File::MemoryMap(size_t, int) {\n return nullptr;\n}\n\nvoid *File::MemoryRemap(void *, size_t, void *, size_t) {\n return nullptr;\n}\n\nvoid File::MemoryUnmap(void *addr, size_t /*len*/) {\n ailego_return_if_false(addr);\n UnmapViewOfFile(addr);\n}\n\nbool File::MemoryFlush(void *addr, size_t /*len*/) {\n ailego_false_if_false(addr);\n return (!!FlushViewOfFile(addr, 0));\n}\n\nbool File::MemoryLock(void *addr, size_t len) {\n ailego_false_if_false(addr && len);\n return (!!VirtualLock(addr, len));\n}\n\nbool File::MemoryUnlock(void *addr, size_t len) {\n ailego_false_if_false(addr && len);\n return (!!VirtualUnlock(addr, len));\n}\n\nstatic inline int getpagesize(void) {\n SYSTEM_INFO info;\n GetSystemInfo(&info);\n return info.dwPageSize;\n}\n#endif\n\nvoid File::MemoryWarmup(void *addr, size_t len) {\n static int page_size = getpagesize();\n\n if (addr && len) {\n uint8_t *p = reinterpret_cast(addr);\n uint8_t *end = p + len;\n volatile uint8_t tmp = 0;\n\n while (p < end) {\n tmp ^= *p;\n p += page_size;\n }\n }\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/io/file_lock.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"file_lock.h\"\n\n#if !defined(_WIN64) && !defined(_WIN32)\n#include \n#else\n#include \n#endif\n\nnamespace zvec {\nnamespace ailego {\n\n#if !defined(_WIN64) && !defined(_WIN32)\nbool FileLock::Lock(int fd) {\n return (flock(fd, LOCK_EX) == 0);\n}\n\nbool FileLock::TryLock(int fd) {\n return (flock(fd, LOCK_EX | LOCK_NB) == 0);\n}\n\nbool FileLock::LockShared(int fd) {\n return (flock(fd, LOCK_SH) == 0);\n}\n\nbool FileLock::TryLockShared(int fd) {\n return (flock(fd, LOCK_SH | LOCK_NB) == 0);\n}\n\nbool FileLock::Unlock(int fd) {\n return (flock(fd, LOCK_UN) == 0);\n}\n\n#else\nbool FileLock::Lock(HANDLE handle) {\n OVERLAPPED ol = {0};\n return (!!LockFileEx(handle, LOCKFILE_EXCLUSIVE_LOCK, 0, MAXDWORD, MAXDWORD,\n &ol));\n}\n\nbool FileLock::TryLock(HANDLE handle) {\n OVERLAPPED ol = {0};\n return (!!LockFileEx(handle,\n LOCKFILE_EXCLUSIVE_LOCK | LOCKFILE_FAIL_IMMEDIATELY, 0,\n MAXDWORD, MAXDWORD, &ol));\n}\n\nbool FileLock::LockShared(HANDLE handle) {\n OVERLAPPED ol = {0};\n return (!!LockFileEx(handle, 0, 0, MAXDWORD, MAXDWORD, &ol));\n}\n\nbool FileLock::TryLockShared(HANDLE handle) {\n OVERLAPPED ol = {0};\n return (!!LockFileEx(handle, LOCKFILE_FAIL_IMMEDIATELY, 0, MAXDWORD, MAXDWORD,\n &ol));\n}\n\nbool FileLock::Unlock(HANDLE handle) {\n OVERLAPPED ol = {0};\n return (!!UnlockFileEx(handle, 0, MAXDWORD, MAXDWORD, &ol));\n}\n\n#endif // !_WIN64 && !_WIN32\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/io/file_lock.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! File Utility\n */\nclass FileLock {\n public:\n //! Constructor\n FileLock(const File &file) : native_handle_(file.native_handle()) {}\n\n //! Constructor\n FileLock(File::NativeHandle handle) : native_handle_(handle) {}\n\n //! Locking\n bool lock(void) const {\n return FileLock::Lock(native_handle_);\n }\n\n //! Try locking\n bool try_lock(void) const {\n return FileLock::TryLock(native_handle_);\n }\n\n //! Locking (shared)\n bool lock_shared(void) const {\n return FileLock::LockShared(native_handle_);\n }\n\n //! Try locking (shared)\n bool try_lock_shared(void) const {\n return FileLock::TryLockShared(native_handle_);\n }\n\n //! Unlocking\n bool unlock(void) const {\n return FileLock::Unlock(native_handle_);\n }\n\n //! Locking\n static bool Lock(File::NativeHandle handle);\n\n //! Try locking\n static bool TryLock(File::NativeHandle handle);\n\n //! Locking (shared)\n static bool LockShared(File::NativeHandle handle);\n\n //! Try locking (shared)\n static bool TryLockShared(File::NativeHandle handle);\n\n //! Unlocking\n static bool Unlock(File::NativeHandle handle);\n\n private:\n //! Disable them\n FileLock(const FileLock &) = delete;\n FileLock(FileLock &&) = delete;\n FileLock &operator=(const FileLock &) = delete;\n FileLock &operator=(FileLock &&) = delete;\n\n //! Members\n File::NativeHandle native_handle_;\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/io/file_writer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \"file.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n/*! File Stream Writer\n */\nclass FileWriter {\n public:\n //! Constructor\n FileWriter(void) {}\n\n //! Constructor\n FileWriter(FileWriter &&rhs) : file_(std::move(rhs.file_)) {}\n\n //! Destructor\n ~FileWriter(void) {}\n\n //! Assignment\n FileWriter &operator=(FileWriter &&rhs) {\n file_ = std::move(rhs.file_);\n return *this;\n }\n\n //! Output to writer\n FileWriter &operator<<(const char *str) {\n size_t len = std::strlen(str);\n if (file_.write(str, len) != len) {\n throw std::ios_base::failure(\"Write error\");\n }\n return *this;\n }\n\n //! Output to writer\n FileWriter &operator<<(const std::string &str) {\n if (file_.write(str.data(), str.size()) != str.size()) {\n throw std::ios_base::failure(\"Write error\");\n }\n return *this;\n }\n\n //! Output to writer\n FileWriter &operator<<(char c) {\n if (file_.write(&c, 1) != 1) {\n throw std::ios_base::failure(\"Write error\");\n }\n return *this;\n }\n\n //! Test if the file is valid\n bool is_valid(void) const {\n return file_.is_valid();\n }\n\n //! Create a local file\n bool create(const char *path) {\n return file_.create(path, 0, false);\n }\n\n //! Open a local file\n bool open(const char *path) {\n return file_.open(path, false, false);\n }\n\n //! Close the local file\n void close(void) {\n file_.close();\n }\n\n //! Write data into the file\n size_t write(const void *data, size_t len) {\n return file_.write(data, len);\n }\n\n //! Synchronize memory with physical storage\n bool flush(void) {\n return file_.flush();\n }\n\n //! Output with format\n void print(const char *format, va_list args) {\n char buf[8192];\n std::vsnprintf(buf, sizeof(buf), format, args);\n (*this) << buf;\n }\n\n //! Output with format\n#if defined(__GNUC__)\n void print(const char *format, ...) __attribute__((format(printf, 2, 3))) {\n#else\n void print(const char *format, ...) {\n#endif\n va_list args;\n va_start(args, format);\n this->print(format, args);\n va_end(args);\n }\n\n private:\n //! Disable them\n FileWriter(const FileWriter &) = delete;\n FileWriter &operator=(const FileWriter &) = delete;\n\n //! Members\n File file_;\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/logger/logger.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\nconst int Logger::LEVEL_DEBUG = 0;\nconst int Logger::LEVEL_INFO = 1;\nconst int Logger::LEVEL_WARN = 2;\nconst int Logger::LEVEL_ERROR = 3;\nconst int Logger::LEVEL_FATAL = 4;\n\n/*! Console Logger\n */\nstruct ConsoleLogger : public Logger {\n //! Initialize Logger\n int init(const Params &) override {\n return 0;\n }\n\n //! Cleanup Logger\n int cleanup(void) override {\n return 0;\n }\n\n //! Log Message\n void log(int level, const char *file, int line, const char *format,\n va_list args) override {\n char buffer[8192];\n std::ostringstream stream;\n\n ailego::Realtime::Localtime(buffer, sizeof(buffer));\n stream << '[' << LevelString(level) << ' ' << buffer << ' '\n << std::this_thread::get_id() << ' ' << ailego::File::BaseName(file)\n << ':' << line << \"] \";\n\n vsnprintf(buffer, sizeof(buffer), format, args);\n stream << buffer << '\\n';\n\n if (level <= LEVEL_INFO) {\n std::cout << stream.str() << std::flush;\n } else {\n std::cerr << stream.str() << std::flush;\n }\n }\n};\n\n//! Logger Level\nint LoggerBroker::logger_level_ = Logger::LEVEL_WARN;\n\n//! Logger\nLogger::Pointer LoggerBroker::logger_(new ConsoleLogger);\n\n//! Register Console Logger in Factory\nFACTORY_REGISTER_LOGGER(ConsoleLogger);\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/cosine_distance_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n/*! Cosine Distance Matrix\n */\ntemplate \nstruct CosineDistanceMatrix;\n\n/*! Cosine Distance Matrix (M=1, N=1)\n */\ntemplate \nstruct CosineDistanceMatrix<\n T, 1, 1, typename std::enable_if::value>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n constexpr size_t extra_dim = sizeof(float) / sizeof(ValueType);\n size_t d = dim - extra_dim;\n\n float ip;\n InnerProductMatrix::Compute(m, q, d, &ip);\n\n *out = 1 - ip;\n }\n};\n\n/*! Cosine Distance Matrix\n */\ntemplate \nstruct CosineDistanceMatrix<\n T, M, N,\n typename std::enable_if::value && M >= 2 &&\n N >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType * /*m*/, const ValueType * /*q*/,\n size_t /*dim*/, float *out) {\n // ailego_assert(m && q && dim && out);\n\n *out = 0.0f;\n }\n};\n\n/*! Cosine Distance Matrix (N=1)\n */\ntemplate \nstruct CosineDistanceMatrix<\n T, M, 1,\n typename std::enable_if::value && M >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType * /*m*/, const ValueType * /*q*/,\n size_t /*dim*/, float *out) {\n // ailego_assert(m && q && dim && out);\n\n *out = 0.0f;\n }\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/distance.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \"distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n/*! Distance module\n */\nstruct Distance {\n //! Compute the hamming distance between two vectors (BINARY)\n static float Hamming(const uint32_t *lhs, const uint32_t *rhs, size_t dim) {\n float result;\n HammingDistanceMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n#if defined(AILEGO_M64)\n //! Compute the hamming distance between two vectors (BINARY)\n static float Hamming(const uint64_t *lhs, const uint64_t *rhs, size_t dim) {\n float result;\n HammingDistanceMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n#else\n //! Compute the hamming distance between two vectors (BINARY)\n static float Hamming(const uint64_t *lhs, const uint64_t *rhs, size_t dim) {\n float result;\n HammingDistanceMatrix::Compute(\n reinterpret_cast(lhs),\n reinterpret_cast(rhs), dim, &result);\n return result;\n }\n#endif\n\n //! Compute the squared euclidean distance between two vectors (FP32)\n static float SquaredEuclidean(const float *lhs, const float *rhs,\n size_t dim) {\n float result;\n SquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim,\n &result);\n return result;\n }\n\n //! Compute the squared euclidean distance between two vectors (FP16)\n static float SquaredEuclidean(const Float16 *lhs, const Float16 *rhs,\n size_t dim) {\n float result;\n SquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim,\n &result);\n return result;\n }\n\n //! Compute the squared euclidean distance between two vectors (INT8)\n static float SquaredEuclidean(const int8_t *lhs, const int8_t *rhs,\n size_t dim) {\n float result;\n SquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim,\n &result);\n return result;\n }\n\n //! Compute the squared euclidean distance between two vectors (INT4)\n static float SquaredEuclidean(const uint8_t *lhs, const uint8_t *rhs,\n size_t dim) {\n float result;\n SquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim,\n &result);\n return result;\n }\n\n //! Compute the euclidean distance between two vectors (FP32)\n static float Euclidean(const float *lhs, const float *rhs, size_t dim) {\n float result;\n EuclideanDistanceMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the euclidean distance between two vectors (FP16)\n static float Euclidean(const Float16 *lhs, const Float16 *rhs, size_t dim) {\n float result;\n EuclideanDistanceMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the euclidean distance between two vectors (INT8)\n static float Euclidean(const int8_t *lhs, const int8_t *rhs, size_t dim) {\n float result;\n EuclideanDistanceMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the euclidean distance between two vectors (INT4)\n static float Euclidean(const uint8_t *lhs, const uint8_t *rhs, size_t dim) {\n float result;\n EuclideanDistanceMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the inner product between two vectors (FP32)\n static float InnerProduct(const float *lhs, const float *rhs, size_t dim) {\n float result;\n InnerProductMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the inner product between two vectors (FP16)\n static float InnerProduct(const Float16 *lhs, const Float16 *rhs,\n size_t dim) {\n float result;\n InnerProductMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the inner product between two vectors (INT8)\n static float InnerProduct(const int8_t *lhs, const int8_t *rhs, size_t dim) {\n float result;\n InnerProductMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the minus inner product between two vectors (INT4)\n static float InnerProduct(const uint8_t *lhs, const uint8_t *rhs,\n size_t dim) {\n float result;\n InnerProductMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the minus inner product between two vectors (FP32)\n static float MinusInnerProduct(const float *lhs, const float *rhs,\n size_t dim) {\n float result;\n MinusInnerProductMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the minus inner product between two vectors (FP16)\n static float MinusInnerProduct(const Float16 *lhs, const Float16 *rhs,\n size_t dim) {\n float result;\n MinusInnerProductMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the minus inner product between two vectors (INT8)\n static float MinusInnerProduct(const int8_t *lhs, const int8_t *rhs,\n size_t dim) {\n float result;\n MinusInnerProductMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the minus inner product between two vectors (INT4)\n static float MinusInnerProduct(const uint8_t *lhs, const uint8_t *rhs,\n size_t dim) {\n float result;\n MinusInnerProductMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the mips squared L2 distance between two vectors\n //! (FP32, RepeatedQuadraticInjection)\n static float MipsSquaredEuclidean(const float *lhs, const float *rhs,\n size_t dim, size_t m, float eta) {\n float result;\n MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, m,\n eta, &result);\n return result;\n }\n\n //! Compute the mips squared L2 distance between two vectors\n //! (FP16, RepeatedQuadraticInjection)\n static float MipsSquaredEuclidean(const Float16 *lhs, const Float16 *rhs,\n size_t dim, size_t m, float eta) {\n float result;\n MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, m,\n eta, &result);\n return result;\n }\n\n //! Compute the mips squared L2 distance between two vectors\n //! (INT8, RepeatedQuadraticInjection)\n static float MipsSquaredEuclidean(const int8_t *lhs, const int8_t *rhs,\n size_t dim, size_t m, float eta) {\n float result;\n MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, m,\n eta, &result);\n return result;\n }\n\n //! Compute the mips squared L2 distance between two vectors\n //! (INT4, RepeatedQuadraticInjection)\n static float MipsSquaredEuclidean(const uint8_t *lhs, const uint8_t *rhs,\n size_t dim, size_t m, float eta) {\n float result;\n MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, m,\n eta, &result);\n return result;\n }\n\n //! Compute the mips squared L2 distance between two vectors\n //! (FP32, SphericalInjection)\n static float MipsSquaredEuclidean(const float *lhs, const float *rhs,\n size_t dim, float eta) {\n float result;\n MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim, eta,\n &result);\n return result;\n }\n\n //! Compute the mips squared L2 distance between two vectors\n //! (FP16, SphericalInjection)\n static float MipsSquaredEuclidean(const Float16 *lhs, const Float16 *rhs,\n size_t dim, float eta) {\n float result;\n MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim,\n eta, &result);\n return result;\n }\n\n //! Compute the mips squared L2 distance between two vectors\n //! (INT8, SphericalInjection)\n static float MipsSquaredEuclidean(const int8_t *lhs, const int8_t *rhs,\n size_t dim, float eta) {\n float result;\n MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim,\n eta, &result);\n return result;\n }\n\n //! Compute the mips squared L2 distance between two vectors\n //! (INT4, SphericalInjection)\n static float MipsSquaredEuclidean(const uint8_t *lhs, const uint8_t *rhs,\n size_t dim, float eta) {\n float result;\n MipsSquaredEuclideanDistanceMatrix::Compute(lhs, rhs, dim,\n eta, &result);\n return result;\n }\n\n //! Compute the cosine distance between two vectors (FP32)\n static float Cosine(const float *lhs, const float *rhs, size_t dim) {\n float result;\n CosineDistanceMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the cosine distance between two vectors (FP16)\n static float Cosine(const Float16 *lhs, const Float16 *rhs, size_t dim) {\n float result;\n CosineDistanceMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n\n //! Compute the cosine distance between two vectors (FP16)\n static float Cosine(const int8_t *lhs, const int8_t *rhs, size_t dim) {\n float result;\n CosineDistanceMatrix::Compute(lhs, rhs, dim, &result);\n return result;\n }\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/distance_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \"cosine_distance_matrix.h\"\n#include \"euclidean_distance_matrix.h\"\n#include \"hamming_distance_matrix.h\"\n#include \"inner_product_matrix.h\"\n#include \"mips_euclidean_distance_matrix.h\"\n" }, { "path": "src/ailego/math/distance_matrix_accum_fp16.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_fp16.i\"\n#include \"matrix_utility.i\"\n\n#if !defined(__FMA__)\n#define _mm_fmadd_ps(a, b, c) _mm_add_ps(_mm_mul_ps((a), (b)), (c))\n#define _mm256_fmadd_ps(a, b, c) _mm256_add_ps(_mm256_mul_ps((a), (b)), (c))\n#endif // !__FMA__\n\n#if defined(__AVX512F__) && !defined(__AVX512DQ__)\n#define _mm512_and_ps(a, b) \\\n _mm512_castsi512_ps( \\\n _mm512_and_epi32(_mm512_castps_si512(a), _mm512_castps_si512(b)))\n#define _mm512_mask_and_ps(src, k, a, b) \\\n _mm512_castsi512_ps(_mm512_mask_and_epi32(_mm512_castps_si512(src), (k), \\\n _mm512_castps_si512(a), \\\n _mm512_castps_si512(b)))\n#endif // __AVX512DQ__\n\n//! Compute the distance between matrix and query (FP16, M=1, N=1)\n#define ACCUM_FP16_1X1_AVX(m, q, dim, out, _MASK, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const Float16 *qe = q + dim; \\\n const Float16 *qe_aligned = q + ((dim >> 4) << 4); \\\n if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \\\n for (; q != qe_aligned; m += 16, q += 16) { \\\n MATRIX_FP16_ITER_1X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n if (qe >= qe_aligned + 8) { \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)m)); \\\n __m256 ymm_q = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)q)); \\\n ACCUM_FP32_STEP_AVX(ymm_m, ymm_q, ymm_sum_0_0) \\\n m += 8; \\\n q += 8; \\\n } \\\n } else { \\\n for (; q != qe_aligned; m += 16, q += 16) { \\\n MATRIX_FP16_ITER_1X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n if (qe >= qe_aligned + 8) { \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)m)); \\\n __m256 ymm_q = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q)); \\\n ACCUM_FP32_STEP_AVX(ymm_m, ymm_q, ymm_sum_0_0) \\\n m += 8; \\\n q += 8; \\\n } \\\n } \\\n MATRIX_FP16_MASK_AVX(m, q, (qe - q), _MASK, ymm_sum, ACCUM_FP32_STEP_AVX) \\\n *out = _NORM(HorizontalAdd_FP32_V256(ymm_sum_0_0));\n\n//! Compute the distance between matrix and query (FP16, M=2, N=1)\n#define ACCUM_FP16_2X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const Float16 *qe_aligned = q + ((dim >> 2) << 2); \\\n const Float16 *qe = q + dim; \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_aligned; m += 8, q += 4) { \\\n MATRIX_FP16_ITER_2X1_AVX(m, q, ymm_sum, _mm_load_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (; q != qe_aligned; m += 8, q += 4) { \\\n MATRIX_FP16_ITER_2X1_AVX(m, q, ymm_sum, _mm_loadu_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n if (qe >= qe_aligned + 2) { \\\n __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64(*(const __m64 *)(m))); \\\n __m128 xmm_q = _mm_cvtph_ps( \\\n _mm_shufflelo_epi16(_mm_broadcast_si32(q), _MM_SHUFFLE(1, 1, 0, 0))); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n m += 4; \\\n q += 2; \\\n } \\\n xmm_sum_0_0 = \\\n _mm_add_ps(xmm_sum_0_0, _mm_movehl_ps(xmm_sum_0_0, xmm_sum_0_0)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_cvtph_ps( \\\n _mm_shufflelo_epi16(_mm_broadcast_si32(m), _MM_SHUFFLE(0, 0, 1, 0))); \\\n __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi16(*(const short *)(q))); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0));\n\n//! Compute the distance between matrix and query (FP16, M=2, N=2)\n#define ACCUM_FP16_2X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const Float16 *qe_aligned = q + ((dim >> 2) << 3); \\\n const Float16 *qe = q + (dim << 1); \\\n if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \\\n for (; q != qe_aligned; m += 8, q += 8) { \\\n MATRIX_FP16_ITER_2X2_AVX(m, q, ymm_sum, _mm_load_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (; q != qe_aligned; m += 8, q += 8) { \\\n MATRIX_FP16_ITER_2X2_AVX(m, q, ymm_sum, _mm_loadu_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \\\n _mm256_extractf128_ps(ymm_sum_0_1, 1)); \\\n if (qe >= qe_aligned + 4) { \\\n __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64(*(const __m64 *)(m))); \\\n __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi64(*(const __m64 *)(q))); \\\n __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(2, 2, 0, 0)); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_0) \\\n xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 1, 1)); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_1) \\\n m += 4; \\\n q += 4; \\\n } \\\n xmm_sum_0_0 = _mm_add_ps(_mm_movelh_ps(xmm_sum_0_0, xmm_sum_0_1), \\\n _mm_movehl_ps(xmm_sum_0_1, xmm_sum_0_0)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_cvtph_ps( \\\n _mm_shufflelo_epi16(_mm_broadcast_si32(m), _MM_SHUFFLE(1, 0, 1, 0))); \\\n __m128 xmm_q = _mm_cvtph_ps( \\\n _mm_shufflelo_epi16(_mm_broadcast_si32(q), _MM_SHUFFLE(1, 1, 0, 0))); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=4, N=1)\n#define ACCUM_FP16_4X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const Float16 *qe = q + dim; \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const Float16 *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \\\n m += 8, q += 2) { \\\n MATRIX_FP16_ITER_4X1_AVX(m, q, ymm_sum, _mm_load_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \\\n m += 8, q += 2) { \\\n MATRIX_FP16_ITER_4X1_AVX(m, q, ymm_sum, _mm_loadu_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64(*(const __m64 *)(m))); \\\n __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi16(*(const short *)(q))); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=4, N=2)\n#define ACCUM_FP16_4X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const Float16 *qe = q + (dim << 1); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const Float16 *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \\\n m += 8, q += 4) { \\\n MATRIX_FP16_ITER_4X2_AVX(m, q, ymm_sum, _mm_load_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \\\n m += 8, q += 4) { \\\n MATRIX_FP16_ITER_4X2_AVX(m, q, ymm_sum, _mm_loadu_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \\\n _mm256_extractf128_ps(ymm_sum_0_1, 1)); \\\n if (q != qe) { \\\n __m128 xmm_q_0 = _mm_cvtph_ps(_mm_set1_epi16(*(const short *)(q + 0))); \\\n __m128 xmm_q_1 = _mm_cvtph_ps(_mm_set1_epi16(*(const short *)(q + 1))); \\\n __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64(*(const __m64 *)(m))); \\\n MATRIX_VAR_PROC(1, 2, 0, xmm_m, xmm_q, xmm_sum, ACCUM_FP32_STEP_SSE) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=4, N=4)\n#define ACCUM_FP16_4X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const Float16 *qe = q + (dim << 2); \\\n if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \\\n for (const Float16 *qe_aligned = q + ((dim >> 1) << 3); q != qe_aligned; \\\n m += 8, q += 8) { \\\n MATRIX_FP16_ITER_4X4_AVX(m, q, ymm_sum, _mm_load_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe_aligned = q + ((dim >> 1) << 3); q != qe_aligned; \\\n m += 8, q += 8) { \\\n MATRIX_FP16_ITER_4X4_AVX(m, q, ymm_sum, _mm_loadu_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \\\n _mm256_extractf128_ps(ymm_sum_0_1, 1)); \\\n __m128 xmm_sum_0_2 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_2), \\\n _mm256_extractf128_ps(ymm_sum_0_2, 1)); \\\n __m128 xmm_sum_0_3 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_3), \\\n _mm256_extractf128_ps(ymm_sum_0_3, 1)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_cvtph_ps(_mm_set1_epi64(*(const __m64 *)(m))); \\\n __m128 xmm_q = _mm_cvtph_ps(_mm_set1_epi64(*(const __m64 *)(q))); \\\n __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(0, 0, 0, 0)); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_0) \\\n xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(1, 1, 1, 1)); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_1) \\\n xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(2, 2, 2, 2)); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_2) \\\n xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 3, 3)); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_3) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=8, N=1)\n#define ACCUM_FP16_8X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const Float16 *qe = q + dim; q != qe; m += 8, ++q) { \\\n MATRIX_FP16_ITER_8X1_AVX(m, q, ymm_sum, _mm_load_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + dim; q != qe; m += 8, ++q) { \\\n MATRIX_FP16_ITER_8X1_AVX(m, q, ymm_sum, _mm_loadu_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=8, N=2)\n#define ACCUM_FP16_8X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \\\n MATRIX_FP16_ITER_8X2_AVX(m, q, ymm_sum, _mm_load_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \\\n MATRIX_FP16_ITER_8X2_AVX(m, q, ymm_sum, _mm_loadu_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=8, N=4)\n#define ACCUM_FP16_8X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \\\n MATRIX_FP16_ITER_8X4_AVX(m, q, ymm_sum, _mm_load_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \\\n MATRIX_FP16_ITER_8X4_AVX(m, q, ymm_sum, _mm_loadu_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=8, N=8)\n#define ACCUM_FP16_8X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 8, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \\\n MATRIX_FP16_ITER_8X8_AVX(m, q, ymm_sum, _mm_load_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \\\n MATRIX_FP16_ITER_8X8_AVX(m, q, ymm_sum, _mm_loadu_si128, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=16, N=1)\n#define ACCUM_FP16_16X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP16_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP16_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=16, N=2)\n#define ACCUM_FP16_16X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP16_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP16_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=16, N=4)\n#define ACCUM_FP16_16X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP16_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP16_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=16, N=8)\n#define ACCUM_FP16_16X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP16_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP16_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=16, N=16)\n#define ACCUM_FP16_16X16_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 16, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP16_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP16_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=1)\n#define ACCUM_FP16_32X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP16_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP16_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=2)\n#define ACCUM_FP16_32X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP16_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP16_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=4)\n#define ACCUM_FP16_32X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP16_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP16_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=8)\n#define ACCUM_FP16_32X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP16_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP16_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=16)\n#define ACCUM_FP16_32X16_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP16_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP16_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=32)\n#define ACCUM_FP16_32X32_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 32, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP16_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP16_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=1, N=1)\n#define ACCUM_FP16_1X1_AVX512(m, q, dim, out, _MASK, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m512, zmm_sum, _mm512_setzero_ps()) \\\n const Float16 *qe = q + dim; \\\n const Float16 *qe_aligned = q + ((dim >> 5) << 5); \\\n if (((uintptr_t)m & 0x3f) == 0 && ((uintptr_t)q & 0x3f) == 0) { \\\n for (; q != qe_aligned; m += 32, q += 32) { \\\n MATRIX_FP16_ITER_1X1_AVX512(m, q, zmm_sum, _mm512_load_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n if (qe >= qe_aligned + 16) { \\\n __m512 zmm_m = _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)m)); \\\n __m512 zmm_q = _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)q)); \\\n ACCUM_FP32_STEP_AVX512(zmm_m, zmm_q, zmm_sum_0_0) \\\n m += 16; \\\n q += 16; \\\n } \\\n } else { \\\n for (; q != qe_aligned; m += 32, q += 32) { \\\n MATRIX_FP16_ITER_1X1_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n if (qe >= qe_aligned + 16) { \\\n __m512 zmm_m = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)m)); \\\n __m512 zmm_q = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)q)); \\\n ACCUM_FP32_STEP_AVX512(zmm_m, zmm_q, zmm_sum_0_0) \\\n m += 16; \\\n q += 16; \\\n } \\\n } \\\n __m256 ymm_sum_0_0 = _mm256_add_ps(_mm512_castps512_ps256(zmm_sum_0_0), \\\n _mm256_castpd_ps(_mm512_extractf64x4_pd( \\\n _mm512_castps_pd(zmm_sum_0_0), 1))); \\\n if (qe >= q + 8) { \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)m)); \\\n __m256 ymm_q = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q)); \\\n ACCUM_FP32_STEP_AVX(ymm_m, ymm_q, ymm_sum_0_0) \\\n m += 8; \\\n q += 8; \\\n } \\\n MATRIX_FP16_MASK_AVX(m, q, (qe - q), _MASK, ymm_sum, ACCUM_FP32_STEP_AVX) \\\n *out = _NORM(HorizontalAdd_FP32_V256(ymm_sum_0_0));\n\n//! Compute the distance between matrix and query (FP16, M=16, N=1)\n#define ACCUM_FP16_16X1_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP16_ITER_16X1_AVX512(m, q, zmm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP16_ITER_16X1_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=16, N=2)\n#define ACCUM_FP16_16X2_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP16_ITER_16X2_AVX512(m, q, zmm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP16_ITER_16X2_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=16, N=4)\n#define ACCUM_FP16_16X4_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP16_ITER_16X4_AVX512(m, q, zmm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP16_ITER_16X4_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=16, N=8)\n#define ACCUM_FP16_16X8_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 8, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP16_ITER_16X8_AVX512(m, q, zmm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP16_ITER_16X8_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 8, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 8, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=16, N=16)\n#define ACCUM_FP16_16X16_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 16, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \\\n for (const Float16 *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP16_ITER_16X16_AVX512(m, q, zmm_sum, _mm256_load_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP16_ITER_16X16_AVX512(m, q, zmm_sum, _mm256_loadu_si256, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 16, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 16, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=1)\n#define ACCUM_FP16_32X1_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const Float16 *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP16_ITER_32X1_AVX512(m, q, zmm_sum, _mm512_load_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP16_ITER_32X1_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=2)\n#define ACCUM_FP16_32X2_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP16_ITER_32X2_AVX512(m, q, zmm_sum, _mm512_load_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP16_ITER_32X2_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=4)\n#define ACCUM_FP16_32X4_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP16_ITER_32X4_AVX512(m, q, zmm_sum, _mm512_load_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP16_ITER_32X4_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=8)\n#define ACCUM_FP16_32X8_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP16_ITER_32X8_AVX512(m, q, zmm_sum, _mm512_load_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP16_ITER_32X8_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=16)\n#define ACCUM_FP16_32X16_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 16, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const Float16 *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP16_ITER_32X16_AVX512(m, q, zmm_sum, _mm512_load_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP16_ITER_32X16_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 16, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 16, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP16, M=32, N=32)\n#define ACCUM_FP16_32X32_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 32, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const Float16 *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP16_ITER_32X32_AVX512(m, q, zmm_sum, _mm512_load_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const Float16 *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP16_ITER_32X32_AVX512(m, q, zmm_sum, _mm512_loadu_si512, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 32, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 32, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)\n//! Compute the distance between matrix and query (FP16, M=1, N=1)\n#define ACCUM_FP16_1X1_NEON(m, q, dim, out, _MASK, _NORM) \\\n MATRIX_VAR_INIT(1, 1, float16x8_t, v_sum, vdupq_n_f16(0)) \\\n const Float16 *qe = q + dim; \\\n const Float16 *qe_aligned = q + ((dim >> 3) << 3); \\\n for (; q != qe_aligned; m += 8, q += 8) { \\\n MATRIX_FP16_ITER_1X1_NEON(m, q, v_sum, ACCUM_FP16_STEP_NEON) \\\n } \\\n if (qe >= qe_aligned + 4) { \\\n float16x8_t v_m = \\\n vcombine_f16(vld1_f16((const float16_t *)m), \\\n vreinterpret_f16_u64(vdup_n_u64((uint64_t)(_MASK)))); \\\n float16x8_t v_q = \\\n vcombine_f16(vld1_f16((const float16_t *)q), \\\n vreinterpret_f16_u64(vdup_n_u64((uint64_t)(_MASK)))); \\\n ACCUM_FP16_STEP_NEON(v_m, v_q, v_sum_0_0) \\\n m += 4; \\\n q += 4; \\\n } \\\n float result = vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(v_sum_0_0)), \\\n vcvt_high_f32_f16(v_sum_0_0))); \\\n switch (qe - q) { \\\n case 3: \\\n ACCUM_FP16_STEP_GENERAL(m[2], q[2], result) \\\n /* FALLTHRU */ \\\n case 2: \\\n ACCUM_FP16_STEP_GENERAL(m[1], q[1], result) \\\n /* FALLTHRU */ \\\n case 1: \\\n ACCUM_FP16_STEP_GENERAL(m[0], q[0], result) \\\n } \\\n *out = _NORM(result);\n\n#else\n//! Compute the distance between matrix and query (FP16, M=1, N=1)\n#define ACCUM_FP16_1X1_NEON(m, q, dim, out, _MASK, _NORM) \\\n MATRIX_VAR_INIT(1, 1, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n const Float16 *qe = q + dim; \\\n const Float16 *qe_aligned = q + ((dim >> 3) << 3); \\\n for (; q != qe_aligned; m += 8, q += 8) { \\\n MATRIX_FP16_ITER_1X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n if (qe >= qe_aligned + 4) { \\\n float32x4_t v_m = vcvt_f32_f16(vld1_f16((const float16_t *)m)); \\\n float32x4_t v_q = vcvt_f32_f16(vld1_f16((const float16_t *)q)); \\\n ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum_0_0) \\\n m += 4; \\\n q += 4; \\\n } \\\n float result = vaddvq_f32(v_sum_0_0); \\\n switch (qe - q) { \\\n case 3: \\\n ACCUM_FP16_STEP_GENERAL(m[2], q[2], result) \\\n /* FALLTHRU */ \\\n case 2: \\\n ACCUM_FP16_STEP_GENERAL(m[1], q[1], result) \\\n /* FALLTHRU */ \\\n case 1: \\\n ACCUM_FP16_STEP_GENERAL(m[0], q[0], result) \\\n } \\\n *out = _NORM(result);\n\n#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC\n" }, { "path": "src/ailego/math/distance_matrix_accum_fp32.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_fp32.i\"\n#include \"matrix_utility.i\"\n\n#if !defined(__FMA__)\n#define _mm_fmadd_ps(a, b, c) _mm_add_ps(_mm_mul_ps((a), (b)), (c))\n#define _mm256_fmadd_ps(a, b, c) _mm256_add_ps(_mm256_mul_ps((a), (b)), (c))\n#endif // !__FMA__\n\n#if defined(__AVX512F__) && !defined(__AVX512DQ__)\n#define _mm512_and_ps(a, b) \\\n _mm512_castsi512_ps( \\\n _mm512_and_epi32(_mm512_castps_si512(a), _mm512_castps_si512(b)))\n#define _mm512_mask_and_ps(src, k, a, b) \\\n _mm512_castsi512_ps(_mm512_mask_and_epi32(_mm512_castps_si512(src), (k), \\\n _mm512_castps_si512(a), \\\n _mm512_castps_si512(b)))\n#endif // __AVX512DQ__\n\n#if defined(__ARM_NEON) && !defined(__aarch64__)\nstatic inline float32_t vaddvq_f32(float32x4_t v) {\n float32x2_t s = vadd_f32(vget_low_f32(v), vget_high_f32(v));\n return vget_lane_f32(vpadd_f32(s, s), 0);\n}\n\nstatic inline int32_t vaddvq_s32(int32x4_t v) {\n int32x2_t s = vadd_s32(vget_low_s32(v), vget_high_s32(v));\n return vget_lane_s32(vpadd_s32(s, s), 0);\n}\n#endif //__ARM_NEON && !__aarch64__\n\n#if defined(__aarch64__)\n#define ACCUM_FP32_2X1_NEON ACCUM_FP32_2X1_NEON_A64\n#else\n#define ACCUM_FP32_2X1_NEON ACCUM_FP32_2X1_NEON_A32\n#endif // __aarch64__\n\n//! Compute the distance between matrix and query (FP32, M=2, N=1)\n#define ACCUM_FP32_2X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128, xmm_sum, _mm_setzero_ps()) \\\n const float *qe_aligned = q + ((dim >> 2) << 2); \\\n const float *qe = q + dim; \\\n if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \\\n for (; q != qe_aligned; m += 8, q += 4) { \\\n MATRIX_FP32_ITER_2X1_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n if (qe >= qe_aligned + 2) { \\\n __m128 xmm_m = _mm_load_ps(m); \\\n __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n m += 4; \\\n q += 2; \\\n } \\\n } else { \\\n for (; q != qe_aligned; m += 8, q += 4) { \\\n MATRIX_FP32_ITER_2X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n if (qe >= qe_aligned + 2) { \\\n __m128 xmm_m = _mm_loadu_ps(m); \\\n __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n m += 4; \\\n q += 2; \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_ps(xmm_sum_0_0, xmm_sum_0_1); \\\n xmm_sum_0_0 = \\\n _mm_add_ps(xmm_sum_0_0, _mm_movehl_ps(xmm_sum_0_0, xmm_sum_0_0)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_set_ps(0.0f, 0.0f, m[1], m[0]); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0));\n\n//! Compute the distance between matrix and query (FP32, M=2, N=2)\n#define ACCUM_FP32_2X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128, xmm_sum, _mm_setzero_ps()) \\\n const float *qe = q + (dim << 1); \\\n if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \\\n m += 4, q += 4) { \\\n MATRIX_FP32_ITER_2X2_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \\\n m += 4, q += 4) { \\\n MATRIX_FP32_ITER_2X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_ps(_mm_movelh_ps(xmm_sum_0_0, xmm_sum_0_1), \\\n _mm_movehl_ps(xmm_sum_0_1, xmm_sum_0_0)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_set_ps(m[1], m[0], m[1], m[0]); \\\n __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=4, N=1)\n#define ACCUM_FP32_4X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128, xmm_sum, _mm_setzero_ps()) \\\n const float *qe = q + dim; \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \\\n m += 8, q += 2) { \\\n MATRIX_FP32_ITER_4X1_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_load_ps(m); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n } else { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \\\n m += 8, q += 2) { \\\n MATRIX_FP32_ITER_4X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_loadu_ps(m); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_ps(xmm_sum_0_0, xmm_sum_0_1); \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=4, N=2)\n#define ACCUM_FP32_4X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 4, q += 2) { \\\n MATRIX_FP32_ITER_4X2_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 4, q += 2) { \\\n MATRIX_FP32_ITER_4X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=4, N=4)\n#define ACCUM_FP32_4X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 4, q += 4) { \\\n MATRIX_FP32_ITER_4X4_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 4, q += 4) { \\\n MATRIX_FP32_ITER_4X4_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=8, N=1)\n#define ACCUM_FP32_8X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + dim; q != qe; m += 8, ++q) { \\\n MATRIX_FP32_ITER_8X1_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + dim; q != qe; m += 8, ++q) { \\\n MATRIX_FP32_ITER_8X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=8, N=2)\n#define ACCUM_FP32_8X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \\\n MATRIX_FP32_ITER_8X2_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \\\n MATRIX_FP32_ITER_8X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=8, N=4)\n#define ACCUM_FP32_8X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \\\n MATRIX_FP32_ITER_8X4_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \\\n MATRIX_FP32_ITER_8X4_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=8, N=8)\n#define ACCUM_FP32_8X8_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \\\n MATRIX_FP32_ITER_8X8_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \\\n MATRIX_FP32_ITER_8X8_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=1)\n#define ACCUM_FP32_16X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP32_ITER_16X1_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP32_ITER_16X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=2)\n#define ACCUM_FP32_16X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP32_ITER_16X2_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP32_ITER_16X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=4)\n#define ACCUM_FP32_16X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP32_ITER_16X4_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP32_ITER_16X4_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=8)\n#define ACCUM_FP32_16X8_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP32_ITER_16X8_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP32_ITER_16X8_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=16)\n#define ACCUM_FP32_16X16_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP32_ITER_16X16_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP32_ITER_16X16_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=1)\n#define ACCUM_FP32_32X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 1, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP32_ITER_32X1_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP32_ITER_32X1_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=2)\n#define ACCUM_FP32_32X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 2, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP32_ITER_32X2_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP32_ITER_32X2_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=4)\n#define ACCUM_FP32_32X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 4, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP32_ITER_32X4_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP32_ITER_32X4_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=8)\n#define ACCUM_FP32_32X8_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 8, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP32_ITER_32X8_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP32_ITER_32X8_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=16)\n#define ACCUM_FP32_32X16_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 16, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP32_ITER_32X16_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP32_ITER_32X16_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=32)\n#define ACCUM_FP32_32X32_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 32, __m128, xmm_sum, _mm_setzero_ps()) \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP32_ITER_32X32_SSE(m, q, xmm_sum, _mm_load_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP32_ITER_32X32_SSE(m, q, xmm_sum, _mm_loadu_ps, \\\n ACCUM_FP32_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=2, N=1)\n#define ACCUM_FP32_2X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const float *qe_aligned = q + ((dim >> 2) << 2); \\\n const float *qe = q + dim; \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_aligned; m += 8, q += 4) { \\\n MATRIX_FP32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (; q != qe_aligned; m += 8, q += 4) { \\\n MATRIX_FP32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n if (qe >= qe_aligned + 2) { \\\n __m128 xmm_m = _mm_loadu_ps(m); \\\n __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n m += 4; \\\n q += 2; \\\n } \\\n xmm_sum_0_0 = \\\n _mm_add_ps(xmm_sum_0_0, _mm_movehl_ps(xmm_sum_0_0, xmm_sum_0_0)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_set_ps(0.0f, 0.0f, m[1], m[0]); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0));\n\n//! Compute the distance between matrix and query (FP32, M=2, N=2)\n#define ACCUM_FP32_2X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const float *qe_aligned = q + ((dim >> 2) << 3); \\\n const float *qe = q + (dim << 1); \\\n if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \\\n for (; q != qe_aligned; m += 8, q += 8) { \\\n MATRIX_FP32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (; q != qe_aligned; m += 8, q += 8) { \\\n MATRIX_FP32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \\\n _mm256_extractf128_ps(ymm_sum_0_1, 1)); \\\n if (qe >= qe_aligned + 4) { \\\n __m128 xmm_q = _mm_loadu_ps(q); \\\n __m128 xmm_m = _mm_loadu_ps(m); \\\n __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(2, 2, 0, 0)); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_0) \\\n xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 1, 1)); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_p, xmm_sum_0_1) \\\n m += 4; \\\n q += 4; \\\n } \\\n xmm_sum_0_0 = _mm_add_ps(_mm_movelh_ps(xmm_sum_0_0, xmm_sum_0_1), \\\n _mm_movehl_ps(xmm_sum_0_1, xmm_sum_0_0)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_set_ps(m[1], m[0], m[1], m[0]); \\\n __m128 xmm_q = _mm_set_ps(q[1], q[1], q[0], q[0]); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=4, N=1)\n#define ACCUM_FP32_4X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const float *qe = q + dim; \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \\\n m += 8, q += 2) { \\\n MATRIX_FP32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 1); q != qe_aligned; \\\n m += 8, q += 2) { \\\n MATRIX_FP32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_loadu_ps(m); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=4, N=2)\n#define ACCUM_FP32_4X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const float *qe = q + (dim << 1); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \\\n m += 8, q += 4) { \\\n MATRIX_FP32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 2); q != qe_aligned; \\\n m += 8, q += 4) { \\\n MATRIX_FP32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \\\n _mm256_extractf128_ps(ymm_sum_0_1, 1)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_loadu_ps(m); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_1) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=4, N=4)\n#define ACCUM_FP32_4X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const float *qe = q + (dim << 2); \\\n if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 3); q != qe_aligned; \\\n m += 8, q += 8) { \\\n MATRIX_FP32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe_aligned = q + ((dim >> 1) << 3); q != qe_aligned; \\\n m += 8, q += 8) { \\\n MATRIX_FP32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n __m128 xmm_sum_0_0 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_0), \\\n _mm256_extractf128_ps(ymm_sum_0_0, 1)); \\\n __m128 xmm_sum_0_1 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_1), \\\n _mm256_extractf128_ps(ymm_sum_0_1, 1)); \\\n __m128 xmm_sum_0_2 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_2), \\\n _mm256_extractf128_ps(ymm_sum_0_2, 1)); \\\n __m128 xmm_sum_0_3 = _mm_add_ps(_mm256_castps256_ps128(ymm_sum_0_3), \\\n _mm256_extractf128_ps(ymm_sum_0_3, 1)); \\\n if (q != qe) { \\\n __m128 xmm_m = _mm_loadu_ps(m); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_1) \\\n xmm_q = _mm_broadcast_ss(q + 2); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_2) \\\n xmm_q = _mm_broadcast_ss(q + 3); \\\n ACCUM_FP32_STEP_SSE(xmm_m, xmm_q, xmm_sum_0_3) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=8, N=1)\n#define ACCUM_FP32_8X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + dim; q != qe; m += 8, ++q) { \\\n MATRIX_FP32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + dim; q != qe; m += 8, ++q) { \\\n MATRIX_FP32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=8, N=2)\n#define ACCUM_FP32_8X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \\\n MATRIX_FP32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \\\n MATRIX_FP32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=8, N=4)\n#define ACCUM_FP32_8X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \\\n MATRIX_FP32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \\\n MATRIX_FP32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=8, N=8)\n#define ACCUM_FP32_8X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 8, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \\\n MATRIX_FP32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \\\n MATRIX_FP32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=1)\n#define ACCUM_FP32_16X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=2)\n#define ACCUM_FP32_16X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=4)\n#define ACCUM_FP32_16X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=8)\n#define ACCUM_FP32_16X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=16)\n#define ACCUM_FP32_16X16_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 16, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=1)\n#define ACCUM_FP32_32X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=2)\n#define ACCUM_FP32_32X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=4)\n#define ACCUM_FP32_32X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=8)\n#define ACCUM_FP32_32X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=16)\n#define ACCUM_FP32_32X16_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=32)\n#define ACCUM_FP32_32X32_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 32, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_ps, \\\n ACCUM_FP32_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=1)\n#define ACCUM_FP32_16X1_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP32_ITER_16X1_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP32_ITER_16X1_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=2)\n#define ACCUM_FP32_16X2_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP32_ITER_16X2_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP32_ITER_16X2_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=4)\n#define ACCUM_FP32_16X4_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP32_ITER_16X4_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP32_ITER_16X4_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=8)\n#define ACCUM_FP32_16X8_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 8, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP32_ITER_16X8_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP32_ITER_16X8_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 8, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 8, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=16, N=16)\n#define ACCUM_FP32_16X16_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 16, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP32_ITER_16X16_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP32_ITER_16X16_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(1, 16, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 16, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=1)\n#define ACCUM_FP32_32X1_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP32_ITER_32X1_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP32_ITER_32X1_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=2)\n#define ACCUM_FP32_32X2_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP32_ITER_32X2_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP32_ITER_32X2_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=4)\n#define ACCUM_FP32_32X4_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP32_ITER_32X4_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP32_ITER_32X4_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=8)\n#define ACCUM_FP32_32X8_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP32_ITER_32X8_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP32_ITER_32X8_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=16)\n#define ACCUM_FP32_32X16_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 16, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP32_ITER_32X16_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP32_ITER_32X16_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 16, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 16, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=32, N=32)\n#define ACCUM_FP32_32X32_AVX512(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 32, __m512, zmm_sum, _mm512_setzero_ps()) \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP32_ITER_32X32_AVX512(m, q, zmm_sum, _mm512_load_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } else { \\\n for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP32_ITER_32X32_AVX512(m, q, zmm_sum, _mm512_loadu_ps, \\\n ACCUM_FP32_STEP_AVX512) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x3f) == 0) { \\\n MATRIX_VAR_STORE(2, 32, 16, zmm_sum, out, _mm512_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 32, 16, zmm_sum, out, _mm512_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (FP32, M=2, N=1) on A64\n#define ACCUM_FP32_2X1_NEON_A64(m, q, dim, out, _NORM) \\\n float32x4_t v_sum = vdupq_n_f32(0); \\\n const float *qe_aligned = q + ((dim >> 1) << 1); \\\n const float *qe = q + dim; \\\n for (; q != qe_aligned; m += 4, q += 2) { \\\n MATRIX_FP32_ITER_2X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n v_sum = vaddq_f32( \\\n vreinterpretq_f32_u64(vdupq_laneq_u64(vreinterpretq_u64_f32(v_sum), 1)), \\\n v_sum); \\\n if (q != qe) { \\\n float32x4_t v_m = vreinterpretq_f32_u64( \\\n vdupq_lane_u64(vld1_u64((const uint64_t *)m), 0)); \\\n float32x4_t v_q = vld1q_dup_f32(q); \\\n ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum) \\\n } \\\n vst1_f32(out, _NORM(vget_low_f32(v_sum)));\n\n//! Compute the distance between matrix and query (FP32, M=2, N=1) on A32\n#define ACCUM_FP32_2X1_NEON_A32(m, q, dim, out, _NORM) \\\n float32x4_t v_sum = vdupq_n_f32(0); \\\n const float *qe_aligned = q + ((dim >> 1) << 1); \\\n const float *qe = q + dim; \\\n for (; q != qe_aligned; m += 4, q += 2) { \\\n MATRIX_FP32_ITER_2X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n float32x2_t sum = vadd_f32(vget_low_f32(v_sum), vget_high_f32(v_sum)); \\\n v_sum = vcombine_f32(sum, sum); \\\n if (q != qe) { \\\n float32x4_t v_m = vreinterpretq_f32_u64( \\\n vdupq_lane_u64(vld1_u64((const uint64_t *)m), 0)); \\\n float32x4_t v_q = vld1q_dup_f32(q); \\\n ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum) \\\n } \\\n vst1_f32(out, _NORM(vget_low_f32(v_sum)));\n\n//! Compute the distance between matrix and query (FP32, M=2, N=2)\n#define ACCUM_FP32_2X2_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n const float *qe_aligned = q + ((dim >> 1) << 2); \\\n const float *qe = q + (dim << 1); \\\n for (; q != qe_aligned; m += 4, q += 4) { \\\n MATRIX_FP32_ITER_2X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n v_sum_0_0 = vaddq_f32( \\\n vcombine_f32(vget_low_f32(v_sum_0_0), vget_low_f32(v_sum_0_1)), \\\n vcombine_f32(vget_high_f32(v_sum_0_0), vget_high_f32(v_sum_0_1))); \\\n if (q != qe) { \\\n float32x2_t v_m_0 = vld1_f32(m); \\\n float32x2_t v_q_0 = vld1_f32(q); \\\n float32x4_t v_m = vcombine_f32(v_m_0, v_m_0); \\\n float32x4_t v_q = \\\n vcombine_f32(vdup_lane_f32(v_q_0, 0), vdup_lane_f32(v_q_0, 1)); \\\n ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum_0_0) \\\n } \\\n MATRIX_VAR_STORE(1, 1, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=4, N=1)\n#define ACCUM_FP32_4X1_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n const float *qe_aligned = q + ((dim >> 1) << 1); \\\n const float *qe = q + dim; \\\n for (; q != qe_aligned; m += 8, q += 2) { \\\n MATRIX_FP32_ITER_4X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n if (q != qe) { \\\n float32x4_t v_m = vld1q_f32(m); \\\n float32x4_t v_q = vld1q_dup_f32(q); \\\n ACCUM_FP32_STEP_NEON(v_m, v_q, v_sum_0_0) \\\n } \\\n v_sum_0_0 = vaddq_f32(v_sum_0_0, v_sum_0_1); \\\n MATRIX_VAR_STORE(1, 1, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=4, N=2)\n#define ACCUM_FP32_4X2_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 1); q != qe; m += 4, q += 2) { \\\n MATRIX_FP32_ITER_4X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(1, 2, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=4, N=4)\n#define ACCUM_FP32_4X4_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 2); q != qe; m += 4, q += 4) { \\\n MATRIX_FP32_ITER_4X4_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(1, 4, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=8, N=1)\n#define ACCUM_FP32_8X1_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + dim; q != qe; m += 8, ++q) { \\\n MATRIX_FP32_ITER_8X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(2, 1, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=8, N=2)\n#define ACCUM_FP32_8X2_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 1); q != qe; m += 8, q += 2) { \\\n MATRIX_FP32_ITER_8X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(2, 2, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=8, N=4)\n#define ACCUM_FP32_8X4_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 2); q != qe; m += 8, q += 4) { \\\n MATRIX_FP32_ITER_8X4_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(2, 4, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=8, N=8)\n#define ACCUM_FP32_8X8_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 3); q != qe; m += 8, q += 8) { \\\n MATRIX_FP32_ITER_8X8_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(2, 8, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=16, N=1)\n#define ACCUM_FP32_16X1_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + dim; q != qe; m += 16, ++q) { \\\n MATRIX_FP32_ITER_16X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(4, 1, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=16, N=2)\n#define ACCUM_FP32_16X2_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 1); q != qe; m += 16, q += 2) { \\\n MATRIX_FP32_ITER_16X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(4, 2, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=16, N=4)\n#define ACCUM_FP32_16X4_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 2); q != qe; m += 16, q += 4) { \\\n MATRIX_FP32_ITER_16X4_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(4, 4, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=16, N=8)\n#define ACCUM_FP32_16X8_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 3); q != qe; m += 16, q += 8) { \\\n MATRIX_FP32_ITER_16X8_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(4, 8, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=16, N=16)\n#define ACCUM_FP32_16X16_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 4); q != qe; m += 16, q += 16) { \\\n MATRIX_FP32_ITER_16X16_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(4, 16, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=32, N=1)\n#define ACCUM_FP32_32X1_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 1, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + dim; q != qe; m += 32, ++q) { \\\n MATRIX_FP32_ITER_32X1_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(8, 1, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=32, N=2)\n#define ACCUM_FP32_32X2_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 1); q != qe; m += 32, q += 2) { \\\n MATRIX_FP32_ITER_32X2_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(8, 2, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=32, N=4)\n#define ACCUM_FP32_32X4_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 4, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 2); q != qe; m += 32, q += 4) { \\\n MATRIX_FP32_ITER_32X4_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(8, 4, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=32, N=8)\n#define ACCUM_FP32_32X8_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 8, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 3); q != qe; m += 32, q += 8) { \\\n MATRIX_FP32_ITER_32X8_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(8, 8, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=32, N=16)\n#define ACCUM_FP32_32X16_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 16, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 4); q != qe; m += 32, q += 16) { \\\n MATRIX_FP32_ITER_32X16_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(8, 16, 4, v_sum, out, vst1q_f32, _NORM)\n\n//! Compute the distance between matrix and query (FP32, M=32, N=32)\n#define ACCUM_FP32_32X32_NEON(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 32, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n for (const float *qe = q + (dim << 5); q != qe; m += 32, q += 32) { \\\n MATRIX_FP32_ITER_32X32_NEON(m, q, v_sum, ACCUM_FP32_STEP_NEON) \\\n } \\\n MATRIX_VAR_STORE(8, 32, 4, v_sum, out, vst1q_f32, _NORM)\n" }, { "path": "src/ailego/math/distance_matrix_accum_int4.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_int32.i\"\n#include \"matrix_utility.i\"\n\n//! Compute the distance between matrix and query (INT4, M=2, N=1)\n#define ACCUM_INT4_2X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe_aligned = qi + ((dim >> 5) << 2); \\\n const uint32_t *qe = qi + (dim >> 3); \\\n if (((uintptr_t)mi & 0xf) == 0 && ((uintptr_t)qi & 0xf) == 0) { \\\n for (; qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n if (qe >= qe_aligned + 2) { \\\n __m128i xmm_mi = _mm_load_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n mi += 4; \\\n qi += 2; \\\n } \\\n } else { \\\n for (; qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n if (qe >= qe_aligned + 2) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n mi += 4; \\\n qi += 2; \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_0_1); \\\n xmm_sum_0_0 = _mm_add_epi32( \\\n xmm_sum_0_0, _mm_shuffle_epi32(xmm_sum_0_0, _MM_SHUFFLE(0, 0, 3, 2))); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_set_epi32(0, 0, mi[1], mi[0]); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0));\n\n//! Compute the distance between matrix and query (INT4, M=2, N=2)\n#define ACCUM_INT4_2X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + ((dim >> 3) << 1); \\\n if (((uintptr_t)mi & 0xf) == 0 && ((uintptr_t)qi & 0xf) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \\\n qi != qe_aligned; mi += 4, qi += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \\\n qi != qe_aligned; mi += 4, qi += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \\\n _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_set_epi32(mi[1], mi[0], mi[1], mi[0]); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=4, N=1)\n#define ACCUM_INT4_4X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + (dim >> 3); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 1); \\\n qi != qe_aligned; mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_load_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 1); \\\n qi != qe_aligned; mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_1_0); \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=4, N=2)\n#define ACCUM_INT4_4X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 4, qi += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 4, qi += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=4, N=4)\n#define ACCUM_INT4_4X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 4, qi += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 4, qi += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=8, N=1)\n#define ACCUM_INT4_8X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 8, ++qi) { \\\n MATRIX_INT32_ITER_8X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 8, ++qi) { \\\n MATRIX_INT32_ITER_8X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=8, N=2)\n#define ACCUM_INT4_8X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=8, N=4)\n#define ACCUM_INT4_8X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=8, N=8)\n#define ACCUM_INT4_8X8_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=1)\n#define ACCUM_INT4_16X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 16, ++qi) { \\\n MATRIX_INT32_ITER_16X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 16, ++qi) { \\\n MATRIX_INT32_ITER_16X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=2)\n#define ACCUM_INT4_16X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 16, qi += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 16, qi += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=4)\n#define ACCUM_INT4_16X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 16, qi += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 16, qi += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=8)\n#define ACCUM_INT4_16X8_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 16, qi += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 16, qi += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=16)\n#define ACCUM_INT4_16X16_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \\\n mi += 16, qi += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \\\n mi += 16, qi += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=1)\n#define ACCUM_INT4_32X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 32, ++qi) { \\\n MATRIX_INT32_ITER_32X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 32, ++qi) { \\\n MATRIX_INT32_ITER_32X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=2)\n#define ACCUM_INT4_32X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 32, qi += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 32, qi += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=4)\n#define ACCUM_INT4_32X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 32, qi += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 32, qi += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=8)\n#define ACCUM_INT4_32X8_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 8, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 32, qi += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 32, qi += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=16)\n#define ACCUM_INT4_32X16_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 16, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \\\n mi += 32, qi += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \\\n mi += 32, qi += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=32)\n#define ACCUM_INT4_32X32_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 32, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 5); qi != qe; \\\n mi += 32, qi += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 5); qi != qe; \\\n mi += 32, qi += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT4_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=2, N=1)\n#define ACCUM_INT4_2X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe_aligned = qi + ((dim >> 5) << 2); \\\n const uint32_t *qe = qi + (dim >> 3); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (; qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (; qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0 = _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n if (qe >= qe_aligned + 2) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0) \\\n mi += 4; \\\n qi += 2; \\\n } \\\n xmm_sum_0 = _mm_add_epi32( \\\n xmm_sum_0, _mm_shuffle_epi32(xmm_sum_0, _MM_SHUFFLE(0, 0, 3, 2))); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_set_epi32(0, 0, mi[1], mi[0]); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0));\n\n//! Compute the distance between matrix and query (INT4, M=2, N=2)\n#define ACCUM_INT4_2X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe_aligned = qi + ((dim >> 5) << 3); \\\n const uint32_t *qe = qi + ((dim >> 3) << 1); \\\n if (((uintptr_t)mi & 0x1f) == 0 && ((uintptr_t)qi & 0x1f) == 0) { \\\n for (; qi != qe_aligned; mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (; qi != qe_aligned; mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n __m128i xmm_sum_0_1 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \\\n _mm256_extracti128_si256(ymm_sum_0_1, 1)); \\\n if (qe >= qe_aligned + 4) { \\\n __m128i xmm_qi = _mm_loadu_si128((const __m128i *)(qi)); \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_pi, xmm_sum_0_0) \\\n xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_pi, xmm_sum_0_1) \\\n mi += 4; \\\n qi += 4; \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \\\n _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_set_epi32(mi[1], mi[0], mi[1], mi[0]); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=4, N=1)\n#define ACCUM_INT4_4X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + (dim >> 3); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 1); \\\n qi != qe_aligned; mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 1); \\\n qi != qe_aligned; mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=4, N=2)\n#define ACCUM_INT4_4X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + ((dim >> 3) << 1); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \\\n qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \\\n qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n __m128i xmm_sum_0_1 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \\\n _mm256_extracti128_si256(ymm_sum_0_1, 1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_1) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=4, N=4)\n#define ACCUM_INT4_4X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + ((dim >> 3) << 2); \\\n if (((uintptr_t)mi & 0x1f) == 0 && ((uintptr_t)qi & 0x1f) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 3); \\\n qi != qe_aligned; mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 4) << 3); \\\n qi != qe_aligned; mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n __m128i xmm_sum_0_1 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \\\n _mm256_extracti128_si256(ymm_sum_0_1, 1)); \\\n __m128i xmm_sum_0_2 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_2), \\\n _mm256_extracti128_si256(ymm_sum_0_2, 1)); \\\n __m128i xmm_sum_0_3 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_3), \\\n _mm256_extracti128_si256(ymm_sum_0_3, 1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_1) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_2) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n ACCUM_INT4_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_3) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=8, N=1)\n#define ACCUM_INT4_8X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 8, ++qi) { \\\n MATRIX_INT32_ITER_8X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 8, ++qi) { \\\n MATRIX_INT32_ITER_8X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=8, N=2)\n#define ACCUM_INT4_8X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=8, N=4)\n#define ACCUM_INT4_8X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=8, N=8)\n#define ACCUM_INT4_8X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=1)\n#define ACCUM_INT4_16X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 16, ++qi) { \\\n MATRIX_INT32_ITER_16X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 16, ++qi) { \\\n MATRIX_INT32_ITER_16X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=2)\n#define ACCUM_INT4_16X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 16, qi += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 16, qi += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=4)\n#define ACCUM_INT4_16X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 16, qi += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 16, qi += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=8)\n#define ACCUM_INT4_16X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 16, qi += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 16, qi += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=16, N=16)\n#define ACCUM_INT4_16X16_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \\\n mi += 16, qi += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \\\n mi += 16, qi += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=1)\n#define ACCUM_INT4_32X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 32, ++qi) { \\\n MATRIX_INT32_ITER_32X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 3); qi != qe; mi += 32, ++qi) { \\\n MATRIX_INT32_ITER_32X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=2)\n#define ACCUM_INT4_32X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 32, qi += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 1); qi != qe; \\\n mi += 32, qi += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=4)\n#define ACCUM_INT4_32X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 32, qi += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 2); qi != qe; \\\n mi += 32, qi += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=8)\n#define ACCUM_INT4_32X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 32, qi += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 3); qi != qe; \\\n mi += 32, qi += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=16)\n#define ACCUM_INT4_32X16_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \\\n mi += 32, qi += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 4); qi != qe; \\\n mi += 32, qi += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT4, M=32, N=32)\n#define ACCUM_INT4_32X32_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 32, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 5); qi != qe; \\\n mi += 32, qi += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 3) << 5); qi != qe; \\\n mi += 32, qi += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT4_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n" }, { "path": "src/ailego/math/distance_matrix_accum_int8.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_int32.i\"\n#include \"matrix_utility.i\"\n\n//! Compute the distance between matrix and query (INT8, M=2, N=1)\n#define ACCUM_INT8_2X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \\\n const uint32_t *qe = qi + (dim >> 2); \\\n if (((uintptr_t)mi & 0xf) == 0 && ((uintptr_t)qi & 0xf) == 0) { \\\n for (; qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n if (qe >= qe_aligned + 2) { \\\n __m128i xmm_mi = _mm_load_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n mi += 4; \\\n qi += 2; \\\n } \\\n } else { \\\n for (; qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n if (qe >= qe_aligned + 2) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n mi += 4; \\\n qi += 2; \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_0_1); \\\n xmm_sum_0_0 = _mm_add_epi32( \\\n xmm_sum_0_0, _mm_shuffle_epi32(xmm_sum_0_0, _MM_SHUFFLE(0, 0, 3, 2))); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_set_epi32(0, 0, mi[1], mi[0]); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0));\n\n//! Compute the distance between matrix and query (INT8, M=2, N=2)\n#define ACCUM_INT8_2X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + ((dim >> 2) << 1); \\\n if (((uintptr_t)mi & 0xf) == 0 && ((uintptr_t)qi & 0xf) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 2); \\\n qi != qe_aligned; mi += 4, qi += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 2); \\\n qi != qe_aligned; mi += 4, qi += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \\\n _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_set_epi32(mi[1], mi[0], mi[1], mi[0]); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=4, N=1)\n#define ACCUM_INT8_4X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + (dim >> 2); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 1); \\\n qi != qe_aligned; mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_load_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 1); \\\n qi != qe_aligned; mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_1_0); \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=4, N=2)\n#define ACCUM_INT8_4X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 4, qi += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 4, qi += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=4, N=4)\n#define ACCUM_INT8_4X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 4, qi += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 4, qi += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=8, N=1)\n#define ACCUM_INT8_8X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 8, ++qi) { \\\n MATRIX_INT32_ITER_8X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 8, ++qi) { \\\n MATRIX_INT32_ITER_8X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=8, N=2)\n#define ACCUM_INT8_8X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=8, N=4)\n#define ACCUM_INT8_8X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=8, N=8)\n#define ACCUM_INT8_8X8_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=1)\n#define ACCUM_INT8_16X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 16, ++qi) { \\\n MATRIX_INT32_ITER_16X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 16, ++qi) { \\\n MATRIX_INT32_ITER_16X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=2)\n#define ACCUM_INT8_16X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 16, qi += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 16, qi += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=4)\n#define ACCUM_INT8_16X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 16, qi += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 16, qi += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=8)\n#define ACCUM_INT8_16X8_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 16, qi += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 16, qi += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=16)\n#define ACCUM_INT8_16X16_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \\\n mi += 16, qi += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \\\n mi += 16, qi += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=1)\n#define ACCUM_INT8_32X1_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 32, ++qi) { \\\n MATRIX_INT32_ITER_32X1_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 32, ++qi) { \\\n MATRIX_INT32_ITER_32X1_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=2)\n#define ACCUM_INT8_32X2_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 32, qi += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 32, qi += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=4)\n#define ACCUM_INT8_32X4_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 32, qi += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 32, qi += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=8)\n#define ACCUM_INT8_32X8_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 8, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 32, qi += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 32, qi += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=16)\n#define ACCUM_INT8_32X16_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 16, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \\\n mi += 32, qi += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \\\n mi += 32, qi += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=32)\n#define ACCUM_INT8_32X32_SSE(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(8, 32, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0xf) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 5); qi != qe; \\\n mi += 32, qi += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(mi, qi, xmm_sum, _mm_load_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 5); qi != qe; \\\n mi += 32, qi += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(mi, qi, xmm_sum, _mm_loadu_si128, \\\n ACCUM_INT8_STEP_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=2, N=1)\n#define ACCUM_INT8_2X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe_aligned = qi + ((dim >> 4) << 2); \\\n const uint32_t *qe = qi + (dim >> 2); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (; qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (; qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0 = _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n if (qe >= qe_aligned + 2) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0) \\\n mi += 4; \\\n qi += 2; \\\n } \\\n xmm_sum_0 = _mm_add_epi32( \\\n xmm_sum_0, _mm_shuffle_epi32(xmm_sum_0, _MM_SHUFFLE(0, 0, 3, 2))); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_set_epi32(0, 0, mi[1], mi[0]); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0));\n\n//! Compute the distance between matrix and query (INT8, M=2, N=2)\n#define ACCUM_INT8_2X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe_aligned = qi + ((dim >> 4) << 3); \\\n const uint32_t *qe = qi + ((dim >> 2) << 1); \\\n if (((uintptr_t)mi & 0x1f) == 0 && ((uintptr_t)qi & 0x1f) == 0) { \\\n for (; qi != qe_aligned; mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (; qi != qe_aligned; mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n __m128i xmm_sum_0_1 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \\\n _mm256_extracti128_si256(ymm_sum_0_1, 1)); \\\n if (qe >= qe_aligned + 4) { \\\n __m128i xmm_qi = _mm_loadu_si128((const __m128i *)(qi)); \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_pi, xmm_sum_0_0) \\\n xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_pi, xmm_sum_0_1) \\\n mi += 4; \\\n qi += 4; \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \\\n _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_set_epi32(mi[1], mi[0], mi[1], mi[0]); \\\n __m128i xmm_qi = _mm_set_epi32(qi[1], qi[1], qi[0], qi[0]); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=4, N=1)\n#define ACCUM_INT8_4X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + (dim >> 2); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 1); \\\n qi != qe_aligned; mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 1); \\\n qi != qe_aligned; mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=4, N=2)\n#define ACCUM_INT8_4X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + ((dim >> 2) << 1); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 2); \\\n qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 2); \\\n qi != qe_aligned; mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n __m128i xmm_sum_0_1 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \\\n _mm256_extracti128_si256(ymm_sum_0_1, 1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_1) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=4, N=4)\n#define ACCUM_INT8_4X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n const uint32_t *qe = qi + ((dim >> 2) << 2); \\\n if (((uintptr_t)mi & 0x1f) == 0 && ((uintptr_t)qi & 0x1f) == 0) { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 3); \\\n qi != qe_aligned; mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe_aligned = qi + ((dim >> 3) << 3); \\\n qi != qe_aligned; mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n __m128i xmm_sum_0_1 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \\\n _mm256_extracti128_si256(ymm_sum_0_1, 1)); \\\n __m128i xmm_sum_0_2 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_2), \\\n _mm256_extracti128_si256(ymm_sum_0_2, 1)); \\\n __m128i xmm_sum_0_3 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_3), \\\n _mm256_extracti128_si256(ymm_sum_0_3, 1)); \\\n if (qi != qe) { \\\n __m128i xmm_mi = _mm_loadu_si128((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_0) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_1) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_2) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n ACCUM_INT8_STEP_SSE(xmm_mi, xmm_qi, xmm_sum_0_3) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=8, N=1)\n#define ACCUM_INT8_8X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 8, ++qi) { \\\n MATRIX_INT32_ITER_8X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 8, ++qi) { \\\n MATRIX_INT32_ITER_8X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=8, N=2)\n#define ACCUM_INT8_8X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 8, qi += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=8, N=4)\n#define ACCUM_INT8_8X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 8, qi += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=8, N=8)\n#define ACCUM_INT8_8X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 8, qi += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=1)\n#define ACCUM_INT8_16X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 16, ++qi) { \\\n MATRIX_INT32_ITER_16X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 16, ++qi) { \\\n MATRIX_INT32_ITER_16X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=2)\n#define ACCUM_INT8_16X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 16, qi += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 16, qi += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=4)\n#define ACCUM_INT8_16X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 16, qi += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 16, qi += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=8)\n#define ACCUM_INT8_16X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 16, qi += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 16, qi += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=16, N=16)\n#define ACCUM_INT8_16X16_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(2, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \\\n mi += 16, qi += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \\\n mi += 16, qi += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=1)\n#define ACCUM_INT8_32X1_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 32, ++qi) { \\\n MATRIX_INT32_ITER_32X1_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + (dim >> 2); qi != qe; mi += 32, ++qi) { \\\n MATRIX_INT32_ITER_32X1_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=2)\n#define ACCUM_INT8_32X2_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 32, qi += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 1); qi != qe; \\\n mi += 32, qi += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=4)\n#define ACCUM_INT8_32X4_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 32, qi += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 2); qi != qe; \\\n mi += 32, qi += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=8)\n#define ACCUM_INT8_32X8_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 32, qi += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 3); qi != qe; \\\n mi += 32, qi += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=16)\n#define ACCUM_INT8_32X16_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \\\n mi += 32, qi += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 4); qi != qe; \\\n mi += 32, qi += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (INT8, M=32, N=32)\n#define ACCUM_INT8_32X32_AVX(m, q, dim, out, _NORM) \\\n MATRIX_VAR_INIT(4, 32, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qi = reinterpret_cast(q); \\\n const uint32_t *mi = reinterpret_cast(m); \\\n if (((uintptr_t)mi & 0x1f) == 0) { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 5); qi != qe; \\\n mi += 32, qi += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(mi, qi, ymm_sum, _mm256_load_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } else { \\\n for (const uint32_t *qe = qi + ((dim >> 2) << 5); qi != qe; \\\n mi += 32, qi += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(mi, qi, ymm_sum, _mm256_loadu_si256, \\\n ACCUM_INT8_STEP_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n" }, { "path": "src/ailego/math/distance_matrix_euclidean_utility.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n//! Calculate sum of squared difference (GENERAL)\n#define SSD_FP32_GENERAL(m, q, sum) \\\n { \\\n float x = m - q; \\\n sum += (x * x); \\\n }\n\n//! Calculate sum of squared difference (SSE)\n#define SSD_FP32_SSE(xmm_m, xmm_q, xmm_sum) \\\n { \\\n __m128 xmm_d = _mm_sub_ps(xmm_m, xmm_q); \\\n xmm_sum = _mm_fmadd_ps(xmm_d, xmm_d, xmm_sum); \\\n }\n\n//! Calculate sum of squared difference (AVX)\n#define SSD_FP32_AVX(ymm_m, ymm_q, ymm_sum) \\\n { \\\n __m256 ymm_d = _mm256_sub_ps(ymm_m, ymm_q); \\\n ymm_sum = _mm256_fmadd_ps(ymm_d, ymm_d, ymm_sum); \\\n }\n\n//! Calculate sum of squared difference (NEON)\n#define SSD_FP32_NEON(v_m, v_q, v_sum) \\\n { \\\n float32x4_t v_d = vsubq_f32(v_m, v_q); \\\n v_sum = vfmaq_f32(v_sum, v_d, v_d); \\\n }\n\n//! Calculate sum of squared difference (GENERAL)\n#define SSD_FP16_GENERAL(m, q, sum) \\\n { \\\n float x = m - q; \\\n sum += (x * x); \\\n }\n\n//! Calculate sum of squared difference (NEON)\n#define SSD_FP16_NEON(v_m, v_q, v_sum) \\\n { \\\n float16x8_t v_d = vsubq_f16(v_m, v_q); \\\n v_sum = vfmaq_f16(v_sum, v_d, v_d); \\\n }\n\n//! Calculate sum of squared difference (AVX512)\n#define SSD_FP32_AVX512(zmm_m, zmm_q, zmm_sum) \\\n { \\\n __m512 zmm_d = _mm512_sub_ps(zmm_m, zmm_q); \\\n zmm_sum = _mm512_fmadd_ps(zmm_d, zmm_d, zmm_sum); \\\n }\n\n//! Calculate sum of squared difference (GENERAL)\n#define SSD_INT4_GENERAL(m, q, sum) \\\n sum += Int4SquaredDiffTable[(((m) << 4) & 0xf0) | (((q) >> 0) & 0xf)] + \\\n Int4SquaredDiffTable[(((m) >> 0) & 0xf0) | (((q) >> 4) & 0xf)];\n\n\n#if defined(__SSE4_1__)\nstatic const __m128i MASK_INT4_SSE = _mm_set1_epi32(0xf0f0f0f0);\nstatic const __m128i ONES_INT16_SSE = _mm_set1_epi32(0x00010001);\n#endif // __SSE4_1__\n\n//! Compute the square root of value (SSE)\n#define SQRT_FP32_SSE(v, ...) _mm_sqrt_ps(_mm_cvtepi32_ps(v))\n\n#if defined(__AVX2__)\nstatic const __m256i MASK_INT4_AVX = _mm256_set1_epi32(0xf0f0f0f0);\nstatic const __m256i ONES_INT16_AVX = _mm256_set1_epi32(0x00010001);\n#endif // __AVX2__\n\n//! Calculate sum of squared difference (SSE)\n#define SSD_INT4_SSE(xmm_m, xmm_q, xmm_sum) \\\n { \\\n __m128i xmm_lhs = \\\n _mm_and_si128(_mm_slli_epi32((xmm_m), 4), MASK_INT4_SSE); \\\n __m128i xmm_rhs = \\\n _mm_and_si128(_mm_slli_epi32((xmm_q), 4), MASK_INT4_SSE); \\\n xmm_lhs = _mm_srli_epi32(_mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs), \\\n _mm_min_epi8(xmm_lhs, xmm_rhs)), \\\n 4); \\\n xmm_sum = _mm_add_epi32( \\\n _mm_madd_epi16(_mm_maddubs_epi16(xmm_lhs, xmm_lhs), ONES_INT16_SSE), \\\n xmm_sum); \\\n xmm_lhs = _mm_and_si128((xmm_m), MASK_INT4_SSE); \\\n xmm_rhs = _mm_and_si128((xmm_q), MASK_INT4_SSE); \\\n xmm_lhs = _mm_srli_epi32(_mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs), \\\n _mm_min_epi8(xmm_lhs, xmm_rhs)), \\\n 4); \\\n xmm_sum = _mm_add_epi32( \\\n _mm_madd_epi16(_mm_maddubs_epi16(xmm_lhs, xmm_lhs), ONES_INT16_SSE), \\\n xmm_sum); \\\n }\n\n//! Compute the distance between matrix and query\n#define SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) \\\n { \\\n __m128i xmm_lhs_0 = \\\n _mm_and_si128(_mm_slli_epi32((xmm_lhs), 4), MASK_INT4_SSE); \\\n __m128i xmm_rhs_0 = \\\n _mm_and_si128(_mm_slli_epi32((xmm_rhs), 4), MASK_INT4_SSE); \\\n __m128i xmm_lhs_1 = _mm_and_si128((xmm_lhs), MASK_INT4_SSE); \\\n __m128i xmm_rhs_1 = _mm_and_si128((xmm_rhs), MASK_INT4_SSE); \\\n xmm_lhs_0 = \\\n _mm_srli_epi32(_mm_sub_epi8(_mm_max_epi8(xmm_lhs_0, xmm_rhs_0), \\\n _mm_min_epi8(xmm_lhs_0, xmm_rhs_0)), \\\n 4); \\\n xmm_rhs_0 = \\\n _mm_srli_epi32(_mm_sub_epi8(_mm_max_epi8(xmm_lhs_1, xmm_rhs_1), \\\n _mm_min_epi8(xmm_lhs_1, xmm_rhs_1)), \\\n 4); \\\n xmm_lhs_0 = _mm_madd_epi16(_mm_maddubs_epi16(xmm_lhs_0, xmm_lhs_0), \\\n ONES_INT16_SSE); \\\n xmm_rhs_0 = _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_0, xmm_rhs_0), \\\n ONES_INT16_SSE); \\\n xmm_sum = _mm_add_epi32(_mm_add_epi32(xmm_lhs_0, xmm_rhs_0), xmm_sum); \\\n }\n\n//! Calculate sum of squared difference (AVX)\n#define SSD_INT4_AVX(ymm_m, ymm_q, ymm_sum) \\\n { \\\n __m256i ymm_lhs = \\\n _mm256_and_si256(_mm256_slli_epi32((ymm_m), 4), MASK_INT4_AVX); \\\n __m256i ymm_rhs = \\\n _mm256_and_si256(_mm256_slli_epi32((ymm_q), 4), MASK_INT4_AVX); \\\n ymm_lhs = \\\n _mm256_srli_epi32(_mm256_sub_epi8(_mm256_max_epi8(ymm_lhs, ymm_rhs), \\\n _mm256_min_epi8(ymm_lhs, ymm_rhs)), \\\n 4); \\\n ymm_sum = _mm256_add_epi32( \\\n _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_lhs, ymm_lhs), \\\n ONES_INT16_AVX), \\\n ymm_sum); \\\n ymm_lhs = _mm256_and_si256((ymm_m), MASK_INT4_AVX); \\\n ymm_rhs = _mm256_and_si256((ymm_q), MASK_INT4_AVX); \\\n ymm_lhs = \\\n _mm256_srli_epi32(_mm256_sub_epi8(_mm256_max_epi8(ymm_lhs, ymm_rhs), \\\n _mm256_min_epi8(ymm_lhs, ymm_rhs)), \\\n 4); \\\n ymm_sum = _mm256_add_epi32( \\\n _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_lhs, ymm_lhs), \\\n ONES_INT16_AVX), \\\n ymm_sum); \\\n }\n\n//! Compute the distance between matrix and query\n#define SSD_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum) \\\n { \\\n __m256i ymm_lhs_0 = \\\n _mm256_and_si256(_mm256_slli_epi32((ymm_lhs), 4), MASK_INT4_AVX); \\\n __m256i ymm_rhs_0 = \\\n _mm256_and_si256(_mm256_slli_epi32((ymm_rhs), 4), MASK_INT4_AVX); \\\n __m256i ymm_lhs_1 = _mm256_and_si256((ymm_lhs), MASK_INT4_AVX); \\\n __m256i ymm_rhs_1 = _mm256_and_si256((ymm_rhs), MASK_INT4_AVX); \\\n ymm_lhs_0 = _mm256_srli_epi32( \\\n _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_0, ymm_rhs_0), \\\n _mm256_min_epi8(ymm_lhs_0, ymm_rhs_0)), \\\n 4); \\\n ymm_rhs_0 = _mm256_srli_epi32( \\\n _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_1, ymm_rhs_1), \\\n _mm256_min_epi8(ymm_lhs_1, ymm_rhs_1)), \\\n 4); \\\n ymm_lhs_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_lhs_0, ymm_lhs_0), \\\n ONES_INT16_AVX); \\\n ymm_rhs_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_0, ymm_rhs_0), \\\n ONES_INT16_AVX); \\\n ymm_sum = \\\n _mm256_add_epi32(_mm256_add_epi32(ymm_lhs_0, ymm_rhs_0), ymm_sum); \\\n }\n\n//! Calculate sum of squared difference (GENERAL)\n#define SSD_INT8_GENERAL(m, q, sum) \\\n { \\\n int32_t x = m - q; \\\n sum += static_cast(x * x); \\\n }\n\n//! Calculate sum of squared difference (SSE)\n#define SSD_INT8_SSE(xmm_m, xmm_q, xmm_sum) \\\n { \\\n xmm_sum = _mm_add_epi32( \\\n _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_m), \\\n _mm_sign_epi8(xmm_m, xmm_m)), \\\n ONES_INT16_SSE), \\\n xmm_sum); \\\n xmm_sum = _mm_add_epi32( \\\n _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_q), \\\n _mm_sign_epi8(xmm_q, xmm_q)), \\\n ONES_INT16_SSE), \\\n xmm_sum); \\\n xmm_sum = _mm_sub_epi32( \\\n xmm_sum, \\\n _mm_slli_epi32( \\\n _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_q), \\\n _mm_sign_epi8(xmm_m, xmm_q)), \\\n ONES_INT16_SSE), \\\n 1)); \\\n }\n\n//! Calculate sum of squared difference (AVX)\n#define SSD_INT8_AVX(ymm_m, ymm_q, ymm_sum) \\\n { \\\n ymm_sum = _mm256_add_epi32( \\\n _mm256_madd_epi16( \\\n _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_m), \\\n _mm256_sign_epi8(ymm_m, ymm_m)), \\\n ONES_INT16_AVX), \\\n ymm_sum); \\\n ymm_sum = _mm256_add_epi32( \\\n _mm256_madd_epi16( \\\n _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_q), \\\n _mm256_sign_epi8(ymm_q, ymm_q)), \\\n ONES_INT16_AVX), \\\n ymm_sum); \\\n ymm_sum = _mm256_sub_epi32( \\\n ymm_sum, _mm256_slli_epi32( \\\n _mm256_madd_epi16( \\\n _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_q), \\\n _mm256_sign_epi8(ymm_m, ymm_q)), \\\n ONES_INT16_AVX), \\\n 1)); \\\n }\n\n//! Compute the square root of value (AVX)\n#define SQRT_FP32_AVX(v, ...) _mm256_sqrt_ps(_mm256_cvtepi32_ps(v))\n\n//! Compute the square root of value (AVX512)\n#define SQRT_FP32_AVX512(v, ...) _mm512_sqrt_ps(_mm512_cvtepi32_ps(v))\n\n#define ACCUM_FP32_STEP_SSE SSD_FP32_SSE\n#define ACCUM_FP32_STEP_AVX SSD_FP32_AVX\n\n#define ACCUM_FP32_STEP_AVX512 SSD_FP32_AVX512\n#define ACCUM_FP16_STEP_GENERAL SSD_FP16_GENERAL\n\n#define ACCUM_FP16_STEP_NEON SSD_FP16_NEON\n#define ACCUM_FP32_STEP_NEON SSD_FP32_NEON\n\n#define ACCUM_INT4_STEP_SSE SSD_INT4_SSE\n#define ACCUM_INT4_STEP_AVX SSD_INT4_AVX\n#define ACCUM_INT8_STEP_SSE SSD_INT8_SSE\n#define ACCUM_INT8_STEP_AVX SSD_INT8_AVX" }, { "path": "src/ailego/math/distance_matrix_fp16.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"matrix_define.i\"\n#include \n#if !defined(__AVX__)\n#define _mm_broadcast_si32(a) _mm_castps_si128(_mm_load1_ps((const float *)(a)))\n#else\n#define _mm_broadcast_si32(a) \\\n _mm_castps_si128(_mm_broadcast_ss((const float *)(a)))\n#define _mm256_broadcast_si32(a) \\\n _mm256_castps_si256(_mm256_broadcast_ss((const float *)(a)))\n#endif // !__AVX__\n\n//! Mask process of computing distance (FP16)\n#define MATRIX_FP16_MASK_AVX(lhs, rhs, cnt, _MASK, _RES, _PROC) \\\n switch (cnt) { \\\n case 7: { \\\n __m256 ymm_lhs = _mm256_cvtph_ps(_mm_set_epi16( \\\n (short)(_MASK), *((const short *)(lhs) + 6), \\\n *((const short *)(lhs) + 5), *((const short *)(lhs) + 4), \\\n *((const short *)(lhs) + 3), *((const short *)(lhs) + 2), \\\n *((const short *)(lhs) + 1), *((const short *)(lhs)))); \\\n __m256 ymm_rhs = _mm256_cvtph_ps(_mm_set_epi16( \\\n (short)(_MASK), *((const short *)(rhs) + 6), \\\n *((const short *)(rhs) + 5), *((const short *)(rhs) + 4), \\\n *((const short *)(rhs) + 3), *((const short *)(rhs) + 2), \\\n *((const short *)(rhs) + 1), *((const short *)(rhs)))); \\\n _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \\\n break; \\\n } \\\n case 6: { \\\n __m256 ymm_lhs = _mm256_cvtph_ps( \\\n _mm_set_epi32((int)(_MASK), *((const int *)(lhs) + 2), \\\n *((const int *)(lhs) + 1), *((const int *)(lhs)))); \\\n __m256 ymm_rhs = _mm256_cvtph_ps( \\\n _mm_set_epi32((int)(_MASK), *((const int *)(rhs) + 2), \\\n *((const int *)(rhs) + 1), *((const int *)(rhs)))); \\\n _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \\\n break; \\\n } \\\n case 5: { \\\n __m256 ymm_lhs = _mm256_cvtph_ps(_mm_set_epi16( \\\n (short)(_MASK), (short)(_MASK), (short)(_MASK), \\\n *((const short *)(lhs) + 4), *((const short *)(lhs) + 3), \\\n *((const short *)(lhs) + 2), *((const short *)(lhs) + 1), \\\n *((const short *)(lhs)))); \\\n __m256 ymm_rhs = _mm256_cvtph_ps(_mm_set_epi16( \\\n (short)(_MASK), (short)(_MASK), (short)(_MASK), \\\n *((const short *)(rhs) + 4), *((const short *)(rhs) + 3), \\\n *((const short *)(rhs) + 2), *((const short *)(rhs) + 1), \\\n *((const short *)(rhs)))); \\\n _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \\\n break; \\\n } \\\n case 4: { \\\n __m256 ymm_lhs = _mm256_cvtph_ps( \\\n _mm_set_epi64((__m64)(_MASK), *((const __m64 *)(lhs)))); \\\n __m256 ymm_rhs = _mm256_cvtph_ps( \\\n _mm_set_epi64((__m64)(_MASK), *((const __m64 *)(rhs)))); \\\n _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \\\n break; \\\n } \\\n case 3: { \\\n __m256 ymm_lhs = _mm256_cvtph_ps(_mm_set_epi16( \\\n (short)(_MASK), (short)(_MASK), (short)(_MASK), (short)(_MASK), \\\n (short)(_MASK), *((const short *)(lhs) + 2), \\\n *((const short *)(lhs) + 1), *((const short *)(lhs)))); \\\n __m256 ymm_rhs = _mm256_cvtph_ps(_mm_set_epi16( \\\n (short)(_MASK), (short)(_MASK), (short)(_MASK), (short)(_MASK), \\\n (short)(_MASK), *((const short *)(rhs) + 2), \\\n *((const short *)(rhs) + 1), *((const short *)(rhs)))); \\\n _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \\\n break; \\\n } \\\n case 2: { \\\n __m256 ymm_lhs = _mm256_cvtph_ps(_mm_set_epi32( \\\n (int)(_MASK), (int)(_MASK), (int)(_MASK), *((const int *)(lhs)))); \\\n __m256 ymm_rhs = _mm256_cvtph_ps(_mm_set_epi32( \\\n (int)(_MASK), (int)(_MASK), (int)(_MASK), *((const int *)(rhs)))); \\\n _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \\\n break; \\\n } \\\n case 1: { \\\n __m256 ymm_lhs = _mm256_cvtph_ps( \\\n _mm_set_epi16(*((const short *)(lhs)), (short)(_MASK), \\\n (short)(_MASK), (short)(_MASK), (short)(_MASK), \\\n (short)(_MASK), (short)(_MASK), (short)(_MASK))); \\\n __m256 ymm_rhs = _mm256_cvtph_ps( \\\n _mm_set_epi16(*((const short *)(rhs)), (short)(_MASK), \\\n (short)(_MASK), (short)(_MASK), (short)(_MASK), \\\n (short)(_MASK), (short)(_MASK), (short)(_MASK))); \\\n _PROC(ymm_lhs, ymm_rhs, _RES##_0_0) \\\n break; \\\n } \\\n }\n\n//! Iterative process of computing distance (FP16, M=1, N=1)\n#define MATRIX_FP16_ITER_1X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)m); \\\n __m256i ymm_qi = _LOAD((const __m256i *)q); \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_q = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_qi)); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0); \\\n ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n ymm_q = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_qi, 1)); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0); \\\n }\n\n//! Iterative process of computing distance (FP16, M=2, N=1)\n#define MATRIX_FP16_ITER_2X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \\\n __m256 ymm_q = _mm256_cvtph_ps(_mm_shufflehi_epi16( \\\n _mm_shufflelo_epi16(_mm_set1_epi64(*(const __m64 *)(q)), \\\n _MM_SHUFFLE(1, 1, 0, 0)), \\\n _MM_SHUFFLE(3, 3, 2, 2))); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (FP16, M=2, N=2)\n#define MATRIX_FP16_ITER_2X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_q = _mm256_cvtph_ps(_LOAD((const __m128i *)(q))); \\\n __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \\\n __m256 ymm_p = _mm256_moveldup_ps(ymm_q); \\\n _PROC(ymm_m, ymm_p, _RES##_0_0) \\\n ymm_p = _mm256_movehdup_ps(ymm_q); \\\n _PROC(ymm_m, ymm_p, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP16, M=4, N=1)\n#define MATRIX_FP16_ITER_4X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \\\n __m256 ymm_q = _mm256_cvtph_ps( \\\n _mm_shufflehi_epi16(_mm_shufflelo_epi16(_mm_broadcast_si32(q), 0), \\\n _MM_SHUFFLE(1, 1, 1, 1))); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (FP16, M=4, N=2)\n#define MATRIX_FP16_ITER_4X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_qi = _mm_set1_epi64(*(const __m64 *)(q)); \\\n __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \\\n __m256 ymm_q_0 = _mm256_cvtph_ps(_mm_shufflehi_epi16( \\\n _mm_shufflelo_epi16(xmm_qi, _MM_SHUFFLE(0, 0, 0, 0)), \\\n _MM_SHUFFLE(2, 2, 2, 2))); \\\n __m256 ymm_q_1 = _mm256_cvtph_ps(_mm_shufflehi_epi16( \\\n _mm_shufflelo_epi16(xmm_qi, _MM_SHUFFLE(1, 1, 1, 1)), \\\n _MM_SHUFFLE(3, 3, 3, 3))); \\\n MATRIX_VAR_PROC(1, 2, 0, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=4, N=4)\n#define MATRIX_FP16_ITER_4X4_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \\\n __m256 ymm_q = _mm256_cvtph_ps(_LOAD((const __m128i *)(q))); \\\n __m256 ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(0, 0, 0, 0)); \\\n _PROC(ymm_m, ymm_p, _RES##_0_0) \\\n ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(1, 1, 1, 1)); \\\n _PROC(ymm_m, ymm_p, _RES##_0_1) \\\n ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(2, 2, 2, 2)); \\\n _PROC(ymm_m, ymm_p, _RES##_0_2) \\\n ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(3, 3, 3, 3)); \\\n _PROC(ymm_m, ymm_p, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (FP16, M=8, N=1)\n#define MATRIX_FP16_ITER_8X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \\\n __m256 ymm_q = _mm256_cvtph_ps(_mm_set1_epi16(*(const short *)(q))); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (FP16, M=8, N=2)\n#define MATRIX_FP16_ITER_8X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \\\n __m256 ymm_q_0 = _mm256_set1_ps(xmm_p[0]); \\\n __m256 ymm_q_1 = _mm256_set1_ps(xmm_p[1]); \\\n MATRIX_VAR_PROC(1, 2, 0, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=8, N=4)\n#define MATRIX_FP16_ITER_8X4_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \\\n __m256 ymm_q = _mm256_set1_ps(xmm_p[0]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n ymm_q = _mm256_set1_ps(xmm_p[1]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_1) \\\n ymm_q = _mm256_set1_ps(xmm_p[2]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_2) \\\n ymm_q = _mm256_set1_ps(xmm_p[3]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (FP16, M=8, N=8)\n#define MATRIX_FP16_ITER_8X8_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _mm256_cvtph_ps(_LOAD((const __m128i *)(m))); \\\n __m256 ymm_p = _mm256_cvtph_ps(_LOAD((const __m128i *)(q))); \\\n __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_1) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_2) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_3) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_4) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_5) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_6) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_7) \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=1)\n#define MATRIX_FP16_ITER_16X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_q = _mm256_cvtph_ps(_mm_set1_epi16(*(const short *)q)); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=2)\n#define MATRIX_FP16_ITER_16X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \\\n __m256 ymm_q_0 = _mm256_set1_ps(xmm_p[0]); \\\n __m256 ymm_q_1 = _mm256_set1_ps(xmm_p[1]); \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n MATRIX_VAR_PROC(1, 2, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n MATRIX_VAR_PROC(1, 2, 1, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=4)\n#define MATRIX_FP16_ITER_16X4_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \\\n __m256 ymm_q = _mm256_set1_ps(xmm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(xmm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(xmm_p[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(xmm_p[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=8)\n#define MATRIX_FP16_ITER_16X8_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \\\n __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(2, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(2, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(2, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(2, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=16)\n#define MATRIX_FP16_ITER_16X16_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \\\n __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(2, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(2, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(2, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(2, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 8))); \\\n ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 8, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 9, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(2, 1, 10, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(2, 1, 11, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(2, 1, 12, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(2, 1, 13, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(2, 1, 14, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(2, 1, 15, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=1)\n#define MATRIX_FP16_ITER_32X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n ymm_mi = _LOAD((const __m256i *)(m + 16)); \\\n __m256 ymm_m_2 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_3 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n __m256 ymm_q = _mm256_cvtph_ps(_mm_set1_epi16(*(const short *)q)); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=2)\n#define MATRIX_FP16_ITER_32X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \\\n __m256 ymm_q_0 = _mm256_set1_ps(xmm_p[0]); \\\n __m256 ymm_q_1 = _mm256_set1_ps(xmm_p[1]); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n MATRIX_VAR_PROC(1, 2, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n MATRIX_VAR_PROC(1, 2, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(m + 16)); \\\n ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n MATRIX_VAR_PROC(1, 2, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n MATRIX_VAR_PROC(1, 2, 3, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=4)\n#define MATRIX_FP16_ITER_32X4_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \\\n __m256 ymm_q_0 = _mm256_set1_ps(xmm_p[0]); \\\n __m256 ymm_q_1 = _mm256_set1_ps(xmm_p[1]); \\\n __m256 ymm_q_2 = _mm256_set1_ps(xmm_p[2]); \\\n __m256 ymm_q_3 = _mm256_set1_ps(xmm_p[3]); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n MATRIX_VAR_PROC(1, 4, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n MATRIX_VAR_PROC(1, 4, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(m + 16)); \\\n ymm_m = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n MATRIX_VAR_PROC(1, 4, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_m = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n MATRIX_VAR_PROC(1, 4, 3, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=8)\n#define MATRIX_FP16_ITER_32X8_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n ymm_mi = _LOAD((const __m256i *)(m + 16)); \\\n __m256 ymm_m_2 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_3 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \\\n __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=16)\n#define MATRIX_FP16_ITER_32X16_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n ymm_mi = _LOAD((const __m256i *)(m + 16)); \\\n __m256 ymm_m_2 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_3 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \\\n __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 8))); \\\n ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(4, 1, 8, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(4, 1, 9, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(4, 1, 10, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(4, 1, 11, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(4, 1, 12, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(4, 1, 13, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(4, 1, 14, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(4, 1, 15, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=32)\n#define MATRIX_FP16_ITER_32X32_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(m)); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n ymm_mi = _LOAD((const __m256i *)(m + 16)); \\\n __m256 ymm_m_2 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_3 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \\\n __m256 ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 8))); \\\n ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(4, 1, 8, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(4, 1, 9, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(4, 1, 10, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(4, 1, 11, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(4, 1, 12, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(4, 1, 13, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(4, 1, 14, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(4, 1, 15, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 16))); \\\n ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(4, 1, 16, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(4, 1, 17, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(4, 1, 18, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(4, 1, 19, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(4, 1, 20, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(4, 1, 21, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(4, 1, 22, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(4, 1, 23, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q + 24))); \\\n ymm_q = _mm256_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(4, 1, 24, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(4, 1, 25, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(4, 1, 26, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(4, 1, 27, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(4, 1, 28, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(4, 1, 29, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(4, 1, 30, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(4, 1, 31, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=1, N=1)\n#define MATRIX_FP16_ITER_1X1_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512i zmm_mi = _LOAD((const __m512i *)m); \\\n __m512i zmm_qi = _LOAD((const __m512i *)q); \\\n __m512 zmm_m = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \\\n __m512 zmm_q = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_qi)); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0); \\\n zmm_m = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \\\n zmm_q = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_qi, 1)); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0); \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=1)\n#define MATRIX_FP16_ITER_16X1_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \\\n __m512 zmm_q = _mm512_cvtph_ps(_mm256_set1_epi16(*(const short *)q)); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=2)\n#define MATRIX_FP16_ITER_16X2_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \\\n __m512 zmm_q_0 = _mm512_set1_ps(xmm_p[0]); \\\n __m512 zmm_q_1 = _mm512_set1_ps(xmm_p[1]); \\\n MATRIX_VAR_PROC(1, 2, 0, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=4)\n#define MATRIX_FP16_ITER_16X4_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \\\n __m512 zmm_q = _mm512_set1_ps(xmm_p[0]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n zmm_q = _mm512_set1_ps(xmm_p[1]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_1) \\\n zmm_q = _mm512_set1_ps(xmm_p[2]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_2) \\\n zmm_q = _mm512_set1_ps(xmm_p[3]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=8)\n#define MATRIX_FP16_ITER_16X8_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \\\n __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \\\n __m512 zmm_q = _mm512_set1_ps(ymm_p[0]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n zmm_q = _mm512_set1_ps(ymm_p[1]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_1) \\\n zmm_q = _mm512_set1_ps(ymm_p[2]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_2) \\\n zmm_q = _mm512_set1_ps(ymm_p[3]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_3) \\\n zmm_q = _mm512_set1_ps(ymm_p[4]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_4) \\\n zmm_q = _mm512_set1_ps(ymm_p[5]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_5) \\\n zmm_q = _mm512_set1_ps(ymm_p[6]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_6) \\\n zmm_q = _mm512_set1_ps(ymm_p[7]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_7) \\\n }\n\n//! Iterative process of computing distance (FP16, M=16, N=16)\n#define MATRIX_FP16_ITER_16X16_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _mm512_cvtph_ps(_LOAD((const __m256i *)(m))); \\\n __m512 zmm_p = _mm512_cvtph_ps(_LOAD((const __m256i *)(q))); \\\n __m512 zmm_q = _mm512_set1_ps(zmm_p[0]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n zmm_q = _mm512_set1_ps(zmm_p[1]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_1) \\\n zmm_q = _mm512_set1_ps(zmm_p[2]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_2) \\\n zmm_q = _mm512_set1_ps(zmm_p[3]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_3) \\\n zmm_q = _mm512_set1_ps(zmm_p[4]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_4) \\\n zmm_q = _mm512_set1_ps(zmm_p[5]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_5) \\\n zmm_q = _mm512_set1_ps(zmm_p[6]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_6) \\\n zmm_q = _mm512_set1_ps(zmm_p[7]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_7) \\\n zmm_q = _mm512_set1_ps(zmm_p[8]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_8) \\\n zmm_q = _mm512_set1_ps(zmm_p[9]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_9) \\\n zmm_q = _mm512_set1_ps(zmm_p[10]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_10) \\\n zmm_q = _mm512_set1_ps(zmm_p[11]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_11) \\\n zmm_q = _mm512_set1_ps(zmm_p[12]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_12) \\\n zmm_q = _mm512_set1_ps(zmm_p[13]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_13) \\\n zmm_q = _mm512_set1_ps(zmm_p[14]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_14) \\\n zmm_q = _mm512_set1_ps(zmm_p[15]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_15) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=1)\n#define MATRIX_FP16_ITER_32X1_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512i zmm_mi = _LOAD((const __m512i *)(m)); \\\n __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \\\n __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \\\n __m512 zmm_q = _mm512_cvtph_ps(_mm256_set1_epi16(*(const short *)q)); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=2)\n#define MATRIX_FP16_ITER_32X2_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512i zmm_mi = _LOAD((const __m512i *)(m)); \\\n __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \\\n __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_broadcast_si32(q)); \\\n __m512 zmm_q = _mm512_set1_ps(xmm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(xmm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=4)\n#define MATRIX_FP16_ITER_32X4_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512i zmm_mi = _LOAD((const __m512i *)(m)); \\\n __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \\\n __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \\\n __m128 xmm_p = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)(q))); \\\n __m512 zmm_q = _mm512_set1_ps(xmm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(xmm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(xmm_p[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(xmm_p[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=8)\n#define MATRIX_FP16_ITER_32X8_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512i zmm_mi = _LOAD((const __m512i *)(m)); \\\n __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \\\n __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \\\n __m256 ymm_p = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(q))); \\\n __m512 zmm_q = _mm512_set1_ps(ymm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(ymm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(ymm_p[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(ymm_p[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(ymm_p[4]); \\\n MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(ymm_p[5]); \\\n MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(ymm_p[6]); \\\n MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(ymm_p[7]); \\\n MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=16)\n#define MATRIX_FP16_ITER_32X16_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512i zmm_mi = _LOAD((const __m512i *)(m)); \\\n __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \\\n __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \\\n __m512 zmm_p = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(q))); \\\n __m512 zmm_q = _mm512_set1_ps(zmm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[4]); \\\n MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[5]); \\\n MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[6]); \\\n MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[7]); \\\n MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[8]); \\\n MATRIX_VAR_PROC(2, 1, 8, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[9]); \\\n MATRIX_VAR_PROC(2, 1, 9, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[10]); \\\n MATRIX_VAR_PROC(2, 1, 10, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[11]); \\\n MATRIX_VAR_PROC(2, 1, 11, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[12]); \\\n MATRIX_VAR_PROC(2, 1, 12, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[13]); \\\n MATRIX_VAR_PROC(2, 1, 13, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[14]); \\\n MATRIX_VAR_PROC(2, 1, 14, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[15]); \\\n MATRIX_VAR_PROC(2, 1, 15, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP16, M=32, N=32)\n#define MATRIX_FP16_ITER_32X32_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512i zmm_mi = _LOAD((const __m512i *)(m)); \\\n __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \\\n __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \\\n __m512 zmm_p = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(q))); \\\n __m512 zmm_q = _mm512_set1_ps(zmm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[4]); \\\n MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[5]); \\\n MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[6]); \\\n MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[7]); \\\n MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[8]); \\\n MATRIX_VAR_PROC(2, 1, 8, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[9]); \\\n MATRIX_VAR_PROC(2, 1, 9, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[10]); \\\n MATRIX_VAR_PROC(2, 1, 10, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[11]); \\\n MATRIX_VAR_PROC(2, 1, 11, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[12]); \\\n MATRIX_VAR_PROC(2, 1, 12, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[13]); \\\n MATRIX_VAR_PROC(2, 1, 13, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[14]); \\\n MATRIX_VAR_PROC(2, 1, 14, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[15]); \\\n MATRIX_VAR_PROC(2, 1, 15, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_p = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(q + 16))); \\\n zmm_q = _mm512_set1_ps(zmm_p[0]); \\\n MATRIX_VAR_PROC(2, 1, 16, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[1]); \\\n MATRIX_VAR_PROC(2, 1, 17, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[2]); \\\n MATRIX_VAR_PROC(2, 1, 18, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[3]); \\\n MATRIX_VAR_PROC(2, 1, 19, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[4]); \\\n MATRIX_VAR_PROC(2, 1, 20, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[5]); \\\n MATRIX_VAR_PROC(2, 1, 21, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[6]); \\\n MATRIX_VAR_PROC(2, 1, 22, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[7]); \\\n MATRIX_VAR_PROC(2, 1, 23, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[8]); \\\n MATRIX_VAR_PROC(2, 1, 24, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[9]); \\\n MATRIX_VAR_PROC(2, 1, 25, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[10]); \\\n MATRIX_VAR_PROC(2, 1, 26, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[11]); \\\n MATRIX_VAR_PROC(2, 1, 27, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[12]); \\\n MATRIX_VAR_PROC(2, 1, 28, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[13]); \\\n MATRIX_VAR_PROC(2, 1, 29, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[14]); \\\n MATRIX_VAR_PROC(2, 1, 30, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(zmm_p[15]); \\\n MATRIX_VAR_PROC(2, 1, 31, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)\n//! Iterative process of computing distance (FP16, M=1, N=1)\n#define MATRIX_FP16_ITER_1X1_NEON(m, q, _RES, _PROC) \\\n { \\\n float16x8_t v_m = vld1q_f16((const float16_t *)m); \\\n float16x8_t v_q = vld1q_f16((const float16_t *)q); \\\n _PROC(v_m, v_q, _RES##_0_0) \\\n }\n\n#else\n//! Iterative process of computing distance (FP16, M=1, N=1)\n#define MATRIX_FP16_ITER_1X1_NEON(m, q, _RES, _PROC) \\\n { \\\n float16x8_t v_m = vld1q_f16((const float16_t *)m); \\\n float16x8_t v_q = vld1q_f16((const float16_t *)q); \\\n float32x4_t v_m_0 = vcvt_f32_f16(vget_low_f16(v_m)); \\\n float32x4_t v_q_0 = vcvt_f32_f16(vget_low_f16(v_q)); \\\n _PROC(v_m_0, v_q_0, _RES##_0_0) \\\n v_m_0 = vcvt_high_f32_f16(v_m); \\\n v_q_0 = vcvt_high_f32_f16(v_q); \\\n _PROC(v_m_0, v_q_0, _RES##_0_0) \\\n }\n\n#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC" }, { "path": "src/ailego/math/distance_matrix_fp32.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"matrix_define.i\"\n\n#if !defined(__AVX__)\n#undef _mm_permute_ps\n#define _mm_permute_ps(a, b) _mm_shuffle_ps((a), (a), (b))\n#define _mm_broadcast_ss(a) _mm_load1_ps(a)\n#endif // !__AVX__\n\n#if defined(__AVX__) && defined(__GNUC__)\n#define _mm256_set_m128(a, b) \\\n _mm256_insertf128_ps(_mm256_castps128_ps256(b), (a), 1)\n#endif // __AVX__\n\n#if defined(__ARM_NEON) && !defined(__aarch64__)\n#define vdupq_laneq_f32(a, b) vdupq_n_f32(vgetq_lane_f32(a, b))\n#endif // __ARM_NEON && __aarch64__\n\n//! Iterative process of computing distance (FP32, M=2, N=1)\n#define MATRIX_FP32_ITER_2X1_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_q = _LOAD(q); \\\n __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(1, 1, 0, 0)); \\\n _PROC(xmm_m_0, xmm_p, _RES##_0_0) \\\n xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 2, 2)); \\\n _PROC(xmm_m_1, xmm_p, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=2, N=2)\n#define MATRIX_FP32_ITER_2X2_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_q = _LOAD(q); \\\n __m128 xmm_m = _LOAD(m); \\\n __m128 xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(2, 2, 0, 0)); \\\n _PROC(xmm_m, xmm_p, _RES##_0_0) \\\n xmm_p = _mm_permute_ps(xmm_q, _MM_SHUFFLE(3, 3, 1, 1)); \\\n _PROC(xmm_m, xmm_p, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=4, N=1)\n#define MATRIX_FP32_ITER_4X1_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_q = _mm_broadcast_ss(q + 0); \\\n _PROC(xmm_m_0, xmm_q, _RES##_0_0) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n _PROC(xmm_m_1, xmm_q, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=4, N=2)\n#define MATRIX_FP32_ITER_4X2_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m = _LOAD(m); \\\n __m128 xmm_q = _mm_broadcast_ss(q + 0); \\\n _PROC(xmm_m, xmm_q, _RES##_0_0) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n _PROC(xmm_m, xmm_q, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=4, N=4)\n#define MATRIX_FP32_ITER_4X4_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m = _LOAD(m); \\\n __m128 xmm_q = _mm_broadcast_ss(q + 0); \\\n _PROC(xmm_m, xmm_q, _RES##_0_0) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n _PROC(xmm_m, xmm_q, _RES##_0_1) \\\n xmm_q = _mm_broadcast_ss(q + 2); \\\n _PROC(xmm_m, xmm_q, _RES##_0_2) \\\n xmm_q = _mm_broadcast_ss(q + 3); \\\n _PROC(xmm_m, xmm_q, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=1)\n#define MATRIX_FP32_ITER_8X1_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n _PROC(xmm_m_0, xmm_q, _RES##_0_0) \\\n _PROC(xmm_m_1, xmm_q, _RES##_1_0) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=2)\n#define MATRIX_FP32_ITER_8X2_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_q = _mm_broadcast_ss(q + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=4)\n#define MATRIX_FP32_ITER_8X4_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_q = _mm_broadcast_ss(q + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=8)\n#define MATRIX_FP32_ITER_8X8_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n MATRIX_VAR_PROC(2, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(2, 1, 4, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(2, 1, 5, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(2, 1, 6, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(2, 1, 7, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=1)\n#define MATRIX_FP32_ITER_16X1_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=2)\n#define MATRIX_FP32_ITER_16X2_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n __m128 xmm_q = _mm_broadcast_ss(q + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=4)\n#define MATRIX_FP32_ITER_16X4_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n __m128 xmm_q = _mm_broadcast_ss(q + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=8)\n#define MATRIX_FP32_ITER_16X8_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=16)\n#define MATRIX_FP32_ITER_16X16_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 8); \\\n MATRIX_VAR_PROC(4, 1, 8, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 9); \\\n MATRIX_VAR_PROC(4, 1, 9, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 10); \\\n MATRIX_VAR_PROC(4, 1, 10, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 11); \\\n MATRIX_VAR_PROC(4, 1, 11, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 12); \\\n MATRIX_VAR_PROC(4, 1, 12, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 13); \\\n MATRIX_VAR_PROC(4, 1, 13, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 14); \\\n MATRIX_VAR_PROC(4, 1, 14, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 15); \\\n MATRIX_VAR_PROC(4, 1, 15, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=1)\n#define MATRIX_FP32_ITER_32X1_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n _PROC(xmm_m_0, xmm_q, _RES##_0_0) \\\n _PROC(xmm_m_1, xmm_q, _RES##_1_0) \\\n _PROC(xmm_m_2, xmm_q, _RES##_2_0) \\\n _PROC(xmm_m_3, xmm_q, _RES##_3_0) \\\n xmm_m_0 = _LOAD(m + 16); \\\n xmm_m_1 = _LOAD(m + 20); \\\n xmm_m_2 = _LOAD(m + 24); \\\n xmm_m_3 = _LOAD(m + 28); \\\n _PROC(xmm_m_0, xmm_q, _RES##_4_0) \\\n _PROC(xmm_m_1, xmm_q, _RES##_5_0) \\\n _PROC(xmm_m_2, xmm_q, _RES##_6_0) \\\n _PROC(xmm_m_3, xmm_q, _RES##_7_0) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=2)\n#define MATRIX_FP32_ITER_32X2_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_q_0 = _mm_broadcast_ss(q + 0); \\\n __m128 xmm_q_1 = _mm_broadcast_ss(q + 1); \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n MATRIX_VAR_PROC(1, 2, 0, xmm_m_0, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 1, xmm_m_1, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 2, xmm_m_2, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 3, xmm_m_3, xmm_q, _RES, _PROC) \\\n xmm_m_0 = _LOAD(m + 16); \\\n xmm_m_1 = _LOAD(m + 20); \\\n xmm_m_2 = _LOAD(m + 24); \\\n xmm_m_3 = _LOAD(m + 28); \\\n MATRIX_VAR_PROC(1, 2, 4, xmm_m_0, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 5, xmm_m_1, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 6, xmm_m_2, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 7, xmm_m_3, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=4)\n#define MATRIX_FP32_ITER_32X4_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_q_0 = _mm_broadcast_ss(q + 0); \\\n __m128 xmm_q_1 = _mm_broadcast_ss(q + 1); \\\n __m128 xmm_q_2 = _mm_broadcast_ss(q + 2); \\\n __m128 xmm_q_3 = _mm_broadcast_ss(q + 3); \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n MATRIX_VAR_PROC(1, 4, 0, xmm_m_0, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 1, xmm_m_1, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 2, xmm_m_2, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 3, xmm_m_3, xmm_q, _RES, _PROC) \\\n xmm_m_0 = _LOAD(m + 16); \\\n xmm_m_1 = _LOAD(m + 20); \\\n xmm_m_2 = _LOAD(m + 24); \\\n xmm_m_3 = _LOAD(m + 28); \\\n MATRIX_VAR_PROC(1, 4, 4, xmm_m_0, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 5, xmm_m_1, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 6, xmm_m_2, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 7, xmm_m_3, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=8)\n#define MATRIX_FP32_ITER_32X8_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_q_0 = _mm_broadcast_ss(q + 0); \\\n __m128 xmm_q_1 = _mm_broadcast_ss(q + 1); \\\n __m128 xmm_q_2 = _mm_broadcast_ss(q + 2); \\\n __m128 xmm_q_3 = _mm_broadcast_ss(q + 3); \\\n __m128 xmm_q_4 = _mm_broadcast_ss(q + 4); \\\n __m128 xmm_q_5 = _mm_broadcast_ss(q + 5); \\\n __m128 xmm_q_6 = _mm_broadcast_ss(q + 6); \\\n __m128 xmm_q_7 = _mm_broadcast_ss(q + 7); \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n MATRIX_VAR_PROC(1, 8, 0, xmm_m_0, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 1, xmm_m_1, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 2, xmm_m_2, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 3, xmm_m_3, xmm_q, _RES, _PROC) \\\n xmm_m_0 = _LOAD(m + 16); \\\n xmm_m_1 = _LOAD(m + 20); \\\n xmm_m_2 = _LOAD(m + 24); \\\n xmm_m_3 = _LOAD(m + 28); \\\n MATRIX_VAR_PROC(1, 8, 4, xmm_m_0, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 5, xmm_m_1, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 6, xmm_m_2, xmm_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 7, xmm_m_3, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=16)\n#define MATRIX_FP32_ITER_32X16_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n __m128 xmm_m_4 = _LOAD(m + 16); \\\n __m128 xmm_m_5 = _LOAD(m + 20); \\\n __m128 xmm_m_6 = _LOAD(m + 24); \\\n __m128 xmm_m_7 = _LOAD(m + 28); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n MATRIX_VAR_PROC(8, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(8, 1, 1, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(8, 1, 2, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(8, 1, 3, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(8, 1, 4, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(8, 1, 5, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(8, 1, 6, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(8, 1, 7, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 8); \\\n MATRIX_VAR_PROC(8, 1, 8, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 9); \\\n MATRIX_VAR_PROC(8, 1, 9, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 10); \\\n MATRIX_VAR_PROC(8, 1, 10, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 11); \\\n MATRIX_VAR_PROC(8, 1, 11, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 12); \\\n MATRIX_VAR_PROC(8, 1, 12, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 13); \\\n MATRIX_VAR_PROC(8, 1, 13, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 14); \\\n MATRIX_VAR_PROC(8, 1, 14, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 15); \\\n MATRIX_VAR_PROC(8, 1, 15, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=32)\n#define MATRIX_FP32_ITER_32X32_SSE(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m128 xmm_m_0 = _LOAD(m + 0); \\\n __m128 xmm_m_1 = _LOAD(m + 4); \\\n __m128 xmm_m_2 = _LOAD(m + 8); \\\n __m128 xmm_m_3 = _LOAD(m + 12); \\\n __m128 xmm_m_4 = _LOAD(m + 16); \\\n __m128 xmm_m_5 = _LOAD(m + 20); \\\n __m128 xmm_m_6 = _LOAD(m + 24); \\\n __m128 xmm_m_7 = _LOAD(m + 28); \\\n __m128 xmm_q = _mm_broadcast_ss(q); \\\n MATRIX_VAR_PROC(8, 1, 0, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(8, 1, 1, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(8, 1, 2, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(8, 1, 3, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(8, 1, 4, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(8, 1, 5, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(8, 1, 6, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(8, 1, 7, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 8); \\\n MATRIX_VAR_PROC(8, 1, 8, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 9); \\\n MATRIX_VAR_PROC(8, 1, 9, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 10); \\\n MATRIX_VAR_PROC(8, 1, 10, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 11); \\\n MATRIX_VAR_PROC(8, 1, 11, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 12); \\\n MATRIX_VAR_PROC(8, 1, 12, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 13); \\\n MATRIX_VAR_PROC(8, 1, 13, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 14); \\\n MATRIX_VAR_PROC(8, 1, 14, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 15); \\\n MATRIX_VAR_PROC(8, 1, 15, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 16); \\\n MATRIX_VAR_PROC(8, 1, 16, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 17); \\\n MATRIX_VAR_PROC(8, 1, 17, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 18); \\\n MATRIX_VAR_PROC(8, 1, 18, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 19); \\\n MATRIX_VAR_PROC(8, 1, 19, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 20); \\\n MATRIX_VAR_PROC(8, 1, 20, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 21); \\\n MATRIX_VAR_PROC(8, 1, 21, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 22); \\\n MATRIX_VAR_PROC(8, 1, 22, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 23); \\\n MATRIX_VAR_PROC(8, 1, 23, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 24); \\\n MATRIX_VAR_PROC(8, 1, 24, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 25); \\\n MATRIX_VAR_PROC(8, 1, 25, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 26); \\\n MATRIX_VAR_PROC(8, 1, 26, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 27); \\\n MATRIX_VAR_PROC(8, 1, 27, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 28); \\\n MATRIX_VAR_PROC(8, 1, 28, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 29); \\\n MATRIX_VAR_PROC(8, 1, 29, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 30); \\\n MATRIX_VAR_PROC(8, 1, 30, xmm_m, xmm_q, _RES, _PROC) \\\n xmm_q = _mm_broadcast_ss(q + 31); \\\n MATRIX_VAR_PROC(8, 1, 31, xmm_m, xmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=2, N=1)\n#define MATRIX_FP32_ITER_2X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _LOAD(m); \\\n __m256 ymm_q = \\\n _mm256_set_ps(q[3], q[3], q[2], q[2], q[1], q[1], q[0], q[0]); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (FP32, M=2, N=2)\n#define MATRIX_FP32_ITER_2X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_q = _LOAD(q); \\\n __m256 ymm_m = _LOAD(m); \\\n __m256 ymm_p = _mm256_moveldup_ps(ymm_q); \\\n _PROC(ymm_m, ymm_p, _RES##_0_0) \\\n ymm_p = _mm256_movehdup_ps(ymm_q); \\\n _PROC(ymm_m, ymm_p, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=4, N=1)\n#define MATRIX_FP32_ITER_4X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _LOAD(m); \\\n __m256 ymm_q = \\\n _mm256_set_m128(_mm_broadcast_ss(q + 1), _mm_broadcast_ss(q)); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (FP32, M=4, N=2)\n#define MATRIX_FP32_ITER_4X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _LOAD(m); \\\n __m256 ymm_q = \\\n _mm256_set_m128(_mm_broadcast_ss(q + 2), _mm_broadcast_ss(q + 0)); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n ymm_q = _mm256_set_m128(_mm_broadcast_ss(q + 3), _mm_broadcast_ss(q + 1)); \\\n _PROC(ymm_m, ymm_q, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=4, N=4)\n#define MATRIX_FP32_ITER_4X4_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_q = _LOAD(q); \\\n __m256 ymm_m = _LOAD(m); \\\n __m256 ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(0, 0, 0, 0)); \\\n _PROC(ymm_m, ymm_p, _RES##_0_0) \\\n ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(1, 1, 1, 1)); \\\n _PROC(ymm_m, ymm_p, _RES##_0_1) \\\n ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(2, 2, 2, 2)); \\\n _PROC(ymm_m, ymm_p, _RES##_0_2) \\\n ymm_p = _mm256_permute_ps(ymm_q, _MM_SHUFFLE(3, 3, 3, 3)); \\\n _PROC(ymm_m, ymm_p, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=1)\n#define MATRIX_FP32_ITER_8X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _LOAD(m); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=2)\n#define MATRIX_FP32_ITER_8X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _LOAD(m); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n _PROC(ymm_m, ymm_q, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=4)\n#define MATRIX_FP32_ITER_8X4_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _LOAD(m); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n _PROC(ymm_m, ymm_q, _RES##_0_1) \\\n ymm_q = _mm256_broadcast_ss(q + 2); \\\n _PROC(ymm_m, ymm_q, _RES##_0_2) \\\n ymm_q = _mm256_broadcast_ss(q + 3); \\\n _PROC(ymm_m, ymm_q, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=8)\n#define MATRIX_FP32_ITER_8X8_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m = _LOAD(m); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n _PROC(ymm_m, ymm_q, _RES##_0_0) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n _PROC(ymm_m, ymm_q, _RES##_0_1) \\\n ymm_q = _mm256_broadcast_ss(q + 2); \\\n _PROC(ymm_m, ymm_q, _RES##_0_2) \\\n ymm_q = _mm256_broadcast_ss(q + 3); \\\n _PROC(ymm_m, ymm_q, _RES##_0_3) \\\n ymm_q = _mm256_broadcast_ss(q + 4); \\\n _PROC(ymm_m, ymm_q, _RES##_0_4) \\\n ymm_q = _mm256_broadcast_ss(q + 5); \\\n _PROC(ymm_m, ymm_q, _RES##_0_5) \\\n ymm_q = _mm256_broadcast_ss(q + 6); \\\n _PROC(ymm_m, ymm_q, _RES##_0_6) \\\n ymm_q = _mm256_broadcast_ss(q + 7); \\\n _PROC(ymm_m, ymm_q, _RES##_0_7) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=1)\n#define MATRIX_FP32_ITER_16X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=2)\n#define MATRIX_FP32_ITER_16X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=4)\n#define MATRIX_FP32_ITER_16X4_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=8)\n#define MATRIX_FP32_ITER_16X8_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(2, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(2, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(2, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(2, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=16)\n#define MATRIX_FP32_ITER_16X16_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(2, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(2, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(2, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(2, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 8); \\\n MATRIX_VAR_PROC(2, 1, 8, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 9); \\\n MATRIX_VAR_PROC(2, 1, 9, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 10); \\\n MATRIX_VAR_PROC(2, 1, 10, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 11); \\\n MATRIX_VAR_PROC(2, 1, 11, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 12); \\\n MATRIX_VAR_PROC(2, 1, 12, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 13); \\\n MATRIX_VAR_PROC(2, 1, 13, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 14); \\\n MATRIX_VAR_PROC(2, 1, 14, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 15); \\\n MATRIX_VAR_PROC(2, 1, 15, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=1)\n#define MATRIX_FP32_ITER_32X1_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_m_2 = _LOAD(m + 16); \\\n __m256 ymm_m_3 = _LOAD(m + 24); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=2)\n#define MATRIX_FP32_ITER_32X2_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_m_2 = _LOAD(m + 16); \\\n __m256 ymm_m_3 = _LOAD(m + 24); \\\n __m256 ymm_q = _mm256_broadcast_ss(q + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=4)\n#define MATRIX_FP32_ITER_32X4_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_m_2 = _LOAD(m + 16); \\\n __m256 ymm_m_3 = _LOAD(m + 24); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=8)\n#define MATRIX_FP32_ITER_32X8_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_m_2 = _LOAD(m + 16); \\\n __m256 ymm_m_3 = _LOAD(m + 24); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=16)\n#define MATRIX_FP32_ITER_32X16_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_m_2 = _LOAD(m + 16); \\\n __m256 ymm_m_3 = _LOAD(m + 24); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 8); \\\n MATRIX_VAR_PROC(4, 1, 8, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 9); \\\n MATRIX_VAR_PROC(4, 1, 9, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 10); \\\n MATRIX_VAR_PROC(4, 1, 10, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 11); \\\n MATRIX_VAR_PROC(4, 1, 11, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 12); \\\n MATRIX_VAR_PROC(4, 1, 12, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 13); \\\n MATRIX_VAR_PROC(4, 1, 13, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 14); \\\n MATRIX_VAR_PROC(4, 1, 14, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 15); \\\n MATRIX_VAR_PROC(4, 1, 15, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=32)\n#define MATRIX_FP32_ITER_32X32_AVX(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m256 ymm_m_0 = _LOAD(m + 0); \\\n __m256 ymm_m_1 = _LOAD(m + 8); \\\n __m256 ymm_m_2 = _LOAD(m + 16); \\\n __m256 ymm_m_3 = _LOAD(m + 24); \\\n __m256 ymm_q = _mm256_broadcast_ss(q); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 8); \\\n MATRIX_VAR_PROC(4, 1, 8, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 9); \\\n MATRIX_VAR_PROC(4, 1, 9, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 10); \\\n MATRIX_VAR_PROC(4, 1, 10, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 11); \\\n MATRIX_VAR_PROC(4, 1, 11, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 12); \\\n MATRIX_VAR_PROC(4, 1, 12, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 13); \\\n MATRIX_VAR_PROC(4, 1, 13, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 14); \\\n MATRIX_VAR_PROC(4, 1, 14, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 15); \\\n MATRIX_VAR_PROC(4, 1, 15, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 16); \\\n MATRIX_VAR_PROC(4, 1, 16, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 17); \\\n MATRIX_VAR_PROC(4, 1, 17, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 18); \\\n MATRIX_VAR_PROC(4, 1, 18, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 19); \\\n MATRIX_VAR_PROC(4, 1, 19, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 20); \\\n MATRIX_VAR_PROC(4, 1, 20, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 21); \\\n MATRIX_VAR_PROC(4, 1, 21, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 22); \\\n MATRIX_VAR_PROC(4, 1, 22, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 23); \\\n MATRIX_VAR_PROC(4, 1, 23, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 24); \\\n MATRIX_VAR_PROC(4, 1, 24, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 25); \\\n MATRIX_VAR_PROC(4, 1, 25, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 26); \\\n MATRIX_VAR_PROC(4, 1, 26, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 27); \\\n MATRIX_VAR_PROC(4, 1, 27, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 28); \\\n MATRIX_VAR_PROC(4, 1, 28, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 29); \\\n MATRIX_VAR_PROC(4, 1, 29, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 30); \\\n MATRIX_VAR_PROC(4, 1, 30, ymm_m, ymm_q, _RES, _PROC) \\\n ymm_q = _mm256_broadcast_ss(q + 31); \\\n MATRIX_VAR_PROC(4, 1, 31, ymm_m, ymm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=1)\n#define MATRIX_FP32_ITER_16X1_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _LOAD(m); \\\n __m512 zmm_q = _mm512_set1_ps(*q); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=2)\n#define MATRIX_FP32_ITER_16X2_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _LOAD(m); \\\n __m512 zmm_q = _mm512_set1_ps(q[0]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n zmm_q = _mm512_set1_ps(q[1]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=4)\n#define MATRIX_FP32_ITER_16X4_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _LOAD(m); \\\n __m512 zmm_q = _mm512_set1_ps(q[0]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n zmm_q = _mm512_set1_ps(q[1]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_1) \\\n zmm_q = _mm512_set1_ps(q[2]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_2) \\\n zmm_q = _mm512_set1_ps(q[3]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=8)\n#define MATRIX_FP32_ITER_16X8_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _LOAD(m); \\\n __m512 zmm_q = _mm512_set1_ps(q[0]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n zmm_q = _mm512_set1_ps(q[1]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_1) \\\n zmm_q = _mm512_set1_ps(q[2]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_2) \\\n zmm_q = _mm512_set1_ps(q[3]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_3) \\\n zmm_q = _mm512_set1_ps(q[4]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_4) \\\n zmm_q = _mm512_set1_ps(q[5]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_5) \\\n zmm_q = _mm512_set1_ps(q[6]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_6) \\\n zmm_q = _mm512_set1_ps(q[7]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_7) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=16)\n#define MATRIX_FP32_ITER_16X16_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m = _LOAD(m); \\\n __m512 zmm_q = _mm512_set1_ps(q[0]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n zmm_q = _mm512_set1_ps(q[1]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_1) \\\n zmm_q = _mm512_set1_ps(q[2]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_2) \\\n zmm_q = _mm512_set1_ps(q[3]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_3) \\\n zmm_q = _mm512_set1_ps(q[4]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_4) \\\n zmm_q = _mm512_set1_ps(q[5]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_5) \\\n zmm_q = _mm512_set1_ps(q[6]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_6) \\\n zmm_q = _mm512_set1_ps(q[7]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_7) \\\n zmm_q = _mm512_set1_ps(q[8]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_8) \\\n zmm_q = _mm512_set1_ps(q[9]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_9) \\\n zmm_q = _mm512_set1_ps(q[10]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_10) \\\n zmm_q = _mm512_set1_ps(q[11]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_11) \\\n zmm_q = _mm512_set1_ps(q[12]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_12) \\\n zmm_q = _mm512_set1_ps(q[13]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_13) \\\n zmm_q = _mm512_set1_ps(q[14]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_14) \\\n zmm_q = _mm512_set1_ps(q[15]); \\\n _PROC(zmm_m, zmm_q, _RES##_0_15) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=1)\n#define MATRIX_FP32_ITER_32X1_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_q = _mm512_set1_ps(*q); \\\n __m512 zmm_m = _LOAD(m); \\\n _PROC(zmm_m, zmm_q, _RES##_0_0) \\\n zmm_m = _LOAD(m + 16); \\\n _PROC(zmm_m, zmm_q, _RES##_1_0) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=2)\n#define MATRIX_FP32_ITER_32X2_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m_0 = _LOAD(m + 0); \\\n __m512 zmm_m_1 = _LOAD(m + 16); \\\n __m512 zmm_q = _mm512_set1_ps(q[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=4)\n#define MATRIX_FP32_ITER_32X4_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m_0 = _LOAD(m + 0); \\\n __m512 zmm_m_1 = _LOAD(m + 16); \\\n __m512 zmm_q = _mm512_set1_ps(q[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=8)\n#define MATRIX_FP32_ITER_32X8_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m_0 = _LOAD(m + 0); \\\n __m512 zmm_m_1 = _LOAD(m + 16); \\\n __m512 zmm_q = _mm512_set1_ps(q[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[4]); \\\n MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[5]); \\\n MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[6]); \\\n MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[7]); \\\n MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=16)\n#define MATRIX_FP32_ITER_32X16_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m_0 = _LOAD(m + 0); \\\n __m512 zmm_m_1 = _LOAD(m + 16); \\\n __m512 zmm_q = _mm512_set1_ps(q[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[4]); \\\n MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[5]); \\\n MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[6]); \\\n MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[7]); \\\n MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[8]); \\\n MATRIX_VAR_PROC(2, 1, 8, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[9]); \\\n MATRIX_VAR_PROC(2, 1, 9, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[10]); \\\n MATRIX_VAR_PROC(2, 1, 10, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[11]); \\\n MATRIX_VAR_PROC(2, 1, 11, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[12]); \\\n MATRIX_VAR_PROC(2, 1, 12, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[13]); \\\n MATRIX_VAR_PROC(2, 1, 13, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[14]); \\\n MATRIX_VAR_PROC(2, 1, 14, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[15]); \\\n MATRIX_VAR_PROC(2, 1, 15, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=32)\n#define MATRIX_FP32_ITER_32X32_AVX512(m, q, _RES, _LOAD, _PROC) \\\n { \\\n __m512 zmm_m_0 = _LOAD(m + 0); \\\n __m512 zmm_m_1 = _LOAD(m + 16); \\\n __m512 zmm_q = _mm512_set1_ps(q[0]); \\\n MATRIX_VAR_PROC(2, 1, 0, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[1]); \\\n MATRIX_VAR_PROC(2, 1, 1, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[2]); \\\n MATRIX_VAR_PROC(2, 1, 2, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[3]); \\\n MATRIX_VAR_PROC(2, 1, 3, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[4]); \\\n MATRIX_VAR_PROC(2, 1, 4, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[5]); \\\n MATRIX_VAR_PROC(2, 1, 5, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[6]); \\\n MATRIX_VAR_PROC(2, 1, 6, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[7]); \\\n MATRIX_VAR_PROC(2, 1, 7, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[8]); \\\n MATRIX_VAR_PROC(2, 1, 8, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[9]); \\\n MATRIX_VAR_PROC(2, 1, 9, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[10]); \\\n MATRIX_VAR_PROC(2, 1, 10, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[11]); \\\n MATRIX_VAR_PROC(2, 1, 11, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[12]); \\\n MATRIX_VAR_PROC(2, 1, 12, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[13]); \\\n MATRIX_VAR_PROC(2, 1, 13, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[14]); \\\n MATRIX_VAR_PROC(2, 1, 14, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[15]); \\\n MATRIX_VAR_PROC(2, 1, 15, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[16]); \\\n MATRIX_VAR_PROC(2, 1, 16, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[17]); \\\n MATRIX_VAR_PROC(2, 1, 17, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[18]); \\\n MATRIX_VAR_PROC(2, 1, 18, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[19]); \\\n MATRIX_VAR_PROC(2, 1, 19, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[20]); \\\n MATRIX_VAR_PROC(2, 1, 20, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[21]); \\\n MATRIX_VAR_PROC(2, 1, 21, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[22]); \\\n MATRIX_VAR_PROC(2, 1, 22, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[23]); \\\n MATRIX_VAR_PROC(2, 1, 23, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[24]); \\\n MATRIX_VAR_PROC(2, 1, 24, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[25]); \\\n MATRIX_VAR_PROC(2, 1, 25, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[26]); \\\n MATRIX_VAR_PROC(2, 1, 26, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[27]); \\\n MATRIX_VAR_PROC(2, 1, 27, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[28]); \\\n MATRIX_VAR_PROC(2, 1, 28, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[29]); \\\n MATRIX_VAR_PROC(2, 1, 29, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[30]); \\\n MATRIX_VAR_PROC(2, 1, 30, zmm_m, zmm_q, _RES, _PROC) \\\n zmm_q = _mm512_set1_ps(q[31]); \\\n MATRIX_VAR_PROC(2, 1, 31, zmm_m, zmm_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=2, N=1)\n#define MATRIX_FP32_ITER_2X1_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m = vld1q_f32(m); \\\n float32x2_t v_q = vld1_f32(q); \\\n float32x4_t v_p = \\\n vcombine_f32(vdup_lane_f32(v_q, 0), vdup_lane_f32(v_q, 1)); \\\n _PROC(v_m, v_p, _RES) \\\n }\n\n//! Iterative process of computing distance (FP32, M=2, N=2)\n#define MATRIX_FP32_ITER_2X2_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_q = vld1q_f32(q); \\\n float32x4_t v_m = vld1q_f32(m); \\\n float32x2_t v_q_0 = vget_low_f32(v_q); \\\n float32x2_t v_q_1 = vget_high_f32(v_q); \\\n v_q = vcombine_f32(vdup_lane_f32(v_q_0, 0), vdup_lane_f32(v_q_1, 0)); \\\n _PROC(v_m, v_q, _RES##_0_0) \\\n v_q = vcombine_f32(vdup_lane_f32(v_q_0, 1), vdup_lane_f32(v_q_1, 1)); \\\n _PROC(v_m, v_q, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=4, N=1)\n#define MATRIX_FP32_ITER_4X1_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x2_t v_p = vld1_f32(q); \\\n float32x4_t v_q = vdupq_lane_f32(v_p, 0); \\\n _PROC(v_m_0, v_q, _RES##_0_0) \\\n v_q = vdupq_lane_f32(v_p, 1); \\\n _PROC(v_m_1, v_q, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=4, N=2)\n#define MATRIX_FP32_ITER_4X2_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m = vld1q_f32(m); \\\n float32x2_t v_p = vld1_f32(q); \\\n float32x4_t v_q = vdupq_lane_f32(v_p, 0); \\\n _PROC(v_m, v_q, _RES##_0_0) \\\n v_q = vdupq_lane_f32(v_p, 1); \\\n _PROC(v_m, v_q, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (FP32, M=4, N=4)\n#define MATRIX_FP32_ITER_4X4_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m = vld1q_f32(m); \\\n float32x4_t v_p = vld1q_f32(q); \\\n float32x4_t v_q = vdupq_laneq_f32(v_p, 0); \\\n _PROC(v_m, v_q, _RES##_0_0) \\\n v_q = vdupq_laneq_f32(v_p, 1); \\\n _PROC(v_m, v_q, _RES##_0_1) \\\n v_q = vdupq_laneq_f32(v_p, 2); \\\n _PROC(v_m, v_q, _RES##_0_2) \\\n v_q = vdupq_laneq_f32(v_p, 3); \\\n _PROC(v_m, v_q, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=1)\n#define MATRIX_FP32_ITER_8X1_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_q = vld1q_dup_f32(q); \\\n _PROC(v_m_0, v_q, _RES##_0_0) \\\n _PROC(v_m_1, v_q, _RES##_1_0) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=2)\n#define MATRIX_FP32_ITER_8X2_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x2_t v_p = vld1_f32(q); \\\n float32x4_t v_q = vdupq_lane_f32(v_p, 0); \\\n MATRIX_VAR_PROC(2, 1, 0, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_lane_f32(v_p, 1); \\\n MATRIX_VAR_PROC(2, 1, 1, v_m, v_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=4)\n#define MATRIX_FP32_ITER_8X4_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_p = vld1q_f32(q); \\\n float32x4_t v_q = vdupq_laneq_f32(v_p, 0); \\\n MATRIX_VAR_PROC(2, 1, 0, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_laneq_f32(v_p, 1); \\\n MATRIX_VAR_PROC(2, 1, 1, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_laneq_f32(v_p, 2); \\\n MATRIX_VAR_PROC(2, 1, 2, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_laneq_f32(v_p, 3); \\\n MATRIX_VAR_PROC(2, 1, 3, v_m, v_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=8, N=8)\n#define MATRIX_FP32_ITER_8X8_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_p = vld1q_f32(q + 0); \\\n float32x4_t v_q = vdupq_laneq_f32(v_p, 0); \\\n MATRIX_VAR_PROC(2, 1, 0, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_laneq_f32(v_p, 1); \\\n MATRIX_VAR_PROC(2, 1, 1, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_laneq_f32(v_p, 2); \\\n MATRIX_VAR_PROC(2, 1, 2, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_laneq_f32(v_p, 3); \\\n MATRIX_VAR_PROC(2, 1, 3, v_m, v_q, _RES, _PROC) \\\n v_p = vld1q_f32(q + 4); \\\n v_q = vdupq_laneq_f32(v_p, 0); \\\n MATRIX_VAR_PROC(2, 1, 4, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_laneq_f32(v_p, 1); \\\n MATRIX_VAR_PROC(2, 1, 5, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_laneq_f32(v_p, 2); \\\n MATRIX_VAR_PROC(2, 1, 6, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_laneq_f32(v_p, 3); \\\n MATRIX_VAR_PROC(2, 1, 7, v_m, v_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=1)\n#define MATRIX_FP32_ITER_16X1_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n float32x4_t v_q = vld1q_dup_f32(q); \\\n MATRIX_VAR_PROC(4, 1, 0, v_m, v_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=2)\n#define MATRIX_FP32_ITER_16X2_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n float32x2_t v_p = vld1_f32(q); \\\n float32x4_t v_q = vdupq_lane_f32(v_p, 0); \\\n MATRIX_VAR_PROC(4, 1, 0, v_m, v_q, _RES, _PROC) \\\n v_q = vdupq_lane_f32(v_p, 1); \\\n MATRIX_VAR_PROC(4, 1, 1, v_m, v_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=4)\n#define MATRIX_FP32_ITER_16X4_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n float32x4_t v_q = vld1q_f32(q); \\\n float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(4, 1, 0, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(4, 1, 1, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(4, 1, 2, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(4, 1, 3, v_m, v_p, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=8)\n#define MATRIX_FP32_ITER_16X8_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n float32x4_t v_q = vld1q_f32(q + 0); \\\n float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(4, 1, 0, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(4, 1, 1, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(4, 1, 2, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(4, 1, 3, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 4); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(4, 1, 4, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(4, 1, 5, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(4, 1, 6, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(4, 1, 7, v_m, v_p, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=16, N=16)\n#define MATRIX_FP32_ITER_16X16_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n float32x4_t v_q = vld1q_f32(q + 0); \\\n float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(4, 1, 0, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(4, 1, 1, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(4, 1, 2, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(4, 1, 3, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 4); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(4, 1, 4, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(4, 1, 5, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(4, 1, 6, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(4, 1, 7, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 8); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(4, 1, 8, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(4, 1, 9, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(4, 1, 10, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(4, 1, 11, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 12); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(4, 1, 12, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(4, 1, 13, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(4, 1, 14, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(4, 1, 15, v_m, v_p, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=1)\n#define MATRIX_FP32_ITER_32X1_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_q = vld1q_dup_f32(q); \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n _PROC(v_m_0, v_q, _RES##_0_0) \\\n _PROC(v_m_1, v_q, _RES##_1_0) \\\n _PROC(v_m_2, v_q, _RES##_2_0) \\\n _PROC(v_m_3, v_q, _RES##_3_0) \\\n v_m_0 = vld1q_f32(m + 16); \\\n v_m_1 = vld1q_f32(m + 20); \\\n v_m_2 = vld1q_f32(m + 24); \\\n v_m_3 = vld1q_f32(m + 28); \\\n _PROC(v_m_0, v_q, _RES##_4_0) \\\n _PROC(v_m_1, v_q, _RES##_5_0) \\\n _PROC(v_m_2, v_q, _RES##_6_0) \\\n _PROC(v_m_3, v_q, _RES##_7_0) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=2)\n#define MATRIX_FP32_ITER_32X2_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x2_t v_p = vld1_f32(q); \\\n float32x4_t v_q_0 = vdupq_lane_f32(v_p, 0); \\\n float32x4_t v_q_1 = vdupq_lane_f32(v_p, 1); \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n MATRIX_VAR_PROC(1, 2, 0, v_m_0, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 1, v_m_1, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 2, v_m_2, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 3, v_m_3, v_q, _RES, _PROC) \\\n v_m_0 = vld1q_f32(m + 16); \\\n v_m_1 = vld1q_f32(m + 20); \\\n v_m_2 = vld1q_f32(m + 24); \\\n v_m_3 = vld1q_f32(m + 28); \\\n MATRIX_VAR_PROC(1, 2, 4, v_m_0, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 5, v_m_1, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 6, v_m_2, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 7, v_m_3, v_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=4)\n#define MATRIX_FP32_ITER_32X4_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_p = vld1q_f32(q); \\\n float32x4_t v_q_0 = vdupq_laneq_f32(v_p, 0); \\\n float32x4_t v_q_1 = vdupq_laneq_f32(v_p, 1); \\\n float32x4_t v_q_2 = vdupq_laneq_f32(v_p, 2); \\\n float32x4_t v_q_3 = vdupq_laneq_f32(v_p, 3); \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n MATRIX_VAR_PROC(1, 4, 0, v_m_0, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 1, v_m_1, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 2, v_m_2, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 3, v_m_3, v_q, _RES, _PROC) \\\n v_m_0 = vld1q_f32(m + 16); \\\n v_m_1 = vld1q_f32(m + 20); \\\n v_m_2 = vld1q_f32(m + 24); \\\n v_m_3 = vld1q_f32(m + 28); \\\n MATRIX_VAR_PROC(1, 4, 4, v_m_0, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 5, v_m_1, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 6, v_m_2, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 7, v_m_3, v_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=8)\n#define MATRIX_FP32_ITER_32X8_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_p_0 = vld1q_f32(q + 0); \\\n float32x4_t v_p_1 = vld1q_f32(q + 4); \\\n float32x4_t v_q_0 = vdupq_laneq_f32(v_p_0, 0); \\\n float32x4_t v_q_1 = vdupq_laneq_f32(v_p_0, 1); \\\n float32x4_t v_q_2 = vdupq_laneq_f32(v_p_0, 2); \\\n float32x4_t v_q_3 = vdupq_laneq_f32(v_p_0, 3); \\\n float32x4_t v_q_4 = vdupq_laneq_f32(v_p_1, 0); \\\n float32x4_t v_q_5 = vdupq_laneq_f32(v_p_1, 1); \\\n float32x4_t v_q_6 = vdupq_laneq_f32(v_p_1, 2); \\\n float32x4_t v_q_7 = vdupq_laneq_f32(v_p_1, 3); \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n MATRIX_VAR_PROC(1, 8, 0, v_m_0, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 1, v_m_1, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 2, v_m_2, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 3, v_m_3, v_q, _RES, _PROC) \\\n v_m_0 = vld1q_f32(m + 16); \\\n v_m_1 = vld1q_f32(m + 20); \\\n v_m_2 = vld1q_f32(m + 24); \\\n v_m_3 = vld1q_f32(m + 28); \\\n MATRIX_VAR_PROC(1, 8, 4, v_m_0, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 5, v_m_1, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 6, v_m_2, v_q, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 7, v_m_3, v_q, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=16)\n#define MATRIX_FP32_ITER_32X16_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n float32x4_t v_m_4 = vld1q_f32(m + 16); \\\n float32x4_t v_m_5 = vld1q_f32(m + 20); \\\n float32x4_t v_m_6 = vld1q_f32(m + 24); \\\n float32x4_t v_m_7 = vld1q_f32(m + 28); \\\n float32x4_t v_q = vld1q_f32(q + 0); \\\n float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 0, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 1, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 2, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 3, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 4); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 4, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 5, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 6, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 7, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 8); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 8, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 9, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 10, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 11, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 12); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 12, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 13, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 14, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 15, v_m, v_p, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (FP32, M=32, N=32)\n#define MATRIX_FP32_ITER_32X32_NEON(m, q, _RES, _PROC) \\\n { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n float32x4_t v_m_2 = vld1q_f32(m + 8); \\\n float32x4_t v_m_3 = vld1q_f32(m + 12); \\\n float32x4_t v_m_4 = vld1q_f32(m + 16); \\\n float32x4_t v_m_5 = vld1q_f32(m + 20); \\\n float32x4_t v_m_6 = vld1q_f32(m + 24); \\\n float32x4_t v_m_7 = vld1q_f32(m + 28); \\\n float32x4_t v_q = vld1q_f32(q + 0); \\\n float32x4_t v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 0, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 1, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 2, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 3, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 4); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 4, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 5, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 6, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 7, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 8); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 8, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 9, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 10, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 11, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 12); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 12, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 13, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 14, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 15, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 16); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 16, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 17, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 18, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 19, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 20); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 20, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 21, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 22, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 23, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 24); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 24, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 25, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 26, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 27, v_m, v_p, _RES, _PROC) \\\n v_q = vld1q_f32(q + 28); \\\n v_p = vdupq_laneq_f32(v_q, 0); \\\n MATRIX_VAR_PROC(8, 1, 28, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 1); \\\n MATRIX_VAR_PROC(8, 1, 29, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 2); \\\n MATRIX_VAR_PROC(8, 1, 30, v_m, v_p, _RES, _PROC) \\\n v_p = vdupq_laneq_f32(v_q, 3); \\\n MATRIX_VAR_PROC(8, 1, 31, v_m, v_p, _RES, _PROC) \\\n }\n" }, { "path": "src/ailego/math/distance_matrix_inner_product_utility.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#if defined(__SSE4_1__)\n//! Four-bits Convert Table\nstatic const AILEGO_ALIGNED(32) int8_t Int4ConvertTable[32] = {\n 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1,\n 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1};\n\n#define NEGZEROS_FP32_SSE _mm_set1_ps(-0.0f)\n#define MASK_INT4_SSE _mm_set1_epi32(0x0f0f0f0f)\n#define ONES_INT16_SSE _mm_set1_epi32(0x00010001)\n#define INT4_LOOKUP_SSE _mm_load_si128((const __m128i *)Int4ConvertTable)\n#endif // __SSE4_1__\n\n#if defined(__AVX__)\n// #define NEGZEROS_FP32_AVX _mm256_set1_ps(-0.0f)\n#define MASK_INT4_AVX _mm256_set1_epi32(0x0f0f0f0f)\n#define ONES_INT16_AVX _mm256_set1_epi32(0x00010001)\n#define INT4_LOOKUP_AVX _mm256_load_si256((const __m256i *)Int4ConvertTable)\n#endif // __AVX__\n\n#if defined(__AVX512F__) && !defined(__AVX512DQ__)\n#define _mm512_xor_ps(a, b) \\\n _mm512_castsi512_ps( \\\n _mm512_xor_epi32(_mm512_castps_si512(a), _mm512_castps_si512(b)))\n#endif // __AVX512DQ__\n\n//! Reverse sign of value (GENERAL)\n#define NEGATE_FP32_GENERAL(v) -(v)\n\n//! Calculate Fused-Multiply-Add (SSE)\n#define FMA_FP32_SSE(xmm_m, xmm_q, xmm_sum) \\\n xmm_sum = _mm_fmadd_ps(xmm_m, xmm_q, xmm_sum);\n\n//! Calculate Fused-Multiply-Add (AVX)\n#define FMA_FP32_AVX(ymm_m, ymm_q, ymm_sum) \\\n ymm_sum = _mm256_fmadd_ps(ymm_m, ymm_q, ymm_sum);\n\n//! Calculate Fused-Multiply-Add (AVX512)\n#define FMA_FP32_AVX512(zmm_m, zmm_q, zmm_sum) \\\n zmm_sum = _mm512_fmadd_ps(zmm_m, zmm_q, zmm_sum);\n\n//! Calculate Fused-Multiply-Add (AVX512FP16)\n#define FMA_FP16_AVX512FP16(zmm_m, zmm_q, zmm_sum) \\\n zmm_sum = _mm512_fmadd_ph(zmm_m, zmm_q, zmm_sum);\n\n//! Calculate Fused-Multiply-Add (GENERAL)\n#define FMA_FP16_GENERAL(m, q, sum) sum += (m * q);\n\n//! Calculate Fused-Multiply-Add (GENERAL)\n#define FMA_FP32_GENERAL(m, q, sum) sum += (m * q);\n\n//! Calculate Fused-Multiply-Add (NEON)\n#define FMA_FP16_NEON(v_m, v_q, v_sum) v_sum = vfmaq_f16(v_sum, v_m, v_q);\n\n//! Calculate Fused-Multiply-Add (NEON)\n#define FMA_FP32_NEON(v_m, v_q, v_sum) v_sum = vfmaq_f32(v_sum, v_m, v_q);\n\n//! Calculate Fused-Multiply-Add (GENERAL)\n#define FMA_INT4_GENERAL(m, q, sum) \\\n sum += Int4MulTable[(((m) << 4) & 0xf0) | (((q) >> 0) & 0xf)] + \\\n Int4MulTable[(((m) >> 0) & 0xf0) | (((q) >> 4) & 0xf)];\n\n//! Calculate Fused-Multiply-Add (GENERAL)\n#define FMA_INT8_GENERAL(m, q, sum) sum += static_cast(m * q);\n\n//! Calculate Fused-Multiply-Add (SSE)\n#define FMA_INT8_SSE(xmm_m, xmm_q, xmm_sum) \\\n xmm_sum = _mm_add_epi32( \\\n _mm_madd_epi16( \\\n _mm_maddubs_epi16(_mm_abs_epi8(xmm_q), _mm_sign_epi8(xmm_m, xmm_q)), \\\n ONES_INT16_SSE), \\\n xmm_sum);\n\n//! Calculate Fused-Multiply-Add (AVX)\n#define FMA_INT8_AVX(ymm_m, ymm_q, ymm_sum) \\\n ymm_sum = _mm256_add_epi32( \\\n _mm256_madd_epi16(_mm256_maddubs_epi16(_mm256_abs_epi8(ymm_q), \\\n _mm256_sign_epi8(ymm_m, ymm_q)), \\\n ONES_INT16_AVX), \\\n ymm_sum);\n\n//! Calculate Fused-Multiply-Add (SSE)\n#define FMA_INT4_SSE(xmm_m, xmm_q, xmm_sum) \\\n { \\\n __m128i xmm_lhs = _mm_shuffle_epi8(INT4_LOOKUP_SSE, \\\n _mm_and_si128((xmm_m), MASK_INT4_SSE)); \\\n __m128i xmm_rhs = _mm_shuffle_epi8(INT4_LOOKUP_SSE, \\\n _mm_and_si128((xmm_q), MASK_INT4_SSE)); \\\n xmm_sum = _mm_add_epi32( \\\n _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_rhs), \\\n _mm_sign_epi8(xmm_lhs, xmm_rhs)), \\\n ONES_INT16_SSE), \\\n xmm_sum); \\\n xmm_lhs = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, \\\n _mm_and_si128(_mm_srli_epi32((xmm_m), 4), MASK_INT4_SSE)); \\\n xmm_rhs = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, \\\n _mm_and_si128(_mm_srli_epi32((xmm_q), 4), MASK_INT4_SSE)); \\\n xmm_sum = _mm_add_epi32( \\\n _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_rhs), \\\n _mm_sign_epi8(xmm_lhs, xmm_rhs)), \\\n ONES_INT16_SSE), \\\n xmm_sum); \\\n }\n\n//! Calculate Fused-Multiply-Add (AVX)\n#define FMA_INT4_AVX(ymm_m, ymm_q, ymm_sum) \\\n { \\\n __m256i ymm_lhs = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, _mm256_and_si256((ymm_m), MASK_INT4_AVX)); \\\n __m256i ymm_rhs = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, _mm256_and_si256((ymm_q), MASK_INT4_AVX)); \\\n ymm_sum = _mm256_add_epi32( \\\n _mm256_madd_epi16( \\\n _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_rhs), \\\n _mm256_sign_epi8(ymm_lhs, ymm_rhs)), \\\n ONES_INT16_AVX), \\\n ymm_sum); \\\n ymm_lhs = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, \\\n _mm256_and_si256(_mm256_srli_epi32((ymm_m), 4), MASK_INT4_AVX)); \\\n ymm_rhs = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, \\\n _mm256_and_si256(_mm256_srli_epi32((ymm_q), 4), MASK_INT4_AVX)); \\\n ymm_sum = _mm256_add_epi32( \\\n _mm256_madd_epi16( \\\n _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_rhs), \\\n _mm256_sign_epi8(ymm_lhs, ymm_rhs)), \\\n ONES_INT16_AVX), \\\n ymm_sum); \\\n }\n\n//! Compute the distance between matrix and query\n#define FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum) \\\n { \\\n __m128i xmm_lhs_0 = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, _mm_and_si128((xmm_lhs), MASK_INT4_SSE)); \\\n __m128i xmm_rhs_0 = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, _mm_and_si128((xmm_rhs), MASK_INT4_SSE)); \\\n __m128i xmm_lhs_1 = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, \\\n _mm_and_si128(_mm_srli_epi32((xmm_lhs), 4), MASK_INT4_SSE)); \\\n __m128i xmm_rhs_1 = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, \\\n _mm_and_si128(_mm_srli_epi32((xmm_rhs), 4), MASK_INT4_SSE)); \\\n xmm_lhs_0 = _mm_sign_epi8(xmm_lhs_0, xmm_rhs_0); \\\n xmm_lhs_1 = _mm_sign_epi8(xmm_lhs_1, xmm_rhs_1); \\\n xmm_rhs_0 = _mm_abs_epi8(xmm_rhs_0); \\\n xmm_rhs_1 = _mm_abs_epi8(xmm_rhs_1); \\\n xmm_lhs_0 = _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_0, xmm_lhs_0), \\\n ONES_INT16_SSE); \\\n xmm_lhs_1 = _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_1, xmm_lhs_1), \\\n ONES_INT16_SSE); \\\n xmm_sum = _mm_add_epi32(_mm_add_epi32(xmm_lhs_0, xmm_lhs_1), xmm_sum); \\\n }\n\n//! Compute the distance between matrix and query\n#define FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum) \\\n { \\\n __m256i ymm_lhs_0 = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, _mm256_and_si256((ymm_lhs), MASK_INT4_AVX)); \\\n __m256i ymm_rhs_0 = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, _mm256_and_si256((ymm_rhs), MASK_INT4_AVX)); \\\n __m256i ymm_lhs_1 = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, \\\n _mm256_and_si256(_mm256_srli_epi32((ymm_lhs), 4), MASK_INT4_AVX)); \\\n __m256i ymm_rhs_1 = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, \\\n _mm256_and_si256(_mm256_srli_epi32((ymm_rhs), 4), MASK_INT4_AVX)); \\\n ymm_lhs_0 = _mm256_sign_epi8(ymm_lhs_0, ymm_rhs_0); \\\n ymm_lhs_1 = _mm256_sign_epi8(ymm_lhs_1, ymm_rhs_1); \\\n ymm_rhs_0 = _mm256_abs_epi8(ymm_rhs_0); \\\n ymm_rhs_1 = _mm256_abs_epi8(ymm_rhs_1); \\\n ymm_lhs_0 = _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_0, ymm_lhs_0), \\\n ONES_INT16_AVX); \\\n ymm_lhs_1 = _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_1, ymm_lhs_1), \\\n ONES_INT16_AVX); \\\n ymm_sum = \\\n _mm256_add_epi32(_mm256_add_epi32(ymm_lhs_0, ymm_lhs_1), ymm_sum); \\\n }\n\n#define ACCUM_FP16_STEP_GENERAL FMA_FP16_GENERAL\n#define ACCUM_FP16_STEP_NEON FMA_FP16_NEON\n\n#define ACCUM_FP32_STEP_SSE FMA_FP32_SSE\n#define ACCUM_FP32_STEP_AVX FMA_FP32_AVX\n#define ACCUM_FP32_STEP_AVX512 FMA_FP32_AVX512\n#define ACCUM_FP32_STEP_NEON FMA_FP32_NEON\n\n#define ACCUM_INT4_STEP_SSE FMA_INT4_SSE\n#define ACCUM_INT4_STEP_AVX FMA_INT4_AVX\n\n#define ACCUM_INT8_STEP_SSE FMA_INT8_SSE\n#define ACCUM_INT8_STEP_AVX FMA_INT8_AVX\n" }, { "path": "src/ailego/math/distance_matrix_int32.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"matrix_define.i\"\n\n#if defined(__AVX__) && defined(__GNUC__)\n#define _mm256_set_m128i(a, b) \\\n _mm256_inserti128_si256(_mm256_castsi128_si256(b), (a), 1)\n#endif // __AVX__\n\n#if !defined(__AVX__)\n#define _mm_broadcast_si32(a) _mm_castps_si128(_mm_load1_ps((const float *)(a)))\n#else\n#define _mm_broadcast_si32(a) \\\n _mm_castps_si128(_mm_broadcast_ss((const float *)(a)))\n#define _mm256_broadcast_si32(a) \\\n _mm256_castps_si256(_mm256_broadcast_ss((const float *)(a)))\n#endif // !__AVX__\n\n//! Iterative process of computing distance (INT32, M=2, N=1)\n#define MATRIX_INT32_ITER_2X1_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_qi = _LOAD((const __m128i *)(qi)); \\\n __m128i xmm_mi = _LOAD((const __m128i *)(mi)); \\\n __m128i xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(1, 1, 0, 0)); \\\n _PROC(xmm_mi, xmm_pi, _RES##_0_0) \\\n xmm_mi = _LOAD((const __m128i *)(mi + 4)); \\\n xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(3, 3, 2, 2)); \\\n _PROC(xmm_mi, xmm_pi, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (INT32, M=2, N=2)\n#define MATRIX_INT32_ITER_2X2_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_qi = _LOAD((const __m128i *)(qi)); \\\n __m128i xmm_mi = _LOAD((const __m128i *)(mi)); \\\n __m128i xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \\\n _PROC(xmm_mi, xmm_pi, _RES##_0_0) \\\n xmm_pi = _mm_shuffle_epi32(xmm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \\\n _PROC(xmm_mi, xmm_pi, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (INT32, M=4, N=1)\n#define MATRIX_INT32_ITER_4X1_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n _PROC(xmm_mi_0, xmm_qi, _RES##_0_0) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n _PROC(xmm_mi_1, xmm_qi, _RES##_1_0) \\\n }\n\n//! Iterative process of computing distance (INT32, M=4, N=2)\n#define MATRIX_INT32_ITER_4X2_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \\\n __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \\\n __m128i xmm_mi = _LOAD((const __m128i *)(mi)); \\\n MATRIX_VAR_PROC(1, 2, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=4, N=4)\n#define MATRIX_INT32_ITER_4X4_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi = _LOAD((const __m128i *)(mi)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n _PROC(xmm_mi, xmm_qi, _RES##_0_0) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n _PROC(xmm_mi, xmm_qi, _RES##_0_1) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n _PROC(xmm_mi, xmm_qi, _RES##_0_2) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n _PROC(xmm_mi, xmm_qi, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (INT32, M=8, N=1)\n#define MATRIX_INT32_ITER_8X1_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n MATRIX_VAR_PROC(2, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=8, N=2)\n#define MATRIX_INT32_ITER_8X2_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \\\n __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \\\n __m128i xmm_mi = _LOAD((const __m128i *)(mi + 0)); \\\n MATRIX_VAR_PROC(1, 2, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_mi = _LOAD((const __m128i *)(mi + 4)); \\\n MATRIX_VAR_PROC(1, 2, 1, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=8, N=4)\n#define MATRIX_INT32_ITER_8X4_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=8, N=8)\n#define MATRIX_INT32_ITER_8X8_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(2, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(2, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(2, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(2, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=1)\n#define MATRIX_INT32_ITER_16X1_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=2)\n#define MATRIX_INT32_ITER_16X2_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=4)\n#define MATRIX_INT32_ITER_16X4_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=8)\n#define MATRIX_INT32_ITER_16X8_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=16)\n#define MATRIX_INT32_ITER_16X16_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 8); \\\n MATRIX_VAR_PROC(4, 1, 8, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 9); \\\n MATRIX_VAR_PROC(4, 1, 9, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 10); \\\n MATRIX_VAR_PROC(4, 1, 10, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 11); \\\n MATRIX_VAR_PROC(4, 1, 11, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 12); \\\n MATRIX_VAR_PROC(4, 1, 12, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 13); \\\n MATRIX_VAR_PROC(4, 1, 13, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 14); \\\n MATRIX_VAR_PROC(4, 1, 14, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 15); \\\n MATRIX_VAR_PROC(4, 1, 15, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=1)\n#define MATRIX_INT32_ITER_32X1_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_qi = _mm_broadcast_si32(qi); \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n _PROC(xmm_mi_0, xmm_qi, _RES##_0_0) \\\n _PROC(xmm_mi_1, xmm_qi, _RES##_1_0) \\\n _PROC(xmm_mi_2, xmm_qi, _RES##_2_0) \\\n _PROC(xmm_mi_3, xmm_qi, _RES##_3_0) \\\n xmm_mi_0 = _LOAD((const __m128i *)(mi + 16)); \\\n xmm_mi_1 = _LOAD((const __m128i *)(mi + 20)); \\\n xmm_mi_2 = _LOAD((const __m128i *)(mi + 24)); \\\n xmm_mi_3 = _LOAD((const __m128i *)(mi + 28)); \\\n _PROC(xmm_mi_0, xmm_qi, _RES##_4_0) \\\n _PROC(xmm_mi_1, xmm_qi, _RES##_5_0) \\\n _PROC(xmm_mi_2, xmm_qi, _RES##_6_0) \\\n _PROC(xmm_mi_3, xmm_qi, _RES##_7_0) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=2)\n#define MATRIX_INT32_ITER_32X2_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \\\n __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n MATRIX_VAR_PROC(1, 2, 0, xmm_mi_0, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 1, xmm_mi_1, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 2, xmm_mi_2, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 3, xmm_mi_3, xmm_qi, _RES, _PROC) \\\n xmm_mi_0 = _LOAD((const __m128i *)(mi + 16)); \\\n xmm_mi_1 = _LOAD((const __m128i *)(mi + 20)); \\\n xmm_mi_2 = _LOAD((const __m128i *)(mi + 24)); \\\n xmm_mi_3 = _LOAD((const __m128i *)(mi + 28)); \\\n MATRIX_VAR_PROC(1, 2, 4, xmm_mi_0, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 5, xmm_mi_1, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 6, xmm_mi_2, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 2, 7, xmm_mi_3, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=4)\n#define MATRIX_INT32_ITER_32X4_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \\\n __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \\\n __m128i xmm_qi_2 = _mm_broadcast_si32(qi + 2); \\\n __m128i xmm_qi_3 = _mm_broadcast_si32(qi + 3); \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n MATRIX_VAR_PROC(1, 4, 0, xmm_mi_0, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 1, xmm_mi_1, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 2, xmm_mi_2, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 3, xmm_mi_3, xmm_qi, _RES, _PROC) \\\n xmm_mi_0 = _LOAD((const __m128i *)(mi + 16)); \\\n xmm_mi_1 = _LOAD((const __m128i *)(mi + 20)); \\\n xmm_mi_2 = _LOAD((const __m128i *)(mi + 24)); \\\n xmm_mi_3 = _LOAD((const __m128i *)(mi + 28)); \\\n MATRIX_VAR_PROC(1, 4, 4, xmm_mi_0, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 5, xmm_mi_1, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 6, xmm_mi_2, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 4, 7, xmm_mi_3, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=8)\n#define MATRIX_INT32_ITER_32X8_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_qi_0 = _mm_broadcast_si32(qi + 0); \\\n __m128i xmm_qi_1 = _mm_broadcast_si32(qi + 1); \\\n __m128i xmm_qi_2 = _mm_broadcast_si32(qi + 2); \\\n __m128i xmm_qi_3 = _mm_broadcast_si32(qi + 3); \\\n __m128i xmm_qi_4 = _mm_broadcast_si32(qi + 4); \\\n __m128i xmm_qi_5 = _mm_broadcast_si32(qi + 5); \\\n __m128i xmm_qi_6 = _mm_broadcast_si32(qi + 6); \\\n __m128i xmm_qi_7 = _mm_broadcast_si32(qi + 7); \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n MATRIX_VAR_PROC(1, 8, 0, xmm_mi_0, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 1, xmm_mi_1, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 2, xmm_mi_2, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 3, xmm_mi_3, xmm_qi, _RES, _PROC) \\\n xmm_mi_0 = _LOAD((const __m128i *)(mi + 16)); \\\n xmm_mi_1 = _LOAD((const __m128i *)(mi + 20)); \\\n xmm_mi_2 = _LOAD((const __m128i *)(mi + 24)); \\\n xmm_mi_3 = _LOAD((const __m128i *)(mi + 28)); \\\n MATRIX_VAR_PROC(1, 8, 4, xmm_mi_0, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 5, xmm_mi_1, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 6, xmm_mi_2, xmm_qi, _RES, _PROC) \\\n MATRIX_VAR_PROC(1, 8, 7, xmm_mi_3, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=16)\n#define MATRIX_INT32_ITER_32X16_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n __m128i xmm_mi_4 = _LOAD((const __m128i *)(mi + 16)); \\\n __m128i xmm_mi_5 = _LOAD((const __m128i *)(mi + 20)); \\\n __m128i xmm_mi_6 = _LOAD((const __m128i *)(mi + 24)); \\\n __m128i xmm_mi_7 = _LOAD((const __m128i *)(mi + 28)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(8, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(8, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(8, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(8, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(8, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(8, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(8, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(8, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 8); \\\n MATRIX_VAR_PROC(8, 1, 8, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 9); \\\n MATRIX_VAR_PROC(8, 1, 9, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 10); \\\n MATRIX_VAR_PROC(8, 1, 10, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 11); \\\n MATRIX_VAR_PROC(8, 1, 11, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 12); \\\n MATRIX_VAR_PROC(8, 1, 12, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 13); \\\n MATRIX_VAR_PROC(8, 1, 13, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 14); \\\n MATRIX_VAR_PROC(8, 1, 14, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 15); \\\n MATRIX_VAR_PROC(8, 1, 15, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=32)\n#define MATRIX_INT32_ITER_32X32_SSE(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m128i xmm_mi_0 = _LOAD((const __m128i *)(mi + 0)); \\\n __m128i xmm_mi_1 = _LOAD((const __m128i *)(mi + 4)); \\\n __m128i xmm_mi_2 = _LOAD((const __m128i *)(mi + 8)); \\\n __m128i xmm_mi_3 = _LOAD((const __m128i *)(mi + 12)); \\\n __m128i xmm_mi_4 = _LOAD((const __m128i *)(mi + 16)); \\\n __m128i xmm_mi_5 = _LOAD((const __m128i *)(mi + 20)); \\\n __m128i xmm_mi_6 = _LOAD((const __m128i *)(mi + 24)); \\\n __m128i xmm_mi_7 = _LOAD((const __m128i *)(mi + 28)); \\\n __m128i xmm_qi = _mm_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(8, 1, 0, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(8, 1, 1, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(8, 1, 2, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(8, 1, 3, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(8, 1, 4, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(8, 1, 5, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(8, 1, 6, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(8, 1, 7, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 8); \\\n MATRIX_VAR_PROC(8, 1, 8, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 9); \\\n MATRIX_VAR_PROC(8, 1, 9, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 10); \\\n MATRIX_VAR_PROC(8, 1, 10, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 11); \\\n MATRIX_VAR_PROC(8, 1, 11, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 12); \\\n MATRIX_VAR_PROC(8, 1, 12, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 13); \\\n MATRIX_VAR_PROC(8, 1, 13, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 14); \\\n MATRIX_VAR_PROC(8, 1, 14, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 15); \\\n MATRIX_VAR_PROC(8, 1, 15, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 16); \\\n MATRIX_VAR_PROC(8, 1, 16, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 17); \\\n MATRIX_VAR_PROC(8, 1, 17, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 18); \\\n MATRIX_VAR_PROC(8, 1, 18, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 19); \\\n MATRIX_VAR_PROC(8, 1, 19, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 20); \\\n MATRIX_VAR_PROC(8, 1, 20, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 21); \\\n MATRIX_VAR_PROC(8, 1, 21, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 22); \\\n MATRIX_VAR_PROC(8, 1, 22, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 23); \\\n MATRIX_VAR_PROC(8, 1, 23, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 24); \\\n MATRIX_VAR_PROC(8, 1, 24, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 25); \\\n MATRIX_VAR_PROC(8, 1, 25, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 26); \\\n MATRIX_VAR_PROC(8, 1, 26, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 27); \\\n MATRIX_VAR_PROC(8, 1, 27, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 28); \\\n MATRIX_VAR_PROC(8, 1, 28, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 29); \\\n MATRIX_VAR_PROC(8, 1, 29, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 30); \\\n MATRIX_VAR_PROC(8, 1, 30, xmm_mi, xmm_qi, _RES, _PROC) \\\n xmm_qi = _mm_broadcast_si32(qi + 31); \\\n MATRIX_VAR_PROC(8, 1, 31, xmm_mi, xmm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=2, N=1)\n#define MATRIX_INT32_ITER_2X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_qi = _mm256_set_epi32(qi[3], qi[3], qi[2], qi[2], qi[1], \\\n qi[1], qi[0], qi[0]); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (INT32, M=2, N=2)\n#define MATRIX_INT32_ITER_2X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi = _LOAD((const __m256i *)(qi)); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_0) \\\n ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (INT32, M=4, N=1)\n#define MATRIX_INT32_ITER_4X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_qi = \\\n _mm256_set_m128i(_mm_broadcast_si32(qi + 1), _mm_broadcast_si32(qi)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (INT32, M=4, N=2)\n#define MATRIX_INT32_ITER_4X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_qi = _mm256_set_m128i(_mm_broadcast_si32(qi + 2), \\\n _mm_broadcast_si32(qi + 0)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n ymm_qi = _mm256_set_m128i(_mm_broadcast_si32(qi + 3), \\\n _mm_broadcast_si32(qi + 1)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (INT32, M=4, N=4)\n#define MATRIX_INT32_ITER_4X4_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi = _LOAD((const __m256i *)(qi)); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(0, 0, 0, 0)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_0) \\\n ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(1, 1, 1, 1)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_1) \\\n ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(2, 2, 2, 2)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_2) \\\n ymm_pi = _mm256_shuffle_epi32(ymm_qi, _MM_SHUFFLE(3, 3, 3, 3)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (INT32, M=8, N=1)\n#define MATRIX_INT32_ITER_8X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n }\n\n//! Iterative process of computing distance (INT32, M=8, N=2)\n#define MATRIX_INT32_ITER_8X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi_0 = _mm256_broadcast_si32(qi + 0); \\\n __m256i ymm_qi_1 = _mm256_broadcast_si32(qi + 1); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=8, N=4)\n#define MATRIX_INT32_ITER_8X4_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_1) \\\n ymm_qi = _mm256_broadcast_si32(qi + 2); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_2) \\\n ymm_qi = _mm256_broadcast_si32(qi + 3); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (INT32, M=8, N=8)\n#define MATRIX_INT32_ITER_8X8_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_1) \\\n ymm_qi = _mm256_broadcast_si32(qi + 2); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_2) \\\n ymm_qi = _mm256_broadcast_si32(qi + 3); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_3) \\\n ymm_qi = _mm256_broadcast_si32(qi + 4); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_4) \\\n ymm_qi = _mm256_broadcast_si32(qi + 5); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_5) \\\n ymm_qi = _mm256_broadcast_si32(qi + 6); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_6) \\\n ymm_qi = _mm256_broadcast_si32(qi + 7); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_7) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=1)\n#define MATRIX_INT32_ITER_16X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=2)\n#define MATRIX_INT32_ITER_16X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=4)\n#define MATRIX_INT32_ITER_16X4_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=8)\n#define MATRIX_INT32_ITER_16X8_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(2, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(2, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(2, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(2, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=16, N=16)\n#define MATRIX_INT32_ITER_16X16_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(2, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(2, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(2, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(2, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 8); \\\n MATRIX_VAR_PROC(2, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 9); \\\n MATRIX_VAR_PROC(2, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 10); \\\n MATRIX_VAR_PROC(2, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 11); \\\n MATRIX_VAR_PROC(2, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 12); \\\n MATRIX_VAR_PROC(2, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 13); \\\n MATRIX_VAR_PROC(2, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 14); \\\n MATRIX_VAR_PROC(2, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 15); \\\n MATRIX_VAR_PROC(2, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=1)\n#define MATRIX_INT32_ITER_32X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=2)\n#define MATRIX_INT32_ITER_32X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=4)\n#define MATRIX_INT32_ITER_32X4_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=8)\n#define MATRIX_INT32_ITER_32X8_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=16)\n#define MATRIX_INT32_ITER_32X16_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 8); \\\n MATRIX_VAR_PROC(4, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 9); \\\n MATRIX_VAR_PROC(4, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 10); \\\n MATRIX_VAR_PROC(4, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 11); \\\n MATRIX_VAR_PROC(4, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 12); \\\n MATRIX_VAR_PROC(4, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 13); \\\n MATRIX_VAR_PROC(4, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 14); \\\n MATRIX_VAR_PROC(4, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 15); \\\n MATRIX_VAR_PROC(4, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT32, M=32, N=32)\n#define MATRIX_INT32_ITER_32X32_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 16)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 24)); \\\n __m256i ymm_qi = _mm256_broadcast_si32(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 8); \\\n MATRIX_VAR_PROC(4, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 9); \\\n MATRIX_VAR_PROC(4, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 10); \\\n MATRIX_VAR_PROC(4, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 11); \\\n MATRIX_VAR_PROC(4, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 12); \\\n MATRIX_VAR_PROC(4, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 13); \\\n MATRIX_VAR_PROC(4, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 14); \\\n MATRIX_VAR_PROC(4, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 15); \\\n MATRIX_VAR_PROC(4, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 16); \\\n MATRIX_VAR_PROC(4, 1, 16, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 17); \\\n MATRIX_VAR_PROC(4, 1, 17, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 18); \\\n MATRIX_VAR_PROC(4, 1, 18, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 19); \\\n MATRIX_VAR_PROC(4, 1, 19, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 20); \\\n MATRIX_VAR_PROC(4, 1, 20, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 21); \\\n MATRIX_VAR_PROC(4, 1, 21, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 22); \\\n MATRIX_VAR_PROC(4, 1, 22, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 23); \\\n MATRIX_VAR_PROC(4, 1, 23, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 24); \\\n MATRIX_VAR_PROC(4, 1, 24, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 25); \\\n MATRIX_VAR_PROC(4, 1, 25, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 26); \\\n MATRIX_VAR_PROC(4, 1, 26, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 27); \\\n MATRIX_VAR_PROC(4, 1, 27, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 28); \\\n MATRIX_VAR_PROC(4, 1, 28, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 29); \\\n MATRIX_VAR_PROC(4, 1, 29, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 30); \\\n MATRIX_VAR_PROC(4, 1, 30, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si32(qi + 31); \\\n MATRIX_VAR_PROC(4, 1, 31, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n" }, { "path": "src/ailego/math/distance_matrix_int64.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"matrix_define.i\"\n\n#if defined(__AVX__)\n#define _mm256_broadcast_si64(a) \\\n _mm256_castpd_si256(_mm256_broadcast_sd((const double *)(a)))\n#endif // __AVX__\n\n//! Iterative process of computing distance (INT64, M=2, N=1)\n#define MATRIX_INT64_ITER_2X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi = _LOAD((const __m256i *)(qi)); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_pi = \\\n _mm256_permute4x64_epi64(ymm_qi, _MM_SHUFFLE(1, 1, 0, 0)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n ymm_pi = _mm256_permute4x64_epi64(ymm_qi, _MM_SHUFFLE(3, 3, 2, 2)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (INT64, M=2, N=2)\n#define MATRIX_INT64_ITER_2X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi = _LOAD((const __m256i *)(qi)); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_pi = \\\n _mm256_permute4x64_epi64(ymm_qi, _MM_SHUFFLE(2, 2, 0, 0)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_0) \\\n ymm_pi = _mm256_permute4x64_epi64(ymm_qi, _MM_SHUFFLE(3, 3, 1, 1)); \\\n _PROC(ymm_mi, ymm_pi, _RES##_0_1) \\\n }\n\n//! Iterative process of computing distance (INT64, M=4, N=1)\n#define MATRIX_INT64_ITER_4X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n ymm_qi = _mm256_broadcast_si64(qi + 1); \\\n _PROC(ymm_mi, ymm_qi, _RES##_1_0) \\\n }\n\n//! Iterative process of computing distance (INT64, M=4, N=2)\n#define MATRIX_INT64_ITER_4X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \\\n __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=4, N=4)\n#define MATRIX_INT64_ITER_4X4_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi)); \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n ymm_qi = _mm256_broadcast_si64(qi + 1); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_1) \\\n ymm_qi = _mm256_broadcast_si64(qi + 2); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_2) \\\n ymm_qi = _mm256_broadcast_si64(qi + 3); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_3) \\\n }\n\n//! Iterative process of computing distance (INT64, M=8, N=1)\n#define MATRIX_INT64_ITER_8X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_1_0) \\\n }\n\n//! Iterative process of computing distance (INT64, M=8, N=2)\n#define MATRIX_INT64_ITER_8X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \\\n __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \\\n MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n MATRIX_VAR_PROC(1, 2, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=8, N=4)\n#define MATRIX_INT64_ITER_8X4_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=8, N=8)\n#define MATRIX_INT64_ITER_8X8_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \\\n MATRIX_VAR_PROC(2, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 1); \\\n MATRIX_VAR_PROC(2, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 2); \\\n MATRIX_VAR_PROC(2, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 3); \\\n MATRIX_VAR_PROC(2, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 4); \\\n MATRIX_VAR_PROC(2, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 5); \\\n MATRIX_VAR_PROC(2, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 6); \\\n MATRIX_VAR_PROC(2, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 7); \\\n MATRIX_VAR_PROC(2, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=16, N=1)\n#define MATRIX_INT64_ITER_16X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_1_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 8)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_2_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 12)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_3_0) \\\n }\n\n//! Iterative process of computing distance (INT64, M=16, N=2)\n#define MATRIX_INT64_ITER_16X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \\\n __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \\\n MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n MATRIX_VAR_PROC(1, 2, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 8)); \\\n MATRIX_VAR_PROC(1, 2, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 12)); \\\n MATRIX_VAR_PROC(1, 2, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=16, N=4)\n#define MATRIX_INT64_ITER_16X4_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=16, N=8)\n#define MATRIX_INT64_ITER_16X8_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=16, N=16)\n#define MATRIX_INT64_ITER_16X16_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \\\n MATRIX_VAR_PROC(4, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 1); \\\n MATRIX_VAR_PROC(4, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 2); \\\n MATRIX_VAR_PROC(4, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 3); \\\n MATRIX_VAR_PROC(4, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 4); \\\n MATRIX_VAR_PROC(4, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 5); \\\n MATRIX_VAR_PROC(4, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 6); \\\n MATRIX_VAR_PROC(4, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 7); \\\n MATRIX_VAR_PROC(4, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 8); \\\n MATRIX_VAR_PROC(4, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 9); \\\n MATRIX_VAR_PROC(4, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 10); \\\n MATRIX_VAR_PROC(4, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 11); \\\n MATRIX_VAR_PROC(4, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 12); \\\n MATRIX_VAR_PROC(4, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 13); \\\n MATRIX_VAR_PROC(4, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 14); \\\n MATRIX_VAR_PROC(4, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 15); \\\n MATRIX_VAR_PROC(4, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=32, N=1)\n#define MATRIX_INT64_ITER_32X1_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_0_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_1_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 8)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_2_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 12)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_3_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 16)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_4_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 20)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_5_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 24)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_6_0) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 28)); \\\n _PROC(ymm_mi, ymm_qi, _RES##_7_0) \\\n }\n\n//! Iterative process of computing distance (INT64, M=32, N=2)\n#define MATRIX_INT64_ITER_32X2_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \\\n __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \\\n MATRIX_VAR_PROC(1, 2, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n MATRIX_VAR_PROC(1, 2, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 8)); \\\n MATRIX_VAR_PROC(1, 2, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 12)); \\\n MATRIX_VAR_PROC(1, 2, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 16)); \\\n MATRIX_VAR_PROC(1, 2, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 20)); \\\n MATRIX_VAR_PROC(1, 2, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 24)); \\\n MATRIX_VAR_PROC(1, 2, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 28)); \\\n MATRIX_VAR_PROC(1, 2, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=32, N=4)\n#define MATRIX_INT64_ITER_32X4_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \\\n __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \\\n __m256i ymm_qi_2 = _mm256_broadcast_si64(qi + 2); \\\n __m256i ymm_qi_3 = _mm256_broadcast_si64(qi + 3); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \\\n MATRIX_VAR_PROC(1, 4, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n MATRIX_VAR_PROC(1, 4, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 8)); \\\n MATRIX_VAR_PROC(1, 4, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 12)); \\\n MATRIX_VAR_PROC(1, 4, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 16)); \\\n MATRIX_VAR_PROC(1, 4, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 20)); \\\n MATRIX_VAR_PROC(1, 4, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 24)); \\\n MATRIX_VAR_PROC(1, 4, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 28)); \\\n MATRIX_VAR_PROC(1, 4, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=32, N=8)\n#define MATRIX_INT64_ITER_32X8_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_qi_0 = _mm256_broadcast_si64(qi + 0); \\\n __m256i ymm_qi_1 = _mm256_broadcast_si64(qi + 1); \\\n __m256i ymm_qi_2 = _mm256_broadcast_si64(qi + 2); \\\n __m256i ymm_qi_3 = _mm256_broadcast_si64(qi + 3); \\\n __m256i ymm_qi_4 = _mm256_broadcast_si64(qi + 4); \\\n __m256i ymm_qi_5 = _mm256_broadcast_si64(qi + 5); \\\n __m256i ymm_qi_6 = _mm256_broadcast_si64(qi + 6); \\\n __m256i ymm_qi_7 = _mm256_broadcast_si64(qi + 7); \\\n __m256i ymm_mi = _LOAD((const __m256i *)(mi + 0)); \\\n MATRIX_VAR_PROC(1, 8, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 4)); \\\n MATRIX_VAR_PROC(1, 8, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 8)); \\\n MATRIX_VAR_PROC(1, 8, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 12)); \\\n MATRIX_VAR_PROC(1, 8, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 16)); \\\n MATRIX_VAR_PROC(1, 8, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 20)); \\\n MATRIX_VAR_PROC(1, 8, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 24)); \\\n MATRIX_VAR_PROC(1, 8, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_mi = _LOAD((const __m256i *)(mi + 28)); \\\n MATRIX_VAR_PROC(1, 8, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=32, N=16)\n#define MATRIX_INT64_ITER_32X16_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \\\n __m256i ymm_mi_4 = _LOAD((const __m256i *)(mi + 16)); \\\n __m256i ymm_mi_5 = _LOAD((const __m256i *)(mi + 20)); \\\n __m256i ymm_mi_6 = _LOAD((const __m256i *)(mi + 24)); \\\n __m256i ymm_mi_7 = _LOAD((const __m256i *)(mi + 28)); \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \\\n MATRIX_VAR_PROC(8, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 1); \\\n MATRIX_VAR_PROC(8, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 2); \\\n MATRIX_VAR_PROC(8, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 3); \\\n MATRIX_VAR_PROC(8, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 4); \\\n MATRIX_VAR_PROC(8, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 5); \\\n MATRIX_VAR_PROC(8, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 6); \\\n MATRIX_VAR_PROC(8, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 7); \\\n MATRIX_VAR_PROC(8, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 8); \\\n MATRIX_VAR_PROC(8, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 9); \\\n MATRIX_VAR_PROC(8, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 10); \\\n MATRIX_VAR_PROC(8, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 11); \\\n MATRIX_VAR_PROC(8, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 12); \\\n MATRIX_VAR_PROC(8, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 13); \\\n MATRIX_VAR_PROC(8, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 14); \\\n MATRIX_VAR_PROC(8, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 15); \\\n MATRIX_VAR_PROC(8, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n\n//! Iterative process of computing distance (INT64, M=32, N=32)\n#define MATRIX_INT64_ITER_32X32_AVX(mi, qi, _RES, _LOAD, _PROC) \\\n { \\\n __m256i ymm_mi_0 = _LOAD((const __m256i *)(mi + 0)); \\\n __m256i ymm_mi_1 = _LOAD((const __m256i *)(mi + 4)); \\\n __m256i ymm_mi_2 = _LOAD((const __m256i *)(mi + 8)); \\\n __m256i ymm_mi_3 = _LOAD((const __m256i *)(mi + 12)); \\\n __m256i ymm_mi_4 = _LOAD((const __m256i *)(mi + 16)); \\\n __m256i ymm_mi_5 = _LOAD((const __m256i *)(mi + 20)); \\\n __m256i ymm_mi_6 = _LOAD((const __m256i *)(mi + 24)); \\\n __m256i ymm_mi_7 = _LOAD((const __m256i *)(mi + 28)); \\\n __m256i ymm_qi = _mm256_broadcast_si64(qi + 0); \\\n MATRIX_VAR_PROC(8, 1, 0, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 1); \\\n MATRIX_VAR_PROC(8, 1, 1, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 2); \\\n MATRIX_VAR_PROC(8, 1, 2, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 3); \\\n MATRIX_VAR_PROC(8, 1, 3, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 4); \\\n MATRIX_VAR_PROC(8, 1, 4, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 5); \\\n MATRIX_VAR_PROC(8, 1, 5, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 6); \\\n MATRIX_VAR_PROC(8, 1, 6, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 7); \\\n MATRIX_VAR_PROC(8, 1, 7, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 8); \\\n MATRIX_VAR_PROC(8, 1, 8, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 9); \\\n MATRIX_VAR_PROC(8, 1, 9, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 10); \\\n MATRIX_VAR_PROC(8, 1, 10, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 11); \\\n MATRIX_VAR_PROC(8, 1, 11, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 12); \\\n MATRIX_VAR_PROC(8, 1, 12, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 13); \\\n MATRIX_VAR_PROC(8, 1, 13, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 14); \\\n MATRIX_VAR_PROC(8, 1, 14, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 15); \\\n MATRIX_VAR_PROC(8, 1, 15, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 16); \\\n MATRIX_VAR_PROC(8, 1, 16, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 17); \\\n MATRIX_VAR_PROC(8, 1, 17, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 18); \\\n MATRIX_VAR_PROC(8, 1, 18, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 19); \\\n MATRIX_VAR_PROC(8, 1, 19, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 20); \\\n MATRIX_VAR_PROC(8, 1, 20, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 21); \\\n MATRIX_VAR_PROC(8, 1, 21, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 22); \\\n MATRIX_VAR_PROC(8, 1, 22, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 23); \\\n MATRIX_VAR_PROC(8, 1, 23, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 24); \\\n MATRIX_VAR_PROC(8, 1, 24, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 25); \\\n MATRIX_VAR_PROC(8, 1, 25, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 26); \\\n MATRIX_VAR_PROC(8, 1, 26, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 27); \\\n MATRIX_VAR_PROC(8, 1, 27, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 28); \\\n MATRIX_VAR_PROC(8, 1, 28, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 29); \\\n MATRIX_VAR_PROC(8, 1, 29, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 30); \\\n MATRIX_VAR_PROC(8, 1, 30, ymm_mi, ymm_qi, _RES, _PROC) \\\n ymm_qi = _mm256_broadcast_si64(qi + 31); \\\n MATRIX_VAR_PROC(8, 1, 31, ymm_mi, ymm_qi, _RES, _PROC) \\\n }\n" }, { "path": "src/ailego/math/distance_matrix_mips_utility.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n//! Calculate Fused-Multiply-Add (AVX512)\n#define FMA_FP32_AVX512(zmm_m, zmm_q, zmm_sum) \\\n zmm_sum = _mm512_fmadd_ps(zmm_m, zmm_q, zmm_sum);\n\n#define FMA_MASK_FP32_AVX512(zmm_m, zmm_q, zmm_sum, mask) \\\n zmm_sum = _mm512_mask3_fmadd_ps(zmm_m, zmm_q, zmm_sum, mask);\n\n#define HorizontalAdd_FP16_NEON(v) \\\n vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(v)), vcvt_high_f32_f16(v)))\n\n#define HorizontalAdd_FP32_V512_TO_V256(zmm) \\\n _mm256_add_ps( \\\n _mm512_castps512_ps256(zmm), \\\n _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(zmm), 1)))\n\n//! Calculate Fused-Multiply-Add (AVX, FP16)\n#define FMA_FP16_GENERAL(lhs, rhs, sum, norm1, norm2) \\\n { \\\n float v1 = lhs; \\\n float v2 = rhs; \\\n sum += v1 * v2; \\\n norm1 += v1 * v1; \\\n norm2 += v2 * v2; \\\n }\n\n//! Calculate Fused-Multiply-Add (GENERAL)\n#define FMA_FP32_GENERAL(lhs, rhs, sum, norm1, norm2) \\\n { \\\n sum += (lhs) * (rhs); \\\n norm1 += (lhs) * (lhs); \\\n norm2 += (rhs) * (rhs); \\\n }\n\n#if defined(__SSE4_1__)\n//! Four-bits Convert Table\nstatic const AILEGO_ALIGNED(32) int8_t Int4ConvertTable[32] = {\n 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1,\n 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1};\n#endif // __SSE4_1__\n\n#if defined(__SSE4_1__)\nstatic const __m128i MASK_INT4_SSE = _mm_set1_epi32(0x0f0f0f0f);\nstatic const __m128i ONES_INT16_SSE = _mm_set1_epi32(0x00010001);\nstatic const __m128i INT4_LOOKUP_SSE =\n _mm_load_si128((const __m128i *)Int4ConvertTable);\n#endif // __SSE4_1__\n\n#if defined(__AVX2__)\nstatic const __m256i MASK_INT4_AVX = _mm256_set1_epi32(0x0f0f0f0f);\nstatic const __m256i ONES_INT16_AVX = _mm256_set1_epi32(0x00010001);\nstatic const __m256i INT4_LOOKUP_AVX =\n _mm256_load_si256((const __m256i *)Int4ConvertTable);\n#endif // __AVX2__\n\n//! Calculate Fused-Multiply-Add (GENERAL)\n#define FMA_INT4_GENERAL(lhs, rhs, sum, norm1, norm2) \\\n { \\\n sum += Int4MulTable[(((lhs) << 4) & 0xf0) | (((rhs) >> 0) & 0xf)] + \\\n Int4MulTable[(((lhs) >> 0) & 0xf0) | (((rhs) >> 4) & 0xf)]; \\\n norm1 += static_cast( \\\n ((int8_t)((lhs) << 4) >> 4) * ((int8_t)((lhs) << 4) >> 4) + \\\n ((int8_t)((lhs) & 0xf0) >> 4) * ((int8_t)((lhs) & 0xf0) >> 4)); \\\n norm2 += static_cast( \\\n ((int8_t)((rhs) << 4) >> 4) * ((int8_t)((rhs) << 4) >> 4) + \\\n ((int8_t)((rhs) & 0xf0) >> 4) * ((int8_t)((rhs) & 0xf0) >> 4)); \\\n }\n\n\n//! Compute the distance between matrix and query (SSE)\n#define FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum_0, xmm_sum_norm1, \\\n xmm_sum_norm2) \\\n { \\\n __m128i xmm_lhs_0 = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, _mm_and_si128((xmm_lhs), MASK_INT4_SSE)); \\\n __m128i xmm_rhs_0 = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, _mm_and_si128((xmm_rhs), MASK_INT4_SSE)); \\\n __m128i xmm_lhs_1 = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, \\\n _mm_and_si128(_mm_srli_epi32((xmm_lhs), 4), MASK_INT4_SSE)); \\\n __m128i xmm_rhs_1 = _mm_shuffle_epi8( \\\n INT4_LOOKUP_SSE, \\\n _mm_and_si128(_mm_srli_epi32((xmm_rhs), 4), MASK_INT4_SSE)); \\\n FMA_INT8_SSE(xmm_lhs_0, xmm_rhs_0, xmm_sum_0); \\\n FMA_INT8_SSE(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1); \\\n FMA_INT8_SSE(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2); \\\n FMA_INT8_SSE(xmm_lhs_1, xmm_rhs_1, xmm_sum_0); \\\n FMA_INT8_SSE(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1); \\\n FMA_INT8_SSE(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2); \\\n }\n\n//! Calculate Fused-Multiply-Add (GENERAL)\n#define FMA_INT8_GENERAL(lhs, rhs, sum, norm1, norm2) \\\n { \\\n sum += static_cast(lhs * rhs); \\\n norm1 += static_cast(lhs * lhs); \\\n norm2 += static_cast(rhs * rhs); \\\n }\n\n//! Calculate Fused-Multiply-Add (SSE)\n#define FMA_INT8_SSE(xmm_lhs, xmm_rhs, xmm_sum) \\\n xmm_sum = _mm_add_epi32( \\\n _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_rhs), \\\n _mm_sign_epi8(xmm_lhs, xmm_rhs)), \\\n ONES_INT16_SSE), \\\n xmm_sum)\n\n//! Calculate Fused-Multiply-Add (AVX)\n#define FMA_INT8_AVX(ymm_lhs, ymm_rhs, ymm_sum) \\\n ymm_sum = _mm256_add_epi32( \\\n _mm256_madd_epi16( \\\n _mm256_maddubs_epi16(_mm256_abs_epi8(ymm_rhs), \\\n _mm256_sign_epi8(ymm_lhs, ymm_rhs)), \\\n ONES_INT16_AVX), \\\n ymm_sum)\n\n#define FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_rhs, ymm_sum) \\\n ymm_sum = _mm256_add_epi32( \\\n _mm256_set_m128i( \\\n _mm_setzero_si128(), \\\n _mm_madd_epi16(_mm_maddubs_epi16(_mm_abs_epi8(xmm_rhs), \\\n _mm_sign_epi8(xmm_lhs, xmm_rhs)), \\\n ONES_INT16_SSE)), \\\n ymm_sum)\n\n//! Compute the distance between matrix and query (AVX)\n#define FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum_0, ymm_sum1, \\\n ymm_sum_norm1, ymm_sum_norm2) \\\n { \\\n __m256i ymm_lhs_0 = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, _mm256_and_si256((ymm_lhs), MASK_INT4_AVX)); \\\n __m256i ymm_rhs_0 = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, _mm256_and_si256((ymm_rhs), MASK_INT4_AVX)); \\\n __m256i ymm_lhs_1 = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, \\\n _mm256_and_si256(_mm256_srli_epi32((ymm_lhs), 4), MASK_INT4_AVX)); \\\n __m256i ymm_rhs_1 = _mm256_shuffle_epi8( \\\n INT4_LOOKUP_AVX, \\\n _mm256_and_si256(_mm256_srli_epi32((ymm_rhs), 4), MASK_INT4_AVX)); \\\n FMA_INT8_AVX(ymm_lhs_0, ymm_rhs_0, ymm_sum_0); \\\n FMA_INT8_AVX(ymm_lhs_1, ymm_rhs_1, ymm_sum_1); \\\n FMA_INT8_AVX(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1); \\\n FMA_INT8_AVX(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1); \\\n FMA_INT8_AVX(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2); \\\n FMA_INT8_AVX(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2); \\\n }\n\n" }, { "path": "src/ailego/math/distance_matrix_popcnt.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_int32.i\"\n#include \"distance_matrix_int64.i\"\n#include \"matrix_utility.i\"\n\n//! Calculate population count (UINT32 Permute 1 SSE)\n#define POPCNT_UINT32_PERMUTE1_SSE(v, ...) \\\n _mm_add_epi16(_mm_srli_epi16(v, 8), _mm_and_si128(v, _mm_set1_epi16(0xff)))\n\n//! Calculate population count (UINT32 Permute 2 SSE)\n#define POPCNT_UINT32_PERMUTE2_SSE(v, ...) \\\n _mm_add_epi32(_mm_srli_epi32(v, 16), _mm_and_si128(v, _mm_set1_epi32(0xffff)))\n\n//! Calculate population count (UINT32 Permute 1 AVX)\n#define POPCNT_UINT32_PERMUTE1_AVX(v, ...) \\\n _mm256_add_epi16(_mm256_srli_epi16(v, 8), \\\n _mm256_and_si256(v, _mm256_set1_epi16(0xff)))\n\n//! Calculate population count (UINT32 Permute 2 AVX)\n#define POPCNT_UINT32_PERMUTE2_AVX(v, ...) \\\n _mm256_add_epi32(_mm256_srli_epi32(v, 16), \\\n _mm256_and_si256(v, _mm256_set1_epi32(0xffff)))\n\n//! Calculate population count (UINT64 Permute AVX)\n#define POPCNT_UINT64_PERMUTE_AVX(v, ...) \\\n _mm256_sad_epu8(v, _mm256_setzero_si256())\n\n//! Compute the distance between matrix and query (UINT32, M=2, N=1)\n#define POPCNT_UINT32_2X1_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + ((cnt >> 2) << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 2) << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 2) << 2) : qe_0); \\\n const uint32_t *qe_3 = q + cnt; \\\n if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n if (qe_3 >= qe_0 + 2) { \\\n __m128i xmm_m = _mm_load_si128((const __m128i *)(m)); \\\n __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n m += 4; \\\n q += 2; \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n if (qe_3 >= qe_0 + 2) { \\\n __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \\\n __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n m += 4; \\\n q += 2; \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_0_1); \\\n xmm_sum_0_0 = _mm_add_epi32( \\\n xmm_sum_0_0, _mm_shuffle_epi32(xmm_sum_0_0, _MM_SHUFFLE(0, 0, 3, 2))); \\\n if (q != qe_3) { \\\n __m128i xmm_m = _mm_set_epi32(0, 0, m[1], m[0]); \\\n __m128i xmm_q = _mm_broadcast_si32(q); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0_0));\n\n//! Compute the distance between matrix and query (UINT32, M=2, N=2)\n#define POPCNT_UINT32_2X2_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + ((cnt >> 1) << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 2) : qe_0); \\\n const uint32_t *qe_3 = q + (cnt << 1); \\\n if (((uintptr_t)m & 0xf) == 0 && ((uintptr_t)q & 0xf) == 0) { \\\n for (; q != qe_1; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_2X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \\\n _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \\\n if (q != qe_3) { \\\n __m128i xmm_m = _mm_set_epi32(m[1], m[0], m[1], m[0]); \\\n __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=4, N=1)\n#define POPCNT_UINT32_4X1_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + ((cnt >> 1) << 1); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 1) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 1) : qe_0); \\\n const uint32_t *qe_3 = q + cnt; \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n if (q != qe_3) { \\\n __m128i xmm_m = _mm_load_si128((const __m128i *)(m)); \\\n __m128i xmm_q = _mm_broadcast_si32(q); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n if (q != qe_3) { \\\n __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \\\n __m128i xmm_q = _mm_broadcast_si32(q); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(xmm_sum_0_0, xmm_sum_1_0); \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=4, N=2)\n#define POPCNT_UINT32_4X2_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 1); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 4, q += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 4, q += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 4, q += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 4, q += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 4, q += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 4, q += 2) { \\\n MATRIX_INT32_ITER_4X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=4, N=4)\n#define POPCNT_UINT32_4X4_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 4, q += 4) { \\\n MATRIX_INT32_ITER_4X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=8, N=1)\n#define POPCNT_UINT32_8X1_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + cnt; \\\n const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=8, N=2)\n#define POPCNT_UINT32_8X2_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 1); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=8, N=4)\n#define POPCNT_UINT32_8X4_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=8, N=8)\n#define POPCNT_UINT32_8X8_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 3); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=1)\n#define POPCNT_UINT32_16X1_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + cnt; \\\n const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=2)\n#define POPCNT_UINT32_16X2_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 1); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=4)\n#define POPCNT_UINT32_16X4_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=8)\n#define POPCNT_UINT32_16X8_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 3); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=16)\n#define POPCNT_UINT32_16X16_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 4); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 4) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=1)\n#define POPCNT_UINT32_32X1_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 1, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + cnt; \\\n const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 1, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 1, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=2)\n#define POPCNT_UINT32_32X2_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 2, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 1); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 2, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 2, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=4)\n#define POPCNT_UINT32_32X4_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 4, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 4, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 4, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=8)\n#define POPCNT_UINT32_32X8_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 8, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 3); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 8, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 8, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 8, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=16)\n#define POPCNT_UINT32_32X16_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 16, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 4); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 4) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 16, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 16, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 16, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 16, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 16, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=32)\n#define POPCNT_UINT32_32X32_SSE(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 32, __m128i, xmm_sum, _mm_setzero_si128()) \\\n const uint32_t *qe_0 = q + (cnt << 5); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 5) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 5) : qe_0); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; q != qe_1; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 32, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 32, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_load_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP1_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 32, xmm_sum, POPCNT_UINT32_PERMUTE1_SSE) \\\n for (; q != qe_2; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP2_SSE) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 32, xmm_sum, POPCNT_UINT32_PERMUTE2_SSE) \\\n for (; q != qe_0; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_SSE(m, q, xmm_sum, _mm_loadu_si128, \\\n POPCNT_UINT32_STEP3_SSE) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 32, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=2, N=1)\n#define POPCNT_UINT32_2X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + ((cnt >> 2) << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 2) << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 2) << 2) : qe_0); \\\n const uint32_t *qe_3 = q + cnt; \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0 = _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n if (qe_3 >= qe_0 + 2) { \\\n __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \\\n __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0) \\\n m += 4; \\\n q += 2; \\\n } \\\n xmm_sum_0 = _mm_add_epi32( \\\n xmm_sum_0, _mm_shuffle_epi32(xmm_sum_0, _MM_SHUFFLE(0, 0, 3, 2))); \\\n if (q != qe_3) { \\\n __m128i xmm_m = _mm_set_epi32(0, 0, m[1], m[0]); \\\n __m128i xmm_q = _mm_broadcast_si32(q); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(xmm_sum_0));\n\n//! Compute the distance between matrix and query (UINT32, M=2, N=2)\n#define POPCNT_UINT32_2X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + ((cnt >> 2) << 3); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 2) << 3) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 2) << 3) : qe_0); \\\n const uint32_t *qe_3 = q + (cnt << 1); \\\n if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n __m128i xmm_sum_0_1 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \\\n _mm256_extracti128_si256(ymm_sum_0_1, 1)); \\\n if (qe_3 >= qe_0 + 4) { \\\n __m128i xmm_q = _mm_loadu_si128((const __m128i *)(q)); \\\n __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \\\n __m128i xmm_p = _mm_shuffle_epi32(xmm_q, _MM_SHUFFLE(2, 2, 0, 0)); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_p, xmm_sum_0_0) \\\n xmm_p = _mm_shuffle_epi32(xmm_q, _MM_SHUFFLE(3, 3, 1, 1)); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_p, xmm_sum_0_1) \\\n m += 4; \\\n q += 4; \\\n } \\\n xmm_sum_0_0 = _mm_add_epi32(_mm_unpacklo_epi64(xmm_sum_0_0, xmm_sum_0_1), \\\n _mm_unpackhi_epi64(xmm_sum_0_0, xmm_sum_0_1)); \\\n if (q != qe_3) { \\\n __m128i xmm_m = _mm_set_epi32(m[1], m[0], m[1], m[0]); \\\n __m128i xmm_q = _mm_set_epi32(q[1], q[1], q[0], q[0]); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=4, N=1)\n#define POPCNT_UINT32_4X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + ((cnt >> 1) << 1); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 1) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 1) : qe_0); \\\n const uint32_t *qe_3 = q + cnt; \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n if (q != qe_3) { \\\n __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \\\n __m128i xmm_q = _mm_broadcast_si32(q); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=4, N=2)\n#define POPCNT_UINT32_4X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + ((cnt >> 1) << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 2) : qe_0); \\\n const uint32_t *qe_3 = q + (cnt << 1); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n __m128i xmm_sum_0_1 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \\\n _mm256_extracti128_si256(ymm_sum_0_1, 1)); \\\n if (q != qe_3) { \\\n __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \\\n __m128i xmm_q = _mm_broadcast_si32(q); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n xmm_q = _mm_broadcast_si32(q + 1); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_1) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=4, N=4)\n#define POPCNT_UINT32_4X4_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + ((cnt >> 1) << 3); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 3) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + ((4095 >> 1) << 3) : qe_0); \\\n const uint32_t *qe_3 = q + (cnt << 2); \\\n if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n __m128i xmm_sum_0_0 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_0), \\\n _mm256_extracti128_si256(ymm_sum_0_0, 1)); \\\n __m128i xmm_sum_0_1 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_1), \\\n _mm256_extracti128_si256(ymm_sum_0_1, 1)); \\\n __m128i xmm_sum_0_2 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_2), \\\n _mm256_extracti128_si256(ymm_sum_0_2, 1)); \\\n __m128i xmm_sum_0_3 = \\\n _mm_add_epi32(_mm256_castsi256_si128(ymm_sum_0_3), \\\n _mm256_extracti128_si256(ymm_sum_0_3, 1)); \\\n if (q != qe_3) { \\\n __m128i xmm_m = _mm_loadu_si128((const __m128i *)(m)); \\\n __m128i xmm_q = _mm_broadcast_si32(q); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_0) \\\n xmm_q = _mm_broadcast_si32(q + 1); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_1) \\\n xmm_q = _mm_broadcast_si32(q + 2); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_2) \\\n xmm_q = _mm_broadcast_si32(q + 3); \\\n POPCNT_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum_0_3) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, xmm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=8, N=1)\n#define POPCNT_UINT32_8X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + cnt; \\\n const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, ++q) { \\\n MATRIX_INT32_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=8, N=2)\n#define POPCNT_UINT32_8X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 1); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT32_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=8, N=4)\n#define POPCNT_UINT32_8X4_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT32_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=8, N=8)\n#define POPCNT_UINT32_8X8_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 3); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT32_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=1)\n#define POPCNT_UINT32_16X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + cnt; \\\n const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, ++q) { \\\n MATRIX_INT32_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=2)\n#define POPCNT_UINT32_16X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 1); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, q += 2) { \\\n MATRIX_INT32_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=4)\n#define POPCNT_UINT32_16X4_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, q += 4) { \\\n MATRIX_INT32_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=8)\n#define POPCNT_UINT32_16X8_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 3); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, q += 8) { \\\n MATRIX_INT32_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=16, N=16)\n#define POPCNT_UINT32_16X16_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 4); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 4) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 16, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 16, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 16, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 16, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 16, q += 16) { \\\n MATRIX_INT32_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=1)\n#define POPCNT_UINT32_32X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + cnt; \\\n const uint32_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + 4095 : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, ++q) { \\\n MATRIX_INT32_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=2)\n#define POPCNT_UINT32_32X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 1); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 1) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 2) { \\\n MATRIX_INT32_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=4)\n#define POPCNT_UINT32_32X4_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 2); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 2) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 4) { \\\n MATRIX_INT32_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=8)\n#define POPCNT_UINT32_32X8_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 3); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 3) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 8) { \\\n MATRIX_INT32_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=16)\n#define POPCNT_UINT32_32X16_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 4); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 4) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 16) { \\\n MATRIX_INT32_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT32, M=32, N=32)\n#define POPCNT_UINT32_32X32_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 32, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint32_t *qe_0 = q + (cnt << 5); \\\n const uint32_t *qe_1 = (cnt > 31 ? q + (31 << 5) : qe_0); \\\n const uint32_t *qe_2 = (cnt > 4095 ? q + (4095 << 5) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 32, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 32, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 32, ymm_sum, POPCNT_UINT32_PERMUTE1_AVX) \\\n for (; q != qe_2; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP2_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 32, ymm_sum, POPCNT_UINT32_PERMUTE2_AVX) \\\n for (; q != qe_0; m += 32, q += 32) { \\\n MATRIX_INT32_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT32_STEP3_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0x1f) == 0) { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 32, 8, ymm_sum, out, _mm256_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=2, N=1)\n#define POPCNT_UINT64_2X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + ((cnt >> 2) << 2); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + ((31 >> 2) << 2) : qe_0); \\\n const uint64_t *qe_2 = q + cnt; \\\n if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT64_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT64_ITER_2X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n if (qe_2 >= qe_0 + 2) { \\\n __m256i ymm_m = _mm256_load_si256((const __m256i *)(m)); \\\n __m256i ymm_q = _mm256_set_epi64x(q[1], q[1], q[0], q[0]); \\\n POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \\\n m += 4; \\\n q += 2; \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT64_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT64_ITER_2X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n if (qe_2 >= qe_0 + 2) { \\\n __m256i ymm_m = _mm256_loadu_si256((const __m256i *)(m)); \\\n __m256i ymm_q = _mm256_set_epi64x(q[1], q[1], q[0], q[0]); \\\n POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \\\n m += 4; \\\n q += 2; \\\n } \\\n } \\\n ymm_sum_0_0 = _mm256_add_epi64(ymm_sum_0_0, ymm_sum_0_1); \\\n ymm_sum_0_0 = _mm256_add_epi64( \\\n ymm_sum_0_0, \\\n _mm256_permute4x64_epi64(ymm_sum_0_0, _MM_SHUFFLE(0, 0, 3, 2))); \\\n if (q != qe_2) { \\\n __m256i ymm_m = _mm256_set_epi64x(0, 0, m[1], m[0]); \\\n __m256i ymm_q = _mm256_broadcast_si64(q); \\\n POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \\\n } \\\n _mm_storel_pi((__m64 *)out, _NORM(ymm_sum_0_0));\n\n//! Compute the distance between matrix and query (UINT64, M=2, N=2)\n#define POPCNT_UINT64_2X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + ((cnt >> 1) << 2); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 2) : qe_0); \\\n const uint64_t *qe_2 = q + (cnt << 1); \\\n if (((uintptr_t)m & 0x1f) == 0 && ((uintptr_t)q & 0x1f) == 0) { \\\n for (; q != qe_1; m += 4, q += 4) { \\\n MATRIX_INT64_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 4, q += 4) { \\\n MATRIX_INT64_ITER_2X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 4, q += 4) { \\\n MATRIX_INT64_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 4, q += 4) { \\\n MATRIX_INT64_ITER_2X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n ymm_sum_0_0 = _mm256_add_epi64( \\\n _mm256_inserti128_si256(ymm_sum_0_0, \\\n _mm256_castsi256_si128(ymm_sum_0_1), 1), \\\n _mm256_inserti128_si256(ymm_sum_0_1, \\\n _mm256_extractf128_si256(ymm_sum_0_0, 1), 0)); \\\n if (q != qe_2) { \\\n __m256i ymm_m = _mm256_set_epi64x(m[1], m[0], m[1], m[0]); \\\n __m256i ymm_q = _mm256_set_epi64x(q[1], q[1], q[0], q[0]); \\\n POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=4, N=1)\n#define POPCNT_UINT64_4X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + ((cnt >> 1) << 1); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + ((31 >> 1) << 1) : qe_0); \\\n const uint64_t *qe_2 = q + cnt; \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT64_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT64_ITER_4X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n if (q != qe_2) { \\\n __m256i ymm_m = _mm256_load_si256((const __m256i *)(m)); \\\n __m256i ymm_q = _mm256_broadcast_si64(q); \\\n POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT64_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT64_ITER_4X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n if (q != qe_2) { \\\n __m256i ymm_m = _mm256_loadu_si256((const __m256i *)(m)); \\\n __m256i ymm_q = _mm256_broadcast_si64(q); \\\n POPCNT_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum_0_0) \\\n } \\\n } \\\n ymm_sum_0_0 = _mm256_add_epi64(ymm_sum_0_0, ymm_sum_1_0); \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=4, N=2)\n#define POPCNT_UINT64_4X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 1); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 4, q += 2) { \\\n MATRIX_INT64_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 4, q += 2) { \\\n MATRIX_INT64_ITER_4X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 4, q += 2) { \\\n MATRIX_INT64_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 4, q += 2) { \\\n MATRIX_INT64_ITER_4X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 2, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 2, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=4, N=4)\n#define POPCNT_UINT64_4X4_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(1, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 2); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 4, q += 4) { \\\n MATRIX_INT64_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 4, q += 4) { \\\n MATRIX_INT64_ITER_4X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 4, q += 4) { \\\n MATRIX_INT64_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(1, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 4, q += 4) { \\\n MATRIX_INT64_ITER_4X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(1, 4, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(1, 4, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=8, N=1)\n#define POPCNT_UINT64_8X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + cnt; \\\n const uint64_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, ++q) { \\\n MATRIX_INT64_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, ++q) { \\\n MATRIX_INT64_ITER_8X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, ++q) { \\\n MATRIX_INT64_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, ++q) { \\\n MATRIX_INT64_ITER_8X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=8, N=2)\n#define POPCNT_UINT64_8X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 1); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT64_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT64_ITER_8X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 2) { \\\n MATRIX_INT64_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 2) { \\\n MATRIX_INT64_ITER_8X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 2, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 2, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=8, N=4)\n#define POPCNT_UINT64_8X4_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 2); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT64_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT64_ITER_8X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 4) { \\\n MATRIX_INT64_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 4) { \\\n MATRIX_INT64_ITER_8X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 4, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 4, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=8, N=8)\n#define POPCNT_UINT64_8X8_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(2, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 3); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT64_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT64_ITER_8X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 8, q += 8) { \\\n MATRIX_INT64_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(2, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 8, q += 8) { \\\n MATRIX_INT64_ITER_8X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(2, 8, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(2, 8, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=16, N=1)\n#define POPCNT_UINT64_16X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + cnt; \\\n const uint64_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, ++q) { \\\n MATRIX_INT64_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, ++q) { \\\n MATRIX_INT64_ITER_16X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, ++q) { \\\n MATRIX_INT64_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, ++q) { \\\n MATRIX_INT64_ITER_16X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=16, N=2)\n#define POPCNT_UINT64_16X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 1); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, q += 2) { \\\n MATRIX_INT64_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, q += 2) { \\\n MATRIX_INT64_ITER_16X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 2) { \\\n MATRIX_INT64_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, q += 2) { \\\n MATRIX_INT64_ITER_16X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 2, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 2, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=16, N=4)\n#define POPCNT_UINT64_16X4_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 2); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, q += 4) { \\\n MATRIX_INT64_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, q += 4) { \\\n MATRIX_INT64_ITER_16X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 4) { \\\n MATRIX_INT64_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, q += 4) { \\\n MATRIX_INT64_ITER_16X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 4, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 4, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=16, N=8)\n#define POPCNT_UINT64_16X8_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 3); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, q += 8) { \\\n MATRIX_INT64_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, q += 8) { \\\n MATRIX_INT64_ITER_16X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 8) { \\\n MATRIX_INT64_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, q += 8) { \\\n MATRIX_INT64_ITER_16X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 8, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 8, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=16, N=16)\n#define POPCNT_UINT64_16X16_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(4, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 4); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 16, q += 16) { \\\n MATRIX_INT64_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, q += 16) { \\\n MATRIX_INT64_ITER_16X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 16, q += 16) { \\\n MATRIX_INT64_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(4, 16, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 16, q += 16) { \\\n MATRIX_INT64_ITER_16X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(4, 16, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(4, 16, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=32, N=1)\n#define POPCNT_UINT64_32X1_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 1, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + cnt; \\\n const uint64_t *qe_1 = (cnt > 31 ? q + 31 : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, ++q) { \\\n MATRIX_INT64_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, ++q) { \\\n MATRIX_INT64_ITER_32X1_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, ++q) { \\\n MATRIX_INT64_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 1, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, ++q) { \\\n MATRIX_INT64_ITER_32X1_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 1, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 1, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=32, N=2)\n#define POPCNT_UINT64_32X2_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 2, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 1); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 1) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 2) { \\\n MATRIX_INT64_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 2) { \\\n MATRIX_INT64_ITER_32X2_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 2) { \\\n MATRIX_INT64_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 2, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 2) { \\\n MATRIX_INT64_ITER_32X2_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 2, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 2, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=32, N=4)\n#define POPCNT_UINT64_32X4_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 4, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 2); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 2) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 4) { \\\n MATRIX_INT64_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 4) { \\\n MATRIX_INT64_ITER_32X4_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 4) { \\\n MATRIX_INT64_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 4, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 4) { \\\n MATRIX_INT64_ITER_32X4_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 4, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 4, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=32, N=8)\n#define POPCNT_UINT64_32X8_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 8, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 3); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 3) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 8) { \\\n MATRIX_INT64_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 8) { \\\n MATRIX_INT64_ITER_32X8_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 8) { \\\n MATRIX_INT64_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 8, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 8) { \\\n MATRIX_INT64_ITER_32X8_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 8, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 8, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=32, N=16)\n#define POPCNT_UINT64_32X16_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 16, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 4); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 4) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 16) { \\\n MATRIX_INT64_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 16, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 16) { \\\n MATRIX_INT64_ITER_32X16_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 16) { \\\n MATRIX_INT64_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 16, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 16) { \\\n MATRIX_INT64_ITER_32X16_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 16, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 16, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n\n//! Compute the distance between matrix and query (UINT64, M=32, N=32)\n#define POPCNT_UINT64_32X32_AVX(m, q, cnt, out, _NORM) \\\n MATRIX_VAR_INIT(8, 32, __m256i, ymm_sum, _mm256_setzero_si256()) \\\n const uint64_t *qe_0 = q + (cnt << 5); \\\n const uint64_t *qe_1 = (cnt > 31 ? q + (31 << 5) : qe_0); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; q != qe_1; m += 32, q += 32) { \\\n MATRIX_INT64_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 32, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 32) { \\\n MATRIX_INT64_ITER_32X32_AVX(m, q, ymm_sum, _mm256_load_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } else { \\\n for (; q != qe_1; m += 32, q += 32) { \\\n MATRIX_INT64_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP1_AVX) \\\n } \\\n MATRIX_VAR_PERMUTE(8, 32, ymm_sum, POPCNT_UINT64_PERMUTE_AVX) \\\n for (; q != qe_0; m += 32, q += 32) { \\\n MATRIX_INT64_ITER_32X32_AVX(m, q, ymm_sum, _mm256_loadu_si256, \\\n POPCNT_UINT64_STEP2_AVX) \\\n } \\\n } \\\n if (((uintptr_t)out & 0xf) == 0) { \\\n MATRIX_VAR_STORE(8, 32, 4, ymm_sum, out, _mm_store_ps, _NORM) \\\n } else { \\\n MATRIX_VAR_STORE(8, 32, 4, ymm_sum, out, _mm_storeu_ps, _NORM) \\\n }\n" }, { "path": "src/ailego/math/distance_utility.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Four-bits Squared Difference Table\n */\nstatic const AILEGO_ALIGNED(64) uint8_t Int4SquaredDiffTable[256] = {\n 0, 1, 4, 9, 16, 25, 36, 49, 64, 49, 36, 25, 16, 9, 4, 1,\n 1, 0, 1, 4, 9, 16, 25, 36, 81, 64, 49, 36, 25, 16, 9, 4,\n 4, 1, 0, 1, 4, 9, 16, 25, 100, 81, 64, 49, 36, 25, 16, 9,\n 9, 4, 1, 0, 1, 4, 9, 16, 121, 100, 81, 64, 49, 36, 25, 16,\n 16, 9, 4, 1, 0, 1, 4, 9, 144, 121, 100, 81, 64, 49, 36, 25,\n 25, 16, 9, 4, 1, 0, 1, 4, 169, 144, 121, 100, 81, 64, 49, 36,\n 36, 25, 16, 9, 4, 1, 0, 1, 196, 169, 144, 121, 100, 81, 64, 49,\n 49, 36, 25, 16, 9, 4, 1, 0, 225, 196, 169, 144, 121, 100, 81, 64,\n 64, 81, 100, 121, 144, 169, 196, 225, 0, 1, 4, 9, 16, 25, 36, 49,\n 49, 64, 81, 100, 121, 144, 169, 196, 1, 0, 1, 4, 9, 16, 25, 36,\n 36, 49, 64, 81, 100, 121, 144, 169, 4, 1, 0, 1, 4, 9, 16, 25,\n 25, 36, 49, 64, 81, 100, 121, 144, 9, 4, 1, 0, 1, 4, 9, 16,\n 16, 25, 36, 49, 64, 81, 100, 121, 16, 9, 4, 1, 0, 1, 4, 9,\n 9, 16, 25, 36, 49, 64, 81, 100, 25, 16, 9, 4, 1, 0, 1, 4,\n 4, 9, 16, 25, 36, 49, 64, 81, 36, 25, 16, 9, 4, 1, 0, 1,\n 1, 4, 9, 16, 25, 36, 49, 64, 49, 36, 25, 16, 9, 4, 1, 0,\n};\n\n/*! Four-bits Integer Multiplication Table\n */\nstatic const AILEGO_ALIGNED(64) int8_t Int4MulTable[256] = {\n 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1,\n 0, 2, 4, 6, 8, 10, 12, 14, -16, -14, -12, -10, -8, -6, -4, -2,\n 0, 3, 6, 9, 12, 15, 18, 21, -24, -21, -18, -15, -12, -9, -6, -3,\n 0, 4, 8, 12, 16, 20, 24, 28, -32, -28, -24, -20, -16, -12, -8, -4,\n 0, 5, 10, 15, 20, 25, 30, 35, -40, -35, -30, -25, -20, -15, -10, -5,\n 0, 6, 12, 18, 24, 30, 36, 42, -48, -42, -36, -30, -24, -18, -12, -6,\n 0, 7, 14, 21, 28, 35, 42, 49, -56, -49, -42, -35, -28, -21, -14, -7,\n 0, -8, -16, -24, -32, -40, -48, -56, 64, 56, 48, 40, 32, 24, 16, 8,\n 0, -7, -14, -21, -28, -35, -42, -49, 56, 49, 42, 35, 28, 21, 14, 7,\n 0, -6, -12, -18, -24, -30, -36, -42, 48, 42, 36, 30, 24, 18, 12, 6,\n 0, -5, -10, -15, -20, -25, -30, -35, 40, 35, 30, 25, 20, 15, 10, 5,\n 0, -4, -8, -12, -16, -20, -24, -28, 32, 28, 24, 20, 16, 12, 8, 4,\n 0, -3, -6, -9, -12, -15, -18, -21, 24, 21, 18, 15, 12, 9, 6, 3,\n 0, -2, -4, -6, -8, -10, -12, -14, 16, 14, 12, 10, 8, 6, 4, 2,\n 0, -1, -2, -3, -4, -5, -6, -7, 8, 7, 6, 5, 4, 3, 2, 1,\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/euclidean_distance_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \"distance_utility.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n/*! Squared Euclidean Distance Matrix\n */\ntemplate \nstruct SquaredEuclideanDistanceMatrix;\n\n/*! Squared Euclidean Distance Matrix (M=1, N=1)\n */\ntemplate \nstruct SquaredEuclideanDistanceMatrix<\n T, 1, 1, typename std::enable_if::value>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n float sum = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n sum += MathHelper::SquaredDifference(m[i], q[i]);\n }\n *out = sum;\n }\n};\n\ntemplate <>\nstruct SquaredEuclideanDistanceMatrix {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct SquaredEuclideanDistanceMatrix {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct SquaredEuclideanDistanceMatrix {\n //! Type of value\n using ValueType = Float16;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct SquaredEuclideanDistanceMatrix {\n //! Type of value\n using ValueType = float;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\n/*! Squared Euclidean Distance Matrix\n */\ntemplate \nstruct SquaredEuclideanDistanceMatrix<\n T, M, N,\n typename std::enable_if::value && sizeof(T) >= 2 &&\n M >= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n if (dim > 0) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = MathHelper::SquaredDifference(m_val, q[j]);\n r += M;\n }\n }\n m += M;\n q += N;\n }\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r += MathHelper::SquaredDifference(m_val, q[j]);\n r += M;\n }\n }\n m += M;\n q += N;\n }\n }\n};\n\n/*! Squared Euclidean Distance Matrix (N=1)\n */\ntemplate \nstruct SquaredEuclideanDistanceMatrix<\n T, M, 1,\n typename std::enable_if::value && sizeof(T) >= 2 &&\n M >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n const ValueType *q_end = q + dim;\n if (q != q_end) {\n ValueType q_val = *q++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = MathHelper::SquaredDifference(m[i], q_val);\n }\n m += M;\n }\n\n while (q != q_end) {\n ValueType q_val = *q++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += MathHelper::SquaredDifference(m[i], q_val);\n }\n m += M;\n }\n }\n};\n\n/*! Squared Euclidean Distance Matrix (INT8)\n */\ntemplate \nstruct SquaredEuclideanDistanceMatrix<\n int8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 3) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n\n dim >>= 2;\n if (dim > 0) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = SquaredDifference(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r += SquaredDifference(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n }\n\n protected:\n //! Calculate the squared difference\n static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) {\n volatile int32_t sum = MathHelper::SquaredDifference(\n (int8_t)(lhs >> 0), (int8_t)(rhs >> 0)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 8), (int8_t)(rhs >> 8)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 16), (int8_t)(rhs >> 16)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 24), (int8_t)(rhs >> 24));\n return static_cast(sum);\n }\n};\n\n/*! Squared Euclidean Distance Matrix (INT8, N=1)\n */\ntemplate \nstruct SquaredEuclideanDistanceMatrix= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 3) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 2);\n\n if (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = SquaredDifference(m_it[i], q_val);\n }\n m_it += M;\n }\n\n while (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += SquaredDifference(m_it[i], q_val);\n }\n m_it += M;\n }\n }\n\n protected:\n //! Calculate the squared difference\n static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) {\n volatile int32_t sum = MathHelper::SquaredDifference(\n (int8_t)(lhs >> 0), (int8_t)(rhs >> 0)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 8), (int8_t)(rhs >> 8)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 16), (int8_t)(rhs >> 16)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 24), (int8_t)(rhs >> 24));\n return static_cast(sum);\n }\n};\n\n/*! Squared Euclidean Distance Matrix (INT4)\n */\ntemplate \nstruct SquaredEuclideanDistanceMatrix<\n uint8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 7) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n\n dim >>= 3;\n if (dim > 0) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = SquaredDifference(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r += SquaredDifference(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n }\n\n protected:\n //! Calculate the squared difference\n static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4SquaredDiffTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n};\n\n/*! Squared Euclidean Distance Matrix (INT4, N=1)\n */\ntemplate \nstruct SquaredEuclideanDistanceMatrix= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 7) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 3);\n\n if (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = SquaredDifference(m_it[i], q_val);\n }\n m_it += M;\n }\n\n while (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += SquaredDifference(m_it[i], q_val);\n }\n m_it += M;\n }\n }\n\n protected:\n //! Calculate the squared difference\n static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4SquaredDiffTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n};\n\n/*! Euclidean Distance Matrix\n */\ntemplate ::value ||\n std::is_same::value) &&\n M >= 1 && N >= 1>::type>\nstruct EuclideanDistanceMatrix {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out);\n for (size_t i = 0; i < N * M; ++i) {\n float val = *out;\n *out++ = std::sqrt(val);\n }\n }\n};\n\n/*! Euclidean Distance Matrix (M=1, N=1)\n */\ntemplate \nstruct EuclideanDistanceMatrix<\n T, 1, 1, typename std::enable_if::value>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n float sum = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n sum += MathHelper::SquaredDifference(m[i], q[i]);\n }\n *out = std::sqrt(sum);\n }\n};\n\ntemplate <>\nstruct EuclideanDistanceMatrix {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct EuclideanDistanceMatrix {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct EuclideanDistanceMatrix {\n //! Type of value\n using ValueType = Float16;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct EuclideanDistanceMatrix {\n //! Type of value\n using ValueType = float;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\n\n//--------------------------------------------------\n// Sparse\n//--------------------------------------------------\n/*! Squared Euclidean Distance Sparse Matrix\n */\ntemplate \nstruct SquaredEuclideanSparseDistanceMatrix {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n static float ComputeSquaredEuclideanSparseDistanceInSegment(\n uint32_t m_sparse_count, const uint16_t *m_sparse_index,\n const ValueType *m_sparse_value, uint32_t q_sparse_count,\n const uint16_t *q_sparse_index, const ValueType *q_sparse_value);\n\n //! Compute the distance between matrix and query\n static inline void Compute(const void *m_sparse_data_in,\n const void *q_sparse_data_in, float *out) {\n ailego_assert(out);\n\n const uint8_t *m_sparse_data =\n reinterpret_cast(m_sparse_data_in);\n const uint8_t *q_sparse_data =\n reinterpret_cast(q_sparse_data_in);\n\n const uint32_t m_sparse_count =\n *reinterpret_cast(m_sparse_data);\n const uint32_t q_sparse_count =\n *reinterpret_cast(q_sparse_data);\n\n const uint32_t m_seg_count =\n *reinterpret_cast(m_sparse_data + sizeof(uint32_t));\n const uint32_t q_seg_count =\n *reinterpret_cast(q_sparse_data + sizeof(uint32_t));\n\n const uint32_t *m_seg_id = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t));\n const uint32_t *q_seg_id = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t));\n\n const uint32_t *m_seg_vec_cnt = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t) + m_seg_count * sizeof(uint32_t));\n const uint32_t *q_seg_vec_cnt = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t) + q_seg_count * sizeof(uint32_t));\n\n const uint16_t *m_sparse_index = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t) +\n m_seg_count * 2 * sizeof(uint32_t));\n const uint16_t *q_sparse_index = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t) +\n q_seg_count * 2 * sizeof(uint32_t));\n\n const ValueType *m_sparse_value = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t) +\n m_seg_count * 2 * sizeof(uint32_t) + m_sparse_count * sizeof(uint16_t));\n const ValueType *q_sparse_value = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t) +\n q_seg_count * 2 * sizeof(uint32_t) + q_sparse_count * sizeof(uint16_t));\n\n float sum = 0.0f;\n\n size_t m_s = 0;\n size_t q_s = 0;\n\n size_t m_count = 0;\n size_t q_count = 0;\n\n while (m_s < m_seg_count && q_s < q_seg_count) {\n if (m_seg_id[m_s] == q_seg_id[q_s]) {\n sum += ComputeSquaredEuclideanSparseDistanceInSegment(\n m_seg_vec_cnt[m_s], m_sparse_index + m_count,\n m_sparse_value + m_count, q_seg_vec_cnt[q_s],\n q_sparse_index + q_count, q_sparse_value + q_count);\n\n m_count += m_seg_vec_cnt[m_s];\n q_count += q_seg_vec_cnt[q_s];\n\n ++m_s;\n ++q_s;\n } else if (m_seg_id[m_s] < q_seg_id[q_s]) {\n for (size_t i = 0; i < m_seg_vec_cnt[m_s]; i++) {\n float value = (m_sparse_value + m_count)[i];\n sum += value * value;\n }\n\n m_count += m_seg_vec_cnt[m_s];\n\n ++m_s;\n } else {\n for (size_t i = 0; i < q_seg_vec_cnt[q_s]; i++) {\n float value = (q_sparse_value + q_count)[i];\n sum += value * value;\n }\n\n q_count += q_seg_vec_cnt[q_s];\n ++q_s;\n }\n }\n\n for (; m_s < m_seg_count; m_s++) {\n for (size_t i = 0; i < m_seg_vec_cnt[m_s]; i++) {\n float diff = (m_sparse_value + m_count)[i];\n sum += diff * diff;\n }\n\n m_count += m_seg_vec_cnt[m_s];\n }\n\n for (; q_s < q_seg_count; q_s++) {\n for (size_t i = 0; i < q_seg_vec_cnt[q_s]; i++) {\n float diff = (q_sparse_value + q_count)[i];\n sum += diff * diff;\n }\n\n q_count += q_seg_vec_cnt[q_s];\n }\n\n *out = sum;\n }\n};\n\ntemplate \nfloat SquaredEuclideanSparseDistanceMatrix::\n ComputeSquaredEuclideanSparseDistanceInSegment(\n uint32_t m_sparse_count, const uint16_t *m_sparse_index,\n const ValueType *m_sparse_value, uint32_t q_sparse_count,\n const uint16_t *q_sparse_index, const ValueType *q_sparse_value) {\n float sum = 0.0f;\n\n size_t m_i = 0;\n size_t q_i = 0;\n\n while (m_i < m_sparse_count && q_i < q_sparse_count) {\n if (m_sparse_index[m_i] == q_sparse_index[q_i]) {\n float diff = m_sparse_value[m_i] - q_sparse_value[q_i];\n sum += diff * diff;\n ++m_i;\n ++q_i;\n } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) {\n float diff = m_sparse_value[m_i];\n sum += diff * diff;\n ++m_i;\n } else {\n float diff = q_sparse_value[q_i];\n sum += diff * diff;\n\n ++q_i;\n }\n }\n\n for (; m_i < m_sparse_count; m_i++) {\n float diff = m_sparse_value[m_i];\n sum += diff * diff;\n }\n\n for (; q_i < q_sparse_count; q_i++) {\n float diff = q_sparse_value[q_i];\n sum += diff * diff;\n }\n\n return sum;\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp16_avx.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX__)\n\nfloat SquaredEuclideanDistanceFp16AVX(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n float score{0.0f};\n\n ACCUM_FP16_1X1_AVX(lhs, rhs, size, &score, 0ull, )\n\n return score;\n}\n\n#endif // __AVX__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp16_avx512.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX512F__)\nfloat SquaredEuclideanDistanceFp16AVX512(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n float score{0.0f};\n\n ACCUM_FP16_1X1_AVX512(lhs, rhs, size, &score, 0ull, )\n\n return score;\n}\n#endif\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp16_avx512fp16.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX512FP16__)\n//! Squared Euclidean Distance\nfloat SquaredEuclideanDistanceFp16AVX512FP16(const Float16 *lhs,\n const Float16 *rhs, size_t size) {\n const Float16 *last = lhs + size;\n const Float16 *last_aligned = lhs + ((size >> 6) << 6);\n\n __m512h zmm_sum_0 = _mm512_setzero_ph();\n __m512h zmm_sum_1 = _mm512_setzero_ph();\n\n if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n __m512h zmm_d_0 =\n _mm512_sub_ph(_mm512_load_ph(lhs + 0), _mm512_load_ph(rhs + 0));\n __m512h zmm_d_1 =\n _mm512_sub_ph(_mm512_load_ph(lhs + 32), _mm512_load_ph(rhs + 32));\n zmm_sum_0 = _mm512_fmadd_ph(zmm_d_0, zmm_d_0, zmm_sum_0);\n zmm_sum_1 = _mm512_fmadd_ph(zmm_d_1, zmm_d_1, zmm_sum_1);\n }\n\n if (last >= last_aligned + 32) {\n __m512h zmm_d = _mm512_sub_ph(_mm512_load_ph(lhs), _mm512_load_ph(rhs));\n zmm_sum_0 = _mm512_fmadd_ph(zmm_d, zmm_d, zmm_sum_0);\n lhs += 32;\n rhs += 32;\n }\n } else {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n __m512h zmm_d_0 =\n _mm512_sub_ph(_mm512_loadu_ph(lhs + 0), _mm512_loadu_ph(rhs + 0));\n __m512h zmm_d_1 =\n _mm512_sub_ph(_mm512_loadu_ph(lhs + 32), _mm512_loadu_ph(rhs + 32));\n zmm_sum_0 = _mm512_fmadd_ph(zmm_d_0, zmm_d_0, zmm_sum_0);\n zmm_sum_1 = _mm512_fmadd_ph(zmm_d_1, zmm_d_1, zmm_sum_1);\n }\n\n if (last >= last_aligned + 32) {\n __m512h zmm_d = _mm512_sub_ph(_mm512_loadu_ph(lhs), _mm512_loadu_ph(rhs));\n zmm_sum_0 = _mm512_fmadd_ph(zmm_d, zmm_d, zmm_sum_0);\n lhs += 32;\n rhs += 32;\n }\n }\n\n zmm_sum_0 = _mm512_add_ph(zmm_sum_0, zmm_sum_1);\n if (lhs != last) {\n __mmask32 mask = (__mmask32)((1 << (last - lhs)) - 1);\n __m512i zmm_undefined = _mm512_undefined_epi32();\n __m512h zmm_undefined_ph = _mm512_undefined_ph();\n __m512h zmm_d = _mm512_mask_sub_ph(\n zmm_undefined_ph, mask,\n _mm512_castsi512_ph(_mm512_mask_loadu_epi16(zmm_undefined, mask, lhs)),\n _mm512_castsi512_ph(_mm512_mask_loadu_epi16(zmm_undefined, mask, rhs)));\n zmm_sum_0 = _mm512_mask3_fmadd_ph(zmm_d, zmm_d, zmm_sum_0, mask);\n }\n\n return HorizontalAdd_FP16_V512(zmm_sum_0);\n}\n#endif\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp16_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__ARM_NEON)\nfloat SquaredEuclideanDistanceFp16NEON(const Float16 *lhs, const Float16 *rhs,\n size_t size);\n#endif\n\n#if defined(__AVX512FP16__)\nfloat SquaredEuclideanDistanceFp16AVX512FP16(const Float16 *lhs,\n const Float16 *rhs, size_t size);\n#endif\n\n#if defined(__AVX512F__)\nfloat SquaredEuclideanDistanceFp16AVX512(const Float16 *lhs, const Float16 *rhs,\n size_t size);\n#endif\n\n#if defined(__AVX__)\nfloat SquaredEuclideanDistanceFp16AVX(const Float16 *lhs, const Float16 *rhs,\n size_t size);\n#endif\n\nfloat SquaredEuclideanDistanceFp16Scalar(const Float16 *lhs, const Float16 *rhs,\n size_t size);\n\n//! Compute the distance between matrix and query (FP16, M=1, N=1)\nvoid SquaredEuclideanDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n *out = SquaredEuclideanDistanceFp16NEON(m, q, dim);\n#else\n#if defined(__AVX512FP16__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) {\n *out = SquaredEuclideanDistanceFp16AVX512FP16(m, q, dim);\n return;\n }\n#endif\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = SquaredEuclideanDistanceFp16AVX512(m, q, dim);\n return;\n }\n#endif\n\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = SquaredEuclideanDistanceFp16AVX(m, q, dim);\n return;\n }\n#endif\n *out = SquaredEuclideanDistanceFp16Scalar(m, q, dim);\n\n#endif //__ARM_NEON\n}\n\n//! Compute the distance between matrix and query (FP16, M=1, N=1)\nvoid EuclideanDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim, float *out) {\n SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out);\n *out = std::sqrt(*out);\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp16_neon.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__ARM_NEON)\nfloat SquaredEuclideanDistanceFp16NEON(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n float score{0.0f};\n\n ACCUM_FP16_1X1_NEON(lhs, rhs, size, &score, 0ull, )\n\n return score;\n}\n#endif\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp32_avx.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX__)\nfloat SquaredEuclideanDistanceFp32SSEInternal(const float *lhs,\n const float *rhs, size_t size);\n\nfloat SquaredEuclideanDistanceFp32AVXInternal(const float *lhs,\n const float *rhs, size_t size) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 4) << 4);\n\n __m256 ymm_sum_0 = _mm256_setzero_ps();\n __m256 ymm_sum_1 = _mm256_setzero_ps();\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m256 ymm_d_0 =\n _mm256_sub_ps(_mm256_load_ps(lhs + 0), _mm256_load_ps(rhs + 0));\n __m256 ymm_d_1 =\n _mm256_sub_ps(_mm256_load_ps(lhs + 8), _mm256_load_ps(rhs + 8));\n ymm_sum_0 = _mm256_fmadd_ps(ymm_d_0, ymm_d_0, ymm_sum_0);\n ymm_sum_1 = _mm256_fmadd_ps(ymm_d_1, ymm_d_1, ymm_sum_1);\n }\n\n if (last >= last_aligned + 8) {\n __m256 ymm_d = _mm256_sub_ps(_mm256_load_ps(lhs), _mm256_load_ps(rhs));\n ymm_sum_0 = _mm256_fmadd_ps(ymm_d, ymm_d, ymm_sum_0);\n lhs += 8;\n rhs += 8;\n }\n } else {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m256 ymm_d_0 =\n _mm256_sub_ps(_mm256_loadu_ps(lhs + 0), _mm256_loadu_ps(rhs + 0));\n __m256 ymm_d_1 =\n _mm256_sub_ps(_mm256_loadu_ps(lhs + 8), _mm256_loadu_ps(rhs + 8));\n ymm_sum_0 = _mm256_fmadd_ps(ymm_d_0, ymm_d_0, ymm_sum_0);\n ymm_sum_1 = _mm256_fmadd_ps(ymm_d_1, ymm_d_1, ymm_sum_1);\n }\n\n if (last >= last_aligned + 8) {\n __m256 ymm_d = _mm256_sub_ps(_mm256_loadu_ps(lhs), _mm256_loadu_ps(rhs));\n ymm_sum_0 = _mm256_fmadd_ps(ymm_d, ymm_d, ymm_sum_0);\n lhs += 8;\n rhs += 8;\n }\n }\n float result = HorizontalAdd_FP32_V256(_mm256_add_ps(ymm_sum_0, ymm_sum_1));\n\n switch (last - lhs) {\n case 7:\n SSD_FP32_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n SSD_FP32_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n SSD_FP32_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n SSD_FP32_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n SSD_FP32_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n SSD_FP32_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n SSD_FP32_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat SquaredEuclideanDistanceFp32AVX(const float *lhs, const float *rhs,\n size_t size) {\n if (size > 7) {\n return SquaredEuclideanDistanceFp32AVXInternal(lhs, rhs, size);\n }\n\n return SquaredEuclideanDistanceFp32SSEInternal(lhs, rhs, size);\n}\n\n#endif // __AVX__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp32_avx512.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX512F__)\nfloat SquaredEuclideanDistanceFp32SSEInternal(const float *lhs,\n const float *rhs, size_t size);\n\nfloat SquaredEuclideanDistanceFp32AVXInternal(const float *lhs,\n const float *rhs, size_t size);\n\nfloat SquaredEuclideanDistanceFp32AVX512Internal(const float *lhs,\n const float *rhs,\n size_t size) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 5) << 5);\n\n __m512 zmm_sum_0 = _mm512_setzero_ps();\n __m512 zmm_sum_1 = _mm512_setzero_ps();\n\n if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m512 zmm_d_0 =\n _mm512_sub_ps(_mm512_load_ps(lhs + 0), _mm512_load_ps(rhs + 0));\n __m512 zmm_d_1 =\n _mm512_sub_ps(_mm512_load_ps(lhs + 16), _mm512_load_ps(rhs + 16));\n zmm_sum_0 = _mm512_fmadd_ps(zmm_d_0, zmm_d_0, zmm_sum_0);\n zmm_sum_1 = _mm512_fmadd_ps(zmm_d_1, zmm_d_1, zmm_sum_1);\n }\n\n if (last >= last_aligned + 16) {\n __m512 zmm_d = _mm512_sub_ps(_mm512_load_ps(lhs), _mm512_load_ps(rhs));\n zmm_sum_0 = _mm512_fmadd_ps(zmm_d, zmm_d, zmm_sum_0);\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m512 zmm_d_0 =\n _mm512_sub_ps(_mm512_loadu_ps(lhs + 0), _mm512_loadu_ps(rhs + 0));\n __m512 zmm_d_1 =\n _mm512_sub_ps(_mm512_loadu_ps(lhs + 16), _mm512_loadu_ps(rhs + 16));\n zmm_sum_0 = _mm512_fmadd_ps(zmm_d_0, zmm_d_0, zmm_sum_0);\n zmm_sum_1 = _mm512_fmadd_ps(zmm_d_1, zmm_d_1, zmm_sum_1);\n }\n\n if (last >= last_aligned + 16) {\n __m512 zmm_d = _mm512_sub_ps(_mm512_loadu_ps(lhs), _mm512_loadu_ps(rhs));\n zmm_sum_0 = _mm512_fmadd_ps(zmm_d, zmm_d, zmm_sum_0);\n lhs += 16;\n rhs += 16;\n }\n }\n\n zmm_sum_0 = _mm512_add_ps(zmm_sum_0, zmm_sum_1);\n if (lhs != last) {\n __mmask16 mask = (__mmask16)((1 << (last - lhs)) - 1);\n __m512 zmm_undefined = _mm512_undefined_ps();\n __m512 zmm_d = _mm512_mask_sub_ps(\n zmm_undefined, mask, _mm512_mask_loadu_ps(zmm_undefined, mask, lhs),\n _mm512_mask_loadu_ps(zmm_undefined, mask, rhs));\n zmm_sum_0 = _mm512_mask3_fmadd_ps(zmm_d, zmm_d, zmm_sum_0, mask);\n }\n return HorizontalAdd_FP32_V512(zmm_sum_0);\n}\n\nfloat SquaredEuclideanDistanceFp32AVX512(const float *lhs, const float *rhs,\n size_t size) {\n if (size > 15) {\n return SquaredEuclideanDistanceFp32AVX512Internal(lhs, rhs, size);\n }\n\n if (size > 7) {\n return SquaredEuclideanDistanceFp32AVXInternal(lhs, rhs, size);\n }\n\n return SquaredEuclideanDistanceFp32SSEInternal(lhs, rhs, size);\n}\n\n#endif\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp32_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__ARM_NEON)\nvoid SquaredEuclideanDistanceFp32NEON(const float *lhs, const float *rhs,\n size_t size, float *out);\n#endif\n\n#if defined(__AVX512F__)\nfloat SquaredEuclideanDistanceFp32AVX512(const float *lhs, const float *rhs,\n size_t size);\n#endif\n\n#if defined(__AVX__)\nfloat SquaredEuclideanDistanceFp32AVX(const float *lhs, const float *rhs,\n size_t size);\n#endif\n\n#if defined(__SSE__)\nfloat SquaredEuclideanDistanceFp32SSE(const float *lhs, const float *rhs,\n size_t size);\n#endif\n\nfloat SquaredEuclideanDistanceFp32Scalar(const float *lhs, const float *rhs,\n size_t size);\n\n//-----------------------------------------------------------\n// SquaredEuclideanDistance\n//-----------------------------------------------------------\n//! Compute the distance between matrix and query (FP32, M=1, N=1)\nvoid SquaredEuclideanDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n SquaredEuclideanDistanceFp32NEON(m, q, dim, out);\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = SquaredEuclideanDistanceFp32AVX512(m, q, dim);\n return;\n }\n#endif // __AVX512F__\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = SquaredEuclideanDistanceFp32AVX(m, q, dim);\n return;\n }\n#endif // __AVX__\n\n#if defined(__SSE__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) {\n *out = SquaredEuclideanDistanceFp32SSE(m, q, dim);\n return;\n }\n#endif // __SSE__\n *out = SquaredEuclideanDistanceFp32Scalar(m, q, dim);\n#endif // __ARM_NEON\n}\n\n//-----------------------------------------------------------\n// EuclideanDistance\n//-----------------------------------------------------------\n//! Compute the distance between matrix and query (FP32, M=1, N=1)\nvoid EuclideanDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim, float *out) {\n SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out);\n *out = std::sqrt(*out);\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp32_neon.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__ARM_NEON)\n//! Squared Euclidean Distance\nvoid SquaredEuclideanDistanceFp32NEON(const float *lhs, const float *rhs,\n size_t size, float *out) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 3) << 3);\n\n float32x4_t v_sum_0 = vdupq_n_f32(0);\n float32x4_t v_sum_1 = vdupq_n_f32(0);\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n float32x4_t v_d_0 = vsubq_f32(vld1q_f32(lhs + 0), vld1q_f32(rhs + 0));\n float32x4_t v_d_1 = vsubq_f32(vld1q_f32(lhs + 4), vld1q_f32(rhs + 4));\n v_sum_0 = vfmaq_f32(v_sum_0, v_d_0, v_d_0);\n v_sum_1 = vfmaq_f32(v_sum_1, v_d_1, v_d_1);\n }\n if (last >= last_aligned + 4) {\n float32x4_t v_d = vsubq_f32(vld1q_f32(lhs), vld1q_f32(rhs));\n v_sum_0 = vfmaq_f32(v_sum_0, v_d, v_d);\n lhs += 4;\n rhs += 4;\n }\n\n float result = vaddvq_f32(vaddq_f32(v_sum_0, v_sum_1));\n switch (last - lhs) {\n case 3:\n SSD_FP32_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n SSD_FP32_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n SSD_FP32_GENERAL(lhs[0], rhs[0], result)\n }\n *out = result;\n}\n\n#endif // __ARM_NEON\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_fp32_sse.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__SSE__)\nfloat SquaredEuclideanDistanceFp32SSEInternal(const float *lhs,\n const float *rhs, size_t size) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 3) << 3);\n\n __m128 xmm_sum_0 = _mm_setzero_ps();\n __m128 xmm_sum_1 = _mm_setzero_ps();\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m128 xmm_d_0 = _mm_sub_ps(_mm_load_ps(lhs + 0), _mm_load_ps(rhs + 0));\n __m128 xmm_d_1 = _mm_sub_ps(_mm_load_ps(lhs + 4), _mm_load_ps(rhs + 4));\n xmm_sum_0 = _mm_fmadd_ps(xmm_d_0, xmm_d_0, xmm_sum_0);\n xmm_sum_1 = _mm_fmadd_ps(xmm_d_1, xmm_d_1, xmm_sum_1);\n }\n\n if (last >= last_aligned + 4) {\n __m128 xmm_d = _mm_sub_ps(_mm_load_ps(lhs), _mm_load_ps(rhs));\n xmm_sum_0 = _mm_fmadd_ps(xmm_d, xmm_d, xmm_sum_0);\n lhs += 4;\n rhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m128 xmm_d_0 = _mm_sub_ps(_mm_loadu_ps(lhs + 0), _mm_loadu_ps(rhs + 0));\n __m128 xmm_d_1 = _mm_sub_ps(_mm_loadu_ps(lhs + 4), _mm_loadu_ps(rhs + 4));\n xmm_sum_0 = _mm_fmadd_ps(xmm_d_0, xmm_d_0, xmm_sum_0);\n xmm_sum_1 = _mm_fmadd_ps(xmm_d_1, xmm_d_1, xmm_sum_1);\n }\n\n if (last >= last_aligned + 4) {\n __m128 xmm_d = _mm_sub_ps(_mm_loadu_ps(lhs), _mm_loadu_ps(rhs));\n xmm_sum_0 = _mm_fmadd_ps(xmm_d, xmm_d, xmm_sum_0);\n lhs += 4;\n rhs += 4;\n }\n }\n float result = HorizontalAdd_FP32_V128(_mm_add_ps(xmm_sum_0, xmm_sum_1));\n\n switch (last - lhs) {\n case 3:\n SSD_FP32_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n SSD_FP32_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n SSD_FP32_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat SquaredEuclideanDistanceFp32SSE(const float *lhs, const float *rhs,\n size_t size) {\n return SquaredEuclideanDistanceFp32SSEInternal(lhs, rhs, size);\n}\n\n#endif // __SSE__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_int4_avx2.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int4.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX2__)\nfloat SquaredEuclideanDistanceInt4SSEInternal(const uint8_t *lhs,\n const uint8_t *rhs, size_t size);\n\ninline float SquaredEuclideanDistanceInt4AVX2Internal(const uint8_t *lhs,\n const uint8_t *rhs,\n size_t size) {\n const uint8_t *last = lhs + size;\n const uint8_t *last_aligned = lhs + ((size >> 5) << 5);\n\n __m256i ymm_sum = _mm256_setzero_si256();\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m256i ymm_lhs = _mm256_load_si256((const __m256i *)(lhs));\n __m256i ymm_rhs = _mm256_load_si256((const __m256i *)(rhs));\n SSD_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum)\n }\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs);\n __m128i xmm_sum = _mm_setzero_si128();\n SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum)\n ymm_sum = _mm256_add_epi32(_mm256_set_m128i(_mm_setzero_si128(), xmm_sum),\n ymm_sum);\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)(lhs));\n __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)(rhs));\n SSD_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum)\n }\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs);\n __m128i xmm_sum = _mm_setzero_si128();\n SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum)\n ymm_sum = _mm256_add_epi32(_mm256_set_m128i(_mm_setzero_si128(), xmm_sum),\n ymm_sum);\n lhs += 16;\n rhs += 16;\n }\n }\n float result = static_cast(HorizontalAdd_INT32_V256(ymm_sum));\n\n switch (last - lhs) {\n case 15:\n SSD_INT4_GENERAL(lhs[14], rhs[14], result)\n /* FALLTHRU */\n case 14:\n SSD_INT4_GENERAL(lhs[13], rhs[13], result)\n /* FALLTHRU */\n case 13:\n SSD_INT4_GENERAL(lhs[12], rhs[12], result)\n /* FALLTHRU */\n case 12:\n SSD_INT4_GENERAL(lhs[11], rhs[11], result)\n /* FALLTHRU */\n case 11:\n SSD_INT4_GENERAL(lhs[10], rhs[10], result)\n /* FALLTHRU */\n case 10:\n SSD_INT4_GENERAL(lhs[9], rhs[9], result)\n /* FALLTHRU */\n case 9:\n SSD_INT4_GENERAL(lhs[8], rhs[8], result)\n /* FALLTHRU */\n case 8:\n SSD_INT4_GENERAL(lhs[7], rhs[7], result)\n /* FALLTHRU */\n case 7:\n SSD_INT4_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n SSD_INT4_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n SSD_INT4_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n SSD_INT4_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n SSD_INT4_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n SSD_INT4_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n SSD_INT4_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat SquaredEuclideanDistanceInt4AVX2(const uint8_t *lhs, const uint8_t *rhs,\n size_t size) {\n if (size > 63) {\n return SquaredEuclideanDistanceInt4AVX2Internal(lhs, rhs, size >> 1);\n }\n\n return SquaredEuclideanDistanceInt4SSEInternal(lhs, rhs, size >> 1);\n}\n\n#endif // __AVX2__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_int4_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX2__)\nfloat SquaredEuclideanDistanceInt4AVX2(const uint8_t *lhs, const uint8_t *rhs,\n size_t size);\n#endif\n\n#if defined(__SSE4_1__)\nfloat SquaredEuclideanDistanceInt4SSE(const uint8_t *lhs, const uint8_t *rhs,\n size_t size);\n#endif\n\nfloat SquaredEuclideanDistanceInt4Scalar(const uint8_t *lhs, const uint8_t *rhs,\n size_t size);\n\n//! Compute the distance between matrix and query (INT4, M=1, N=1)\nvoid SquaredEuclideanDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim,\n float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = SquaredEuclideanDistanceInt4AVX2(m, q, dim);\n return;\n }\n#endif // __AVX2__\n\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = SquaredEuclideanDistanceInt4SSE(m, q, dim);\n return;\n }\n#endif\n\n *out = SquaredEuclideanDistanceInt4Scalar(m, q, dim);\n}\n\n//! Compute the distance between matrix and query (INT4, M=1, N=1)\nvoid EuclideanDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim, float *out) {\n SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out);\n *out = std::sqrt(*out);\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_int4_sse.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int4.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__SSE4_1__)\nfloat SquaredEuclideanDistanceInt4SSEInternal(const uint8_t *lhs,\n const uint8_t *rhs, size_t size) {\n const uint8_t *last = lhs + size;\n const uint8_t *last_aligned = lhs + ((size >> 4) << 4);\n\n __m128i xmm_sum = _mm_setzero_si128();\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)(lhs));\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)(rhs));\n SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum)\n }\n } else {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)(lhs));\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)(rhs));\n SSD_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum)\n }\n }\n float result = static_cast(HorizontalAdd_INT32_V128(xmm_sum));\n\n switch (last - lhs) {\n case 15:\n SSD_INT4_GENERAL(lhs[14], rhs[14], result)\n /* FALLTHRU */\n case 14:\n SSD_INT4_GENERAL(lhs[13], rhs[13], result)\n /* FALLTHRU */\n case 13:\n SSD_INT4_GENERAL(lhs[12], rhs[12], result)\n /* FALLTHRU */\n case 12:\n SSD_INT4_GENERAL(lhs[11], rhs[11], result)\n /* FALLTHRU */\n case 11:\n SSD_INT4_GENERAL(lhs[10], rhs[10], result)\n /* FALLTHRU */\n case 10:\n SSD_INT4_GENERAL(lhs[9], rhs[9], result)\n /* FALLTHRU */\n case 9:\n SSD_INT4_GENERAL(lhs[8], rhs[8], result)\n /* FALLTHRU */\n case 8:\n SSD_INT4_GENERAL(lhs[7], rhs[7], result)\n /* FALLTHRU */\n case 7:\n SSD_INT4_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n SSD_INT4_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n SSD_INT4_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n SSD_INT4_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n SSD_INT4_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n SSD_INT4_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n SSD_INT4_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat SquaredEuclideanDistanceInt4SSE(const uint8_t *lhs, const uint8_t *rhs,\n size_t size) {\n return SquaredEuclideanDistanceInt4SSEInternal(lhs, rhs, size >> 1);\n}\n\n#endif // __SSE4_1__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_int8_avx2.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int8.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX2__)\nfloat SquaredEuclideanDistanceInt8SSEInternal(const int8_t *lhs,\n const int8_t *rhs, size_t size);\n\nfloat SquaredEuclideanDistanceInt8AVX2Internal(const int8_t *lhs,\n const int8_t *rhs, size_t size) {\n const int8_t *last = lhs + size;\n const int8_t *last_aligned = lhs + ((size >> 6) << 6);\n float result = 0.0;\n\n __m256i ymm_sum_0 = _mm256_setzero_si256();\n __m256i ymm_sum_1 = _mm256_setzero_si256();\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n __m256i ymm_lhs_0 = _mm256_load_si256((const __m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_load_si256((const __m256i *)(lhs + 32));\n __m256i ymm_rhs_0 = _mm256_load_si256((const __m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_load_si256((const __m256i *)(rhs + 32));\n\n __m256i ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_0, ymm_rhs_0),\n _mm256_min_epi8(ymm_lhs_0, ymm_rhs_0));\n ymm_lhs_0 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d));\n ymm_rhs_0 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1));\n ymm_sum_0 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs_0, ymm_lhs_0), ymm_sum_0);\n ymm_sum_1 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs_0, ymm_rhs_0), ymm_sum_1);\n\n ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_1, ymm_rhs_1),\n _mm256_min_epi8(ymm_lhs_1, ymm_rhs_1));\n ymm_lhs_1 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d));\n ymm_rhs_1 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1));\n ymm_sum_0 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs_1, ymm_lhs_1), ymm_sum_0);\n ymm_sum_1 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs_1, ymm_rhs_1), ymm_sum_1);\n }\n\n if (last >= last_aligned + 32) {\n __m256i ymm_lhs = _mm256_load_si256((const __m256i *)lhs);\n __m256i ymm_rhs = _mm256_load_si256((const __m256i *)rhs);\n __m256i ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs, ymm_rhs),\n _mm256_min_epi8(ymm_lhs, ymm_rhs));\n ymm_lhs = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d));\n ymm_rhs = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1));\n ymm_sum_0 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs, ymm_lhs), ymm_sum_0);\n ymm_sum_1 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs, ymm_rhs), ymm_sum_1);\n lhs += 32;\n rhs += 32;\n }\n } else {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n __m256i ymm_lhs_0 = _mm256_loadu_si256((const __m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_loadu_si256((const __m256i *)(lhs + 32));\n __m256i ymm_rhs_0 = _mm256_loadu_si256((const __m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_loadu_si256((const __m256i *)(rhs + 32));\n\n __m256i ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_0, ymm_rhs_0),\n _mm256_min_epi8(ymm_lhs_0, ymm_rhs_0));\n ymm_lhs_0 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d));\n ymm_rhs_0 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1));\n ymm_sum_0 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs_0, ymm_lhs_0), ymm_sum_0);\n ymm_sum_1 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs_0, ymm_rhs_0), ymm_sum_1);\n\n ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs_1, ymm_rhs_1),\n _mm256_min_epi8(ymm_lhs_1, ymm_rhs_1));\n ymm_lhs_1 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d));\n ymm_rhs_1 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1));\n ymm_sum_0 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs_1, ymm_lhs_1), ymm_sum_0);\n ymm_sum_1 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs_1, ymm_rhs_1), ymm_sum_1);\n }\n\n if (last >= last_aligned + 32) {\n __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)lhs);\n __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)rhs);\n __m256i ymm_d = _mm256_sub_epi8(_mm256_max_epi8(ymm_lhs, ymm_rhs),\n _mm256_min_epi8(ymm_lhs, ymm_rhs));\n ymm_lhs = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(ymm_d));\n ymm_rhs = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(ymm_d, 1));\n ymm_sum_0 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_lhs, ymm_lhs), ymm_sum_0);\n ymm_sum_1 =\n _mm256_add_epi32(_mm256_madd_epi16(ymm_rhs, ymm_rhs), ymm_sum_1);\n lhs += 32;\n rhs += 32;\n }\n }\n result = static_cast(\n HorizontalAdd_INT32_V256(_mm256_add_epi32(ymm_sum_0, ymm_sum_1)));\n\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs);\n __m128i xmm_sum = _mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs),\n _mm_min_epi8(xmm_lhs, xmm_rhs));\n xmm_lhs = _mm_cvtepu8_epi16(xmm_sum);\n xmm_rhs = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_sum, xmm_sum));\n xmm_sum = _mm_add_epi32(_mm_madd_epi16(xmm_lhs, xmm_lhs),\n _mm_madd_epi16(xmm_rhs, xmm_rhs));\n result += static_cast(HorizontalAdd_INT32_V128(xmm_sum));\n lhs += 16;\n rhs += 16;\n }\n switch (last - lhs) {\n case 15:\n SSD_INT8_GENERAL(lhs[14], rhs[14], result)\n /* FALLTHRU */\n case 14:\n SSD_INT8_GENERAL(lhs[13], rhs[13], result)\n /* FALLTHRU */\n case 13:\n SSD_INT8_GENERAL(lhs[12], rhs[12], result)\n /* FALLTHRU */\n case 12:\n SSD_INT8_GENERAL(lhs[11], rhs[11], result)\n /* FALLTHRU */\n case 11:\n SSD_INT8_GENERAL(lhs[10], rhs[10], result)\n /* FALLTHRU */\n case 10:\n SSD_INT8_GENERAL(lhs[9], rhs[9], result)\n /* FALLTHRU */\n case 9:\n SSD_INT8_GENERAL(lhs[8], rhs[8], result)\n /* FALLTHRU */\n case 8:\n SSD_INT8_GENERAL(lhs[7], rhs[7], result)\n /* FALLTHRU */\n case 7:\n SSD_INT8_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n SSD_INT8_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n SSD_INT8_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n SSD_INT8_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n SSD_INT8_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n SSD_INT8_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n SSD_INT8_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat SquaredEuclideanDistanceInt8AVX2(const int8_t *lhs, const int8_t *rhs,\n size_t size) {\n if (size > 31) {\n return SquaredEuclideanDistanceInt8AVX2Internal(lhs, rhs, size);\n }\n\n return SquaredEuclideanDistanceInt8SSEInternal(lhs, rhs, size);\n}\n#endif // __AVX2__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_int8_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX2__)\nfloat SquaredEuclideanDistanceInt8AVX2(const int8_t *lhs, const int8_t *rhs,\n size_t size);\n#endif\n\n#if defined(__SSE4_1__)\nfloat SquaredEuclideanDistanceInt8SSE(const int8_t *lhs, const int8_t *rhs,\n size_t size);\n#endif\n\nfloat SquaredEuclideanDistanceInt8Scalar(const int8_t *lhs, const int8_t *rhs,\n size_t size);\n\n//! Compute the distance between matrix and query (INT8, M=1, N=1)\nvoid SquaredEuclideanDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim,\n float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = SquaredEuclideanDistanceInt8AVX2(m, q, dim);\n return;\n }\n#endif // __AVX2__\n\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = SquaredEuclideanDistanceInt8SSE(m, q, dim);\n return;\n }\n#endif\n\n *out = SquaredEuclideanDistanceInt8Scalar(m, q, dim);\n}\n\n//! Compute the distance between matrix and query (INT8, M=1, N=1)\nvoid EuclideanDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim, float *out) {\n SquaredEuclideanDistanceMatrix::Compute(m, q, dim, out);\n *out = std::sqrt(*out);\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_int8_sse.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int8.i\"\n#include \"distance_matrix_euclidean_utility.i\"\n#include \"euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__SSE4_1__)\ninline float SquaredEuclideanDistanceInt8SSEInternal(const int8_t *lhs,\n const int8_t *rhs,\n size_t size) {\n const int8_t *last = lhs + size;\n const int8_t *last_aligned = lhs + ((size >> 5) << 5);\n\n __m128i xmm_sum_0 = _mm_setzero_si128();\n __m128i xmm_sum_1 = _mm_setzero_si128();\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m128i xmm_lhs_0 = _mm_load_si128((const __m128i *)(lhs + 0));\n __m128i xmm_lhs_1 = _mm_load_si128((const __m128i *)(lhs + 16));\n __m128i xmm_rhs_0 = _mm_load_si128((const __m128i *)(rhs + 0));\n __m128i xmm_rhs_1 = _mm_load_si128((const __m128i *)(rhs + 16));\n\n __m128i xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs_0, xmm_rhs_0),\n _mm_min_epi8(xmm_lhs_0, xmm_rhs_0));\n xmm_lhs_0 = _mm_cvtepu8_epi16(xmm_d);\n xmm_rhs_0 = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d));\n xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs_1, xmm_rhs_1),\n _mm_min_epi8(xmm_lhs_1, xmm_rhs_1));\n xmm_lhs_1 = _mm_cvtepu8_epi16(xmm_d);\n xmm_rhs_1 = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d));\n\n xmm_sum_0 =\n _mm_add_epi32(_mm_madd_epi16(xmm_lhs_0, xmm_lhs_0), xmm_sum_0);\n xmm_sum_1 =\n _mm_add_epi32(_mm_madd_epi16(xmm_rhs_0, xmm_rhs_0), xmm_sum_1);\n xmm_sum_0 =\n _mm_add_epi32(_mm_madd_epi16(xmm_lhs_1, xmm_lhs_1), xmm_sum_0);\n xmm_sum_1 =\n _mm_add_epi32(_mm_madd_epi16(xmm_rhs_1, xmm_rhs_1), xmm_sum_1);\n }\n\n if (last >= last_aligned + 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs);\n __m128i xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs),\n _mm_min_epi8(xmm_lhs, xmm_rhs));\n xmm_lhs = _mm_cvtepu8_epi16(xmm_d);\n xmm_rhs = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d));\n xmm_sum_0 = _mm_add_epi32(_mm_madd_epi16(xmm_lhs, xmm_lhs), xmm_sum_0);\n xmm_sum_1 = _mm_add_epi32(_mm_madd_epi16(xmm_rhs, xmm_rhs), xmm_sum_1);\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m128i xmm_lhs_0 = _mm_loadu_si128((const __m128i *)(lhs + 0));\n __m128i xmm_lhs_1 = _mm_loadu_si128((const __m128i *)(lhs + 16));\n __m128i xmm_rhs_0 = _mm_loadu_si128((const __m128i *)(rhs + 0));\n __m128i xmm_rhs_1 = _mm_loadu_si128((const __m128i *)(rhs + 16));\n\n __m128i xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs_0, xmm_rhs_0),\n _mm_min_epi8(xmm_lhs_0, xmm_rhs_0));\n xmm_lhs_0 = _mm_cvtepu8_epi16(xmm_d);\n xmm_rhs_0 = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d));\n xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs_1, xmm_rhs_1),\n _mm_min_epi8(xmm_lhs_1, xmm_rhs_1));\n xmm_lhs_1 = _mm_cvtepu8_epi16(xmm_d);\n xmm_rhs_1 = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d));\n\n xmm_sum_0 =\n _mm_add_epi32(_mm_madd_epi16(xmm_lhs_0, xmm_lhs_0), xmm_sum_0);\n xmm_sum_1 =\n _mm_add_epi32(_mm_madd_epi16(xmm_rhs_0, xmm_rhs_0), xmm_sum_1);\n xmm_sum_0 =\n _mm_add_epi32(_mm_madd_epi16(xmm_lhs_1, xmm_lhs_1), xmm_sum_0);\n xmm_sum_1 =\n _mm_add_epi32(_mm_madd_epi16(xmm_rhs_1, xmm_rhs_1), xmm_sum_1);\n }\n\n if (last >= last_aligned + 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs);\n __m128i xmm_d = _mm_sub_epi8(_mm_max_epi8(xmm_lhs, xmm_rhs),\n _mm_min_epi8(xmm_lhs, xmm_rhs));\n xmm_lhs = _mm_cvtepu8_epi16(xmm_d);\n xmm_rhs = _mm_cvtepu8_epi16(_mm_unpackhi_epi64(xmm_d, xmm_d));\n xmm_sum_0 = _mm_add_epi32(_mm_madd_epi16(xmm_lhs, xmm_lhs), xmm_sum_0);\n xmm_sum_1 = _mm_add_epi32(_mm_madd_epi16(xmm_rhs, xmm_rhs), xmm_sum_1);\n lhs += 16;\n rhs += 16;\n }\n }\n float result = static_cast(\n HorizontalAdd_INT32_V128(_mm_add_epi32(xmm_sum_0, xmm_sum_1)));\n\n switch (last - lhs) {\n case 15:\n SSD_INT8_GENERAL(lhs[14], rhs[14], result)\n /* FALLTHRU */\n case 14:\n SSD_INT8_GENERAL(lhs[13], rhs[13], result)\n /* FALLTHRU */\n case 13:\n SSD_INT8_GENERAL(lhs[12], rhs[12], result)\n /* FALLTHRU */\n case 12:\n SSD_INT8_GENERAL(lhs[11], rhs[11], result)\n /* FALLTHRU */\n case 11:\n SSD_INT8_GENERAL(lhs[10], rhs[10], result)\n /* FALLTHRU */\n case 10:\n SSD_INT8_GENERAL(lhs[9], rhs[9], result)\n /* FALLTHRU */\n case 9:\n SSD_INT8_GENERAL(lhs[8], rhs[8], result)\n /* FALLTHRU */\n case 8:\n SSD_INT8_GENERAL(lhs[7], rhs[7], result)\n /* FALLTHRU */\n case 7:\n SSD_INT8_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n SSD_INT8_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n SSD_INT8_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n SSD_INT8_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n SSD_INT8_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n SSD_INT8_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n SSD_INT8_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\n//! Squared Euclidean Distance\nfloat SquaredEuclideanDistanceInt8SSE(const int8_t *lhs, const int8_t *rhs,\n size_t size) {\n return SquaredEuclideanDistanceInt8SSEInternal(lhs, rhs, size);\n}\n\n#endif // __SSE4_1__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/euclidean_distance_matrix_scalar.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \"distance_utility.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\ntemplate \ninline float SquaredEuclideanDistanceScalar(const T *m, const T *q,\n size_t dim) {\n ailego_assert(m && q && dim);\n\n float sum = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n sum += MathHelper::SquaredDifference(m[i], q[i]);\n }\n\n return sum;\n}\n\ntemplate \ninline float EuclideanDistanceScalar(const T *m, const T *q, size_t dim) {\n ailego_assert(m && q && dim);\n\n float sum = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n sum += MathHelper::SquaredDifference(m[i], q[i]);\n }\n\n return std::sqrt(sum);\n}\n\nfloat SquaredEuclideanDistanceInt4Scalar(const uint8_t *m, const uint8_t *q,\n size_t dim) {\n ailego_assert(m && q && dim && !(dim & 1));\n\n float sum = 0.0;\n for (size_t i = 0; i < (dim >> 1); ++i) {\n uint8_t m_val = m[i];\n uint8_t q_val = q[i];\n sum += Int4SquaredDiffTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] +\n Int4SquaredDiffTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)];\n }\n\n return sum;\n}\n\n\nfloat EuclideanDistanceInt4Scalar(const uint8_t *m, const uint8_t *q,\n size_t dim) {\n ailego_assert(m && q && dim && !(dim & 1));\n\n float sum = 0.0;\n for (size_t i = 0; i < (dim >> 1); ++i) {\n uint8_t m_val = m[i];\n uint8_t q_val = q[i];\n sum += Int4SquaredDiffTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] +\n Int4SquaredDiffTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)];\n }\n\n return std::sqrt(sum);\n}\n\n\nfloat SquaredEuclideanDistanceInt8Scalar(const int8_t *m, const int8_t *q,\n size_t dim) {\n return SquaredEuclideanDistanceScalar(m, q, dim);\n}\n\nfloat EuclideanDistanceInt8Scalar(const int8_t *m, const int8_t *q,\n size_t dim) {\n return EuclideanDistanceScalar(m, q, dim);\n}\n\nfloat SquaredEuclideanDistanceFp16Scalar(const ailego::Float16 *m,\n const ailego::Float16 *q, size_t dim) {\n return SquaredEuclideanDistanceScalar(m, q, dim);\n}\n\nfloat EuclideanDistanceFp16Scalar(const ailego::Float16 *m,\n const ailego::Float16 *q, size_t dim) {\n return EuclideanDistanceScalar(m, q, dim);\n}\n\nfloat SquaredEuclideanDistanceFp32Scalar(const float *m, const float *q,\n size_t dim) {\n return SquaredEuclideanDistanceScalar(m, q, dim);\n}\n\nfloat EuclideanDistanceFp32Scalar(const float *m, const float *q, size_t dim) {\n return EuclideanDistanceScalar(m, q, dim);\n}\n\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/hamming_distance_matrix.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"hamming_distance_matrix.h\"\n#include \n#include \n#include \"distance_matrix_popcnt.i\"\n\nnamespace zvec {\nnamespace ailego {\n\n#define POPCNT_UINT32_STEP1_SSE HAMMING_UINT32_STEP1_SSE\n#define POPCNT_UINT32_STEP2_SSE HAMMING_UINT32_STEP2_SSE\n#define POPCNT_UINT32_STEP3_SSE HAMMING_UINT32_STEP3_SSE\n#define POPCNT_UINT32_STEP1_AVX HAMMING_UINT32_STEP1_AVX\n#define POPCNT_UINT32_STEP2_AVX HAMMING_UINT32_STEP2_AVX\n#define POPCNT_UINT32_STEP3_AVX HAMMING_UINT32_STEP3_AVX\n#define POPCNT_UINT64_STEP1_AVX HAMMING_UINT64_STEP1_AVX\n#define POPCNT_UINT64_STEP2_AVX HAMMING_UINT64_STEP2_AVX\n\n//! Calculate population count (Step 1 SSE)\n#define HAMMING_UINT32_STEP1_SSE(xmm_m, xmm_q, xmm_sum) \\\n xmm_sum = _mm_add_epi8( \\\n VerticalPopCount_INT8_V128(_mm_xor_si128(xmm_m, xmm_q)), xmm_sum);\n\n//! Calculate population count (Step 2 SSE)\n#define HAMMING_UINT32_STEP2_SSE(xmm_m, xmm_q, xmm_sum) \\\n xmm_sum = _mm_add_epi16( \\\n VerticalPopCount_INT16_V128(_mm_xor_si128(xmm_m, xmm_q)), xmm_sum);\n\n//! Calculate population count (Step 3 SSE)\n#define HAMMING_UINT32_STEP3_SSE(xmm_m, xmm_q, xmm_sum) \\\n xmm_sum = _mm_add_epi32( \\\n VerticalPopCount_INT32_V128(_mm_xor_si128(xmm_m, xmm_q)), xmm_sum);\n\n//! Calculate population count (Step 1 AVX)\n#define HAMMING_UINT32_STEP1_AVX(ymm_m, ymm_q, ymm_sum) \\\n ymm_sum = _mm256_add_epi8( \\\n VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum);\n\n//! Calculate population count (Step 2 AVX)\n#define HAMMING_UINT32_STEP2_AVX(ymm_m, ymm_q, ymm_sum) \\\n ymm_sum = _mm256_add_epi16( \\\n VerticalPopCount_INT16_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum);\n\n//! Calculate population count (Step 3 AVX)\n#define HAMMING_UINT32_STEP3_AVX(ymm_m, ymm_q, ymm_sum) \\\n ymm_sum = _mm256_add_epi32( \\\n VerticalPopCount_INT32_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum);\n\n//! Calculate population count (Step 1 AVX)\n#define HAMMING_UINT64_STEP1_AVX(ymm_m, ymm_q, ymm_sum) \\\n ymm_sum = _mm256_add_epi8( \\\n VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum);\n\n//! Calculate population count (Step 2 AVX)\n#define HAMMING_UINT64_STEP2_AVX(ymm_m, ymm_q, ymm_sum) \\\n ymm_sum = _mm256_add_epi64( \\\n VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_m, ymm_q)), ymm_sum);\n\n#if defined(__AVX512VL__) && defined(__AVX512DQ__)\n#define CONVERT_UINT64_TO_FP32(v, ...) _mm256_cvtepu64_ps(v)\n#elif defined(__AVX2__)\nstatic const __m256i CONVERT_UINT32_MASK_AVX =\n _mm256_set_epi32(0, 0, 0, 0, 6, 4, 2, 0);\n\n#define CONVERT_UINT64_TO_FP32(v, ...) \\\n _mm_cvtepi32_ps(_mm256_castsi256_si128( \\\n _mm256_permutevar8x32_epi32(v, CONVERT_UINT32_MASK_AVX)))\n#endif // __AVX512VL__ && __AVX512DQ__\n\n#define SQRT_UINT64_TO_FP32(v, ...) _mm_sqrt_ps(CONVERT_UINT64_TO_FP32(v))\n#define SQRT_UINT32_TO_FP32_SSE(v, ...) _mm_sqrt_ps(_mm_cvtepi32_ps(v))\n#define SQRT_UINT32_TO_FP32_AVX(v, ...) _mm256_sqrt_ps(_mm256_cvtepi32_ps(v))\n\n#if defined(__AVX2__)\nstatic inline size_t HammingDistanceAVX(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n __m256i ymm_sum_0 = _mm256_setzero_si256();\n __m256i ymm_sum_1 = _mm256_setzero_si256();\n\n const uint32_t *lhs_0 = lhs + ((size >> 4) << 4);\n const uint32_t *lhs_1 = (size > 496 ? lhs + 496 : lhs_0);\n const uint32_t *lhs_2 = lhs + size;\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != lhs_1; lhs += 16, rhs += 16) {\n __m256i ymm_lhs_0 = _mm256_load_si256((__m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_load_si256((__m256i *)(lhs + 8));\n __m256i ymm_rhs_0 = _mm256_load_si256((__m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_load_si256((__m256i *)(rhs + 8));\n\n ymm_sum_0 = _mm256_add_epi8(\n VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_lhs_0, ymm_rhs_0)),\n ymm_sum_0);\n ymm_sum_1 = _mm256_add_epi8(\n VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_lhs_1, ymm_rhs_1)),\n ymm_sum_1);\n }\n ymm_sum_0 = _mm256_sad_epu8(ymm_sum_0, POPCNT_ZERO_AVX);\n ymm_sum_1 = _mm256_sad_epu8(ymm_sum_1, POPCNT_ZERO_AVX);\n\n for (; lhs != lhs_0; lhs += 16, rhs += 16) {\n __m256i ymm_lhs_0 = _mm256_load_si256((__m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_load_si256((__m256i *)(lhs + 8));\n __m256i ymm_rhs_0 = _mm256_load_si256((__m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_load_si256((__m256i *)(rhs + 8));\n\n ymm_sum_0 = _mm256_add_epi64(\n VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs_0, ymm_rhs_0)),\n ymm_sum_0);\n ymm_sum_1 = _mm256_add_epi64(\n VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs_1, ymm_rhs_1)),\n ymm_sum_1);\n }\n\n if (lhs_2 >= lhs + 8) {\n __m256i ymm_lhs = _mm256_load_si256((__m256i *)(lhs));\n __m256i ymm_rhs = _mm256_load_si256((__m256i *)(rhs));\n ymm_sum_0 = _mm256_add_epi64(\n VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs, ymm_rhs)),\n ymm_sum_0);\n lhs += 8;\n rhs += 8;\n }\n } else {\n for (; lhs != lhs_1; lhs += 16, rhs += 16) {\n __m256i ymm_lhs_0 = _mm256_loadu_si256((__m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_loadu_si256((__m256i *)(lhs + 8));\n __m256i ymm_rhs_0 = _mm256_loadu_si256((__m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_loadu_si256((__m256i *)(rhs + 8));\n\n ymm_sum_0 = _mm256_add_epi8(\n VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_lhs_0, ymm_rhs_0)),\n ymm_sum_0);\n ymm_sum_1 = _mm256_add_epi8(\n VerticalPopCount_INT8_V256(_mm256_xor_si256(ymm_lhs_1, ymm_rhs_1)),\n ymm_sum_1);\n }\n ymm_sum_0 = _mm256_sad_epu8(ymm_sum_0, POPCNT_ZERO_AVX);\n ymm_sum_1 = _mm256_sad_epu8(ymm_sum_1, POPCNT_ZERO_AVX);\n\n for (; lhs != lhs_0; lhs += 16, rhs += 16) {\n __m256i ymm_lhs_0 = _mm256_loadu_si256((__m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_loadu_si256((__m256i *)(lhs + 8));\n __m256i ymm_rhs_0 = _mm256_loadu_si256((__m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_loadu_si256((__m256i *)(rhs + 8));\n\n ymm_sum_0 = _mm256_add_epi64(\n VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs_0, ymm_rhs_0)),\n ymm_sum_0);\n ymm_sum_1 = _mm256_add_epi64(\n VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs_1, ymm_rhs_1)),\n ymm_sum_1);\n }\n\n if (lhs_2 >= lhs + 8) {\n __m256i ymm_lhs = _mm256_loadu_si256((__m256i *)(lhs));\n __m256i ymm_rhs = _mm256_loadu_si256((__m256i *)(rhs));\n ymm_sum_0 = _mm256_add_epi64(\n VerticalPopCount_INT64_V256(_mm256_xor_si256(ymm_lhs, ymm_rhs)),\n ymm_sum_0);\n lhs += 8;\n rhs += 8;\n }\n }\n\n size_t count =\n (size_t)HorizontalAdd_INT64_V256(_mm256_add_epi64(ymm_sum_0, ymm_sum_1));\n switch (lhs_2 - lhs) {\n case 7:\n count += ailego_popcount32(lhs[6] ^ rhs[6]);\n /* FALLTHRU */\n case 6:\n count += ailego_popcount32(lhs[5] ^ rhs[5]);\n /* FALLTHRU */\n case 5:\n count += ailego_popcount32(lhs[4] ^ rhs[4]);\n /* FALLTHRU */\n case 4:\n count += ailego_popcount32(lhs[3] ^ rhs[3]);\n /* FALLTHRU */\n case 3:\n count += ailego_popcount32(lhs[2] ^ rhs[2]);\n /* FALLTHRU */\n case 2:\n count += ailego_popcount32(lhs[1] ^ rhs[1]);\n /* FALLTHRU */\n case 1:\n count += ailego_popcount32(lhs[0] ^ rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t HammingDistanceAVX(const uint64_t *lhs,\n const uint64_t *rhs, size_t size) {\n return HammingDistanceAVX(reinterpret_cast(lhs),\n reinterpret_cast(rhs),\n (size << 1));\n}\n#endif // __AVX2__\n\n#if defined(AILEGO_M64)\nstatic inline size_t HammingDistance(const uint32_t *lhs, const uint32_t *rhs,\n size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n count += ailego_popcount64(*(uint64_t *)(&lhs[6]) ^ *(uint64_t *)(&rhs[6]));\n count += ailego_popcount64(*(uint64_t *)(&lhs[4]) ^ *(uint64_t *)(&rhs[4]));\n count += ailego_popcount64(*(uint64_t *)(&lhs[2]) ^ *(uint64_t *)(&rhs[2]));\n count += ailego_popcount64(*(uint64_t *)(&lhs[0]) ^ *(uint64_t *)(&rhs[0]));\n }\n switch (last - last_aligned) {\n case 7:\n count += ailego_popcount32(lhs[6] ^ rhs[6]);\n /* FALLTHRU */\n case 6:\n count += ailego_popcount32(lhs[5] ^ rhs[5]);\n /* FALLTHRU */\n case 5:\n count += ailego_popcount32(lhs[4] ^ rhs[4]);\n /* FALLTHRU */\n case 4:\n count += ailego_popcount32(lhs[3] ^ rhs[3]);\n /* FALLTHRU */\n case 3:\n count += ailego_popcount32(lhs[2] ^ rhs[2]);\n /* FALLTHRU */\n case 2:\n count += ailego_popcount32(lhs[1] ^ rhs[1]);\n /* FALLTHRU */\n case 1:\n count += ailego_popcount32(lhs[0] ^ rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t HammingDistance(const uint64_t *lhs, const uint64_t *rhs,\n size_t size) {\n const uint64_t *last = lhs + size;\n const uint64_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n count += ailego_popcount64(lhs[3] ^ rhs[3]);\n count += ailego_popcount64(lhs[2] ^ rhs[2]);\n count += ailego_popcount64(lhs[1] ^ rhs[1]);\n count += ailego_popcount64(lhs[0] ^ rhs[0]);\n }\n switch (last - last_aligned) {\n case 3:\n count += ailego_popcount64(lhs[2] ^ rhs[2]);\n /* FALLTHRU */\n case 2:\n count += ailego_popcount64(lhs[1] ^ rhs[1]);\n /* FALLTHRU */\n case 1:\n count += ailego_popcount64(lhs[0] ^ rhs[0]);\n }\n return count;\n}\n#else\nstatic inline size_t HammingDistance(const uint32_t *lhs, const uint32_t *rhs,\n size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n count += ailego_popcount32(lhs[3] ^ rhs[3]);\n count += ailego_popcount32(lhs[2] ^ rhs[2]);\n count += ailego_popcount32(lhs[1] ^ rhs[1]);\n count += ailego_popcount32(lhs[0] ^ rhs[0]);\n }\n switch (last - last_aligned) {\n case 3:\n count += ailego_popcount32(lhs[2] ^ rhs[2]);\n /* FALLTHRU */\n case 2:\n count += ailego_popcount32(lhs[1] ^ rhs[1]);\n /* FALLTHRU */\n case 1:\n count += ailego_popcount32(lhs[0] ^ rhs[0]);\n }\n return count;\n}\n#endif // AILEGO_M64\n\n//! Compute the distance between matrix and query (UINT32, M=1, N=1)\nvoid HammingDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim, float *out) {\n size_t cnt = (dim >> 5);\n#if defined(__AVX2__)\n if (cnt > 63) {\n *out = static_cast(HammingDistanceAVX(m, q, cnt));\n return;\n }\n#endif\n *out = static_cast(HammingDistance(m, q, cnt));\n}\n\n#if defined(AILEGO_M64)\n//! Compute the distance between matrix and query (UINT64, M=1, N=1)\nvoid HammingDistanceMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim, float *out) {\n size_t cnt = (dim >> 6);\n#if defined(__AVX2__)\n if (cnt > 31) {\n *out = static_cast(HammingDistanceAVX(m, q, cnt));\n return;\n }\n#endif\n *out = static_cast(HammingDistance(m, q, cnt));\n}\n\n#endif // AILEGO_M64\n\n//! Compute the distance between matrix and query (UINT32, M=1, N=1)\nvoid HammingSquareRootDistanceMatrix::Compute(\n const ValueType *m, const ValueType *q, size_t dim, float *out) {\n size_t cnt = (dim >> 5);\n#if defined(__AVX2__)\n if (cnt > 63) {\n *out = std::sqrt(static_cast(HammingDistanceAVX(m, q, cnt)));\n return;\n }\n#endif\n *out = std::sqrt(static_cast(HammingDistance(m, q, cnt)));\n}\n\n\n#if defined(AILEGO_M64)\n//! Compute the distance between matrix and query (UINT64, M=1, N=1)\nvoid HammingSquareRootDistanceMatrix::Compute(\n const ValueType *m, const ValueType *q, size_t dim, float *out) {\n size_t cnt = (dim >> 6);\n#if defined(__AVX2__)\n if (cnt > 31) {\n *out = std::sqrt(static_cast(HammingDistanceAVX(m, q, cnt)));\n return;\n }\n#endif\n *out = std::sqrt(static_cast(HammingDistance(m, q, cnt)));\n}\n\n#endif // AILEGO_M64\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/hamming_distance_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Hamming Distance Matrix\n */\ntemplate // NOTE: useless 'typename=void' to avoid clang\n // compile error\nstruct HammingDistanceMatrix;\n\n/*! Hamming Distance Matrix (UINT32)\n */\ntemplate \nstruct HammingDistanceMatrix {\n //! Type of value\n using ValueType = uint32_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && !(dim & 31) && out);\n\n size_t cnt = (dim >> 5);\n if (cnt > 0) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = static_cast(ailego_popcount32(m_val ^ q[j]));\n r += M;\n }\n }\n m += M;\n q += N;\n }\n\n for (size_t k = 1; k < cnt; ++k) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r += static_cast(ailego_popcount32(m_val ^ q[j]));\n r += M;\n }\n }\n m += M;\n q += N;\n }\n }\n};\n\n/*! Hamming Distance Matrix (UINT32, M=1, N=1)\n */\ntemplate <>\nstruct HammingDistanceMatrix {\n //! Type of value\n using ValueType = uint32_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\n#if defined(AILEGO_M64)\n/*! Hamming Distance Matrix (UINT64)\n */\ntemplate \nstruct HammingDistanceMatrix {\n //! Type of value\n using ValueType = uint64_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && !(dim & 63) && out);\n\n size_t cnt = (dim >> 6);\n if (cnt > 0) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = static_cast(ailego_popcount64(m_val ^ q[j]));\n r += M;\n }\n }\n m += M;\n q += N;\n }\n\n for (size_t k = 1; k < cnt; ++k) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r += static_cast(ailego_popcount64(m_val ^ q[j]));\n r += M;\n }\n }\n m += M;\n q += N;\n }\n }\n};\n\n/*! Hamming Distance Matrix (UINT64, M=1, N=1)\n */\ntemplate <>\nstruct HammingDistanceMatrix {\n //! Type of value\n using ValueType = uint64_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\n#endif // AILEGO_M64\n\n/*! Hamming Square Root Distance Matrix\n */\ntemplate \nstruct HammingSquareRootDistanceMatrix {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n HammingDistanceMatrix::Compute(m, q, dim, out);\n for (size_t i = 0; i < N * M; ++i) {\n float val = *out;\n *out++ = std::sqrt(val);\n }\n }\n};\n\n/*! Hamming Square Root Distance Matrix (UINT32, M=1, N=1)\n */\ntemplate <>\nstruct HammingSquareRootDistanceMatrix {\n //! Type of value\n using ValueType = uint32_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\n\n#if defined(AILEGO_M64)\n/*! Hamming Square Root Distance Matrix (UINT64, M=1, N=1)\n */\ntemplate <>\nstruct HammingSquareRootDistanceMatrix {\n //! Type of value\n using ValueType = uint64_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\n#endif // AILEGO_M64\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"distance_utility.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n/*! Inner Product Matrix\n */\ntemplate \nstruct InnerProductMatrix;\n\n/*! Inner Product Matrix\n */\ntemplate \nstruct MinusInnerProductMatrix;\n\n/*! Inner Product Matrix (M=1, N=1)\n */\ntemplate \nstruct InnerProductMatrix<\n T, 1, 1, typename std::enable_if::value>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n float sum = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n sum += static_cast(m[i] * q[i]);\n }\n *out = sum;\n }\n};\n\n/*! Minus Inner Product Matrix (M=1, N=1)\n */\ntemplate \nstruct MinusInnerProductMatrix<\n T, 1, 1, typename std::enable_if::value>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n float sum = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n sum += static_cast(m[i] * q[i]);\n }\n *out = -sum;\n }\n};\n\ntemplate <>\nstruct InnerProductMatrix {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct InnerProductMatrix {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct InnerProductMatrix {\n //! Type of value\n using ValueType = Float16;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct InnerProductMatrix {\n //! Type of value\n using ValueType = float;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct MinusInnerProductMatrix {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct MinusInnerProductMatrix {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct MinusInnerProductMatrix {\n //! Type of value\n using ValueType = Float16;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\ntemplate <>\nstruct MinusInnerProductMatrix {\n //! Type of value\n using ValueType = float;\n\n //! Compute the distance between matrix and query\n static void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out);\n};\n\n/*! Inner Product Matrix\n */\ntemplate \nstruct InnerProductMatrix<\n T, M, N,\n typename std::enable_if::value && sizeof(T) >= 2 &&\n M >= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n if (dim > 0) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = static_cast(m_val * q[j]);\n r += M;\n }\n }\n m += M;\n q += N;\n }\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r += m_val * q[j];\n r += M;\n }\n }\n m += M;\n q += N;\n }\n }\n};\n\n/*! Inner Product Matrix (N=1)\n */\ntemplate \nstruct InnerProductMatrix<\n T, M, 1,\n typename std::enable_if::value && sizeof(T) >= 2 &&\n M >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n const ValueType *q_end = q + dim;\n if (q != q_end) {\n ValueType q_val = *q++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = static_cast(m[i] * q_val);\n }\n m += M;\n }\n\n while (q != q_end) {\n ValueType q_val = *q++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += m[i] * q_val;\n }\n m += M;\n }\n }\n};\n\n/*! Inner Product Matrix (INT8)\n */\ntemplate \nstruct InnerProductMatrix= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 3) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n\n dim >>= 2;\n if (dim > 0) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = FusedMultiplyAdd(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r += FusedMultiplyAdd(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n volatile int32_t sum = ((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) +\n (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) +\n (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) +\n (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24));\n\n return static_cast(sum);\n }\n};\n\n/*! Inner Product Matrix (INT8, N=1)\n */\ntemplate \nstruct InnerProductMatrix= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 3) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 2);\n\n if (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = FusedMultiplyAdd(m_it[i], q_val);\n }\n m_it += M;\n }\n\n while (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += FusedMultiplyAdd(m_it[i], q_val);\n }\n m_it += M;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n volatile int32_t sum = ((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) +\n (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) +\n (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) +\n (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24));\n\n return static_cast(sum);\n }\n};\n\n/*! Inner Product Matrix (INT4)\n */\ntemplate \nstruct InnerProductMatrix= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 7) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n\n dim >>= 3;\n if (dim > 0) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = FusedMultiplyAdd(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r += FusedMultiplyAdd(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n};\n\n/*! Inner Product Matrix (INT4, N=1)\n */\ntemplate \nstruct InnerProductMatrix= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 7) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 3);\n\n if (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = FusedMultiplyAdd(m_it[i], q_val);\n }\n m_it += M;\n }\n\n while (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += FusedMultiplyAdd(m_it[i], q_val);\n }\n m_it += M;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n};\n\n\n/*! Minus Inner Product Matrix\n */\ntemplate \nstruct MinusInnerProductMatrix<\n T, M, N,\n typename std::enable_if::value && sizeof(T) >= 2 &&\n M >= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n if (dim > 0) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = -static_cast(m_val * q[j]);\n r += M;\n }\n }\n m += M;\n q += N;\n }\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n ValueType m_val = m[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r -= m_val * q[j];\n r += M;\n }\n }\n m += M;\n q += N;\n }\n }\n};\n\n/*! Minus Inner Product Matrix (N=1)\n */\ntemplate \nstruct MinusInnerProductMatrix<\n T, M, 1,\n typename std::enable_if::value && sizeof(T) >= 2 &&\n M >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && out);\n\n const ValueType *q_end = q + dim;\n if (q != q_end) {\n ValueType q_val = *q++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = -static_cast(m[i] * q_val);\n }\n m += M;\n }\n\n while (q != q_end) {\n ValueType q_val = *q++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) -= m[i] * q_val;\n }\n m += M;\n }\n }\n};\n\n/*! Minus Inner Product Matrix (INT8)\n */\ntemplate \nstruct MinusInnerProductMatrix<\n int8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 3) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n\n dim >>= 2;\n if (dim > 0) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = -FusedMultiplyAdd(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r -= FusedMultiplyAdd(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n volatile int32_t sum = ((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) +\n (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) +\n (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) +\n (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24));\n\n return static_cast(sum);\n }\n};\n\n/*! Minus Inner Product Matrix (INT8, N=1)\n */\ntemplate \nstruct MinusInnerProductMatrix= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 3) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 2);\n\n if (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = -FusedMultiplyAdd(m_it[i], q_val);\n }\n m_it += M;\n }\n\n while (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) -= FusedMultiplyAdd(m_it[i], q_val);\n }\n m_it += M;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n volatile int32_t sum = ((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) +\n (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) +\n (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) +\n (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24));\n\n return static_cast(sum);\n }\n};\n\n/*! Minus Inner Product Matrix (INT4)\n */\ntemplate \nstruct MinusInnerProductMatrix<\n uint8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 7) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n\n dim >>= 3;\n if (dim > 0) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r = -FusedMultiplyAdd(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n uint32_t m_val = m_it[i];\n float *r = out + i;\n\n for (size_t j = 0; j < N; ++j) {\n *r -= FusedMultiplyAdd(m_val, q_it[j]);\n r += M;\n }\n }\n m_it += M;\n q_it += N;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n};\n\n/*! Minus Inner Product Matrix (INT4, N=1)\n */\ntemplate \nstruct MinusInnerProductMatrix= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the distance between matrix and query\n static inline void Compute(const ValueType *m, const ValueType *q, size_t dim,\n float *out) {\n ailego_assert(m && q && dim && !(dim & 7) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 3);\n\n if (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = -FusedMultiplyAdd(m_it[i], q_val);\n }\n m_it += M;\n }\n\n while (q_it != q_end) {\n uint32_t q_val = *q_it++;\n\n for (size_t i = 0; i < M; ++i) {\n *(out + i) -= FusedMultiplyAdd(m_it[i], q_val);\n }\n m_it += M;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n};\n\n//--------------------------------------------------\n// Sparse\n//--------------------------------------------------\nstruct SparseSegmentInfo {\n public:\n uint32_t seg_id_{-1U};\n uint32_t vec_cnt_{0};\n\n public:\n SparseSegmentInfo() : seg_id_{-1U}, vec_cnt_{0} {}\n\n SparseSegmentInfo(uint32_t seg_id, uint32_t vec_cnt)\n : seg_id_{seg_id}, vec_cnt_{vec_cnt} {}\n};\n\nconstexpr static uint32_t SEGMENT_ID_BITS = 16;\nconstexpr static uint32_t SEGMENT_ID_MASK = 0xFFFF;\n\ntemplate \nstruct MinusInnerProductSparseMatrix {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n static inline float ComputeInnerProductSparseInSegment(\n uint32_t m_sparse_count, const uint16_t *m_sparse_index,\n const ValueType *m_sparse_value, uint32_t q_sparse_count,\n const uint16_t *q_sparse_index, const ValueType *q_sparse_value);\n\n //! Compute the distance between matrix and query\n static inline void Compute(const void *m_sparse_data_in,\n const void *q_sparse_data_in, float *out);\n\n static inline void transform_sparse_format(uint32_t sparse_count,\n const uint32_t *sparse_index,\n const void *sparse_value,\n std::string &buffer);\n};\n\ntemplate <>\nstruct MinusInnerProductSparseMatrix {\n //! Type of value\n using ValueType = Float16;\n\n static float ComputeInnerProductSparseInSegment(\n uint32_t m_sparse_count, const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value, uint32_t q_sparse_count,\n const uint16_t *q_sparse_index, const Float16 *q_sparse_value);\n\n //! Compute the distance between matrix and query\n static void Compute(const void *m_sparse_data_in,\n const void *q_sparse_data_in, float *out);\n\n static void transform_sparse_format(uint32_t sparse_count,\n const uint32_t *sparse_index,\n const void *sparse_value,\n std::string &buffer) {\n uint32_t unit_size = sizeof(ValueType);\n\n uint32_t seg_count = 0;\n if (sparse_count == 0) {\n buffer.reserve(sizeof(uint32_t) + sizeof(uint32_t));\n\n buffer.append(reinterpret_cast(&sparse_count),\n sizeof(uint32_t));\n\n buffer.append(reinterpret_cast(&seg_count),\n sizeof(uint32_t));\n\n return;\n }\n\n std::vector seg_infos;\n\n uint32_t cur_seg_id = -1U;\n uint32_t cur_vec_cnt = 0;\n\n for (size_t i = 0; i < sparse_count; ++i) {\n uint32_t seg_id = sparse_index[i] >> SEGMENT_ID_BITS;\n if (cur_seg_id == -1U) {\n cur_seg_id = seg_id;\n cur_vec_cnt++;\n } else {\n if (seg_id == cur_seg_id) {\n cur_vec_cnt++;\n } else if (seg_id > cur_seg_id) {\n seg_infos.emplace_back(cur_seg_id, cur_vec_cnt);\n\n cur_seg_id = seg_id;\n cur_vec_cnt = 1;\n } else {\n // std::abort();\n }\n }\n }\n\n if (cur_vec_cnt > 0) {\n seg_infos.emplace_back(cur_seg_id, cur_vec_cnt);\n }\n\n uint32_t buffer_len = 2 * sizeof(uint32_t) +\n seg_infos.size() * 2 * sizeof(uint32_t) +\n sparse_count * (sizeof(uint16_t) + sizeof(ValueType));\n\n buffer.reserve(buffer_len);\n\n buffer.append(reinterpret_cast(&sparse_count),\n sizeof(uint32_t));\n\n seg_count = seg_infos.size();\n buffer.append(reinterpret_cast(&seg_count), sizeof(uint32_t));\n\n for (size_t i = 0; i < seg_count; ++i) {\n uint32_t seg_id = seg_infos[i].seg_id_;\n buffer.append(reinterpret_cast(&seg_id), sizeof(uint32_t));\n }\n\n for (size_t i = 0; i < seg_count; ++i) {\n uint32_t vec_cnt = seg_infos[i].vec_cnt_;\n buffer.append(reinterpret_cast(&vec_cnt), sizeof(uint32_t));\n }\n\n for (size_t i = 0; i < sparse_count; ++i) {\n uint16_t temp_dim = sparse_index[i] & SEGMENT_ID_MASK;\n buffer.append(reinterpret_cast(&temp_dim),\n sizeof(uint16_t));\n }\n\n const char *sparse_value_ptr = reinterpret_cast(sparse_value);\n for (size_t i = 0; i < sparse_count; ++i) {\n buffer.append(sparse_value_ptr, unit_size);\n sparse_value_ptr += unit_size;\n }\n }\n};\n\ntemplate <>\nstruct MinusInnerProductSparseMatrix {\n //! Type of value\n using ValueType = float;\n\n static float ComputeInnerProductSparseInSegment(\n uint32_t m_sparse_count, const uint16_t *m_sparse_index,\n const float *m_sparse_value, uint32_t q_sparse_count,\n const uint16_t *q_sparse_index, const float *q_sparse_value);\n\n //! Compute the distance between matrix and query\n static void Compute(const void *m_sparse_data_in,\n const void *q_sparse_data_in, float *out);\n\n static void transform_sparse_format(uint32_t sparse_count,\n const uint32_t *sparse_index,\n const void *sparse_value,\n std::string &buffer) {\n uint32_t unit_size = sizeof(ValueType);\n\n uint32_t seg_count = 0;\n if (sparse_count == 0) {\n buffer.reserve(sizeof(uint32_t) + sizeof(uint32_t));\n\n buffer.append(reinterpret_cast(&sparse_count),\n sizeof(uint32_t));\n\n buffer.append(reinterpret_cast(&seg_count),\n sizeof(uint32_t));\n\n return;\n }\n\n std::vector seg_infos;\n\n uint32_t cur_seg_id = -1U;\n uint32_t cur_vec_cnt = 0;\n\n for (size_t i = 0; i < sparse_count; ++i) {\n uint32_t seg_id = sparse_index[i] >> SEGMENT_ID_BITS;\n if (cur_seg_id == -1U) {\n cur_seg_id = seg_id;\n cur_vec_cnt++;\n } else {\n if (seg_id == cur_seg_id) {\n cur_vec_cnt++;\n } else if (seg_id > cur_seg_id) {\n seg_infos.emplace_back(cur_seg_id, cur_vec_cnt);\n\n cur_seg_id = seg_id;\n cur_vec_cnt = 1;\n } else {\n // std::abort();\n }\n }\n }\n\n if (cur_vec_cnt > 0) {\n seg_infos.emplace_back(cur_seg_id, cur_vec_cnt);\n }\n\n uint32_t buffer_len = 2 * sizeof(uint32_t) +\n seg_infos.size() * 2 * sizeof(uint32_t) +\n sparse_count * (sizeof(uint16_t) + sizeof(ValueType));\n\n buffer.reserve(buffer_len);\n\n buffer.append(reinterpret_cast(&sparse_count),\n sizeof(uint32_t));\n\n seg_count = seg_infos.size();\n buffer.append(reinterpret_cast(&seg_count), sizeof(uint32_t));\n\n for (size_t i = 0; i < seg_count; ++i) {\n uint32_t seg_id = seg_infos[i].seg_id_;\n buffer.append(reinterpret_cast(&seg_id), sizeof(uint32_t));\n }\n\n for (size_t i = 0; i < seg_count; ++i) {\n uint32_t vec_cnt = seg_infos[i].vec_cnt_;\n buffer.append(reinterpret_cast(&vec_cnt), sizeof(uint32_t));\n }\n\n for (size_t i = 0; i < sparse_count; ++i) {\n uint16_t temp_dim = sparse_index[i] & SEGMENT_ID_MASK;\n buffer.append(reinterpret_cast(&temp_dim),\n sizeof(uint16_t));\n }\n\n const char *sparse_value_ptr = reinterpret_cast(sparse_value);\n for (size_t i = 0; i < sparse_count; ++i) {\n buffer.append(sparse_value_ptr, unit_size);\n sparse_value_ptr += unit_size;\n }\n }\n};\n\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_fp16_avx.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__AVX__)\nfloat InnerProductFp16AVX(const Float16 *lhs, const Float16 *rhs, size_t size) {\n float score{0.0f};\n\n ACCUM_FP16_1X1_AVX(lhs, rhs, size, &score, 0ull, )\n\n return score;\n}\n\nfloat MinusInnerProductFp16AVX(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n float score{0.0f};\n\n ACCUM_FP16_1X1_AVX(lhs, rhs, size, &score, 0ull, NEGATE_FP32_GENERAL)\n\n return score;\n}\n#endif\n\n//--------------------------------------------------\n// Sparse\n//--------------------------------------------------\n#if defined(__AVX__)\nconst static __m128i SHUFFLE_MASK256[256] = {\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, -127, -127),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 5,\n 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 7,\n 6, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 7,\n 6, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 7,\n 6, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 7, 6, 5, 4, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9,\n 8, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9,\n 8, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9,\n 8, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 5, 4, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9,\n 8, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9,\n 8, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 9,\n 8, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 5,\n 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 5,\n 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 11, 10),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 5, 4,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 7, 6, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 11,\n 10, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 13, 12),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 5, 4,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 7, 6, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 11, 10),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 11, 10, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 11, 10, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 11, 10, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 11, 10, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 13,\n 12, 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11, 10,\n 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6,\n 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3,\n 2),\n _mm_set_epi8(-127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 15, 14),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 5, 4,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 7, 6, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 11, 10),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 11, 10, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 11, 10, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 11, 10, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 11, 10, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11, 10,\n 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6,\n 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 3,\n 2),\n _mm_set_epi8(-127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 13, 12),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 13, 12, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 13, 12, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 13, 12, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 13, 12, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 13, 12, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6,\n 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 3,\n 2),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127, 15,\n 14, 13, 12, 11, 10),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 11, 10, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 11, 10, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 11, 10, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 5,\n 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 5,\n 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 11, 10, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7,\n 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7,\n 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7,\n 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, 3,\n 2),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9,\n 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9,\n 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9,\n 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, 3,\n 2),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9,\n 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 3,\n 2),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5,\n 4),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2),\n _mm_set_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n};\n\nconstexpr uint32_t MAX_SPARSE_BUFFER_LENGTH = 65536;\n\nfloat InnerProductSparseInSegmentFp16AVX(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const Float16 *q_sparse_value) {\n float sum = 0.0f;\n\n // handle if the first dim is zero\n bool m_zero = false;\n Float16 m_zero_value{0.0f};\n if (m_sparse_count > 0 && m_sparse_index[0] == 0) {\n m_sparse_count--;\n m_sparse_index++;\n m_zero_value = *m_sparse_value++;\n m_zero = true;\n }\n\n bool q_zero = false;\n Float16 q_zero_value{0.0f};\n if (q_sparse_count > 0 && q_sparse_index[0] == 0) {\n q_sparse_count--;\n q_sparse_index++;\n q_zero_value = *q_sparse_value++;\n q_zero = true;\n }\n\n if (m_zero && q_zero) {\n sum = m_zero_value * q_zero_value;\n }\n\n size_t i1 = 0, i2 = 0;\n size_t end1 = m_sparse_count / 8 * 8;\n size_t end2 = q_sparse_count / 8 * 8;\n\n uint16_t fixed_buffer_1[MAX_SPARSE_BUFFER_LENGTH];\n uint16_t fixed_buffer_2[MAX_SPARSE_BUFFER_LENGTH];\n\n Float16 *val_start_1 = reinterpret_cast(fixed_buffer_1);\n Float16 *val_start_2 = reinterpret_cast(fixed_buffer_2);\n\n Float16 *val_1 = val_start_1;\n Float16 *val_2 = val_start_2;\n\n if (i1 < end1 && i2 < end2) {\n while (m_sparse_index[i1 + 7] < q_sparse_index[i2]) {\n i1 += 8;\n if (i1 >= end1) goto do_scalar;\n }\n\n while (q_sparse_index[i2 + 7] < m_sparse_index[i1]) {\n i2 += 8;\n if (i2 >= end2) goto do_scalar;\n }\n\n __m128i mm_index_m =\n _mm_loadu_si128(reinterpret_cast(&m_sparse_index[i1]));\n __m128i mm_index_q =\n _mm_loadu_si128(reinterpret_cast(&q_sparse_index[i2]));\n\n while (true) {\n#ifdef DEBUG_PRINT\n std::cout << \"index 1: \" << std::endl;\n print_data16(&mm_index_m);\n\n std::cout << \"index 2: \" << std::endl;\n print_data16(&mm_index_q);\n#endif\n\n __m128i mm_cmp_res =\n _mm_cmpistrm(mm_index_q, mm_index_m,\n _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK);\n\n#ifdef DEBUG_PRINT\n std::cout << \"cmp res: \" << std::endl;\n print_data16(&mm_cmp_res);\n#endif\n\n int r = _mm_extract_epi32(mm_cmp_res, 0);\n\n if (r) {\n int r1 = r;\n\n __m128i v = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_value[i1]));\n __m128i vs = _mm_shuffle_epi8(v, SHUFFLE_MASK256[r1]);\n\n _mm_storeu_si128(reinterpret_cast<__m128i *>(val_1), vs);\n val_1 += _mm_popcnt_u32(r1);\n\n mm_cmp_res = _mm_cmpistrm(\n mm_index_m, mm_index_q,\n _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK);\n r = _mm_extract_epi32(mm_cmp_res, 0);\n\n r1 = r;\n\n v = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_value[i2]));\n vs = _mm_shuffle_epi8(v, SHUFFLE_MASK256[r1]);\n\n _mm_storeu_si128(reinterpret_cast<__m128i *>(val_2), vs);\n val_2 += _mm_popcnt_u32(r1);\n }\n\n const uint16_t id1_max = m_sparse_index[i1 + 7];\n\n if (id1_max <= q_sparse_index[i2 + 7]) {\n i1 += 8;\n if (i1 >= end1) goto do_scalar;\n mm_index_m = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_index[i1]));\n }\n\n if (id1_max >= q_sparse_index[i2 + 7]) {\n i2 += 8;\n if (i2 >= end2) goto do_scalar;\n mm_index_q = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_index[i2]));\n }\n }\n }\n\ndo_scalar:\n while (i1 < m_sparse_count && i2 < q_sparse_count) {\n if (m_sparse_index[i1] == q_sparse_index[i2]) {\n *val_1++ = m_sparse_value[i1];\n *val_2++ = q_sparse_value[i2];\n\n ++i1;\n ++i2;\n } else if (m_sparse_index[i1] < q_sparse_index[i2]) {\n ++i1;\n } else {\n ++i2;\n }\n }\n\n size_t res_num = val_1 - val_start_1;\n\n size_t res_num8 = res_num / 8 * 8;\n\n if (res_num8) {\n __m256 sum256 = _mm256_setzero_ps();\n\n for (size_t k = 0; k < res_num8; k += 8) {\n __m256 ymm_1 =\n _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(val_start_1 + k)));\n __m256 ymm_2 =\n _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(val_start_2 + k)));\n ACCUM_FP32_STEP_AVX(ymm_1, ymm_2, sum256);\n }\n\n sum += HorizontalAdd_FP32_V256(sum256);\n }\n\n for (size_t k = res_num8; k < res_num; ++k)\n sum += val_start_1[k] * val_start_2[k];\n\n return sum;\n}\n\n#endif // __AVX__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/inner_product_matrix_fp16_avx512.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX512F__)\nfloat InnerProductFp16AVX512(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n float score{0.0f};\n\n ACCUM_FP16_1X1_AVX512(lhs, rhs, size, &score, 0ull, )\n\n return score;\n}\n\nfloat MinusInnerProductFp16AVX512(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n float score{0.0f};\n\n ACCUM_FP16_1X1_AVX512(lhs, rhs, size, &score, 0ull, NEGATE_FP32_GENERAL)\n\n return score;\n}\n#endif //__AVX512F__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_fp16_avx512fp16.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX512FP16__)\n//! Inner Product\nfloat InnerProductFp16AVX512FP16(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n const Float16 *last = lhs + size;\n const Float16 *last_aligned = lhs + ((size >> 6) << 6);\n\n __m512h zmm_sum_0 = _mm512_setzero_ph();\n __m512h zmm_sum_1 = _mm512_setzero_ph();\n\n if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n FMA_FP16_AVX512FP16(_mm512_load_ph(lhs + 0), _mm512_load_ph(rhs + 0),\n zmm_sum_0)\n\n FMA_FP16_AVX512FP16(_mm512_load_ph(lhs + 32), _mm512_load_ph(rhs + 32),\n zmm_sum_1)\n }\n\n if (last >= last_aligned + 32) {\n FMA_FP16_AVX512FP16(_mm512_load_ph(lhs), _mm512_load_ph(rhs), zmm_sum_0)\n lhs += 32;\n rhs += 32;\n }\n } else {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n FMA_FP16_AVX512FP16(_mm512_loadu_ph(lhs + 0), _mm512_loadu_ph(rhs + 0),\n zmm_sum_0)\n\n FMA_FP16_AVX512FP16(_mm512_loadu_ph(lhs + 32), _mm512_loadu_ph(rhs + 32),\n zmm_sum_1)\n }\n\n if (last >= last_aligned + 32) {\n FMA_FP16_AVX512FP16(_mm512_loadu_ph(lhs), _mm512_loadu_ph(rhs), zmm_sum_0)\n lhs += 32;\n rhs += 32;\n }\n }\n\n zmm_sum_0 = _mm512_add_ph(zmm_sum_0, zmm_sum_1);\n\n if (lhs != last) {\n __mmask32 mask = (__mmask32)((1 << (last - lhs)) - 1);\n __m512i zmm_undefined = _mm512_undefined_epi32();\n zmm_sum_0 = _mm512_mask3_fmadd_ph(\n _mm512_castsi512_ph(_mm512_mask_loadu_epi16(zmm_undefined, mask, lhs)),\n _mm512_castsi512_ph(_mm512_mask_loadu_epi16(zmm_undefined, mask, rhs)),\n zmm_sum_0, mask);\n }\n\n return HorizontalAdd_FP16_V512(zmm_sum_0);\n}\n\nfloat MinusInnerProductFp16AVX512FP16(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n return -1 * InnerProductFp16AVX512FP16(lhs, rhs, size);\n}\n#endif\n\n// sparse\n#if defined(__AVX512FP16__)\nconstexpr uint32_t MAX_SPARSE_BUFFER_LENGTH = 65536;\n\nfloat InnerProductSparseInSegmentFp16AVX512FP16(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const Float16 *q_sparse_value) {\n const static __m128i SHUFFLE_MASK256[256] = {\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, -127, -127),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 7, 6, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 7, 6, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 7, 6, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 7, 6, 5, 4,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 9, 8, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 9, 8, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 9, 8, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 5, 4,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 9, 8, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 9, 8, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 9, 8, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 9, 8, 7, 6,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 9, 8, 7, 6, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 11, 10),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 5, 4,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 7, 6,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 7, 6, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 11, 10, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 13, 12),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 5, 4,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 7, 6,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 7, 6, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 9, 8,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 9, 8, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 13, 12, 11, 10),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 11, 10, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 11, 10, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 11, 10, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 5, 4, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 11, 10, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5,\n 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5,\n 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 13, 12, 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5,\n 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5,\n 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 13, 12, 11,\n 10, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7,\n 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7,\n 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7,\n 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3,\n 2),\n _mm_set_epi8(-127, -127, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, 15, 14),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 5, 4,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 7, 6,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 7, 6, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 9, 8,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 9, 8, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 11, 10),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 11, 10, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 11, 10, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 11, 10, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 5, 4, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 11, 10, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5,\n 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5,\n 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5,\n 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5,\n 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 11,\n 10, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7,\n 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7,\n 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7,\n 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 3,\n 2),\n _mm_set_epi8(-127, -127, 15, 14, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 13, 12),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 13, 12, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 13, 12, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 13, 12, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 5, 4, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 13, 12, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5,\n 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5,\n 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 13, 12, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5,\n 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5,\n 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7,\n 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7,\n 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7,\n 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 3,\n 2),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n 15, 14, 13, 12, 11, 10),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 11, 10, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 11, 10, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 11, 10, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 5, 4, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 5, 4, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 11, 10, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 7, 6, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 7, 6, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13,\n 12, 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 9, 8, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 9, 8, 3, 2),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 3, 2,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 9, 8, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4,\n 3, 2),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, 15, 14, 13, 12, 11, 10,\n 9, 8, 7, 6),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6,\n 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6,\n 3, 2),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6,\n 5, 4),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 1, 0),\n _mm_set_epi8(-127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2),\n _mm_set_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n };\n\n float sum = 0.0f;\n\n // handle if the first dim is zero\n bool m_zero = false;\n Float16 m_zero_value{0.0f};\n if (m_sparse_count > 0 && m_sparse_index[0] == 0) {\n m_sparse_count--;\n m_sparse_index++;\n m_zero_value = *m_sparse_value++;\n m_zero = true;\n }\n\n bool q_zero = false;\n Float16 q_zero_value{0.0f};\n if (q_sparse_count > 0 && q_sparse_index[0] == 0) {\n q_sparse_count--;\n q_sparse_index++;\n q_zero_value = *q_sparse_value++;\n q_zero = true;\n }\n\n if (m_zero && q_zero) {\n sum = m_zero_value * q_zero_value;\n }\n\n size_t i1 = 0, i2 = 0;\n size_t end1 = m_sparse_count / 8 * 8;\n size_t end2 = q_sparse_count / 8 * 8;\n\n uint16_t fixed_buffer_1[MAX_SPARSE_BUFFER_LENGTH];\n uint16_t fixed_buffer_2[MAX_SPARSE_BUFFER_LENGTH];\n\n Float16 *val_start_1 = reinterpret_cast(fixed_buffer_1);\n Float16 *val_start_2 = reinterpret_cast(fixed_buffer_2);\n\n Float16 *val_1 = val_start_1;\n Float16 *val_2 = val_start_2;\n\n if (i1 < end1 && i2 < end2) {\n while (m_sparse_index[i1 + 7] < q_sparse_index[i2]) {\n i1 += 8;\n if (i1 >= end1) goto do_scalar;\n }\n\n while (q_sparse_index[i2 + 7] < m_sparse_index[i1]) {\n i2 += 8;\n if (i2 >= end2) goto do_scalar;\n }\n\n __m128i mm_index_m =\n _mm_loadu_si128(reinterpret_cast(&m_sparse_index[i1]));\n __m128i mm_index_q =\n _mm_loadu_si128(reinterpret_cast(&q_sparse_index[i2]));\n\n while (true) {\n#ifdef DEBUG_PRINT\n std::cout << \"index 1: \" << std::endl;\n print_data16(&mm_index_m);\n\n std::cout << \"index 2: \" << std::endl;\n print_data16(&mm_index_q);\n#endif\n\n __m128i mm_cmp_res =\n _mm_cmpistrm(mm_index_q, mm_index_m,\n _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK);\n\n#ifdef DEBUG_PRINT\n std::cout << \"cmp res: \" << std::endl;\n print_data16(&mm_cmp_res);\n#endif\n\n int r = _mm_extract_epi32(mm_cmp_res, 0);\n\n if (r) {\n int r1 = r;\n\n __m128i v = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_value[i1]));\n __m128h vs = _mm_castsi128_ph(_mm_shuffle_epi8(v, SHUFFLE_MASK256[r1]));\n\n _mm_storeu_ph(val_1, vs);\n val_1 += _mm_popcnt_u32(r1);\n\n mm_cmp_res = _mm_cmpistrm(\n mm_index_m, mm_index_q,\n _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK);\n r = _mm_extract_epi32(mm_cmp_res, 0);\n\n r1 = r;\n\n v = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_value[i2]));\n vs = _mm_castsi128_ph(_mm_shuffle_epi8(v, SHUFFLE_MASK256[r1]));\n\n _mm_storeu_ph(val_2, vs);\n val_2 += _mm_popcnt_u32(r1);\n }\n\n const uint16_t id1_max = m_sparse_index[i1 + 7];\n\n if (id1_max <= q_sparse_index[i2 + 7]) {\n i1 += 8;\n if (i1 >= end1) goto do_scalar;\n mm_index_m = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_index[i1]));\n }\n\n if (id1_max >= q_sparse_index[i2 + 7]) {\n i2 += 8;\n if (i2 >= end2) goto do_scalar;\n mm_index_q = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_index[i2]));\n }\n }\n }\n\ndo_scalar:\n while (i1 < m_sparse_count && i2 < q_sparse_count) {\n if (m_sparse_index[i1] == q_sparse_index[i2]) {\n *val_1++ = m_sparse_value[i1];\n *val_2++ = q_sparse_value[i2];\n\n ++i1;\n ++i2;\n } else if (m_sparse_index[i1] < q_sparse_index[i2]) {\n ++i1;\n } else {\n ++i2;\n }\n }\n\n size_t res_num = val_1 - val_start_1;\n\n size_t res_num8 = res_num / 8 * 8;\n\n if (res_num8) {\n __m128h sum128 = _mm_set1_ph(0);\n\n for (size_t k = 0; k < res_num8; k += 8) {\n sum128 = _mm_add_ph(sum128, _mm_mul_ph(_mm_loadu_ph(val_start_1 + k),\n _mm_loadu_ph(val_start_2 + k)));\n }\n\n Float16 __attribute__((aligned(16))) tmp_res[8];\n _mm_store_ph(tmp_res, sum128);\n sum += (tmp_res[0] + tmp_res[1] + tmp_res[2] + tmp_res[3] + tmp_res[4] +\n tmp_res[5] + tmp_res[6] + tmp_res[7]);\n }\n\n for (size_t k = res_num8; k < res_num; ++k)\n sum += val_start_1[k] * val_start_2[k];\n\n return sum;\n}\n\n#endif // __AVX512FP16__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_fp16_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__ARM_NEON)\nfloat InnerProductFp16NEON(const Float16 *lhs, const Float16 *rhs, size_t size);\nfloat MinusInnerProductFp16NEON(const Float16 *lhs, const Float16 *rhs,\n size_t size);\n#endif\n\n#if defined(__AVX__)\nfloat InnerProductFp16AVX(const Float16 *lhs, const Float16 *rhs, size_t size);\nfloat MinusInnerProductFp16AVX(const Float16 *lhs, const Float16 *rhs,\n size_t size);\n#endif\n\n#if defined(__AVX512F__)\nfloat InnerProductFp16AVX512(const Float16 *lhs, const Float16 *rhs,\n size_t size);\nfloat MinusInnerProductFp16AVX512(const Float16 *lhs, const Float16 *rhs,\n size_t size);\n#endif\n\n#if defined(__AVX512FP16__)\nfloat InnerProductFp16AVX512FP16(const Float16 *lhs, const Float16 *rhs,\n size_t size);\nfloat MinusInnerProductFp16AVX512FP16(const Float16 *lhs, const Float16 *rhs,\n size_t size);\n#endif\n\nfloat InnerProductFp16Scalar(const Float16 *lhs, const Float16 *rhs,\n size_t size);\nfloat MinusInnerProductFp16Scalar(const Float16 *lhs, const Float16 *rhs,\n size_t size);\n\n//! Compute the distance between matrix and query (FP16, M=1, N=1)\nvoid InnerProductMatrix::Compute(const ValueType *m,\n const ValueType *q, size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n *out = InnerProductFp16NEON(m, q, dim);\n#else\n#if defined(__AVX512FP16__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) {\n *out = InnerProductFp16AVX512FP16(m, q, dim);\n return;\n }\n#endif //__AVX512FP16__\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = InnerProductFp16AVX512(m, q, dim);\n return;\n }\n#endif //__AVX512F__\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = InnerProductFp16AVX(m, q, dim);\n return;\n }\n#endif //__AVX__\n *out = InnerProductFp16Scalar(m, q, dim);\n\n#endif //__ARM_NEON\n}\n\n//! Compute the distance between matrix and query (FP16, M=1, N=1)\nvoid MinusInnerProductMatrix::Compute(const ValueType *m,\n const ValueType *q,\n size_t dim, float *out) {\n#if defined(__ARM_NEON)\n *out = MinusInnerProductFp16NEON(m, q, dim);\n#else\n#if defined(__AVX512FP16__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) {\n *out = MinusInnerProductFp16AVX512FP16(m, q, dim);\n return;\n }\n#endif //__AVX512FP16__\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = MinusInnerProductFp16AVX512(m, q, dim);\n return;\n }\n#endif //__AVX512F__\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = MinusInnerProductFp16AVX(m, q, dim);\n return;\n }\n#endif //__AVX__\n\n *out = MinusInnerProductFp16Scalar(m, q, dim);\n\n#endif //__ARM_NEON\n}\n\n//--------------------------------------------------\n// Sparse\n//--------------------------------------------------\n#if defined(__AVX512FP16__)\nfloat InnerProductSparseInSegmentFp16AVX512FP16(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const Float16 *q_sparse_value);\n#endif //__AVX512FP16__\n\n#if defined(__AVX__)\nfloat InnerProductSparseInSegmentFp16AVX(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const Float16 *q_sparse_value);\n#endif //__AVX__\n\nfloat InnerProductSparseInSegmentFp16Scalar(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const Float16 *q_sparse_value);\n\nfloat MinusInnerProductSparseFp16Scalar(const void *m_sparse_data_in,\n const void *q_sparse_data_in);\n\n//! Compute the distance between matrix and query\nvoid MinusInnerProductSparseMatrix::Compute(\n const void *m_sparse_data_in, const void *q_sparse_data_in, float *out) {\n *out = MinusInnerProductSparseFp16Scalar(m_sparse_data_in, q_sparse_data_in);\n}\n\nfloat ComputeInnerProductSparseInSegmentFp16(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const Float16 *q_sparse_value) {\n#if defined(__AVX512FP16__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) {\n return InnerProductSparseInSegmentFp16AVX512FP16(\n m_sparse_count, m_sparse_index, m_sparse_value, q_sparse_count,\n q_sparse_index, q_sparse_value);\n }\n#endif\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n return InnerProductSparseInSegmentFp16AVX(m_sparse_count, m_sparse_index,\n m_sparse_value, q_sparse_count,\n q_sparse_index, q_sparse_value);\n }\n#endif\n return InnerProductSparseInSegmentFp16Scalar(m_sparse_count, m_sparse_index,\n m_sparse_value, q_sparse_count,\n q_sparse_index, q_sparse_value);\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_fp16_neon.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__ARM_NEON)\nfloat InnerProductFp16NEON(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n float score;\n\n ACCUM_FP16_1X1_NEON(lhs, rhs, size, &score, 0ull, )\n\n return score;\n}\n\nfloat MinusInnerProductFp16NEON(const Float16 *lhs, const Float16 *rhs,\n size_t size) {\n float score;\n\n ACCUM_FP16_1X1_NEON(lhs, rhs, size, &score, 0ull, NEGATE_FP32_GENERAL)\n\n return score;\n}\n#endif\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/inner_product_matrix_fp32_avx.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__AVX__)\nfloat InnerProductFp32SSEInternal(const float *lhs, const float *rhs,\n size_t size);\n\n//! Inner Product\nfloat InnerProductFp32AVXInternal(const float *lhs, const float *rhs,\n size_t size) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 4) << 4);\n\n __m256 ymm_sum_0 = _mm256_setzero_ps();\n __m256 ymm_sum_1 = _mm256_setzero_ps();\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m256 ymm_lhs_0 = _mm256_load_ps(lhs + 0);\n __m256 ymm_lhs_1 = _mm256_load_ps(lhs + 8);\n __m256 ymm_rhs_0 = _mm256_load_ps(rhs + 0);\n __m256 ymm_rhs_1 = _mm256_load_ps(rhs + 8);\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1);\n }\n\n if (last >= last_aligned + 8) {\n ymm_sum_0 =\n _mm256_fmadd_ps(_mm256_load_ps(lhs), _mm256_load_ps(rhs), ymm_sum_0);\n lhs += 8;\n rhs += 8;\n }\n } else {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m256 ymm_lhs_0 = _mm256_loadu_ps(lhs + 0);\n __m256 ymm_lhs_1 = _mm256_loadu_ps(lhs + 8);\n __m256 ymm_rhs_0 = _mm256_loadu_ps(rhs + 0);\n __m256 ymm_rhs_1 = _mm256_loadu_ps(rhs + 8);\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1);\n }\n\n if (last >= last_aligned + 8) {\n ymm_sum_0 = _mm256_fmadd_ps(_mm256_loadu_ps(lhs), _mm256_loadu_ps(rhs),\n ymm_sum_0);\n lhs += 8;\n rhs += 8;\n }\n }\n float result = HorizontalAdd_FP32_V256(_mm256_add_ps(ymm_sum_0, ymm_sum_1));\n\n switch (last - lhs) {\n case 7:\n FMA_FP32_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n FMA_FP32_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n FMA_FP32_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n FMA_FP32_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n FMA_FP32_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n FMA_FP32_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n FMA_FP32_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat InnerProductFp32AVX(const float *lhs, const float *rhs, size_t size) {\n if (size > 7) {\n return InnerProductFp32AVXInternal(lhs, rhs, size);\n }\n\n return InnerProductFp32SSEInternal(lhs, rhs, size);\n}\n\nfloat MinusInnerProductFp32AVX(const float *lhs, const float *rhs,\n size_t size) {\n return -1 * InnerProductFp32AVX(lhs, rhs, size);\n}\n\n#endif // __AVX__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_fp32_avx512.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__AVX512F__)\nfloat InnerProductFp32AVXInternal(const float *lhs, const float *rhs,\n size_t size);\n\nfloat InnerProductFp32SSEInternal(const float *lhs, const float *rhs,\n size_t size);\n\n//! Inner Product\nfloat InnerProductFp32AVX512Internal(const float *lhs, const float *rhs,\n size_t size) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 5) << 5);\n\n __m512 zmm_sum_0 = _mm512_setzero_ps();\n __m512 zmm_sum_1 = _mm512_setzero_ps();\n\n if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n FMA_FP32_AVX512(_mm512_load_ps(lhs + 0), _mm512_load_ps(rhs + 0),\n zmm_sum_0)\n\n FMA_FP32_AVX512(_mm512_load_ps(lhs + 16), _mm512_load_ps(rhs + 16),\n zmm_sum_1)\n }\n\n if (last >= last_aligned + 16) {\n FMA_FP32_AVX512(_mm512_load_ps(lhs), _mm512_load_ps(rhs), zmm_sum_0)\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n FMA_FP32_AVX512(_mm512_loadu_ps(lhs + 0), _mm512_loadu_ps(rhs + 0),\n zmm_sum_0)\n\n FMA_FP32_AVX512(_mm512_loadu_ps(lhs + 16), _mm512_loadu_ps(rhs + 16),\n zmm_sum_1)\n }\n\n if (last >= last_aligned + 16) {\n FMA_FP32_AVX512(_mm512_loadu_ps(lhs), _mm512_loadu_ps(rhs), zmm_sum_0)\n lhs += 16;\n rhs += 16;\n }\n }\n\n zmm_sum_0 = _mm512_add_ps(zmm_sum_0, zmm_sum_1);\n if (lhs != last) {\n __mmask16 mask = (__mmask16)((1 << (last - lhs)) - 1);\n __m512 zmm_undefined = _mm512_undefined_ps();\n zmm_sum_0 = _mm512_mask3_fmadd_ps(\n _mm512_mask_loadu_ps(zmm_undefined, mask, lhs),\n _mm512_mask_loadu_ps(zmm_undefined, mask, rhs), zmm_sum_0, mask);\n }\n return HorizontalAdd_FP32_V512(zmm_sum_0);\n}\n\nfloat InnerProductFp32AVX512(const float *lhs, const float *rhs, size_t size) {\n if (size > 15) {\n return InnerProductFp32AVX512Internal(lhs, rhs, size);\n }\n\n if (size > 7) {\n return InnerProductFp32AVXInternal(lhs, rhs, size);\n }\n\n return InnerProductFp32SSEInternal(lhs, rhs, size);\n}\n\nfloat MinusInnerProductFp32AVX512(const float *lhs, const float *rhs,\n size_t size) {\n return -1 * InnerProductFp32AVX512(lhs, rhs, size);\n}\n\n#endif\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_fp32_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__ARM_NEON)\nfloat InnerProductFp32NEON(const float *lhs, const float *rhs, size_t size);\nfloat MinusInnerProductFp32NEON(const float *lhs, const float *rhs,\n size_t size);\n#endif\n\n#if defined(__AVX512F__)\nfloat InnerProductFp32AVX512(const float *lhs, const float *rhs, size_t size);\nfloat MinusInnerProductFp32AVX512(const float *lhs, const float *rhs,\n size_t size);\n#endif\n\n#if defined(__AVX__)\nfloat InnerProductFp32AVX(const float *lhs, const float *rhs, size_t size);\nfloat MinusInnerProductFp32AVX(const float *lhs, const float *rhs, size_t size);\n#endif\n\n#if defined(__SSE__)\nfloat InnerProductFp32SSE(const float *lhs, const float *rhs, size_t size);\nfloat MinusInnerProductFp32SSE(const float *lhs, const float *rhs, size_t size);\n#endif\n\nfloat InnerProductFp32Scalar(const float *lhs, const float *rhs, size_t size);\nfloat MinusInnerProductFp32Scalar(const float *lhs, const float *rhs,\n size_t size);\n\n//! Compute the distance between matrix and query (FP32, M=1, N=1)\nvoid InnerProductMatrix::Compute(const float *m, const float *q,\n size_t dim, float *out) {\n#if defined(__ARM_NEON)\n *out = InnerProductFp32NEON(m, q, dim);\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = InnerProductFp32AVX512(m, q, dim);\n return;\n }\n#endif // __AVX512F__\n\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = InnerProductFp32AVX(m, q, dim);\n return;\n }\n#endif // __AVX__\n\n#if defined(__SSE__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) {\n *out = InnerProductFp32SSE(m, q, dim);\n return;\n }\n#endif // __SSE__\n *out = InnerProductFp32Scalar(m, q, dim);\n#endif // __ARM_NEON\n}\n\n//! Compute the distance between matrix and query (FP32, M=1, N=1)\nvoid MinusInnerProductMatrix::Compute(const float *m,\n const float *q, size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n *out = MinusInnerProductFp32NEON(m, q, dim);\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = MinusInnerProductFp32AVX512(m, q, dim);\n return;\n }\n#endif // __AVX512F__\n\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = MinusInnerProductFp32AVX(m, q, dim);\n return;\n }\n#endif // __AVX__\n\n#if defined(__SSE__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) {\n *out = MinusInnerProductFp32SSE(m, q, dim);\n return;\n }\n#endif // __SSE__\n *out = MinusInnerProductFp32Scalar(m, q, dim);\n#endif // __ARM_NEON\n}\n\n//--------------------------------------------------\n// Sparse\n//--------------------------------------------------\n#if defined(__SSE4_1__)\nfloat InnerProductSparseInSegmentFp32SSE(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value);\n#endif\nfloat InnerProductSparseInSegmentFp32Scalar(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value);\n\nfloat MinusInnerProductSparseFp32Scalar(const void *m_sparse_data_in,\n const void *q_sparse_data_in);\n\nvoid MinusInnerProductSparseMatrix::Compute(const void *m_sparse_data_in,\n const void *q_sparse_data_in,\n float *out) {\n *out = MinusInnerProductSparseFp32Scalar(m_sparse_data_in, q_sparse_data_in);\n}\n\nfloat ComputeInnerProductSparseInSegmentFp32(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value) {\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n return InnerProductSparseInSegmentFp32SSE(m_sparse_count, m_sparse_index,\n m_sparse_value, q_sparse_count,\n q_sparse_index, q_sparse_value);\n }\n#endif\n return InnerProductSparseInSegmentFp32Scalar(m_sparse_count, m_sparse_index,\n m_sparse_value, q_sparse_count,\n q_sparse_index, q_sparse_value);\n}\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_fp32_neon.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__ARM_NEON)\nfloat InnerProductFp32NEON(const float *lhs, const float *rhs, size_t size) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 3) << 3);\n\n float32x4_t v_sum_0 = vdupq_n_f32(0);\n float32x4_t v_sum_1 = vdupq_n_f32(0);\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n v_sum_0 = vfmaq_f32(v_sum_0, vld1q_f32(lhs + 0), vld1q_f32(rhs + 0));\n v_sum_1 = vfmaq_f32(v_sum_1, vld1q_f32(lhs + 4), vld1q_f32(rhs + 4));\n }\n if (last >= last_aligned + 4) {\n v_sum_0 = vfmaq_f32(v_sum_0, vld1q_f32(lhs), vld1q_f32(rhs));\n lhs += 4;\n rhs += 4;\n }\n\n float result = vaddvq_f32(vaddq_f32(v_sum_0, v_sum_1));\n switch (last - lhs) {\n case 3:\n FMA_FP32_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n FMA_FP32_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n FMA_FP32_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat MinusInnerProductFp32NEON(const float *lhs, const float *rhs,\n size_t size) {\n return -1 * InnerProductFp32NEON(lhs, rhs, size);\n}\n\n#endif // __ARM_NEON\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_fp32_sse.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__SSE__)\nfloat InnerProductFp32SSEInternal(const float *lhs, const float *rhs,\n size_t size) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 3) << 3);\n\n __m128 xmm_sum_0 = _mm_setzero_ps();\n __m128 xmm_sum_1 = _mm_setzero_ps();\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m128 xmm_lhs_0 = _mm_load_ps(lhs + 0);\n __m128 xmm_lhs_1 = _mm_load_ps(lhs + 4);\n __m128 xmm_rhs_0 = _mm_load_ps(rhs + 0);\n __m128 xmm_rhs_1 = _mm_load_ps(rhs + 4);\n xmm_sum_0 = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum_0);\n xmm_sum_1 = _mm_fmadd_ps(xmm_lhs_1, xmm_rhs_1, xmm_sum_1);\n }\n\n if (last >= last_aligned + 4) {\n xmm_sum_0 = _mm_fmadd_ps(_mm_load_ps(lhs), _mm_load_ps(rhs), xmm_sum_0);\n lhs += 4;\n rhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m128 xmm_lhs_0 = _mm_loadu_ps(lhs + 0);\n __m128 xmm_lhs_1 = _mm_loadu_ps(lhs + 4);\n __m128 xmm_rhs_0 = _mm_loadu_ps(rhs + 0);\n __m128 xmm_rhs_1 = _mm_loadu_ps(rhs + 4);\n xmm_sum_0 = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum_0);\n xmm_sum_1 = _mm_fmadd_ps(xmm_lhs_1, xmm_rhs_1, xmm_sum_1);\n }\n\n if (last >= last_aligned + 4) {\n xmm_sum_0 = _mm_fmadd_ps(_mm_loadu_ps(lhs), _mm_loadu_ps(rhs), xmm_sum_0);\n lhs += 4;\n rhs += 4;\n }\n }\n float result = HorizontalAdd_FP32_V128(_mm_add_ps(xmm_sum_0, xmm_sum_1));\n\n switch (last - lhs) {\n case 3:\n FMA_FP32_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n FMA_FP32_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n FMA_FP32_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat InnerProductFp32SSE(const float *lhs, const float *rhs, size_t size) {\n return InnerProductFp32SSEInternal(lhs, rhs, size);\n}\n\nfloat MinusInnerProductFp32SSE(const float *lhs, const float *rhs,\n size_t size) {\n return -1 * InnerProductFp32SSE(lhs, rhs, size);\n}\n\n#endif // __SSE__\n\n//--------------------------------------------------\n// Sparse\n//--------------------------------------------------\n#if defined(__SSE4_1__)\nconst static __m128i SHUFFLE_MASK16[16] = {\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, -127, -127),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 7, 6, 5, 4, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 13, 12),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5,\n 4),\n _mm_set_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n};\n\nconstexpr uint32_t MAX_SPARSE_BUFFER_LENGTH = 65536;\n\nfloat InnerProductSparseInSegmentFp32SSE(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value) {\n float sum = 0.0f;\n\n // handle if the first dim is zero\n bool m_zero = false;\n float m_zero_value = 0.0f;\n if (m_sparse_count > 0 && m_sparse_index[0] == 0) {\n m_sparse_count--;\n m_sparse_index++;\n m_zero_value = *m_sparse_value++;\n m_zero = true;\n }\n\n bool q_zero = false;\n float q_zero_value = 0.0f;\n if (q_sparse_count > 0 && q_sparse_index[0] == 0) {\n q_sparse_count--;\n q_sparse_index++;\n q_zero_value = *q_sparse_value++;\n q_zero = true;\n }\n\n if (m_zero && q_zero) {\n sum = m_zero_value * q_zero_value;\n }\n\n size_t i1 = 0, i2 = 0;\n size_t end1 = m_sparse_count / 8 * 8;\n size_t end2 = q_sparse_count / 8 * 8;\n\n // std::vector mem1;\n // std::vector mem2;\n\n float fixed_buffer_1[MAX_SPARSE_BUFFER_LENGTH];\n float fixed_buffer_2[MAX_SPARSE_BUFFER_LENGTH];\n\n float *val_start_1 = fixed_buffer_1;\n float *val_start_2 = fixed_buffer_2;\n\n // uint32_t max_count = std::max(m_sparse_count, q_sparse_count);\n\n // if (MAX_SPARSE_BUFFER_LENGTH < max_count) {\n // mem1.reserve(max_count);\n // mem2.reserve(max_count);\n\n // val_start_1 = mem1.data();\n // val_start_2 = mem2.data();\n // }\n\n float *val_1 = val_start_1;\n float *val_2 = val_start_2;\n\n if (i1 < end1 && i2 < end2) {\n while (m_sparse_index[i1 + 7] < q_sparse_index[i2]) {\n i1 += 8;\n if (i1 >= end1) goto do_scalar;\n }\n\n while (q_sparse_index[i2 + 7] < m_sparse_index[i1]) {\n i2 += 8;\n if (i2 >= end2) goto do_scalar;\n }\n\n __m128i mm_index_m =\n _mm_loadu_si128(reinterpret_cast(&m_sparse_index[i1]));\n __m128i mm_index_q =\n _mm_loadu_si128(reinterpret_cast(&q_sparse_index[i2]));\n\n while (true) {\n#ifdef DEBUG_PRINT\n std::cout << \"index 1: \" << std::endl;\n print_data16(&mm_index_m);\n\n std::cout << \"index 2: \" << std::endl;\n print_data16(&mm_index_q);\n#endif\n\n __m128i mm_cmp_res =\n _mm_cmpistrm(mm_index_q, mm_index_m,\n _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK);\n\n#ifdef DEBUG_PRINT\n std::cout << \"cmp res: \" << std::endl;\n print_data16(&mm_cmp_res);\n#endif\n\n int r = _mm_extract_epi32(mm_cmp_res, 0);\n\n if (r) {\n int r1 = r & 15;\n\n __m128i v = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_value[i1]));\n __m128 vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r1]));\n\n _mm_storeu_ps(val_1, vs);\n val_1 += _mm_popcnt_u32(r1);\n\n int r2 = (r >> 4) & 15;\n v = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_value[i1 + 4]));\n vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r2]));\n _mm_storeu_ps(val_1, vs);\n val_1 += _mm_popcnt_u32(r2);\n\n mm_cmp_res = _mm_cmpistrm(\n mm_index_m, mm_index_q,\n _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK);\n r = _mm_extract_epi32(mm_cmp_res, 0);\n\n r1 = r & 15;\n\n v = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_value[i2]));\n vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r1]));\n _mm_storeu_ps(val_2, vs);\n val_2 += _mm_popcnt_u32(r1);\n\n r2 = (r >> 4) & 15;\n v = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_value[i2 + 4]));\n vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r2]));\n _mm_storeu_ps(val_2, vs);\n val_2 += _mm_popcnt_u32(r2);\n }\n\n const uint16_t id1_max = m_sparse_index[i1 + 7];\n\n if (id1_max <= q_sparse_index[i2 + 7]) {\n i1 += 8;\n if (i1 >= end1) goto do_scalar;\n mm_index_m = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_index[i1]));\n }\n\n if (id1_max >= q_sparse_index[i2 + 7]) {\n i2 += 8;\n if (i2 >= end2) goto do_scalar;\n mm_index_q = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_index[i2]));\n }\n }\n }\n\ndo_scalar:\n while (i1 < m_sparse_count && i2 < q_sparse_count) {\n if (m_sparse_index[i1] == q_sparse_index[i2]) {\n *val_1++ = m_sparse_value[i1];\n *val_2++ = q_sparse_value[i2];\n\n ++i1;\n ++i2;\n } else if (m_sparse_index[i1] < q_sparse_index[i2]) {\n ++i1;\n } else {\n ++i2;\n }\n }\n\n size_t res_num = val_1 - val_start_1;\n\n // if (res_num != val_2 - val_start_2) {\n // std::cerr << \"size mismatch!\" << std::endl;\n // }\n\n size_t res_num4 = res_num / 4 * 4;\n\n if (res_num4) {\n __m128 sum128 = _mm_set1_ps(0);\n\n for (size_t k = 0; k < res_num4; k += 4) {\n sum128 = _mm_add_ps(sum128, _mm_mul_ps(_mm_loadu_ps(val_start_1 + k),\n _mm_loadu_ps(val_start_2 + k)));\n }\n\n float __attribute__((aligned(16))) tmp_res[4];\n _mm_store_ps(tmp_res, sum128);\n sum += (tmp_res[0] + tmp_res[1] + tmp_res[2] + tmp_res[3]);\n }\n\n for (size_t k = res_num4; k < res_num; ++k)\n sum += val_start_1[k] * val_start_2[k];\n\n return sum;\n}\n#endif // __SSE4_1__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_int4_avx2.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int4.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__AVX2__)\nfloat InnerProductInt4SSEInternal(const uint8_t *lhs, const uint8_t *rhs,\n size_t size);\n\n//! Inner Product\nfloat InnerProductInt4AVX2Internal(const uint8_t *lhs, const uint8_t *rhs,\n size_t size) {\n const uint8_t *last = lhs + size;\n const uint8_t *last_aligned = lhs + ((size >> 5) << 5);\n __m256i ymm_sum = _mm256_setzero_si256();\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m256i ymm_lhs = _mm256_load_si256((const __m256i *)(lhs));\n __m256i ymm_rhs = _mm256_load_si256((const __m256i *)(rhs));\n FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum)\n }\n\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs);\n __m128i xmm_sum = _mm_setzero_si128();\n FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum)\n ymm_sum = _mm256_add_epi32(_mm256_set_m128i(_mm_setzero_si128(), xmm_sum),\n ymm_sum);\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)(lhs));\n __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)(rhs));\n FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum)\n }\n\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs);\n __m128i xmm_sum = _mm_setzero_si128();\n FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum)\n ymm_sum = _mm256_add_epi32(_mm256_set_m128i(_mm_setzero_si128(), xmm_sum),\n ymm_sum);\n lhs += 16;\n rhs += 16;\n }\n }\n float result = static_cast(HorizontalAdd_INT32_V256(ymm_sum));\n\n switch (last - lhs) {\n case 15:\n FMA_INT4_GENERAL(lhs[14], rhs[14], result)\n /* FALLTHRU */\n case 14:\n FMA_INT4_GENERAL(lhs[13], rhs[13], result)\n /* FALLTHRU */\n case 13:\n FMA_INT4_GENERAL(lhs[12], rhs[12], result)\n /* FALLTHRU */\n case 12:\n FMA_INT4_GENERAL(lhs[11], rhs[11], result)\n /* FALLTHRU */\n case 11:\n FMA_INT4_GENERAL(lhs[10], rhs[10], result)\n /* FALLTHRU */\n case 10:\n FMA_INT4_GENERAL(lhs[9], rhs[9], result)\n /* FALLTHRU */\n case 9:\n FMA_INT4_GENERAL(lhs[8], rhs[8], result)\n /* FALLTHRU */\n case 8:\n FMA_INT4_GENERAL(lhs[7], rhs[7], result)\n /* FALLTHRU */\n case 7:\n FMA_INT4_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n FMA_INT4_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n FMA_INT4_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n FMA_INT4_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n FMA_INT4_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n FMA_INT4_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n FMA_INT4_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat InnerProductInt4AVX2(const uint8_t *lhs, const uint8_t *rhs,\n size_t size) {\n if (size > 63) {\n return InnerProductInt4AVX2Internal(lhs, rhs, size >> 1);\n }\n\n return InnerProductInt4SSEInternal(lhs, rhs, size >> 1);\n}\n\nfloat MinusInnerProductInt4AVX2(const uint8_t *lhs, const uint8_t *rhs,\n size_t size) {\n return -InnerProductInt4AVX2(lhs, rhs, size);\n}\n\n#endif // __AVX2__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/inner_product_matrix_int4_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__AVX2__)\nfloat InnerProductInt4AVX2(const uint8_t *lhs, const uint8_t *rhs, size_t size);\nfloat MinusInnerProductInt4AVX2(const uint8_t *lhs, const uint8_t *rhs,\n size_t size);\n#endif\n\n#if defined(__SSE4_1__)\nfloat InnerProductInt4SSE(const uint8_t *lhs, const uint8_t *rhs, size_t size);\nfloat MinusInnerProductInt4SSE(const uint8_t *lhs, const uint8_t *rhs,\n size_t size);\n#endif\n\nfloat InnerProductInt4Scalar(const uint8_t *m, const uint8_t *q, size_t dim);\nfloat MinusInnerProductInt4Scalar(const uint8_t *m, const uint8_t *q,\n size_t dim);\n\n//! Compute the distance between matrix and query (INT4, M=1, N=1)\nvoid InnerProductMatrix::Compute(const uint8_t *m,\n const uint8_t *q, size_t dim,\n float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = InnerProductInt4AVX2(m, q, dim);\n return;\n }\n#endif // __AVX2__\n\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = InnerProductInt4SSE(m, q, dim);\n return;\n }\n#endif //__SSE4_1__\n *out = InnerProductInt4Scalar(m, q, dim);\n}\n\n//! Compute the distance between matrix and query (INT4, M=1, N=1)\nvoid MinusInnerProductMatrix::Compute(const uint8_t *m,\n const uint8_t *q,\n size_t dim, float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = MinusInnerProductInt4AVX2(m, q, dim);\n return;\n }\n#endif // __AVX2__\n\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = MinusInnerProductInt4SSE(m, q, dim);\n return;\n }\n#endif //__SSE4_1__\n *out = MinusInnerProductInt4Scalar(m, q, dim);\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/inner_product_matrix_int4_sse.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int4.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__SSE4_1__)\nfloat InnerProductInt4SSEInternal(const uint8_t *lhs, const uint8_t *rhs,\n size_t size) {\n const uint8_t *last = lhs + size;\n const uint8_t *last_aligned = lhs + ((size >> 4) << 4);\n __m128i xmm_sum = _mm_setzero_si128();\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)(lhs));\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)(rhs));\n FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum)\n }\n } else {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)(lhs));\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)(rhs));\n FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum)\n }\n }\n float result = static_cast(HorizontalAdd_INT32_V128(xmm_sum));\n\n switch (last - lhs) {\n case 15:\n FMA_INT4_GENERAL(lhs[14], rhs[14], result)\n /* FALLTHRU */\n case 14:\n FMA_INT4_GENERAL(lhs[13], rhs[13], result)\n /* FALLTHRU */\n case 13:\n FMA_INT4_GENERAL(lhs[12], rhs[12], result)\n /* FALLTHRU */\n case 12:\n FMA_INT4_GENERAL(lhs[11], rhs[11], result)\n /* FALLTHRU */\n case 11:\n FMA_INT4_GENERAL(lhs[10], rhs[10], result)\n /* FALLTHRU */\n case 10:\n FMA_INT4_GENERAL(lhs[9], rhs[9], result)\n /* FALLTHRU */\n case 9:\n FMA_INT4_GENERAL(lhs[8], rhs[8], result)\n /* FALLTHRU */\n case 8:\n FMA_INT4_GENERAL(lhs[7], rhs[7], result)\n /* FALLTHRU */\n case 7:\n FMA_INT4_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n FMA_INT4_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n FMA_INT4_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n FMA_INT4_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n FMA_INT4_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n FMA_INT4_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n FMA_INT4_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat InnerProductInt4SSE(const uint8_t *lhs, const uint8_t *rhs, size_t size) {\n return InnerProductInt4SSEInternal(lhs, rhs, size >> 1);\n}\n\nfloat MinusInnerProductInt4SSE(const uint8_t *lhs, const uint8_t *rhs,\n size_t size) {\n return -InnerProductInt4SSE(lhs, rhs, size);\n}\n\n#endif // __SSE4_1__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/inner_product_matrix_int8_avx2.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int8.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__AVX2__)\nfloat InnerProductInt8SSEInternal(const int8_t *lhs, const int8_t *rhs,\n size_t size);\n\ninline float InnerProductInt8AVX2Internal(const int8_t *lhs, const int8_t *rhs,\n size_t size) {\n const int8_t *last = lhs + size;\n const int8_t *last_aligned = lhs + ((size >> 6) << 6);\n float result = 0.0;\n\n __m256i ymm_sum_0 = _mm256_setzero_si256();\n __m256i ymm_sum_1 = _mm256_setzero_si256();\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n __m256i ymm_lhs_0 = _mm256_load_si256((const __m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_load_si256((const __m256i *)(lhs + 32));\n __m256i ymm_rhs_0 = _mm256_load_si256((const __m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_load_si256((const __m256i *)(rhs + 32));\n\n ymm_lhs_0 = _mm256_sign_epi8(ymm_lhs_0, ymm_rhs_0);\n ymm_lhs_1 = _mm256_sign_epi8(ymm_lhs_1, ymm_rhs_1);\n ymm_rhs_0 = _mm256_abs_epi8(ymm_rhs_0);\n ymm_rhs_1 = _mm256_abs_epi8(ymm_rhs_1);\n\n ymm_sum_0 = _mm256_add_epi32(\n _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_0, ymm_lhs_0),\n ONES_INT16_AVX),\n ymm_sum_0);\n ymm_sum_1 = _mm256_add_epi32(\n _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_1, ymm_lhs_1),\n ONES_INT16_AVX),\n ymm_sum_1);\n }\n\n if (last >= last_aligned + 32) {\n __m256i ymm_lhs = _mm256_load_si256((const __m256i *)lhs);\n __m256i ymm_rhs = _mm256_load_si256((const __m256i *)rhs);\n ymm_lhs = _mm256_sign_epi8(ymm_lhs, ymm_rhs);\n ymm_rhs = _mm256_abs_epi8(ymm_rhs);\n ymm_sum_0 = _mm256_add_epi32(\n _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs, ymm_lhs),\n ONES_INT16_AVX),\n ymm_sum_0);\n lhs += 32;\n rhs += 32;\n }\n\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs);\n xmm_lhs = _mm_sign_epi8(xmm_lhs, xmm_rhs);\n xmm_rhs = _mm_abs_epi8(xmm_rhs);\n ymm_sum_0 = _mm256_add_epi32(\n _mm256_set_m128i(_mm_setzero_si128(),\n _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs, xmm_lhs),\n ONES_INT16_SSE)),\n ymm_sum_0);\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n __m256i ymm_lhs_0 = _mm256_loadu_si256((const __m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_loadu_si256((const __m256i *)(lhs + 32));\n __m256i ymm_rhs_0 = _mm256_loadu_si256((const __m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_loadu_si256((const __m256i *)(rhs + 32));\n\n ymm_lhs_0 = _mm256_sign_epi8(ymm_lhs_0, ymm_rhs_0);\n ymm_lhs_1 = _mm256_sign_epi8(ymm_lhs_1, ymm_rhs_1);\n ymm_rhs_0 = _mm256_abs_epi8(ymm_rhs_0);\n ymm_rhs_1 = _mm256_abs_epi8(ymm_rhs_1);\n\n ymm_sum_0 = _mm256_add_epi32(\n _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_0, ymm_lhs_0),\n ONES_INT16_AVX),\n ymm_sum_0);\n ymm_sum_1 = _mm256_add_epi32(\n _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs_1, ymm_lhs_1),\n ONES_INT16_AVX),\n ymm_sum_1);\n }\n\n if (last >= last_aligned + 32) {\n __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)lhs);\n __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)rhs);\n ymm_lhs = _mm256_sign_epi8(ymm_lhs, ymm_rhs);\n ymm_rhs = _mm256_abs_epi8(ymm_rhs);\n ymm_sum_0 = _mm256_add_epi32(\n _mm256_madd_epi16(_mm256_maddubs_epi16(ymm_rhs, ymm_lhs),\n ONES_INT16_AVX),\n ymm_sum_0);\n lhs += 32;\n rhs += 32;\n }\n\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs);\n xmm_lhs = _mm_sign_epi8(xmm_lhs, xmm_rhs);\n xmm_rhs = _mm_abs_epi8(xmm_rhs);\n ymm_sum_0 = _mm256_add_epi32(\n _mm256_set_m128i(_mm_setzero_si128(),\n _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs, xmm_lhs),\n ONES_INT16_SSE)),\n ymm_sum_0);\n lhs += 16;\n rhs += 16;\n }\n }\n result = static_cast(\n HorizontalAdd_INT32_V256(_mm256_add_epi32(ymm_sum_0, ymm_sum_1)));\n\n switch (last - lhs) {\n case 15:\n FMA_INT8_GENERAL(lhs[14], rhs[14], result)\n /* FALLTHRU */\n case 14:\n FMA_INT8_GENERAL(lhs[13], rhs[13], result)\n /* FALLTHRU */\n case 13:\n FMA_INT8_GENERAL(lhs[12], rhs[12], result)\n /* FALLTHRU */\n case 12:\n FMA_INT8_GENERAL(lhs[11], rhs[11], result)\n /* FALLTHRU */\n case 11:\n FMA_INT8_GENERAL(lhs[10], rhs[10], result)\n /* FALLTHRU */\n case 10:\n FMA_INT8_GENERAL(lhs[9], rhs[9], result)\n /* FALLTHRU */\n case 9:\n FMA_INT8_GENERAL(lhs[8], rhs[8], result)\n /* FALLTHRU */\n case 8:\n FMA_INT8_GENERAL(lhs[7], rhs[7], result)\n /* FALLTHRU */\n case 7:\n FMA_INT8_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n FMA_INT8_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n FMA_INT8_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n FMA_INT8_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n FMA_INT8_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n FMA_INT8_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n FMA_INT8_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat InnerProductInt8AVX2(const int8_t *lhs, const int8_t *rhs, size_t size) {\n if (size > 31) {\n return InnerProductInt8AVX2Internal(lhs, rhs, size);\n }\n\n return InnerProductInt8SSEInternal(lhs, rhs, size);\n}\n\nfloat MinusInnerProductInt8AVX2(const int8_t *lhs, const int8_t *rhs,\n size_t size) {\n return -InnerProductInt8AVX2(lhs, rhs, size);\n}\n\n#endif // __AVX2__\n\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/inner_product_matrix_int8_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__AVX2__)\nfloat InnerProductInt8AVX2(const int8_t *lhs, const int8_t *rhs, size_t size);\nfloat MinusInnerProductInt8AVX2(const int8_t *lhs, const int8_t *rhs,\n size_t size);\n#endif\n\n#if defined(__SSE4_1__)\nfloat InnerProductInt8SSE(const int8_t *lhs, const int8_t *rhs, size_t size);\nfloat MinusInnerProductInt8SSE(const int8_t *lhs, const int8_t *rhs,\n size_t size);\n#endif\n\nfloat InnerProductInt8Scalar(const int8_t *m, const int8_t *q, size_t dim);\nfloat MinusInnerProductInt8Scalar(const int8_t *m, const int8_t *q, size_t dim);\n\n//! Compute the distance between matrix and query (INT8, M=1, N=1)\nvoid InnerProductMatrix::Compute(const int8_t *m, const int8_t *q,\n size_t dim, float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = InnerProductInt8AVX2(m, q, dim);\n return;\n }\n#endif // __AVX2__\n\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = InnerProductInt8SSE(m, q, dim);\n return;\n }\n\n#endif //__SSE4_1__\n\n *out = InnerProductInt8Scalar(m, q, dim);\n}\n\n//! Compute the distance between matrix and query (INT8, M=1, N=1)\nvoid MinusInnerProductMatrix::Compute(const int8_t *m,\n const int8_t *q, size_t dim,\n float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = MinusInnerProductInt8AVX2(m, q, dim);\n return;\n }\n#endif // __AVX2__\n\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = MinusInnerProductInt8SSE(m, q, dim);\n return;\n }\n#endif //__SSE4_1__\n\n *out = MinusInnerProductInt8Scalar(m, q, dim);\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/inner_product_matrix_int8_sse.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int8.i\"\n#include \"distance_matrix_inner_product_utility.i\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n#if defined(__SSE4_1__)\n//! Inner Product\nfloat InnerProductInt8SSEInternal(const int8_t *lhs, const int8_t *rhs,\n size_t size) {\n const int8_t *last = lhs + size;\n const int8_t *last_aligned = lhs + ((size >> 5) << 5);\n\n __m128i xmm_sum_0 = _mm_setzero_si128();\n __m128i xmm_sum_1 = _mm_setzero_si128();\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m128i xmm_lhs_0 = _mm_load_si128((const __m128i *)(lhs + 0));\n __m128i xmm_lhs_1 = _mm_load_si128((const __m128i *)(lhs + 16));\n __m128i xmm_rhs_0 = _mm_load_si128((const __m128i *)(rhs + 0));\n __m128i xmm_rhs_1 = _mm_load_si128((const __m128i *)(rhs + 16));\n\n xmm_lhs_0 = _mm_sign_epi8(xmm_lhs_0, xmm_rhs_0);\n xmm_lhs_1 = _mm_sign_epi8(xmm_lhs_1, xmm_rhs_1);\n xmm_rhs_0 = _mm_abs_epi8(xmm_rhs_0);\n xmm_rhs_1 = _mm_abs_epi8(xmm_rhs_1);\n xmm_sum_0 =\n _mm_add_epi32(_mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_0, xmm_lhs_0),\n ONES_INT16_SSE),\n xmm_sum_0);\n xmm_sum_1 =\n _mm_add_epi32(_mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_1, xmm_lhs_1),\n ONES_INT16_SSE),\n xmm_sum_1);\n }\n\n if (last >= last_aligned + 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs);\n\n xmm_lhs = _mm_sign_epi8(xmm_lhs, xmm_rhs);\n xmm_rhs = _mm_abs_epi8(xmm_rhs);\n xmm_sum_0 = _mm_add_epi32(\n _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs, xmm_lhs), ONES_INT16_SSE),\n xmm_sum_0);\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m128i xmm_lhs_0 = _mm_loadu_si128((const __m128i *)(lhs + 0));\n __m128i xmm_lhs_1 = _mm_loadu_si128((const __m128i *)(lhs + 16));\n __m128i xmm_rhs_0 = _mm_loadu_si128((const __m128i *)(rhs + 0));\n __m128i xmm_rhs_1 = _mm_loadu_si128((const __m128i *)(rhs + 16));\n\n xmm_lhs_0 = _mm_sign_epi8(xmm_lhs_0, xmm_rhs_0);\n xmm_lhs_1 = _mm_sign_epi8(xmm_lhs_1, xmm_rhs_1);\n xmm_rhs_0 = _mm_abs_epi8(xmm_rhs_0);\n xmm_rhs_1 = _mm_abs_epi8(xmm_rhs_1);\n xmm_sum_0 =\n _mm_add_epi32(_mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_0, xmm_lhs_0),\n ONES_INT16_SSE),\n xmm_sum_0);\n xmm_sum_1 =\n _mm_add_epi32(_mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs_1, xmm_lhs_1),\n ONES_INT16_SSE),\n xmm_sum_1);\n }\n\n if (last >= last_aligned + 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs);\n\n xmm_lhs = _mm_sign_epi8(xmm_lhs, xmm_rhs);\n xmm_rhs = _mm_abs_epi8(xmm_rhs);\n xmm_sum_0 = _mm_add_epi32(\n _mm_madd_epi16(_mm_maddubs_epi16(xmm_rhs, xmm_lhs), ONES_INT16_SSE),\n xmm_sum_0);\n lhs += 16;\n rhs += 16;\n }\n }\n float result = static_cast(\n HorizontalAdd_INT32_V128(_mm_add_epi32(xmm_sum_0, xmm_sum_1)));\n\n switch (last - lhs) {\n case 15:\n FMA_INT8_GENERAL(lhs[14], rhs[14], result)\n /* FALLTHRU */\n case 14:\n FMA_INT8_GENERAL(lhs[13], rhs[13], result)\n /* FALLTHRU */\n case 13:\n FMA_INT8_GENERAL(lhs[12], rhs[12], result)\n /* FALLTHRU */\n case 12:\n FMA_INT8_GENERAL(lhs[11], rhs[11], result)\n /* FALLTHRU */\n case 11:\n FMA_INT8_GENERAL(lhs[10], rhs[10], result)\n /* FALLTHRU */\n case 10:\n FMA_INT8_GENERAL(lhs[9], rhs[9], result)\n /* FALLTHRU */\n case 9:\n FMA_INT8_GENERAL(lhs[8], rhs[8], result)\n /* FALLTHRU */\n case 8:\n FMA_INT8_GENERAL(lhs[7], rhs[7], result)\n /* FALLTHRU */\n case 7:\n FMA_INT8_GENERAL(lhs[6], rhs[6], result)\n /* FALLTHRU */\n case 6:\n FMA_INT8_GENERAL(lhs[5], rhs[5], result)\n /* FALLTHRU */\n case 5:\n FMA_INT8_GENERAL(lhs[4], rhs[4], result)\n /* FALLTHRU */\n case 4:\n FMA_INT8_GENERAL(lhs[3], rhs[3], result)\n /* FALLTHRU */\n case 3:\n FMA_INT8_GENERAL(lhs[2], rhs[2], result)\n /* FALLTHRU */\n case 2:\n FMA_INT8_GENERAL(lhs[1], rhs[1], result)\n /* FALLTHRU */\n case 1:\n FMA_INT8_GENERAL(lhs[0], rhs[0], result)\n }\n return result;\n}\n\nfloat InnerProductInt8SSE(const int8_t *lhs, const int8_t *rhs, size_t size) {\n return InnerProductInt8SSEInternal(lhs, rhs, size);\n}\n\nfloat MinusInnerProductInt8SSE(const int8_t *lhs, const int8_t *rhs,\n size_t size) {\n return -InnerProductInt8SSEInternal(lhs, rhs, size);\n}\n\n#endif // __SSE4_1__\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/inner_product_matrix_scalar.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"distance_utility.h\"\n#include \"inner_product_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\ntemplate \ninline float InnerProductScalar(const T *m, const T *q, size_t dim) {\n ailego_assert(m && q && dim);\n\n float sum = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n sum += static_cast(m[i] * q[i]);\n }\n return sum;\n}\n\ntemplate \ninline float MinusInnerProductScalar(const T *m, const T *q, size_t dim) {\n ailego_assert(m && q && dim);\n\n float sum = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n sum += static_cast(m[i] * q[i]);\n }\n return -sum;\n}\n\nfloat InnerProductInt4Scalar(const uint8_t *m, const uint8_t *q, size_t dim) {\n ailego_assert(m && q && dim && !(dim & 1));\n\n float sum = 0.0;\n for (size_t i = 0; i < (dim >> 1); ++i) {\n uint8_t m_val = m[i];\n uint8_t q_val = q[i];\n sum += Int4MulTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] +\n Int4MulTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)];\n }\n\n return sum;\n}\n\nfloat MinusInnerProductInt4Scalar(const uint8_t *m, const uint8_t *q,\n size_t dim) {\n ailego_assert(m && q && dim && !(dim & 1));\n\n float sum = 0.0;\n for (size_t i = 0; i < (dim >> 1); ++i) {\n uint8_t m_val = m[i];\n uint8_t q_val = q[i];\n sum -= Int4MulTable[((m_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] +\n Int4MulTable[((m_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)];\n }\n return sum;\n}\n\nfloat InnerProductInt8Scalar(const int8_t *m, const int8_t *q, size_t dim) {\n return InnerProductScalar(m, q, dim);\n}\n\nfloat MinusInnerProductInt8Scalar(const int8_t *m, const int8_t *q,\n size_t dim) {\n return MinusInnerProductScalar(m, q, dim);\n}\n\nfloat InnerProductFp16Scalar(const ailego::Float16 *m, const ailego::Float16 *q,\n size_t dim) {\n return InnerProductScalar(m, q, dim);\n}\n\nfloat MinusInnerProductFp16Scalar(const ailego::Float16 *m,\n const ailego::Float16 *q, size_t dim) {\n return MinusInnerProductScalar(m, q, dim);\n}\n\nfloat InnerProductFp32Scalar(const float *m, const float *q, size_t dim) {\n return InnerProductScalar(m, q, dim);\n}\n\nfloat MinusInnerProductFp32Scalar(const float *m, const float *q, size_t dim) {\n return MinusInnerProductScalar(m, q, dim);\n}\n\n//--------------------------------------------------\n// Sparse\n//--------------------------------------------------\nfloat ComputeInnerProductSparseInSegmentFp32(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value);\n\nfloat ComputeInnerProductSparseInSegmentFp16(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const Float16 *q_sparse_value);\n\ntemplate \nfloat ComputeInnerProductSparseInSegment(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const T *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const T *q_sparse_value);\n\ntemplate <>\nfloat ComputeInnerProductSparseInSegment(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value) {\n return ComputeInnerProductSparseInSegmentFp32(m_sparse_count, m_sparse_index,\n m_sparse_value, q_sparse_count,\n q_sparse_index, q_sparse_value);\n}\n\ntemplate <>\nfloat ComputeInnerProductSparseInSegment(\n uint32_t m_sparse_count, const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value, uint32_t q_sparse_count,\n const uint16_t *q_sparse_index, const Float16 *q_sparse_value) {\n return ComputeInnerProductSparseInSegmentFp16(m_sparse_count, m_sparse_index,\n m_sparse_value, q_sparse_count,\n q_sparse_index, q_sparse_value);\n}\n\ntemplate \nfloat ComputeSegments(const void *m_sparse_data_in,\n const void *q_sparse_data_in) {\n ailego_assert(m_sparse_data_in && q_sparse_data_in);\n\n float sum{0.0f};\n\n const uint8_t *m_sparse_data =\n reinterpret_cast(m_sparse_data_in);\n const uint8_t *q_sparse_data =\n reinterpret_cast(q_sparse_data_in);\n\n const uint32_t m_sparse_count =\n *reinterpret_cast(m_sparse_data);\n const uint32_t q_sparse_count =\n *reinterpret_cast(q_sparse_data);\n\n if (m_sparse_count == 0 || q_sparse_count == 0) {\n return 0.0f;\n }\n\n const uint32_t m_seg_count =\n *reinterpret_cast(m_sparse_data + sizeof(uint32_t));\n const uint32_t q_seg_count =\n *reinterpret_cast(q_sparse_data + sizeof(uint32_t));\n\n const uint32_t *m_seg_id =\n reinterpret_cast(m_sparse_data + 2 * sizeof(uint32_t));\n const uint32_t *q_seg_id =\n reinterpret_cast(q_sparse_data + 2 * sizeof(uint32_t));\n\n const uint32_t *m_seg_vec_cnt = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t) + m_seg_count * sizeof(uint32_t));\n const uint32_t *q_seg_vec_cnt = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t) + q_seg_count * sizeof(uint32_t));\n\n const uint16_t *m_sparse_index =\n reinterpret_cast(m_sparse_data + 2 * sizeof(uint32_t) +\n m_seg_count * 2 * sizeof(uint32_t));\n const uint16_t *q_sparse_index =\n reinterpret_cast(q_sparse_data + 2 * sizeof(uint32_t) +\n q_seg_count * 2 * sizeof(uint32_t));\n\n const T *m_sparse_value = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t) +\n m_seg_count * 2 * sizeof(uint32_t) + m_sparse_count * sizeof(uint16_t));\n const T *q_sparse_value = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t) +\n q_seg_count * 2 * sizeof(uint32_t) + q_sparse_count * sizeof(uint16_t));\n\n size_t m_s = 0;\n size_t q_s = 0;\n\n size_t m_count = 0;\n size_t q_count = 0;\n\n while (m_s < m_seg_count && q_s < q_seg_count) {\n if (m_seg_id[m_s] == q_seg_id[q_s]) {\n sum += ComputeInnerProductSparseInSegment(\n m_seg_vec_cnt[m_s], m_sparse_index + m_count,\n m_sparse_value + m_count, q_seg_vec_cnt[q_s],\n q_sparse_index + q_count, q_sparse_value + q_count);\n\n m_count += m_seg_vec_cnt[m_s];\n q_count += q_seg_vec_cnt[q_s];\n\n ++m_s;\n ++q_s;\n } else if (m_seg_id[m_s] < q_seg_id[q_s]) {\n m_count += m_seg_vec_cnt[m_s];\n\n ++m_s;\n } else {\n q_count += q_seg_vec_cnt[q_s];\n\n ++q_s;\n }\n }\n\n return -sum;\n}\n\nfloat MinusInnerProductSparseFp16Scalar(const void *m_sparse_data_in,\n const void *q_sparse_data_in) {\n return ComputeSegments(m_sparse_data_in, q_sparse_data_in);\n}\n\nfloat MinusInnerProductSparseFp32Scalar(const void *m_sparse_data_in,\n const void *q_sparse_data_in) {\n return ComputeSegments(m_sparse_data_in, q_sparse_data_in);\n}\n\nfloat InnerProductSparseInSegmentFp16Scalar(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const Float16 *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const Float16 *q_sparse_value) {\n float sum = 0.0f;\n\n size_t m_i = 0;\n size_t q_i = 0;\n while (m_i < m_sparse_count && q_i < q_sparse_count) {\n if (m_sparse_index[m_i] == q_sparse_index[q_i]) {\n sum += m_sparse_value[m_i] * q_sparse_value[q_i];\n\n ++m_i;\n ++q_i;\n } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) {\n ++m_i;\n } else {\n ++q_i;\n }\n }\n\n return sum;\n}\n\nfloat InnerProductSparseInSegmentFp32Scalar(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value) {\n float sum = 0.0f;\n\n size_t m_i = 0;\n size_t q_i = 0;\n while (m_i < m_sparse_count && q_i < q_sparse_count) {\n if (m_sparse_index[m_i] == q_sparse_index[q_i]) {\n sum += m_sparse_value[m_i] * q_sparse_value[q_i];\n\n ++m_i;\n ++q_i;\n } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) {\n ++m_i;\n } else {\n ++q_i;\n }\n }\n\n return sum;\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/matrix_define.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#define MATRIX_VAR_INIT_1X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_0 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_1X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_1X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_1X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_1X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_1X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_1X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_7 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_1X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_1X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_15 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_2X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_1X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_1_0 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_2X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_2X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_1 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_2X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_2X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_3 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_2X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_2X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_7 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_2X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_2X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_15 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_2X32(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_2X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_31 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_4X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_2X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_2_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_0 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_4X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_4X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_1 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_4X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_4X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_3 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_4X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_4X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_7 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_4X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_4X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_15 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_4X32(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_4X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_31 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_8X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_4X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_4_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_0 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_8X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_8X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_1 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_8X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_8X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_3 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_8X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_8X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_4 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_5 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_6 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_7 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_7 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_8X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_8X8(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_8 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_9 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_10 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_11 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_12 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_13 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_14 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_15 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_15 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_8X32(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_8X16(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_16 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_17 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_18 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_19 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_20 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_21 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_22 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_23 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_24 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_25 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_26 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_27 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_28 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_29 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_30 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_31 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_31 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_16X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_8X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_8_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_9_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_10_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_11_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_12_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_13_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_14_0 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_15_0 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_16X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_16X1(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_8_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_9_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_10_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_11_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_12_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_13_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_14_1 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_15_1 = (_VAR_INIT);\n\n#define MATRIX_VAR_INIT_16X4(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_16X2(_VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n _VAR_TYPE _VAR_NAME##_0_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_8_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_9_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_10_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_11_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_12_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_13_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_14_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_15_2 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_0_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_1_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_2_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_3_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_4_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_5_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_6_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_7_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_8_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_9_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_10_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_11_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_12_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_13_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_14_3 = (_VAR_INIT); \\\n _VAR_TYPE _VAR_NAME##_15_3 = (_VAR_INIT);\n\n#define MATRIX_VAR_STORE_1X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n _STORE((_ARRAY) + (_STEP) * (0), _NORM((_VAR##_0_0), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_1X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_1X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (1), _NORM((_VAR##_0_1), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_1X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_1X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (2), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (3), _NORM((_VAR##_0_3), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_1X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_1X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (4), _NORM((_VAR##_0_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (5), _NORM((_VAR##_0_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (6), _NORM((_VAR##_0_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (7), _NORM((_VAR##_0_7), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_1X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_1X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_0_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_0_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_0_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_0_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_0_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_0_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_0_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_0_15), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_2X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_1X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (1), _NORM((_VAR##_1_0), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_2X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_2X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (2), _NORM((_VAR##_0_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (3), _NORM((_VAR##_1_1), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_2X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_2X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (4), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (5), _NORM((_VAR##_1_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (6), _NORM((_VAR##_0_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (7), _NORM((_VAR##_1_3), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_2X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_2X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_0_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_1_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_0_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_1_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_0_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_1_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_0_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_1_7), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_2X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_2X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (16), _NORM((_VAR##_0_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (17), _NORM((_VAR##_1_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (18), _NORM((_VAR##_0_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (19), _NORM((_VAR##_1_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (20), _NORM((_VAR##_0_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (21), _NORM((_VAR##_1_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (22), _NORM((_VAR##_0_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (23), _NORM((_VAR##_1_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (24), _NORM((_VAR##_0_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (25), _NORM((_VAR##_1_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (26), _NORM((_VAR##_0_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (27), _NORM((_VAR##_1_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (28), _NORM((_VAR##_0_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (29), _NORM((_VAR##_1_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (30), _NORM((_VAR##_0_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (31), _NORM((_VAR##_1_15), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_2X32(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_2X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (32), _NORM((_VAR##_0_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (33), _NORM((_VAR##_1_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (34), _NORM((_VAR##_0_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (35), _NORM((_VAR##_1_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (36), _NORM((_VAR##_0_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (37), _NORM((_VAR##_1_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (38), _NORM((_VAR##_0_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (39), _NORM((_VAR##_1_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (40), _NORM((_VAR##_0_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (41), _NORM((_VAR##_1_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (42), _NORM((_VAR##_0_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (43), _NORM((_VAR##_1_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (44), _NORM((_VAR##_0_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (45), _NORM((_VAR##_1_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (46), _NORM((_VAR##_0_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (47), _NORM((_VAR##_1_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (48), _NORM((_VAR##_0_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (49), _NORM((_VAR##_1_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (50), _NORM((_VAR##_0_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (51), _NORM((_VAR##_1_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (52), _NORM((_VAR##_0_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (53), _NORM((_VAR##_1_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (54), _NORM((_VAR##_0_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (55), _NORM((_VAR##_1_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (56), _NORM((_VAR##_0_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (57), _NORM((_VAR##_1_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (58), _NORM((_VAR##_0_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (59), _NORM((_VAR##_1_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (60), _NORM((_VAR##_0_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (61), _NORM((_VAR##_1_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (62), _NORM((_VAR##_0_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (63), _NORM((_VAR##_1_31), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_4X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_2X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (2), _NORM((_VAR##_2_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (3), _NORM((_VAR##_3_0), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_4X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_4X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (4), _NORM((_VAR##_0_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (5), _NORM((_VAR##_1_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (6), _NORM((_VAR##_2_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (7), _NORM((_VAR##_3_1), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_4X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_4X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_1_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_2_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_3_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_0_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_1_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_2_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_3_3), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_4X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_4X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (16), _NORM((_VAR##_0_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (17), _NORM((_VAR##_1_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (18), _NORM((_VAR##_2_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (19), _NORM((_VAR##_3_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (20), _NORM((_VAR##_0_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (21), _NORM((_VAR##_1_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (22), _NORM((_VAR##_2_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (23), _NORM((_VAR##_3_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (24), _NORM((_VAR##_0_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (25), _NORM((_VAR##_1_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (26), _NORM((_VAR##_2_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (27), _NORM((_VAR##_3_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (28), _NORM((_VAR##_0_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (29), _NORM((_VAR##_1_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (30), _NORM((_VAR##_2_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (31), _NORM((_VAR##_3_7), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_4X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_4X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (32), _NORM((_VAR##_0_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (33), _NORM((_VAR##_1_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (34), _NORM((_VAR##_2_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (35), _NORM((_VAR##_3_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (36), _NORM((_VAR##_0_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (37), _NORM((_VAR##_1_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (38), _NORM((_VAR##_2_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (39), _NORM((_VAR##_3_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (40), _NORM((_VAR##_0_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (41), _NORM((_VAR##_1_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (42), _NORM((_VAR##_2_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (43), _NORM((_VAR##_3_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (44), _NORM((_VAR##_0_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (45), _NORM((_VAR##_1_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (46), _NORM((_VAR##_2_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (47), _NORM((_VAR##_3_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (48), _NORM((_VAR##_0_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (49), _NORM((_VAR##_1_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (50), _NORM((_VAR##_2_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (51), _NORM((_VAR##_3_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (52), _NORM((_VAR##_0_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (53), _NORM((_VAR##_1_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (54), _NORM((_VAR##_2_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (55), _NORM((_VAR##_3_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (56), _NORM((_VAR##_0_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (57), _NORM((_VAR##_1_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (58), _NORM((_VAR##_2_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (59), _NORM((_VAR##_3_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (60), _NORM((_VAR##_0_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (61), _NORM((_VAR##_1_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (62), _NORM((_VAR##_2_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (63), _NORM((_VAR##_3_15), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_4X32(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_4X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (64), _NORM((_VAR##_0_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (65), _NORM((_VAR##_1_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (66), _NORM((_VAR##_2_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (67), _NORM((_VAR##_3_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (68), _NORM((_VAR##_0_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (69), _NORM((_VAR##_1_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (70), _NORM((_VAR##_2_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (71), _NORM((_VAR##_3_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (72), _NORM((_VAR##_0_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (73), _NORM((_VAR##_1_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (74), _NORM((_VAR##_2_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (75), _NORM((_VAR##_3_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (76), _NORM((_VAR##_0_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (77), _NORM((_VAR##_1_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (78), _NORM((_VAR##_2_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (79), _NORM((_VAR##_3_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (80), _NORM((_VAR##_0_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (81), _NORM((_VAR##_1_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (82), _NORM((_VAR##_2_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (83), _NORM((_VAR##_3_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (84), _NORM((_VAR##_0_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (85), _NORM((_VAR##_1_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (86), _NORM((_VAR##_2_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (87), _NORM((_VAR##_3_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (88), _NORM((_VAR##_0_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (89), _NORM((_VAR##_1_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (90), _NORM((_VAR##_2_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (91), _NORM((_VAR##_3_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (92), _NORM((_VAR##_0_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (93), _NORM((_VAR##_1_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (94), _NORM((_VAR##_2_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (95), _NORM((_VAR##_3_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (96), _NORM((_VAR##_0_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (97), _NORM((_VAR##_1_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (98), _NORM((_VAR##_2_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (99), _NORM((_VAR##_3_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (100), _NORM((_VAR##_0_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (101), _NORM((_VAR##_1_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (102), _NORM((_VAR##_2_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (103), _NORM((_VAR##_3_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (104), _NORM((_VAR##_0_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (105), _NORM((_VAR##_1_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (106), _NORM((_VAR##_2_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (107), _NORM((_VAR##_3_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (108), _NORM((_VAR##_0_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (109), _NORM((_VAR##_1_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (110), _NORM((_VAR##_2_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (111), _NORM((_VAR##_3_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (112), _NORM((_VAR##_0_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (113), _NORM((_VAR##_1_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (114), _NORM((_VAR##_2_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (115), _NORM((_VAR##_3_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (116), _NORM((_VAR##_0_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (117), _NORM((_VAR##_1_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (118), _NORM((_VAR##_2_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (119), _NORM((_VAR##_3_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (120), _NORM((_VAR##_0_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (121), _NORM((_VAR##_1_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (122), _NORM((_VAR##_2_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (123), _NORM((_VAR##_3_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (124), _NORM((_VAR##_0_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (125), _NORM((_VAR##_1_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (126), _NORM((_VAR##_2_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (127), _NORM((_VAR##_3_31), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_8X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_4X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (4), _NORM((_VAR##_4_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (5), _NORM((_VAR##_5_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (6), _NORM((_VAR##_6_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (7), _NORM((_VAR##_7_0), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_8X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_8X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_0_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_1_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_2_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_3_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_4_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_5_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_6_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_7_1), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_8X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_8X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (16), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (17), _NORM((_VAR##_1_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (18), _NORM((_VAR##_2_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (19), _NORM((_VAR##_3_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (20), _NORM((_VAR##_4_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (21), _NORM((_VAR##_5_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (22), _NORM((_VAR##_6_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (23), _NORM((_VAR##_7_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (24), _NORM((_VAR##_0_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (25), _NORM((_VAR##_1_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (26), _NORM((_VAR##_2_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (27), _NORM((_VAR##_3_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (28), _NORM((_VAR##_4_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (29), _NORM((_VAR##_5_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (30), _NORM((_VAR##_6_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (31), _NORM((_VAR##_7_3), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_8X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_8X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (32), _NORM((_VAR##_0_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (33), _NORM((_VAR##_1_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (34), _NORM((_VAR##_2_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (35), _NORM((_VAR##_3_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (36), _NORM((_VAR##_4_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (37), _NORM((_VAR##_5_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (38), _NORM((_VAR##_6_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (39), _NORM((_VAR##_7_4), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (40), _NORM((_VAR##_0_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (41), _NORM((_VAR##_1_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (42), _NORM((_VAR##_2_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (43), _NORM((_VAR##_3_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (44), _NORM((_VAR##_4_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (45), _NORM((_VAR##_5_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (46), _NORM((_VAR##_6_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (47), _NORM((_VAR##_7_5), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (48), _NORM((_VAR##_0_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (49), _NORM((_VAR##_1_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (50), _NORM((_VAR##_2_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (51), _NORM((_VAR##_3_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (52), _NORM((_VAR##_4_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (53), _NORM((_VAR##_5_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (54), _NORM((_VAR##_6_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (55), _NORM((_VAR##_7_6), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (56), _NORM((_VAR##_0_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (57), _NORM((_VAR##_1_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (58), _NORM((_VAR##_2_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (59), _NORM((_VAR##_3_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (60), _NORM((_VAR##_4_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (61), _NORM((_VAR##_5_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (62), _NORM((_VAR##_6_7), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (63), _NORM((_VAR##_7_7), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_8X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_8X8(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (64), _NORM((_VAR##_0_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (65), _NORM((_VAR##_1_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (66), _NORM((_VAR##_2_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (67), _NORM((_VAR##_3_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (68), _NORM((_VAR##_4_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (69), _NORM((_VAR##_5_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (70), _NORM((_VAR##_6_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (71), _NORM((_VAR##_7_8), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (72), _NORM((_VAR##_0_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (73), _NORM((_VAR##_1_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (74), _NORM((_VAR##_2_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (75), _NORM((_VAR##_3_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (76), _NORM((_VAR##_4_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (77), _NORM((_VAR##_5_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (78), _NORM((_VAR##_6_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (79), _NORM((_VAR##_7_9), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (80), _NORM((_VAR##_0_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (81), _NORM((_VAR##_1_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (82), _NORM((_VAR##_2_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (83), _NORM((_VAR##_3_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (84), _NORM((_VAR##_4_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (85), _NORM((_VAR##_5_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (86), _NORM((_VAR##_6_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (87), _NORM((_VAR##_7_10), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (88), _NORM((_VAR##_0_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (89), _NORM((_VAR##_1_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (90), _NORM((_VAR##_2_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (91), _NORM((_VAR##_3_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (92), _NORM((_VAR##_4_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (93), _NORM((_VAR##_5_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (94), _NORM((_VAR##_6_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (95), _NORM((_VAR##_7_11), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (96), _NORM((_VAR##_0_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (97), _NORM((_VAR##_1_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (98), _NORM((_VAR##_2_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (99), _NORM((_VAR##_3_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (100), _NORM((_VAR##_4_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (101), _NORM((_VAR##_5_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (102), _NORM((_VAR##_6_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (103), _NORM((_VAR##_7_12), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (104), _NORM((_VAR##_0_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (105), _NORM((_VAR##_1_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (106), _NORM((_VAR##_2_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (107), _NORM((_VAR##_3_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (108), _NORM((_VAR##_4_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (109), _NORM((_VAR##_5_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (110), _NORM((_VAR##_6_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (111), _NORM((_VAR##_7_13), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (112), _NORM((_VAR##_0_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (113), _NORM((_VAR##_1_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (114), _NORM((_VAR##_2_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (115), _NORM((_VAR##_3_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (116), _NORM((_VAR##_4_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (117), _NORM((_VAR##_5_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (118), _NORM((_VAR##_6_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (119), _NORM((_VAR##_7_14), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (120), _NORM((_VAR##_0_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (121), _NORM((_VAR##_1_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (122), _NORM((_VAR##_2_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (123), _NORM((_VAR##_3_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (124), _NORM((_VAR##_4_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (125), _NORM((_VAR##_5_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (126), _NORM((_VAR##_6_15), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (127), _NORM((_VAR##_7_15), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_8X32(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_8X16(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (128), _NORM((_VAR##_0_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (129), _NORM((_VAR##_1_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (130), _NORM((_VAR##_2_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (131), _NORM((_VAR##_3_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (132), _NORM((_VAR##_4_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (133), _NORM((_VAR##_5_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (134), _NORM((_VAR##_6_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (135), _NORM((_VAR##_7_16), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (136), _NORM((_VAR##_0_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (137), _NORM((_VAR##_1_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (138), _NORM((_VAR##_2_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (139), _NORM((_VAR##_3_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (140), _NORM((_VAR##_4_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (141), _NORM((_VAR##_5_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (142), _NORM((_VAR##_6_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (143), _NORM((_VAR##_7_17), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (144), _NORM((_VAR##_0_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (145), _NORM((_VAR##_1_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (146), _NORM((_VAR##_2_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (147), _NORM((_VAR##_3_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (148), _NORM((_VAR##_4_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (149), _NORM((_VAR##_5_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (150), _NORM((_VAR##_6_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (151), _NORM((_VAR##_7_18), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (152), _NORM((_VAR##_0_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (153), _NORM((_VAR##_1_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (154), _NORM((_VAR##_2_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (155), _NORM((_VAR##_3_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (156), _NORM((_VAR##_4_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (157), _NORM((_VAR##_5_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (158), _NORM((_VAR##_6_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (159), _NORM((_VAR##_7_19), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (160), _NORM((_VAR##_0_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (161), _NORM((_VAR##_1_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (162), _NORM((_VAR##_2_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (163), _NORM((_VAR##_3_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (164), _NORM((_VAR##_4_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (165), _NORM((_VAR##_5_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (166), _NORM((_VAR##_6_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (167), _NORM((_VAR##_7_20), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (168), _NORM((_VAR##_0_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (169), _NORM((_VAR##_1_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (170), _NORM((_VAR##_2_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (171), _NORM((_VAR##_3_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (172), _NORM((_VAR##_4_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (173), _NORM((_VAR##_5_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (174), _NORM((_VAR##_6_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (175), _NORM((_VAR##_7_21), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (176), _NORM((_VAR##_0_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (177), _NORM((_VAR##_1_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (178), _NORM((_VAR##_2_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (179), _NORM((_VAR##_3_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (180), _NORM((_VAR##_4_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (181), _NORM((_VAR##_5_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (182), _NORM((_VAR##_6_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (183), _NORM((_VAR##_7_22), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (184), _NORM((_VAR##_0_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (185), _NORM((_VAR##_1_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (186), _NORM((_VAR##_2_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (187), _NORM((_VAR##_3_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (188), _NORM((_VAR##_4_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (189), _NORM((_VAR##_5_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (190), _NORM((_VAR##_6_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (191), _NORM((_VAR##_7_23), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (192), _NORM((_VAR##_0_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (193), _NORM((_VAR##_1_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (194), _NORM((_VAR##_2_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (195), _NORM((_VAR##_3_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (196), _NORM((_VAR##_4_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (197), _NORM((_VAR##_5_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (198), _NORM((_VAR##_6_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (199), _NORM((_VAR##_7_24), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (200), _NORM((_VAR##_0_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (201), _NORM((_VAR##_1_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (202), _NORM((_VAR##_2_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (203), _NORM((_VAR##_3_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (204), _NORM((_VAR##_4_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (205), _NORM((_VAR##_5_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (206), _NORM((_VAR##_6_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (207), _NORM((_VAR##_7_25), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (208), _NORM((_VAR##_0_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (209), _NORM((_VAR##_1_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (210), _NORM((_VAR##_2_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (211), _NORM((_VAR##_3_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (212), _NORM((_VAR##_4_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (213), _NORM((_VAR##_5_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (214), _NORM((_VAR##_6_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (215), _NORM((_VAR##_7_26), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (216), _NORM((_VAR##_0_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (217), _NORM((_VAR##_1_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (218), _NORM((_VAR##_2_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (219), _NORM((_VAR##_3_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (220), _NORM((_VAR##_4_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (221), _NORM((_VAR##_5_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (222), _NORM((_VAR##_6_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (223), _NORM((_VAR##_7_27), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (224), _NORM((_VAR##_0_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (225), _NORM((_VAR##_1_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (226), _NORM((_VAR##_2_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (227), _NORM((_VAR##_3_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (228), _NORM((_VAR##_4_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (229), _NORM((_VAR##_5_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (230), _NORM((_VAR##_6_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (231), _NORM((_VAR##_7_28), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (232), _NORM((_VAR##_0_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (233), _NORM((_VAR##_1_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (234), _NORM((_VAR##_2_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (235), _NORM((_VAR##_3_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (236), _NORM((_VAR##_4_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (237), _NORM((_VAR##_5_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (238), _NORM((_VAR##_6_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (239), _NORM((_VAR##_7_29), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (240), _NORM((_VAR##_0_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (241), _NORM((_VAR##_1_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (242), _NORM((_VAR##_2_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (243), _NORM((_VAR##_3_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (244), _NORM((_VAR##_4_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (245), _NORM((_VAR##_5_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (246), _NORM((_VAR##_6_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (247), _NORM((_VAR##_7_30), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (248), _NORM((_VAR##_0_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (249), _NORM((_VAR##_1_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (250), _NORM((_VAR##_2_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (251), _NORM((_VAR##_3_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (252), _NORM((_VAR##_4_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (253), _NORM((_VAR##_5_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (254), _NORM((_VAR##_6_31), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (255), _NORM((_VAR##_7_31), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_16X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_8X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (8), _NORM((_VAR##_8_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (9), _NORM((_VAR##_9_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (10), _NORM((_VAR##_10_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (11), _NORM((_VAR##_11_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (12), _NORM((_VAR##_12_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (13), _NORM((_VAR##_13_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (14), _NORM((_VAR##_14_0), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (15), _NORM((_VAR##_15_0), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_16X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_16X1(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (16), _NORM((_VAR##_0_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (17), _NORM((_VAR##_1_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (18), _NORM((_VAR##_2_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (19), _NORM((_VAR##_3_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (20), _NORM((_VAR##_4_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (21), _NORM((_VAR##_5_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (22), _NORM((_VAR##_6_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (23), _NORM((_VAR##_7_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (24), _NORM((_VAR##_8_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (25), _NORM((_VAR##_9_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (26), _NORM((_VAR##_10_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (27), _NORM((_VAR##_11_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (28), _NORM((_VAR##_12_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (29), _NORM((_VAR##_13_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (30), _NORM((_VAR##_14_1), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (31), _NORM((_VAR##_15_1), ##__VA_ARGS__));\n\n#define MATRIX_VAR_STORE_16X4(_STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_16X2(_STEP, _VAR, _ARRAY, _STORE, _NORM, ##__VA_ARGS__) \\\n _STORE((_ARRAY) + (_STEP) * (32), _NORM((_VAR##_0_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (33), _NORM((_VAR##_1_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (34), _NORM((_VAR##_2_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (35), _NORM((_VAR##_3_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (36), _NORM((_VAR##_4_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (37), _NORM((_VAR##_5_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (38), _NORM((_VAR##_6_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (39), _NORM((_VAR##_7_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (40), _NORM((_VAR##_8_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (41), _NORM((_VAR##_9_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (42), _NORM((_VAR##_10_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (43), _NORM((_VAR##_11_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (44), _NORM((_VAR##_12_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (45), _NORM((_VAR##_13_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (46), _NORM((_VAR##_14_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (47), _NORM((_VAR##_15_2), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (48), _NORM((_VAR##_0_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (49), _NORM((_VAR##_1_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (50), _NORM((_VAR##_2_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (51), _NORM((_VAR##_3_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (52), _NORM((_VAR##_4_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (53), _NORM((_VAR##_5_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (54), _NORM((_VAR##_6_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (55), _NORM((_VAR##_7_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (56), _NORM((_VAR##_8_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (57), _NORM((_VAR##_9_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (58), _NORM((_VAR##_10_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (59), _NORM((_VAR##_11_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (60), _NORM((_VAR##_12_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (61), _NORM((_VAR##_13_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (62), _NORM((_VAR##_14_3), ##__VA_ARGS__)); \\\n _STORE((_ARRAY) + (_STEP) * (63), _NORM((_VAR##_15_3), ##__VA_ARGS__));\n\n#define MATRIX_VAR_PERMUTE_1X1(_VAR, _PERMUTE, ...) \\\n (_VAR##_0_0) = _PERMUTE((_VAR##_0_0), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_1X2(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_1X1(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_1X4(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_1X2(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \\\n (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_1X8(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_1X4(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_4) = _PERMUTE((_VAR##_0_4), ##__VA_ARGS__); \\\n (_VAR##_0_5) = _PERMUTE((_VAR##_0_5), ##__VA_ARGS__); \\\n (_VAR##_0_6) = _PERMUTE((_VAR##_0_6), ##__VA_ARGS__); \\\n (_VAR##_0_7) = _PERMUTE((_VAR##_0_7), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_1X16(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_1X8(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_8) = _PERMUTE((_VAR##_0_8), ##__VA_ARGS__); \\\n (_VAR##_0_9) = _PERMUTE((_VAR##_0_9), ##__VA_ARGS__); \\\n (_VAR##_0_10) = _PERMUTE((_VAR##_0_10), ##__VA_ARGS__); \\\n (_VAR##_0_11) = _PERMUTE((_VAR##_0_11), ##__VA_ARGS__); \\\n (_VAR##_0_12) = _PERMUTE((_VAR##_0_12), ##__VA_ARGS__); \\\n (_VAR##_0_13) = _PERMUTE((_VAR##_0_13), ##__VA_ARGS__); \\\n (_VAR##_0_14) = _PERMUTE((_VAR##_0_14), ##__VA_ARGS__); \\\n (_VAR##_0_15) = _PERMUTE((_VAR##_0_15), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_2X1(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_1X1(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_1_0) = _PERMUTE((_VAR##_1_0), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_2X2(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_2X1(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__); \\\n (_VAR##_1_1) = _PERMUTE((_VAR##_1_1), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_2X4(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_2X2(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \\\n (_VAR##_1_2) = _PERMUTE((_VAR##_1_2), ##__VA_ARGS__); \\\n (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__); \\\n (_VAR##_1_3) = _PERMUTE((_VAR##_1_3), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_2X8(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_2X4(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_4) = _PERMUTE((_VAR##_0_4), ##__VA_ARGS__); \\\n (_VAR##_1_4) = _PERMUTE((_VAR##_1_4), ##__VA_ARGS__); \\\n (_VAR##_0_5) = _PERMUTE((_VAR##_0_5), ##__VA_ARGS__); \\\n (_VAR##_1_5) = _PERMUTE((_VAR##_1_5), ##__VA_ARGS__); \\\n (_VAR##_0_6) = _PERMUTE((_VAR##_0_6), ##__VA_ARGS__); \\\n (_VAR##_1_6) = _PERMUTE((_VAR##_1_6), ##__VA_ARGS__); \\\n (_VAR##_0_7) = _PERMUTE((_VAR##_0_7), ##__VA_ARGS__); \\\n (_VAR##_1_7) = _PERMUTE((_VAR##_1_7), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_2X16(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_2X8(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_8) = _PERMUTE((_VAR##_0_8), ##__VA_ARGS__); \\\n (_VAR##_1_8) = _PERMUTE((_VAR##_1_8), ##__VA_ARGS__); \\\n (_VAR##_0_9) = _PERMUTE((_VAR##_0_9), ##__VA_ARGS__); \\\n (_VAR##_1_9) = _PERMUTE((_VAR##_1_9), ##__VA_ARGS__); \\\n (_VAR##_0_10) = _PERMUTE((_VAR##_0_10), ##__VA_ARGS__); \\\n (_VAR##_1_10) = _PERMUTE((_VAR##_1_10), ##__VA_ARGS__); \\\n (_VAR##_0_11) = _PERMUTE((_VAR##_0_11), ##__VA_ARGS__); \\\n (_VAR##_1_11) = _PERMUTE((_VAR##_1_11), ##__VA_ARGS__); \\\n (_VAR##_0_12) = _PERMUTE((_VAR##_0_12), ##__VA_ARGS__); \\\n (_VAR##_1_12) = _PERMUTE((_VAR##_1_12), ##__VA_ARGS__); \\\n (_VAR##_0_13) = _PERMUTE((_VAR##_0_13), ##__VA_ARGS__); \\\n (_VAR##_1_13) = _PERMUTE((_VAR##_1_13), ##__VA_ARGS__); \\\n (_VAR##_0_14) = _PERMUTE((_VAR##_0_14), ##__VA_ARGS__); \\\n (_VAR##_1_14) = _PERMUTE((_VAR##_1_14), ##__VA_ARGS__); \\\n (_VAR##_0_15) = _PERMUTE((_VAR##_0_15), ##__VA_ARGS__); \\\n (_VAR##_1_15) = _PERMUTE((_VAR##_1_15), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_2X32(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_2X16(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_16) = _PERMUTE((_VAR##_0_16), ##__VA_ARGS__); \\\n (_VAR##_1_16) = _PERMUTE((_VAR##_1_16), ##__VA_ARGS__); \\\n (_VAR##_0_17) = _PERMUTE((_VAR##_0_17), ##__VA_ARGS__); \\\n (_VAR##_1_17) = _PERMUTE((_VAR##_1_17), ##__VA_ARGS__); \\\n (_VAR##_0_18) = _PERMUTE((_VAR##_0_18), ##__VA_ARGS__); \\\n (_VAR##_1_18) = _PERMUTE((_VAR##_1_18), ##__VA_ARGS__); \\\n (_VAR##_0_19) = _PERMUTE((_VAR##_0_19), ##__VA_ARGS__); \\\n (_VAR##_1_19) = _PERMUTE((_VAR##_1_19), ##__VA_ARGS__); \\\n (_VAR##_0_20) = _PERMUTE((_VAR##_0_20), ##__VA_ARGS__); \\\n (_VAR##_1_20) = _PERMUTE((_VAR##_1_20), ##__VA_ARGS__); \\\n (_VAR##_0_21) = _PERMUTE((_VAR##_0_21), ##__VA_ARGS__); \\\n (_VAR##_1_21) = _PERMUTE((_VAR##_1_21), ##__VA_ARGS__); \\\n (_VAR##_0_22) = _PERMUTE((_VAR##_0_22), ##__VA_ARGS__); \\\n (_VAR##_1_22) = _PERMUTE((_VAR##_1_22), ##__VA_ARGS__); \\\n (_VAR##_0_23) = _PERMUTE((_VAR##_0_23), ##__VA_ARGS__); \\\n (_VAR##_1_23) = _PERMUTE((_VAR##_1_23), ##__VA_ARGS__); \\\n (_VAR##_0_24) = _PERMUTE((_VAR##_0_24), ##__VA_ARGS__); \\\n (_VAR##_1_24) = _PERMUTE((_VAR##_1_24), ##__VA_ARGS__); \\\n (_VAR##_0_25) = _PERMUTE((_VAR##_0_25), ##__VA_ARGS__); \\\n (_VAR##_1_25) = _PERMUTE((_VAR##_1_25), ##__VA_ARGS__); \\\n (_VAR##_0_26) = _PERMUTE((_VAR##_0_26), ##__VA_ARGS__); \\\n (_VAR##_1_26) = _PERMUTE((_VAR##_1_26), ##__VA_ARGS__); \\\n (_VAR##_0_27) = _PERMUTE((_VAR##_0_27), ##__VA_ARGS__); \\\n (_VAR##_1_27) = _PERMUTE((_VAR##_1_27), ##__VA_ARGS__); \\\n (_VAR##_0_28) = _PERMUTE((_VAR##_0_28), ##__VA_ARGS__); \\\n (_VAR##_1_28) = _PERMUTE((_VAR##_1_28), ##__VA_ARGS__); \\\n (_VAR##_0_29) = _PERMUTE((_VAR##_0_29), ##__VA_ARGS__); \\\n (_VAR##_1_29) = _PERMUTE((_VAR##_1_29), ##__VA_ARGS__); \\\n (_VAR##_0_30) = _PERMUTE((_VAR##_0_30), ##__VA_ARGS__); \\\n (_VAR##_1_30) = _PERMUTE((_VAR##_1_30), ##__VA_ARGS__); \\\n (_VAR##_0_31) = _PERMUTE((_VAR##_0_31), ##__VA_ARGS__); \\\n (_VAR##_1_31) = _PERMUTE((_VAR##_1_31), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_4X1(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_2X1(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_2_0) = _PERMUTE((_VAR##_2_0), ##__VA_ARGS__); \\\n (_VAR##_3_0) = _PERMUTE((_VAR##_3_0), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_4X2(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_4X1(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__); \\\n (_VAR##_1_1) = _PERMUTE((_VAR##_1_1), ##__VA_ARGS__); \\\n (_VAR##_2_1) = _PERMUTE((_VAR##_2_1), ##__VA_ARGS__); \\\n (_VAR##_3_1) = _PERMUTE((_VAR##_3_1), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_4X4(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_4X2(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \\\n (_VAR##_1_2) = _PERMUTE((_VAR##_1_2), ##__VA_ARGS__); \\\n (_VAR##_2_2) = _PERMUTE((_VAR##_2_2), ##__VA_ARGS__); \\\n (_VAR##_3_2) = _PERMUTE((_VAR##_3_2), ##__VA_ARGS__); \\\n (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__); \\\n (_VAR##_1_3) = _PERMUTE((_VAR##_1_3), ##__VA_ARGS__); \\\n (_VAR##_2_3) = _PERMUTE((_VAR##_2_3), ##__VA_ARGS__); \\\n (_VAR##_3_3) = _PERMUTE((_VAR##_3_3), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_4X8(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_4X4(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_4) = _PERMUTE((_VAR##_0_4), ##__VA_ARGS__); \\\n (_VAR##_1_4) = _PERMUTE((_VAR##_1_4), ##__VA_ARGS__); \\\n (_VAR##_2_4) = _PERMUTE((_VAR##_2_4), ##__VA_ARGS__); \\\n (_VAR##_3_4) = _PERMUTE((_VAR##_3_4), ##__VA_ARGS__); \\\n (_VAR##_0_5) = _PERMUTE((_VAR##_0_5), ##__VA_ARGS__); \\\n (_VAR##_1_5) = _PERMUTE((_VAR##_1_5), ##__VA_ARGS__); \\\n (_VAR##_2_5) = _PERMUTE((_VAR##_2_5), ##__VA_ARGS__); \\\n (_VAR##_3_5) = _PERMUTE((_VAR##_3_5), ##__VA_ARGS__); \\\n (_VAR##_0_6) = _PERMUTE((_VAR##_0_6), ##__VA_ARGS__); \\\n (_VAR##_1_6) = _PERMUTE((_VAR##_1_6), ##__VA_ARGS__); \\\n (_VAR##_2_6) = _PERMUTE((_VAR##_2_6), ##__VA_ARGS__); \\\n (_VAR##_3_6) = _PERMUTE((_VAR##_3_6), ##__VA_ARGS__); \\\n (_VAR##_0_7) = _PERMUTE((_VAR##_0_7), ##__VA_ARGS__); \\\n (_VAR##_1_7) = _PERMUTE((_VAR##_1_7), ##__VA_ARGS__); \\\n (_VAR##_2_7) = _PERMUTE((_VAR##_2_7), ##__VA_ARGS__); \\\n (_VAR##_3_7) = _PERMUTE((_VAR##_3_7), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_4X16(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_4X8(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_8) = _PERMUTE((_VAR##_0_8), ##__VA_ARGS__); \\\n (_VAR##_1_8) = _PERMUTE((_VAR##_1_8), ##__VA_ARGS__); \\\n (_VAR##_2_8) = _PERMUTE((_VAR##_2_8), ##__VA_ARGS__); \\\n (_VAR##_3_8) = _PERMUTE((_VAR##_3_8), ##__VA_ARGS__); \\\n (_VAR##_0_9) = _PERMUTE((_VAR##_0_9), ##__VA_ARGS__); \\\n (_VAR##_1_9) = _PERMUTE((_VAR##_1_9), ##__VA_ARGS__); \\\n (_VAR##_2_9) = _PERMUTE((_VAR##_2_9), ##__VA_ARGS__); \\\n (_VAR##_3_9) = _PERMUTE((_VAR##_3_9), ##__VA_ARGS__); \\\n (_VAR##_0_10) = _PERMUTE((_VAR##_0_10), ##__VA_ARGS__); \\\n (_VAR##_1_10) = _PERMUTE((_VAR##_1_10), ##__VA_ARGS__); \\\n (_VAR##_2_10) = _PERMUTE((_VAR##_2_10), ##__VA_ARGS__); \\\n (_VAR##_3_10) = _PERMUTE((_VAR##_3_10), ##__VA_ARGS__); \\\n (_VAR##_0_11) = _PERMUTE((_VAR##_0_11), ##__VA_ARGS__); \\\n (_VAR##_1_11) = _PERMUTE((_VAR##_1_11), ##__VA_ARGS__); \\\n (_VAR##_2_11) = _PERMUTE((_VAR##_2_11), ##__VA_ARGS__); \\\n (_VAR##_3_11) = _PERMUTE((_VAR##_3_11), ##__VA_ARGS__); \\\n (_VAR##_0_12) = _PERMUTE((_VAR##_0_12), ##__VA_ARGS__); \\\n (_VAR##_1_12) = _PERMUTE((_VAR##_1_12), ##__VA_ARGS__); \\\n (_VAR##_2_12) = _PERMUTE((_VAR##_2_12), ##__VA_ARGS__); \\\n (_VAR##_3_12) = _PERMUTE((_VAR##_3_12), ##__VA_ARGS__); \\\n (_VAR##_0_13) = _PERMUTE((_VAR##_0_13), ##__VA_ARGS__); \\\n (_VAR##_1_13) = _PERMUTE((_VAR##_1_13), ##__VA_ARGS__); \\\n (_VAR##_2_13) = _PERMUTE((_VAR##_2_13), ##__VA_ARGS__); \\\n (_VAR##_3_13) = _PERMUTE((_VAR##_3_13), ##__VA_ARGS__); \\\n (_VAR##_0_14) = _PERMUTE((_VAR##_0_14), ##__VA_ARGS__); \\\n (_VAR##_1_14) = _PERMUTE((_VAR##_1_14), ##__VA_ARGS__); \\\n (_VAR##_2_14) = _PERMUTE((_VAR##_2_14), ##__VA_ARGS__); \\\n (_VAR##_3_14) = _PERMUTE((_VAR##_3_14), ##__VA_ARGS__); \\\n (_VAR##_0_15) = _PERMUTE((_VAR##_0_15), ##__VA_ARGS__); \\\n (_VAR##_1_15) = _PERMUTE((_VAR##_1_15), ##__VA_ARGS__); \\\n (_VAR##_2_15) = _PERMUTE((_VAR##_2_15), ##__VA_ARGS__); \\\n (_VAR##_3_15) = _PERMUTE((_VAR##_3_15), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_4X32(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_4X16(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_16) = _PERMUTE((_VAR##_0_16), ##__VA_ARGS__); \\\n (_VAR##_1_16) = _PERMUTE((_VAR##_1_16), ##__VA_ARGS__); \\\n (_VAR##_2_16) = _PERMUTE((_VAR##_2_16), ##__VA_ARGS__); \\\n (_VAR##_3_16) = _PERMUTE((_VAR##_3_16), ##__VA_ARGS__); \\\n (_VAR##_0_17) = _PERMUTE((_VAR##_0_17), ##__VA_ARGS__); \\\n (_VAR##_1_17) = _PERMUTE((_VAR##_1_17), ##__VA_ARGS__); \\\n (_VAR##_2_17) = _PERMUTE((_VAR##_2_17), ##__VA_ARGS__); \\\n (_VAR##_3_17) = _PERMUTE((_VAR##_3_17), ##__VA_ARGS__); \\\n (_VAR##_0_18) = _PERMUTE((_VAR##_0_18), ##__VA_ARGS__); \\\n (_VAR##_1_18) = _PERMUTE((_VAR##_1_18), ##__VA_ARGS__); \\\n (_VAR##_2_18) = _PERMUTE((_VAR##_2_18), ##__VA_ARGS__); \\\n (_VAR##_3_18) = _PERMUTE((_VAR##_3_18), ##__VA_ARGS__); \\\n (_VAR##_0_19) = _PERMUTE((_VAR##_0_19), ##__VA_ARGS__); \\\n (_VAR##_1_19) = _PERMUTE((_VAR##_1_19), ##__VA_ARGS__); \\\n (_VAR##_2_19) = _PERMUTE((_VAR##_2_19), ##__VA_ARGS__); \\\n (_VAR##_3_19) = _PERMUTE((_VAR##_3_19), ##__VA_ARGS__); \\\n (_VAR##_0_20) = _PERMUTE((_VAR##_0_20), ##__VA_ARGS__); \\\n (_VAR##_1_20) = _PERMUTE((_VAR##_1_20), ##__VA_ARGS__); \\\n (_VAR##_2_20) = _PERMUTE((_VAR##_2_20), ##__VA_ARGS__); \\\n (_VAR##_3_20) = _PERMUTE((_VAR##_3_20), ##__VA_ARGS__); \\\n (_VAR##_0_21) = _PERMUTE((_VAR##_0_21), ##__VA_ARGS__); \\\n (_VAR##_1_21) = _PERMUTE((_VAR##_1_21), ##__VA_ARGS__); \\\n (_VAR##_2_21) = _PERMUTE((_VAR##_2_21), ##__VA_ARGS__); \\\n (_VAR##_3_21) = _PERMUTE((_VAR##_3_21), ##__VA_ARGS__); \\\n (_VAR##_0_22) = _PERMUTE((_VAR##_0_22), ##__VA_ARGS__); \\\n (_VAR##_1_22) = _PERMUTE((_VAR##_1_22), ##__VA_ARGS__); \\\n (_VAR##_2_22) = _PERMUTE((_VAR##_2_22), ##__VA_ARGS__); \\\n (_VAR##_3_22) = _PERMUTE((_VAR##_3_22), ##__VA_ARGS__); \\\n (_VAR##_0_23) = _PERMUTE((_VAR##_0_23), ##__VA_ARGS__); \\\n (_VAR##_1_23) = _PERMUTE((_VAR##_1_23), ##__VA_ARGS__); \\\n (_VAR##_2_23) = _PERMUTE((_VAR##_2_23), ##__VA_ARGS__); \\\n (_VAR##_3_23) = _PERMUTE((_VAR##_3_23), ##__VA_ARGS__); \\\n (_VAR##_0_24) = _PERMUTE((_VAR##_0_24), ##__VA_ARGS__); \\\n (_VAR##_1_24) = _PERMUTE((_VAR##_1_24), ##__VA_ARGS__); \\\n (_VAR##_2_24) = _PERMUTE((_VAR##_2_24), ##__VA_ARGS__); \\\n (_VAR##_3_24) = _PERMUTE((_VAR##_3_24), ##__VA_ARGS__); \\\n (_VAR##_0_25) = _PERMUTE((_VAR##_0_25), ##__VA_ARGS__); \\\n (_VAR##_1_25) = _PERMUTE((_VAR##_1_25), ##__VA_ARGS__); \\\n (_VAR##_2_25) = _PERMUTE((_VAR##_2_25), ##__VA_ARGS__); \\\n (_VAR##_3_25) = _PERMUTE((_VAR##_3_25), ##__VA_ARGS__); \\\n (_VAR##_0_26) = _PERMUTE((_VAR##_0_26), ##__VA_ARGS__); \\\n (_VAR##_1_26) = _PERMUTE((_VAR##_1_26), ##__VA_ARGS__); \\\n (_VAR##_2_26) = _PERMUTE((_VAR##_2_26), ##__VA_ARGS__); \\\n (_VAR##_3_26) = _PERMUTE((_VAR##_3_26), ##__VA_ARGS__); \\\n (_VAR##_0_27) = _PERMUTE((_VAR##_0_27), ##__VA_ARGS__); \\\n (_VAR##_1_27) = _PERMUTE((_VAR##_1_27), ##__VA_ARGS__); \\\n (_VAR##_2_27) = _PERMUTE((_VAR##_2_27), ##__VA_ARGS__); \\\n (_VAR##_3_27) = _PERMUTE((_VAR##_3_27), ##__VA_ARGS__); \\\n (_VAR##_0_28) = _PERMUTE((_VAR##_0_28), ##__VA_ARGS__); \\\n (_VAR##_1_28) = _PERMUTE((_VAR##_1_28), ##__VA_ARGS__); \\\n (_VAR##_2_28) = _PERMUTE((_VAR##_2_28), ##__VA_ARGS__); \\\n (_VAR##_3_28) = _PERMUTE((_VAR##_3_28), ##__VA_ARGS__); \\\n (_VAR##_0_29) = _PERMUTE((_VAR##_0_29), ##__VA_ARGS__); \\\n (_VAR##_1_29) = _PERMUTE((_VAR##_1_29), ##__VA_ARGS__); \\\n (_VAR##_2_29) = _PERMUTE((_VAR##_2_29), ##__VA_ARGS__); \\\n (_VAR##_3_29) = _PERMUTE((_VAR##_3_29), ##__VA_ARGS__); \\\n (_VAR##_0_30) = _PERMUTE((_VAR##_0_30), ##__VA_ARGS__); \\\n (_VAR##_1_30) = _PERMUTE((_VAR##_1_30), ##__VA_ARGS__); \\\n (_VAR##_2_30) = _PERMUTE((_VAR##_2_30), ##__VA_ARGS__); \\\n (_VAR##_3_30) = _PERMUTE((_VAR##_3_30), ##__VA_ARGS__); \\\n (_VAR##_0_31) = _PERMUTE((_VAR##_0_31), ##__VA_ARGS__); \\\n (_VAR##_1_31) = _PERMUTE((_VAR##_1_31), ##__VA_ARGS__); \\\n (_VAR##_2_31) = _PERMUTE((_VAR##_2_31), ##__VA_ARGS__); \\\n (_VAR##_3_31) = _PERMUTE((_VAR##_3_31), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_8X1(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_4X1(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_4_0) = _PERMUTE((_VAR##_4_0), ##__VA_ARGS__); \\\n (_VAR##_5_0) = _PERMUTE((_VAR##_5_0), ##__VA_ARGS__); \\\n (_VAR##_6_0) = _PERMUTE((_VAR##_6_0), ##__VA_ARGS__); \\\n (_VAR##_7_0) = _PERMUTE((_VAR##_7_0), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_8X2(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_8X1(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__); \\\n (_VAR##_1_1) = _PERMUTE((_VAR##_1_1), ##__VA_ARGS__); \\\n (_VAR##_2_1) = _PERMUTE((_VAR##_2_1), ##__VA_ARGS__); \\\n (_VAR##_3_1) = _PERMUTE((_VAR##_3_1), ##__VA_ARGS__); \\\n (_VAR##_4_1) = _PERMUTE((_VAR##_4_1), ##__VA_ARGS__); \\\n (_VAR##_5_1) = _PERMUTE((_VAR##_5_1), ##__VA_ARGS__); \\\n (_VAR##_6_1) = _PERMUTE((_VAR##_6_1), ##__VA_ARGS__); \\\n (_VAR##_7_1) = _PERMUTE((_VAR##_7_1), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_8X4(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_8X2(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \\\n (_VAR##_1_2) = _PERMUTE((_VAR##_1_2), ##__VA_ARGS__); \\\n (_VAR##_2_2) = _PERMUTE((_VAR##_2_2), ##__VA_ARGS__); \\\n (_VAR##_3_2) = _PERMUTE((_VAR##_3_2), ##__VA_ARGS__); \\\n (_VAR##_4_2) = _PERMUTE((_VAR##_4_2), ##__VA_ARGS__); \\\n (_VAR##_5_2) = _PERMUTE((_VAR##_5_2), ##__VA_ARGS__); \\\n (_VAR##_6_2) = _PERMUTE((_VAR##_6_2), ##__VA_ARGS__); \\\n (_VAR##_7_2) = _PERMUTE((_VAR##_7_2), ##__VA_ARGS__); \\\n (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__); \\\n (_VAR##_1_3) = _PERMUTE((_VAR##_1_3), ##__VA_ARGS__); \\\n (_VAR##_2_3) = _PERMUTE((_VAR##_2_3), ##__VA_ARGS__); \\\n (_VAR##_3_3) = _PERMUTE((_VAR##_3_3), ##__VA_ARGS__); \\\n (_VAR##_4_3) = _PERMUTE((_VAR##_4_3), ##__VA_ARGS__); \\\n (_VAR##_5_3) = _PERMUTE((_VAR##_5_3), ##__VA_ARGS__); \\\n (_VAR##_6_3) = _PERMUTE((_VAR##_6_3), ##__VA_ARGS__); \\\n (_VAR##_7_3) = _PERMUTE((_VAR##_7_3), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_8X8(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_8X4(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_4) = _PERMUTE((_VAR##_0_4), ##__VA_ARGS__); \\\n (_VAR##_1_4) = _PERMUTE((_VAR##_1_4), ##__VA_ARGS__); \\\n (_VAR##_2_4) = _PERMUTE((_VAR##_2_4), ##__VA_ARGS__); \\\n (_VAR##_3_4) = _PERMUTE((_VAR##_3_4), ##__VA_ARGS__); \\\n (_VAR##_4_4) = _PERMUTE((_VAR##_4_4), ##__VA_ARGS__); \\\n (_VAR##_5_4) = _PERMUTE((_VAR##_5_4), ##__VA_ARGS__); \\\n (_VAR##_6_4) = _PERMUTE((_VAR##_6_4), ##__VA_ARGS__); \\\n (_VAR##_7_4) = _PERMUTE((_VAR##_7_4), ##__VA_ARGS__); \\\n (_VAR##_0_5) = _PERMUTE((_VAR##_0_5), ##__VA_ARGS__); \\\n (_VAR##_1_5) = _PERMUTE((_VAR##_1_5), ##__VA_ARGS__); \\\n (_VAR##_2_5) = _PERMUTE((_VAR##_2_5), ##__VA_ARGS__); \\\n (_VAR##_3_5) = _PERMUTE((_VAR##_3_5), ##__VA_ARGS__); \\\n (_VAR##_4_5) = _PERMUTE((_VAR##_4_5), ##__VA_ARGS__); \\\n (_VAR##_5_5) = _PERMUTE((_VAR##_5_5), ##__VA_ARGS__); \\\n (_VAR##_6_5) = _PERMUTE((_VAR##_6_5), ##__VA_ARGS__); \\\n (_VAR##_7_5) = _PERMUTE((_VAR##_7_5), ##__VA_ARGS__); \\\n (_VAR##_0_6) = _PERMUTE((_VAR##_0_6), ##__VA_ARGS__); \\\n (_VAR##_1_6) = _PERMUTE((_VAR##_1_6), ##__VA_ARGS__); \\\n (_VAR##_2_6) = _PERMUTE((_VAR##_2_6), ##__VA_ARGS__); \\\n (_VAR##_3_6) = _PERMUTE((_VAR##_3_6), ##__VA_ARGS__); \\\n (_VAR##_4_6) = _PERMUTE((_VAR##_4_6), ##__VA_ARGS__); \\\n (_VAR##_5_6) = _PERMUTE((_VAR##_5_6), ##__VA_ARGS__); \\\n (_VAR##_6_6) = _PERMUTE((_VAR##_6_6), ##__VA_ARGS__); \\\n (_VAR##_7_6) = _PERMUTE((_VAR##_7_6), ##__VA_ARGS__); \\\n (_VAR##_0_7) = _PERMUTE((_VAR##_0_7), ##__VA_ARGS__); \\\n (_VAR##_1_7) = _PERMUTE((_VAR##_1_7), ##__VA_ARGS__); \\\n (_VAR##_2_7) = _PERMUTE((_VAR##_2_7), ##__VA_ARGS__); \\\n (_VAR##_3_7) = _PERMUTE((_VAR##_3_7), ##__VA_ARGS__); \\\n (_VAR##_4_7) = _PERMUTE((_VAR##_4_7), ##__VA_ARGS__); \\\n (_VAR##_5_7) = _PERMUTE((_VAR##_5_7), ##__VA_ARGS__); \\\n (_VAR##_6_7) = _PERMUTE((_VAR##_6_7), ##__VA_ARGS__); \\\n (_VAR##_7_7) = _PERMUTE((_VAR##_7_7), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_8X16(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_8X8(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_8) = _PERMUTE((_VAR##_0_8), ##__VA_ARGS__); \\\n (_VAR##_1_8) = _PERMUTE((_VAR##_1_8), ##__VA_ARGS__); \\\n (_VAR##_2_8) = _PERMUTE((_VAR##_2_8), ##__VA_ARGS__); \\\n (_VAR##_3_8) = _PERMUTE((_VAR##_3_8), ##__VA_ARGS__); \\\n (_VAR##_4_8) = _PERMUTE((_VAR##_4_8), ##__VA_ARGS__); \\\n (_VAR##_5_8) = _PERMUTE((_VAR##_5_8), ##__VA_ARGS__); \\\n (_VAR##_6_8) = _PERMUTE((_VAR##_6_8), ##__VA_ARGS__); \\\n (_VAR##_7_8) = _PERMUTE((_VAR##_7_8), ##__VA_ARGS__); \\\n (_VAR##_0_9) = _PERMUTE((_VAR##_0_9), ##__VA_ARGS__); \\\n (_VAR##_1_9) = _PERMUTE((_VAR##_1_9), ##__VA_ARGS__); \\\n (_VAR##_2_9) = _PERMUTE((_VAR##_2_9), ##__VA_ARGS__); \\\n (_VAR##_3_9) = _PERMUTE((_VAR##_3_9), ##__VA_ARGS__); \\\n (_VAR##_4_9) = _PERMUTE((_VAR##_4_9), ##__VA_ARGS__); \\\n (_VAR##_5_9) = _PERMUTE((_VAR##_5_9), ##__VA_ARGS__); \\\n (_VAR##_6_9) = _PERMUTE((_VAR##_6_9), ##__VA_ARGS__); \\\n (_VAR##_7_9) = _PERMUTE((_VAR##_7_9), ##__VA_ARGS__); \\\n (_VAR##_0_10) = _PERMUTE((_VAR##_0_10), ##__VA_ARGS__); \\\n (_VAR##_1_10) = _PERMUTE((_VAR##_1_10), ##__VA_ARGS__); \\\n (_VAR##_2_10) = _PERMUTE((_VAR##_2_10), ##__VA_ARGS__); \\\n (_VAR##_3_10) = _PERMUTE((_VAR##_3_10), ##__VA_ARGS__); \\\n (_VAR##_4_10) = _PERMUTE((_VAR##_4_10), ##__VA_ARGS__); \\\n (_VAR##_5_10) = _PERMUTE((_VAR##_5_10), ##__VA_ARGS__); \\\n (_VAR##_6_10) = _PERMUTE((_VAR##_6_10), ##__VA_ARGS__); \\\n (_VAR##_7_10) = _PERMUTE((_VAR##_7_10), ##__VA_ARGS__); \\\n (_VAR##_0_11) = _PERMUTE((_VAR##_0_11), ##__VA_ARGS__); \\\n (_VAR##_1_11) = _PERMUTE((_VAR##_1_11), ##__VA_ARGS__); \\\n (_VAR##_2_11) = _PERMUTE((_VAR##_2_11), ##__VA_ARGS__); \\\n (_VAR##_3_11) = _PERMUTE((_VAR##_3_11), ##__VA_ARGS__); \\\n (_VAR##_4_11) = _PERMUTE((_VAR##_4_11), ##__VA_ARGS__); \\\n (_VAR##_5_11) = _PERMUTE((_VAR##_5_11), ##__VA_ARGS__); \\\n (_VAR##_6_11) = _PERMUTE((_VAR##_6_11), ##__VA_ARGS__); \\\n (_VAR##_7_11) = _PERMUTE((_VAR##_7_11), ##__VA_ARGS__); \\\n (_VAR##_0_12) = _PERMUTE((_VAR##_0_12), ##__VA_ARGS__); \\\n (_VAR##_1_12) = _PERMUTE((_VAR##_1_12), ##__VA_ARGS__); \\\n (_VAR##_2_12) = _PERMUTE((_VAR##_2_12), ##__VA_ARGS__); \\\n (_VAR##_3_12) = _PERMUTE((_VAR##_3_12), ##__VA_ARGS__); \\\n (_VAR##_4_12) = _PERMUTE((_VAR##_4_12), ##__VA_ARGS__); \\\n (_VAR##_5_12) = _PERMUTE((_VAR##_5_12), ##__VA_ARGS__); \\\n (_VAR##_6_12) = _PERMUTE((_VAR##_6_12), ##__VA_ARGS__); \\\n (_VAR##_7_12) = _PERMUTE((_VAR##_7_12), ##__VA_ARGS__); \\\n (_VAR##_0_13) = _PERMUTE((_VAR##_0_13), ##__VA_ARGS__); \\\n (_VAR##_1_13) = _PERMUTE((_VAR##_1_13), ##__VA_ARGS__); \\\n (_VAR##_2_13) = _PERMUTE((_VAR##_2_13), ##__VA_ARGS__); \\\n (_VAR##_3_13) = _PERMUTE((_VAR##_3_13), ##__VA_ARGS__); \\\n (_VAR##_4_13) = _PERMUTE((_VAR##_4_13), ##__VA_ARGS__); \\\n (_VAR##_5_13) = _PERMUTE((_VAR##_5_13), ##__VA_ARGS__); \\\n (_VAR##_6_13) = _PERMUTE((_VAR##_6_13), ##__VA_ARGS__); \\\n (_VAR##_7_13) = _PERMUTE((_VAR##_7_13), ##__VA_ARGS__); \\\n (_VAR##_0_14) = _PERMUTE((_VAR##_0_14), ##__VA_ARGS__); \\\n (_VAR##_1_14) = _PERMUTE((_VAR##_1_14), ##__VA_ARGS__); \\\n (_VAR##_2_14) = _PERMUTE((_VAR##_2_14), ##__VA_ARGS__); \\\n (_VAR##_3_14) = _PERMUTE((_VAR##_3_14), ##__VA_ARGS__); \\\n (_VAR##_4_14) = _PERMUTE((_VAR##_4_14), ##__VA_ARGS__); \\\n (_VAR##_5_14) = _PERMUTE((_VAR##_5_14), ##__VA_ARGS__); \\\n (_VAR##_6_14) = _PERMUTE((_VAR##_6_14), ##__VA_ARGS__); \\\n (_VAR##_7_14) = _PERMUTE((_VAR##_7_14), ##__VA_ARGS__); \\\n (_VAR##_0_15) = _PERMUTE((_VAR##_0_15), ##__VA_ARGS__); \\\n (_VAR##_1_15) = _PERMUTE((_VAR##_1_15), ##__VA_ARGS__); \\\n (_VAR##_2_15) = _PERMUTE((_VAR##_2_15), ##__VA_ARGS__); \\\n (_VAR##_3_15) = _PERMUTE((_VAR##_3_15), ##__VA_ARGS__); \\\n (_VAR##_4_15) = _PERMUTE((_VAR##_4_15), ##__VA_ARGS__); \\\n (_VAR##_5_15) = _PERMUTE((_VAR##_5_15), ##__VA_ARGS__); \\\n (_VAR##_6_15) = _PERMUTE((_VAR##_6_15), ##__VA_ARGS__); \\\n (_VAR##_7_15) = _PERMUTE((_VAR##_7_15), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_8X32(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_8X16(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_16) = _PERMUTE((_VAR##_0_16), ##__VA_ARGS__); \\\n (_VAR##_1_16) = _PERMUTE((_VAR##_1_16), ##__VA_ARGS__); \\\n (_VAR##_2_16) = _PERMUTE((_VAR##_2_16), ##__VA_ARGS__); \\\n (_VAR##_3_16) = _PERMUTE((_VAR##_3_16), ##__VA_ARGS__); \\\n (_VAR##_4_16) = _PERMUTE((_VAR##_4_16), ##__VA_ARGS__); \\\n (_VAR##_5_16) = _PERMUTE((_VAR##_5_16), ##__VA_ARGS__); \\\n (_VAR##_6_16) = _PERMUTE((_VAR##_6_16), ##__VA_ARGS__); \\\n (_VAR##_7_16) = _PERMUTE((_VAR##_7_16), ##__VA_ARGS__); \\\n (_VAR##_0_17) = _PERMUTE((_VAR##_0_17), ##__VA_ARGS__); \\\n (_VAR##_1_17) = _PERMUTE((_VAR##_1_17), ##__VA_ARGS__); \\\n (_VAR##_2_17) = _PERMUTE((_VAR##_2_17), ##__VA_ARGS__); \\\n (_VAR##_3_17) = _PERMUTE((_VAR##_3_17), ##__VA_ARGS__); \\\n (_VAR##_4_17) = _PERMUTE((_VAR##_4_17), ##__VA_ARGS__); \\\n (_VAR##_5_17) = _PERMUTE((_VAR##_5_17), ##__VA_ARGS__); \\\n (_VAR##_6_17) = _PERMUTE((_VAR##_6_17), ##__VA_ARGS__); \\\n (_VAR##_7_17) = _PERMUTE((_VAR##_7_17), ##__VA_ARGS__); \\\n (_VAR##_0_18) = _PERMUTE((_VAR##_0_18), ##__VA_ARGS__); \\\n (_VAR##_1_18) = _PERMUTE((_VAR##_1_18), ##__VA_ARGS__); \\\n (_VAR##_2_18) = _PERMUTE((_VAR##_2_18), ##__VA_ARGS__); \\\n (_VAR##_3_18) = _PERMUTE((_VAR##_3_18), ##__VA_ARGS__); \\\n (_VAR##_4_18) = _PERMUTE((_VAR##_4_18), ##__VA_ARGS__); \\\n (_VAR##_5_18) = _PERMUTE((_VAR##_5_18), ##__VA_ARGS__); \\\n (_VAR##_6_18) = _PERMUTE((_VAR##_6_18), ##__VA_ARGS__); \\\n (_VAR##_7_18) = _PERMUTE((_VAR##_7_18), ##__VA_ARGS__); \\\n (_VAR##_0_19) = _PERMUTE((_VAR##_0_19), ##__VA_ARGS__); \\\n (_VAR##_1_19) = _PERMUTE((_VAR##_1_19), ##__VA_ARGS__); \\\n (_VAR##_2_19) = _PERMUTE((_VAR##_2_19), ##__VA_ARGS__); \\\n (_VAR##_3_19) = _PERMUTE((_VAR##_3_19), ##__VA_ARGS__); \\\n (_VAR##_4_19) = _PERMUTE((_VAR##_4_19), ##__VA_ARGS__); \\\n (_VAR##_5_19) = _PERMUTE((_VAR##_5_19), ##__VA_ARGS__); \\\n (_VAR##_6_19) = _PERMUTE((_VAR##_6_19), ##__VA_ARGS__); \\\n (_VAR##_7_19) = _PERMUTE((_VAR##_7_19), ##__VA_ARGS__); \\\n (_VAR##_0_20) = _PERMUTE((_VAR##_0_20), ##__VA_ARGS__); \\\n (_VAR##_1_20) = _PERMUTE((_VAR##_1_20), ##__VA_ARGS__); \\\n (_VAR##_2_20) = _PERMUTE((_VAR##_2_20), ##__VA_ARGS__); \\\n (_VAR##_3_20) = _PERMUTE((_VAR##_3_20), ##__VA_ARGS__); \\\n (_VAR##_4_20) = _PERMUTE((_VAR##_4_20), ##__VA_ARGS__); \\\n (_VAR##_5_20) = _PERMUTE((_VAR##_5_20), ##__VA_ARGS__); \\\n (_VAR##_6_20) = _PERMUTE((_VAR##_6_20), ##__VA_ARGS__); \\\n (_VAR##_7_20) = _PERMUTE((_VAR##_7_20), ##__VA_ARGS__); \\\n (_VAR##_0_21) = _PERMUTE((_VAR##_0_21), ##__VA_ARGS__); \\\n (_VAR##_1_21) = _PERMUTE((_VAR##_1_21), ##__VA_ARGS__); \\\n (_VAR##_2_21) = _PERMUTE((_VAR##_2_21), ##__VA_ARGS__); \\\n (_VAR##_3_21) = _PERMUTE((_VAR##_3_21), ##__VA_ARGS__); \\\n (_VAR##_4_21) = _PERMUTE((_VAR##_4_21), ##__VA_ARGS__); \\\n (_VAR##_5_21) = _PERMUTE((_VAR##_5_21), ##__VA_ARGS__); \\\n (_VAR##_6_21) = _PERMUTE((_VAR##_6_21), ##__VA_ARGS__); \\\n (_VAR##_7_21) = _PERMUTE((_VAR##_7_21), ##__VA_ARGS__); \\\n (_VAR##_0_22) = _PERMUTE((_VAR##_0_22), ##__VA_ARGS__); \\\n (_VAR##_1_22) = _PERMUTE((_VAR##_1_22), ##__VA_ARGS__); \\\n (_VAR##_2_22) = _PERMUTE((_VAR##_2_22), ##__VA_ARGS__); \\\n (_VAR##_3_22) = _PERMUTE((_VAR##_3_22), ##__VA_ARGS__); \\\n (_VAR##_4_22) = _PERMUTE((_VAR##_4_22), ##__VA_ARGS__); \\\n (_VAR##_5_22) = _PERMUTE((_VAR##_5_22), ##__VA_ARGS__); \\\n (_VAR##_6_22) = _PERMUTE((_VAR##_6_22), ##__VA_ARGS__); \\\n (_VAR##_7_22) = _PERMUTE((_VAR##_7_22), ##__VA_ARGS__); \\\n (_VAR##_0_23) = _PERMUTE((_VAR##_0_23), ##__VA_ARGS__); \\\n (_VAR##_1_23) = _PERMUTE((_VAR##_1_23), ##__VA_ARGS__); \\\n (_VAR##_2_23) = _PERMUTE((_VAR##_2_23), ##__VA_ARGS__); \\\n (_VAR##_3_23) = _PERMUTE((_VAR##_3_23), ##__VA_ARGS__); \\\n (_VAR##_4_23) = _PERMUTE((_VAR##_4_23), ##__VA_ARGS__); \\\n (_VAR##_5_23) = _PERMUTE((_VAR##_5_23), ##__VA_ARGS__); \\\n (_VAR##_6_23) = _PERMUTE((_VAR##_6_23), ##__VA_ARGS__); \\\n (_VAR##_7_23) = _PERMUTE((_VAR##_7_23), ##__VA_ARGS__); \\\n (_VAR##_0_24) = _PERMUTE((_VAR##_0_24), ##__VA_ARGS__); \\\n (_VAR##_1_24) = _PERMUTE((_VAR##_1_24), ##__VA_ARGS__); \\\n (_VAR##_2_24) = _PERMUTE((_VAR##_2_24), ##__VA_ARGS__); \\\n (_VAR##_3_24) = _PERMUTE((_VAR##_3_24), ##__VA_ARGS__); \\\n (_VAR##_4_24) = _PERMUTE((_VAR##_4_24), ##__VA_ARGS__); \\\n (_VAR##_5_24) = _PERMUTE((_VAR##_5_24), ##__VA_ARGS__); \\\n (_VAR##_6_24) = _PERMUTE((_VAR##_6_24), ##__VA_ARGS__); \\\n (_VAR##_7_24) = _PERMUTE((_VAR##_7_24), ##__VA_ARGS__); \\\n (_VAR##_0_25) = _PERMUTE((_VAR##_0_25), ##__VA_ARGS__); \\\n (_VAR##_1_25) = _PERMUTE((_VAR##_1_25), ##__VA_ARGS__); \\\n (_VAR##_2_25) = _PERMUTE((_VAR##_2_25), ##__VA_ARGS__); \\\n (_VAR##_3_25) = _PERMUTE((_VAR##_3_25), ##__VA_ARGS__); \\\n (_VAR##_4_25) = _PERMUTE((_VAR##_4_25), ##__VA_ARGS__); \\\n (_VAR##_5_25) = _PERMUTE((_VAR##_5_25), ##__VA_ARGS__); \\\n (_VAR##_6_25) = _PERMUTE((_VAR##_6_25), ##__VA_ARGS__); \\\n (_VAR##_7_25) = _PERMUTE((_VAR##_7_25), ##__VA_ARGS__); \\\n (_VAR##_0_26) = _PERMUTE((_VAR##_0_26), ##__VA_ARGS__); \\\n (_VAR##_1_26) = _PERMUTE((_VAR##_1_26), ##__VA_ARGS__); \\\n (_VAR##_2_26) = _PERMUTE((_VAR##_2_26), ##__VA_ARGS__); \\\n (_VAR##_3_26) = _PERMUTE((_VAR##_3_26), ##__VA_ARGS__); \\\n (_VAR##_4_26) = _PERMUTE((_VAR##_4_26), ##__VA_ARGS__); \\\n (_VAR##_5_26) = _PERMUTE((_VAR##_5_26), ##__VA_ARGS__); \\\n (_VAR##_6_26) = _PERMUTE((_VAR##_6_26), ##__VA_ARGS__); \\\n (_VAR##_7_26) = _PERMUTE((_VAR##_7_26), ##__VA_ARGS__); \\\n (_VAR##_0_27) = _PERMUTE((_VAR##_0_27), ##__VA_ARGS__); \\\n (_VAR##_1_27) = _PERMUTE((_VAR##_1_27), ##__VA_ARGS__); \\\n (_VAR##_2_27) = _PERMUTE((_VAR##_2_27), ##__VA_ARGS__); \\\n (_VAR##_3_27) = _PERMUTE((_VAR##_3_27), ##__VA_ARGS__); \\\n (_VAR##_4_27) = _PERMUTE((_VAR##_4_27), ##__VA_ARGS__); \\\n (_VAR##_5_27) = _PERMUTE((_VAR##_5_27), ##__VA_ARGS__); \\\n (_VAR##_6_27) = _PERMUTE((_VAR##_6_27), ##__VA_ARGS__); \\\n (_VAR##_7_27) = _PERMUTE((_VAR##_7_27), ##__VA_ARGS__); \\\n (_VAR##_0_28) = _PERMUTE((_VAR##_0_28), ##__VA_ARGS__); \\\n (_VAR##_1_28) = _PERMUTE((_VAR##_1_28), ##__VA_ARGS__); \\\n (_VAR##_2_28) = _PERMUTE((_VAR##_2_28), ##__VA_ARGS__); \\\n (_VAR##_3_28) = _PERMUTE((_VAR##_3_28), ##__VA_ARGS__); \\\n (_VAR##_4_28) = _PERMUTE((_VAR##_4_28), ##__VA_ARGS__); \\\n (_VAR##_5_28) = _PERMUTE((_VAR##_5_28), ##__VA_ARGS__); \\\n (_VAR##_6_28) = _PERMUTE((_VAR##_6_28), ##__VA_ARGS__); \\\n (_VAR##_7_28) = _PERMUTE((_VAR##_7_28), ##__VA_ARGS__); \\\n (_VAR##_0_29) = _PERMUTE((_VAR##_0_29), ##__VA_ARGS__); \\\n (_VAR##_1_29) = _PERMUTE((_VAR##_1_29), ##__VA_ARGS__); \\\n (_VAR##_2_29) = _PERMUTE((_VAR##_2_29), ##__VA_ARGS__); \\\n (_VAR##_3_29) = _PERMUTE((_VAR##_3_29), ##__VA_ARGS__); \\\n (_VAR##_4_29) = _PERMUTE((_VAR##_4_29), ##__VA_ARGS__); \\\n (_VAR##_5_29) = _PERMUTE((_VAR##_5_29), ##__VA_ARGS__); \\\n (_VAR##_6_29) = _PERMUTE((_VAR##_6_29), ##__VA_ARGS__); \\\n (_VAR##_7_29) = _PERMUTE((_VAR##_7_29), ##__VA_ARGS__); \\\n (_VAR##_0_30) = _PERMUTE((_VAR##_0_30), ##__VA_ARGS__); \\\n (_VAR##_1_30) = _PERMUTE((_VAR##_1_30), ##__VA_ARGS__); \\\n (_VAR##_2_30) = _PERMUTE((_VAR##_2_30), ##__VA_ARGS__); \\\n (_VAR##_3_30) = _PERMUTE((_VAR##_3_30), ##__VA_ARGS__); \\\n (_VAR##_4_30) = _PERMUTE((_VAR##_4_30), ##__VA_ARGS__); \\\n (_VAR##_5_30) = _PERMUTE((_VAR##_5_30), ##__VA_ARGS__); \\\n (_VAR##_6_30) = _PERMUTE((_VAR##_6_30), ##__VA_ARGS__); \\\n (_VAR##_7_30) = _PERMUTE((_VAR##_7_30), ##__VA_ARGS__); \\\n (_VAR##_0_31) = _PERMUTE((_VAR##_0_31), ##__VA_ARGS__); \\\n (_VAR##_1_31) = _PERMUTE((_VAR##_1_31), ##__VA_ARGS__); \\\n (_VAR##_2_31) = _PERMUTE((_VAR##_2_31), ##__VA_ARGS__); \\\n (_VAR##_3_31) = _PERMUTE((_VAR##_3_31), ##__VA_ARGS__); \\\n (_VAR##_4_31) = _PERMUTE((_VAR##_4_31), ##__VA_ARGS__); \\\n (_VAR##_5_31) = _PERMUTE((_VAR##_5_31), ##__VA_ARGS__); \\\n (_VAR##_6_31) = _PERMUTE((_VAR##_6_31), ##__VA_ARGS__); \\\n (_VAR##_7_31) = _PERMUTE((_VAR##_7_31), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_16X1(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_8X1(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_8_0) = _PERMUTE((_VAR##_8_0), ##__VA_ARGS__); \\\n (_VAR##_9_0) = _PERMUTE((_VAR##_9_0), ##__VA_ARGS__); \\\n (_VAR##_10_0) = _PERMUTE((_VAR##_10_0), ##__VA_ARGS__); \\\n (_VAR##_11_0) = _PERMUTE((_VAR##_11_0), ##__VA_ARGS__); \\\n (_VAR##_12_0) = _PERMUTE((_VAR##_12_0), ##__VA_ARGS__); \\\n (_VAR##_13_0) = _PERMUTE((_VAR##_13_0), ##__VA_ARGS__); \\\n (_VAR##_14_0) = _PERMUTE((_VAR##_14_0), ##__VA_ARGS__); \\\n (_VAR##_15_0) = _PERMUTE((_VAR##_15_0), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_16X2(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_16X1(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_1) = _PERMUTE((_VAR##_0_1), ##__VA_ARGS__); \\\n (_VAR##_1_1) = _PERMUTE((_VAR##_1_1), ##__VA_ARGS__); \\\n (_VAR##_2_1) = _PERMUTE((_VAR##_2_1), ##__VA_ARGS__); \\\n (_VAR##_3_1) = _PERMUTE((_VAR##_3_1), ##__VA_ARGS__); \\\n (_VAR##_4_1) = _PERMUTE((_VAR##_4_1), ##__VA_ARGS__); \\\n (_VAR##_5_1) = _PERMUTE((_VAR##_5_1), ##__VA_ARGS__); \\\n (_VAR##_6_1) = _PERMUTE((_VAR##_6_1), ##__VA_ARGS__); \\\n (_VAR##_7_1) = _PERMUTE((_VAR##_7_1), ##__VA_ARGS__); \\\n (_VAR##_8_1) = _PERMUTE((_VAR##_8_1), ##__VA_ARGS__); \\\n (_VAR##_9_1) = _PERMUTE((_VAR##_9_1), ##__VA_ARGS__); \\\n (_VAR##_10_1) = _PERMUTE((_VAR##_10_1), ##__VA_ARGS__); \\\n (_VAR##_11_1) = _PERMUTE((_VAR##_11_1), ##__VA_ARGS__); \\\n (_VAR##_12_1) = _PERMUTE((_VAR##_12_1), ##__VA_ARGS__); \\\n (_VAR##_13_1) = _PERMUTE((_VAR##_13_1), ##__VA_ARGS__); \\\n (_VAR##_14_1) = _PERMUTE((_VAR##_14_1), ##__VA_ARGS__); \\\n (_VAR##_15_1) = _PERMUTE((_VAR##_15_1), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PERMUTE_16X4(_VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_16X2(_VAR, _PERMUTE, ##__VA_ARGS__) \\\n (_VAR##_0_2) = _PERMUTE((_VAR##_0_2), ##__VA_ARGS__); \\\n (_VAR##_1_2) = _PERMUTE((_VAR##_1_2), ##__VA_ARGS__); \\\n (_VAR##_2_2) = _PERMUTE((_VAR##_2_2), ##__VA_ARGS__); \\\n (_VAR##_3_2) = _PERMUTE((_VAR##_3_2), ##__VA_ARGS__); \\\n (_VAR##_4_2) = _PERMUTE((_VAR##_4_2), ##__VA_ARGS__); \\\n (_VAR##_5_2) = _PERMUTE((_VAR##_5_2), ##__VA_ARGS__); \\\n (_VAR##_6_2) = _PERMUTE((_VAR##_6_2), ##__VA_ARGS__); \\\n (_VAR##_7_2) = _PERMUTE((_VAR##_7_2), ##__VA_ARGS__); \\\n (_VAR##_8_2) = _PERMUTE((_VAR##_8_2), ##__VA_ARGS__); \\\n (_VAR##_9_2) = _PERMUTE((_VAR##_9_2), ##__VA_ARGS__); \\\n (_VAR##_10_2) = _PERMUTE((_VAR##_10_2), ##__VA_ARGS__); \\\n (_VAR##_11_2) = _PERMUTE((_VAR##_11_2), ##__VA_ARGS__); \\\n (_VAR##_12_2) = _PERMUTE((_VAR##_12_2), ##__VA_ARGS__); \\\n (_VAR##_13_2) = _PERMUTE((_VAR##_13_2), ##__VA_ARGS__); \\\n (_VAR##_14_2) = _PERMUTE((_VAR##_14_2), ##__VA_ARGS__); \\\n (_VAR##_15_2) = _PERMUTE((_VAR##_15_2), ##__VA_ARGS__); \\\n (_VAR##_0_3) = _PERMUTE((_VAR##_0_3), ##__VA_ARGS__); \\\n (_VAR##_1_3) = _PERMUTE((_VAR##_1_3), ##__VA_ARGS__); \\\n (_VAR##_2_3) = _PERMUTE((_VAR##_2_3), ##__VA_ARGS__); \\\n (_VAR##_3_3) = _PERMUTE((_VAR##_3_3), ##__VA_ARGS__); \\\n (_VAR##_4_3) = _PERMUTE((_VAR##_4_3), ##__VA_ARGS__); \\\n (_VAR##_5_3) = _PERMUTE((_VAR##_5_3), ##__VA_ARGS__); \\\n (_VAR##_6_3) = _PERMUTE((_VAR##_6_3), ##__VA_ARGS__); \\\n (_VAR##_7_3) = _PERMUTE((_VAR##_7_3), ##__VA_ARGS__); \\\n (_VAR##_8_3) = _PERMUTE((_VAR##_8_3), ##__VA_ARGS__); \\\n (_VAR##_9_3) = _PERMUTE((_VAR##_9_3), ##__VA_ARGS__); \\\n (_VAR##_10_3) = _PERMUTE((_VAR##_10_3), ##__VA_ARGS__); \\\n (_VAR##_11_3) = _PERMUTE((_VAR##_11_3), ##__VA_ARGS__); \\\n (_VAR##_12_3) = _PERMUTE((_VAR##_12_3), ##__VA_ARGS__); \\\n (_VAR##_13_3) = _PERMUTE((_VAR##_13_3), ##__VA_ARGS__); \\\n (_VAR##_14_3) = _PERMUTE((_VAR##_14_3), ##__VA_ARGS__); \\\n (_VAR##_15_3) = _PERMUTE((_VAR##_15_3), ##__VA_ARGS__);\n\n#define MATRIX_VAR_PROC_2X1(_K, _LHS, _RHS, _RES, _PROCESS) \\\n _PROCESS((_LHS##_0), (_RHS), (_RES##_0_##_K)) \\\n _PROCESS((_LHS##_1), (_RHS), (_RES##_1_##_K))\n\n#define MATRIX_VAR_PROC_4X1(_K, _LHS, _RHS, _RES, _PROCESS) \\\n MATRIX_VAR_PROC_2X1(_K, _LHS, _RHS, _RES, _PROCESS) \\\n _PROCESS((_LHS##_2), (_RHS), (_RES##_2_##_K)) \\\n _PROCESS((_LHS##_3), (_RHS), (_RES##_3_##_K))\n\n#define MATRIX_VAR_PROC_8X1(_K, _LHS, _RHS, _RES, _PROCESS) \\\n MATRIX_VAR_PROC_4X1(_K, _LHS, _RHS, _RES, _PROCESS) \\\n _PROCESS((_LHS##_4), (_RHS), (_RES##_4_##_K)) \\\n _PROCESS((_LHS##_5), (_RHS), (_RES##_5_##_K)) \\\n _PROCESS((_LHS##_6), (_RHS), (_RES##_6_##_K)) \\\n _PROCESS((_LHS##_7), (_RHS), (_RES##_7_##_K))\n\n#define MATRIX_VAR_PROC_16X1(_K, _LHS, _RHS, _RES, _PROCESS) \\\n MATRIX_VAR_PROC_8X1(_K, _LHS, _RHS, _RES, _PROCESS) \\\n _PROCESS((_LHS##_8), (_RHS), (_RES##_8_##_K)) \\\n _PROCESS((_LHS##_9), (_RHS), (_RES##_9_##_K)) \\\n _PROCESS((_LHS##_10), (_RHS), (_RES##_10_##_K)) \\\n _PROCESS((_LHS##_11), (_RHS), (_RES##_11_##_K)) \\\n _PROCESS((_LHS##_12), (_RHS), (_RES##_12_##_K)) \\\n _PROCESS((_LHS##_13), (_RHS), (_RES##_13_##_K)) \\\n _PROCESS((_LHS##_14), (_RHS), (_RES##_14_##_K)) \\\n _PROCESS((_LHS##_15), (_RHS), (_RES##_15_##_K))\n\n#define MATRIX_VAR_PROC_1X2(_K, _LHS, _RHS, _RES, _PROCESS) \\\n _PROCESS((_LHS), (_RHS##_0), (_RES##_##_K##_0)) \\\n _PROCESS((_LHS), (_RHS##_1), (_RES##_##_K##_1))\n\n#define MATRIX_VAR_PROC_1X4(_K, _LHS, _RHS, _RES, _PROCESS) \\\n MATRIX_VAR_PROC_1X2(_K, _LHS, _RHS, _RES, _PROCESS) \\\n _PROCESS((_LHS), (_RHS##_2), (_RES##_##_K##_2)) \\\n _PROCESS((_LHS), (_RHS##_3), (_RES##_##_K##_3))\n\n#define MATRIX_VAR_PROC_1X8(_K, _LHS, _RHS, _RES, _PROCESS) \\\n MATRIX_VAR_PROC_1X4(_K, _LHS, _RHS, _RES, _PROCESS) \\\n _PROCESS((_LHS), (_RHS##_4), (_RES##_##_K##_4)) \\\n _PROCESS((_LHS), (_RHS##_5), (_RES##_##_K##_5)) \\\n _PROCESS((_LHS), (_RHS##_6), (_RES##_##_K##_6)) \\\n _PROCESS((_LHS), (_RHS##_7), (_RES##_##_K##_7))\n\n#define MATRIX_VAR_PROC_1X16(_K, _LHS, _RHS, _RES, _PROCESS) \\\n MATRIX_VAR_PROC_1X8(_K, _LHS, _RHS, _RES, _PROCESS) \\\n _PROCESS((_LHS), (_RHS##_8), (_RES##_##_K##_8)) \\\n _PROCESS((_LHS), (_RHS##_9), (_RES##_##_K##_9)) \\\n _PROCESS((_LHS), (_RHS##_10), (_RES##_##_K##_10)) \\\n _PROCESS((_LHS), (_RHS##_11), (_RES##_##_K##_11)) \\\n _PROCESS((_LHS), (_RHS##_12), (_RES##_##_K##_12)) \\\n _PROCESS((_LHS), (_RHS##_13), (_RES##_##_K##_13)) \\\n _PROCESS((_LHS), (_RHS##_14), (_RES##_##_K##_14)) \\\n _PROCESS((_LHS), (_RHS##_15), (_RES##_##_K##_15))\n\n#define MATRIX_VAR_INIT(_M, _N, _VAR_TYPE, _VAR_NAME, _VAR_INIT) \\\n MATRIX_VAR_INIT_##_M##X##_N(_VAR_TYPE, _VAR_NAME, _VAR_INIT)\n\n#define MATRIX_VAR_STORE(_M, _N, _STEP, _VAR, _ARRAY, _STORE, _NORM, ...) \\\n MATRIX_VAR_STORE_##_M##X##_N(_STEP, _VAR, _ARRAY, _STORE, _NORM, \\\n ##__VA_ARGS__)\n\n#define MATRIX_VAR_PERMUTE(_M, _N, _VAR, _PERMUTE, ...) \\\n MATRIX_VAR_PERMUTE_##_M##X##_N(_VAR, _PERMUTE, ##__VA_ARGS__)\n\n#define MATRIX_VAR_PROC(_M, _N, _K, _LHS, _RHS, _RES, _PROCESS) \\\n MATRIX_VAR_PROC_##_M##X##_N(_K, _LHS, _RHS, _RES, _PROCESS)\n" }, { "path": "src/ailego/math/matrix_utility.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n\nnamespace zvec {\nnamespace ailego {\n\n//! Absolute value of a float\nstatic inline float FastAbs(float x) {\n uint32_t *p = reinterpret_cast(&x);\n *p &= 0x7fffffffu;\n return *reinterpret_cast(p);\n}\n\n#if defined(__SSE__)\nstatic inline float HorizontalMax_FP32_V128(__m128 v) {\n __m128 x1 = _mm_shuffle_ps(v, v, _MM_SHUFFLE(0, 0, 3, 2));\n __m128 x2 = _mm_max_ps(v, x1);\n __m128 x3 = _mm_shuffle_ps(x2, x2, _MM_SHUFFLE(0, 0, 0, 1));\n __m128 x4 = _mm_max_ps(x2, x3);\n return _mm_cvtss_f32(x4);\n}\n\nstatic inline float HorizontalAdd_FP32_V128(__m128 v) {\n#ifdef __SSE3__\n __m128 x1 = _mm_hadd_ps(v, v);\n __m128 x2 = _mm_hadd_ps(x1, x1);\n return _mm_cvtss_f32(x2);\n#else\n __m128 x1 = _mm_movehl_ps(v, v);\n __m128 x2 = _mm_add_ps(v, x1);\n __m128 x3 = _mm_shuffle_ps(x2, x2, 1);\n __m128 x4 = _mm_add_ss(x2, x3);\n return _mm_cvtss_f32(x4);\n#endif\n}\n#endif // __SSE__\n\n#if defined(__SSE2__)\nstatic inline int32_t HorizontalAdd_INT32_V128(__m128i v) {\n#ifdef __SSE3__\n __m128i x1 = _mm_hadd_epi32(v, v);\n __m128i x2 = _mm_hadd_epi32(x1, x1);\n return _mm_cvtsi128_si32(x2);\n#else\n __m128i x1 = _mm_shuffle_epi32(v, _MM_SHUFFLE(0, 0, 3, 2));\n __m128i x2 = _mm_add_epi32(v, x1);\n __m128i x3 = _mm_shuffle_epi32(x2, _MM_SHUFFLE(0, 0, 0, 1));\n __m128i x4 = _mm_add_epi32(x2, x3);\n return _mm_cvtsi128_si32(x4);\n#endif\n}\n\nstatic inline int64_t HorizontalAdd_INT64_V128(__m128i v) {\n#ifdef __SSE4_1__\n return (_mm_extract_epi64(v, 0) + _mm_extract_epi64(v, 1));\n#else\n return _mm_cvtsi128_si64(\n _mm_add_epi64(_mm_shuffle_epi32(v, _MM_SHUFFLE(0, 0, 3, 2)), v));\n#endif\n}\n#endif // __SSE2__\n\n#if defined(__SSSE3__)\nstatic const __m128i POPCNT_LOOKUP_SSE =\n _mm_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4);\n\nstatic inline __m128i VerticalPopCount_INT8_V128(__m128i v) {\n#if defined(__AVX512VL__) && defined(__AVX512BITALG__)\n return _mm_popcnt_epi8(v);\n#else\n const __m128i low_mask = _mm_set1_epi8(0x0f);\n __m128i lo = _mm_shuffle_epi8(POPCNT_LOOKUP_SSE, _mm_and_si128(v, low_mask));\n __m128i hi = _mm_shuffle_epi8(POPCNT_LOOKUP_SSE,\n _mm_and_si128(_mm_srli_epi32(v, 4), low_mask));\n return _mm_add_epi8(lo, hi);\n#endif // __AVX512VL__ && __AVX512BITALG__\n}\n\nstatic inline __m128i VerticalPopCount_INT16_V128(__m128i v) {\n#if defined(__AVX512VL__) && defined(__AVX512BITALG__)\n return _mm_popcnt_epi16(v);\n#else\n __m128i total = VerticalPopCount_INT8_V128(v);\n return _mm_add_epi16(_mm_srli_epi16(total, 8),\n _mm_and_si128(total, _mm_set1_epi16(0xff)));\n#endif // __AVX512VL__ && __AVX512BITALG__\n}\n\nstatic inline __m128i VerticalPopCount_INT32_V128(__m128i v) {\n#if defined(__AVX512VL__) && defined(__AVX512VPOPCNTDQ__)\n return _mm_popcnt_epi32(v);\n#else\n __m128i total =\n _mm_madd_epi16(VerticalPopCount_INT8_V128(v), _mm_set1_epi16(1));\n return _mm_add_epi32(_mm_srli_epi32(total, 8),\n _mm_and_si128(total, _mm_set1_epi32(0xff)));\n#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__\n}\n\nstatic inline __m128i VerticalPopCount_INT64_V128(__m128i v) {\n#if defined(__AVX512VL__) && defined(__AVX512VPOPCNTDQ__)\n return _mm_popcnt_epi64(v);\n#else\n return _mm_sad_epu8(VerticalPopCount_INT8_V128(v), _mm_setzero_si128());\n#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__\n}\n#endif // __SSSE3__\n\n#if defined(__SSE4_1__)\nstatic inline int16_t HorizontalMax_UINT8_V128(__m128i v) {\n v = _mm_max_epu8(v, _mm_shuffle_epi32(v, _MM_SHUFFLE(3, 2, 3, 2)));\n v = _mm_max_epu8(v, _mm_shuffle_epi32(v, _MM_SHUFFLE(1, 1, 1, 1)));\n v = _mm_max_epu8(v, _mm_shufflelo_epi16(v, _MM_SHUFFLE(1, 1, 1, 1)));\n v = _mm_max_epu8(v, _mm_srli_epi16(v, 8));\n return static_cast(_mm_cvtsi128_si32(v));\n}\n#endif // __SSE4_1__\n\n#if defined(__AVX__)\nstatic inline float HorizontalMax_FP32_V256(__m256 v) {\n __m256 x1 = _mm256_permute_ps(v, _MM_SHUFFLE(0, 0, 3, 2));\n __m256 x2 = _mm256_max_ps(v, x1);\n __m256 x3 = _mm256_permute_ps(x2, _MM_SHUFFLE(0, 0, 0, 1));\n __m256 x4 = _mm256_max_ps(x2, x3);\n __m128 x5 = _mm256_extractf128_ps(x4, 1);\n __m128 x6 = _mm_max_ss(_mm256_castps256_ps128(x4), x5);\n return _mm_cvtss_f32(x6);\n}\n\nstatic inline float HorizontalAdd_FP32_V256(__m256 v) {\n __m256 x1 = _mm256_hadd_ps(v, v);\n __m256 x2 = _mm256_hadd_ps(x1, x1);\n __m128 x3 = _mm256_extractf128_ps(x2, 1);\n __m128 x4 = _mm_add_ss(_mm256_castps256_ps128(x2), x3);\n return _mm_cvtss_f32(x4);\n}\n#endif // __AVX__\n\n#if defined(__AVX2__)\n#define POPCNT_MASK1_INT8_AVX _mm256_set1_epi8(0x0f)\n#define POPCNT_MASK1_INT16_AVX _mm256_set1_epi16(1)\n#define POPCNT_MASK2_INT16_AVX _mm256_set1_epi16(0xff)\n#define POPCNT_MASK1_INT32_AVX _mm256_set1_epi32(0xff)\n#define POPCNT_ZERO_AVX _mm256_setzero_si256()\n#define POPCNT_LOOKUP_AVX _mm256_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4)\n\nstatic inline __m256i VerticalPopCount_INT8_V256(__m256i v) {\n#if defined(__AVX512VL__) && defined(__AVX512BITALG__)\n return _mm256_popcnt_epi8(v);\n#else\n __m256i lo = _mm256_shuffle_epi8(POPCNT_LOOKUP_AVX,\n _mm256_and_si256(v, POPCNT_MASK1_INT8_AVX));\n __m256i hi = _mm256_shuffle_epi8(\n POPCNT_LOOKUP_AVX,\n _mm256_and_si256(_mm256_srli_epi32(v, 4), POPCNT_MASK1_INT8_AVX));\n return _mm256_add_epi8(lo, hi);\n#endif // __AVX512VL__ && __AVX512BITALG__\n}\n\nstatic inline __m256i VerticalPopCount_INT16_V256(__m256i v) {\n#if defined(__AVX512VL__) && defined(__AVX512BITALG__)\n return _mm256_popcnt_epi16(v);\n#else\n __m256i total = VerticalPopCount_INT8_V256(v);\n return _mm256_add_epi16(_mm256_srli_epi16(total, 8),\n _mm256_and_si256(total, POPCNT_MASK2_INT16_AVX));\n#endif // __AVX512VL__ && __AVX512BITALG__\n}\n\nstatic inline __m256i VerticalPopCount_INT32_V256(__m256i v) {\n#if defined(__AVX512VL__) && defined(__AVX512VPOPCNTDQ__)\n return _mm256_popcnt_epi32(v);\n#else\n __m256i total =\n _mm256_madd_epi16(VerticalPopCount_INT8_V256(v), POPCNT_MASK1_INT16_AVX);\n return _mm256_add_epi32(_mm256_srli_epi32(total, 8),\n _mm256_and_si256(total, POPCNT_MASK1_INT32_AVX));\n#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__\n}\n\nstatic inline __m256i VerticalPopCount_INT64_V256(__m256i v) {\n#if defined(__AVX512VL__) && defined(__AVX512VPOPCNTDQ__)\n return _mm256_popcnt_epi64(v);\n#else\n return _mm256_sad_epu8(VerticalPopCount_INT8_V256(v), POPCNT_ZERO_AVX);\n#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__\n}\n\nstatic inline int16_t HorizontalMax_UINT8_V256(__m256i v) {\n v = _mm256_max_epu8(v, _mm256_shuffle_epi32(v, _MM_SHUFFLE(3, 2, 3, 2)));\n v = _mm256_max_epu8(v, _mm256_shuffle_epi32(v, _MM_SHUFFLE(1, 1, 1, 1)));\n v = _mm256_max_epu8(v, _mm256_shufflelo_epi16(v, _MM_SHUFFLE(1, 1, 1, 1)));\n __m128i x =\n _mm_max_epu8(_mm256_castsi256_si128(v), _mm256_extractf128_si256(v, 1));\n x = _mm_max_epu8(x, _mm_srli_epi16(x, 8));\n return static_cast(_mm_cvtsi128_si32(x));\n}\n\nstatic inline int32_t HorizontalAdd_INT32_V256(__m256i v) {\n __m256i x1 = _mm256_hadd_epi32(v, v);\n __m256i x2 = _mm256_hadd_epi32(x1, x1);\n __m128i x3 = _mm256_extractf128_si256(x2, 1);\n __m128i x4 = _mm_add_epi32(_mm256_castsi256_si128(x2), x3);\n return _mm_cvtsi128_si32(x4);\n}\n\nstatic inline int64_t HorizontalAdd_INT64_V256(__m256i v) {\n __m256i x1 = _mm256_shuffle_epi32(v, _MM_SHUFFLE(1, 0, 3, 2));\n __m256i x2 = _mm256_add_epi64(v, x1);\n __m128i x3 = _mm256_extractf128_si256(x2, 1);\n __m128i x4 = _mm_add_epi64(_mm256_extractf128_si256(x2, 0), x3);\n return _mm_cvtsi128_si64(x4);\n}\n#endif // __AVX2__\n\n#if defined(__AVX512F__)\nstatic inline float HorizontalMax_FP32_V512(__m512 v) {\n __m256 low = _mm512_castps512_ps256(v);\n __m256 high =\n _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(v), 1));\n return HorizontalMax_FP32_V256(_mm256_max_ps(low, high));\n}\n\nstatic inline float HorizontalAdd_FP32_V512(__m512 v) {\n __m256 low = _mm512_castps512_ps256(v);\n __m256 high =\n _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(v), 1));\n return HorizontalAdd_FP32_V256(_mm256_add_ps(low, high));\n}\n#endif // __AVX512F__\n\n#if defined(__AVX512FP16__)\nstatic inline float HorizontalMax_FP16_V512(__m512h v) {\n __m512 low = _mm512_cvtxph_ps(_mm512_castph512_ph256(v));\n __m512 high = _mm512_cvtxph_ps(\n _mm256_castpd_ph(_mm512_extractf64x4_pd(_mm512_castph_pd(v), 1)));\n return HorizontalMax_FP32_V512(_mm512_max_ps(low, high));\n}\n\nstatic inline float HorizontalAdd_FP16_V512(__m512h v) {\n __m512 low = _mm512_cvtxph_ps(_mm512_castph512_ph256(v));\n __m512 high = _mm512_cvtxph_ps(\n _mm256_castpd_ph(_mm512_extractf64x4_pd(_mm512_castph_pd(v), 1)));\n\n return HorizontalAdd_FP32_V512(_mm512_add_ps(low, high));\n}\n#endif // __AVX512FP16__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \"distance_utility.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n/*! Compute the Mips SphericalInjection Squared Euclidean Distance with the two\n * vectors's InnerProduct and each squared l2-normlized value, and the e2 is\n * 1.0 / max_squared_l2_norm\n */\nstatic float inline ComputeSphericalInjection(double ip, double u2, double v2,\n double e2) {\n if (e2 == 0.0) {\n // Implies *localized* spherical injection.\n return static_cast(2.0 - 2.0 * ip / std::max(u2, v2));\n }\n auto v = (1.0 - e2 * u2) * (1.0 - e2 * v2);\n auto score = v > 0.0 ? (1.0 - e2 * ip - std::sqrt(v)) : (1.0 - e2 * ip);\n return static_cast(score * 2.0);\n}\n\n/*! Mips Squared Euclidean Distance Matrix\n */\ntemplate \nstruct MipsSquaredEuclideanDistanceMatrix;\n\n/*! Mips Squared Euclidean Distance Matrix (M=1, N=1)\n */\ntemplate \nstruct MipsSquaredEuclideanDistanceMatrix {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n // Compute the distance between matrix and query by SphericalInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out) {\n ailego_assert(p && q && dim && out);\n\n float sum = 0.0;\n float u2 = 0.0;\n float v2 = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n u2 += p[i] * p[i];\n v2 += q[i] * q[i];\n sum += static_cast(p[i] * q[i]);\n }\n *out = ComputeSphericalInjection(sum, u2, v2, e2);\n }\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out) {\n ailego_assert(p && q && dim && out);\n\n float sum = 0.0;\n float u2 = 0.0;\n float v2 = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n u2 += p[i] * p[i];\n v2 += q[i] * q[i];\n sum += MathHelper::SquaredDifference(p[i], q[i]);\n }\n\n sum *= e2;\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n *out = sum;\n }\n};\n\ntemplate <>\nstruct MipsSquaredEuclideanDistanceMatrix {\n //! Type of value\n using ValueType = uint8_t;\n\n // Compute the distance between matrix and query by SphericalInjection\n static void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out);\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out);\n};\n\ntemplate <>\nstruct MipsSquaredEuclideanDistanceMatrix {\n //! Type of value\n using ValueType = int8_t;\n\n // Compute the distance between matrix and query by SphericalInjection\n static void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out);\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out);\n};\n\ntemplate <>\nstruct MipsSquaredEuclideanDistanceMatrix {\n //! Type of value\n using ValueType = Float16;\n\n // Compute the distance between matrix and query by SphericalInjection\n static void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out);\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out);\n};\n\ntemplate <>\nstruct MipsSquaredEuclideanDistanceMatrix {\n //! Type of value\n using ValueType = float;\n\n // Compute the distance between matrix and query by SphericalInjection\n static void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out);\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out);\n};\n\n/*! Mips Squared Euclidean Distance Matrix (M >= 2, N >= 2)\n */\ntemplate \nstruct MipsSquaredEuclideanDistanceMatrix<\n T, M, N, typename std::enable_if= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n // Compute the distance between matrix and query by SphericalInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out) {\n ailego_assert(p && q && dim && out);\n if (dim == 0) {\n return;\n }\n\n std::array u2;\n std::array v2;\n for (size_t i = 0; i < M; ++i) {\n const ValueType p_val = p[i];\n u2[i] = static_cast(p_val * p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r = static_cast(p_val * q[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] = static_cast(q[i] * q[i]);\n }\n p += M;\n q += N;\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n const ValueType p_val = p[i];\n u2[i] += static_cast(p_val * p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r += static_cast(p_val * q[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] += static_cast(q[i] * q[i]);\n }\n p += M;\n q += N;\n }\n\n // Compute the injection\n for (size_t i = 0; i < M; ++i) {\n float *r = out + i;\n const float u2_val = u2[i];\n for (size_t j = 0; j < N; ++j) {\n *r = ComputeSphericalInjection(*r, u2_val, v2[j], e2);\n r += M;\n }\n }\n }\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out) {\n ailego_assert(p && q && dim && out);\n if (dim == 0) {\n return;\n }\n\n std::array u2;\n std::array v2;\n for (size_t i = 0; i < M; ++i) {\n const ValueType p_val = p[i];\n u2[i] = static_cast(p_val * p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r = MathHelper::SquaredDifference(p_val, q[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] = static_cast(q[i] * q[i]);\n }\n p += M;\n q += N;\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n const ValueType p_val = p[i];\n u2[i] += static_cast(p_val * p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r += MathHelper::SquaredDifference(p_val, q[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] += static_cast(q[i] * q[i]);\n }\n p += M;\n q += N;\n }\n\n // Compute the injections\n float *r = out;\n for (size_t i = 0; i < M; ++i) {\n u2[i] *= e2;\n for (size_t j = 0; j < N; ++j) {\n (*r++) *= e2;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] *= e2;\n }\n for (size_t k = 0; k < m; ++k) {\n for (size_t i = 0; i < M; ++i) {\n r = out + i;\n float u2_val = u2[i];\n u2[i] = u2_val * u2_val;\n for (size_t j = 0; j < N; ++j) {\n *r += (u2_val - v2[j]) * (u2_val - v2[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] = v2[i] * v2[i];\n }\n }\n }\n};\n\n/*! Mips Squared Euclidean Distance Matrix (N=1)\n */\ntemplate \nstruct MipsSquaredEuclideanDistanceMatrix<\n T, M, 1, typename std::enable_if= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n // Compute the distance between matrix and query by SphericalInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out) {\n ailego_assert(p && q && dim && out);\n const ValueType *q_end = q + dim;\n if (q == q_end) {\n return;\n }\n\n std::array u2;\n ValueType q_val = *q++;\n float v2 = static_cast(q_val * q_val);\n for (size_t i = 0; i < M; ++i) {\n u2[i] = static_cast(p[i] * p[i]);\n out[i] = static_cast(p[i] * q_val);\n }\n p += M;\n\n while (q != q_end) {\n q_val = *q++;\n v2 += static_cast(q_val * q_val);\n for (size_t i = 0; i < M; ++i) {\n u2[i] += static_cast(p[i] * p[i]);\n out[i] += static_cast(p[i] * q_val);\n }\n p += M;\n }\n\n // Compute the injection\n for (size_t i = 0; i < M; ++i) {\n out[i] = ComputeSphericalInjection(out[i], u2[i], v2, e2);\n }\n }\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out) {\n ailego_assert(p && q && dim && out);\n const ValueType *q_end = q + dim;\n if (q == q_end) {\n return;\n }\n\n std::array u2;\n ValueType q_val = *q++;\n float v2 = static_cast(q_val * q_val);\n for (size_t i = 0; i < M; ++i) {\n u2[i] = static_cast(p[i] * p[i]);\n out[i] = MathHelper::SquaredDifference(p[i], q_val);\n }\n p += M;\n\n while (q != q_end) {\n q_val = *q++;\n v2 += static_cast(q_val * q_val);\n for (size_t i = 0; i < M; ++i) {\n u2[i] += static_cast(p[i] * p[i]);\n out[i] += MathHelper::SquaredDifference(p[i], q_val);\n }\n p += M;\n }\n\n // Compute the injections\n for (size_t i = 0; i < M; ++i) {\n out[i] *= e2;\n u2[i] *= e2;\n }\n v2 *= e2;\n for (size_t k = 0; k < m; ++k) {\n for (size_t i = 0; i < M; ++i) {\n const float u_val = u2[i];\n u2[i] = u_val * u_val;\n out[i] += (u_val - v2) * (u_val - v2);\n }\n v2 = v2 * v2;\n }\n }\n};\n\n/*! Mips Squared Euclidean Distance Matrix (INT8, M >=2, N >= 2)\n */\ntemplate \nstruct MipsSquaredEuclideanDistanceMatrix<\n int8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n // Compute the distance between matrix and query by SphericalInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out) {\n ailego_assert(p && q && dim && !(dim & 3) && out);\n dim >>= 2;\n if (dim == 0) {\n return;\n }\n\n std::array u2;\n std::array v2;\n const uint32_t *p_it = reinterpret_cast(p);\n const uint32_t *q_it = reinterpret_cast(q);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = p_it[i];\n u2[i] = Squared(p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r = FusedMultiplyAdd(p_val, q_it[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] = Squared(q_it[i]);\n }\n p_it += M;\n q_it += N;\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = p_it[i];\n u2[i] += Squared(p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r += FusedMultiplyAdd(p_val, q_it[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] += Squared(q_it[i]);\n }\n p_it += M;\n q_it += N;\n }\n\n // Compute the injection\n for (size_t i = 0; i < M; ++i) {\n float *r = out + i;\n const float u2_val = u2[i];\n for (size_t j = 0; j < N; ++j) {\n *r = ComputeSphericalInjection(*r, u2_val, v2[j], e2);\n r += M;\n }\n }\n }\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out) {\n ailego_assert(p && q && dim && !(dim & 3) && out);\n dim >>= 2;\n if (dim == 0) {\n return;\n }\n\n std::array u2;\n std::array v2;\n const uint32_t *p_it = reinterpret_cast(p);\n const uint32_t *q_it = reinterpret_cast(q);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = p_it[i];\n u2[i] = Squared(p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r = SquaredDifference(p_val, q_it[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] = Squared(q_it[i]);\n }\n p_it += M;\n q_it += N;\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = p_it[i];\n u2[i] += Squared(p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r += SquaredDifference(p_val, q_it[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] += Squared(q_it[i]);\n }\n p_it += M;\n q_it += N;\n }\n\n // Compute the injections\n float *r = out;\n for (size_t i = 0; i < M; ++i) {\n u2[i] *= e2;\n for (size_t j = 0; j < N; ++j) {\n (*r++) *= e2;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] *= e2;\n }\n for (size_t k = 0; k < m; ++k) {\n for (size_t i = 0; i < M; ++i) {\n r = out + i;\n float u2_val = u2[i];\n u2[i] = u2_val * u2_val;\n for (size_t j = 0; j < N; ++j) {\n *r += (u2_val - v2[j]) * (u2_val - v2[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] = v2[i] * v2[i];\n }\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n return static_cast((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) +\n (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) +\n (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) +\n (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24));\n }\n\n //! Calculate the squared difference\n static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) {\n return static_cast(MathHelper::SquaredDifference(\n (int8_t)(lhs >> 0), (int8_t)(rhs >> 0)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 8), (int8_t)(rhs >> 8)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 16), (int8_t)(rhs >> 16)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 24), (int8_t)(rhs >> 24)));\n }\n\n //! Calculate sum of squared values\n static inline float Squared(uint32_t v) {\n return static_cast((int8_t)(v >> 0) * (int8_t)(v >> 0) +\n (int8_t)(v >> 8) * (int8_t)(v >> 8) +\n (int8_t)(v >> 16) * (int8_t)(v >> 16) +\n (int8_t)(v >> 24) * (int8_t)(v >> 24));\n }\n};\n\n/*! Mips Squared Euclidean Distance Matrix (INT8, N=1)\n */\ntemplate \nstruct MipsSquaredEuclideanDistanceMatrix<\n int8_t, M, 1, typename std::enable_if= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n // Compute the distance between matrix and query by SphericalInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out) {\n ailego_assert(p && q && dim && !(dim & 3) && out);\n const uint32_t *p_it = reinterpret_cast(p);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 2);\n if (q_it == q_end) {\n return;\n }\n\n std::array u2;\n uint32_t q_val = *q_it++;\n float v2 = Squared(q_val);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = *p_it++;\n u2[i] = Squared(p_val);\n out[i] = FusedMultiplyAdd(p_val, q_val);\n }\n\n while (q_it != q_end) {\n q_val = *q_it++;\n v2 += Squared(q_val);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = *p_it++;\n u2[i] += Squared(p_val);\n out[i] += FusedMultiplyAdd(p_val, q_val);\n }\n }\n\n // Compute the injection\n for (size_t i = 0; i < M; ++i) {\n out[i] = ComputeSphericalInjection(out[i], u2[i], v2, e2);\n }\n }\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out) {\n ailego_assert(p && q && dim && !(dim & 3) && out);\n const uint32_t *p_it = reinterpret_cast(p);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 2);\n if (q_it == q_end) {\n return;\n }\n\n std::array u2;\n uint32_t q_val = *q_it++;\n float v2 = Squared(q_val);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = *p_it++;\n u2[i] = Squared(p_val);\n out[i] = SquaredDifference(p_val, q_val);\n }\n\n while (q_it != q_end) {\n q_val = *q_it++;\n v2 += Squared(q_val);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = *p_it++;\n u2[i] += Squared(p_val);\n out[i] += SquaredDifference(p_val, q_val);\n }\n }\n\n // Compute the injections\n for (size_t i = 0; i < M; ++i) {\n out[i] *= e2;\n u2[i] *= e2;\n }\n v2 *= e2;\n for (size_t k = 0; k < m; ++k) {\n for (size_t i = 0; i < M; ++i) {\n const float u_val = u2[i];\n u2[i] = u_val * u_val;\n out[i] += (u_val - v2) * (u_val - v2);\n }\n v2 = v2 * v2;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n return static_cast((int8_t)(lhs >> 0) * (int8_t)(rhs >> 0) +\n (int8_t)(lhs >> 8) * (int8_t)(rhs >> 8) +\n (int8_t)(lhs >> 16) * (int8_t)(rhs >> 16) +\n (int8_t)(lhs >> 24) * (int8_t)(rhs >> 24));\n }\n\n //! Calculate the squared difference\n static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) {\n return static_cast(MathHelper::SquaredDifference(\n (int8_t)(lhs >> 0), (int8_t)(rhs >> 0)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 8), (int8_t)(rhs >> 8)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 16), (int8_t)(rhs >> 16)) +\n MathHelper::SquaredDifference(\n (int8_t)(lhs >> 24), (int8_t)(rhs >> 24)));\n }\n\n //! Calculate sum of squared values\n static inline float Squared(uint32_t v) {\n return static_cast((int8_t)(v >> 0) * (int8_t)(v >> 0) +\n (int8_t)(v >> 8) * (int8_t)(v >> 8) +\n (int8_t)(v >> 16) * (int8_t)(v >> 16) +\n (int8_t)(v >> 24) * (int8_t)(v >> 24));\n }\n};\n\n/*! Mips Squared Euclidean Distance Matrix (INT4, M >=2, N >= 2)\n */\ntemplate \nstruct MipsSquaredEuclideanDistanceMatrix<\n uint8_t, M, N, typename std::enable_if= 2 && N >= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n // Compute the distance between matrix and query by SphericalInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out) {\n ailego_assert(p && q && dim && !(dim & 7) && out);\n dim >>= 3;\n if (dim == 0) {\n return;\n }\n\n std::array u2;\n std::array v2;\n const uint32_t *p_it = reinterpret_cast(p);\n const uint32_t *q_it = reinterpret_cast(q);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = p_it[i];\n u2[i] = Squared(p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r = FusedMultiplyAdd(p_val, q_it[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] = Squared(q_it[i]);\n }\n p_it += M;\n q_it += N;\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = p_it[i];\n u2[i] += Squared(p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r += FusedMultiplyAdd(p_val, q_it[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] += Squared(q_it[i]);\n }\n p_it += M;\n q_it += N;\n }\n\n // Compute the injection\n for (size_t i = 0; i < M; ++i) {\n float *r = out + i;\n const float u2_val = u2[i];\n for (size_t j = 0; j < N; ++j) {\n *r = ComputeSphericalInjection(*r, u2_val, v2[j], e2);\n r += M;\n }\n }\n }\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out) {\n ailego_assert(p && q && dim && !(dim & 7) && out);\n dim >>= 3;\n if (dim == 0) {\n return;\n }\n\n std::array u2;\n std::array v2;\n const uint32_t *p_it = reinterpret_cast(p);\n const uint32_t *q_it = reinterpret_cast(q);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = p_it[i];\n u2[i] = Squared(p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r = SquaredDifference(p_val, q_it[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] = Squared(q_it[i]);\n }\n p_it += M;\n q_it += N;\n\n for (size_t k = 1; k < dim; ++k) {\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = p_it[i];\n u2[i] += Squared(p_val);\n float *r = out + i;\n for (size_t j = 0; j < N; ++j) {\n *r += SquaredDifference(p_val, q_it[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] += Squared(q_it[i]);\n }\n p_it += M;\n q_it += N;\n }\n\n // Compute the injections\n float *r = out;\n for (size_t i = 0; i < M; ++i) {\n u2[i] *= e2;\n for (size_t j = 0; j < N; ++j) {\n (*r++) *= e2;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] *= e2;\n }\n for (size_t k = 0; k < m; ++k) {\n for (size_t i = 0; i < M; ++i) {\n r = out + i;\n float u2_val = u2[i];\n u2[i] = u2_val * u2_val;\n for (size_t j = 0; j < N; ++j) {\n *r += (u2_val - v2[j]) * (u2_val - v2[j]);\n r += M;\n }\n }\n for (size_t i = 0; i < N; ++i) {\n v2[i] = v2[i] * v2[i];\n }\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n\n //! Calculate the squared difference\n static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4SquaredDiffTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n\n //! Calculate sum of squared values\n static inline float Squared(uint32_t u) {\n float sum = 0.0f;\n for (size_t i = 0; i < 32; i += 8) {\n uint8_t v = (uint8_t)(u >> i);\n int8_t lo = (int8_t)(v << 4) >> 4;\n int8_t hi = (int8_t)(v & 0xf0) >> 4;\n sum += hi * hi + lo * lo;\n }\n return sum;\n }\n};\n\n/*! Mips Squared Euclidean Distance Matrix (INT4, N=1)\n */\ntemplate \nstruct MipsSquaredEuclideanDistanceMatrix<\n uint8_t, M, 1, typename std::enable_if= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n // Compute the distance between matrix and query by SphericalInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n float e2, float *out) {\n ailego_assert(p && q && dim && !(dim & 7) && out);\n const uint32_t *p_it = reinterpret_cast(p);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 3);\n if (q_it == q_end) {\n return;\n }\n\n std::array u2;\n uint32_t q_val = *q_it++;\n float v2 = Squared(q_val);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = *p_it++;\n u2[i] = Squared(p_val);\n out[i] = FusedMultiplyAdd(p_val, q_val);\n }\n\n while (q_it != q_end) {\n q_val = *q_it++;\n v2 += Squared(q_val);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = *p_it++;\n u2[i] += Squared(p_val);\n out[i] += FusedMultiplyAdd(p_val, q_val);\n }\n }\n\n // Compute the injection\n for (size_t i = 0; i < M; ++i) {\n out[i] = ComputeSphericalInjection(out[i], u2[i], v2, e2);\n }\n }\n\n // Compute the distance between matrix and query by RepeatedQuadraticInjection\n static inline void Compute(const ValueType *p, const ValueType *q, size_t dim,\n size_t m, float e2, float *out) {\n ailego_assert(p && q && dim && !(dim & 7) && out);\n const uint32_t *p_it = reinterpret_cast(p);\n const uint32_t *q_it = reinterpret_cast(q);\n const uint32_t *q_end = q_it + (dim >> 3);\n if (q_it == q_end) {\n return;\n }\n\n std::array u2;\n uint32_t q_val = *q_it++;\n float v2 = Squared(q_val);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = *p_it++;\n u2[i] = Squared(p_val);\n out[i] = SquaredDifference(p_val, q_val);\n }\n\n while (q_it != q_end) {\n q_val = *q_it++;\n v2 += Squared(q_val);\n for (size_t i = 0; i < M; ++i) {\n const uint32_t p_val = *p_it++;\n u2[i] += Squared(p_val);\n out[i] += SquaredDifference(p_val, q_val);\n }\n }\n\n // Compute the injections\n for (size_t i = 0; i < M; ++i) {\n out[i] *= e2;\n u2[i] *= e2;\n }\n v2 *= e2;\n for (size_t k = 0; k < m; ++k) {\n for (size_t i = 0; i < M; ++i) {\n const float u_val = u2[i];\n u2[i] = u_val * u_val;\n out[i] += (u_val - v2) * (u_val - v2);\n }\n v2 = v2 * v2;\n }\n }\n\n protected:\n //! Calculate Fused-Multiply-Add\n static inline float FusedMultiplyAdd(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4MulTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4MulTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4MulTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4MulTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4MulTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4MulTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4MulTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4MulTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n\n //! Calculate the squared difference\n static inline float SquaredDifference(uint32_t lhs, uint32_t rhs) {\n return static_cast(\n Int4SquaredDiffTable[((lhs << 4) & 0xf0) | ((rhs >> 0) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 0) & 0xf0) | ((rhs >> 4) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 4) & 0xf0) | ((rhs >> 8) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 8) & 0xf0) | ((rhs >> 12) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 12) & 0xf0) | ((rhs >> 16) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 16) & 0xf0) | ((rhs >> 20) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 20) & 0xf0) | ((rhs >> 24) & 0xf)] +\n Int4SquaredDiffTable[((lhs >> 24) & 0xf0) | ((rhs >> 28) & 0xf)]);\n }\n\n //! Calculate sum of squared values\n static inline float Squared(uint32_t u) {\n float sum = 0.0f;\n for (size_t i = 0; i < 32; i += 8) {\n uint8_t v = (uint8_t)(u >> i);\n int8_t lo = (int8_t)(v << 4) >> 4;\n int8_t hi = (int8_t)(v & 0xf0) >> 4;\n sum += hi * hi + lo * lo;\n }\n return sum;\n }\n};\n\n//--------------------------------------------------\n// Sparse\n//--------------------------------------------------\n/*! Mips Squared Euclidean Sparse Distance Matrix\n */\ntemplate \nstruct MipsSquaredEuclideanSparseDistanceMatrix {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n static float ComputeInnerProductSparseInSegment(\n uint32_t m_sparse_count, const uint16_t *m_sparse_index,\n const ValueType *m_sparse_value, uint32_t q_sparse_count,\n const uint16_t *q_sparse_index, const ValueType *q_sparse_value);\n\n // Compute the distance between matrix and query by SphericalInjection\n static inline void Compute(const void *m_sparse_data_in,\n const void *q_sparse_data_in, float *out) {\n ailego_assert(m_sparse_data_in && q_sparse_data_in && out);\n\n const uint8_t *m_sparse_data =\n reinterpret_cast(m_sparse_data_in);\n const uint8_t *q_sparse_data =\n reinterpret_cast(q_sparse_data_in);\n\n const uint32_t m_sparse_count =\n *reinterpret_cast(m_sparse_data);\n const uint32_t q_sparse_count =\n *reinterpret_cast(q_sparse_data);\n\n if (m_sparse_count == 0 && q_sparse_count == 0) {\n *out = 0;\n return;\n }\n\n if (m_sparse_count == 0 || q_sparse_count == 0) {\n *out = 2;\n return;\n }\n\n const uint32_t m_seg_count =\n *reinterpret_cast(m_sparse_data + sizeof(uint32_t));\n const uint32_t q_seg_count =\n *reinterpret_cast(q_sparse_data + sizeof(uint32_t));\n\n const uint32_t *m_seg_id = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t));\n const uint32_t *q_seg_id = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t));\n\n const uint32_t *m_seg_vec_cnt = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t) + m_seg_count * sizeof(uint32_t));\n const uint32_t *q_seg_vec_cnt = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t) + q_seg_count * sizeof(uint32_t));\n\n const uint16_t *m_sparse_index = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t) +\n m_seg_count * 2 * sizeof(uint32_t));\n const uint16_t *q_sparse_index = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t) +\n q_seg_count * 2 * sizeof(uint32_t));\n\n const ValueType *m_sparse_value = reinterpret_cast(\n m_sparse_data + 2 * sizeof(uint32_t) +\n m_seg_count * 2 * sizeof(uint32_t) + m_sparse_count * sizeof(uint16_t));\n const ValueType *q_sparse_value = reinterpret_cast(\n q_sparse_data + 2 * sizeof(uint32_t) +\n q_seg_count * 2 * sizeof(uint32_t) + q_sparse_count * sizeof(uint16_t));\n\n float ip = 0.0f;\n\n size_t m_s = 0;\n size_t q_s = 0;\n\n size_t m_count = 0;\n size_t q_count = 0;\n\n while (m_s < m_seg_count && q_s < q_seg_count) {\n if (m_seg_id[m_s] == q_seg_id[q_s]) {\n ip += ComputeInnerProductSparseInSegment(\n m_seg_vec_cnt[m_s], m_sparse_index + m_count,\n m_sparse_value + m_count, q_seg_vec_cnt[q_s],\n q_sparse_index + q_count, q_sparse_value + q_count);\n\n m_count += m_seg_vec_cnt[m_s];\n q_count += q_seg_vec_cnt[q_s];\n\n ++m_s;\n ++q_s;\n } else if (m_seg_id[m_s] < q_seg_id[q_s]) {\n m_count += m_seg_vec_cnt[m_s];\n\n ++m_s;\n } else {\n q_count += q_seg_vec_cnt[q_s];\n\n ++q_s;\n }\n }\n\n float l2_m{0.0f};\n SquaredNorm2Matrix::Compute(m_sparse_value, m_sparse_count,\n &l2_m);\n\n float l2_q{0.0f};\n SquaredNorm2Matrix::Compute(q_sparse_value, q_sparse_count,\n &l2_q);\n\n *out = ComputeSphericalInjection(ip, l2_m, l2_q, 0.0f);\n }\n};\n\ntemplate \nfloat MipsSquaredEuclideanSparseDistanceMatrix<\n T>::ComputeInnerProductSparseInSegment(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const ValueType *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const ValueType *q_sparse_value) {\n float sum = 0.0f;\n\n size_t m_i = 0;\n size_t q_i = 0;\n while (m_i < m_sparse_count && q_i < q_sparse_count) {\n if (m_sparse_index[m_i] == q_sparse_index[q_i]) {\n sum += m_sparse_value[m_i] * q_sparse_value[q_i];\n\n ++m_i;\n ++q_i;\n } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) {\n ++m_i;\n } else {\n ++q_i;\n }\n }\n\n return sum;\n}\n\ntemplate <>\nfloat MipsSquaredEuclideanSparseDistanceMatrix<\n float>::ComputeInnerProductSparseInSegment(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const ValueType *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const ValueType *q_sparse_value);\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX__) && defined(__F16C__)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormFp16AVX(const Float16 *lhs, const Float16 *rhs,\n size_t size, float *sql, float *sqr) {\n __m256 ymm_sum_0 = _mm256_setzero_ps();\n __m256 ymm_sum_1 = _mm256_setzero_ps();\n __m256 ymm_sum_norm1 = _mm256_setzero_ps();\n __m256 ymm_sum_norm2 = _mm256_setzero_ps();\n\n const Float16 *last = lhs + size;\n const Float16 *last_aligned = lhs + ((size >> 4) << 4);\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m256i ymm_lhs = _mm256_load_si256((const __m256i *)lhs);\n __m256i ymm_rhs = _mm256_load_si256((const __m256i *)rhs);\n __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_lhs));\n __m256 ymm_lhs_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_lhs, 1));\n __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_rhs));\n __m256 ymm_rhs_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_rhs, 1));\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2);\n }\n if (last >= last_aligned + 8) {\n __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)lhs));\n __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)rhs));\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n lhs += 8;\n rhs += 8;\n }\n } else {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)lhs);\n __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)rhs);\n __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_lhs));\n __m256 ymm_lhs_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_lhs, 1));\n __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_rhs));\n __m256 ymm_rhs_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_rhs, 1));\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2);\n }\n if (last >= last_aligned + 8) {\n __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)lhs));\n __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)rhs));\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n lhs += 8;\n rhs += 8;\n }\n }\n\n float result = HorizontalAdd_FP32_V256(_mm256_add_ps(ymm_sum_0, ymm_sum_1));\n float norm1 = HorizontalAdd_FP32_V256(ymm_sum_norm1);\n float norm2 = HorizontalAdd_FP32_V256(ymm_sum_norm2);\n switch (last - lhs) {\n case 7:\n FMA_FP16_GENERAL(lhs[6], rhs[6], result, norm1, norm2);\n /* FALLTHRU */\n case 6:\n FMA_FP16_GENERAL(lhs[5], rhs[5], result, norm1, norm2);\n /* FALLTHRU */\n case 5:\n FMA_FP16_GENERAL(lhs[4], rhs[4], result, norm1, norm2);\n /* FALLTHRU */\n case 4:\n FMA_FP16_GENERAL(lhs[3], rhs[3], result, norm1, norm2);\n /* FALLTHRU */\n case 3:\n FMA_FP16_GENERAL(lhs[2], rhs[2], result, norm1, norm2);\n /* FALLTHRU */\n case 2:\n FMA_FP16_GENERAL(lhs[1], rhs[1], result, norm1, norm2);\n /* FALLTHRU */\n case 1:\n FMA_FP16_GENERAL(lhs[0], rhs[0], result, norm1, norm2);\n }\n\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionFp16AVX(const Float16 *lhs,\n const Float16 *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormFp16AVX(lhs, rhs, size, &u2, &v2);\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX(\n const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormFp16AVX(lhs, rhs, size, &u2, &v2);\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n\n#endif // __AVX__ && __F16C__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX512F__)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormFp16AVX512(const Float16 *lhs,\n const Float16 *rhs, size_t size,\n float *sql, float *sqr) {\n __m512 zmm_sum_0 = _mm512_setzero_ps();\n __m512 zmm_sum_1 = _mm512_setzero_ps();\n __m512 zmm_sum_norm1 = _mm512_setzero_ps();\n __m512 zmm_sum_norm2 = _mm512_setzero_ps();\n\n const Float16 *last = lhs + size;\n const Float16 *last_aligned = lhs + ((size >> 5) << 5);\n if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m512i zmm_lhs = _mm512_load_si512((const __m512i *)lhs);\n __m512i zmm_rhs = _mm512_load_si512((const __m512i *)rhs);\n __m512 zmm_lhs_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_lhs));\n __m512 zmm_lhs_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_lhs, 1));\n __m512 zmm_rhs_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_rhs));\n __m512 zmm_rhs_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_rhs, 1));\n FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0)\n FMA_FP32_AVX512(zmm_lhs_1, zmm_rhs_1, zmm_sum_1)\n FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_lhs_1, zmm_lhs_1, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2)\n FMA_FP32_AVX512(zmm_rhs_1, zmm_rhs_1, zmm_sum_norm2)\n }\n if (last >= last_aligned + 16) {\n __m512 zmm_lhs_0 =\n _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)lhs));\n __m512 zmm_rhs_0 =\n _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)rhs));\n FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0)\n FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2)\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m512i zmm_lhs = _mm512_loadu_si512((const __m512i *)lhs);\n __m512i zmm_rhs = _mm512_loadu_si512((const __m512i *)rhs);\n __m512 zmm_lhs_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_lhs));\n __m512 zmm_lhs_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_lhs, 1));\n __m512 zmm_rhs_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_rhs));\n __m512 zmm_rhs_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_rhs, 1));\n FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0)\n FMA_FP32_AVX512(zmm_lhs_1, zmm_rhs_1, zmm_sum_1)\n FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_lhs_1, zmm_lhs_1, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2)\n FMA_FP32_AVX512(zmm_rhs_1, zmm_rhs_1, zmm_sum_norm2)\n }\n if (last >= last_aligned + 16) {\n __m512 zmm_lhs_0 =\n _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)lhs));\n __m512 zmm_rhs_0 =\n _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)rhs));\n FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0)\n FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2)\n lhs += 16;\n rhs += 16;\n }\n }\n\n __m256 ymm_sum_0 =\n HorizontalAdd_FP32_V512_TO_V256(_mm512_add_ps(zmm_sum_0, zmm_sum_1));\n __m256 ymm_sum_norm1 = HorizontalAdd_FP32_V512_TO_V256(zmm_sum_norm1);\n __m256 ymm_sum_norm2 = HorizontalAdd_FP32_V512_TO_V256(zmm_sum_norm2);\n if (last >= lhs + 8) {\n __m256 ymm_lhs_0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)lhs));\n __m256 ymm_rhs_0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)rhs));\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n lhs += 8;\n rhs += 8;\n }\n\n float result = HorizontalAdd_FP32_V256(ymm_sum_0);\n float norm1 = HorizontalAdd_FP32_V256(ymm_sum_norm1);\n float norm2 = HorizontalAdd_FP32_V256(ymm_sum_norm2);\n switch (last - lhs) {\n case 7:\n FMA_FP16_GENERAL(lhs[6], rhs[6], result, norm1, norm2);\n /* FALLTHRU */\n case 6:\n FMA_FP16_GENERAL(lhs[5], rhs[5], result, norm1, norm2);\n /* FALLTHRU */\n case 5:\n FMA_FP16_GENERAL(lhs[4], rhs[4], result, norm1, norm2);\n /* FALLTHRU */\n case 4:\n FMA_FP16_GENERAL(lhs[3], rhs[3], result, norm1, norm2);\n /* FALLTHRU */\n case 3:\n FMA_FP16_GENERAL(lhs[2], rhs[2], result, norm1, norm2);\n /* FALLTHRU */\n case 2:\n FMA_FP16_GENERAL(lhs[1], rhs[1], result, norm1, norm2);\n /* FALLTHRU */\n case 1:\n FMA_FP16_GENERAL(lhs[0], rhs[0], result, norm1, norm2);\n }\n\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionFp16AVX512(const Float16 *lhs,\n const Float16 *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormFp16AVX512(lhs, rhs, size, &u2, &v2);\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX512(\n const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormFp16AVX512(lhs, rhs, size, &u2, &v2);\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n\n#endif // __AVX512F__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__ARM_NEON)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp16NEON(\n const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionFp16NEON(const Float16 *lhs,\n const Float16 *rhs,\n size_t size, float e2);\n#endif\n\n#if defined(__AVX512F__)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX512(\n const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionFp16AVX512(const Float16 *lhs,\n const Float16 *rhs,\n size_t size, float e2);\n#endif\n\n#if defined(__AVX__)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX(\n const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionFp16AVX(const Float16 *lhs,\n const Float16 *rhs,\n size_t size, float e2);\n#endif\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp16Scalar(\n const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionFp16Scalar(\n const ailego::Float16 *p, const ailego::Float16 *q, size_t dim, float e2);\n\n\n//! Compute the distance between matrix and query by SphericalInjection\nvoid MipsSquaredEuclideanDistanceMatrix::Compute(\n const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {\n#if defined(__ARM_NEON)\n *out = MipsEuclideanDistanceSphericalInjectionFp16NEON(p, q, dim, e2);\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = MipsEuclideanDistanceSphericalInjectionFp16AVX512(p, q, dim, e2);\n return;\n }\n#endif\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = MipsEuclideanDistanceSphericalInjectionFp16AVX(p, q, dim, e2);\n return;\n }\n#endif //__AVX__\n *out = MipsEuclideanDistanceSphericalInjectionFp16Scalar(p, q, dim, e2);\n return;\n#endif //__ARM_NEON\n}\n\n//! Compute the distance between matrix and query by RepeatedQuadraticInjection\nvoid MipsSquaredEuclideanDistanceMatrix::Compute(\n const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,\n float *out) {\n#if defined(__ARM_NEON)\n *out =\n MipsEuclideanDistanceRepeatedQuadraticInjectionFp16NEON(p, q, dim, m, e2);\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX512(p, q, dim,\n m, e2);\n return;\n }\n#endif\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp16AVX(p, q, dim, m,\n e2);\n return;\n }\n#endif //__AVX__\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp16Scalar(p, q, dim, m,\n e2);\n return;\n#endif //__ARM_NEON\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp16.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__ARM_NEON) && defined(__aarch64__)\n#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormFp16NEON(const Float16 *lhs, const Float16 *rhs,\n size_t size, float *sql, float *sqr) {\n const Float16 *last = lhs + size;\n const Float16 *last_aligned = lhs + ((size >> 3) << 3);\n float16x8_t v_sum = vdupq_n_f16(0);\n float16x8_t v_sum_norm1 = vdupq_n_f16(0);\n float16x8_t v_sum_norm2 = vdupq_n_f16(0);\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n float16x8_t v_lhs = vld1q_f16((const float16_t *)lhs);\n float16x8_t v_rhs = vld1q_f16((const float16_t *)rhs);\n v_sum = vfmaq_f16(v_sum, v_lhs, v_rhs);\n v_sum_norm1 = vfmaq_f16(v_sum_norm1, v_lhs, v_lhs);\n v_sum_norm2 = vfmaq_f16(v_sum_norm2, v_rhs, v_rhs);\n }\n if (last >= last_aligned + 4) {\n float16x8_t v_lhs = vcombine_f16(vld1_f16((const float16_t *)lhs),\n vreinterpret_f16_u64(vdup_n_u64(0ul)));\n float16x8_t v_rhs = vcombine_f16(vld1_f16((const float16_t *)rhs),\n vreinterpret_f16_u64(vdup_n_u64(0ul)));\n v_sum = vfmaq_f16(v_sum, v_lhs, v_rhs);\n v_sum_norm1 = vfmaq_f16(v_sum_norm1, v_lhs, v_lhs);\n v_sum_norm2 = vfmaq_f16(v_sum_norm2, v_rhs, v_rhs);\n lhs += 4;\n rhs += 4;\n }\n\n float result = HorizontalAdd_FP16_NEON(v_sum);\n float norm1 = HorizontalAdd_FP16_NEON(v_sum_norm1);\n float norm2 = HorizontalAdd_FP16_NEON(v_sum_norm2);\n\n switch (last - lhs) {\n case 3:\n FMA_FP16_GENERAL(lhs[2], rhs[2], result, norm1, norm2);\n /* FALLTHRU */\n case 2:\n FMA_FP16_GENERAL(lhs[1], rhs[1], result, norm1, norm2);\n /* FALLTHRU */\n case 1:\n FMA_FP16_GENERAL(lhs[0], rhs[0], result, norm1, norm2);\n }\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n#else\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormFp16NEON(const Float16 *lhs, const Float16 *rhs,\n size_t size, float *sql, float *sqr) {\n const Float16 *last = lhs + size;\n const Float16 *last_aligned = lhs + ((size >> 3) << 3);\n float32x4_t v_sum_0 = vdupq_n_f32(0);\n float32x4_t v_sum_1 = vdupq_n_f32(0);\n float32x4_t v_sum_norm1 = vdupq_n_f32(0);\n float32x4_t v_sum_norm2 = vdupq_n_f32(0);\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n float16x8_t v_lhs = vld1q_f16((const float16_t *)lhs);\n float16x8_t v_rhs = vld1q_f16((const float16_t *)rhs);\n float32x4_t v_lhs_0 = vcvt_f32_f16(vget_low_f16(v_lhs));\n float32x4_t v_rhs_0 = vcvt_f32_f16(vget_low_f16(v_rhs));\n float32x4_t v_lhs_1 = vcvt_high_f32_f16(v_lhs);\n float32x4_t v_rhs_1 = vcvt_high_f32_f16(v_rhs);\n v_sum_0 = vfmaq_f32(v_sum_0, v_lhs_0, v_rhs_0);\n v_sum_1 = vfmaq_f32(v_sum_1, v_lhs_1, v_rhs_1);\n v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_0, v_lhs_0);\n v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_1, v_lhs_1);\n v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_0, v_rhs_0);\n v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_1, v_rhs_1);\n }\n if (last >= last_aligned + 4) {\n float32x4_t v_lhs_0 = vcvt_f32_f16(vld1_f16((const float16_t *)lhs));\n float32x4_t v_rhs_0 = vcvt_f32_f16(vld1_f16((const float16_t *)rhs));\n v_sum_0 = vfmaq_f32(v_sum_0, v_lhs_0, v_rhs_0);\n v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_0, v_lhs_0);\n v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_0, v_rhs_0);\n lhs += 4;\n rhs += 4;\n }\n\n float result = vaddvq_f32(vaddq_f32(v_sum_0, v_sum_1));\n float norm1 = vaddvq_f32(v_sum_norm1);\n float norm2 = vaddvq_f32(v_sum_norm2);\n switch (last - lhs) {\n case 3:\n FMA_FP16_GENERAL(lhs[2], rhs[2], result, norm1, norm2);\n /* FALLTHRU */\n case 2:\n FMA_FP16_GENERAL(lhs[1], rhs[1], result, norm1, norm2);\n /* FALLTHRU */\n case 1:\n FMA_FP16_GENERAL(lhs[0], rhs[0], result, norm1, norm2);\n }\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\n#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC\n\nfloat MipsEuclideanDistanceSphericalInjectionFp16NEON(const Float16 *lhs,\n const Float16 *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormFp16NEON(lhs, rhs, size, &u2, &v2);\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp16NEON(\n const Float16 *lhs, const Float16 *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormFp16NEON(lhs, rhs, size, &u2, &v2);\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n#endif // __ARM_NEON && __aarch64__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__SSE__)\nfloat InnerProductAndSquaredNormFp32SSE(const float *lhs, const float *rhs,\n size_t size, float *sql, float *sqr);\n#endif\n\n#if defined(__AVX__)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormFp32AVX(const float *lhs, const float *rhs,\n size_t size, float *sql, float *sqr) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 4) << 4);\n\n __m256 ymm_sum_0 = _mm256_setzero_ps();\n __m256 ymm_sum_1 = _mm256_setzero_ps();\n __m256 ymm_sum_norm1 = _mm256_setzero_ps();\n __m256 ymm_sum_norm2 = _mm256_setzero_ps();\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m256 ymm_lhs_0 = _mm256_load_ps(lhs + 0);\n __m256 ymm_lhs_1 = _mm256_load_ps(lhs + 8);\n __m256 ymm_rhs_0 = _mm256_load_ps(rhs + 0);\n __m256 ymm_rhs_1 = _mm256_load_ps(rhs + 8);\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2);\n }\n\n if (last >= last_aligned + 8) {\n __m256 ymm_lhs_0 = _mm256_load_ps(lhs);\n __m256 ymm_rhs_0 = _mm256_load_ps(rhs);\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n lhs += 8;\n rhs += 8;\n }\n } else {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m256 ymm_lhs_0 = _mm256_loadu_ps(lhs + 0);\n __m256 ymm_lhs_1 = _mm256_loadu_ps(lhs + 8);\n __m256 ymm_rhs_0 = _mm256_loadu_ps(rhs + 0);\n __m256 ymm_rhs_1 = _mm256_loadu_ps(rhs + 8);\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_rhs_1, ymm_sum_1);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2);\n }\n\n if (last >= last_aligned + 8) {\n __m256 ymm_lhs_0 = _mm256_loadu_ps(lhs);\n __m256 ymm_rhs_0 = _mm256_loadu_ps(rhs);\n ymm_sum_0 = _mm256_fmadd_ps(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n ymm_sum_norm1 = _mm256_fmadd_ps(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_fmadd_ps(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n lhs += 8;\n rhs += 8;\n }\n }\n float result = HorizontalAdd_FP32_V256(_mm256_add_ps(ymm_sum_0, ymm_sum_1));\n float norm1 = HorizontalAdd_FP32_V256(ymm_sum_norm1);\n float norm2 = HorizontalAdd_FP32_V256(ymm_sum_norm2);\n\n switch (last - lhs) {\n case 7:\n FMA_FP32_GENERAL(lhs[6], rhs[6], result, norm1, norm2)\n /* FALLTHRU */\n case 6:\n FMA_FP32_GENERAL(lhs[5], rhs[5], result, norm1, norm2)\n /* FALLTHRU */\n case 5:\n FMA_FP32_GENERAL(lhs[4], rhs[4], result, norm1, norm2)\n /* FALLTHRU */\n case 4:\n FMA_FP32_GENERAL(lhs[3], rhs[3], result, norm1, norm2)\n /* FALLTHRU */\n case 3:\n FMA_FP32_GENERAL(lhs[2], rhs[2], result, norm1, norm2)\n /* FALLTHRU */\n case 2:\n FMA_FP32_GENERAL(lhs[1], rhs[1], result, norm1, norm2)\n /* FALLTHRU */\n case 1:\n FMA_FP32_GENERAL(lhs[0], rhs[0], result, norm1, norm2)\n }\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionFp32AVX(const float *lhs,\n const float *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n if (size > 7) {\n sum = InnerProductAndSquaredNormFp32AVX(lhs, rhs, size, &u2, &v2);\n } else {\n sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2);\n }\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX(\n const float *lhs, const float *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n if (size > 7) {\n sum = InnerProductAndSquaredNormFp32AVX(lhs, rhs, size, &u2, &v2);\n } else {\n sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2);\n }\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n#endif // __AVX__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__SSE__)\nfloat InnerProductAndSquaredNormFp32SSE(const float *lhs, const float *rhs,\n size_t size, float *sql, float *sqr);\n#endif\n\n#if defined(__AVX__)\nfloat InnerProductAndSquaredNormFp32AVX(const float *lhs, const float *rhs,\n size_t size, float *sql, float *sqr);\n#endif\n\n#if defined(__AVX512F__)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormFp32AVX512(const float *lhs, const float *rhs,\n size_t size, float *sql,\n float *sqr) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 5) << 5);\n\n __m512 zmm_sum_0 = _mm512_setzero_ps();\n __m512 zmm_sum_1 = _mm512_setzero_ps();\n __m512 zmm_sum_norm1 = _mm512_setzero_ps();\n __m512 zmm_sum_norm2 = _mm512_setzero_ps();\n\n if (((uintptr_t)lhs & 0x3f) == 0 && ((uintptr_t)rhs & 0x3f) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m512 zmm_lhs_0 = _mm512_load_ps(lhs + 0);\n __m512 zmm_lhs_1 = _mm512_load_ps(lhs + 16);\n __m512 zmm_rhs_0 = _mm512_load_ps(rhs + 0);\n __m512 zmm_rhs_1 = _mm512_load_ps(rhs + 16);\n FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0)\n FMA_FP32_AVX512(zmm_lhs_1, zmm_rhs_1, zmm_sum_1)\n FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_lhs_1, zmm_lhs_1, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2)\n FMA_FP32_AVX512(zmm_rhs_1, zmm_rhs_1, zmm_sum_norm2)\n }\n\n if (last >= last_aligned + 16) {\n __m512 zmm_lhs_0 = _mm512_load_ps(lhs);\n __m512 zmm_rhs_0 = _mm512_load_ps(rhs);\n FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0)\n FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2)\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m512 zmm_lhs_0 = _mm512_loadu_ps(lhs + 0);\n __m512 zmm_lhs_1 = _mm512_loadu_ps(lhs + 16);\n __m512 zmm_rhs_0 = _mm512_loadu_ps(rhs + 0);\n __m512 zmm_rhs_1 = _mm512_loadu_ps(rhs + 16);\n FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0)\n FMA_FP32_AVX512(zmm_lhs_1, zmm_rhs_1, zmm_sum_1)\n FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_lhs_1, zmm_lhs_1, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2)\n FMA_FP32_AVX512(zmm_rhs_1, zmm_rhs_1, zmm_sum_norm2)\n }\n\n if (last >= last_aligned + 16) {\n __m512 zmm_lhs_0 = _mm512_loadu_ps(lhs);\n __m512 zmm_rhs_0 = _mm512_loadu_ps(rhs);\n FMA_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0)\n FMA_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1)\n FMA_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2)\n lhs += 16;\n rhs += 16;\n }\n }\n\n zmm_sum_0 = _mm512_add_ps(zmm_sum_0, zmm_sum_1);\n if (lhs != last) {\n __mmask16 mask = (__mmask16)((1 << (last - lhs)) - 1);\n __m512 zmm_undefined = _mm512_undefined_ps();\n __m512 zmm_lhs_0 = _mm512_mask_loadu_ps(zmm_undefined, mask, lhs);\n __m512 zmm_rhs_0 = _mm512_mask_loadu_ps(zmm_undefined, mask, rhs);\n FMA_MASK_FP32_AVX512(zmm_lhs_0, zmm_rhs_0, zmm_sum_0, mask);\n FMA_MASK_FP32_AVX512(zmm_lhs_0, zmm_lhs_0, zmm_sum_norm1, mask);\n FMA_MASK_FP32_AVX512(zmm_rhs_0, zmm_rhs_0, zmm_sum_norm2, mask);\n }\n\n *sql = HorizontalAdd_FP32_V512(zmm_sum_norm1);\n *sqr = HorizontalAdd_FP32_V512(zmm_sum_norm2);\n return HorizontalAdd_FP32_V512(zmm_sum_0);\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionFp32AVX512(const float *lhs,\n const float *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n if (size > 15) {\n sum = InnerProductAndSquaredNormFp32AVX512(lhs, rhs, size, &u2, &v2);\n } else if (size > 7) {\n sum = InnerProductAndSquaredNormFp32AVX(lhs, rhs, size, &u2, &v2);\n } else {\n sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2);\n }\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX512(\n const float *lhs, const float *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n if (size > 15) {\n sum = InnerProductAndSquaredNormFp32AVX512(lhs, rhs, size, &u2, &v2);\n } else if (size > 7) {\n sum = InnerProductAndSquaredNormFp32AVX(lhs, rhs, size, &u2, &v2);\n } else {\n sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2);\n }\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n#endif // __AVX512F__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__ARM_NEON)\nfloat InnerProductAndSquaredNormFp32NEON(const float *lhs, const float *rhs,\n size_t size, float *sql, float *sqr);\n#endif\n\n#if defined(__AVX512F__)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX512(\n const float *lhs, const float *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionFp32AVX512(const float *lhs,\n const float *rhs,\n size_t size, float e2);\n#endif\n\n#if defined(__AVX__)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX(\n const float *lhs, const float *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionFp32AVX(const float *lhs,\n const float *rhs,\n size_t size, float e2);\n#endif\n\n#if defined(__SSE__)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp32SSE(\n const float *lhs, const float *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionFp32SSE(const float *lhs,\n const float *rhs,\n size_t size, float e2);\n#endif\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp32Scalar(\n const float *p, const float *q, size_t dim, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionFp32Scalar(const float *p,\n const float *q,\n size_t dim, float e2);\n\nfloat MipsInnerProductSparseInSegment(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value);\n\n//! Compute the distance between matrix and query by SphericalInjection\nvoid MipsSquaredEuclideanDistanceMatrix::Compute(\n const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {\n#if __ARM_NEON\n float u2{0.0f};\n float v2{0.0f};\n float sum = InnerProductAndSquaredNormFp32NEON(p, q, dim, &u2, &v2);\n\n *out = ComputeSphericalInjection(sum, u2, v2, e2);\n return;\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = MipsEuclideanDistanceSphericalInjectionFp32AVX512(p, q, dim, e2);\n return;\n }\n#endif //__AVX512F__\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = MipsEuclideanDistanceSphericalInjectionFp32AVX(p, q, dim, e2);\n return;\n }\n#endif // __AVX__\n#if defined(__SSE__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) {\n *out = MipsEuclideanDistanceSphericalInjectionFp32SSE(p, q, dim, e2);\n return;\n }\n#endif // __SSE__\n *out = MipsEuclideanDistanceSphericalInjectionFp32Scalar(p, q, dim, e2);\n return;\n#endif //__ARM_NEON\n}\n\n//! Compute the distance between matrix and query by RepeatedQuadraticInjection\nvoid MipsSquaredEuclideanDistanceMatrix::Compute(\n const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,\n float *out) {\n#if defined(__ARM_NEON)\n float u2{0.0f};\n float v2{0.0f};\n float sum = InnerProductAndSquaredNormFp32NEON(p, q, dim, &u2, &v2);\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n *out = sum;\n return;\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX512(p, q, dim,\n m, e2);\n return;\n }\n#endif //__AVX512F__\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp32AVX(p, q, dim, m,\n e2);\n return;\n }\n#endif // __AVX__\n\n#if defined(__SSE__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) {\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp32SSE(p, q, dim, m,\n e2);\n return;\n }\n#endif //__SSE__\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionFp32Scalar(p, q, dim, m,\n e2);\n\n return;\n#endif //__ARM_NEON\n}\n\n// Sparse\n#if defined(__SSE4_1__)\nfloat MipsInnerProductSparseInSegmentSSE(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value);\n#endif\n\ntemplate <>\nfloat MipsSquaredEuclideanSparseDistanceMatrix::\n ComputeInnerProductSparseInSegment(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const ValueType *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const ValueType *q_sparse_value) {\n#if defined(__SSE4_1__)\n return MipsInnerProductSparseInSegmentSSE(m_sparse_count, m_sparse_index,\n m_sparse_value, q_sparse_count,\n q_sparse_index, q_sparse_value);\n#else\n return MipsInnerProductSparseInSegment(m_sparse_count, m_sparse_index,\n m_sparse_value, q_sparse_count,\n q_sparse_index, q_sparse_value);\n#endif\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__ARM_NEON)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormFp32NEON(const float *lhs, const float *rhs,\n size_t size, float *sql, float *sqr) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 3) << 3);\n\n float32x4_t v_sum_0 = vdupq_n_f32(0);\n float32x4_t v_sum_1 = vdupq_n_f32(0);\n float32x4_t v_sum_norm1 = vdupq_n_f32(0);\n float32x4_t v_sum_norm2 = vdupq_n_f32(0);\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n float32x4_t v_lhs_0 = vld1q_f32(lhs + 0);\n float32x4_t v_lhs_1 = vld1q_f32(lhs + 4);\n float32x4_t v_rhs_0 = vld1q_f32(rhs + 0);\n float32x4_t v_rhs_1 = vld1q_f32(rhs + 4);\n v_sum_0 = vfmaq_f32(v_sum_0, v_lhs_0, v_rhs_0);\n v_sum_1 = vfmaq_f32(v_sum_1, v_lhs_1, v_rhs_1);\n v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_0, v_lhs_0);\n v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_1, v_lhs_1);\n v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_0, v_rhs_0);\n v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_1, v_rhs_1);\n }\n if (last >= last_aligned + 4) {\n float32x4_t v_lhs_0 = vld1q_f32(lhs);\n float32x4_t v_rhs_0 = vld1q_f32(rhs);\n v_sum_0 = vfmaq_f32(v_sum_0, v_lhs_0, v_rhs_0);\n v_sum_norm1 = vfmaq_f32(v_sum_norm1, v_lhs_0, v_lhs_0);\n v_sum_norm2 = vfmaq_f32(v_sum_norm2, v_rhs_0, v_rhs_0);\n lhs += 4;\n rhs += 4;\n }\n\n float result = vaddvq_f32(vaddq_f32(v_sum_0, v_sum_1));\n float norm1 = vaddvq_f32(v_sum_norm1);\n float norm2 = vaddvq_f32(v_sum_norm2);\n switch (last - lhs) {\n case 3:\n FMA_FP32_GENERAL(lhs[2], rhs[2], result, norm1, norm2)\n /* FALLTHRU */\n case 2:\n FMA_FP32_GENERAL(lhs[1], rhs[1], result, norm1, norm2)\n /* FALLTHRU */\n case 1:\n FMA_FP32_GENERAL(lhs[0], rhs[0], result, norm1, norm2)\n }\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\n#endif //__ARM_NEON\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_fp32.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__SSE__)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormFp32SSE(const float *lhs, const float *rhs,\n size_t size, float *sql, float *sqr) {\n const float *last = lhs + size;\n const float *last_aligned = lhs + ((size >> 3) << 3);\n\n __m128 xmm_sum = _mm_setzero_ps();\n __m128 xmm_sum_norm1 = _mm_setzero_ps();\n __m128 xmm_sum_norm2 = _mm_setzero_ps();\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m128 xmm_lhs_0 = _mm_load_ps(lhs + 0);\n __m128 xmm_lhs_1 = _mm_load_ps(lhs + 4);\n __m128 xmm_rhs_0 = _mm_load_ps(rhs + 0);\n __m128 xmm_rhs_1 = _mm_load_ps(rhs + 4);\n xmm_sum = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum);\n xmm_sum = _mm_fmadd_ps(xmm_lhs_1, xmm_rhs_1, xmm_sum);\n xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1);\n xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1);\n xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2);\n xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2);\n }\n\n if (last >= last_aligned + 4) {\n __m128 xmm_lhs_0 = _mm_load_ps(lhs);\n __m128 xmm_rhs_0 = _mm_load_ps(rhs);\n xmm_sum = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum);\n xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1);\n xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2);\n lhs += 4;\n rhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m128 xmm_lhs_0 = _mm_loadu_ps(lhs + 0);\n __m128 xmm_lhs_1 = _mm_loadu_ps(lhs + 4);\n __m128 xmm_rhs_0 = _mm_loadu_ps(rhs + 0);\n __m128 xmm_rhs_1 = _mm_loadu_ps(rhs + 4);\n xmm_sum = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum);\n xmm_sum = _mm_fmadd_ps(xmm_lhs_1, xmm_rhs_1, xmm_sum);\n xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1);\n xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1);\n xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2);\n xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2);\n }\n\n if (last >= last_aligned + 4) {\n __m128 xmm_lhs_0 = _mm_loadu_ps(lhs);\n __m128 xmm_rhs_0 = _mm_loadu_ps(rhs);\n xmm_sum = _mm_fmadd_ps(xmm_lhs_0, xmm_rhs_0, xmm_sum);\n xmm_sum_norm1 = _mm_fmadd_ps(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1);\n xmm_sum_norm2 = _mm_fmadd_ps(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2);\n lhs += 4;\n rhs += 4;\n }\n }\n float result = HorizontalAdd_FP32_V128(xmm_sum);\n float norm1 = HorizontalAdd_FP32_V128(xmm_sum_norm1);\n float norm2 = HorizontalAdd_FP32_V128(xmm_sum_norm2);\n\n switch (last - lhs) {\n case 3:\n FMA_FP32_GENERAL(lhs[2], rhs[2], result, norm1, norm2)\n /* FALLTHRU */\n case 2:\n FMA_FP32_GENERAL(lhs[1], rhs[1], result, norm1, norm2)\n /* FALLTHRU */\n case 1:\n FMA_FP32_GENERAL(lhs[0], rhs[0], result, norm1, norm2)\n }\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionFp32SSE(const float *lhs,\n const float *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2);\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp32SSE(\n const float *lhs, const float *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormFp32SSE(lhs, rhs, size, &u2, &v2);\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n\n#endif // __SSE__\n\n// #if 1\n#if defined(__SSE4_1__)\nconst static __m128i SHUFFLE_MASK16[16] = {\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, -127, -127, -127, -127),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 7, 6, 5, 4, 3,\n 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 11, 10, 9, 8,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, -127, -127,\n -127, -127, 15, 14, 13, 12),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 3, 2, 1, 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 7, 6, 5, 4),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 7, 6, 5, 4, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, -127, -127, -127, -127, 15, 14, 13, 12,\n 11, 10, 9, 8),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 3, 2, 1,\n 0),\n _mm_set_epi8(-127, -127, -127, -127, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5,\n 4),\n _mm_set_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0),\n};\n\nconstexpr uint32_t MAX_SPARSE_BUFFER_LENGTH = 65536;\n\nfloat MipsInnerProductSparseInSegmentSSE(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value) {\n float sum = 0.0f;\n\n // size_t alloc_size = 0;\n\n size_t i1 = 0, i2 = 0;\n size_t end1 = m_sparse_count / 8 * 8;\n size_t end2 = q_sparse_count / 8 * 8;\n\n // std::vector mem1;\n // std::vector mem2;\n\n float fixed_buffer_1[MAX_SPARSE_BUFFER_LENGTH];\n float fixed_buffer_2[MAX_SPARSE_BUFFER_LENGTH];\n\n float *val_start_1 = fixed_buffer_1;\n float *val_start_2 = fixed_buffer_2;\n\n // uint32_t max_count = std::max(m_sparse_count, q_sparse_count);\n\n // if (MAX_SPARSE_BUFFER_LENGTH < max_count) {\n // mem1.reserve(max_count);\n // mem2.reserve(max_count);\n\n // val_start_1 = mem1.data();\n // val_start_2 = mem2.data();\n // }\n\n float *val_1 = val_start_1;\n float *val_2 = val_start_2;\n\n if (i1 < end1 && i2 < end2) {\n while (m_sparse_index[i1 + 7] < q_sparse_index[i2]) {\n i1 += 8;\n if (i1 >= end1) goto do_scalar;\n }\n\n while (q_sparse_index[i2 + 7] < m_sparse_index[i1]) {\n i2 += 8;\n if (i2 >= end2) goto do_scalar;\n }\n\n __m128i mm_index_m =\n _mm_loadu_si128(reinterpret_cast(&m_sparse_index[i1]));\n __m128i mm_index_q =\n _mm_loadu_si128(reinterpret_cast(&q_sparse_index[i2]));\n\n while (true) {\n#ifdef DEBUG_PRINT\n std::cout << \"index 1: \" << std::endl;\n print_data16(&mm_index_m);\n\n std::cout << \"index 2: \" << std::endl;\n print_data16(&mm_index_q);\n#endif\n\n __m128i mm_cmp_res =\n _mm_cmpistrm(mm_index_q, mm_index_m,\n _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK);\n\n#ifdef DEBUG_PRINT\n std::cout << \"cmp res: \" << std::endl;\n print_data16(&mm_cmp_res);\n#endif\n\n int r = _mm_extract_epi32(mm_cmp_res, 0);\n\n if (r) {\n int r1 = r & 15;\n\n __m128i v = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_value[i1]));\n __m128 vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r1]));\n\n _mm_storeu_ps(val_1, vs);\n val_1 += _mm_popcnt_u32(r1);\n\n int r2 = (r >> 4) & 15;\n v = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_value[i1 + 4]));\n vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r2]));\n _mm_storeu_ps(val_1, vs);\n val_1 += _mm_popcnt_u32(r2);\n\n mm_cmp_res = _mm_cmpistrm(\n mm_index_m, mm_index_q,\n _SIDD_UWORD_OPS | _SIDD_CMP_EQUAL_ANY | _SIDD_BIT_MASK);\n r = _mm_extract_epi32(mm_cmp_res, 0);\n\n r1 = r & 15;\n\n v = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_value[i2]));\n vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r1]));\n _mm_storeu_ps(val_2, vs);\n val_2 += _mm_popcnt_u32(r1);\n\n r2 = (r >> 4) & 15;\n v = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_value[i2 + 4]));\n vs = _mm_castsi128_ps(_mm_shuffle_epi8(v, SHUFFLE_MASK16[r2]));\n _mm_storeu_ps(val_2, vs);\n val_2 += _mm_popcnt_u32(r2);\n }\n\n const uint16_t id1_max = m_sparse_index[i1 + 7];\n\n if (id1_max <= q_sparse_index[i2 + 7]) {\n i1 += 8;\n if (i1 >= end1) goto do_scalar;\n mm_index_m = _mm_loadu_si128(\n reinterpret_cast(&m_sparse_index[i1]));\n }\n\n if (id1_max >= q_sparse_index[i2 + 7]) {\n i2 += 8;\n if (i2 >= end2) goto do_scalar;\n mm_index_q = _mm_loadu_si128(\n reinterpret_cast(&q_sparse_index[i2]));\n }\n }\n }\n\ndo_scalar:\n while (i1 < m_sparse_count && i2 < q_sparse_count) {\n if (m_sparse_index[i1] == q_sparse_index[i2]) {\n *val_1++ = m_sparse_value[i1];\n *val_2++ = q_sparse_value[i2];\n\n ++i1;\n ++i2;\n } else if (m_sparse_index[i1] < q_sparse_index[i2]) {\n ++i1;\n } else {\n ++i2;\n }\n }\n\n size_t res_num = val_1 - val_start_1;\n\n // if (res_num != val_2 - val_start_2) {\n // std::cerr << \"size mismatch!\" << std::endl;\n // }\n\n size_t res_num4 = res_num / 4 * 4;\n\n if (res_num4) {\n __m128 sum128 = _mm_set1_ps(0);\n\n for (size_t k = 0; k < res_num4; k += 4) {\n sum128 = _mm_add_ps(sum128, _mm_mul_ps(_mm_loadu_ps(val_start_1 + k),\n _mm_loadu_ps(val_start_2 + k)));\n }\n\n float __attribute__((aligned(16))) tmp_res[4];\n _mm_store_ps(tmp_res, sum128);\n sum += (tmp_res[0] + tmp_res[1] + tmp_res[2] + tmp_res[3]);\n }\n\n for (size_t k = res_num4; k < res_num; ++k)\n sum += val_start_1[k] * val_start_2[k];\n\n return sum;\n}\n#else\nfloat MipsInnerProductSparseInSegment(uint32_t m_sparse_count,\n const uint16_t *m_sparse_index,\n const float *m_sparse_value,\n uint32_t q_sparse_count,\n const uint16_t *q_sparse_index,\n const float *q_sparse_value) {\n float sum = 0.0f;\n\n size_t m_i = 0;\n size_t q_i = 0;\n while (m_i < m_sparse_count && q_i < q_sparse_count) {\n if (m_sparse_index[m_i] == q_sparse_index[q_i]) {\n sum += m_sparse_value[m_i] * q_sparse_value[q_i];\n\n ++m_i;\n ++q_i;\n } else if (m_sparse_index[m_i] < q_sparse_index[q_i]) {\n ++m_i;\n } else {\n ++q_i;\n }\n }\n\n return sum;\n}\n#endif // __SSE4_1__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int8.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"inner_product_matrix.h\"\n#include \"mips_euclidean_distance_matrix.h\"\n#include \"norm_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX2__)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormInt4AVX2(const uint8_t *lhs, const uint8_t *rhs,\n size_t size, float *sql, float *sqr) {\n const uint8_t *last = lhs + size;\n const uint8_t *last_aligned = lhs + ((size >> 5) << 5);\n __m256i ymm_sum_0 = _mm256_setzero_si256();\n __m256i ymm_sum_1 = _mm256_setzero_si256();\n __m256i ymm_sum_norm1 = _mm256_setzero_si256();\n __m256i ymm_sum_norm2 = _mm256_setzero_si256();\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m256i ymm_lhs = _mm256_load_si256((const __m256i *)(lhs));\n __m256i ymm_rhs = _mm256_load_si256((const __m256i *)(rhs));\n FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum_0, ymm_sum1, ymm_sum_norm1,\n ymm_sum_norm2)\n }\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs);\n __m128i xmm_sum = _mm_setzero_si128();\n __m128i xmm_sum_norm1 = _mm_setzero_si128();\n __m128i xmm_sum_norm2 = _mm_setzero_si128();\n FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum, xmm_sum_norm1, xmm_sum_norm2)\n ymm_sum_0 = _mm256_add_epi32(\n _mm256_set_m128i(_mm_setzero_si128(), xmm_sum), ymm_sum_0);\n ymm_sum_norm1 = _mm256_add_epi32(\n _mm256_set_m128i(_mm_setzero_si128(), xmm_sum_norm1), ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_add_epi32(\n _mm256_set_m128i(_mm_setzero_si128(), xmm_sum_norm2), ymm_sum_norm2);\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)(lhs));\n __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)(rhs));\n FMA_INT4_ITER_AVX(ymm_lhs, ymm_rhs, ymm_sum_0, ymm_sum1, ymm_sum_norm1,\n ymm_sum_norm2)\n }\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs);\n __m128i xmm_sum = _mm_setzero_si128();\n __m128i xmm_sum_norm1 = _mm_setzero_si128();\n __m128i xmm_sum_norm2 = _mm_setzero_si128();\n FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum, xmm_sum_norm1, xmm_sum_norm2)\n ymm_sum_0 = _mm256_add_epi32(\n _mm256_set_m128i(_mm_setzero_si128(), xmm_sum), ymm_sum_0);\n ymm_sum_norm1 = _mm256_add_epi32(\n _mm256_set_m128i(_mm_setzero_si128(), xmm_sum_norm1), ymm_sum_norm1);\n ymm_sum_norm2 = _mm256_add_epi32(\n _mm256_set_m128i(_mm_setzero_si128(), xmm_sum_norm2), ymm_sum_norm2);\n lhs += 16;\n rhs += 16;\n }\n }\n float result = static_cast(\n HorizontalAdd_INT32_V256(_mm256_add_epi32(ymm_sum_0, ymm_sum_1)));\n float norm1 = static_cast(HorizontalAdd_INT32_V256(ymm_sum_norm1));\n float norm2 = static_cast(HorizontalAdd_INT32_V256(ymm_sum_norm2));\n\n switch (last - lhs) {\n case 15:\n FMA_INT4_GENERAL(lhs[14], rhs[14], result, norm1, norm2)\n /* FALLTHRU */\n case 14:\n FMA_INT4_GENERAL(lhs[13], rhs[13], result, norm1, norm2)\n /* FALLTHRU */\n case 13:\n FMA_INT4_GENERAL(lhs[12], rhs[12], result, norm1, norm2)\n /* FALLTHRU */\n case 12:\n FMA_INT4_GENERAL(lhs[11], rhs[11], result, norm1, norm2)\n /* FALLTHRU */\n case 11:\n FMA_INT4_GENERAL(lhs[10], rhs[10], result, norm1, norm2)\n /* FALLTHRU */\n case 10:\n FMA_INT4_GENERAL(lhs[9], rhs[9], result, norm1, norm2)\n /* FALLTHRU */\n case 9:\n FMA_INT4_GENERAL(lhs[8], rhs[8], result, norm1, norm2)\n /* FALLTHRU */\n case 8:\n FMA_INT4_GENERAL(lhs[7], rhs[7], result, norm1, norm2)\n /* FALLTHRU */\n case 7:\n FMA_INT4_GENERAL(lhs[6], rhs[6], result, norm1, norm2)\n /* FALLTHRU */\n case 6:\n FMA_INT4_GENERAL(lhs[5], rhs[5], result, norm1, norm2)\n /* FALLTHRU */\n case 5:\n FMA_INT4_GENERAL(lhs[4], rhs[4], result, norm1, norm2)\n /* FALLTHRU */\n case 4:\n FMA_INT4_GENERAL(lhs[3], rhs[3], result, norm1, norm2)\n /* FALLTHRU */\n case 3:\n FMA_INT4_GENERAL(lhs[2], rhs[2], result, norm1, norm2)\n /* FALLTHRU */\n case 2:\n FMA_INT4_GENERAL(lhs[1], rhs[1], result, norm1, norm2)\n /* FALLTHRU */\n case 1:\n FMA_INT4_GENERAL(lhs[0], rhs[0], result, norm1, norm2)\n }\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionInt4AVX2(const uint8_t *lhs,\n const uint8_t *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormInt4AVX2(lhs, rhs, size >> 1, &u2, &v2);\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt4AVX2(\n const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormInt4AVX2(lhs, rhs, size >> 1, &u2, &v2);\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n#endif // __AVX2__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"inner_product_matrix.h\"\n#include \"mips_euclidean_distance_matrix.h\"\n#include \"norm_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX2__)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt4AVX2(\n const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionInt4AVX2(const uint8_t *lhs,\n const uint8_t *rhs,\n size_t size, float e2);\n#endif\n\n#if defined(__SSE4_1__)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt4SSE(\n const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionInt4SSE(const uint8_t *lhs,\n const uint8_t *rhs,\n size_t size, float e2);\n#endif\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt4Scalar(\n const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionInt4Scalar(const uint8_t *lhs,\n const uint8_t *rhs,\n size_t size, float e2);\n\n//! Compute the distance between matrix and query by SphericalInjection\nvoid MipsSquaredEuclideanDistanceMatrix::Compute(\n const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = MipsEuclideanDistanceSphericalInjectionInt4AVX2(p, q, dim, e2);\n return;\n }\n#endif\n\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = MipsEuclideanDistanceSphericalInjectionInt4SSE(p, q, dim, e2);\n return;\n }\n#endif\n\n *out = MipsEuclideanDistanceSphericalInjectionInt4Scalar(p, q, dim, e2);\n}\n\n//! Compute the distance between matrix and query by RepeatedQuadraticInjection\nvoid MipsSquaredEuclideanDistanceMatrix::Compute(\n const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,\n float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt4AVX2(p, q, dim, m,\n e2);\n return;\n }\n#endif\n\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt4SSE(p, q, dim, m,\n e2);\n return;\n }\n#endif\n\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt4Scalar(p, q, dim, m,\n e2);\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int8.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"inner_product_matrix.h\"\n#include \"mips_euclidean_distance_matrix.h\"\n#include \"norm_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__SSE4_1__)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormInt4SSE(const uint8_t *lhs, const uint8_t *rhs,\n size_t size, float *sql, float *sqr) {\n const uint8_t *last = lhs + size;\n const uint8_t *last_aligned = lhs + ((size >> 4) << 4);\n __m128i xmm_sum = _mm_setzero_si128();\n __m128i xmm_sum_norm1 = _mm_setzero_si128();\n __m128i xmm_sum_norm2 = _mm_setzero_si128();\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)(lhs));\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)(rhs));\n FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum, xmm_sum_norm1, xmm_sum_norm2)\n }\n } else {\n for (; lhs != last_aligned; lhs += 16, rhs += 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)(lhs));\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)(rhs));\n FMA_INT4_ITER_SSE(xmm_lhs, xmm_rhs, xmm_sum, xmm_sum_norm1, xmm_sum_norm2)\n }\n }\n float result = static_cast(HorizontalAdd_INT32_V128(xmm_sum));\n float norm1 = static_cast(HorizontalAdd_INT32_V128(xmm_sum_norm1));\n float norm2 = static_cast(HorizontalAdd_INT32_V128(xmm_sum_norm2));\n\n switch (last - lhs) {\n case 15:\n FMA_INT4_GENERAL(lhs[14], rhs[14], result, norm1, norm2)\n /* FALLTHRU */\n case 14:\n FMA_INT4_GENERAL(lhs[13], rhs[13], result, norm1, norm2)\n /* FALLTHRU */\n case 13:\n FMA_INT4_GENERAL(lhs[12], rhs[12], result, norm1, norm2)\n /* FALLTHRU */\n case 12:\n FMA_INT4_GENERAL(lhs[11], rhs[11], result, norm1, norm2)\n /* FALLTHRU */\n case 11:\n FMA_INT4_GENERAL(lhs[10], rhs[10], result, norm1, norm2)\n /* FALLTHRU */\n case 10:\n FMA_INT4_GENERAL(lhs[9], rhs[9], result, norm1, norm2)\n /* FALLTHRU */\n case 9:\n FMA_INT4_GENERAL(lhs[8], rhs[8], result, norm1, norm2)\n /* FALLTHRU */\n case 8:\n FMA_INT4_GENERAL(lhs[7], rhs[7], result, norm1, norm2)\n /* FALLTHRU */\n case 7:\n FMA_INT4_GENERAL(lhs[6], rhs[6], result, norm1, norm2)\n /* FALLTHRU */\n case 6:\n FMA_INT4_GENERAL(lhs[5], rhs[5], result, norm1, norm2)\n /* FALLTHRU */\n case 5:\n FMA_INT4_GENERAL(lhs[4], rhs[4], result, norm1, norm2)\n /* FALLTHRU */\n case 4:\n FMA_INT4_GENERAL(lhs[3], rhs[3], result, norm1, norm2)\n /* FALLTHRU */\n case 3:\n FMA_INT4_GENERAL(lhs[2], rhs[2], result, norm1, norm2)\n /* FALLTHRU */\n case 2:\n FMA_INT4_GENERAL(lhs[1], rhs[1], result, norm1, norm2)\n /* FALLTHRU */\n case 1:\n FMA_INT4_GENERAL(lhs[0], rhs[0], result, norm1, norm2)\n }\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionInt4SSE(const uint8_t *lhs,\n const uint8_t *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormInt4SSE(lhs, rhs, size >> 1, &u2, &v2);\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt4SSE(\n const uint8_t *lhs, const uint8_t *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormInt4SSE(lhs, rhs, size >> 1, &u2, &v2);\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n\n#endif // __SSE4_1__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int8.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX2__)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormInt8AVX2(const int8_t *lhs, const int8_t *rhs,\n size_t size, float *sql, float *sqr) {\n const int8_t *last = lhs + size;\n const int8_t *last_aligned = lhs + ((size >> 6) << 6);\n\n __m256i ymm_sum_0 = _mm256_setzero_si256();\n __m256i ymm_sum_1 = _mm256_setzero_si256();\n __m256i ymm_sum_norm1 = _mm256_setzero_si256();\n __m256i ymm_sum_norm2 = _mm256_setzero_si256();\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n __m256i ymm_lhs_0 = _mm256_load_si256((const __m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_load_si256((const __m256i *)(lhs + 32));\n __m256i ymm_rhs_0 = _mm256_load_si256((const __m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_load_si256((const __m256i *)(rhs + 32));\n FMA_INT8_AVX(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n FMA_INT8_AVX(ymm_lhs_1, ymm_rhs_1, ymm_sum_1);\n FMA_INT8_AVX(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n FMA_INT8_AVX(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1);\n FMA_INT8_AVX(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n FMA_INT8_AVX(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2);\n }\n\n if (last >= last_aligned + 32) {\n __m256i ymm_lhs = _mm256_load_si256((const __m256i *)lhs);\n __m256i ymm_rhs = _mm256_load_si256((const __m256i *)rhs);\n FMA_INT8_AVX(ymm_lhs, ymm_rhs, ymm_sum_0);\n FMA_INT8_AVX(ymm_lhs, ymm_lhs, ymm_sum_norm1);\n FMA_INT8_AVX(ymm_rhs, ymm_rhs, ymm_sum_norm2);\n lhs += 32;\n rhs += 32;\n }\n\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs);\n FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_rhs, ymm_sum_0);\n FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_lhs, ymm_sum_norm1);\n FMA_INT8_AVX_SSE_HYBRID(xmm_rhs, xmm_rhs, ymm_sum_norm2);\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 64, rhs += 64) {\n __m256i ymm_lhs_0 = _mm256_loadu_si256((const __m256i *)(lhs + 0));\n __m256i ymm_lhs_1 = _mm256_loadu_si256((const __m256i *)(lhs + 32));\n __m256i ymm_rhs_0 = _mm256_loadu_si256((const __m256i *)(rhs + 0));\n __m256i ymm_rhs_1 = _mm256_loadu_si256((const __m256i *)(rhs + 32));\n FMA_INT8_AVX(ymm_lhs_0, ymm_rhs_0, ymm_sum_0);\n FMA_INT8_AVX(ymm_lhs_1, ymm_rhs_1, ymm_sum_1);\n FMA_INT8_AVX(ymm_lhs_0, ymm_lhs_0, ymm_sum_norm1);\n FMA_INT8_AVX(ymm_lhs_1, ymm_lhs_1, ymm_sum_norm1);\n FMA_INT8_AVX(ymm_rhs_0, ymm_rhs_0, ymm_sum_norm2);\n FMA_INT8_AVX(ymm_rhs_1, ymm_rhs_1, ymm_sum_norm2);\n }\n\n if (last >= last_aligned + 32) {\n __m256i ymm_lhs = _mm256_loadu_si256((const __m256i *)lhs);\n __m256i ymm_rhs = _mm256_loadu_si256((const __m256i *)rhs);\n FMA_INT8_AVX(ymm_lhs, ymm_rhs, ymm_sum_0);\n FMA_INT8_AVX(ymm_lhs, ymm_lhs, ymm_sum_norm1);\n FMA_INT8_AVX(ymm_rhs, ymm_rhs, ymm_sum_norm2);\n lhs += 32;\n rhs += 32;\n }\n\n if (last >= lhs + 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs);\n FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_rhs, ymm_sum_0);\n FMA_INT8_AVX_SSE_HYBRID(xmm_lhs, xmm_lhs, ymm_sum_norm1);\n FMA_INT8_AVX_SSE_HYBRID(xmm_rhs, xmm_rhs, ymm_sum_norm2);\n lhs += 16;\n rhs += 16;\n }\n }\n float result = static_cast(\n HorizontalAdd_INT32_V256(_mm256_add_epi32(ymm_sum_0, ymm_sum_1)));\n float norm1 = static_cast(HorizontalAdd_INT32_V256(ymm_sum_norm1));\n float norm2 = static_cast(HorizontalAdd_INT32_V256(ymm_sum_norm2));\n\n switch (last - lhs) {\n case 15:\n FMA_INT8_GENERAL(lhs[14], rhs[14], result, norm1, norm2)\n /* FALLTHRU */\n case 14:\n FMA_INT8_GENERAL(lhs[13], rhs[13], result, norm1, norm2)\n /* FALLTHRU */\n case 13:\n FMA_INT8_GENERAL(lhs[12], rhs[12], result, norm1, norm2)\n /* FALLTHRU */\n case 12:\n FMA_INT8_GENERAL(lhs[11], rhs[11], result, norm1, norm2)\n /* FALLTHRU */\n case 11:\n FMA_INT8_GENERAL(lhs[10], rhs[10], result, norm1, norm2)\n /* FALLTHRU */\n case 10:\n FMA_INT8_GENERAL(lhs[9], rhs[9], result, norm1, norm2)\n /* FALLTHRU */\n case 9:\n FMA_INT8_GENERAL(lhs[8], rhs[8], result, norm1, norm2)\n /* FALLTHRU */\n case 8:\n FMA_INT8_GENERAL(lhs[7], rhs[7], result, norm1, norm2)\n /* FALLTHRU */\n case 7:\n FMA_INT8_GENERAL(lhs[6], rhs[6], result, norm1, norm2)\n /* FALLTHRU */\n case 6:\n FMA_INT8_GENERAL(lhs[5], rhs[5], result, norm1, norm2)\n /* FALLTHRU */\n case 5:\n FMA_INT8_GENERAL(lhs[4], rhs[4], result, norm1, norm2)\n /* FALLTHRU */\n case 4:\n FMA_INT8_GENERAL(lhs[3], rhs[3], result, norm1, norm2)\n /* FALLTHRU */\n case 3:\n FMA_INT8_GENERAL(lhs[2], rhs[2], result, norm1, norm2)\n /* FALLTHRU */\n case 2:\n FMA_INT8_GENERAL(lhs[1], rhs[1], result, norm1, norm2)\n /* FALLTHRU */\n case 1:\n FMA_INT8_GENERAL(lhs[0], rhs[0], result, norm1, norm2)\n }\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionInt8AVX2(const int8_t *lhs,\n const int8_t *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormInt8AVX2(lhs, rhs, size, &u2, &v2);\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt8AVX2(\n const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormInt8AVX2(lhs, rhs, size, &u2, &v2);\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n#endif // __AVX2__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__AVX2__)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt8AVX2(\n const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionInt8AVX2(const int8_t *lhs,\n const int8_t *rhs,\n size_t size, float e2);\n#endif\n\n#if defined(__SSE4_1__)\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt8SSE(\n const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionInt8SSE(const int8_t *lhs,\n const int8_t *rhs,\n size_t size, float e2);\n#endif\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt8Scalar(\n const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2);\nfloat MipsEuclideanDistanceSphericalInjectionInt8Scalar(const int8_t *lhs,\n const int8_t *rhs,\n size_t size, float e2);\n\n//! Compute the distance between matrix and query by SphericalInjection\nvoid MipsSquaredEuclideanDistanceMatrix::Compute(\n const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = MipsEuclideanDistanceSphericalInjectionInt8AVX2(p, q, dim, e2);\n return;\n }\n#endif\n\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = MipsEuclideanDistanceSphericalInjectionInt8SSE(p, q, dim, e2);\n return;\n }\n#endif //__SSE4_1__\n\n *out = MipsEuclideanDistanceSphericalInjectionInt8Scalar(p, q, dim, e2);\n}\n\n//! Compute the distance between matrix and query by RepeatedQuadraticInjection\nvoid MipsSquaredEuclideanDistanceMatrix::Compute(\n const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,\n float *out) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt8AVX2(p, q, dim, m,\n e2);\n return;\n }\n#endif\n#if defined(__SSE4_1__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE4_1) {\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt8SSE(p, q, dim, m,\n e2);\n return;\n }\n#endif //__SSE4_1__\n\n *out = MipsEuclideanDistanceRepeatedQuadraticInjectionInt8Scalar(p, q, dim, m,\n e2);\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"distance_matrix_accum_int8.i\"\n#include \"distance_matrix_mips_utility.i\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__SSE4_1__)\n//! Compute the Inner Product between p and q, and each Squared L2-Norm value\nfloat InnerProductAndSquaredNormInt8SSE(const int8_t *lhs, const int8_t *rhs,\n size_t size, float *sql, float *sqr) {\n const int8_t *last = lhs + size;\n const int8_t *last_aligned = lhs + ((size >> 5) << 5);\n\n __m128i xmm_sum = _mm_setzero_si128();\n __m128i xmm_sum_norm1 = _mm_setzero_si128();\n __m128i xmm_sum_norm2 = _mm_setzero_si128();\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m128i xmm_lhs_0 = _mm_load_si128((const __m128i *)(lhs + 0));\n __m128i xmm_lhs_1 = _mm_load_si128((const __m128i *)(lhs + 16));\n __m128i xmm_rhs_0 = _mm_load_si128((const __m128i *)(rhs + 0));\n __m128i xmm_rhs_1 = _mm_load_si128((const __m128i *)(rhs + 16));\n FMA_INT8_SSE(xmm_lhs_0, xmm_rhs_0, xmm_sum);\n FMA_INT8_SSE(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1);\n FMA_INT8_SSE(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2);\n FMA_INT8_SSE(xmm_lhs_1, xmm_rhs_1, xmm_sum);\n FMA_INT8_SSE(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1);\n FMA_INT8_SSE(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2);\n }\n\n if (last >= last_aligned + 16) {\n __m128i xmm_lhs = _mm_load_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_load_si128((const __m128i *)rhs);\n FMA_INT8_SSE(xmm_lhs, xmm_rhs, xmm_sum);\n FMA_INT8_SSE(xmm_lhs, xmm_lhs, xmm_sum_norm1);\n FMA_INT8_SSE(xmm_rhs, xmm_rhs, xmm_sum_norm2);\n lhs += 16;\n rhs += 16;\n }\n } else {\n for (; lhs != last_aligned; lhs += 32, rhs += 32) {\n __m128i xmm_lhs_0 = _mm_loadu_si128((const __m128i *)(lhs + 0));\n __m128i xmm_lhs_1 = _mm_loadu_si128((const __m128i *)(lhs + 16));\n __m128i xmm_rhs_0 = _mm_loadu_si128((const __m128i *)(rhs + 0));\n __m128i xmm_rhs_1 = _mm_loadu_si128((const __m128i *)(rhs + 16));\n FMA_INT8_SSE(xmm_lhs_0, xmm_rhs_0, xmm_sum);\n FMA_INT8_SSE(xmm_lhs_0, xmm_lhs_0, xmm_sum_norm1);\n FMA_INT8_SSE(xmm_rhs_0, xmm_rhs_0, xmm_sum_norm2);\n FMA_INT8_SSE(xmm_lhs_1, xmm_rhs_1, xmm_sum);\n FMA_INT8_SSE(xmm_lhs_1, xmm_lhs_1, xmm_sum_norm1);\n FMA_INT8_SSE(xmm_rhs_1, xmm_rhs_1, xmm_sum_norm2);\n }\n\n if (last >= last_aligned + 16) {\n __m128i xmm_lhs = _mm_loadu_si128((const __m128i *)lhs);\n __m128i xmm_rhs = _mm_loadu_si128((const __m128i *)rhs);\n FMA_INT8_SSE(xmm_lhs, xmm_rhs, xmm_sum);\n FMA_INT8_SSE(xmm_lhs, xmm_lhs, xmm_sum_norm1);\n FMA_INT8_SSE(xmm_rhs, xmm_rhs, xmm_sum_norm2);\n lhs += 16;\n rhs += 16;\n }\n }\n float result = static_cast(HorizontalAdd_INT32_V128(xmm_sum));\n float norm1 = static_cast(HorizontalAdd_INT32_V128(xmm_sum_norm1));\n float norm2 = static_cast(HorizontalAdd_INT32_V128(xmm_sum_norm2));\n\n switch (last - lhs) {\n case 15:\n FMA_INT8_GENERAL(lhs[14], rhs[14], result, norm1, norm2)\n /* FALLTHRU */\n case 14:\n FMA_INT8_GENERAL(lhs[13], rhs[13], result, norm1, norm2)\n /* FALLTHRU */\n case 13:\n FMA_INT8_GENERAL(lhs[12], rhs[12], result, norm1, norm2)\n /* FALLTHRU */\n case 12:\n FMA_INT8_GENERAL(lhs[11], rhs[11], result, norm1, norm2)\n /* FALLTHRU */\n case 11:\n FMA_INT8_GENERAL(lhs[10], rhs[10], result, norm1, norm2)\n /* FALLTHRU */\n case 10:\n FMA_INT8_GENERAL(lhs[9], rhs[9], result, norm1, norm2)\n /* FALLTHRU */\n case 9:\n FMA_INT8_GENERAL(lhs[8], rhs[8], result, norm1, norm2)\n /* FALLTHRU */\n case 8:\n FMA_INT8_GENERAL(lhs[7], rhs[7], result, norm1, norm2)\n /* FALLTHRU */\n case 7:\n FMA_INT8_GENERAL(lhs[6], rhs[6], result, norm1, norm2)\n /* FALLTHRU */\n case 6:\n FMA_INT8_GENERAL(lhs[5], rhs[5], result, norm1, norm2)\n /* FALLTHRU */\n case 5:\n FMA_INT8_GENERAL(lhs[4], rhs[4], result, norm1, norm2)\n /* FALLTHRU */\n case 4:\n FMA_INT8_GENERAL(lhs[3], rhs[3], result, norm1, norm2)\n /* FALLTHRU */\n case 3:\n FMA_INT8_GENERAL(lhs[2], rhs[2], result, norm1, norm2)\n /* FALLTHRU */\n case 2:\n FMA_INT8_GENERAL(lhs[1], rhs[1], result, norm1, norm2)\n /* FALLTHRU */\n case 1:\n FMA_INT8_GENERAL(lhs[0], rhs[0], result, norm1, norm2)\n }\n *sql = norm1;\n *sqr = norm2;\n return result;\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionInt8SSE(const int8_t *lhs,\n const int8_t *rhs,\n size_t size, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormInt8SSE(lhs, rhs, size, &u2, &v2);\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt8SSE(\n const int8_t *lhs, const int8_t *rhs, size_t size, size_t m, float e2) {\n float u2{0.0f};\n float v2{0.0f};\n float sum{0.0f};\n\n sum = InnerProductAndSquaredNormInt8SSE(lhs, rhs, size, &u2, &v2);\n\n sum = e2 * (u2 + v2 - 2 * sum);\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n\n#endif // __SSE4_1__\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/mips_euclidean_distance_matrix_scalar.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n#include \"distance_utility.h\"\n#include \"mips_euclidean_distance_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n//--------------------------------------------------\n// Dense\n//--------------------------------------------------\n// Compute the distance between matrix and query by SphericalInjection\ntemplate \ninline float MipsEuclideanDistanceSphericalInjectionScalar(const T *p,\n const T *q,\n size_t dim,\n float e2) {\n ailego_assert(p && q && dim);\n\n float sum = 0.0;\n float u2 = 0.0;\n float v2 = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n u2 += p[i] * p[i];\n v2 += q[i] * q[i];\n sum += static_cast(p[i] * q[i]);\n }\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\n// Compute the distance between matrix and query by RepeatedQuadraticInjection\ntemplate \ninline float MipsEuclideanDistanceRepeatedQuadraticInjectionScalar(\n const T *p, const T *q, size_t dim, size_t m, float e2) {\n ailego_assert(p && q && dim);\n\n float sum = 0.0;\n float u2 = 0.0;\n float v2 = 0.0;\n for (size_t i = 0; i < dim; ++i) {\n u2 += p[i] * p[i];\n v2 += q[i] * q[i];\n sum += MathHelper::SquaredDifference(p[i], q[i]);\n }\n\n sum *= e2;\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n\n/*! Mips Squared Euclidean Distance Matrix (INT4, M=1, N=1)\n */\n//! Calculate sum of squared values\nstatic inline float Squared(uint8_t v) {\n return static_cast(((int8_t)(v << 4) >> 4) * ((int8_t)(v << 4) >> 4) +\n ((int8_t)(v & 0xf0) >> 4) *\n ((int8_t)(v & 0xf0) >> 4));\n}\n\n// Compute the distance between matrix and query by SphericalInjection\nfloat MipsEuclideanDistanceSphericalInjectionInt4Scalar(const uint8_t *p,\n const uint8_t *q,\n size_t dim, float e2) {\n ailego_assert(p && q && dim && !(dim & 1));\n\n float sum = 0.0;\n float u2 = 0.0;\n float v2 = 0.0;\n for (size_t i = 0; i < (dim >> 1); ++i) {\n const uint8_t p_val = p[i];\n const uint8_t q_val = q[i];\n u2 += Squared(p_val);\n v2 += Squared(q_val);\n sum += Int4MulTable[((p_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] +\n Int4MulTable[((p_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)];\n }\n\n return ComputeSphericalInjection(sum, u2, v2, e2);\n}\n\n// Compute the distance between matrix and query by RepeatedQuadraticInjection\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt4Scalar(\n const uint8_t *p, const uint8_t *q, size_t dim, size_t m, float e2) {\n ailego_assert(p && q && dim && !(dim & 1));\n\n float sum = 0.0;\n float u2 = 0.0;\n float v2 = 0.0;\n for (size_t i = 0; i < (dim >> 1); ++i) {\n const uint8_t p_val = p[i];\n const uint8_t q_val = q[i];\n u2 += Squared(p_val);\n v2 += Squared(q_val);\n sum += Int4SquaredDiffTable[((p_val << 4) & 0xf0) | ((q_val >> 0) & 0xf)] +\n Int4SquaredDiffTable[((p_val >> 0) & 0xf0) | ((q_val >> 4) & 0xf)];\n }\n sum *= e2;\n u2 *= e2;\n v2 *= e2;\n for (size_t i = 0; i < m; ++i) {\n sum += (u2 - v2) * (u2 - v2);\n u2 = u2 * u2;\n v2 = v2 * v2;\n }\n\n return sum;\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionInt8Scalar(const int8_t *p,\n const int8_t *q,\n size_t dim, float e2) {\n return MipsEuclideanDistanceSphericalInjectionScalar(p, q, dim, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionInt8Scalar(\n const int8_t *p, const int8_t *q, size_t dim, size_t m, float e2) {\n return MipsEuclideanDistanceRepeatedQuadraticInjectionScalar(\n p, q, dim, m, e2);\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionFp16Scalar(\n const ailego::Float16 *p, const ailego::Float16 *q, size_t dim, float e2) {\n return MipsEuclideanDistanceSphericalInjectionScalar(\n p, q, dim, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp16Scalar(\n const ailego::Float16 *p, const ailego::Float16 *q, size_t dim, size_t m,\n float e2) {\n return MipsEuclideanDistanceRepeatedQuadraticInjectionScalar(\n p, q, dim, m, e2);\n}\n\nfloat MipsEuclideanDistanceSphericalInjectionFp32Scalar(const float *p,\n const float *q,\n size_t dim, float e2) {\n return MipsEuclideanDistanceSphericalInjectionScalar(p, q, dim, e2);\n}\n\nfloat MipsEuclideanDistanceRepeatedQuadraticInjectionFp32Scalar(\n const float *p, const float *q, size_t dim, size_t m, float e2) {\n return MipsEuclideanDistanceRepeatedQuadraticInjectionScalar(p, q, dim,\n m, e2);\n}\n\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/norm1_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! L1-Norm Matrix\n */\ntemplate \nstruct Norm1Matrix;\n\n/*! L1-Norm Matrix\n */\ntemplate \nstruct Norm1Matrix::value &&\n sizeof(T) >= 2 && M >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && out);\n\n const ValueType *m_end = m + dim * M;\n if (m != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = MathHelper::Absolute(m[i]);\n }\n m += M;\n }\n while (m != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += MathHelper::Absolute(m[i]);\n }\n m += M;\n }\n }\n};\n\n/*! L1-Norm Matrix (INT8)\n */\ntemplate \nstruct Norm1Matrix= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && !(dim & 3) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *m_end = m_it + (dim >> 2) * M;\n\n if (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = Absolute(m_it[i]);\n }\n m_it += M;\n }\n while (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += Absolute(m_it[i]);\n }\n m_it += M;\n }\n }\n\n protected:\n //! Calculate sum of absolute values\n static inline float Absolute(uint32_t v) {\n return static_cast(\n MathHelper::Absolute((int8_t)(v >> 0)) +\n MathHelper::Absolute((int8_t)(v >> 8)) +\n MathHelper::Absolute((int8_t)(v >> 16)) +\n MathHelper::Absolute((int8_t)(v >> 24)));\n }\n};\n\n/*! L1-Norm Matrix (M=1)\n */\ntemplate \nstruct Norm1Matrix<\n T, 1, typename std::enable_if::value>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && out);\n\n const ValueType *m_end = m + dim;\n if (m != m_end) {\n *out = MathHelper::Absolute(*m++);\n }\n while (m != m_end) {\n *out += MathHelper::Absolute(*m++);\n }\n }\n};\n\n#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))\n/*! L1-Norm Matrix (FP32, M=1)\n */\ntemplate <>\nstruct Norm1Matrix {\n //! Type of value\n using ValueType = float;\n\n //! Compute the L1-norm of vectors\n static void Compute(const ValueType *m, size_t dim, float *out);\n};\n#endif // __SSE__ || (__ARM_NEON && __aarch64__)\n\n#if (defined(__F16C__) && defined(__AVX__)) || \\\n (defined(__ARM_NEON) && defined(__aarch64__))\n/*! L1-Norm Matrix (FP16, M=1)\n */\ntemplate <>\nstruct Norm1Matrix {\n //! Type of value\n using ValueType = Float16;\n\n //! Compute the L1-norm of vectors\n static void Compute(const ValueType *m, size_t dim, float *out);\n};\n#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__)\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/norm1_matrix_fp16.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"ailego/internal/cpu_features.h\"\n#include \"norm1_matrix.h\"\n#include \"norm_matrix_fp16.i\"\n\nnamespace zvec {\nnamespace ailego {\n\n#define NORM_FP32_STEP_GENERAL SA_FP32_GENERAL\n#define NORM_FP32_STEP_SSE SA_FP32_SSE\n#define NORM_FP32_STEP_AVX SA_FP32_AVX\n#define NORM_FP32_STEP_AVX512 SA_FP32_AVX512\n#define NORM_FP32_STEP_NEON SA_FP32_NEON\n#define NORM_FP16_STEP_GENERAL SA_FP16_GENERAL\n#define NORM_FP16_STEP_NEON SA_FP16_NEON\n\n#if defined(__SSE__)\nstatic const __m128 ABS_MASK_FP32_SSE =\n _mm_castsi128_ps(_mm_set1_epi32(0x7fffffffu));\n#endif // __SSE__\n\n#if defined(__AVX__)\nstatic const __m256 ABS_MASK_FP32_AVX =\n _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffffu));\n#endif // __AVX__\n\n#if defined(__AVX512F__)\nstatic const __m512 ABS_MASK_FP32_AVX512 =\n _mm512_castsi512_ps(_mm512_set1_epi32(0x7fffffffu));\n#endif // __AVX512F__\n\n//! Calculate sum of absolute (GENERAL)\n#define SA_FP32_GENERAL(m, sum) sum += FastAbs(m);\n\n//! Calculate sum of absolute (SSE)\n#define SA_FP32_SSE(xmm_m, xmm_sum) \\\n xmm_sum = _mm_add_ps(_mm_and_ps(xmm_m, ABS_MASK_FP32_SSE), xmm_sum);\n\n//! Calculate sum of absolute (AVX)\n#define SA_FP32_AVX(ymm_m, ymm_sum) \\\n ymm_sum = _mm256_add_ps(_mm256_and_ps(ymm_m, ABS_MASK_FP32_AVX), ymm_sum);\n\n//! Calculate sum of absolute (AVX512)\n#define SA_FP32_AVX512(zmm_m, zmm_sum) \\\n zmm_sum = _mm512_add_ps(_mm512_and_ps(zmm_m, ABS_MASK_FP32_AVX512), zmm_sum);\n\n//! Calculate sum of absolute (NEON)\n#define SA_FP32_NEON(v_m, v_sum) v_sum = vaddq_f32(vabsq_f32(v_m), v_sum);\n\n//! Calculate sum of absolute (GENERAL)\n#define SA_FP16_GENERAL(m, sum) sum += Float16::Absolute(m);\n\n//! Calculate sum of absolute (NEON)\n#define SA_FP16_NEON(v_m, v_sum) v_sum = vaddq_f16(vabsq_f16(v_m), v_sum);\n\n#if (defined(__F16C__) && defined(__AVX__)) || \\\n (defined(__ARM_NEON) && defined(__aarch64__))\n//! Compute the L1-norm of vectors (FP16, M=1)\nvoid Norm1Matrix::Compute(const ValueType *m, size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n NORM_FP16_1_NEON(m, dim, out, )\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n NORM_FP16_1_AVX512(m, dim, out, )\n return;\n }\n#endif\n NORM_FP16_1_AVX(m, dim, out, )\n#endif\n}\n#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__)\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/norm1_matrix_fp32.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"ailego/internal/cpu_features.h\"\n#include \"norm1_matrix.h\"\n#include \"norm_matrix_fp32.i\"\n\nnamespace zvec {\nnamespace ailego {\n\n#define NORM_FP32_STEP_GENERAL SA_FP32_GENERAL\n#define NORM_FP32_STEP_SSE SA_FP32_SSE\n#define NORM_FP32_STEP_AVX SA_FP32_AVX\n#define NORM_FP32_STEP_AVX512 SA_FP32_AVX512\n#define NORM_FP32_STEP_NEON SA_FP32_NEON\n\n#if defined(__SSE__)\n#define ABS_MASK_FP32_SSE _mm_castsi128_ps(_mm_set1_epi32(0x7fffffffu))\n#endif // __SSE__\n\n#if defined(__AVX__)\n#define ABS_MASK_FP32_AVX _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffffu))\n#endif // __AVX__\n\n#if defined(__AVX512F__)\n#define ABS_MASK_FP32_AVX512 _mm512_castsi512_ps(_mm512_set1_epi32(0x7fffffffu))\n#endif // __AVX512F__\n\n//! Calculate sum of absolute (GENERAL)\n#define SA_FP32_GENERAL(m, sum) sum += FastAbs(m);\n\n//! Calculate sum of absolute (SSE)\n#define SA_FP32_SSE(xmm_m, xmm_sum) \\\n xmm_sum = _mm_add_ps(_mm_and_ps(xmm_m, ABS_MASK_FP32_SSE), xmm_sum);\n\n//! Calculate sum of absolute (AVX)\n#define SA_FP32_AVX(ymm_m, ymm_sum) \\\n ymm_sum = _mm256_add_ps(_mm256_and_ps(ymm_m, ABS_MASK_FP32_AVX), ymm_sum);\n\n//! Calculate sum of absolute (AVX512)\n#define SA_FP32_AVX512(zmm_m, zmm_sum) \\\n zmm_sum = _mm512_add_ps(_mm512_and_ps(zmm_m, ABS_MASK_FP32_AVX512), zmm_sum);\n\n//! Calculate sum of absolute (NEON)\n#define SA_FP32_NEON(v_m, v_sum) v_sum = vaddq_f32(vabsq_f32(v_m), v_sum);\n\n#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))\n//! Compute the L1-norm of vectors (FP32, M=1)\nvoid Norm1Matrix::Compute(const ValueType *m, size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n NORM_FP32_1_NEON(m, dim, out, )\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n NORM_FP32_1_AVX512(m, dim, out, )\n return;\n }\n#endif\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n NORM_FP32_1_AVX(m, dim, out, )\n return;\n }\n#endif\n NORM_FP32_1_SSE(m, dim, out, )\n#endif\n}\n#endif // __SSE__ || (__ARM_NEON && __aarch64__)\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/norm2_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! L2-Norm Matrix\n */\ntemplate \nstruct Norm2Matrix;\n\n/*! L2-Norm Matrix\n */\ntemplate \nstruct Norm2Matrix::value &&\n sizeof(T) >= 2 && M >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && out);\n\n const ValueType *m_end = m + dim * M;\n if (m != m_end) {\n for (size_t i = 0; i < M; ++i) {\n ValueType v = m[i];\n *(out + i) = static_cast(v * v);\n }\n m += M;\n }\n while (m != m_end) {\n for (size_t i = 0; i < M; ++i) {\n ValueType v = m[i];\n *(out + i) += static_cast(v * v);\n }\n m += M;\n }\n for (size_t i = 0; i < M; ++i) {\n float v = *out;\n *out++ = std::sqrt(v);\n }\n }\n};\n\n/*! L2-Norm Matrix (INT8)\n */\ntemplate \nstruct Norm2Matrix= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && !(dim & 3) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *m_end = m_it + (dim >> 2) * M;\n\n if (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = Squared(m_it[i]);\n }\n m_it += M;\n }\n while (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += Squared(m_it[i]);\n }\n m_it += M;\n }\n for (size_t i = 0; i < M; ++i) {\n float v = *out;\n *out++ = std::sqrt(v);\n }\n }\n\n protected:\n //! Calculate sum of squared values\n static inline float Squared(uint32_t v) {\n return static_cast((int8_t)(v >> 0) * (int8_t)(v >> 0) +\n (int8_t)(v >> 8) * (int8_t)(v >> 8) +\n (int8_t)(v >> 16) * (int8_t)(v >> 16) +\n (int8_t)(v >> 24) * (int8_t)(v >> 24));\n }\n};\n\n/*! L2-Norm Matrix (M=1)\n */\ntemplate \nstruct Norm2Matrix<\n T, 1, typename std::enable_if::value>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && out);\n\n const ValueType *m_end = m + dim;\n if (m != m_end) {\n ValueType v = *m++;\n *out = static_cast(v * v);\n }\n while (m != m_end) {\n ValueType v = *m++;\n *out += static_cast(v * v);\n }\n *out = std::sqrt(*out);\n }\n};\n\n/*! L2-Norm Matrix (M=1, INT4)\n */\ntemplate <>\nstruct Norm2Matrix {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && !(dim & 1) && dim && out);\n\n const uint8_t *m_end = m + (dim >> 1);\n float square = 0.0f;\n while (m != m_end) {\n square += Squared(*m++);\n }\n *out = std::sqrt(square);\n }\n\n protected:\n //! Calculate sum of squared values\n static inline float Squared(uint8_t v) {\n return static_cast(\n ((int8_t)(v << 4) >> 4) * ((int8_t)(v << 4) >> 4) +\n ((int8_t)(v & 0xf0) >> 4) * ((int8_t)(v & 0xf0) >> 4));\n }\n};\n\n/*! L2-Norm Matrix (INT4)\n */\ntemplate \nstruct Norm2Matrix= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && !(dim & 7) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *m_end = m_it + (dim >> 3) * M;\n\n if (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = Squared(m_it[i]);\n }\n m_it += M;\n }\n while (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += Squared(m_it[i]);\n }\n m_it += M;\n }\n for (size_t i = 0; i < M; ++i) {\n float v = *out;\n *out++ = std::sqrt(v);\n }\n }\n\n protected:\n //! Calculate sum of squared values\n static inline float Squared(uint32_t u) {\n float sum = 0.0f;\n for (size_t i = 0; i < 32; i += 8) {\n uint8_t v = (uint8_t)(u >> i);\n int8_t lo = (int8_t)(v << 4) >> 4;\n int8_t hi = (int8_t)(v & 0xf0) >> 4;\n sum += hi * hi + lo * lo;\n }\n return sum;\n }\n};\n\n/*! Squared L2-Norm Matrix\n */\ntemplate \nstruct SquaredNorm2Matrix;\n\n/*! Squared L2-Norm Matrix\n */\ntemplate \nstruct SquaredNorm2Matrix<\n T, M,\n typename std::enable_if::value && sizeof(T) >= 2 &&\n M >= 2>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && out);\n\n const ValueType *m_end = m + dim * M;\n if (m != m_end) {\n for (size_t i = 0; i < M; ++i) {\n ValueType v = m[i];\n *(out + i) = static_cast(v * v);\n }\n m += M;\n }\n while (m != m_end) {\n for (size_t i = 0; i < M; ++i) {\n ValueType v = m[i];\n *(out + i) += static_cast(v * v);\n }\n m += M;\n }\n }\n};\n\n/*! Squared L2-Norm Matrix (INT8)\n */\ntemplate \nstruct SquaredNorm2Matrix= 2>::type> {\n //! Type of value\n using ValueType = int8_t;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && !(dim & 3) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *m_end = m_it + (dim >> 2) * M;\n\n if (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = Squared(m_it[i]);\n }\n m_it += M;\n }\n while (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += Squared(m_it[i]);\n }\n m_it += M;\n }\n }\n\n protected:\n //! Calculate sum of squared values\n static inline float Squared(uint32_t v) {\n return static_cast((int8_t)(v >> 0) * (int8_t)(v >> 0) +\n (int8_t)(v >> 8) * (int8_t)(v >> 8) +\n (int8_t)(v >> 16) * (int8_t)(v >> 16) +\n (int8_t)(v >> 24) * (int8_t)(v >> 24));\n }\n};\n\n/*! Squared L2-Norm Matrix (M=1)\n */\ntemplate \nstruct SquaredNorm2Matrix<\n T, 1, typename std::enable_if::value>::type> {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && out);\n\n const ValueType *m_end = m + dim;\n if (m != m_end) {\n ValueType v = *m++;\n *out = static_cast(v * v);\n }\n while (m != m_end) {\n ValueType v = *m++;\n *out += static_cast(v * v);\n }\n }\n};\n\n/*! L2-Norm Matrix (M=1, INT4)\n */\ntemplate <>\nstruct SquaredNorm2Matrix {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && !(dim & 1) && out);\n\n const uint8_t *m_end = m + (dim >> 1);\n *out = 0.0f;\n while (m != m_end) {\n *out += Squared(*m++);\n }\n }\n\n protected:\n //! Calculate sum of squared values\n static inline float Squared(uint8_t v) {\n return static_cast(\n ((int8_t)(v << 4) >> 4) * ((int8_t)(v << 4) >> 4) +\n ((int8_t)(v & 0xf0) >> 4) * ((int8_t)(v & 0xf0) >> 4));\n }\n};\n\n/*! Squared L2-Norm Matrix (INT4)\n */\ntemplate \nstruct SquaredNorm2Matrix= 2>::type> {\n //! Type of value\n using ValueType = uint8_t;\n\n //! Compute the norm of vectors\n static inline void Compute(const ValueType *m, size_t dim, float *out) {\n ailego_assert(m && dim && !(dim & 7) && out);\n\n const uint32_t *m_it = reinterpret_cast(m);\n const uint32_t *m_end = m_it + (dim >> 3) * M;\n\n if (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) = Squared(m_it[i]);\n }\n m_it += M;\n }\n while (m_it != m_end) {\n for (size_t i = 0; i < M; ++i) {\n *(out + i) += Squared(m_it[i]);\n }\n m_it += M;\n }\n }\n\n protected:\n //! Calculate sum of squared values\n static inline float Squared(uint32_t u) {\n float sum = 0.0f;\n for (size_t i = 0; i < 32; i += 8) {\n uint8_t v = (uint8_t)(u >> i);\n int8_t lo = (int8_t)(v << 4) >> 4;\n int8_t hi = (int8_t)(v & 0xf0) >> 4;\n sum += hi * hi + lo * lo;\n }\n return sum;\n }\n};\n\n#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))\n/*! L2-Norm Matrix (FP32, M=1)\n */\ntemplate <>\nstruct Norm2Matrix {\n //! Type of value\n using ValueType = float;\n\n //! Compute the L2-norm of vectors\n static void Compute(const ValueType *m, size_t dim, float *out);\n};\n\n/*! Squared L2-Norm Matrix (FP32, M=1)\n */\ntemplate <>\nstruct SquaredNorm2Matrix {\n //! Type of value\n using ValueType = float;\n\n //! Compute the squared L2-norm of vectors\n static void Compute(const ValueType *m, size_t dim, float *out);\n};\n#endif // __SSE__ || (__ARM_NEON && __aarch64__)\n\n#if (defined(__F16C__) && defined(__AVX__)) || \\\n (defined(__ARM_NEON) && defined(__aarch64__))\n/*! L2-Norm Matrix (FP16, M=1)\n */\ntemplate <>\nstruct Norm2Matrix {\n //! Type of value\n using ValueType = Float16;\n\n //! Compute the L2-norm of vectors\n static void Compute(const ValueType *m, size_t dim, float *out);\n};\n\n/*! Squared L2-Norm Matrix (FP16, M=1)\n */\ntemplate <>\nstruct SquaredNorm2Matrix {\n //! Type of value\n using ValueType = Float16;\n\n //! Compute the squared L2-norm of vectors\n static void Compute(const ValueType *m, size_t dim, float *out);\n};\n#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__)\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/norm2_matrix_fp16.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"ailego/internal/cpu_features.h\"\n#include \"norm2_matrix.h\"\n#include \"norm_matrix_fp16.i\"\n\nnamespace zvec {\nnamespace ailego {\n\n#define NORM_FP32_STEP_GENERAL SS_FP32_GENERAL\n#define NORM_FP32_STEP_SSE SS_FP32_SSE\n#define NORM_FP32_STEP_AVX SS_FP32_AVX\n#define NORM_FP32_STEP_AVX512 SS_FP32_AVX512\n#define NORM_FP32_STEP_NEON SS_FP32_NEON\n#define NORM_FP16_STEP_GENERAL SS_FP16_GENERAL\n#define NORM_FP16_STEP_NEON SS_FP16_NEON\n\n//! Calculate sum of squared (GENERAL)\n#define SS_FP32_GENERAL(m, sum) sum += (m) * (m);\n\n//! Calculate sum of squared (SSE)\n#define SS_FP32_SSE(xmm_m, xmm_sum) \\\n xmm_sum = _mm_fmadd_ps(xmm_m, xmm_m, xmm_sum);\n\n//! Calculate sum of squared (AVX)\n#define SS_FP32_AVX(ymm_m, ymm_sum) \\\n ymm_sum = _mm256_fmadd_ps(ymm_m, ymm_m, ymm_sum);\n\n//! Calculate sum of squared (AVX512)\n#define SS_FP32_AVX512(zmm_m, zmm_sum) \\\n zmm_sum = _mm512_fmadd_ps(zmm_m, zmm_m, zmm_sum);\n\n//! Calculate sum of squared (NEON)\n#define SS_FP32_NEON(v_m, v_sum) v_sum = vfmaq_f32(v_sum, v_m, v_m);\n\n//! Calculate sum of squared (GENERAL)\n#define SS_FP16_GENERAL(m, sum) sum += (m) * (m);\n\n//! Calculate sum of squared (NEON)\n#define SS_FP16_NEON(v_m, v_sum) v_sum = vfmaq_f16(v_sum, v_m, v_m);\n\n#if (defined(__F16C__) && defined(__AVX__)) || \\\n (defined(__ARM_NEON) && defined(__aarch64__))\n//! Compute the L2-norm of vectors (FP16, M=1)\nvoid Norm2Matrix::Compute(const ValueType *m, size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n NORM_FP16_1_NEON(m, dim, out, std::sqrt)\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n NORM_FP16_1_AVX512(m, dim, out, std::sqrt)\n return;\n }\n#endif\n NORM_FP16_1_AVX(m, dim, out, std::sqrt)\n#endif\n}\n\n//! Compute the L2-norm of vectors (FP16, M=1)\nvoid SquaredNorm2Matrix::Compute(const ValueType *m, size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n NORM_FP16_1_NEON(m, dim, out, )\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n NORM_FP16_1_AVX512(m, dim, out, )\n return;\n }\n#endif\n NORM_FP16_1_AVX(m, dim, out, )\n#endif\n}\n#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__)\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/norm2_matrix_fp32.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \"norm2_matrix.h\"\n#include \"norm_matrix_fp32.i\"\n\nnamespace zvec {\nnamespace ailego {\n\n#define NORM_FP32_STEP_GENERAL SS_FP32_GENERAL\n#define NORM_FP32_STEP_SSE SS_FP32_SSE\n#define NORM_FP32_STEP_AVX SS_FP32_AVX\n#define NORM_FP32_STEP_AVX512 SS_FP32_AVX512\n#define NORM_FP32_STEP_NEON SS_FP32_NEON\n\n//! Calculate sum of squared (GENERAL)\n#define SS_FP32_GENERAL(m, sum) sum += (m) * (m);\n\n//! Calculate sum of squared (SSE)\n#define SS_FP32_SSE(xmm_m, xmm_sum) \\\n xmm_sum = _mm_fmadd_ps(xmm_m, xmm_m, xmm_sum);\n\n//! Calculate sum of squared (AVX)\n#define SS_FP32_AVX(ymm_m, ymm_sum) \\\n ymm_sum = _mm256_fmadd_ps(ymm_m, ymm_m, ymm_sum);\n\n//! Calculate sum of squared (AVX512)\n#define SS_FP32_AVX512(zmm_m, zmm_sum) \\\n zmm_sum = _mm512_fmadd_ps(zmm_m, zmm_m, zmm_sum);\n\n//! Calculate sum of squared (NEON)\n#define SS_FP32_NEON(v_m, v_sum) v_sum = vfmaq_f32(v_sum, v_m, v_m);\n\n#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))\n//! Compute the L2-norm of vectors (FP32, M=1)\nvoid Norm2Matrix::Compute(const ValueType *m, size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n NORM_FP32_1_NEON(m, dim, out, std::sqrt)\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n NORM_FP32_1_AVX512(m, dim, out, std::sqrt)\n return;\n }\n#endif\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n NORM_FP32_1_AVX(m, dim, out, std::sqrt)\n return;\n }\n#endif\n NORM_FP32_1_SSE(m, dim, out, std::sqrt)\n#endif\n}\n\n//! Compute the squared L2-norm of vectors (FP32, M=1)\nvoid SquaredNorm2Matrix::Compute(const ValueType *m, size_t dim,\n float *out) {\n#if defined(__ARM_NEON)\n NORM_FP32_1_NEON(m, dim, out, )\n#else\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n NORM_FP32_1_AVX512(m, dim, out, )\n return;\n }\n#endif\n#if defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n NORM_FP32_1_AVX(m, dim, out, )\n return;\n }\n#endif\n NORM_FP32_1_SSE(m, dim, out, )\n#endif\n}\n#endif // __SSE__ || (__ARM_NEON && __aarch64__)\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math/norm_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \"norm1_matrix.h\"\n#include \"norm2_matrix.h\"\n" }, { "path": "src/ailego/math/norm_matrix_fp16.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"matrix_define.i\"\n#include \"matrix_utility.i\"\n\n#if !defined(__FMA__)\n#define _mm_fmadd_ps(a, b, c) _mm_add_ps(_mm_mul_ps((a), (b)), (c))\n#define _mm256_fmadd_ps(a, b, c) _mm256_add_ps(_mm256_mul_ps((a), (b)), (c))\n#endif // !__FMA__\n\n//! Mask process of computing norm (FP16)\n#define NORM_FP16_MASK_AVX(m, cnt, _RES) \\\n switch (cnt) { \\\n case 7: { \\\n __m256 ymm_m = _mm256_cvtph_ps( \\\n _mm_set_epi16(0, *((const short *)(m) + 6), \\\n *((const short *)(m) + 5), *((const short *)(m) + 4), \\\n *((const short *)(m) + 3), *((const short *)(m) + 2), \\\n *((const short *)(m) + 1), *((const short *)(m)))); \\\n NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \\\n break; \\\n } \\\n case 6: { \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm_set_epi32(0, *((const int *)(m) + 2), \\\n *((const int *)(m) + 1), \\\n *((const int *)(m)))); \\\n NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \\\n break; \\\n } \\\n case 5: { \\\n __m256 ymm_m = _mm256_cvtph_ps( \\\n _mm_set_epi16(0, 0, 0, *((const short *)(m) + 4), \\\n *((const short *)(m) + 3), *((const short *)(m) + 2), \\\n *((const short *)(m) + 1), *((const short *)(m)))); \\\n NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \\\n break; \\\n } \\\n case 4: { \\\n __m256 ymm_m = _mm256_cvtph_ps( \\\n _mm_set_epi64((__m64)(0ull), *((const __m64 *)(m)))); \\\n NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \\\n break; \\\n } \\\n case 3: { \\\n __m256 ymm_m = _mm256_cvtph_ps( \\\n _mm_set_epi16(0, 0, 0, 0, 0, *((const short *)(m) + 2), \\\n *((const short *)(m) + 1), *((const short *)(m)))); \\\n NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \\\n break; \\\n } \\\n case 2: { \\\n __m256 ymm_m = \\\n _mm256_cvtph_ps(_mm_set_epi32(0, 0, 0, *((const int *)(m)))); \\\n NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \\\n break; \\\n } \\\n case 1: { \\\n __m256 ymm_m = _mm256_cvtph_ps( \\\n _mm_set_epi16(0, 0, 0, 0, 0, 0, 0, *((const short *)(m)))); \\\n NORM_FP32_STEP_AVX(ymm_m, _RES##_0_0) \\\n break; \\\n } \\\n }\n\n//! Compute the norm of vectors (FP16, M=1)\n#define NORM_FP16_1_AVX(m, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const Float16 *last = m + dim; \\\n const Float16 *last_aligned = m + ((dim >> 4) << 4); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; m != last_aligned; m += 16) { \\\n __m256i ymm_mi = _mm256_load_si256((const __m256i *)m); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n NORM_FP32_STEP_AVX(ymm_m_0, ymm_sum_0_0) \\\n NORM_FP32_STEP_AVX(ymm_m_1, ymm_sum_0_0) \\\n } \\\n if (last >= last_aligned + 8) { \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)m)); \\\n NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \\\n m += 8; \\\n } \\\n } else { \\\n for (; m != last_aligned; m += 16) { \\\n __m256i ymm_mi = _mm256_loadu_si256((const __m256i *)m); \\\n __m256 ymm_m_0 = _mm256_cvtph_ps(_mm256_castsi256_si128(ymm_mi)); \\\n __m256 ymm_m_1 = _mm256_cvtph_ps(_mm256_extractf128_si256(ymm_mi, 1)); \\\n NORM_FP32_STEP_AVX(ymm_m_0, ymm_sum_0_0) \\\n NORM_FP32_STEP_AVX(ymm_m_1, ymm_sum_0_0) \\\n } \\\n if (last >= last_aligned + 8) { \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)m)); \\\n NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \\\n m += 8; \\\n } \\\n } \\\n NORM_FP16_MASK_AVX(m, (last - m), ymm_sum) \\\n *out = _NORM(HorizontalAdd_FP32_V256(ymm_sum_0_0));\n\n//! Compute the norm of vectors (FP16, M=1)\n#define NORM_FP16_1_AVX512(m, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m512, zmm_sum, _mm512_setzero_ps()) \\\n const Float16 *last = m + dim; \\\n const Float16 *last_aligned = m + ((dim >> 5) << 5); \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (; m != last_aligned; m += 32) { \\\n __m512i zmm_mi = _mm512_load_si512((const __m512i *)m); \\\n __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \\\n __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \\\n NORM_FP32_STEP_AVX512(zmm_m_0, zmm_sum_0_0) \\\n NORM_FP32_STEP_AVX512(zmm_m_1, zmm_sum_0_1) \\\n } \\\n if (last >= last_aligned + 16) { \\\n __m512 zmm_m = _mm512_cvtph_ps(_mm256_load_si256((const __m256i *)m)); \\\n NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \\\n m += 16; \\\n } \\\n } else { \\\n for (; m != last_aligned; m += 32) { \\\n __m512i zmm_mi = _mm512_loadu_si512((const __m512i *)m); \\\n __m512 zmm_m_0 = _mm512_cvtph_ps(_mm512_castsi512_si256(zmm_mi)); \\\n __m512 zmm_m_1 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(zmm_mi, 1)); \\\n NORM_FP32_STEP_AVX512(zmm_m_0, zmm_sum_0_0) \\\n NORM_FP32_STEP_AVX512(zmm_m_1, zmm_sum_0_1) \\\n } \\\n if (last >= last_aligned + 16) { \\\n __m512 zmm_m = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)m)); \\\n NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \\\n m += 16; \\\n } \\\n } \\\n float result = \\\n HorizontalAdd_FP32_V512(_mm512_add_ps(zmm_sum_0_0, zmm_sum_0_1)); \\\n if (m != last) { \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n if (last >= m + 8) { \\\n __m256 ymm_m = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)m)); \\\n NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \\\n m += 8; \\\n } \\\n NORM_FP16_MASK_AVX(m, (last - m), ymm_sum) \\\n result += HorizontalAdd_FP32_V256(ymm_sum_0_0); \\\n } \\\n *out = _NORM(result);\n\n#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)\n//! Compute the norm of vectors (FP16, M=1)\n#define NORM_FP16_1_NEON(m, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, float16x8_t, v_sum, vdupq_n_f16(0)) \\\n const Float16 *last = m + dim; \\\n const Float16 *last_aligned = m + ((dim >> 3) << 3); \\\n for (; m != last_aligned; m += 8) { \\\n float16x8_t v_m = vld1q_f16((const float16_t *)m); \\\n NORM_FP16_STEP_NEON(v_m, v_sum_0_0) \\\n } \\\n if (last >= m + 4) { \\\n float16x8_t v_m = vreinterpretq_f16_u64( \\\n vld1q_lane_u64((const uint64_t *)m, vdupq_n_u64(0), 0)); \\\n NORM_FP16_STEP_NEON(v_m, v_sum_0_0) \\\n m += 4; \\\n } \\\n float result = vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(v_sum_0_0)), \\\n vcvt_high_f32_f16(v_sum_0_0))); \\\n switch (last - m) { \\\n case 3: \\\n NORM_FP16_STEP_GENERAL(m[2], result) \\\n /* FALLTHRU */ \\\n case 2: \\\n NORM_FP16_STEP_GENERAL(m[1], result) \\\n /* FALLTHRU */ \\\n case 1: \\\n NORM_FP16_STEP_GENERAL(m[0], result) \\\n } \\\n *out = _NORM(result);\n\n#else\n//! Compute the norm of vectors (FP16, M=1)\n#define NORM_FP16_1_NEON(m, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n const Float16 *last = m + dim; \\\n const Float16 *last_aligned = m + ((dim >> 3) << 3); \\\n for (; m != last_aligned; m += 8) { \\\n float16x8_t v_m = vld1q_f16((const float16_t *)m); \\\n float32x4_t v_n_0 = vcvt_f32_f16(vget_low_f16(v_m)); \\\n float32x4_t v_n_1 = vcvt_high_f32_f16(v_m); \\\n NORM_FP32_STEP_NEON(v_n_0, v_sum_0_0) \\\n NORM_FP32_STEP_NEON(v_n_1, v_sum_0_1) \\\n } \\\n if (last >= m + 4) { \\\n float32x4_t v_m = vcvt_f32_f16(vld1_f16((const float16_t *)m)); \\\n NORM_FP32_STEP_NEON(v_m, v_sum_0_0) \\\n m += 4; \\\n } \\\n float result = vaddvq_f32(vaddq_f32(v_sum_0_0, v_sum_0_1)); \\\n switch (last - m) { \\\n case 3: \\\n NORM_FP16_STEP_GENERAL(m[2], result) \\\n /* FALLTHRU */ \\\n case 2: \\\n NORM_FP16_STEP_GENERAL(m[1], result) \\\n /* FALLTHRU */ \\\n case 1: \\\n NORM_FP16_STEP_GENERAL(m[0], result) \\\n } \\\n *out = _NORM(result);\n\n#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC" }, { "path": "src/ailego/math/norm_matrix_fp32.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"matrix_define.i\"\n#include \"matrix_utility.i\"\n\n#if !defined(__FMA__)\n#define _mm_fmadd_ps(a, b, c) _mm_add_ps(_mm_mul_ps((a), (b)), (c))\n#define _mm256_fmadd_ps(a, b, c) _mm256_add_ps(_mm256_mul_ps((a), (b)), (c))\n#endif // !__FMA__\n\n//! Mask process of computing norm (FP32)\n#define NORM_FP32_MASK_SSE(m, cnt, _RES) \\\n switch (cnt) { \\\n case 3: { \\\n __m128 xmm_m = _mm_set_ps(0.0f, m[2], m[1], m[0]); \\\n NORM_FP32_STEP_SSE(xmm_m, _RES##_0_0) \\\n break; \\\n } \\\n case 2: { \\\n __m128 xmm_m = _mm_set_ps(0.0f, 0.0f, m[1], m[0]); \\\n NORM_FP32_STEP_SSE(xmm_m, _RES##_0_0) \\\n break; \\\n } \\\n case 1: { \\\n __m128 xmm_m = _mm_set_ps(0.0f, 0.0f, 0.0f, m[0]); \\\n NORM_FP32_STEP_SSE(xmm_m, _RES##_0_0) \\\n break; \\\n } \\\n }\n\n//! Compute the norm of vectors (FP32, M=1)\n#define NORM_FP32_1_SSE(m, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m128, xmm_sum, _mm_setzero_ps()) \\\n const float *last = m + dim; \\\n const float *last_aligned = m + ((dim >> 3) << 3); \\\n if (((uintptr_t)m & 0xf) == 0) { \\\n for (; m != last_aligned; m += 8) { \\\n __m128 xmm_m_0 = _mm_load_ps(m + 0); \\\n __m128 xmm_m_1 = _mm_load_ps(m + 4); \\\n NORM_FP32_STEP_SSE(xmm_m_0, xmm_sum_0_0) \\\n NORM_FP32_STEP_SSE(xmm_m_1, xmm_sum_0_0) \\\n } \\\n if (last >= last_aligned + 4) { \\\n __m128 xmm_m = _mm_load_ps(m); \\\n NORM_FP32_STEP_SSE(xmm_m, xmm_sum_0_0) \\\n m += 4; \\\n } \\\n } else { \\\n for (; m != last_aligned; m += 8) { \\\n __m128 xmm_m_0 = _mm_loadu_ps(m + 0); \\\n __m128 xmm_m_1 = _mm_loadu_ps(m + 4); \\\n NORM_FP32_STEP_SSE(xmm_m_0, xmm_sum_0_0) \\\n NORM_FP32_STEP_SSE(xmm_m_1, xmm_sum_0_0) \\\n } \\\n if (last >= last_aligned + 4) { \\\n __m128 xmm_m = _mm_loadu_ps(m); \\\n NORM_FP32_STEP_SSE(xmm_m, xmm_sum_0_0) \\\n m += 4; \\\n } \\\n } \\\n NORM_FP32_MASK_SSE(m, (last - m), xmm_sum) \\\n *out = _NORM(HorizontalAdd_FP32_V128(xmm_sum_0_0));\n\n//! Compute the norm of vectors (FP32, M=1)\n#define NORM_FP32_1_AVX(m, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 1, __m256, ymm_sum, _mm256_setzero_ps()) \\\n const float *last = m + dim; \\\n const float *last_aligned = m + ((dim >> 4) << 4); \\\n if (((uintptr_t)m & 0x1f) == 0) { \\\n for (; m != last_aligned; m += 16) { \\\n __m256 ymm_m_0 = _mm256_load_ps(m + 0); \\\n __m256 ymm_m_1 = _mm256_load_ps(m + 8); \\\n NORM_FP32_STEP_AVX(ymm_m_0, ymm_sum_0_0) \\\n NORM_FP32_STEP_AVX(ymm_m_1, ymm_sum_0_0) \\\n } \\\n if (last >= last_aligned + 8) { \\\n __m256 ymm_m = _mm256_load_ps(m); \\\n NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \\\n m += 8; \\\n } \\\n } else { \\\n for (; m != last_aligned; m += 16) { \\\n __m256 ymm_m_0 = _mm256_loadu_ps(m + 0); \\\n __m256 ymm_m_1 = _mm256_loadu_ps(m + 8); \\\n NORM_FP32_STEP_AVX(ymm_m_0, ymm_sum_0_0) \\\n NORM_FP32_STEP_AVX(ymm_m_1, ymm_sum_0_0) \\\n } \\\n if (last >= last_aligned + 8) { \\\n __m256 ymm_m = _mm256_loadu_ps(m); \\\n NORM_FP32_STEP_AVX(ymm_m, ymm_sum_0_0) \\\n m += 8; \\\n } \\\n } \\\n float result = HorizontalAdd_FP32_V256(ymm_sum_0_0); \\\n if (m != last) { \\\n __m128 xmm_sum_0_0 = _mm_setzero_ps(); \\\n if (last >= m + 4) { \\\n __m128 xmm_m = _mm_loadu_ps(m); \\\n NORM_FP32_STEP_SSE(xmm_m, xmm_sum_0_0) \\\n m += 4; \\\n } \\\n NORM_FP32_MASK_SSE(m, (last - m), xmm_sum) \\\n result += HorizontalAdd_FP32_V128(xmm_sum_0_0); \\\n } \\\n *out = _NORM(result);\n\n//! Compute the norm of vectors (FP32, M=1)\n#define NORM_FP32_1_AVX512(m, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, __m512, zmm_sum, _mm512_setzero_ps()) \\\n const float *last = m + dim; \\\n const float *last_aligned = m + ((dim >> 5) << 5); \\\n if (((uintptr_t)m & 0x3f) == 0) { \\\n for (; m != last_aligned; m += 32) { \\\n __m512 zmm_m_0 = _mm512_load_ps(m + 0); \\\n NORM_FP32_STEP_AVX512(zmm_m_0, zmm_sum_0_0) \\\n __m512 zmm_m_1 = _mm512_load_ps(m + 16); \\\n NORM_FP32_STEP_AVX512(zmm_m_1, zmm_sum_0_1) \\\n } \\\n if (last >= last_aligned + 16) { \\\n __m512 zmm_m = _mm512_load_ps(m); \\\n NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \\\n m += 16; \\\n } \\\n } else { \\\n for (; m != last_aligned; m += 32) { \\\n __m512 zmm_m_0 = _mm512_loadu_ps(m + 0); \\\n NORM_FP32_STEP_AVX512(zmm_m_0, zmm_sum_0_0) \\\n __m512 zmm_m_1 = _mm512_loadu_ps(m + 16); \\\n NORM_FP32_STEP_AVX512(zmm_m_1, zmm_sum_0_1) \\\n } \\\n if (last >= last_aligned + 16) { \\\n __m512 zmm_m = _mm512_loadu_ps(m); \\\n NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \\\n m += 16; \\\n } \\\n } \\\n if (m != last) { \\\n __mmask16 mask = (__mmask16)((1 << (last - m)) - 1); \\\n __m512 zmm_m = _mm512_mask_loadu_ps(_mm512_setzero_ps(), mask, m); \\\n NORM_FP32_STEP_AVX512(zmm_m, zmm_sum_0_0) \\\n } \\\n float result = \\\n HorizontalAdd_FP32_V512(_mm512_add_ps(zmm_sum_0_0, zmm_sum_0_1)); \\\n *out = _NORM(result);\n\n//! Compute the norm of vectors (FP32, M=1)\n#define NORM_FP32_1_NEON(m, dim, out, _NORM) \\\n MATRIX_VAR_INIT(1, 2, float32x4_t, v_sum, vdupq_n_f32(0)) \\\n const float *last = m + dim; \\\n const float *last_aligned = m + ((dim >> 3) << 3); \\\n for (; m != last_aligned; m += 8) { \\\n float32x4_t v_m_0 = vld1q_f32(m + 0); \\\n float32x4_t v_m_1 = vld1q_f32(m + 4); \\\n NORM_FP32_STEP_NEON(v_m_0, v_sum_0_0) \\\n NORM_FP32_STEP_NEON(v_m_1, v_sum_0_1) \\\n } \\\n if (last >= last_aligned + 4) { \\\n float32x4_t v_m = vld1q_f32(m); \\\n NORM_FP32_STEP_NEON(v_m, v_sum_0_0) \\\n m += 4; \\\n } \\\n float result = vaddvq_f32(vaddq_f32(v_sum_0_0, v_sum_0_1)); \\\n switch (last - m) { \\\n case 3: \\\n NORM_FP32_STEP_GENERAL(m[2], result) \\\n /* FALLTHRU */ \\\n case 2: \\\n NORM_FP32_STEP_GENERAL(m[1], result) \\\n /* FALLTHRU */ \\\n case 1: \\\n NORM_FP32_STEP_GENERAL(m[0], result) \\\n } \\\n *out = _NORM(result);\n" }, { "path": "src/ailego/math/normalizer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"normalizer.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n#if (defined(__ARM_NEON) && defined(__aarch64__))\nstatic inline void NormalizeNEON(float *arr, size_t dim, float norm) {\n float *last = arr + dim;\n float *last_aligned = arr + ((dim >> 3) << 3);\n\n float32x4_t v_norm = vdupq_n_f32(norm);\n for (; arr != last_aligned; arr += 8) {\n vst1q_f32(arr + 0, vdivq_f32(vld1q_f32(arr + 0), v_norm));\n vst1q_f32(arr + 4, vdivq_f32(vld1q_f32(arr + 4), v_norm));\n }\n if (last >= last_aligned + 4) {\n vst1q_f32(arr, vdivq_f32(vld1q_f32(arr), v_norm));\n arr += 4;\n }\n switch (last - arr) {\n case 3:\n arr[2] /= norm;\n /* FALLTHRU */\n case 2:\n arr[1] /= norm;\n /* FALLTHRU */\n case 1:\n arr[0] /= norm;\n }\n}\n\n#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)\nstatic inline void NormalizeNEON(float16_t *arr, size_t dim, float norm) {\n float16_t *last = arr + dim;\n float16_t *last_aligned = arr + ((dim >> 4) << 4);\n\n float16x8_t v_norm = vdupq_n_f16((float16_t)norm);\n for (; arr != last_aligned; arr += 16) {\n vst1q_f16(arr + 0, vdivq_f16(vld1q_f16(arr + 0), v_norm));\n vst1q_f16(arr + 8, vdivq_f16(vld1q_f16(arr + 8), v_norm));\n }\n if (last >= arr + 8) {\n vst1q_f16(arr, vdivq_f16(vld1q_f16(arr), v_norm));\n arr += 8;\n }\n if (last >= arr + 4) {\n vst1_f16(arr, vdiv_f16(vld1_f16(arr), vget_low_f16(v_norm)));\n arr += 4;\n }\n switch (last - arr) {\n case 3:\n arr[2] /= norm;\n /* FALLTHRU */\n case 2:\n arr[1] /= norm;\n /* FALLTHRU */\n case 1:\n arr[0] /= norm;\n }\n}\n#else\nstatic inline void NormalizeNEON(float16_t *arr, size_t dim, float norm) {\n float16_t *last = arr + dim;\n float16_t *last_aligned = arr + ((dim >> 4) << 4);\n\n float32x4_t v_norm = vdupq_n_f32(norm);\n for (; arr != last_aligned; arr += 16) {\n float16x8_t vf16_0 = vld1q_f16(arr + 0);\n float16x8_t vf16_1 = vld1q_f16(arr + 8);\n vf16_0 = vcombine_f16(\n vcvt_f16_f32(vdivq_f32(vcvt_f32_f16(vget_low_f16(vf16_0)), v_norm)),\n vcvt_f16_f32(vdivq_f32(vcvt_high_f32_f16(vf16_0), v_norm)));\n vf16_1 = vcombine_f16(\n vcvt_f16_f32(vdivq_f32(vcvt_f32_f16(vget_low_f16(vf16_1)), v_norm)),\n vcvt_f16_f32(vdivq_f32(vcvt_high_f32_f16(vf16_1), v_norm)));\n vst1q_f16(arr + 0, vf16_0);\n vst1q_f16(arr + 8, vf16_1);\n }\n if (last >= arr + 8) {\n float16x8_t vf16 = vld1q_f16(arr);\n vf16 = vcombine_f16(\n vcvt_f16_f32(vdivq_f32(vcvt_f32_f16(vget_low_f16(vf16)), v_norm)),\n vcvt_f16_f32(vdivq_f32(vcvt_high_f32_f16(vf16), v_norm)));\n vst1q_f16(arr, vf16);\n arr += 8;\n }\n if (last >= arr + 4) {\n vst1_f16(arr, vcvt_f16_f32(vdivq_f32(vcvt_f32_f16(vld1_f16(arr)), v_norm)));\n arr += 4;\n }\n switch (last - arr) {\n case 3:\n arr[2] /= norm;\n /* FALLTHRU */\n case 2:\n arr[1] /= norm;\n /* FALLTHRU */\n case 1:\n arr[0] /= norm;\n }\n}\n#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC\n#endif // __ARM_NEON && __aarch64__\n\n#if defined(__AVX__)\n#if defined(__AVX512F__)\nstatic inline void NormalizeAVX512(float *arr, size_t dim, float norm) {\n float *last = arr + dim;\n float *last_aligned = arr + ((dim >> 4) << 4);\n\n __m512 zmm_norm = _mm512_set1_ps(norm);\n if (((uintptr_t)arr & 0x3f) == 0) {\n for (; arr != last_aligned; arr += 16) {\n _mm512_store_ps(arr, _mm512_div_ps(_mm512_load_ps(arr), zmm_norm));\n }\n if (last >= arr + 8) {\n __m256 ymm_norm = _mm256_set1_ps(norm);\n _mm256_store_ps(arr, _mm256_div_ps(_mm256_load_ps(arr), ymm_norm));\n arr += 8;\n }\n if (last >= arr + 4) {\n __m128 xmm_norm = _mm_set1_ps(norm);\n _mm_store_ps(arr, _mm_div_ps(_mm_load_ps(arr), xmm_norm));\n arr += 4;\n }\n } else {\n for (; arr != last_aligned; arr += 16) {\n _mm512_storeu_ps(arr, _mm512_div_ps(_mm512_loadu_ps(arr), zmm_norm));\n }\n if (last >= arr + 8) {\n __m256 ymm_norm = _mm256_set1_ps(norm);\n _mm256_storeu_ps(arr, _mm256_div_ps(_mm256_loadu_ps(arr), ymm_norm));\n arr += 8;\n }\n if (last >= arr + 4) {\n __m128 xmm_norm = _mm_set1_ps(norm);\n _mm_storeu_ps(arr, _mm_div_ps(_mm_loadu_ps(arr), xmm_norm));\n arr += 4;\n }\n }\n switch (last - arr) {\n case 3:\n arr[2] /= norm;\n /* FALLTHRU */\n case 2:\n arr[1] /= norm;\n /* FALLTHRU */\n case 1:\n arr[0] /= norm;\n }\n}\n#endif // __AVX512F__\n\nstatic inline void NormalizeAVX(float *arr, size_t dim, float norm) {\n float *last = arr + dim;\n float *last_aligned = arr + ((dim >> 4) << 4);\n\n __m256 ymm_norm = _mm256_set1_ps(norm);\n if (((uintptr_t)arr & 0x1f) == 0) {\n for (; arr != last_aligned; arr += 16) {\n _mm256_store_ps(arr + 0,\n _mm256_div_ps(_mm256_load_ps(arr + 0), ymm_norm));\n _mm256_store_ps(arr + 8,\n _mm256_div_ps(_mm256_load_ps(arr + 8), ymm_norm));\n }\n if (last >= arr + 8) {\n _mm256_store_ps(arr, _mm256_div_ps(_mm256_load_ps(arr), ymm_norm));\n arr += 8;\n }\n if (last >= arr + 4) {\n __m128 xmm_norm = _mm_set1_ps(norm);\n _mm_store_ps(arr, _mm_div_ps(_mm_load_ps(arr), xmm_norm));\n arr += 4;\n }\n } else {\n for (; arr != last_aligned; arr += 16) {\n _mm256_storeu_ps(arr + 0,\n _mm256_div_ps(_mm256_loadu_ps(arr + 0), ymm_norm));\n _mm256_storeu_ps(arr + 8,\n _mm256_div_ps(_mm256_loadu_ps(arr + 8), ymm_norm));\n }\n if (last >= arr + 8) {\n _mm256_storeu_ps(arr, _mm256_div_ps(_mm256_loadu_ps(arr), ymm_norm));\n arr += 8;\n }\n if (last >= arr + 4) {\n __m128 xmm_norm = _mm_set1_ps(norm);\n _mm_storeu_ps(arr, _mm_div_ps(_mm_loadu_ps(arr), xmm_norm));\n arr += 4;\n }\n }\n switch (last - arr) {\n case 3:\n arr[2] /= norm;\n /* FALLTHRU */\n case 2:\n arr[1] /= norm;\n /* FALLTHRU */\n case 1:\n arr[0] /= norm;\n }\n}\n#endif // __AVX__\n\n#if defined(__AVX__) && defined(__F16C__)\n#if defined(__AVX512F__)\nstatic inline void NormalizeAVX512(uint16_t *arr, size_t dim, float norm) {\n uint16_t *last = arr + dim;\n uint16_t *last_aligned = arr + ((dim >> 4) << 4);\n\n __m512 zmm_norm = _mm512_set1_ps(norm);\n if (((uintptr_t)arr & 0x1f) == 0) {\n for (; arr != last_aligned; arr += 16) {\n _mm256_store_si256(\n (__m256i *)arr,\n _mm512_cvtps_ph(_mm512_div_ps(_mm512_cvtph_ps(_mm256_load_si256(\n (const __m256i *)arr)),\n zmm_norm),\n _MM_FROUND_NO_EXC));\n }\n if (last >= arr + 8) {\n __m256 ymm_norm = _mm256_set1_ps(norm);\n _mm_store_si128(\n (__m128i *)arr,\n _mm256_cvtps_ph(_mm256_div_ps(_mm256_cvtph_ps(_mm_load_si128(\n (const __m128i *)arr)),\n ymm_norm),\n _MM_FROUND_NO_EXC));\n arr += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 16) {\n _mm256_storeu_si256(\n (__m256i *)arr,\n _mm512_cvtps_ph(_mm512_div_ps(_mm512_cvtph_ps(_mm256_loadu_si256(\n (const __m256i *)arr)),\n zmm_norm),\n _MM_FROUND_NO_EXC));\n }\n if (last >= arr + 8) {\n __m256 ymm_norm = _mm256_set1_ps(norm);\n _mm_storeu_si128(\n (__m128i *)arr,\n _mm256_cvtps_ph(_mm256_div_ps(_mm256_cvtph_ps(_mm_loadu_si128(\n (const __m128i *)arr)),\n ymm_norm),\n _MM_FROUND_NO_EXC));\n arr += 8;\n }\n }\n if (last >= arr + 4) {\n __m128 xmm_norm = _mm_set1_ps(norm);\n _mm_storel_epi64(\n (__m128i *)arr,\n _mm_cvtps_ph(\n _mm_div_ps(_mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)arr)),\n xmm_norm),\n _MM_FROUND_NO_EXC));\n arr += 8;\n }\n switch (last - arr) {\n case 3:\n arr[2] = _cvtss_sh(_cvtsh_ss(arr[2]) / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 2:\n arr[1] = _cvtss_sh(_cvtsh_ss(arr[1]) / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 1:\n arr[0] = _cvtss_sh(_cvtsh_ss(arr[0]) / norm, _MM_FROUND_NO_EXC);\n }\n}\n#endif // __AVX512F__\n\nstatic inline void NormalizeAVX(uint16_t *arr, size_t dim, float norm) {\n uint16_t *last = arr + dim;\n uint16_t *last_aligned = arr + ((dim >> 4) << 4);\n\n __m256 ymm_norm = _mm256_set1_ps(norm);\n if (((uintptr_t)arr & 0xf) == 0) {\n for (; arr != last_aligned; arr += 16) {\n __m128i xmm_0 = _mm_load_si128((const __m128i *)(arr + 0));\n __m128i xmm_1 = _mm_load_si128((const __m128i *)(arr + 8));\n __m256 ymm_0 = _mm256_div_ps(_mm256_cvtph_ps(xmm_0), ymm_norm);\n __m256 ymm_1 = _mm256_div_ps(_mm256_cvtph_ps(xmm_1), ymm_norm);\n _mm_store_si128((__m128i *)(arr + 0),\n _mm256_cvtps_ph(ymm_0, _MM_FROUND_NO_EXC));\n _mm_store_si128((__m128i *)(arr + 8),\n _mm256_cvtps_ph(ymm_1, _MM_FROUND_NO_EXC));\n }\n if (last >= arr + 8) {\n _mm_store_si128(\n (__m128i *)arr,\n _mm256_cvtps_ph(_mm256_div_ps(_mm256_cvtph_ps(_mm_load_si128(\n (const __m128i *)arr)),\n ymm_norm),\n _MM_FROUND_NO_EXC));\n arr += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 16) {\n __m128i xmm_0 = _mm_loadu_si128((const __m128i *)(arr + 0));\n __m128i xmm_1 = _mm_loadu_si128((const __m128i *)(arr + 8));\n __m256 ymm_0 = _mm256_div_ps(_mm256_cvtph_ps(xmm_0), ymm_norm);\n __m256 ymm_1 = _mm256_div_ps(_mm256_cvtph_ps(xmm_1), ymm_norm);\n _mm_storeu_si128((__m128i *)(arr + 0),\n _mm256_cvtps_ph(ymm_0, _MM_FROUND_NO_EXC));\n _mm_storeu_si128((__m128i *)(arr + 8),\n _mm256_cvtps_ph(ymm_1, _MM_FROUND_NO_EXC));\n }\n if (last >= arr + 8) {\n _mm_storeu_si128(\n (__m128i *)arr,\n _mm256_cvtps_ph(_mm256_div_ps(_mm256_cvtph_ps(_mm_loadu_si128(\n (const __m128i *)arr)),\n ymm_norm),\n _MM_FROUND_NO_EXC));\n arr += 8;\n }\n }\n if (last >= arr + 4) {\n __m128 xmm_norm = _mm_set1_ps(norm);\n _mm_storel_epi64(\n (__m128i *)arr,\n _mm_cvtps_ph(\n _mm_div_ps(_mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)arr)),\n xmm_norm),\n _MM_FROUND_NO_EXC));\n arr += 8;\n }\n switch (last - arr) {\n case 3:\n arr[2] = _cvtss_sh(_cvtsh_ss(arr[2]) / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 2:\n arr[1] = _cvtss_sh(_cvtsh_ss(arr[1]) / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 1:\n arr[0] = _cvtss_sh(_cvtsh_ss(arr[0]) / norm, _MM_FROUND_NO_EXC);\n }\n}\n#endif // __AVX__ && __F16C__\n\n#if defined(__SSE__)\nstatic inline void NormalizeSSE(float *arr, size_t dim, float norm) {\n float *last = arr + dim;\n float *last_aligned = arr + ((dim >> 3) << 3);\n\n __m128 xmm_norm = _mm_set1_ps(norm);\n if (((uintptr_t)arr & 0xf) == 0) {\n for (; arr != last_aligned; arr += 8) {\n _mm_store_ps(arr + 0, _mm_div_ps(_mm_load_ps(arr + 0), xmm_norm));\n _mm_store_ps(arr + 4, _mm_div_ps(_mm_load_ps(arr + 4), xmm_norm));\n }\n if (last >= last_aligned + 4) {\n _mm_store_ps(arr, _mm_div_ps(_mm_load_ps(arr), xmm_norm));\n arr += 4;\n }\n } else {\n for (; arr != last_aligned; arr += 8) {\n _mm_storeu_ps(arr + 0, _mm_div_ps(_mm_loadu_ps(arr + 0), xmm_norm));\n _mm_storeu_ps(arr + 4, _mm_div_ps(_mm_loadu_ps(arr + 4), xmm_norm));\n }\n if (last >= last_aligned + 4) {\n _mm_storeu_ps(arr, _mm_div_ps(_mm_loadu_ps(arr), xmm_norm));\n arr += 4;\n }\n }\n switch (last - arr) {\n case 3:\n arr[2] /= norm;\n /* FALLTHRU */\n case 2:\n arr[1] /= norm;\n /* FALLTHRU */\n case 1:\n arr[0] /= norm;\n }\n}\n#endif // __SSE__\n\n#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))\n//! Compute the norm of vector\nvoid Normalizer::Compute(ValueType *arr, size_t dim, float norm) {\n#if defined(__ARM_NEON)\n NormalizeNEON(arr, dim, norm);\n#else\n#if defined(__AVX512F__)\n if (dim > 15) {\n NormalizeAVX512(arr, dim, norm);\n return;\n }\n#endif // __AVX512F__\n#if defined(__AVX__)\n if (dim > 7) {\n NormalizeAVX(arr, dim, norm);\n return;\n }\n#endif // __AVX__\n NormalizeSSE(arr, dim, norm);\n#endif // __ARM_NEON\n}\n#endif // __SSE__ || (__ARM_NEON && __aarch64__)\n\n#if (defined(__F16C__) && defined(__AVX__)) || \\\n (defined(__ARM_NEON) && defined(__aarch64__))\n//! Compute the norm of vector\nvoid Normalizer::Compute(ValueType *arr, size_t dim, float norm) {\n#if defined(__ARM_NEON)\n NormalizeNEON(reinterpret_cast(arr), dim, norm);\n#else\n#if defined(__AVX512F__)\n if (dim > 31) {\n NormalizeAVX512(reinterpret_cast(arr), dim, norm);\n return;\n }\n#endif // __AVX512F__\n NormalizeAVX(reinterpret_cast(arr), dim, norm);\n#endif // __ARM_NEON\n}\n#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__)\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/math/normalizer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \"norm_matrix.h\"\n\nnamespace zvec {\nnamespace ailego {\n\n/*! Normalizer\n */\ntemplate ::value>::type>\nstruct Normalizer {\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n //! Compute the norm of vector\n static inline void Compute(ValueType *arr, size_t dim, float norm) {\n for (size_t i = 0; i < dim; ++i) {\n arr[i] /= norm;\n }\n }\n\n //! Normalize a vector (L1)\n static inline void L1(ValueType *arr, size_t dim, float *norm) {\n Norm1Matrix::Compute(arr, dim, norm);\n if (*norm > 0.0f) {\n Compute(arr, dim, *norm);\n }\n }\n\n //! Normalize a vector (L2)\n static inline void L2(ValueType *arr, size_t dim, float *norm) {\n Norm2Matrix::Compute(arr, dim, norm);\n if (*norm > 0.0f) {\n Compute(arr, dim, *norm);\n }\n }\n};\n\n#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))\n/*! Normalizer (FP32)\n */\ntemplate <>\nstruct Normalizer {\n //! Type of value\n using ValueType = float;\n\n //! Compute the norm of vector\n static void Compute(ValueType *arr, size_t dim, float norm);\n\n //! Normalize a vector (L1)\n static inline void L1(ValueType *arr, size_t dim, float *norm) {\n Norm1Matrix::Compute(arr, dim, norm);\n if (*norm > 0.0f) {\n Compute(arr, dim, *norm);\n }\n }\n\n //! Normalize a vector (L2)\n static inline void L2(ValueType *arr, size_t dim, float *norm) {\n Norm2Matrix::Compute(arr, dim, norm);\n if (*norm > 0.0f) {\n Compute(arr, dim, *norm);\n }\n }\n};\n#endif // __SSE__ || (__ARM_NEON && __aarch64__)\n\n#if (defined(__F16C__) && defined(__AVX__)) || \\\n (defined(__ARM_NEON) && defined(__aarch64__))\n/*! Normalizer (FP16)\n */\ntemplate <>\nstruct Normalizer {\n //! Type of value\n using ValueType = Float16;\n\n //! Compute the norm of vector\n static void Compute(ValueType *arr, size_t dim, float norm);\n\n //! Normalize a vector (L1)\n static inline void L1(ValueType *arr, size_t dim, float *norm) {\n Norm1Matrix::Compute(arr, dim, norm);\n if (*norm > 0.0f) {\n Compute(arr, dim, *norm);\n }\n }\n\n //! Normalize a vector (L2)\n static inline void L2(ValueType *arr, size_t dim, float *norm) {\n Norm2Matrix::Compute(arr, dim, norm);\n if (*norm > 0.0f) {\n Compute(arr, dim, *norm);\n }\n }\n};\n#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__)\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/math_batch/cosine_distance_batch.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \"inner_product_distance_batch.h\"\n\nnamespace zvec::ailego::DistanceBatch {\n\n\ntemplate \nstruct CosineDistanceBatch;\n\ntemplate \nstruct CosineDistanceBatch {\n using ValueType = typename std::remove_cv::type;\n\n static inline void ComputeBatch(const ValueType **vecs,\n const ValueType *query, size_t num_vecs,\n size_t dim, float *results) {\n constexpr size_t extra_dim = sizeof(float) / sizeof(ValueType);\n size_t _dim = dim - extra_dim;\n\n InnerProductDistanceBatch::ComputeBatch(\n vecs, query, num_vecs, _dim, results);\n\n for (size_t i = 0; i < num_vecs; ++i) {\n results[i] = 1 - results[i];\n }\n }\n\n using IPImplType =\n InnerProductDistanceBatch;\n\n static void QueryPreprocess(void *query, size_t dim) {\n return IPImplType::QueryPreprocess(query,\n dim - sizeof(float) / sizeof(ValueType));\n }\n};\n\n\n} // namespace zvec::ailego::DistanceBatch" }, { "path": "src/ailego/math_batch/distance_batch.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \"ailego/math/distance_matrix.h\"\n#include \"cosine_distance_batch.h\"\n#include \"inner_product_distance_batch.h\"\n\nnamespace zvec::ailego {\n\ntemplate <\n template class DistanceType,\n typename ValueType, size_t BatchSize, size_t PrefetchStep, typename = void>\nstruct BaseDistance {\n static inline void _ComputeBatch(const ValueType **m, const ValueType *q,\n size_t num, size_t dim, float *out) {\n for (size_t i = 0; i < num; ++i) {\n DistanceType::Compute(m[i], q, dim, out + i);\n }\n }\n\n // If Distance has ComputeBatch, use it; otherwise fall back to _ComputeBatch.\n static inline void ComputeBatch(const ValueType **m, const ValueType *q,\n size_t num, size_t dim, float *out) {\n if constexpr (std::is_same_v,\n CosineDistanceMatrix>) {\n return DistanceBatch::CosineDistanceBatch<\n ValueType, BatchSize, PrefetchStep>::ComputeBatch(m, q, num, dim,\n out);\n }\n\n _ComputeBatch(m, q, num, dim, out);\n }\n};\n\n} // namespace zvec::ailego" }, { "path": "src/ailego/math_batch/inner_product_distance_batch.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec::ailego::DistanceBatch {\n\ntemplate \nstruct InnerProductDistanceBatch;\n\ntemplate \nstatic void compute_one_to_many_inner_product_fallback(\n const ValueType *query, const ValueType **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n for (size_t j = 0; j < BatchSize; ++j) {\n sums[j] = 0.0;\n InnerProductMatrix::Compute(ptrs[j], query, dim, sums + j);\n ailego_prefetch(&prefetch_ptrs[j]);\n }\n}\n\n// Function template partial specialization is not allowed,\n// therefore the wrapper struct is required.\ntemplate \nstruct InnerProductDistanceBatchImpl {\n using ValueType = typename std::remove_cv::type;\n static void compute_one_to_many(\n const ValueType *query, const ValueType **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n return compute_one_to_many_inner_product_fallback(query, ptrs,\n prefetch_ptrs, dim, sums);\n }\n static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc() {\n return nullptr;\n }\n};\n\ntemplate \nstruct InnerProductDistanceBatch {\n using ValueType = typename std::remove_cv::type;\n\n static inline void ComputeBatch(const ValueType **vecs,\n const ValueType *query, size_t num_vecs,\n size_t dim, float *results) {\n size_t i = 0;\n for (; i + BatchSize <= num_vecs; i += BatchSize) {\n std::array prefetch_ptrs;\n for (size_t j = 0; j < BatchSize; ++j) {\n if (i + j + BatchSize * PrefetchStep < num_vecs) {\n prefetch_ptrs[j] = vecs[i + j + BatchSize * PrefetchStep];\n } else {\n prefetch_ptrs[j] = nullptr;\n }\n }\n InnerProductDistanceBatchImpl::compute_one_to_many(\n query, &vecs[i], prefetch_ptrs, dim, &results[i]);\n }\n for (; i < num_vecs; ++i) { // TODO: unroll by 1, 2, 4, 8, etc.\n std::array prefetch_ptrs{nullptr};\n InnerProductDistanceBatchImpl::compute_one_to_many(\n query, &vecs[i], prefetch_ptrs, dim, &results[i]);\n }\n }\n\n static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc() {\n return InnerProductDistanceBatchImpl::GetQueryPreprocessFunc();\n }\n};\n\ntemplate <>\nstruct InnerProductDistanceBatchImpl {\n using ValueType = ailego::Float16;\n static void compute_one_to_many(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums);\n};\n\ntemplate <>\nstruct InnerProductDistanceBatchImpl {\n using ValueType = float;\n static void compute_one_to_many(const float *query, const float **ptrs,\n std::array &prefetch_ptrs,\n size_t dim, float *sums);\n};\n\ntemplate <>\nstruct InnerProductDistanceBatchImpl {\n using ValueType = int8_t;\n static void compute_one_to_many(const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs,\n size_t dim, float *sums);\n\n static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc();\n};\n\ntemplate <>\nstruct InnerProductDistanceBatchImpl {\n using ValueType = ailego::Float16;\n static void compute_one_to_many(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums);\n};\n\ntemplate <>\nstruct InnerProductDistanceBatchImpl {\n using ValueType = float;\n static void compute_one_to_many(const float *query, const float **ptrs,\n std::array &prefetch_ptrs,\n size_t dim, float *sums);\n};\n\ntemplate <>\nstruct InnerProductDistanceBatchImpl {\n using ValueType = int8_t;\n static void compute_one_to_many(const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs,\n size_t dim, float *sums);\n};\n\n} // namespace zvec::ailego::DistanceBatch\n" }, { "path": "src/ailego/math_batch/inner_product_distance_batch_dispatch.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"inner_product_distance_batch.h\"\n\nnamespace zvec::ailego::DistanceBatch {\n\n#if defined(__AVX512VNNI__)\nvoid compute_one_to_many_inner_product_avx512_vnni_int8_query_preprocess(\n void *query, size_t dim);\n\nvoid compute_one_to_many_inner_product_avx512_vnni_int8_1(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dimensionality,\n float *results);\n\nvoid compute_one_to_many_inner_product_avx512_vnni_int8_12(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dimensionality,\n float *results);\n#endif\n\n#if defined(__AVX512FP16__)\nvoid compute_one_to_many_inner_product_avx512fp16_fp16_1(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results);\n\nvoid compute_one_to_many_inner_product_avx512fp16_fp16_12(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results);\n#endif //__AVX512FP16__\n\n#if defined(__AVX512F__)\nvoid compute_one_to_many_inner_product_avx512f_fp16_1(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results);\n\nvoid compute_one_to_many_inner_product_avx512f_fp16_12(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results);\n#endif //__AVX512F__\n\n#if defined(__AVX2__)\nvoid compute_one_to_many_inner_product_avx2_fp32_1(\n const float *query, const float **ptrs,\n std::array &prefetch_ptrs, size_t dimensionality,\n float *results);\n\nvoid compute_one_to_many_inner_product_avx2_fp16_1(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results);\n\nvoid compute_one_to_many_inner_product_avx2_int8_1(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dimensionality,\n float *results);\n\nvoid compute_one_to_many_inner_product_avx2_fp32_12(\n const float *query, const float **ptrs,\n std::array &prefetch_ptrs, size_t dimensionality,\n float *results);\n\nvoid compute_one_to_many_inner_product_avx2_fp16_12(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results);\n\nvoid compute_one_to_many_inner_product_avx2_int8_12(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dimensionality,\n float *results);\n#endif\n\nvoid InnerProductDistanceBatchImpl::compute_one_to_many(\n const ValueType *query, const ValueType **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n return compute_one_to_many_inner_product_avx2_fp32_1(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs,\n dim, sums);\n}\n\nvoid InnerProductDistanceBatchImpl::compute_one_to_many(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n#if defined(__AVX512FP16__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) {\n return compute_one_to_many_inner_product_avx512fp16_fp16_1(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n return compute_one_to_many_inner_product_avx512f_fp16_1(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n return compute_one_to_many_inner_product_avx2_fp16_1(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs,\n dim, sums);\n}\n\nvoid InnerProductDistanceBatchImpl::compute_one_to_many(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n// #if defined(__AVX512BW__) // TODO: this version is problematic\n// return compute_one_to_many_avx512_int8(\n// query, ptrs, prefetch_ptrs, dim, sums);\n#if defined(__AVX512VNNI__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) {\n return compute_one_to_many_inner_product_avx512_vnni_int8_1(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n return compute_one_to_many_inner_product_avx2_int8_1(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs,\n dim, sums);\n}\n\nDistanceBatchQueryPreprocessFunc\nInnerProductDistanceBatchImpl::GetQueryPreprocessFunc() {\n#if defined(__AVX512VNNI__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) {\n return compute_one_to_many_inner_product_avx512_vnni_int8_query_preprocess;\n }\n#endif\n return nullptr;\n}\n\nvoid InnerProductDistanceBatchImpl::compute_one_to_many(\n const ValueType *query, const ValueType **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n return compute_one_to_many_inner_product_avx2_fp32_12(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs,\n dim, sums);\n}\n\nvoid InnerProductDistanceBatchImpl::compute_one_to_many(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n#if defined(__AVX512FP16__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) {\n return compute_one_to_many_inner_product_avx512fp16_fp16_12(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n#if defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n return compute_one_to_many_inner_product_avx512f_fp16_12(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n return compute_one_to_many_inner_product_avx2_fp16_12(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs,\n dim, sums);\n}\n\nvoid InnerProductDistanceBatchImpl::compute_one_to_many(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n// #if defined(__AVX512BW__) // TODO: this version is problematic\n// return compute_one_to_many_avx512_int8(\n// query, ptrs, prefetch_ptrs, dim, sums);\n#if defined(__AVX512VNNI__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) {\n return compute_one_to_many_inner_product_avx512_vnni_int8_12(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n#if defined(__AVX2__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {\n return compute_one_to_many_inner_product_avx2_int8_12(\n query, ptrs, prefetch_ptrs, dim, sums);\n }\n#endif\n return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs,\n dim, sums);\n}\n\n} // namespace zvec::ailego::DistanceBatch\n" }, { "path": "src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx2.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec::ailego::DistanceBatch {\n\n#if defined(__AVX2__)\n\ntemplate \nstatic std::enable_if_t, void>\ncompute_one_to_many_inner_product_avx2_fp16(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results) {\n __m256 accs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm256_setzero_ps();\n }\n\n size_t dim = 0;\n for (; dim + 16 <= dimensionality; dim += 16) {\n __m256i q =\n _mm256_loadu_si256(reinterpret_cast(query + dim));\n\n __m256 q1 = _mm256_cvtph_ps(_mm256_castsi256_si128(q));\n __m256 q2 = _mm256_cvtph_ps(_mm256_extractf128_si256(q, 1));\n\n __m256 data_regs_1[dp_batch];\n __m256 data_regs_2[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n __m256i m =\n _mm256_loadu_si256(reinterpret_cast(ptrs[i] + dim));\n\n data_regs_1[i] = _mm256_cvtph_ps(_mm256_castsi256_si128(m));\n data_regs_2[i] = _mm256_cvtph_ps(_mm256_extractf128_si256(m, 1));\n }\n\n if (prefetch_ptrs[0]) {\n for (size_t i = 0; i < dp_batch; ++i) {\n ailego_prefetch(prefetch_ptrs[i] + dim);\n }\n }\n\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm256_fmadd_ps(q1, data_regs_1[i], accs[i]);\n accs[i] = _mm256_fmadd_ps(q2, data_regs_2[i], accs[i]);\n }\n }\n\n if (dim + 8 <= dimensionality) {\n __m256 q = _mm256_cvtph_ps(\n _mm_loadu_si128(reinterpret_cast(query + dim)));\n\n __m256 data_regs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n data_regs[i] = _mm256_cvtph_ps(\n _mm_loadu_si128(reinterpret_cast(ptrs[i] + dim)));\n accs[i] = _mm256_fmadd_ps(q, data_regs[i], accs[i]);\n }\n\n dim += 8;\n }\n\n for (size_t i = 0; i < dp_batch; ++i) {\n results[i] = HorizontalAdd_FP32_V256(accs[i]);\n }\n\n for (; dim < dimensionality; ++dim) {\n for (size_t i = 0; i < dp_batch; ++i) {\n results[i] += (*(query + dim)) * (*(ptrs[i] + dim));\n }\n }\n}\n\nvoid compute_one_to_many_inner_product_avx2_fp16_1(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n return compute_one_to_many_inner_product_avx2_fp16(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\nvoid compute_one_to_many_inner_product_avx2_fp16_12(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n return compute_one_to_many_inner_product_avx2_fp16(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\n#endif\n\n} // namespace zvec::ailego::DistanceBatch" }, { "path": "src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec::ailego::DistanceBatch {\n\n#if defined(__AVX512F__)\n\ntemplate \nstatic std::enable_if_t, void>\ncompute_one_to_many_inner_product_avx512f_fp16(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results) {\n __m512 accs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm512_setzero_ps();\n }\n\n size_t dim = 0;\n for (; dim + 32 <= dimensionality; dim += 32) {\n __m512i q =\n _mm512_loadu_si512(reinterpret_cast(query + dim));\n\n __m512 q1 = _mm512_cvtph_ps(_mm512_castsi512_si256(q));\n __m512 q2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(q, 1));\n\n __m512 data_regs_1[dp_batch];\n __m512 data_regs_2[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n __m512i m =\n _mm512_loadu_si512(reinterpret_cast(ptrs[i] + dim));\n\n data_regs_1[i] = _mm512_cvtph_ps(_mm512_castsi512_si256(m));\n data_regs_2[i] = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(m, 1));\n }\n\n if (prefetch_ptrs[0]) {\n for (size_t i = 0; i < dp_batch; ++i) {\n ailego_prefetch(prefetch_ptrs[i] + dim);\n }\n }\n\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm512_fmadd_ps(q1, data_regs_1[i], accs[i]);\n accs[i] = _mm512_fmadd_ps(q2, data_regs_2[i], accs[i]);\n }\n }\n\n if (dim + 16 <= dimensionality) {\n __m512 q = _mm512_cvtph_ps(\n _mm256_loadu_si256(reinterpret_cast(query + dim)));\n\n __m512 data_regs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n data_regs[i] = _mm512_cvtph_ps(\n _mm256_loadu_si256(reinterpret_cast(ptrs[i] + dim)));\n accs[i] = _mm512_fmadd_ps(q, data_regs[i], accs[i]);\n }\n\n dim += 16;\n }\n\n __m256 acc_new[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n acc_new[i] = _mm256_add_ps(\n _mm512_castps512_ps256(accs[i]),\n _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(accs[i]), 1)));\n }\n\n if (dim + 8 <= dimensionality) {\n __m256 q = _mm256_cvtph_ps(\n _mm_loadu_si128(reinterpret_cast(query + dim)));\n\n for (size_t i = 0; i < dp_batch; ++i) {\n __m256 m = _mm256_cvtph_ps(\n _mm_loadu_si128(reinterpret_cast(ptrs[i] + dim)));\n acc_new[i] = _mm256_fmadd_ps(m, q, acc_new[i]);\n }\n\n dim += 8;\n }\n\n for (size_t i = 0; i < dp_batch; ++i) {\n results[i] = HorizontalAdd_FP32_V256(acc_new[i]);\n }\n\n for (; dim < dimensionality; ++dim) {\n for (size_t i = 0; i < dp_batch; ++i) {\n results[i] += (*(query + dim)) * (*(ptrs[i] + dim));\n }\n }\n}\n\nvoid compute_one_to_many_inner_product_avx512f_fp16_1(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n return compute_one_to_many_inner_product_avx512f_fp16(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\nvoid compute_one_to_many_inner_product_avx512f_fp16_12(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n return compute_one_to_many_inner_product_avx512f_fp16(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n#endif\n\n} // namespace zvec::ailego::DistanceBatch\n" }, { "path": "src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512fp16.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec::ailego::DistanceBatch {\n\n#if defined(__AVX512FP16__)\ntemplate \nstatic std::enable_if_t, void>\ncompute_one_to_many_inner_product_avx512fp16_fp16(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results) {\n __m512h accs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm512_setzero_ph();\n }\n\n size_t dim = 0;\n for (; dim + 32 <= dimensionality; dim += 32) {\n __m512h q = _mm512_loadu_ph(query + dim);\n\n __m512h data_regs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n data_regs[i] = _mm512_loadu_ph(ptrs[i] + dim);\n }\n\n if (prefetch_ptrs[0]) {\n for (size_t i = 0; i < dp_batch; ++i) {\n ailego_prefetch(prefetch_ptrs[i] + dim);\n }\n }\n\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm512_fmadd_ph(data_regs[i], q, accs[i]);\n }\n }\n\n if (dim < dimensionality) {\n __mmask32 mask = (__mmask32)((1 << (dimensionality - dim)) - 1);\n\n for (size_t i = 0; i < dp_batch; ++i) {\n __m512i zmm_undefined = _mm512_undefined_epi32();\n\n accs[i] =\n _mm512_mask3_fmadd_ph(_mm512_castsi512_ph(_mm512_mask_loadu_epi16(\n zmm_undefined, mask, query + dim)),\n _mm512_castsi512_ph(_mm512_mask_loadu_epi16(\n zmm_undefined, mask, ptrs[i] + dim)),\n accs[i], mask);\n }\n }\n\n for (size_t i = 0; i < dp_batch; ++i) {\n results[i] = HorizontalAdd_FP16_V512(accs[i]);\n }\n}\n\nvoid compute_one_to_many_inner_product_avx512fp16_fp16_1(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n return compute_one_to_many_inner_product_avx512fp16_fp16(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\nvoid compute_one_to_many_inner_product_avx512fp16_fp16_12(\n const ailego::Float16 *query, const ailego::Float16 **ptrs,\n std::array &prefetch_ptrs, size_t dim,\n float *sums) {\n return compute_one_to_many_inner_product_avx512fp16_fp16(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n#endif\n\n} // namespace zvec::ailego::DistanceBatch\n" }, { "path": "src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec::ailego::DistanceBatch {\n\n#if defined(__AVX2__)\n\ninline float sum4(__m128 v) {\n v = _mm_add_ps(v, _mm_castsi128_ps(_mm_srli_si128(_mm_castps_si128(v), 8)));\n return v[0] + v[1];\n}\n\ninline __m128 sum_top_bottom_avx(__m256 v) {\n const __m128 high = _mm256_extractf128_ps(v, 1);\n const __m128 low = _mm256_castps256_ps128(v);\n return _mm_add_ps(high, low);\n}\n\ntemplate \nstatic std::enable_if_t, void>\ncompute_one_to_many_inner_product_avx2_fp32(\n const ValueType *query, const ValueType **ptrs,\n std::array &prefetch_ptrs,\n size_t dimensionality, float *results) {\n __m256 accs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm256_setzero_ps();\n }\n size_t dim = 0;\n for (; dim + 8 <= dimensionality; dim += 8) {\n __m256 q = _mm256_loadu_ps(query + dim);\n\n __m256 data_regs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n data_regs[i] = _mm256_loadu_ps(ptrs[i] + dim);\n }\n if (prefetch_ptrs[0]) {\n for (size_t i = 0; i < dp_batch; ++i) {\n ailego_prefetch(prefetch_ptrs[i] + dim);\n }\n }\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm256_fnmadd_ps(q, data_regs[i], accs[i]);\n }\n }\n\n __m128 sum128_regs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n sum128_regs[i] = sum_top_bottom_avx(accs[i]);\n }\n if (dim + 4 <= dimensionality) {\n __m128 q = _mm_loadu_ps(query + dim);\n\n __m128 data_regs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n data_regs[i] = _mm_loadu_ps(ptrs[i] + dim);\n }\n if (prefetch_ptrs[0]) {\n for (size_t i = 0; i < dp_batch; ++i) {\n ailego_prefetch(prefetch_ptrs[i] + dim);\n }\n }\n for (size_t i = 0; i < dp_batch; ++i) {\n sum128_regs[i] = _mm_fnmadd_ps(q, data_regs[i], sum128_regs[i]);\n }\n dim += 4;\n }\n if (dim + 2 <= dimensionality) {\n __m128 q = _mm_setzero_ps();\n\n __m128 data_regs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n data_regs[i] = _mm_setzero_ps();\n }\n\n q = _mm_loadh_pi(q, (const __m64 *)(query + dim));\n for (size_t i = 0; i < dp_batch; ++i) {\n data_regs[i] = _mm_loadh_pi(data_regs[i], (const __m64 *)(ptrs[i] + dim));\n }\n for (size_t i = 0; i < dp_batch; ++i) {\n sum128_regs[i] = _mm_fnmadd_ps(q, data_regs[i], sum128_regs[i]);\n }\n dim += 2;\n }\n\n float res[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n res[i] = sum4(sum128_regs[i]);\n }\n if (dim < dimensionality) {\n float q = query[dim];\n for (size_t i = 0; i < dp_batch; ++i) {\n res[i] -= q * ptrs[i][dim];\n }\n }\n for (size_t i = 0; i < dp_batch; ++i) {\n results[i] = -res[i];\n }\n}\n\nvoid compute_one_to_many_inner_product_avx2_fp32_1(\n const float *query, const float **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n return compute_one_to_many_inner_product_avx2_fp32(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\nvoid compute_one_to_many_inner_product_avx2_fp32_12(\n const float *query, const float **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n return compute_one_to_many_inner_product_avx2_fp32(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\n#endif\n\n} // namespace zvec::ailego::DistanceBatch" }, { "path": "src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx2.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n\nnamespace zvec::ailego::DistanceBatch {\n\n#if defined(__AVX2__)\n\ntemplate \nstatic std::enable_if_t, void>\ncompute_one_to_many_inner_product_avx2_int8(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dimensionality,\n float *results) {\n __m256i accs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm256_setzero_si256();\n }\n size_t dim = 0;\n for (; dim + 32 <= dimensionality; dim += 32) {\n __m256i q = _mm256_loadu_si256((const __m256i *)(query + dim));\n\n __m256i data_regs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n data_regs[i] = _mm256_loadu_si256((const __m256i *)(ptrs[i] + dim));\n }\n if (prefetch_ptrs[0]) {\n for (size_t i = 0; i < dp_batch; ++i) {\n ailego_prefetch(prefetch_ptrs[i] + dim);\n }\n }\n __m256i q_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q));\n __m256i q_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q, 1));\n __m256i data_lo[dp_batch];\n __m256i data_hi[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n data_lo[i] = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(data_regs[i]));\n data_hi[i] =\n _mm256_cvtepi8_epi16(_mm256_extracti128_si256(data_regs[i], 1));\n }\n __m256i prod_lo[dp_batch];\n __m256i prod_hi[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n prod_lo[i] = _mm256_madd_epi16(q_lo, data_lo[i]);\n prod_hi[i] = _mm256_madd_epi16(q_hi, data_hi[i]);\n }\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] =\n _mm256_add_epi32(accs[i], _mm256_add_epi32(prod_lo[i], prod_hi[i]));\n }\n }\n\n int temp_results[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n __m128i lo = _mm256_castsi256_si128(accs[i]);\n __m128i hi = _mm256_extracti128_si256(accs[i], 1);\n __m128i sum128 = _mm_add_epi32(lo, hi);\n sum128 = _mm_hadd_epi32(sum128, sum128);\n sum128 = _mm_hadd_epi32(sum128, sum128);\n temp_results[i] = _mm_cvtsi128_si32(sum128);\n }\n for (; dim < dimensionality; ++dim) {\n int8_t q = query[dim];\n for (size_t i = 0; i < dp_batch; ++i) {\n temp_results[i] += q * static_cast(ptrs[i][dim]);\n }\n }\n for (size_t i = 0; i < dp_batch; ++i) {\n results[i] = static_cast(temp_results[i]);\n }\n}\n\nvoid compute_one_to_many_inner_product_avx2_int8_1(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n return compute_one_to_many_inner_product_avx2_int8(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\nvoid compute_one_to_many_inner_product_avx2_int8_12(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n return compute_one_to_many_inner_product_avx2_int8(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\n#endif\n\n} // namespace zvec::ailego::DistanceBatch" }, { "path": "src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512fp16.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n\nnamespace zvec::ailego::DistanceBatch {\n\n#if defined(__AVX512VNNI__)\n\nvoid compute_one_to_many_inner_product_avx512_vnni_int8_query_preprocess(\n void *query, size_t dim) {\n const int8_t *input = reinterpret_cast(query);\n uint8_t *output = reinterpret_cast(query);\n\n // // AVX512 constant: 128 in each byte (cast to int8_t, which becomes -128\n // // in signed representation, but addition works correctly due to two's\n // // complement arithmetic)\n const __m512i offset = _mm512_set1_epi8(static_cast(128));\n //\n size_t i = 0;\n // // Process 64 bytes at a time using AVX512\n for (; i + 64 <= dim; i += 64) {\n __m512i data =\n _mm512_loadu_si512(reinterpret_cast(input + i));\n __m512i result = _mm512_add_epi8(data, offset);\n _mm512_storeu_si512(reinterpret_cast<__m512i *>(output + i), result);\n }\n\n // Handle remaining elements with scalar loop\n for (; i < dim; ++i) {\n output[i] = static_cast(static_cast(input[i]) + 128);\n }\n}\n\n// query is unsigned\ntemplate \nstatic void compute_one_to_many_inner_product_avx512_vnni_int8(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dimensionality,\n float *results) {\n __m512i accs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm512_setzero_si512();\n }\n size_t dim = 0;\n for (; dim + 64 <= dimensionality; dim += 64) {\n __m512i q =\n _mm512_loadu_si512(reinterpret_cast(query + dim));\n\n __m512i data_regs[dp_batch];\n for (size_t i = 0; i < dp_batch; ++i) {\n data_regs[i] =\n _mm512_loadu_si512(reinterpret_cast(ptrs[i] + dim));\n }\n if (prefetch_ptrs[0]) {\n for (size_t i = 0; i < dp_batch; ++i) {\n ailego_prefetch(prefetch_ptrs[i] + dim);\n }\n }\n for (size_t i = 0; i < dp_batch; ++i) {\n accs[i] = _mm512_dpbusd_epi32(accs[i], q, data_regs[i]);\n }\n }\n\n int temp_results[dp_batch]{};\n for (size_t i = 0; i < dp_batch; ++i) {\n temp_results[i] = _mm512_reduce_add_epi32(accs[i]);\n }\n for (; dim < dimensionality; ++dim) {\n uint q = reinterpret_cast(query)[dim];\n for (size_t i = 0; i < dp_batch; ++i) {\n temp_results[i] += q * static_cast(ptrs[i][dim]);\n }\n }\n for (size_t i = 0; i < dp_batch; ++i) {\n results[i] = static_cast(temp_results[i]);\n }\n}\n\n//\n// #elif defined(__AVX512BW__)\n// // TODO: this version is problematic\n// template \n// static std::enable_if_t, void>\n// compute_one_to_many_avx512_int8(\n// const int8_t *query, const int8_t **ptrs,\n// std::array &prefetch_ptrs, size_t\n// dimensionality, float *results) {\n// std::array<__m512i, dp_batch> accs;\n// size_t dim = 0;\n// for (; dim + 64 <= dimensionality; dim += 64) {\n// __m512i q =\n// _mm512_loadu_si512(reinterpret_cast(query + dim));\n// std::array<__m512i, dp_batch> data_regs;\n// for (size_t i = 0; i < dp_batch; ++i) {\n// data_regs[i] =\n// _mm512_loadu_si512(reinterpret_cast(ptrs[i] +\n// dim));\n// }\n// if (prefetch_ptrs[0]) {\n// for (size_t i = 0; i < dp_batch; ++i) {\n// ailego_prefetch(prefetch_ptrs[i] + dim);\n// }\n// }\n// __m512i q_lo = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(q, 0));\n// __m512i q_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(q, 1));\n// std::array<__m512i, dp_batch> data_lo;\n// std::array<__m512i, dp_batch> data_hi;\n// for (size_t i = 0; i < dp_batch; ++i) {\n// data_lo[i] =\n// _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(data_regs[i], 0));\n// data_hi[i] =\n// _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(data_regs[i], 1));\n// }\n// std::array<__m512i, dp_batch> prod_lo;\n// std::array<__m512i, dp_batch> prod_hi;\n// for (size_t i = 0; i < dp_batch; ++i) {\n// prod_lo[i] = _mm512_madd_epi16(q_lo, data_lo[i]);\n// prod_hi[i] = _mm512_madd_epi16(q_hi, data_hi[i]);\n// }\n// for (size_t i = 0; i < dp_batch; ++i) {\n// accs[i] = _mm512_add_epi32(\n// accs[i], _mm512_add_epi32(\n// _mm512_madd_epi16(prod_lo[i], _mm512_set1_epi16(1)),\n// _mm512_madd_epi16(prod_hi[i], _mm512_set1_epi16(1))));\n// }\n// }\n// std::array temp_results;\n// for (size_t i = 0; i < dp_batch; ++i) {\n// temp_results[i] = _mm512_reduce_add_epi32(accs[i]);\n// }\n// for (; dim < dimensionality; ++dim) {\n// int8_t q = query[dim];\n// for (size_t i = 0; i < dp_batch; ++i) {\n// temp_results[i] += q * static_cast(ptrs[i][dim]);\n// }\n// }\n// for (size_t i = 0; i < dp_batch; ++i) {\n// results[i] = static_cast(temp_results[i]);\n// }\n// }\n\nvoid compute_one_to_many_inner_product_avx512_vnni_int8_1(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n return compute_one_to_many_inner_product_avx512_vnni_int8<1>(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\nvoid compute_one_to_many_inner_product_avx512_vnni_int8_12(\n const int8_t *query, const int8_t **ptrs,\n std::array &prefetch_ptrs, size_t dim, float *sums) {\n return compute_one_to_many_inner_product_avx512_vnni_int8<12>(\n query, ptrs, prefetch_ptrs, dim, sums);\n}\n\n\n#endif\n\n} // namespace zvec::ailego::DistanceBatch" }, { "path": "src/ailego/parallel/lock.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#if __cplusplus >= 201703L\n#include \n#endif\n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n// Test if atomic_bool is always lock free.\n// Arm may be always lock free using some compiler flags,\n// see https://stackoverflow.com/a/64253858/486350.\n#if ATOMIC_BOOL_LOCK_FREE == 2\n\n/*! Spin Mutex (The atomic type is always lock-free)\n */\nclass SpinMutex {\n public:\n //! Constructor\n SpinMutex(void) {}\n\n //! Locking\n void lock(void) {\n bool expected = false;\n while (!flag_.compare_exchange_weak(\n expected, true, std::memory_order_acquire, std::memory_order_relaxed)) {\n expected = false;\n // Provide a hint to the processor that the code sequence is a spin-wait\n // loop. This can help improve the performance and power consumption of\n // spin-wait loops.\n ailego_yield();\n }\n }\n\n //! Try locking\n bool try_lock(void) {\n bool expected = false;\n return flag_.compare_exchange_strong(\n expected, true, std::memory_order_acquire, std::memory_order_relaxed);\n }\n\n //! Unlocking\n void unlock(void) {\n flag_.store(false, std::memory_order_release);\n }\n\n private:\n //! Disable them\n SpinMutex(const SpinMutex &) = delete;\n SpinMutex(SpinMutex &&) = delete;\n SpinMutex &operator=(const SpinMutex &) = delete;\n SpinMutex &operator=(SpinMutex &&) = delete;\n\n //! Members\n std::atomic_bool flag_{false};\n};\n#else\n\n/*! Spin Mutex (General)\n */\nclass SpinMutex {\n public:\n //! Constructor\n SpinMutex(void) {}\n\n //! Locking\n void lock(void) {\n while (flag_.test_and_set(std::memory_order_acquire));\n }\n\n //! Try locking\n bool try_lock(void) {\n return (!flag_.test_and_set(std::memory_order_acquire));\n }\n\n //! Unlocking\n void unlock(void) {\n flag_.clear(std::memory_order_release);\n }\n\n private:\n //! Disable them\n SpinMutex(const SpinMutex &) = delete;\n SpinMutex(SpinMutex &&) = delete;\n SpinMutex &operator=(const SpinMutex &) = delete;\n SpinMutex &operator=(SpinMutex &&) = delete;\n\n //! Members\n std::atomic_flag flag_{};\n};\n#endif // ATOMIC_BOOL_LOCK_FREE == 2\n\n#if __cplusplus >= 201703L\n\nusing SharedMutex = std::shared_mutex;\n\n#else\n\n/*! Shared Mutex\n */\nclass SharedMutex {\n public:\n //! Constructor\n SharedMutex(void) {}\n\n //! Locking\n void lock(void) {\n std::unique_lock q(mutex_);\n ++write_count_;\n write_cond_.wait(q, [this]() { return (pending_count_ == 0); });\n --write_count_;\n --pending_count_;\n }\n\n //! Try locking\n bool try_lock(void) {\n std::unique_lock q(mutex_, std::defer_lock);\n if (q.try_lock()) {\n if (pending_count_ == 0) {\n --pending_count_;\n return true;\n }\n }\n return false;\n }\n\n //! Unlocking\n void unlock(void) {\n std::lock_guard q(mutex_);\n ++pending_count_;\n\n if (write_count_ != 0) {\n write_cond_.notify_one();\n } else {\n read_cond_.notify_all();\n }\n }\n\n //! Locking (shared)\n void lock_shared(void) {\n std::unique_lock q(mutex_);\n ++read_count_;\n read_cond_.wait(\n q, [this]() { return (write_count_ == 0 && pending_count_ >= 0); });\n --read_count_;\n ++pending_count_;\n }\n\n //! Try locking (shared)\n bool try_lock_shared(void) {\n std::lock_guard q(mutex_);\n if (write_count_ == 0 && pending_count_ >= 0) {\n ++pending_count_;\n return true;\n }\n return false;\n }\n\n //! Unlocking (shared)\n void unlock_shared(void) {\n std::lock_guard q(mutex_);\n --pending_count_;\n\n if (write_count_ != 0 && pending_count_ == 0) {\n write_cond_.notify_one();\n } else {\n read_cond_.notify_all();\n }\n }\n\n private:\n //! Disable them\n SharedMutex(const SharedMutex &) = delete;\n SharedMutex(SharedMutex &&) = delete;\n SharedMutex &operator=(const SharedMutex &) = delete;\n SharedMutex &operator=(SharedMutex &&) = delete;\n\n //! Members\n int32_t pending_count_{0};\n int32_t read_count_{0};\n int32_t write_count_{0};\n std::mutex mutex_{};\n std::condition_variable read_cond_{};\n std::condition_variable write_cond_{};\n};\n\n#endif // __cplusplus >= 201703L\n\n/*! Write Lock\n */\nclass WriteLock {\n public:\n //! Constructor\n WriteLock(SharedMutex &mutex) : mutex_(mutex) {}\n\n //! Locking\n void lock(void) {\n mutex_.lock();\n }\n\n //! Try locking\n bool try_lock(void) {\n return mutex_.try_lock();\n }\n\n //! Unlocking\n void unlock(void) {\n mutex_.unlock();\n }\n\n private:\n //! Disable them\n WriteLock(void) = delete;\n WriteLock(const WriteLock &) = delete;\n WriteLock(WriteLock &&) = delete;\n WriteLock &operator=(const WriteLock &) = delete;\n WriteLock &operator=(WriteLock &&) = delete;\n\n //! Members\n SharedMutex &mutex_;\n};\n\n/*! Read Lock\n */\nclass ReadLock {\n public:\n //! Constructor\n ReadLock(SharedMutex &mutex) : mutex_(mutex) {}\n\n //! Locking\n void lock(void) {\n mutex_.lock_shared();\n }\n\n //! Try locking\n bool try_lock(void) {\n return mutex_.try_lock_shared();\n }\n\n //! Unlocking\n void unlock(void) {\n mutex_.unlock_shared();\n }\n\n private:\n //! Disable them\n ReadLock(void) = delete;\n ReadLock(const ReadLock &) = delete;\n ReadLock(ReadLock &&) = delete;\n ReadLock &operator=(const ReadLock &) = delete;\n ReadLock &operator=(ReadLock &&) = delete;\n\n //! Members\n SharedMutex &mutex_;\n};\n\n/*\n Atomic Close Lock\n */\n\n#define AILEGO_SAFE_ACCESS(CLOSE_ERR) \\\n counter_.fetch_add(1); \\\n AILEGO_DEFER([this] { counter_.fetch_sub(1); }); \\\n if (!opened_.load()) { \\\n return CLOSE_ERR; \\\n }\n\n#define AILEGO_SAFE_CLOSE \\\n opened_.store(false); \\\n while (counter_.load() > 0) { \\\n std::this_thread::sleep_for(std::chrono::milliseconds(1)); \\\n }\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/parallel/multi_thread_list.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Multi-Thread list\n */\ntemplate \nclass MultiThreadList {\n public:\n MultiThreadList(size_t size_limit = 1000) : size_limit_(size_limit) {}\n bool produce(const T &item) {\n std::unique_lock lk(lock_);\n not_full_.wait(\n lk, [&]() { return (list_.size() < size_limit_) || done_.load(); });\n if (done_.load()) {\n return false;\n }\n list_.emplace_back(item);\n not_empty_.notify_one();\n return true;\n }\n\n bool produce(T &&item) {\n std::unique_lock lk(lock_);\n not_full_.wait(\n lk, [&]() { return (list_.size() < size_limit_) || done_.load(); });\n if (done_.load()) {\n return false;\n }\n list_.emplace_back(std::move(item));\n not_empty_.notify_one();\n return true;\n }\n\n bool consume(T *item) {\n std::unique_lock lk(lock_);\n not_empty_.wait(lk, [&]() {\n return !list_.empty() || done_.load() || consume_stopped_.load();\n });\n if ((list_.empty() && done_.load()) || consume_stopped_.load()) {\n return false;\n }\n *item = std::move(list_.front());\n list_.pop_front();\n not_full_.notify_one();\n return true;\n }\n\n void done() {\n std::unique_lock lk(lock_);\n done_.store(true);\n not_empty_.notify_all();\n not_full_.notify_all();\n }\n\n void reset() {\n done_.store(false);\n list_.clear();\n }\n\n void stop_consume() {\n std::unique_lock lk(lock_);\n consume_stopped_.store(true);\n not_empty_.notify_all();\n }\n\n void resume_consume() {\n consume_stopped_.store(false);\n }\n\n private:\n std::deque list_;\n size_t size_limit_{0};\n std::mutex lock_;\n std::condition_variable not_empty_, not_full_;\n\n std::atomic done_{false};\n std::atomic consume_stopped_{false};\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/parallel/semaphore.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Semaphore\n */\nclass Semaphore {\n public:\n //! Constructor\n Semaphore(void) : Semaphore{1} {}\n\n //! Constructor\n Semaphore(uint32_t count) : count_(count) {}\n\n //! Acquire a permit from this semaphore, suspending until one is available\n void lock(void) {\n while (!this->try_lock()) {\n std::unique_lock latch(mutex_);\n cond_.wait(latch, [this]() { return (count_ > 0); });\n }\n }\n\n //! Try to acquire a permit from this semaphore without suspension\n bool try_lock(void) {\n uint32_t count = count_.load(std::memory_order_acquire);\n return (count > 0 ? count_.compare_exchange_strong(\n count, count - 1, std::memory_order_release,\n std::memory_order_relaxed)\n : false);\n }\n\n //! Release a permit, returning it into this semaphore\n void unlock(void) {\n ++count_;\n std::lock_guard latch(mutex_);\n cond_.notify_one();\n }\n\n private:\n //! Disable them\n Semaphore(const Semaphore &) = delete;\n Semaphore(Semaphore &&) = delete;\n Semaphore &operator=(const Semaphore &) = delete;\n Semaphore &operator=(Semaphore &&) = delete;\n\n //! Members\n std::atomic count_{0};\n std::mutex mutex_{};\n std::condition_variable cond_{};\n};\n\n/*! Binary Semaphores\n */\ntemplate ::type>\nclass BinarySemaphores {\n public:\n using BitwiseType = typename std::conditional<\n N <= 32u,\n typename std::conditional<\n N <= 16u, typename std::conditional::type,\n uint32_t>::type,\n uint64_t>::type;\n\n //! Constructor\n BinarySemaphores(void) : BinarySemaphores{1} {}\n\n //! Constructor\n BinarySemaphores(uint32_t count) {\n if (count == 0 || count > N) {\n count = N;\n }\n count_ = count;\n mask_ = static_cast(BitwiseType(1) << (count - 1));\n mask_ |= static_cast(mask_ - 1);\n flags_.store(mask_);\n }\n\n //! Acquire a permit from this semaphore, suspending until one is available\n int acquire(void) {\n int index = -1;\n while ((index = this->try_acquire()) < 0) {\n std::unique_lock latch(mutex_);\n cond_.wait(latch, [this]() { return (flags_ > 0); });\n }\n return index;\n }\n\n //! Try to acquire a permit from this semaphore without suspension\n int try_acquire(void) {\n BitwiseType flags = flags_.load(std::memory_order_relaxed);\n while (flags > 0) {\n int index = CountTrailingZeros(flags);\n if (flags_.compare_exchange_weak(\n flags, flags & (~(BitwiseType(1) << index)),\n std::memory_order_release, std::memory_order_relaxed)) {\n return index;\n }\n flags = flags_.load(std::memory_order_relaxed);\n }\n return -1;\n }\n\n //! Acquire a specified permit from this semaphore, suspending until index is\n //! available\n int acquire(int index) {\n if (index < 0 || (uint32_t)index >= count_) {\n return -1;\n }\n BitwiseType flags = flags_.load(std::memory_order_relaxed);\n BitwiseType mask = BitwiseType(1) << index;\n while (true) {\n if ((flags & mask) &&\n flags_.compare_exchange_weak(flags, flags & (~mask),\n std::memory_order_release,\n std::memory_order_relaxed)) {\n return index;\n }\n flags = flags_.load(std::memory_order_relaxed);\n }\n }\n\n //! Release a permit, returning it into this semaphore\n void release(int index) {\n flags_.fetch_or((BitwiseType(1) << index) & mask_);\n std::lock_guard latch(mutex_);\n cond_.notify_one();\n }\n\n protected:\n //! Count the trailing zeros (32 bits)\n template \n static inline auto CountTrailingZeros(T val) ->\n typename std::enable_if::type {\n return ailego_ctz32(val);\n }\n\n //! Count the trailing zeros (64 bits)\n template \n static inline auto CountTrailingZeros(T val) ->\n typename std::enable_if::type {\n return ailego_ctz64(val);\n }\n\n private:\n //! Disable them\n BinarySemaphores(const BinarySemaphores &) = delete;\n BinarySemaphores(BinarySemaphores &&) = delete;\n BinarySemaphores &operator=(const BinarySemaphores &) = delete;\n BinarySemaphores &operator=(BinarySemaphores &&) = delete;\n\n //! Members\n uint32_t count_{0};\n BitwiseType mask_{0};\n std::atomic flags_{0};\n std::mutex mutex_{};\n std::condition_variable cond_{};\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/parallel/thread_pool.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n\n#if (defined(__linux) || defined(__linux__)) && !defined(__ANDROID__)\n#include \n\nstatic inline void BindThreads(std::vector &pool) {\n uint32_t hc = std::thread::hardware_concurrency();\n if (hc > 1) {\n cpu_set_t mask;\n\n for (size_t i = 0u; i < pool.size(); ++i) {\n CPU_ZERO(&mask);\n CPU_SET(i % hc, &mask);\n pthread_setaffinity_np(pool[i].native_handle(), sizeof(mask), &mask);\n }\n }\n}\n\nstatic inline void UnbindThreads(std::vector &pool) {\n cpu_set_t mask;\n CPU_ZERO(&mask);\n\n for (size_t i = 0u; i < CPU_SETSIZE; ++i) {\n CPU_SET(i, &mask);\n }\n for (size_t i = 0u; i < pool.size(); ++i) {\n pthread_setaffinity_np(pool[i].native_handle(), sizeof(mask), &mask);\n }\n}\n#else\nstatic inline void BindThreads(std::vector &) {}\nstatic inline void UnbindThreads(std::vector &) {}\n#endif\n\nnamespace zvec {\nnamespace ailego {\n\nThreadPool::ThreadPool(uint32_t size, bool binding) {\n for (uint32_t i = 0u; i < size; ++i) {\n pool_.emplace_back(&ThreadPool::worker, this);\n }\n if (binding) {\n this->bind();\n }\n}\n\nvoid ThreadPool::bind(void) {\n BindThreads(pool_);\n}\n\nvoid ThreadPool::unbind(void) {\n UnbindThreads(pool_);\n}\n\nvoid ThreadPool::worker(void) {\n // Counter of workers\n ++worker_count_;\n\n ThreadPool::Task task;\n while (this->picking(&task)) {\n // Run the task\n task.handle->run();\n task.handle = nullptr;\n\n // Notify task finished\n if (task.control) {\n task.control->notify();\n }\n\n // Notify task group\n if (task.group) {\n task.group->notify();\n task.group = nullptr;\n }\n\n // Decrease count of active works\n std::lock_guard lock(wait_mutex_);\n if (--active_count_ == 0 && pending_count_ == 0) {\n finished_cond_.notify_all();\n }\n }\n\n // Decrease count of workers\n std::lock_guard lock(wait_mutex_);\n if (--worker_count_ == 0) {\n stopped_cond_.notify_all();\n }\n}\n\nbool ThreadPool::picking(ThreadPool::Task *task) {\n std::unique_lock latch(queue_mutex_);\n work_cond_.wait(latch,\n [this]() { return (pending_count_ > 0 || stopping_); });\n if (stopping_) {\n return false;\n }\n\n // Pop a task\n auto &head = queue_.front();\n task->control = head.control;\n task->group = std::move(head.group);\n task->handle = std::move(head.handle);\n queue_.pop();\n\n // Update group control\n if (task->group) {\n task->group->mark_task_actived();\n }\n\n // Counter of active tasks\n std::unique_lock lock(wait_mutex_);\n ++active_count_;\n --pending_count_;\n\n return true;\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/pattern/defer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \"scope_guard.h\"\n\n#define AILEGO_DEFER_NAME_(x, y) x##y\n#define AILEGO_DEFER_NAME(x) AILEGO_DEFER_NAME_(__ailegoDefer_, x)\n\n//! Defer operator\n#define AILEGO_DEFER(...) \\\n auto AILEGO_DEFER_NAME(__LINE__) = ailego::ScopeGuard::Make(__VA_ARGS__)\n" }, { "path": "src/ailego/pattern/scope_guard.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Scope Guard Implementation\n */\ntemplate \nclass ScopeGuardImpl {\n public:\n using Object = CallbackObject;\n using Functor = CallbackFunctor;\n\n //! Constructor\n ScopeGuardImpl(ScopeGuardImpl &&rhs)\n : obj_(rhs.obj_),\n impl_(std::move(rhs.impl_)),\n tuple_(std::move(rhs.tuple_)) {\n rhs.obj_ = nullptr;\n }\n\n //! Constructor\n template \n ScopeGuardImpl(typename Object::Type *obj, const typename Functor::Type &impl,\n TArgs &&...args)\n : obj_(obj), impl_(impl), tuple_(std::forward(args)...) {}\n\n //! Constructor\n template \n ScopeGuardImpl(typename Object::Type *obj, typename Functor::Type &&impl,\n TArgs &&...args)\n : obj_(obj),\n impl_(std::move(impl)),\n tuple_(std::forward(args)...) {}\n\n // Destructor\n ~ScopeGuardImpl(void) {\n if (obj_) {\n Functor::Run(obj_, impl_, tuple_);\n }\n }\n\n protected:\n //! Disable them\n ScopeGuardImpl(void) = delete;\n ScopeGuardImpl(const ScopeGuardImpl &) = delete;\n ScopeGuardImpl &operator=(const ScopeGuardImpl &) = delete;\n\n private:\n //! Members\n typename Object::Type *obj_;\n typename Functor::Type impl_;\n typename Functor::TupleType tuple_;\n};\n\n/*! Scope Guard Implementation (void, TFunc)\n */\ntemplate \nclass ScopeGuardImpl {\n public:\n //! Callback Functor Type\n using Functor = CallbackFunctor;\n\n //! Constructor\n ScopeGuardImpl(ScopeGuardImpl &&rhs)\n : impl_(std::move(rhs.impl_)),\n tuple_(std::move(rhs.tuple_)),\n valid_(rhs.valid_) {\n rhs.valid_ = false;\n }\n\n //! Constructor\n template \n ScopeGuardImpl(const typename Functor::Type &impl, TArgs &&...args)\n : impl_(impl), tuple_(std::forward(args)...), valid_(true) {}\n\n //! Constructor\n template \n ScopeGuardImpl(typename Functor::Type &&impl, TArgs &&...args)\n : impl_(std::move(impl)),\n tuple_(std::forward(args)...),\n valid_(true) {}\n\n // Destructor\n ~ScopeGuardImpl(void) {\n if (valid_) {\n Functor::Run(impl_, tuple_);\n }\n }\n\n protected:\n //! Disable them\n ScopeGuardImpl(void) = delete;\n ScopeGuardImpl(const ScopeGuardImpl &) = delete;\n ScopeGuardImpl &operator=(const ScopeGuardImpl &) = delete;\n\n private:\n //! Members\n typename Functor::Type impl_;\n typename Functor::TupleType tuple_;\n bool valid_;\n};\n\n/*! Scope Guard\n */\nstruct ScopeGuard {\n //! Make a scope guard object (member function pointer)\n template \n static inline auto Make(T *obj, R (T::*impl)(TParams...), TArgs &&...args)\n -> ScopeGuardImpl::Type> {\n return ScopeGuardImpl::Type>(\n obj, impl, std::forward(args)...);\n }\n\n //! Make a scope guard object (constable member function pointer)\n template \n static inline auto Make(const T *obj, R (T::*impl)(TParams...) const,\n TArgs &&...args)\n -> ScopeGuardImpl::Type> {\n return ScopeGuardImpl::Type>(\n obj, impl, std::forward(args)...);\n }\n\n //! Make a scope guard object (volatile member function pointer)\n template \n static inline auto Make(volatile T *obj, R (T::*impl)(TParams...) volatile,\n TArgs &&...args)\n -> ScopeGuardImpl::Type> {\n return ScopeGuardImpl::Type>(\n obj, impl, std::forward(args)...);\n }\n\n //! Make a scope guard object (constable volatile member function pointer)\n template \n static inline auto Make(const volatile T *obj,\n R (T::*impl)(TParams...) const volatile,\n TArgs &&...args)\n -> ScopeGuardImpl::Type> {\n return ScopeGuardImpl::Type>(\n obj, impl, std::forward(args)...);\n }\n\n //! Make a scope guard object (function)\n template <\n typename TFunc, typename... TArgs,\n typename = typename std::enable_if::Value>::type>\n static inline auto Make(TFunc &&impl, TArgs &&...args)\n -> ScopeGuardImpl::Type> {\n return ScopeGuardImpl::Type>(\n std::forward(impl), std::forward(args)...);\n }\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/utility/bit_string_helper.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n\nnamespace zvec {\n\nnamespace ailego {\n\nclass BitStringWriter {\n public:\n BitStringWriter(uint8_t *buffer, size_t buffer_size)\n : buffer_(buffer), buffer_size_(buffer_size), offset_(0) {\n ::memset(buffer_, 0, buffer_size_);\n }\n\n bool write(uint64_t data, int nbit) {\n if (buffer_size_ * 8 < nbit + offset_) {\n return false;\n }\n\n int bits_remain = 8 - (offset_ & 7);\n\n if (nbit <= bits_remain) {\n buffer_[offset_ >> 3] |= data << (offset_ & 7);\n offset_ += nbit;\n } else {\n size_t j = offset_ >> 3;\n buffer_[j++] |= data << (offset_ & 7);\n offset_ += nbit;\n data >>= bits_remain;\n while (data != 0) {\n buffer_[j++] |= data;\n data >>= 8;\n }\n }\n\n return true;\n }\n\n size_t offset() {\n return offset_;\n }\n\n private:\n uint8_t *buffer_;\n size_t buffer_size_;\n size_t offset_;\n};\n\nclass BitStringReader {\n public:\n BitStringReader(const uint8_t *buffer, size_t buffer_size)\n : buffer_(buffer), buffer_size_(buffer_size), offset_(0) {}\n\n bool read(uint64_t &data, int nbit) {\n if (buffer_size_ * 8 < nbit + offset_) {\n return false;\n }\n\n int bits_remain = 8 - (offset_ & 7);\n\n uint64_t result = buffer_[offset_ >> 3] >> (offset_ & 7);\n if (nbit <= bits_remain) {\n result &= (1 << nbit) - 1;\n offset_ += nbit;\n\n data = result;\n } else {\n int temp = bits_remain;\n size_t i = (offset_ >> 3) + 1;\n offset_ += nbit;\n nbit -= bits_remain;\n\n while (nbit > 8) {\n result |= ((uint64_t)buffer_[i++]) << temp;\n temp += 8;\n nbit -= 8;\n }\n\n uint64_t last_byte = buffer_[i];\n\n last_byte &= (1 << nbit) - 1;\n result |= last_byte << temp;\n\n data = result;\n }\n\n return true;\n }\n\n size_t offset() {\n return offset_;\n }\n\n private:\n const uint8_t *buffer_;\n size_t buffer_size_;\n size_t offset_;\n};\n\n} // namespace ailego\n\n} // namespace zvec\n" }, { "path": "src/ailego/utility/bitset_helper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"bitset_helper.h\"\n#include \n\n#ifndef __SSE4_2__\n#define bitset_popcount32 ailego_popcount32\n#define bitset_popcount64 ailego_popcount64\n#else\n#define bitset_popcount32 _mm_popcnt_u32\n#define bitset_popcount64 _mm_popcnt_u64\n#endif // !__SSE4_2__\n\n#if defined(__ARM_NEON)\nstatic inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n vst1q_u32(lhs, vandq_u32(vld1q_u32(lhs), vld1q_u32(rhs)));\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] &= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= rhs[0];\n }\n}\n\nstatic inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs,\n size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n vst1q_u32(lhs, vbicq_u32(vld1q_u32(lhs), vld1q_u32(rhs)));\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] &= ~rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= ~rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= ~rhs[0];\n }\n}\n\nstatic inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n vst1q_u32(lhs, vorrq_u32(vld1q_u32(lhs), vld1q_u32(rhs)));\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] |= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] |= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] |= rhs[0];\n }\n}\n\nstatic inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n vst1q_u32(lhs, veorq_u32(vld1q_u32(lhs), vld1q_u32(rhs)));\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] ^= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] ^= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] ^= rhs[0];\n }\n}\n\nstatic inline void bitset_not(uint32_t *lhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n static const uint32x4_t v_zero = vdupq_n_u32(0);\n\n for (; lhs != last_aligned; lhs += 4) {\n vst1q_u32(lhs, vornq_u32(v_zero, vld1q_u32(lhs)));\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] = ~lhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] = ~lhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] = ~lhs[0];\n }\n}\n\nstatic inline bool bitset_test_all(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4) {\n uint64x2_t vu64 = vld1q_u64((const uint64_t *)lhs);\n if ((vgetq_lane_u64(vu64, 0) & vgetq_lane_u64(vu64, 1)) != (uint64_t)-1) {\n return false;\n }\n }\n switch (last - last_aligned) {\n case 3:\n if (lhs[2] != 0xffffffffu) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0xffffffffu) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0xffffffffu) {\n return false;\n }\n }\n return true;\n}\n\nstatic inline bool bitset_test_any(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4) {\n uint64x2_t vu64 = vld1q_u64((const uint64_t *)lhs);\n if (vgetq_lane_u64(vu64, 0) | vgetq_lane_u64(vu64, 1)) {\n return true;\n }\n }\n switch (last - last_aligned) {\n case 3:\n if (lhs[2] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return true;\n }\n }\n return false;\n}\n\nstatic inline bool bitset_test_none(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4) {\n uint64x2_t vu64 = vld1q_u64((const uint64_t *)lhs);\n if (vgetq_lane_u64(vu64, 0) | vgetq_lane_u64(vu64, 1)) {\n return false;\n }\n }\n switch (last - last_aligned) {\n case 3:\n if (lhs[2] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return false;\n }\n }\n return true;\n}\n\n#elif defined(__AVX2__)\nstatic inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m256i ymm0 = _mm256_load_si256((__m256i *)lhs);\n __m256i ymm1 = _mm256_load_si256((__m256i *)rhs);\n _mm256_store_si256((__m256i *)lhs, _mm256_and_si256(ymm1, ymm0));\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_load_si128((__m128i *)rhs);\n _mm_store_si128((__m128i *)lhs, _mm_and_si128(xmm1, xmm0));\n lhs += 4;\n rhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs);\n __m256i ymm1 = _mm256_loadu_si256((__m256i *)rhs);\n _mm256_storeu_si256((__m256i *)lhs, _mm256_and_si256(ymm1, ymm0));\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs);\n _mm_storeu_si128((__m128i *)lhs, _mm_and_si128(xmm1, xmm0));\n lhs += 4;\n rhs += 4;\n }\n }\n switch (last - lhs) {\n case 3:\n lhs[2] &= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= rhs[0];\n }\n}\n\nstatic inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs,\n size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m256i ymm0 = _mm256_load_si256((__m256i *)lhs);\n __m256i ymm1 = _mm256_load_si256((__m256i *)rhs);\n _mm256_store_si256((__m256i *)lhs, _mm256_andnot_si256(ymm1, ymm0));\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_load_si128((__m128i *)rhs);\n _mm_store_si128((__m128i *)lhs, _mm_andnot_si128(xmm1, xmm0));\n lhs += 4;\n rhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs);\n __m256i ymm1 = _mm256_loadu_si256((__m256i *)rhs);\n _mm256_storeu_si256((__m256i *)lhs, _mm256_andnot_si256(ymm1, ymm0));\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs);\n _mm_storeu_si128((__m128i *)lhs, _mm_andnot_si128(xmm1, xmm0));\n lhs += 4;\n rhs += 4;\n }\n }\n switch (last - lhs) {\n case 3:\n lhs[2] &= ~rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= ~rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= ~rhs[0];\n }\n}\n\nstatic inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m256i ymm0 = _mm256_load_si256((__m256i *)lhs);\n __m256i ymm1 = _mm256_load_si256((__m256i *)rhs);\n _mm256_store_si256((__m256i *)lhs, _mm256_or_si256(ymm1, ymm0));\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_load_si128((__m128i *)rhs);\n _mm_store_si128((__m128i *)lhs, _mm_or_si128(xmm1, xmm0));\n lhs += 4;\n rhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs);\n __m256i ymm1 = _mm256_loadu_si256((__m256i *)rhs);\n _mm256_storeu_si256((__m256i *)lhs, _mm256_or_si256(ymm1, ymm0));\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs);\n _mm_storeu_si128((__m128i *)lhs, _mm_or_si128(xmm1, xmm0));\n lhs += 4;\n rhs += 4;\n }\n }\n switch (last - lhs) {\n case 3:\n lhs[2] |= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] |= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] |= rhs[0];\n }\n}\n\nstatic inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n if (((uintptr_t)lhs & 0x1f) == 0 && ((uintptr_t)rhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m256i ymm0 = _mm256_load_si256((__m256i *)lhs);\n __m256i ymm1 = _mm256_load_si256((__m256i *)rhs);\n _mm256_store_si256((__m256i *)lhs, _mm256_xor_si256(ymm1, ymm0));\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_load_si128((__m128i *)rhs);\n _mm_store_si128((__m128i *)lhs, _mm_xor_si128(xmm1, xmm0));\n lhs += 4;\n rhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs);\n __m256i ymm1 = _mm256_loadu_si256((__m256i *)rhs);\n _mm256_storeu_si256((__m256i *)lhs, _mm256_xor_si256(ymm1, ymm0));\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs);\n _mm_storeu_si128((__m128i *)lhs, _mm_xor_si128(xmm1, xmm0));\n lhs += 4;\n rhs += 4;\n }\n }\n switch (last - lhs) {\n case 3:\n lhs[2] ^= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] ^= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] ^= rhs[0];\n }\n}\n\nstatic inline void bitset_not(uint32_t *lhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n static const __m256i mask_256 = _mm256_set1_epi32(0xffffffffu);\n static const __m128i mask_128 = _mm_set1_epi32(0xffffffffu);\n\n if (((uintptr_t)lhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 8) {\n _mm256_store_si256(\n (__m256i *)lhs,\n _mm256_andnot_si256(_mm256_load_si256((__m256i *)lhs), mask_256));\n }\n if (last >= last_aligned + 4) {\n _mm_store_si128(\n (__m128i *)lhs,\n _mm_andnot_si128(_mm_load_si128((__m128i *)lhs), mask_128));\n lhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8) {\n _mm256_storeu_si256(\n (__m256i *)lhs,\n _mm256_andnot_si256(_mm256_loadu_si256((__m256i *)lhs), mask_256));\n }\n if (last >= last_aligned + 4) {\n _mm_storeu_si128(\n (__m128i *)lhs,\n _mm_andnot_si128(_mm_lddqu_si128((__m128i *)lhs), mask_128));\n lhs += 4;\n }\n }\n switch (last - lhs) {\n case 3:\n lhs[2] = ~lhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] = ~lhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] = ~lhs[0];\n }\n}\n\nstatic inline bool bitset_test_all(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n static const __m256i mask_256 = _mm256_set1_epi32(0xffffffffu);\n static const __m128i mask_128 = _mm_set1_epi32(0xffffffffu);\n\n if (((uintptr_t)lhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 8) {\n __m256i neq =\n _mm256_xor_si256(_mm256_load_si256((__m256i *)lhs), mask_256);\n if (!_mm256_testz_si256(neq, neq)) {\n return false;\n }\n }\n if (last >= last_aligned + 4) {\n __m128i neq = _mm_xor_si128(_mm_load_si128((__m128i *)lhs), mask_128);\n if (!_mm_testz_si128(neq, neq)) {\n return false;\n }\n lhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8) {\n __m256i neq =\n _mm256_xor_si256(_mm256_loadu_si256((__m256i *)lhs), mask_256);\n if (!_mm256_testz_si256(neq, neq)) {\n return false;\n }\n }\n if (last >= last_aligned + 4) {\n __m128i neq = _mm_xor_si128(_mm_lddqu_si128((__m128i *)lhs), mask_128);\n if (!_mm_testz_si128(neq, neq)) {\n return false;\n }\n lhs += 4;\n }\n }\n switch (last - lhs) {\n case 3:\n if (lhs[2] != 0xffffffffu) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0xffffffffu) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0xffffffffu) {\n return false;\n }\n }\n return true;\n}\n\nstatic inline bool bitset_test_any(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n if (((uintptr_t)lhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 8) {\n __m256i ymm0 = _mm256_load_si256((__m256i *)lhs);\n if (!_mm256_testz_si256(ymm0, ymm0)) {\n return true;\n }\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n if (!_mm_testz_si128(xmm0, xmm0)) {\n return true;\n }\n lhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8) {\n __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs);\n if (!_mm256_testz_si256(ymm0, ymm0)) {\n return true;\n }\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n if (!_mm_testz_si128(xmm0, xmm0)) {\n return true;\n }\n lhs += 4;\n }\n }\n switch (last - lhs) {\n case 3:\n if (lhs[2] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return true;\n }\n }\n return false;\n}\n\nstatic inline bool bitset_test_none(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n if (((uintptr_t)lhs & 0x1f) == 0) {\n for (; lhs != last_aligned; lhs += 8) {\n __m256i ymm0 = _mm256_load_si256((__m256i *)lhs);\n if (!_mm256_testz_si256(ymm0, ymm0)) {\n return false;\n }\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n if (!_mm_testz_si128(xmm0, xmm0)) {\n return false;\n }\n lhs += 4;\n }\n } else {\n for (; lhs != last_aligned; lhs += 8) {\n __m256i ymm0 = _mm256_loadu_si256((__m256i *)lhs);\n if (!_mm256_testz_si256(ymm0, ymm0)) {\n return false;\n }\n }\n if (last >= last_aligned + 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n if (!_mm_testz_si128(xmm0, xmm0)) {\n return false;\n }\n lhs += 4;\n }\n }\n switch (last - lhs) {\n case 3:\n if (lhs[2] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return false;\n }\n }\n return true;\n}\n\n#elif defined(__SSE2__)\n#ifndef __SSE3__\n#define _mm_lddqu_si128 _mm_loadu_si128\n#endif // !__SSE3__\n\nstatic inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_load_si128((__m128i *)rhs);\n _mm_store_si128((__m128i *)lhs, _mm_and_si128(xmm1, xmm0));\n }\n } else {\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs);\n _mm_storeu_si128((__m128i *)lhs, _mm_and_si128(xmm1, xmm0));\n }\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] &= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= rhs[0];\n }\n}\n\nstatic inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs,\n size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_load_si128((__m128i *)rhs);\n _mm_store_si128((__m128i *)lhs, _mm_andnot_si128(xmm1, xmm0));\n }\n } else {\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs);\n _mm_storeu_si128((__m128i *)lhs, _mm_andnot_si128(xmm1, xmm0));\n }\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] &= ~rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= ~rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= ~rhs[0];\n }\n}\n\nstatic inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_load_si128((__m128i *)rhs);\n _mm_store_si128((__m128i *)lhs, _mm_or_si128(xmm1, xmm0));\n }\n } else {\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs);\n _mm_storeu_si128((__m128i *)lhs, _mm_or_si128(xmm1, xmm0));\n }\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] |= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] |= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] |= rhs[0];\n }\n}\n\nstatic inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n if (((uintptr_t)lhs & 0xf) == 0 && ((uintptr_t)rhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_load_si128((__m128i *)rhs);\n _mm_store_si128((__m128i *)lhs, _mm_xor_si128(xmm1, xmm0));\n }\n } else {\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n __m128i xmm1 = _mm_lddqu_si128((__m128i *)rhs);\n _mm_storeu_si128((__m128i *)lhs, _mm_xor_si128(xmm1, xmm0));\n }\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] ^= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] ^= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] ^= rhs[0];\n }\n}\n\nstatic inline void bitset_not(uint32_t *lhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n static const __m128i mask = _mm_set1_epi32(0xffffffffu);\n\n if (((uintptr_t)lhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4) {\n _mm_store_si128((__m128i *)lhs,\n _mm_andnot_si128(_mm_load_si128((__m128i *)lhs), mask));\n }\n } else {\n for (; lhs != last_aligned; lhs += 4) {\n _mm_storeu_si128((__m128i *)lhs,\n _mm_andnot_si128(_mm_lddqu_si128((__m128i *)lhs), mask));\n }\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] = ~lhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] = ~lhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] = ~lhs[0];\n }\n}\n\nstatic inline bool bitset_test_all(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n static const __m128i mask = _mm_set1_epi32(0xffffffffu);\n\n#ifndef __SSE4_1__\n if (((uintptr_t)lhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i eq = _mm_cmpeq_epi32(_mm_load_si128((__m128i *)lhs), mask);\n if (_mm_movemask_epi8(eq) != 0xffffu) {\n return false;\n }\n }\n } else {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i eq = _mm_cmpeq_epi32(_mm_lddqu_si128((__m128i *)lhs), mask);\n if (_mm_movemask_epi8(eq) != 0xffffu) {\n return false;\n }\n }\n }\n#else\n if (((uintptr_t)lhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i neq = _mm_xor_si128(_mm_load_si128((__m128i *)lhs), mask);\n if (!_mm_testz_si128(neq, neq)) {\n return false;\n }\n }\n } else {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i neq = _mm_xor_si128(_mm_lddqu_si128((__m128i *)lhs), mask);\n if (!_mm_testz_si128(neq, neq)) {\n return false;\n }\n }\n }\n#endif // !__SSE4_1__\n\n switch (last - last_aligned) {\n case 3:\n if (lhs[2] != 0xffffffffu) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0xffffffffu) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0xffffffffu) {\n return false;\n }\n }\n return true;\n}\n\nstatic inline bool bitset_test_any(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n#ifndef __SSE4_1__\n static const __m128i zero = _mm_setzero_si128();\n\n if (((uintptr_t)lhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i eq = _mm_cmpeq_epi32(_mm_load_si128((__m128i *)lhs), zero);\n if (_mm_movemask_epi8(eq) != 0xffffu) {\n return true;\n }\n }\n } else {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i eq = _mm_cmpeq_epi32(_mm_lddqu_si128((__m128i *)lhs), zero);\n if (_mm_movemask_epi8(eq) != 0xffffu) {\n return true;\n }\n }\n }\n#else\n if (((uintptr_t)lhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n if (!_mm_testz_si128(xmm0, xmm0)) {\n return true;\n }\n }\n } else {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n if (!_mm_testz_si128(xmm0, xmm0)) {\n return true;\n }\n }\n }\n#endif // !__SSE4_1__\n\n switch (last - last_aligned) {\n case 3:\n if (lhs[2] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return true;\n }\n }\n return false;\n}\n\nstatic inline bool bitset_test_none(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n#ifndef __SSE4_1__\n static __m128i zero = _mm_setzero_si128();\n\n if (((uintptr_t)lhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i eq = _mm_cmpeq_epi32(_mm_load_si128((__m128i *)lhs), zero);\n if (_mm_movemask_epi8(eq) != 0xffffu) {\n return false;\n }\n }\n } else {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i eq = _mm_cmpeq_epi32(_mm_lddqu_si128((__m128i *)lhs), zero);\n if (_mm_movemask_epi8(eq) != 0xffffu) {\n return false;\n }\n }\n }\n#else\n if (((uintptr_t)lhs & 0xf) == 0) {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i xmm0 = _mm_load_si128((__m128i *)lhs);\n if (!_mm_testz_si128(xmm0, xmm0)) {\n return false;\n }\n }\n } else {\n for (; lhs != last_aligned; lhs += 4) {\n __m128i xmm0 = _mm_lddqu_si128((__m128i *)lhs);\n if (!_mm_testz_si128(xmm0, xmm0)) {\n return false;\n }\n }\n }\n#endif // !__SSE4_1__\n\n switch (last - last_aligned) {\n case 3:\n if (lhs[2] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return false;\n }\n }\n return true;\n}\n\n#else\n#if defined(AILEGO_M64)\nstatic inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n *(uint64_t *)(&lhs[6]) &= *(uint64_t *)(&rhs[6]);\n *(uint64_t *)(&lhs[4]) &= *(uint64_t *)(&rhs[4]);\n *(uint64_t *)(&lhs[2]) &= *(uint64_t *)(&rhs[2]);\n *(uint64_t *)(&lhs[0]) &= *(uint64_t *)(&rhs[0]);\n }\n switch (last - last_aligned) {\n case 7:\n lhs[6] &= rhs[6];\n /* FALLTHRU */\n case 6:\n lhs[5] &= rhs[5];\n /* FALLTHRU */\n case 5:\n lhs[4] &= rhs[4];\n /* FALLTHRU */\n case 4:\n lhs[3] &= rhs[3];\n /* FALLTHRU */\n case 3:\n lhs[2] &= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= rhs[0];\n }\n}\n\nstatic inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs,\n size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n *(uint64_t *)(&lhs[6]) &= ~(*(uint64_t *)(&rhs[6]));\n *(uint64_t *)(&lhs[4]) &= ~(*(uint64_t *)(&rhs[4]));\n *(uint64_t *)(&lhs[2]) &= ~(*(uint64_t *)(&rhs[2]));\n *(uint64_t *)(&lhs[0]) &= ~(*(uint64_t *)(&rhs[0]));\n }\n switch (last - last_aligned) {\n case 7:\n lhs[6] &= ~rhs[6];\n /* FALLTHRU */\n case 6:\n lhs[5] &= ~rhs[5];\n /* FALLTHRU */\n case 5:\n lhs[4] &= ~rhs[4];\n /* FALLTHRU */\n case 4:\n lhs[3] &= ~rhs[3];\n /* FALLTHRU */\n case 3:\n lhs[2] &= ~rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= ~rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= ~rhs[0];\n }\n}\n\nstatic inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n *(uint64_t *)(&lhs[6]) |= *(uint64_t *)(&rhs[6]);\n *(uint64_t *)(&lhs[4]) |= *(uint64_t *)(&rhs[4]);\n *(uint64_t *)(&lhs[2]) |= *(uint64_t *)(&rhs[2]);\n *(uint64_t *)(&lhs[0]) |= *(uint64_t *)(&rhs[0]);\n }\n switch (last - last_aligned) {\n case 7:\n lhs[6] |= rhs[6];\n /* FALLTHRU */\n case 6:\n lhs[5] |= rhs[5];\n /* FALLTHRU */\n case 5:\n lhs[4] |= rhs[4];\n /* FALLTHRU */\n case 4:\n lhs[3] |= rhs[3];\n /* FALLTHRU */\n case 3:\n lhs[2] |= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] |= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] |= rhs[0];\n }\n}\n\nstatic inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n *(uint64_t *)(&lhs[6]) ^= *(uint64_t *)(&rhs[6]);\n *(uint64_t *)(&lhs[4]) ^= *(uint64_t *)(&rhs[4]);\n *(uint64_t *)(&lhs[2]) ^= *(uint64_t *)(&rhs[2]);\n *(uint64_t *)(&lhs[0]) ^= *(uint64_t *)(&rhs[0]);\n }\n switch (last - last_aligned) {\n case 7:\n lhs[6] ^= rhs[6];\n /* FALLTHRU */\n case 6:\n lhs[5] ^= rhs[5];\n /* FALLTHRU */\n case 5:\n lhs[4] ^= rhs[4];\n /* FALLTHRU */\n case 4:\n lhs[3] ^= rhs[3];\n /* FALLTHRU */\n case 3:\n lhs[2] ^= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] ^= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] ^= rhs[0];\n }\n}\n\nstatic inline void bitset_not(uint32_t *lhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n for (; lhs != last_aligned; lhs += 8) {\n *(uint64_t *)(&lhs[6]) = ~(*(uint64_t *)(&lhs[6]));\n *(uint64_t *)(&lhs[4]) = ~(*(uint64_t *)(&lhs[4]));\n *(uint64_t *)(&lhs[2]) = ~(*(uint64_t *)(&lhs[2]));\n *(uint64_t *)(&lhs[0]) = ~(*(uint64_t *)(&lhs[0]));\n }\n switch (last - last_aligned) {\n case 7:\n lhs[6] = ~lhs[6];\n /* FALLTHRU */\n case 6:\n lhs[5] = ~lhs[5];\n /* FALLTHRU */\n case 5:\n lhs[4] = ~lhs[4];\n /* FALLTHRU */\n case 4:\n lhs[3] = ~lhs[3];\n /* FALLTHRU */\n case 3:\n lhs[2] = ~lhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] = ~lhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] = ~lhs[0];\n }\n}\n\nstatic inline bool bitset_test_all(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n for (; lhs != last_aligned; lhs += 8) {\n if (*(uint64_t *)(&lhs[6]) != (uint64_t)-1) {\n return false;\n }\n if (*(uint64_t *)(&lhs[4]) != (uint64_t)-1) {\n return false;\n }\n if (*(uint64_t *)(&lhs[2]) != (uint64_t)-1) {\n return false;\n }\n if (*(uint64_t *)(&lhs[0]) != (uint64_t)-1) {\n return false;\n }\n }\n switch (last - last_aligned) {\n case 7:\n if (lhs[6] != (uint32_t)-1) {\n return false;\n }\n /* FALLTHRU */\n case 6:\n if (lhs[5] != (uint32_t)-1) {\n return false;\n }\n /* FALLTHRU */\n case 5:\n if (lhs[4] != (uint32_t)-1) {\n return false;\n }\n /* FALLTHRU */\n case 4:\n if (lhs[3] != (uint32_t)-1) {\n return false;\n }\n /* FALLTHRU */\n case 3:\n if (lhs[2] != (uint32_t)-1) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != (uint32_t)-1) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != (uint32_t)-1) {\n return false;\n }\n }\n return true;\n}\n\nstatic inline bool bitset_test_any(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n for (; lhs != last_aligned; lhs += 8) {\n if (*(uint64_t *)(&lhs[6]) != 0u) {\n return true;\n }\n if (*(uint64_t *)(&lhs[4]) != 0u) {\n return true;\n }\n if (*(uint64_t *)(&lhs[2]) != 0u) {\n return true;\n }\n if (*(uint64_t *)(&lhs[0]) != 0u) {\n return true;\n }\n }\n switch (last - last_aligned) {\n case 7:\n if (lhs[6] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 6:\n if (lhs[5] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 5:\n if (lhs[4] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 4:\n if (lhs[3] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 3:\n if (lhs[2] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return true;\n }\n }\n return false;\n}\n\nstatic inline bool bitset_test_none(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n\n for (; lhs != last_aligned; lhs += 8) {\n if (*(uint64_t *)(&lhs[6]) != 0u) {\n return false;\n }\n if (*(uint64_t *)(&lhs[4]) != 0u) {\n return false;\n }\n if (*(uint64_t *)(&lhs[2]) != 0u) {\n return false;\n }\n if (*(uint64_t *)(&lhs[0]) != 0u) {\n return false;\n }\n }\n switch (last - last_aligned) {\n case 7:\n if (lhs[6] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 6:\n if (lhs[5] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 5:\n if (lhs[4] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 4:\n if (lhs[3] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 3:\n if (lhs[2] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return false;\n }\n }\n return true;\n}\n\n#else // AILEGO_M64\nstatic inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n lhs[3] &= rhs[3];\n lhs[2] &= rhs[2];\n lhs[1] &= rhs[1];\n lhs[0] &= rhs[0];\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] &= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= rhs[0];\n }\n}\n\nstatic inline void bitset_andnot(uint32_t *lhs, const uint32_t *rhs,\n size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n lhs[3] &= ~rhs[3];\n lhs[2] &= ~rhs[2];\n lhs[1] &= ~rhs[1];\n lhs[0] &= ~rhs[0];\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] &= ~rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] &= ~rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] &= ~rhs[0];\n }\n}\n\nstatic inline void bitset_or(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n lhs[3] |= rhs[3];\n lhs[2] |= rhs[2];\n lhs[1] |= rhs[1];\n lhs[0] |= rhs[0];\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] |= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] |= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] |= rhs[0];\n }\n}\n\nstatic inline void bitset_xor(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n lhs[3] ^= rhs[3];\n lhs[2] ^= rhs[2];\n lhs[1] ^= rhs[1];\n lhs[0] ^= rhs[0];\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] ^= rhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] ^= rhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] ^= rhs[0];\n }\n}\n\nstatic inline void bitset_not(uint32_t *lhs, size_t size) {\n uint32_t *last = lhs + size;\n uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4) {\n lhs[3] = ~lhs[3];\n lhs[2] = ~lhs[2];\n lhs[1] = ~lhs[1];\n lhs[0] = ~lhs[0];\n }\n switch (last - last_aligned) {\n case 3:\n lhs[2] = ~lhs[2];\n /* FALLTHRU */\n case 2:\n lhs[1] = ~lhs[1];\n /* FALLTHRU */\n case 1:\n lhs[0] = ~lhs[0];\n }\n}\n\nstatic inline bool bitset_test_all(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4) {\n if (lhs[3] != (uint32_t)-1) {\n return false;\n }\n if (lhs[2] != (uint32_t)-1) {\n return false;\n }\n if (lhs[1] != (uint32_t)-1) {\n return false;\n }\n if (lhs[0] != (uint32_t)-1) {\n return false;\n }\n }\n switch (last - last_aligned) {\n case 3:\n if (lhs[2] != (uint32_t)-1) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != (uint32_t)-1) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != (uint32_t)-1) {\n return false;\n }\n }\n return true;\n}\n\nstatic inline bool bitset_test_any(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4) {\n if (lhs[3] != 0u) {\n return true;\n }\n if (lhs[2] != 0u) {\n return true;\n }\n if (lhs[1] != 0u) {\n return true;\n }\n if (lhs[0] != 0u) {\n return true;\n }\n }\n switch (last - last_aligned) {\n case 3:\n if (lhs[2] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return true;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return true;\n }\n }\n return false;\n}\n\nstatic inline bool bitset_test_none(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n\n for (; lhs != last_aligned; lhs += 4) {\n if (lhs[3] != 0u) {\n return false;\n }\n if (lhs[2] != 0u) {\n return false;\n }\n if (lhs[1] != 0u) {\n return false;\n }\n if (lhs[0] != 0u) {\n return false;\n }\n }\n switch (last - last_aligned) {\n case 3:\n if (lhs[2] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 2:\n if (lhs[1] != 0u) {\n return false;\n }\n /* FALLTHRU */\n case 1:\n if (lhs[0] != 0u) {\n return false;\n }\n }\n return true;\n}\n#endif // AILEGO_M64\n#endif // __AVX2__\n\n#if (defined(__ARM_NEON) && defined(__aarch64__))\nstatic inline size_t bitset_cardinality(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n while (lhs != last_aligned) {\n const uint32_t *last_stage =\n (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u;\n\n uint8x16_t v_count = vdupq_n_u8(0);\n for (; lhs != last_stage; lhs += 4) {\n v_count = vaddq_u8(vcntq_u8(vld1q_u8((const uint8_t *)lhs)), v_count);\n }\n\n v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count));\n count += vaddvq_u16(vreinterpretq_u16_u8(v_count));\n }\n\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_xor_cardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n while (lhs != last_aligned) {\n const uint32_t *last_stage =\n (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u;\n\n uint8x16_t v_count = vdupq_n_u8(0);\n for (; lhs != last_stage; lhs += 4, rhs += 4) {\n v_count = vaddq_u8(vcntq_u8(veorq_u8(vld1q_u8((const uint8_t *)lhs),\n vld1q_u8((const uint8_t *)rhs))),\n v_count);\n }\n\n v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count));\n count += vaddvq_u16(vreinterpretq_u16_u8(v_count));\n }\n\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2] ^ rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] ^ rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] ^ rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_and_cardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n while (lhs != last_aligned) {\n const uint32_t *last_stage =\n (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u;\n\n uint8x16_t v_count = vdupq_n_u8(0);\n for (; lhs != last_stage; lhs += 4, rhs += 4) {\n v_count = vaddq_u8(vcntq_u8(vandq_u8(vld1q_u8((const uint8_t *)lhs),\n vld1q_u8((const uint8_t *)rhs))),\n v_count);\n }\n\n v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count));\n count += vaddvq_u16(vreinterpretq_u16_u8(v_count));\n }\n\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2] & rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] & rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] & rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_andnot_cardinality(const uint32_t *lhs,\n const uint32_t *rhs,\n size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n while (lhs != last_aligned) {\n const uint32_t *last_stage =\n (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u;\n\n uint8x16_t v_count = vdupq_n_u8(0);\n for (; lhs != last_stage; lhs += 4, rhs += 4) {\n v_count = vaddq_u8(vcntq_u8(vbicq_u8(vld1q_u8((const uint8_t *)lhs),\n vld1q_u8((const uint8_t *)rhs))),\n v_count);\n }\n\n v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count));\n count += vaddvq_u16(vreinterpretq_u16_u8(v_count));\n }\n\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2] & ~rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] & ~rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] & ~rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_or_cardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n while (lhs != last_aligned) {\n const uint32_t *last_stage =\n (last_aligned <= lhs + 124u) ? last_aligned : lhs + 124u;\n\n uint8x16_t v_count = vdupq_n_u8(0);\n for (; lhs != last_stage; lhs += 4, rhs += 4) {\n v_count = vaddq_u8(vcntq_u8(vorrq_u8(vld1q_u8((const uint8_t *)lhs),\n vld1q_u8((const uint8_t *)rhs))),\n v_count);\n }\n\n v_count = vreinterpretq_u8_u16(vpaddlq_u8(v_count));\n count += vaddvq_u16(vreinterpretq_u16_u8(v_count));\n }\n\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2] | rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] | rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] | rhs[0]);\n }\n return count;\n}\n\n#elif defined(AILEGO_M64)\nstatic inline size_t bitset_cardinality(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 8) {\n count += bitset_popcount64(*(uint64_t *)(&lhs[6]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[4]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[2]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[0]));\n }\n switch (last - last_aligned) {\n case 7:\n count += bitset_popcount32(lhs[6]);\n /* FALLTHRU */\n case 6:\n count += bitset_popcount32(lhs[5]);\n /* FALLTHRU */\n case 5:\n count += bitset_popcount32(lhs[4]);\n /* FALLTHRU */\n case 4:\n count += bitset_popcount32(lhs[3]);\n /* FALLTHRU */\n case 3:\n count += bitset_popcount32(lhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_xor_cardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n count += bitset_popcount64(*(uint64_t *)(&lhs[6]) ^ *(uint64_t *)(&rhs[6]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[4]) ^ *(uint64_t *)(&rhs[4]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[2]) ^ *(uint64_t *)(&rhs[2]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[0]) ^ *(uint64_t *)(&rhs[0]));\n }\n switch (last - last_aligned) {\n case 7:\n count += bitset_popcount32(lhs[6] ^ rhs[6]);\n /* FALLTHRU */\n case 6:\n count += bitset_popcount32(lhs[5] ^ rhs[5]);\n /* FALLTHRU */\n case 5:\n count += bitset_popcount32(lhs[4] ^ rhs[4]);\n /* FALLTHRU */\n case 4:\n count += bitset_popcount32(lhs[3] ^ rhs[3]);\n /* FALLTHRU */\n case 3:\n count += bitset_popcount32(lhs[2] ^ rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] ^ rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] ^ rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_and_cardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n count += bitset_popcount64(*(uint64_t *)(&lhs[6]) & *(uint64_t *)(&rhs[6]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[4]) & *(uint64_t *)(&rhs[4]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[2]) & *(uint64_t *)(&rhs[2]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[0]) & *(uint64_t *)(&rhs[0]));\n }\n switch (last - last_aligned) {\n case 7:\n count += bitset_popcount32(lhs[6] & rhs[6]);\n /* FALLTHRU */\n case 6:\n count += bitset_popcount32(lhs[5] & rhs[5]);\n /* FALLTHRU */\n case 5:\n count += bitset_popcount32(lhs[4] & rhs[4]);\n /* FALLTHRU */\n case 4:\n count += bitset_popcount32(lhs[3] & rhs[3]);\n /* FALLTHRU */\n case 3:\n count += bitset_popcount32(lhs[2] & rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] & rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] & rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_andnot_cardinality(const uint32_t *lhs,\n const uint32_t *rhs,\n size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n count +=\n bitset_popcount64(*(uint64_t *)(&lhs[6]) & ~(*(uint64_t *)(&rhs[6])));\n count +=\n bitset_popcount64(*(uint64_t *)(&lhs[4]) & ~(*(uint64_t *)(&rhs[4])));\n count +=\n bitset_popcount64(*(uint64_t *)(&lhs[2]) & ~(*(uint64_t *)(&rhs[2])));\n count +=\n bitset_popcount64(*(uint64_t *)(&lhs[0]) & ~(*(uint64_t *)(&rhs[0])));\n }\n switch (last - last_aligned) {\n case 7:\n count += bitset_popcount32(lhs[6] & ~rhs[6]);\n /* FALLTHRU */\n case 6:\n count += bitset_popcount32(lhs[5] & ~rhs[5]);\n /* FALLTHRU */\n case 5:\n count += bitset_popcount32(lhs[4] & ~rhs[4]);\n /* FALLTHRU */\n case 4:\n count += bitset_popcount32(lhs[3] & ~rhs[3]);\n /* FALLTHRU */\n case 3:\n count += bitset_popcount32(lhs[2] & ~rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] & ~rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] & ~rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_or_cardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 3) << 3);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 8, rhs += 8) {\n count += bitset_popcount64(*(uint64_t *)(&lhs[6]) | *(uint64_t *)(&rhs[6]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[4]) | *(uint64_t *)(&rhs[4]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[2]) | *(uint64_t *)(&rhs[2]));\n count += bitset_popcount64(*(uint64_t *)(&lhs[0]) | *(uint64_t *)(&rhs[0]));\n }\n switch (last - last_aligned) {\n case 7:\n count += bitset_popcount32(lhs[6] | rhs[6]);\n /* FALLTHRU */\n case 6:\n count += bitset_popcount32(lhs[5] | rhs[5]);\n /* FALLTHRU */\n case 5:\n count += bitset_popcount32(lhs[4] | rhs[4]);\n /* FALLTHRU */\n case 4:\n count += bitset_popcount32(lhs[3] | rhs[3]);\n /* FALLTHRU */\n case 3:\n count += bitset_popcount32(lhs[2] | rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] | rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] | rhs[0]);\n }\n return count;\n}\n\n#else // !__ARM_NEON && !AILEGO_M64\nstatic inline size_t bitset_cardinality(const uint32_t *lhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 4) {\n count += bitset_popcount32(lhs[3]);\n count += bitset_popcount32(lhs[2]);\n count += bitset_popcount32(lhs[1]);\n count += bitset_popcount32(lhs[0]);\n }\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_xor_cardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n count += bitset_popcount32(lhs[3] ^ rhs[3]);\n count += bitset_popcount32(lhs[2] ^ rhs[2]);\n count += bitset_popcount32(lhs[1] ^ rhs[1]);\n count += bitset_popcount32(lhs[0] ^ rhs[0]);\n }\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2] ^ rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] ^ rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] ^ rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_and_cardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n count += bitset_popcount32(lhs[3] & rhs[3]);\n count += bitset_popcount32(lhs[2] & rhs[2]);\n count += bitset_popcount32(lhs[1] & rhs[1]);\n count += bitset_popcount32(lhs[0] & rhs[0]);\n }\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2] & rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] & rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] & rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_andnot_cardinality(const uint32_t *lhs,\n const uint32_t *rhs,\n size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n count += bitset_popcount32(lhs[3] & ~rhs[3]);\n count += bitset_popcount32(lhs[2] & ~rhs[2]);\n count += bitset_popcount32(lhs[1] & ~rhs[1]);\n count += bitset_popcount32(lhs[0] & ~rhs[0]);\n }\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2] & ~rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] & ~rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] & ~rhs[0]);\n }\n return count;\n}\n\nstatic inline size_t bitset_or_cardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n const uint32_t *last = lhs + size;\n const uint32_t *last_aligned = lhs + ((size >> 2) << 2);\n size_t count = 0;\n\n for (; lhs != last_aligned; lhs += 4, rhs += 4) {\n count += bitset_popcount32(lhs[3] | rhs[3]);\n count += bitset_popcount32(lhs[2] | rhs[2]);\n count += bitset_popcount32(lhs[1] | rhs[1]);\n count += bitset_popcount32(lhs[0] | rhs[0]);\n }\n switch (last - last_aligned) {\n case 3:\n count += bitset_popcount32(lhs[2] | rhs[2]);\n /* FALLTHRU */\n case 2:\n count += bitset_popcount32(lhs[1] | rhs[1]);\n /* FALLTHRU */\n case 1:\n count += bitset_popcount32(lhs[0] | rhs[0]);\n }\n return count;\n}\n#endif // __ARM_NEON && __aarch64__\n\nnamespace zvec {\n\nnamespace ailego {\n\nvoid BitsetHelper::BitwiseAnd(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n bitset_and(lhs, rhs, size);\n}\n\nvoid BitsetHelper::BitwiseAndnot(uint32_t *lhs, const uint32_t *rhs,\n size_t size) {\n bitset_andnot(lhs, rhs, size);\n}\n\nvoid BitsetHelper::BitwiseOr(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n bitset_or(lhs, rhs, size);\n}\n\nvoid BitsetHelper::BitwiseXor(uint32_t *lhs, const uint32_t *rhs, size_t size) {\n bitset_xor(lhs, rhs, size);\n}\n\nvoid BitsetHelper::BitwiseNot(uint32_t *arr, size_t size) {\n bitset_not(arr, size);\n}\n\nbool BitsetHelper::TestAll(const uint32_t *arr, size_t size) {\n return bitset_test_all(arr, size);\n}\n\nbool BitsetHelper::TestAny(const uint32_t *arr, size_t size) {\n return bitset_test_any(arr, size);\n}\n\nbool BitsetHelper::TestNone(const uint32_t *arr, size_t size) {\n return bitset_test_none(arr, size);\n}\n\nsize_t BitsetHelper::BitwiseAndCardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n return bitset_and_cardinality(lhs, rhs, size);\n}\n\nsize_t BitsetHelper::BitwiseOrCardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n return bitset_or_cardinality(lhs, rhs, size);\n}\n\nsize_t BitsetHelper::BitwiseAndnotCardinality(const uint32_t *lhs,\n const uint32_t *rhs,\n size_t size) {\n return bitset_andnot_cardinality(lhs, rhs, size);\n}\n\nsize_t BitsetHelper::BitwiseXorCardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size) {\n return bitset_xor_cardinality(lhs, rhs, size);\n}\n\nsize_t BitsetHelper::Cardinality(const uint32_t *arr, size_t size) {\n return bitset_cardinality(arr, size);\n}\n\nbool BitsetHelper::test_all(void) const {\n return bitset_test_all(array_, size_);\n}\n\nbool BitsetHelper::test_any(void) const {\n return bitset_test_any(array_, size_);\n}\n\nbool BitsetHelper::test_none(void) const {\n return bitset_test_none(array_, size_);\n}\n\nsize_t BitsetHelper::cardinality(void) const {\n return bitset_cardinality(array_, size_);\n}\n\n} // namespace ailego\n\n} // namespace zvec" }, { "path": "src/ailego/utility/bitset_helper.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n\nnamespace zvec {\n\nnamespace ailego {\n\n/*! Bitset Helper\n */\nclass BitsetHelper {\n public:\n //! Constructor\n BitsetHelper(void) {}\n\n //! Constructor\n BitsetHelper(void *buf, size_t len)\n : array_(reinterpret_cast(buf)),\n size_(len / sizeof(uint32_t)) {}\n\n //! Mount a buffer as bitset\n void mount(void *buf, size_t len) {\n array_ = reinterpret_cast(buf);\n size_ = len / sizeof(uint32_t);\n }\n\n //! Umount the buffer\n void umount(void) {\n array_ = nullptr;\n size_ = 0u;\n }\n\n // ！Clear the bitset\n void clear(void) {\n memset(array_, 0, sizeof(uint32_t) * size_);\n }\n\n //! Test a bit in bitset\n bool test(size_t num) const {\n ailego_assert_with((size_ << 5) > num, \"overflow argument\");\n return ((array_[num >> 5] & (1u << (num & 0x1f))) != 0);\n }\n\n //! Set a bit in bitset\n void set(size_t num) {\n ailego_assert_with((size_ << 5) > num, \"overflow argument\");\n uint32_t mask = (1u << (num & 0x1f));\n array_[num >> 5] |= mask;\n }\n\n //! Reset a bit in bitset\n void reset(size_t num) {\n ailego_assert_with((size_ << 5) > num, \"overflow argument\");\n uint32_t mask = (1u << (num & 0x1f));\n array_[num >> 5] &= ~mask;\n }\n\n //! Toggle a bit in bitset\n void flip(size_t num) {\n ailego_assert_with((size_ << 5) > num, \"overflow argument\");\n uint32_t mask = (1u << (num & 0x1f));\n array_[num >> 5] ^= mask;\n }\n\n //! Extract the bitset to an array\n void extract(size_t base, std::vector *out) const {\n const uint32_t *iter = array_;\n const uint32_t *last = array_ + size_;\n\n for (; iter != last; ++iter) {\n uint32_t w = *iter;\n\n while (w != 0) {\n uint32_t c = ailego_ctz32(w);\n w &= ~(1u << c);\n out->push_back(base + c);\n }\n base += 32u;\n }\n }\n\n //! Extract the bitset to an array\n void extract(std::vector *out) const {\n this->extract(0, out);\n }\n\n //! Check if all bits are set to true\n bool test_all(void) const;\n\n //! Check if any bits are set to true\n bool test_any(void) const;\n\n //! Check if none of the bits are set to true\n bool test_none(void) const;\n\n //! Compute the cardinality of a bitset\n size_t cardinality(void) const;\n\n //! Calculate the size of buffer if it contains N bits\n static size_t BufferSize(size_t N) {\n return (((N + 0x1f) >> 5) << 2);\n }\n\n //! Calculate the count of bits can be contained\n static size_t BitsCount(size_t len) {\n return ((len >> 2) << 2);\n }\n\n //! Check if all bits are set to true\n static bool TestAll(const uint32_t *arr, size_t size);\n\n //! Check if cube bits are set to true\n static bool TestAny(const uint32_t *arr, size_t size);\n\n //! Check if none of the bits are set to true\n static bool TestNone(const uint32_t *arr, size_t size);\n\n //! Compute the AND cardinality between two bitsets\n static size_t BitwiseAndCardinality(const uint32_t *lhs, const uint32_t *rhs,\n size_t size);\n\n //! Compute the OR cardinality between two bitsets\n static size_t BitwiseOrCardinality(const uint32_t *lhs, const uint32_t *rhs,\n size_t size);\n\n //! Compute the ANDNOT cardinality between two bitsets\n static size_t BitwiseAndnotCardinality(const uint32_t *lhs,\n const uint32_t *rhs, size_t size);\n\n //! Compute the XOR cardinality between two bitsets\n static size_t BitwiseXorCardinality(const uint32_t *lhs, const uint32_t *rhs,\n size_t size);\n\n //! Compute the cardinality of a bitset\n static size_t Cardinality(const uint32_t *arr, size_t size);\n\n //! Perform binary AND\n static void BitwiseAnd(uint32_t *lhs, const uint32_t *rhs, size_t size);\n\n //! Perform binary AND_NOT\n static void BitwiseAndnot(uint32_t *lhs, const uint32_t *rhs, size_t size);\n\n //! Perform binary OR\n static void BitwiseOr(uint32_t *lhs, const uint32_t *rhs, size_t size);\n\n //! Perform binary XOR\n static void BitwiseXor(uint32_t *lhs, const uint32_t *rhs, size_t size);\n\n //! Perform binary NOT\n static void BitwiseNot(uint32_t *arr, size_t size);\n\n private:\n uint32_t *array_{nullptr};\n size_t size_{0u};\n};\n\n} // namespace ailego\n\n} // namespace zvec\n" }, { "path": "src/ailego/utility/concurrency_helper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"concurrency_helper.h\"\n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n// Refer to:\n// https://stackoverflow.com/questions/65551215/get-docker-cpu-memory-limit-inside-container\nConcurrencyHelper::ConcurrencyHelper() {\n std::string cfs_quota_us = \"/sys/fs/cgroup/cpu/cpu.cfs_quota_us\";\n std::string cfs_period_us = \"/sys/fs/cgroup/cpu/cpu.cfs_period_us\";\n\n concurrency_ = std::thread::hardware_concurrency();\n if (FileHelper::IsExist(cfs_quota_us.c_str()) &&\n FileHelper::IsExist(cfs_period_us.c_str())) {\n std::ifstream quota_ifs;\n std::string quota_str{\"\"};\n uint32_t quota_val = 0;\n quota_ifs.open(cfs_quota_us, std::ios::in);\n if (quota_ifs.is_open()) {\n quota_ifs >> quota_str;\n if (quota_str != \"-1\") {\n StringHelper::ToUint32(quota_str, "a_val);\n }\n quota_ifs.close();\n }\n\n if (quota_val > 0) {\n std::ifstream period_ifs;\n std::string period_str{\"\"};\n uint32_t period_val = 0;\n period_ifs.open(cfs_period_us, std::ios::in);\n if (period_ifs.is_open()) {\n period_ifs >> period_str;\n StringHelper::ToUint32(period_str, &period_val);\n period_ifs.close();\n }\n\n if (period_val > 0) {\n concurrency_ = (quota_val + period_val - 1) / period_val;\n }\n }\n }\n}\n\nuint32_t ConcurrencyHelper::container_aware_concurrency() {\n static ConcurrencyHelper concurrency_helper;\n return concurrency_helper.concurrency_;\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/utility/concurrency_helper.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace ailego {\n\nclass ConcurrencyHelper {\n public:\n ConcurrencyHelper();\n\n //! get hardware concurrency from either vm or container\n static uint32_t container_aware_concurrency();\n\n private:\n uint32_t concurrency_{0};\n};\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/utility/dl_helper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"dl_helper.h\"\n#if !defined(_WIN64) && !defined(_WIN32)\n#include \n#else\n#include \n#endif\n\nnamespace zvec {\nnamespace ailego {\n\n#if !defined(_WIN64) && !defined(_WIN32)\nvoid *DLHelper::Load(const char *path, std::string *err) {\n void *handle = dlopen(path, RTLD_NOW);\n if (!handle && err) {\n *err = dlerror();\n }\n return handle;\n}\n\nvoid DLHelper::Unload(void *handle) {\n ailego_return_if_false(handle);\n dlclose(handle);\n}\n\nvoid *DLHelper::Symbol(void *handle, const char *symbol) {\n ailego_null_if_false(handle && symbol);\n return dlsym(handle, symbol);\n}\n\n#else\nvoid *DLHelper::Load(const char *path, std::string *err) {\n HMODULE handle = LoadLibraryA(path);\n if (!handle && err) {\n DWORD error_code = GetLastError();\n LPSTR error_msg = nullptr;\n\n DWORD len = FormatMessageA(\n FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM |\n FORMAT_MESSAGE_IGNORE_INSERTS,\n nullptr, error_code, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),\n (LPSTR)&error_msg, 0, nullptr);\n err->assign(error_msg, len);\n LocalFree(error_msg);\n }\n return handle;\n}\n\nvoid DLHelper::Unload(void *handle) {\n ailego_return_if_false(handle);\n FreeLibrary((HMODULE)handle);\n}\n\nvoid *DLHelper::Symbol(void *handle, const char *symbol) {\n ailego_null_if_false(handle && symbol);\n return GetProcAddress((HMODULE)handle, symbol);\n}\n#endif // !_WIN64 && !_WIN32\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/utility/dl_helper.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Dynamic Library Helper\n */\nstruct DLHelper {\n //! Load library from path\n static void *Load(const char *path, std::string *err);\n\n //! Unload a library\n static void Unload(void *handle);\n\n //! Retrieve a symbol from a library handle\n static void *Symbol(void *handle, const char *symbol);\n\n //! Load library from path\n static void *Load(const std::string &path, std::string *err) {\n return DLHelper::Load(path.c_str(), err);\n }\n\n //! Retrieve a symbol from a library handle\n static void *Symbol(void *handle, const std::string &symbol) {\n return DLHelper::Symbol(handle, symbol.c_str());\n }\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/utility/file_helper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n\n#if defined(_WIN32) || defined(_WIN64)\n#include \n#else\n#if defined(__APPLE__) || defined(__MACH__)\n#include \n#endif\n#include \n#include \n#include \n#include \n#include \n#include \n#endif\n\nnamespace zvec {\nnamespace ailego {\n\nbool FileHelper::GetSelfPath(std::string *path) {\n#if defined(_WIN32) || defined(_WIN64)\n char buf[MAX_PATH];\n DWORD len = GetModuleFileNameA(NULL, buf, MAX_PATH);\n#elif defined(__APPLE__) || defined(__MACH__)\n char buf[PATH_MAX];\n size_t len = 0;\n\n char dirty_buf[PATH_MAX];\n uint32_t size = sizeof(dirty_buf);\n if (_NSGetExecutablePath(dirty_buf, &size) == 0) {\n realpath(dirty_buf, buf);\n len = strlen(buf);\n }\n#elif defined(__FreeBSD__)\n char buf[PATH_MAX];\n size_t len = PATH_MAX;\n int mib[4] = {CTL_KERN, KERN_PROC, KERN_PROC_PATHNAME, -1};\n if (sysctl(mib, 4, &buf, &len, NULL, 0) != 0) {\n len = 0;\n }\n#else\n char buf[PATH_MAX];\n ssize_t len = readlink(\"/proc/self/exe\", buf, PATH_MAX);\n#endif\n\n if (len <= 0) {\n return false;\n }\n path->assign(buf, len);\n return true;\n}\n\nbool FileHelper::GetFilePath(NativeHandle handle, std::string *path) {\n#if defined(_WIN32) || defined(_WIN64)\n char buf[MAX_PATH];\n DWORD len =\n GetFinalPathNameByHandleA(handle, buf, MAX_PATH, FILE_NAME_OPENED);\n#elif defined(__linux) || defined(__linux__)\n char buf[PATH_MAX];\n char src[32];\n snprintf(src, sizeof(src), \"/proc/self/fd/%d\", handle);\n ssize_t len = readlink(src, buf, PATH_MAX);\n#else\n char buf[PATH_MAX];\n size_t len = 0;\n if (fcntl(handle, F_GETPATH, buf) != -1) {\n len = strlen(buf);\n }\n#endif\n\n if (len <= 0) {\n return false;\n }\n path->assign(buf, len);\n return true;\n}\n\n#if !defined(_WIN32) && !defined(_WIN64)\n\nstatic inline char *JoinFilePath(const char *prefix, const char *suffix) {\n size_t prefix_len = strlen(prefix);\n size_t suffix_len = strlen(suffix);\n\n char *path = (char *)malloc(prefix_len + suffix_len + 2);\n if (path) {\n memcpy(path, prefix, prefix_len);\n memcpy(path + prefix_len + 1, suffix, suffix_len);\n path[prefix_len] = '/';\n path[prefix_len + suffix_len + 1] = '\\0';\n }\n return path;\n}\n\nbool FileHelper::GetWorkingDirectory(std::string *path) {\n char buf[PATH_MAX];\n\n if (!getcwd(buf, PATH_MAX)) {\n return false;\n }\n path->assign(buf);\n return !path->empty();\n}\n\nbool FileHelper::GetFileSize(const char *path, size_t *psz) {\n struct stat buf;\n if (stat(path, &buf) != 0) {\n return false;\n }\n *psz = buf.st_size;\n return true;\n}\n\nbool FileHelper::DeleteFile(const char *path) {\n // Delete a file by the path\n return (unlink(path) == 0);\n}\n\nbool FileHelper::RenameFile(const char *oldpath, const char *newpath) {\n return (rename(oldpath, newpath) == 0);\n}\n\nbool FileHelper::MakePath(const char *path) {\n char pathbuf[PATH_MAX];\n char *sp, *pp;\n\n strncpy(pathbuf, path, sizeof(pathbuf) - 1);\n pathbuf[PATH_MAX - 1] = '\\0';\n\n pp = pathbuf;\n while ((sp = strchr(pp, '/')) != nullptr) {\n // Neither root nor double slash in path\n if (sp != pp) {\n *sp = '\\0';\n if (mkdir(pathbuf, 0755) == -1 && errno != EEXIST) {\n return false;\n }\n *sp = '/';\n }\n pp = sp + 1;\n }\n return !(*pp != '\\0' && mkdir(pathbuf, 0755) == -1 && errno != EEXIST);\n}\n\nbool FileHelper::RemoveDirectory(const char *path) {\n DIR *dir = opendir(path);\n if (!dir) {\n return false;\n }\n\n struct dirent *dent;\n while ((dent = readdir(dir)) != nullptr) {\n if (!strcmp(dent->d_name, \".\") || !strcmp(dent->d_name, \"..\")) {\n continue;\n }\n char *fullpath = JoinFilePath(path, dent->d_name);\n if (!fullpath) {\n continue;\n }\n\n if (FileHelper::IsDirectory(fullpath)) {\n FileHelper::RemoveDirectory(fullpath);\n } else {\n FileHelper::DeleteFile(fullpath);\n }\n free(fullpath);\n }\n closedir(dir);\n return (rmdir(path) == 0);\n}\n\nbool FileHelper::IsExist(const char *path) {\n return (access(path, F_OK) == 0);\n}\n\nbool FileHelper::IsRegular(const char *path) {\n struct stat buf;\n if (stat(path, &buf) != 0) {\n return false;\n }\n return ((buf.st_mode & S_IFREG) != 0);\n}\n\nbool FileHelper::IsDirectory(const char *path) {\n struct stat buf;\n if (stat(path, &buf) != 0) {\n return false;\n }\n return ((buf.st_mode & S_IFDIR) != 0);\n}\n\nbool FileHelper::IsSymbolicLink(const char *path) {\n struct stat buf;\n if (stat(path, &buf) != 0) {\n return false;\n }\n return ((buf.st_mode & S_IFLNK) != 0);\n}\n\nbool FileHelper::IsSame(const char *path1, const char *path2) {\n char real_path1[PATH_MAX];\n char real_path2[PATH_MAX];\n if (!realpath(path1, real_path1)) {\n return false;\n }\n if (!realpath(path2, real_path2)) {\n return false;\n }\n return (!strcmp(real_path1, real_path2));\n}\n\n#else\n#undef RemoveDirectory\n#undef DeleteFile\n#undef GetFileSize\n\nstatic inline char *JoinFilePath(const char *prefix, const char *suffix) {\n size_t prefix_len = strlen(prefix);\n size_t suffix_len = strlen(suffix);\n\n char *path = (char *)malloc(prefix_len + suffix_len + 2);\n if (path) {\n memcpy(path, prefix, prefix_len);\n memcpy(path + prefix_len + 1, suffix, suffix_len);\n path[prefix_len] = '\\\\';\n path[prefix_len + suffix_len + 1] = '\\0';\n }\n return path;\n}\n\nbool FileHelper::GetWorkingDirectory(std::string *path) {\n char buf[MAX_PATH];\n DWORD len = GetCurrentDirectoryA(MAX_PATH, buf);\n\n if (len <= 0) {\n return false;\n }\n path->assign(buf, len);\n return true;\n}\n\nbool FileHelper::GetFileSize(const char *path, size_t *psz) {\n HANDLE handle = CreateFileA(path, GENERIC_READ, FILE_SHARE_READ, nullptr,\n OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);\n\n LARGE_INTEGER file_size;\n if (!GetFileSizeEx(handle, &file_size)) {\n return false;\n }\n *psz = (size_t)file_size.QuadPart;\n return true;\n}\n\nbool FileHelper::DeleteFile(const char *path) {\n // Delete a file by the path\n return (DeleteFileA(path));\n}\n\nbool FileHelper::RenameFile(const char *oldpath, const char *newpath) {\n return (MoveFileA(oldpath, newpath));\n}\n\nbool FileHelper::MakePath(const char *path) {\n char pathbuf[MAX_PATH];\n char *sp, *pp;\n\n strncpy(pathbuf, path, sizeof(pathbuf) - 1);\n pathbuf[MAX_PATH - 1] = '\\0';\n\n pp = pathbuf;\n while ((sp = strpbrk(pp, \"/\\\\\")) != nullptr) {\n // Neither root nor double slash in path\n if (sp != pp) {\n *sp = '\\0';\n if (!CreateDirectoryA(pathbuf, nullptr) &&\n GetLastError() != ERROR_ALREADY_EXISTS) {\n return false;\n }\n *sp = '\\\\';\n }\n pp = sp + 1;\n }\n return !(*pp != '\\0' && !CreateDirectoryA(pathbuf, nullptr) &&\n GetLastError() != ERROR_ALREADY_EXISTS);\n}\n\nbool FileHelper::RemoveDirectory(const char *path) {\n char *pathbuf = JoinFilePath(path, \"*.*\");\n ailego_false_if_false(pathbuf);\n\n WIN32_FIND_DATAA file_info;\n HANDLE file = FindFirstFileA(pathbuf, &file_info);\n\n ailego_do_if_false(file != INVALID_HANDLE_VALUE) {\n free(pathbuf);\n FindClose(file);\n return false;\n }\n\n do {\n if (!strcmp(file_info.cFileName, \".\") ||\n !strcmp(file_info.cFileName, \"..\")) {\n continue;\n }\n\n char *fullpath = JoinFilePath(path, file_info.cFileName);\n if (!fullpath) {\n continue;\n }\n\n if (file_info.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {\n FileHelper::RemoveDirectory(fullpath);\n } else {\n FileHelper::DeleteFile(fullpath);\n }\n free(fullpath);\n } while (FindNextFileA(file, &file_info));\n\n free(pathbuf);\n FindClose(file);\n return (!!RemoveDirectoryA(path));\n}\n\nbool FileHelper::IsExist(const char *path) {\n DWORD attr = GetFileAttributesA(path);\n return (attr != INVALID_FILE_ATTRIBUTES);\n}\n\nbool FileHelper::IsRegular(const char *path) {\n DWORD attr = GetFileAttributesA(path);\n return (attr != INVALID_FILE_ATTRIBUTES &&\n !(attr & FILE_ATTRIBUTE_DIRECTORY));\n}\n\nbool FileHelper::IsDirectory(const char *path) {\n DWORD attr = GetFileAttributesA(path);\n return (attr != INVALID_FILE_ATTRIBUTES && (attr & FILE_ATTRIBUTE_DIRECTORY));\n}\n\nbool FileHelper::IsSymbolicLink(const char *path) {\n DWORD attr = GetFileAttributesA(path);\n return (attr != INVALID_FILE_ATTRIBUTES &&\n (attr & FILE_ATTRIBUTE_REPARSE_POINT));\n}\n\nbool FileHelper::IsSame(const char *path1, const char *path2) {\n char real_path1[MAX_PATH];\n char real_path2[MAX_PATH];\n char **part_path1 = nullptr;\n char **part_path2 = nullptr;\n DWORD path1_size =\n GetFullPathNameA(path1, sizeof(real_path1), real_path1, part_path1);\n DWORD path2_size =\n GetFullPathNameA(path2, sizeof(real_path2), real_path2, part_path2);\n\n if ((part_path1 && *part_path1 != 0) || (part_path2 && *part_path2 != 0) ||\n (path1_size != path2_size)) {\n return false;\n }\n return (!strcmp(real_path1, real_path2));\n}\n\n#endif // !_WIN32 && !_WIN64\n\nbool FileHelper::RemovePath(const char *path) {\n if (FileHelper::IsDirectory(path)) {\n return FileHelper::RemoveDirectory(path);\n }\n return FileHelper::DeleteFile(path);\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/utility/float_helper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n\n// #if defined(__F16C__) && defined(__AVX__)\n// #define float16(x) _cvtss_sh((x), _MM_FROUND_NO_EXC)\n// #define float32(x) _cvtsh_ss(x)\n// #endif // __F16C__ && __AVX__\n\n#if defined(__aarch64__)\nstatic inline float float32(uint16_t val) {\n __fp16 *p = reinterpret_cast<__fp16 *>(&val);\n return *p;\n}\n\nstatic inline uint16_t float16(float val) {\n __fp16 f = static_cast<__fp16>(val);\n uint16_t *fp = reinterpret_cast(&f);\n return *fp;\n}\n\nstatic inline void convert_fp16_to_fp32(const uint16_t *arr, size_t size,\n float *out) {\n for (size_t i = 0; i != size; ++i) {\n out[i] = float32(arr[i]);\n }\n}\n\nstatic inline void convert_fp16_to_fp32(const uint16_t *arr, size_t size,\n float norm, float *out) {\n for (size_t i = 0; i != size; ++i) {\n out[i] = float32(arr[i]) / norm;\n }\n}\n\nstatic inline void convert_fp32_to_fp16(const float *arr, size_t size,\n uint16_t *out) {\n for (size_t i = 0; i != size; ++i) {\n out[i] = float16(arr[i]);\n }\n}\n\nstatic inline void convert_fp32_to_fp16(const float *arr, size_t size,\n float norm, uint16_t *out) {\n for (size_t i = 0; i != size; ++i) {\n out[i] = float16(arr[i] / norm);\n }\n}\n#else\n// Refer: https://github.com/Maratyszcza/FP16/blob/master/third-party/half.hpp\nstatic inline float float32(uint16_t val) {\n static const uint32_t mantissa_table[2048] = {\n 0x00000000, 0x33800000, 0x34000000, 0x34400000, 0x34800000, 0x34A00000,\n 0x34C00000, 0x34E00000, 0x35000000, 0x35100000, 0x35200000, 0x35300000,\n 0x35400000, 0x35500000, 0x35600000, 0x35700000, 0x35800000, 0x35880000,\n 0x35900000, 0x35980000, 0x35A00000, 0x35A80000, 0x35B00000, 0x35B80000,\n 0x35C00000, 0x35C80000, 0x35D00000, 0x35D80000, 0x35E00000, 0x35E80000,\n 0x35F00000, 0x35F80000, 0x36000000, 0x36040000, 0x36080000, 0x360C0000,\n 0x36100000, 0x36140000, 0x36180000, 0x361C0000, 0x36200000, 0x36240000,\n 0x36280000, 0x362C0000, 0x36300000, 0x36340000, 0x36380000, 0x363C0000,\n 0x36400000, 0x36440000, 0x36480000, 0x364C0000, 0x36500000, 0x36540000,\n 0x36580000, 0x365C0000, 0x36600000, 0x36640000, 0x36680000, 0x366C0000,\n 0x36700000, 0x36740000, 0x36780000, 0x367C0000, 0x36800000, 0x36820000,\n 0x36840000, 0x36860000, 0x36880000, 0x368A0000, 0x368C0000, 0x368E0000,\n 0x36900000, 0x36920000, 0x36940000, 0x36960000, 0x36980000, 0x369A0000,\n 0x369C0000, 0x369E0000, 0x36A00000, 0x36A20000, 0x36A40000, 0x36A60000,\n 0x36A80000, 0x36AA0000, 0x36AC0000, 0x36AE0000, 0x36B00000, 0x36B20000,\n 0x36B40000, 0x36B60000, 0x36B80000, 0x36BA0000, 0x36BC0000, 0x36BE0000,\n 0x36C00000, 0x36C20000, 0x36C40000, 0x36C60000, 0x36C80000, 0x36CA0000,\n 0x36CC0000, 0x36CE0000, 0x36D00000, 0x36D20000, 0x36D40000, 0x36D60000,\n 0x36D80000, 0x36DA0000, 0x36DC0000, 0x36DE0000, 0x36E00000, 0x36E20000,\n 0x36E40000, 0x36E60000, 0x36E80000, 0x36EA0000, 0x36EC0000, 0x36EE0000,\n 0x36F00000, 0x36F20000, 0x36F40000, 0x36F60000, 0x36F80000, 0x36FA0000,\n 0x36FC0000, 0x36FE0000, 0x37000000, 0x37010000, 0x37020000, 0x37030000,\n 0x37040000, 0x37050000, 0x37060000, 0x37070000, 0x37080000, 0x37090000,\n 0x370A0000, 0x370B0000, 0x370C0000, 0x370D0000, 0x370E0000, 0x370F0000,\n 0x37100000, 0x37110000, 0x37120000, 0x37130000, 0x37140000, 0x37150000,\n 0x37160000, 0x37170000, 0x37180000, 0x37190000, 0x371A0000, 0x371B0000,\n 0x371C0000, 0x371D0000, 0x371E0000, 0x371F0000, 0x37200000, 0x37210000,\n 0x37220000, 0x37230000, 0x37240000, 0x37250000, 0x37260000, 0x37270000,\n 0x37280000, 0x37290000, 0x372A0000, 0x372B0000, 0x372C0000, 0x372D0000,\n 0x372E0000, 0x372F0000, 0x37300000, 0x37310000, 0x37320000, 0x37330000,\n 0x37340000, 0x37350000, 0x37360000, 0x37370000, 0x37380000, 0x37390000,\n 0x373A0000, 0x373B0000, 0x373C0000, 0x373D0000, 0x373E0000, 0x373F0000,\n 0x37400000, 0x37410000, 0x37420000, 0x37430000, 0x37440000, 0x37450000,\n 0x37460000, 0x37470000, 0x37480000, 0x37490000, 0x374A0000, 0x374B0000,\n 0x374C0000, 0x374D0000, 0x374E0000, 0x374F0000, 0x37500000, 0x37510000,\n 0x37520000, 0x37530000, 0x37540000, 0x37550000, 0x37560000, 0x37570000,\n 0x37580000, 0x37590000, 0x375A0000, 0x375B0000, 0x375C0000, 0x375D0000,\n 0x375E0000, 0x375F0000, 0x37600000, 0x37610000, 0x37620000, 0x37630000,\n 0x37640000, 0x37650000, 0x37660000, 0x37670000, 0x37680000, 0x37690000,\n 0x376A0000, 0x376B0000, 0x376C0000, 0x376D0000, 0x376E0000, 0x376F0000,\n 0x37700000, 0x37710000, 0x37720000, 0x37730000, 0x37740000, 0x37750000,\n 0x37760000, 0x37770000, 0x37780000, 0x37790000, 0x377A0000, 0x377B0000,\n 0x377C0000, 0x377D0000, 0x377E0000, 0x377F0000, 0x37800000, 0x37808000,\n 0x37810000, 0x37818000, 0x37820000, 0x37828000, 0x37830000, 0x37838000,\n 0x37840000, 0x37848000, 0x37850000, 0x37858000, 0x37860000, 0x37868000,\n 0x37870000, 0x37878000, 0x37880000, 0x37888000, 0x37890000, 0x37898000,\n 0x378A0000, 0x378A8000, 0x378B0000, 0x378B8000, 0x378C0000, 0x378C8000,\n 0x378D0000, 0x378D8000, 0x378E0000, 0x378E8000, 0x378F0000, 0x378F8000,\n 0x37900000, 0x37908000, 0x37910000, 0x37918000, 0x37920000, 0x37928000,\n 0x37930000, 0x37938000, 0x37940000, 0x37948000, 0x37950000, 0x37958000,\n 0x37960000, 0x37968000, 0x37970000, 0x37978000, 0x37980000, 0x37988000,\n 0x37990000, 0x37998000, 0x379A0000, 0x379A8000, 0x379B0000, 0x379B8000,\n 0x379C0000, 0x379C8000, 0x379D0000, 0x379D8000, 0x379E0000, 0x379E8000,\n 0x379F0000, 0x379F8000, 0x37A00000, 0x37A08000, 0x37A10000, 0x37A18000,\n 0x37A20000, 0x37A28000, 0x37A30000, 0x37A38000, 0x37A40000, 0x37A48000,\n 0x37A50000, 0x37A58000, 0x37A60000, 0x37A68000, 0x37A70000, 0x37A78000,\n 0x37A80000, 0x37A88000, 0x37A90000, 0x37A98000, 0x37AA0000, 0x37AA8000,\n 0x37AB0000, 0x37AB8000, 0x37AC0000, 0x37AC8000, 0x37AD0000, 0x37AD8000,\n 0x37AE0000, 0x37AE8000, 0x37AF0000, 0x37AF8000, 0x37B00000, 0x37B08000,\n 0x37B10000, 0x37B18000, 0x37B20000, 0x37B28000, 0x37B30000, 0x37B38000,\n 0x37B40000, 0x37B48000, 0x37B50000, 0x37B58000, 0x37B60000, 0x37B68000,\n 0x37B70000, 0x37B78000, 0x37B80000, 0x37B88000, 0x37B90000, 0x37B98000,\n 0x37BA0000, 0x37BA8000, 0x37BB0000, 0x37BB8000, 0x37BC0000, 0x37BC8000,\n 0x37BD0000, 0x37BD8000, 0x37BE0000, 0x37BE8000, 0x37BF0000, 0x37BF8000,\n 0x37C00000, 0x37C08000, 0x37C10000, 0x37C18000, 0x37C20000, 0x37C28000,\n 0x37C30000, 0x37C38000, 0x37C40000, 0x37C48000, 0x37C50000, 0x37C58000,\n 0x37C60000, 0x37C68000, 0x37C70000, 0x37C78000, 0x37C80000, 0x37C88000,\n 0x37C90000, 0x37C98000, 0x37CA0000, 0x37CA8000, 0x37CB0000, 0x37CB8000,\n 0x37CC0000, 0x37CC8000, 0x37CD0000, 0x37CD8000, 0x37CE0000, 0x37CE8000,\n 0x37CF0000, 0x37CF8000, 0x37D00000, 0x37D08000, 0x37D10000, 0x37D18000,\n 0x37D20000, 0x37D28000, 0x37D30000, 0x37D38000, 0x37D40000, 0x37D48000,\n 0x37D50000, 0x37D58000, 0x37D60000, 0x37D68000, 0x37D70000, 0x37D78000,\n 0x37D80000, 0x37D88000, 0x37D90000, 0x37D98000, 0x37DA0000, 0x37DA8000,\n 0x37DB0000, 0x37DB8000, 0x37DC0000, 0x37DC8000, 0x37DD0000, 0x37DD8000,\n 0x37DE0000, 0x37DE8000, 0x37DF0000, 0x37DF8000, 0x37E00000, 0x37E08000,\n 0x37E10000, 0x37E18000, 0x37E20000, 0x37E28000, 0x37E30000, 0x37E38000,\n 0x37E40000, 0x37E48000, 0x37E50000, 0x37E58000, 0x37E60000, 0x37E68000,\n 0x37E70000, 0x37E78000, 0x37E80000, 0x37E88000, 0x37E90000, 0x37E98000,\n 0x37EA0000, 0x37EA8000, 0x37EB0000, 0x37EB8000, 0x37EC0000, 0x37EC8000,\n 0x37ED0000, 0x37ED8000, 0x37EE0000, 0x37EE8000, 0x37EF0000, 0x37EF8000,\n 0x37F00000, 0x37F08000, 0x37F10000, 0x37F18000, 0x37F20000, 0x37F28000,\n 0x37F30000, 0x37F38000, 0x37F40000, 0x37F48000, 0x37F50000, 0x37F58000,\n 0x37F60000, 0x37F68000, 0x37F70000, 0x37F78000, 0x37F80000, 0x37F88000,\n 0x37F90000, 0x37F98000, 0x37FA0000, 0x37FA8000, 0x37FB0000, 0x37FB8000,\n 0x37FC0000, 0x37FC8000, 0x37FD0000, 0x37FD8000, 0x37FE0000, 0x37FE8000,\n 0x37FF0000, 0x37FF8000, 0x38000000, 0x38004000, 0x38008000, 0x3800C000,\n 0x38010000, 0x38014000, 0x38018000, 0x3801C000, 0x38020000, 0x38024000,\n 0x38028000, 0x3802C000, 0x38030000, 0x38034000, 0x38038000, 0x3803C000,\n 0x38040000, 0x38044000, 0x38048000, 0x3804C000, 0x38050000, 0x38054000,\n 0x38058000, 0x3805C000, 0x38060000, 0x38064000, 0x38068000, 0x3806C000,\n 0x38070000, 0x38074000, 0x38078000, 0x3807C000, 0x38080000, 0x38084000,\n 0x38088000, 0x3808C000, 0x38090000, 0x38094000, 0x38098000, 0x3809C000,\n 0x380A0000, 0x380A4000, 0x380A8000, 0x380AC000, 0x380B0000, 0x380B4000,\n 0x380B8000, 0x380BC000, 0x380C0000, 0x380C4000, 0x380C8000, 0x380CC000,\n 0x380D0000, 0x380D4000, 0x380D8000, 0x380DC000, 0x380E0000, 0x380E4000,\n 0x380E8000, 0x380EC000, 0x380F0000, 0x380F4000, 0x380F8000, 0x380FC000,\n 0x38100000, 0x38104000, 0x38108000, 0x3810C000, 0x38110000, 0x38114000,\n 0x38118000, 0x3811C000, 0x38120000, 0x38124000, 0x38128000, 0x3812C000,\n 0x38130000, 0x38134000, 0x38138000, 0x3813C000, 0x38140000, 0x38144000,\n 0x38148000, 0x3814C000, 0x38150000, 0x38154000, 0x38158000, 0x3815C000,\n 0x38160000, 0x38164000, 0x38168000, 0x3816C000, 0x38170000, 0x38174000,\n 0x38178000, 0x3817C000, 0x38180000, 0x38184000, 0x38188000, 0x3818C000,\n 0x38190000, 0x38194000, 0x38198000, 0x3819C000, 0x381A0000, 0x381A4000,\n 0x381A8000, 0x381AC000, 0x381B0000, 0x381B4000, 0x381B8000, 0x381BC000,\n 0x381C0000, 0x381C4000, 0x381C8000, 0x381CC000, 0x381D0000, 0x381D4000,\n 0x381D8000, 0x381DC000, 0x381E0000, 0x381E4000, 0x381E8000, 0x381EC000,\n 0x381F0000, 0x381F4000, 0x381F8000, 0x381FC000, 0x38200000, 0x38204000,\n 0x38208000, 0x3820C000, 0x38210000, 0x38214000, 0x38218000, 0x3821C000,\n 0x38220000, 0x38224000, 0x38228000, 0x3822C000, 0x38230000, 0x38234000,\n 0x38238000, 0x3823C000, 0x38240000, 0x38244000, 0x38248000, 0x3824C000,\n 0x38250000, 0x38254000, 0x38258000, 0x3825C000, 0x38260000, 0x38264000,\n 0x38268000, 0x3826C000, 0x38270000, 0x38274000, 0x38278000, 0x3827C000,\n 0x38280000, 0x38284000, 0x38288000, 0x3828C000, 0x38290000, 0x38294000,\n 0x38298000, 0x3829C000, 0x382A0000, 0x382A4000, 0x382A8000, 0x382AC000,\n 0x382B0000, 0x382B4000, 0x382B8000, 0x382BC000, 0x382C0000, 0x382C4000,\n 0x382C8000, 0x382CC000, 0x382D0000, 0x382D4000, 0x382D8000, 0x382DC000,\n 0x382E0000, 0x382E4000, 0x382E8000, 0x382EC000, 0x382F0000, 0x382F4000,\n 0x382F8000, 0x382FC000, 0x38300000, 0x38304000, 0x38308000, 0x3830C000,\n 0x38310000, 0x38314000, 0x38318000, 0x3831C000, 0x38320000, 0x38324000,\n 0x38328000, 0x3832C000, 0x38330000, 0x38334000, 0x38338000, 0x3833C000,\n 0x38340000, 0x38344000, 0x38348000, 0x3834C000, 0x38350000, 0x38354000,\n 0x38358000, 0x3835C000, 0x38360000, 0x38364000, 0x38368000, 0x3836C000,\n 0x38370000, 0x38374000, 0x38378000, 0x3837C000, 0x38380000, 0x38384000,\n 0x38388000, 0x3838C000, 0x38390000, 0x38394000, 0x38398000, 0x3839C000,\n 0x383A0000, 0x383A4000, 0x383A8000, 0x383AC000, 0x383B0000, 0x383B4000,\n 0x383B8000, 0x383BC000, 0x383C0000, 0x383C4000, 0x383C8000, 0x383CC000,\n 0x383D0000, 0x383D4000, 0x383D8000, 0x383DC000, 0x383E0000, 0x383E4000,\n 0x383E8000, 0x383EC000, 0x383F0000, 0x383F4000, 0x383F8000, 0x383FC000,\n 0x38400000, 0x38404000, 0x38408000, 0x3840C000, 0x38410000, 0x38414000,\n 0x38418000, 0x3841C000, 0x38420000, 0x38424000, 0x38428000, 0x3842C000,\n 0x38430000, 0x38434000, 0x38438000, 0x3843C000, 0x38440000, 0x38444000,\n 0x38448000, 0x3844C000, 0x38450000, 0x38454000, 0x38458000, 0x3845C000,\n 0x38460000, 0x38464000, 0x38468000, 0x3846C000, 0x38470000, 0x38474000,\n 0x38478000, 0x3847C000, 0x38480000, 0x38484000, 0x38488000, 0x3848C000,\n 0x38490000, 0x38494000, 0x38498000, 0x3849C000, 0x384A0000, 0x384A4000,\n 0x384A8000, 0x384AC000, 0x384B0000, 0x384B4000, 0x384B8000, 0x384BC000,\n 0x384C0000, 0x384C4000, 0x384C8000, 0x384CC000, 0x384D0000, 0x384D4000,\n 0x384D8000, 0x384DC000, 0x384E0000, 0x384E4000, 0x384E8000, 0x384EC000,\n 0x384F0000, 0x384F4000, 0x384F8000, 0x384FC000, 0x38500000, 0x38504000,\n 0x38508000, 0x3850C000, 0x38510000, 0x38514000, 0x38518000, 0x3851C000,\n 0x38520000, 0x38524000, 0x38528000, 0x3852C000, 0x38530000, 0x38534000,\n 0x38538000, 0x3853C000, 0x38540000, 0x38544000, 0x38548000, 0x3854C000,\n 0x38550000, 0x38554000, 0x38558000, 0x3855C000, 0x38560000, 0x38564000,\n 0x38568000, 0x3856C000, 0x38570000, 0x38574000, 0x38578000, 0x3857C000,\n 0x38580000, 0x38584000, 0x38588000, 0x3858C000, 0x38590000, 0x38594000,\n 0x38598000, 0x3859C000, 0x385A0000, 0x385A4000, 0x385A8000, 0x385AC000,\n 0x385B0000, 0x385B4000, 0x385B8000, 0x385BC000, 0x385C0000, 0x385C4000,\n 0x385C8000, 0x385CC000, 0x385D0000, 0x385D4000, 0x385D8000, 0x385DC000,\n 0x385E0000, 0x385E4000, 0x385E8000, 0x385EC000, 0x385F0000, 0x385F4000,\n 0x385F8000, 0x385FC000, 0x38600000, 0x38604000, 0x38608000, 0x3860C000,\n 0x38610000, 0x38614000, 0x38618000, 0x3861C000, 0x38620000, 0x38624000,\n 0x38628000, 0x3862C000, 0x38630000, 0x38634000, 0x38638000, 0x3863C000,\n 0x38640000, 0x38644000, 0x38648000, 0x3864C000, 0x38650000, 0x38654000,\n 0x38658000, 0x3865C000, 0x38660000, 0x38664000, 0x38668000, 0x3866C000,\n 0x38670000, 0x38674000, 0x38678000, 0x3867C000, 0x38680000, 0x38684000,\n 0x38688000, 0x3868C000, 0x38690000, 0x38694000, 0x38698000, 0x3869C000,\n 0x386A0000, 0x386A4000, 0x386A8000, 0x386AC000, 0x386B0000, 0x386B4000,\n 0x386B8000, 0x386BC000, 0x386C0000, 0x386C4000, 0x386C8000, 0x386CC000,\n 0x386D0000, 0x386D4000, 0x386D8000, 0x386DC000, 0x386E0000, 0x386E4000,\n 0x386E8000, 0x386EC000, 0x386F0000, 0x386F4000, 0x386F8000, 0x386FC000,\n 0x38700000, 0x38704000, 0x38708000, 0x3870C000, 0x38710000, 0x38714000,\n 0x38718000, 0x3871C000, 0x38720000, 0x38724000, 0x38728000, 0x3872C000,\n 0x38730000, 0x38734000, 0x38738000, 0x3873C000, 0x38740000, 0x38744000,\n 0x38748000, 0x3874C000, 0x38750000, 0x38754000, 0x38758000, 0x3875C000,\n 0x38760000, 0x38764000, 0x38768000, 0x3876C000, 0x38770000, 0x38774000,\n 0x38778000, 0x3877C000, 0x38780000, 0x38784000, 0x38788000, 0x3878C000,\n 0x38790000, 0x38794000, 0x38798000, 0x3879C000, 0x387A0000, 0x387A4000,\n 0x387A8000, 0x387AC000, 0x387B0000, 0x387B4000, 0x387B8000, 0x387BC000,\n 0x387C0000, 0x387C4000, 0x387C8000, 0x387CC000, 0x387D0000, 0x387D4000,\n 0x387D8000, 0x387DC000, 0x387E0000, 0x387E4000, 0x387E8000, 0x387EC000,\n 0x387F0000, 0x387F4000, 0x387F8000, 0x387FC000, 0x38000000, 0x38002000,\n 0x38004000, 0x38006000, 0x38008000, 0x3800A000, 0x3800C000, 0x3800E000,\n 0x38010000, 0x38012000, 0x38014000, 0x38016000, 0x38018000, 0x3801A000,\n 0x3801C000, 0x3801E000, 0x38020000, 0x38022000, 0x38024000, 0x38026000,\n 0x38028000, 0x3802A000, 0x3802C000, 0x3802E000, 0x38030000, 0x38032000,\n 0x38034000, 0x38036000, 0x38038000, 0x3803A000, 0x3803C000, 0x3803E000,\n 0x38040000, 0x38042000, 0x38044000, 0x38046000, 0x38048000, 0x3804A000,\n 0x3804C000, 0x3804E000, 0x38050000, 0x38052000, 0x38054000, 0x38056000,\n 0x38058000, 0x3805A000, 0x3805C000, 0x3805E000, 0x38060000, 0x38062000,\n 0x38064000, 0x38066000, 0x38068000, 0x3806A000, 0x3806C000, 0x3806E000,\n 0x38070000, 0x38072000, 0x38074000, 0x38076000, 0x38078000, 0x3807A000,\n 0x3807C000, 0x3807E000, 0x38080000, 0x38082000, 0x38084000, 0x38086000,\n 0x38088000, 0x3808A000, 0x3808C000, 0x3808E000, 0x38090000, 0x38092000,\n 0x38094000, 0x38096000, 0x38098000, 0x3809A000, 0x3809C000, 0x3809E000,\n 0x380A0000, 0x380A2000, 0x380A4000, 0x380A6000, 0x380A8000, 0x380AA000,\n 0x380AC000, 0x380AE000, 0x380B0000, 0x380B2000, 0x380B4000, 0x380B6000,\n 0x380B8000, 0x380BA000, 0x380BC000, 0x380BE000, 0x380C0000, 0x380C2000,\n 0x380C4000, 0x380C6000, 0x380C8000, 0x380CA000, 0x380CC000, 0x380CE000,\n 0x380D0000, 0x380D2000, 0x380D4000, 0x380D6000, 0x380D8000, 0x380DA000,\n 0x380DC000, 0x380DE000, 0x380E0000, 0x380E2000, 0x380E4000, 0x380E6000,\n 0x380E8000, 0x380EA000, 0x380EC000, 0x380EE000, 0x380F0000, 0x380F2000,\n 0x380F4000, 0x380F6000, 0x380F8000, 0x380FA000, 0x380FC000, 0x380FE000,\n 0x38100000, 0x38102000, 0x38104000, 0x38106000, 0x38108000, 0x3810A000,\n 0x3810C000, 0x3810E000, 0x38110000, 0x38112000, 0x38114000, 0x38116000,\n 0x38118000, 0x3811A000, 0x3811C000, 0x3811E000, 0x38120000, 0x38122000,\n 0x38124000, 0x38126000, 0x38128000, 0x3812A000, 0x3812C000, 0x3812E000,\n 0x38130000, 0x38132000, 0x38134000, 0x38136000, 0x38138000, 0x3813A000,\n 0x3813C000, 0x3813E000, 0x38140000, 0x38142000, 0x38144000, 0x38146000,\n 0x38148000, 0x3814A000, 0x3814C000, 0x3814E000, 0x38150000, 0x38152000,\n 0x38154000, 0x38156000, 0x38158000, 0x3815A000, 0x3815C000, 0x3815E000,\n 0x38160000, 0x38162000, 0x38164000, 0x38166000, 0x38168000, 0x3816A000,\n 0x3816C000, 0x3816E000, 0x38170000, 0x38172000, 0x38174000, 0x38176000,\n 0x38178000, 0x3817A000, 0x3817C000, 0x3817E000, 0x38180000, 0x38182000,\n 0x38184000, 0x38186000, 0x38188000, 0x3818A000, 0x3818C000, 0x3818E000,\n 0x38190000, 0x38192000, 0x38194000, 0x38196000, 0x38198000, 0x3819A000,\n 0x3819C000, 0x3819E000, 0x381A0000, 0x381A2000, 0x381A4000, 0x381A6000,\n 0x381A8000, 0x381AA000, 0x381AC000, 0x381AE000, 0x381B0000, 0x381B2000,\n 0x381B4000, 0x381B6000, 0x381B8000, 0x381BA000, 0x381BC000, 0x381BE000,\n 0x381C0000, 0x381C2000, 0x381C4000, 0x381C6000, 0x381C8000, 0x381CA000,\n 0x381CC000, 0x381CE000, 0x381D0000, 0x381D2000, 0x381D4000, 0x381D6000,\n 0x381D8000, 0x381DA000, 0x381DC000, 0x381DE000, 0x381E0000, 0x381E2000,\n 0x381E4000, 0x381E6000, 0x381E8000, 0x381EA000, 0x381EC000, 0x381EE000,\n 0x381F0000, 0x381F2000, 0x381F4000, 0x381F6000, 0x381F8000, 0x381FA000,\n 0x381FC000, 0x381FE000, 0x38200000, 0x38202000, 0x38204000, 0x38206000,\n 0x38208000, 0x3820A000, 0x3820C000, 0x3820E000, 0x38210000, 0x38212000,\n 0x38214000, 0x38216000, 0x38218000, 0x3821A000, 0x3821C000, 0x3821E000,\n 0x38220000, 0x38222000, 0x38224000, 0x38226000, 0x38228000, 0x3822A000,\n 0x3822C000, 0x3822E000, 0x38230000, 0x38232000, 0x38234000, 0x38236000,\n 0x38238000, 0x3823A000, 0x3823C000, 0x3823E000, 0x38240000, 0x38242000,\n 0x38244000, 0x38246000, 0x38248000, 0x3824A000, 0x3824C000, 0x3824E000,\n 0x38250000, 0x38252000, 0x38254000, 0x38256000, 0x38258000, 0x3825A000,\n 0x3825C000, 0x3825E000, 0x38260000, 0x38262000, 0x38264000, 0x38266000,\n 0x38268000, 0x3826A000, 0x3826C000, 0x3826E000, 0x38270000, 0x38272000,\n 0x38274000, 0x38276000, 0x38278000, 0x3827A000, 0x3827C000, 0x3827E000,\n 0x38280000, 0x38282000, 0x38284000, 0x38286000, 0x38288000, 0x3828A000,\n 0x3828C000, 0x3828E000, 0x38290000, 0x38292000, 0x38294000, 0x38296000,\n 0x38298000, 0x3829A000, 0x3829C000, 0x3829E000, 0x382A0000, 0x382A2000,\n 0x382A4000, 0x382A6000, 0x382A8000, 0x382AA000, 0x382AC000, 0x382AE000,\n 0x382B0000, 0x382B2000, 0x382B4000, 0x382B6000, 0x382B8000, 0x382BA000,\n 0x382BC000, 0x382BE000, 0x382C0000, 0x382C2000, 0x382C4000, 0x382C6000,\n 0x382C8000, 0x382CA000, 0x382CC000, 0x382CE000, 0x382D0000, 0x382D2000,\n 0x382D4000, 0x382D6000, 0x382D8000, 0x382DA000, 0x382DC000, 0x382DE000,\n 0x382E0000, 0x382E2000, 0x382E4000, 0x382E6000, 0x382E8000, 0x382EA000,\n 0x382EC000, 0x382EE000, 0x382F0000, 0x382F2000, 0x382F4000, 0x382F6000,\n 0x382F8000, 0x382FA000, 0x382FC000, 0x382FE000, 0x38300000, 0x38302000,\n 0x38304000, 0x38306000, 0x38308000, 0x3830A000, 0x3830C000, 0x3830E000,\n 0x38310000, 0x38312000, 0x38314000, 0x38316000, 0x38318000, 0x3831A000,\n 0x3831C000, 0x3831E000, 0x38320000, 0x38322000, 0x38324000, 0x38326000,\n 0x38328000, 0x3832A000, 0x3832C000, 0x3832E000, 0x38330000, 0x38332000,\n 0x38334000, 0x38336000, 0x38338000, 0x3833A000, 0x3833C000, 0x3833E000,\n 0x38340000, 0x38342000, 0x38344000, 0x38346000, 0x38348000, 0x3834A000,\n 0x3834C000, 0x3834E000, 0x38350000, 0x38352000, 0x38354000, 0x38356000,\n 0x38358000, 0x3835A000, 0x3835C000, 0x3835E000, 0x38360000, 0x38362000,\n 0x38364000, 0x38366000, 0x38368000, 0x3836A000, 0x3836C000, 0x3836E000,\n 0x38370000, 0x38372000, 0x38374000, 0x38376000, 0x38378000, 0x3837A000,\n 0x3837C000, 0x3837E000, 0x38380000, 0x38382000, 0x38384000, 0x38386000,\n 0x38388000, 0x3838A000, 0x3838C000, 0x3838E000, 0x38390000, 0x38392000,\n 0x38394000, 0x38396000, 0x38398000, 0x3839A000, 0x3839C000, 0x3839E000,\n 0x383A0000, 0x383A2000, 0x383A4000, 0x383A6000, 0x383A8000, 0x383AA000,\n 0x383AC000, 0x383AE000, 0x383B0000, 0x383B2000, 0x383B4000, 0x383B6000,\n 0x383B8000, 0x383BA000, 0x383BC000, 0x383BE000, 0x383C0000, 0x383C2000,\n 0x383C4000, 0x383C6000, 0x383C8000, 0x383CA000, 0x383CC000, 0x383CE000,\n 0x383D0000, 0x383D2000, 0x383D4000, 0x383D6000, 0x383D8000, 0x383DA000,\n 0x383DC000, 0x383DE000, 0x383E0000, 0x383E2000, 0x383E4000, 0x383E6000,\n 0x383E8000, 0x383EA000, 0x383EC000, 0x383EE000, 0x383F0000, 0x383F2000,\n 0x383F4000, 0x383F6000, 0x383F8000, 0x383FA000, 0x383FC000, 0x383FE000,\n 0x38400000, 0x38402000, 0x38404000, 0x38406000, 0x38408000, 0x3840A000,\n 0x3840C000, 0x3840E000, 0x38410000, 0x38412000, 0x38414000, 0x38416000,\n 0x38418000, 0x3841A000, 0x3841C000, 0x3841E000, 0x38420000, 0x38422000,\n 0x38424000, 0x38426000, 0x38428000, 0x3842A000, 0x3842C000, 0x3842E000,\n 0x38430000, 0x38432000, 0x38434000, 0x38436000, 0x38438000, 0x3843A000,\n 0x3843C000, 0x3843E000, 0x38440000, 0x38442000, 0x38444000, 0x38446000,\n 0x38448000, 0x3844A000, 0x3844C000, 0x3844E000, 0x38450000, 0x38452000,\n 0x38454000, 0x38456000, 0x38458000, 0x3845A000, 0x3845C000, 0x3845E000,\n 0x38460000, 0x38462000, 0x38464000, 0x38466000, 0x38468000, 0x3846A000,\n 0x3846C000, 0x3846E000, 0x38470000, 0x38472000, 0x38474000, 0x38476000,\n 0x38478000, 0x3847A000, 0x3847C000, 0x3847E000, 0x38480000, 0x38482000,\n 0x38484000, 0x38486000, 0x38488000, 0x3848A000, 0x3848C000, 0x3848E000,\n 0x38490000, 0x38492000, 0x38494000, 0x38496000, 0x38498000, 0x3849A000,\n 0x3849C000, 0x3849E000, 0x384A0000, 0x384A2000, 0x384A4000, 0x384A6000,\n 0x384A8000, 0x384AA000, 0x384AC000, 0x384AE000, 0x384B0000, 0x384B2000,\n 0x384B4000, 0x384B6000, 0x384B8000, 0x384BA000, 0x384BC000, 0x384BE000,\n 0x384C0000, 0x384C2000, 0x384C4000, 0x384C6000, 0x384C8000, 0x384CA000,\n 0x384CC000, 0x384CE000, 0x384D0000, 0x384D2000, 0x384D4000, 0x384D6000,\n 0x384D8000, 0x384DA000, 0x384DC000, 0x384DE000, 0x384E0000, 0x384E2000,\n 0x384E4000, 0x384E6000, 0x384E8000, 0x384EA000, 0x384EC000, 0x384EE000,\n 0x384F0000, 0x384F2000, 0x384F4000, 0x384F6000, 0x384F8000, 0x384FA000,\n 0x384FC000, 0x384FE000, 0x38500000, 0x38502000, 0x38504000, 0x38506000,\n 0x38508000, 0x3850A000, 0x3850C000, 0x3850E000, 0x38510000, 0x38512000,\n 0x38514000, 0x38516000, 0x38518000, 0x3851A000, 0x3851C000, 0x3851E000,\n 0x38520000, 0x38522000, 0x38524000, 0x38526000, 0x38528000, 0x3852A000,\n 0x3852C000, 0x3852E000, 0x38530000, 0x38532000, 0x38534000, 0x38536000,\n 0x38538000, 0x3853A000, 0x3853C000, 0x3853E000, 0x38540000, 0x38542000,\n 0x38544000, 0x38546000, 0x38548000, 0x3854A000, 0x3854C000, 0x3854E000,\n 0x38550000, 0x38552000, 0x38554000, 0x38556000, 0x38558000, 0x3855A000,\n 0x3855C000, 0x3855E000, 0x38560000, 0x38562000, 0x38564000, 0x38566000,\n 0x38568000, 0x3856A000, 0x3856C000, 0x3856E000, 0x38570000, 0x38572000,\n 0x38574000, 0x38576000, 0x38578000, 0x3857A000, 0x3857C000, 0x3857E000,\n 0x38580000, 0x38582000, 0x38584000, 0x38586000, 0x38588000, 0x3858A000,\n 0x3858C000, 0x3858E000, 0x38590000, 0x38592000, 0x38594000, 0x38596000,\n 0x38598000, 0x3859A000, 0x3859C000, 0x3859E000, 0x385A0000, 0x385A2000,\n 0x385A4000, 0x385A6000, 0x385A8000, 0x385AA000, 0x385AC000, 0x385AE000,\n 0x385B0000, 0x385B2000, 0x385B4000, 0x385B6000, 0x385B8000, 0x385BA000,\n 0x385BC000, 0x385BE000, 0x385C0000, 0x385C2000, 0x385C4000, 0x385C6000,\n 0x385C8000, 0x385CA000, 0x385CC000, 0x385CE000, 0x385D0000, 0x385D2000,\n 0x385D4000, 0x385D6000, 0x385D8000, 0x385DA000, 0x385DC000, 0x385DE000,\n 0x385E0000, 0x385E2000, 0x385E4000, 0x385E6000, 0x385E8000, 0x385EA000,\n 0x385EC000, 0x385EE000, 0x385F0000, 0x385F2000, 0x385F4000, 0x385F6000,\n 0x385F8000, 0x385FA000, 0x385FC000, 0x385FE000, 0x38600000, 0x38602000,\n 0x38604000, 0x38606000, 0x38608000, 0x3860A000, 0x3860C000, 0x3860E000,\n 0x38610000, 0x38612000, 0x38614000, 0x38616000, 0x38618000, 0x3861A000,\n 0x3861C000, 0x3861E000, 0x38620000, 0x38622000, 0x38624000, 0x38626000,\n 0x38628000, 0x3862A000, 0x3862C000, 0x3862E000, 0x38630000, 0x38632000,\n 0x38634000, 0x38636000, 0x38638000, 0x3863A000, 0x3863C000, 0x3863E000,\n 0x38640000, 0x38642000, 0x38644000, 0x38646000, 0x38648000, 0x3864A000,\n 0x3864C000, 0x3864E000, 0x38650000, 0x38652000, 0x38654000, 0x38656000,\n 0x38658000, 0x3865A000, 0x3865C000, 0x3865E000, 0x38660000, 0x38662000,\n 0x38664000, 0x38666000, 0x38668000, 0x3866A000, 0x3866C000, 0x3866E000,\n 0x38670000, 0x38672000, 0x38674000, 0x38676000, 0x38678000, 0x3867A000,\n 0x3867C000, 0x3867E000, 0x38680000, 0x38682000, 0x38684000, 0x38686000,\n 0x38688000, 0x3868A000, 0x3868C000, 0x3868E000, 0x38690000, 0x38692000,\n 0x38694000, 0x38696000, 0x38698000, 0x3869A000, 0x3869C000, 0x3869E000,\n 0x386A0000, 0x386A2000, 0x386A4000, 0x386A6000, 0x386A8000, 0x386AA000,\n 0x386AC000, 0x386AE000, 0x386B0000, 0x386B2000, 0x386B4000, 0x386B6000,\n 0x386B8000, 0x386BA000, 0x386BC000, 0x386BE000, 0x386C0000, 0x386C2000,\n 0x386C4000, 0x386C6000, 0x386C8000, 0x386CA000, 0x386CC000, 0x386CE000,\n 0x386D0000, 0x386D2000, 0x386D4000, 0x386D6000, 0x386D8000, 0x386DA000,\n 0x386DC000, 0x386DE000, 0x386E0000, 0x386E2000, 0x386E4000, 0x386E6000,\n 0x386E8000, 0x386EA000, 0x386EC000, 0x386EE000, 0x386F0000, 0x386F2000,\n 0x386F4000, 0x386F6000, 0x386F8000, 0x386FA000, 0x386FC000, 0x386FE000,\n 0x38700000, 0x38702000, 0x38704000, 0x38706000, 0x38708000, 0x3870A000,\n 0x3870C000, 0x3870E000, 0x38710000, 0x38712000, 0x38714000, 0x38716000,\n 0x38718000, 0x3871A000, 0x3871C000, 0x3871E000, 0x38720000, 0x38722000,\n 0x38724000, 0x38726000, 0x38728000, 0x3872A000, 0x3872C000, 0x3872E000,\n 0x38730000, 0x38732000, 0x38734000, 0x38736000, 0x38738000, 0x3873A000,\n 0x3873C000, 0x3873E000, 0x38740000, 0x38742000, 0x38744000, 0x38746000,\n 0x38748000, 0x3874A000, 0x3874C000, 0x3874E000, 0x38750000, 0x38752000,\n 0x38754000, 0x38756000, 0x38758000, 0x3875A000, 0x3875C000, 0x3875E000,\n 0x38760000, 0x38762000, 0x38764000, 0x38766000, 0x38768000, 0x3876A000,\n 0x3876C000, 0x3876E000, 0x38770000, 0x38772000, 0x38774000, 0x38776000,\n 0x38778000, 0x3877A000, 0x3877C000, 0x3877E000, 0x38780000, 0x38782000,\n 0x38784000, 0x38786000, 0x38788000, 0x3878A000, 0x3878C000, 0x3878E000,\n 0x38790000, 0x38792000, 0x38794000, 0x38796000, 0x38798000, 0x3879A000,\n 0x3879C000, 0x3879E000, 0x387A0000, 0x387A2000, 0x387A4000, 0x387A6000,\n 0x387A8000, 0x387AA000, 0x387AC000, 0x387AE000, 0x387B0000, 0x387B2000,\n 0x387B4000, 0x387B6000, 0x387B8000, 0x387BA000, 0x387BC000, 0x387BE000,\n 0x387C0000, 0x387C2000, 0x387C4000, 0x387C6000, 0x387C8000, 0x387CA000,\n 0x387CC000, 0x387CE000, 0x387D0000, 0x387D2000, 0x387D4000, 0x387D6000,\n 0x387D8000, 0x387DA000, 0x387DC000, 0x387DE000, 0x387E0000, 0x387E2000,\n 0x387E4000, 0x387E6000, 0x387E8000, 0x387EA000, 0x387EC000, 0x387EE000,\n 0x387F0000, 0x387F2000, 0x387F4000, 0x387F6000, 0x387F8000, 0x387FA000,\n 0x387FC000, 0x387FE000};\n static const uint32_t exponent_table[64] = {\n 0x00000000, 0x00800000, 0x01000000, 0x01800000, 0x02000000, 0x02800000,\n 0x03000000, 0x03800000, 0x04000000, 0x04800000, 0x05000000, 0x05800000,\n 0x06000000, 0x06800000, 0x07000000, 0x07800000, 0x08000000, 0x08800000,\n 0x09000000, 0x09800000, 0x0A000000, 0x0A800000, 0x0B000000, 0x0B800000,\n 0x0C000000, 0x0C800000, 0x0D000000, 0x0D800000, 0x0E000000, 0x0E800000,\n 0x0F000000, 0x47800000, 0x80000000, 0x80800000, 0x81000000, 0x81800000,\n 0x82000000, 0x82800000, 0x83000000, 0x83800000, 0x84000000, 0x84800000,\n 0x85000000, 0x85800000, 0x86000000, 0x86800000, 0x87000000, 0x87800000,\n 0x88000000, 0x88800000, 0x89000000, 0x89800000, 0x8A000000, 0x8A800000,\n 0x8B000000, 0x8B800000, 0x8C000000, 0x8C800000, 0x8D000000, 0x8D800000,\n 0x8E000000, 0x8E800000, 0x8F000000, 0xC7800000};\n static const uint16_t offset_table[64] = {\n 0, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024,\n 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024,\n 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 0,\n 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024,\n 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024,\n 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024};\n uint16_t hval = static_cast(val >> 10);\n uint32_t bits =\n mantissa_table[offset_table[hval] + (val & 0x3FF)] + exponent_table[hval];\n float *p = reinterpret_cast(&bits);\n return (*p);\n}\n\n// Refer: https://github.com/Maratyszcza/FP16/blob/master/third-party/half.hpp\nstatic inline uint16_t float16(float val) {\n static const uint16_t base_table[512] = {\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000, 0x0000,\n 0x0000, 0x0000, 0x0000, 0x0000, 0x0001, 0x0002, 0x0004, 0x0008, 0x0010,\n 0x0020, 0x0040, 0x0080, 0x0100, 0x0200, 0x0400, 0x0800, 0x0C00, 0x1000,\n 0x1400, 0x1800, 0x1C00, 0x2000, 0x2400, 0x2800, 0x2C00, 0x3000, 0x3400,\n 0x3800, 0x3C00, 0x4000, 0x4400, 0x4800, 0x4C00, 0x5000, 0x5400, 0x5800,\n 0x5C00, 0x6000, 0x6400, 0x6800, 0x6C00, 0x7000, 0x7400, 0x7800, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x7C00,\n 0x7C00, 0x7C00, 0x7C00, 0x7C00, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000,\n 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8000, 0x8001,\n 0x8002, 0x8004, 0x8008, 0x8010, 0x8020, 0x8040, 0x8080, 0x8100, 0x8200,\n 0x8400, 0x8800, 0x8C00, 0x9000, 0x9400, 0x9800, 0x9C00, 0xA000, 0xA400,\n 0xA800, 0xAC00, 0xB000, 0xB400, 0xB800, 0xBC00, 0xC000, 0xC400, 0xC800,\n 0xCC00, 0xD000, 0xD400, 0xD800, 0xDC00, 0xE000, 0xE400, 0xE800, 0xEC00,\n 0xF000, 0xF400, 0xF800, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00,\n 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00, 0xFC00};\n static const uint8_t shift_table[512] = {\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 23, 22, 21, 20, 19,\n 18, 17, 16, 15, 14, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13,\n 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 13, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 23,\n 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 13, 13, 13, 13, 13, 13, 13, 13,\n 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13,\n 13, 13, 13, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,\n 24, 24, 24, 24, 24, 24, 24, 13};\n uint32_t *p = reinterpret_cast(&val);\n uint32_t hbits =\n base_table[*p >> 23] +\n static_cast((*p & 0x7FFFFF) >> shift_table[*p >> 23]);\n hbits += (((*p & 0x7FFFFF) >> (shift_table[*p >> 23] - 1)) |\n (((*p >> 23) & 0xFF) == 102)) &\n ((hbits & 0x7C00) != 0x7C00);\n return static_cast(hbits);\n}\n#if defined(__F16C__) && defined(__AVX512F__)\nstatic inline void convert_fp16_to_fp32_avx512f(const uint16_t *arr,\n size_t size, float *out) {\n const uint16_t *last = arr + size;\n const uint16_t *last_aligned = arr + ((size >> 5) << 5);\n\n if (((uintptr_t)arr & 0x1f) == 0 && ((uintptr_t)out & 0x3f) == 0) {\n for (; arr != last_aligned; arr += 32, out += 32) {\n _mm512_store_ps(out + 0,\n _mm512_cvtph_ps(_mm256_load_si256((__m256i *)(arr + 0))));\n _mm512_store_ps(\n out + 16, _mm512_cvtph_ps(_mm256_load_si256((__m256i *)(arr + 16))));\n }\n\n if (last >= last_aligned + 16) {\n _mm512_store_ps(out, _mm512_cvtph_ps(_mm256_load_si256((__m256i *)arr)));\n arr += 16;\n out += 16;\n }\n if (last >= arr + 8) {\n _mm256_store_ps(out, _mm256_cvtph_ps(_mm_load_si128((__m128i *)arr)));\n arr += 8;\n out += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 32, out += 32) {\n _mm512_storeu_ps(\n out + 0, _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)(arr + 0))));\n _mm512_storeu_ps(\n out + 16, _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)(arr + 16))));\n }\n\n if (last >= last_aligned + 16) {\n _mm512_storeu_ps(out,\n _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)arr)));\n arr += 16;\n out += 16;\n }\n if (last >= arr + 8) {\n _mm256_storeu_ps(out, _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)arr)));\n arr += 8;\n out += 8;\n }\n }\n switch (last - arr) {\n case 7:\n out[6] = float32(arr[6]);\n /* FALLTHRU */\n case 6:\n out[5] = float32(arr[5]);\n /* FALLTHRU */\n case 5:\n out[4] = float32(arr[4]);\n /* FALLTHRU */\n case 4:\n out[3] = float32(arr[3]);\n /* FALLTHRU */\n case 3:\n out[2] = float32(arr[2]);\n /* FALLTHRU */\n case 2:\n out[1] = float32(arr[1]);\n /* FALLTHRU */\n case 1:\n out[0] = float32(arr[0]);\n }\n}\n\nstatic inline void convert_fp16_to_fp32_avx512f(const uint16_t *arr,\n size_t size, float norm,\n float *out) {\n const uint16_t *last = arr + size;\n const uint16_t *last_aligned = arr + ((size >> 5) << 5);\n __m512 zmm_norm = _mm512_set1_ps(norm);\n\n if (((uintptr_t)arr & 0x1f) == 0 && ((uintptr_t)out & 0x3f) == 0) {\n for (; arr != last_aligned; arr += 32, out += 32) {\n __m512 zmm_0 = _mm512_div_ps(\n _mm512_cvtph_ps(_mm256_load_si256((__m256i *)(arr + 0))), zmm_norm);\n __m512 zmm_1 = _mm512_div_ps(\n _mm512_cvtph_ps(_mm256_load_si256((__m256i *)(arr + 16))), zmm_norm);\n _mm512_store_ps(out + 0, zmm_0);\n _mm512_store_ps(out + 16, zmm_1);\n }\n\n if (last >= last_aligned + 16) {\n _mm512_store_ps(\n out, _mm512_div_ps(_mm512_cvtph_ps(_mm256_load_si256((__m256i *)arr)),\n zmm_norm));\n arr += 16;\n out += 16;\n }\n if (last >= arr + 8) {\n _mm256_store_ps(\n out, _mm256_div_ps(_mm256_cvtph_ps(_mm_load_si128((__m128i *)arr)),\n _mm256_set1_ps(norm)));\n arr += 8;\n out += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 32, out += 32) {\n __m512 zmm_0 = _mm512_div_ps(\n _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)(arr + 0))), zmm_norm);\n __m512 zmm_1 = _mm512_div_ps(\n _mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)(arr + 16))), zmm_norm);\n _mm512_storeu_ps(out + 0, zmm_0);\n _mm512_storeu_ps(out + 16, zmm_1);\n }\n\n if (last >= last_aligned + 16) {\n _mm512_storeu_ps(\n out,\n _mm512_div_ps(_mm512_cvtph_ps(_mm256_loadu_si256((__m256i *)arr)),\n zmm_norm));\n arr += 16;\n out += 16;\n }\n if (last >= arr + 8) {\n _mm256_storeu_ps(\n out, _mm256_div_ps(_mm256_cvtph_ps(_mm_loadu_si128((__m128i *)arr)),\n _mm256_set1_ps(norm)));\n arr += 8;\n out += 8;\n }\n }\n switch (last - arr) {\n case 7:\n out[6] = float32(arr[6]) / norm;\n /* FALLTHRU */\n case 6:\n out[5] = float32(arr[5]) / norm;\n /* FALLTHRU */\n case 5:\n out[4] = float32(arr[4]) / norm;\n /* FALLTHRU */\n case 4:\n out[3] = float32(arr[3]) / norm;\n /* FALLTHRU */\n case 3:\n out[2] = float32(arr[2]) / norm;\n /* FALLTHRU */\n case 2:\n out[1] = float32(arr[1]) / norm;\n /* FALLTHRU */\n case 1:\n out[0] = float32(arr[0]) / norm;\n }\n}\n\nstatic inline void convert_fp32_to_fp16_avx512f(const float *arr, size_t size,\n uint16_t *out) {\n const float *last = arr + size;\n const float *last_aligned = arr + ((size >> 5) << 5);\n\n if (((uintptr_t)arr & 0x3f) == 0 && ((uintptr_t)out & 0x1f) == 0) {\n for (; arr != last_aligned; arr += 32, out += 32) {\n _mm256_store_si256(\n (__m256i *)(out + 0),\n _mm512_cvtps_ph(_mm512_load_ps(arr + 0), _MM_FROUND_NO_EXC));\n _mm256_store_si256(\n (__m256i *)(out + 16),\n _mm512_cvtps_ph(_mm512_load_ps(arr + 16), _MM_FROUND_NO_EXC));\n }\n\n if (last >= last_aligned + 16) {\n _mm256_store_si256(\n (__m256i *)(out + 0),\n _mm512_cvtps_ph(_mm512_load_ps(arr + 0), _MM_FROUND_NO_EXC));\n arr += 16;\n out += 16;\n }\n if (last >= arr + 8) {\n _mm_store_si128(\n (__m128i *)(out + 0),\n _mm256_cvtps_ph(_mm256_load_ps(arr + 0), _MM_FROUND_NO_EXC));\n arr += 8;\n out += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 32, out += 32) {\n _mm256_storeu_si256(\n (__m256i *)(out + 0),\n _mm512_cvtps_ph(_mm512_loadu_ps(arr + 0), _MM_FROUND_NO_EXC));\n _mm256_storeu_si256(\n (__m256i *)(out + 16),\n _mm512_cvtps_ph(_mm512_loadu_ps(arr + 16), _MM_FROUND_NO_EXC));\n }\n\n if (last >= last_aligned + 16) {\n _mm256_storeu_si256(\n (__m256i *)(out + 0),\n _mm512_cvtps_ph(_mm512_loadu_ps(arr + 0), _MM_FROUND_NO_EXC));\n arr += 16;\n out += 16;\n }\n if (last >= arr + 8) {\n _mm_storeu_si128(\n (__m128i *)(out + 0),\n _mm256_cvtps_ph(_mm256_loadu_ps(arr + 0), _MM_FROUND_NO_EXC));\n arr += 8;\n out += 8;\n }\n }\n switch (last - arr) {\n case 7:\n out[6] = float16(arr[6]);\n /* FALLTHRU */\n case 6:\n out[5] = float16(arr[5]);\n /* FALLTHRU */\n case 5:\n out[4] = float16(arr[4]);\n /* FALLTHRU */\n case 4:\n out[3] = float16(arr[3]);\n /* FALLTHRU */\n case 3:\n out[2] = float16(arr[2]);\n /* FALLTHRU */\n case 2:\n out[1] = float16(arr[1]);\n /* FALLTHRU */\n case 1:\n out[0] = float16(arr[0]);\n }\n}\n\nstatic inline void convert_fp32_to_fp16_avx512f(const float *arr, size_t size,\n float norm, uint16_t *out) {\n const float *last = arr + size;\n const float *last_aligned = arr + ((size >> 5) << 5);\n __m512 zmm_norm = _mm512_set1_ps(norm);\n\n if (((uintptr_t)arr & 0x3f) == 0 && ((uintptr_t)out & 0x1f) == 0) {\n for (; arr != last_aligned; arr += 32, out += 32) {\n __m512 zmm_0 = _mm512_div_ps(_mm512_load_ps(arr + 0), zmm_norm);\n __m512 zmm_1 = _mm512_div_ps(_mm512_load_ps(arr + 16), zmm_norm);\n _mm256_store_si256((__m256i *)(out + 0),\n _mm512_cvtps_ph(zmm_0, _MM_FROUND_NO_EXC));\n _mm256_store_si256((__m256i *)(out + 16),\n _mm512_cvtps_ph(zmm_1, _MM_FROUND_NO_EXC));\n }\n\n if (last >= last_aligned + 16) {\n _mm256_store_si256(\n (__m256i *)out,\n _mm512_cvtps_ph(_mm512_div_ps(_mm512_load_ps(arr), zmm_norm),\n _MM_FROUND_NO_EXC));\n arr += 16;\n out += 16;\n }\n if (last >= arr + 8) {\n _mm_store_si128((__m128i *)out,\n _mm256_cvtps_ph(_mm256_div_ps(_mm256_load_ps(arr),\n _mm256_set1_ps(norm)),\n _MM_FROUND_NO_EXC));\n arr += 8;\n out += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 32, out += 32) {\n __m512 zmm_0 = _mm512_div_ps(_mm512_loadu_ps(arr + 0), zmm_norm);\n __m512 zmm_1 = _mm512_div_ps(_mm512_loadu_ps(arr + 16), zmm_norm);\n _mm256_storeu_si256((__m256i *)(out + 0),\n _mm512_cvtps_ph(zmm_0, _MM_FROUND_NO_EXC));\n _mm256_storeu_si256((__m256i *)(out + 16),\n _mm512_cvtps_ph(zmm_1, _MM_FROUND_NO_EXC));\n }\n\n if (last >= last_aligned + 16) {\n _mm256_storeu_si256(\n (__m256i *)out,\n _mm512_cvtps_ph(_mm512_div_ps(_mm512_loadu_ps(arr), zmm_norm),\n _MM_FROUND_NO_EXC));\n arr += 16;\n out += 16;\n }\n if (last >= arr + 8) {\n _mm_storeu_si128((__m128i *)out,\n _mm256_cvtps_ph(_mm256_div_ps(_mm256_loadu_ps(arr),\n _mm256_set1_ps(norm)),\n _MM_FROUND_NO_EXC));\n arr += 8;\n out += 8;\n }\n }\n switch (last - arr) {\n case 7:\n out[6] = float16(arr[6] / norm);\n /* FALLTHRU */\n case 6:\n out[5] = float16(arr[5] / norm);\n /* FALLTHRU */\n case 5:\n out[4] = float16(arr[4] / norm);\n /* FALLTHRU */\n case 4:\n out[3] = float16(arr[3] / norm);\n /* FALLTHRU */\n case 3:\n out[2] = float16(arr[2] / norm);\n /* FALLTHRU */\n case 2:\n out[1] = float16(arr[1] / norm);\n /* FALLTHRU */\n case 1:\n out[0] = float16(arr[0] / norm);\n }\n}\n#endif //__F16C__ && __AVX512F__\n\n#if defined(__F16C__) && defined(__AVX__)\nstatic inline void convert_fp16_to_fp32_avx(const uint16_t *arr, size_t size,\n float *out) {\n const uint16_t *last = arr + size;\n const uint16_t *last_aligned = arr + ((size >> 4) << 4);\n\n if (((uintptr_t)arr & 0xf) == 0 && ((uintptr_t)out & 0x1f) == 0) {\n for (; arr != last_aligned; arr += 16, out += 16) {\n _mm256_store_ps(out + 0,\n _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 0))));\n _mm256_store_ps(out + 8,\n _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 8))));\n }\n\n if (last >= last_aligned + 8) {\n _mm256_store_ps(out + 0,\n _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 0))));\n arr += 8;\n out += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 16, out += 16) {\n _mm256_storeu_ps(out + 0,\n _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 0))));\n _mm256_storeu_ps(out + 8,\n _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 8))));\n }\n\n if (last >= last_aligned + 8) {\n _mm256_storeu_ps(out + 0,\n _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 0))));\n arr += 8;\n out += 8;\n }\n }\n switch (last - arr) {\n case 7:\n out[6] = _cvtsh_ss(arr[6]);\n /* FALLTHRU */\n case 6:\n out[5] = _cvtsh_ss(arr[5]);\n /* FALLTHRU */\n case 5:\n out[4] = _cvtsh_ss(arr[4]);\n /* FALLTHRU */\n case 4:\n out[3] = _cvtsh_ss(arr[3]);\n /* FALLTHRU */\n case 3:\n out[2] = _cvtsh_ss(arr[2]);\n /* FALLTHRU */\n case 2:\n out[1] = _cvtsh_ss(arr[1]);\n /* FALLTHRU */\n case 1:\n out[0] = _cvtsh_ss(arr[0]);\n }\n}\n\nstatic inline void convert_fp16_to_fp32_avx(const uint16_t *arr, size_t size,\n float norm, float *out) {\n const uint16_t *last = arr + size;\n const uint16_t *last_aligned = arr + ((size >> 4) << 4);\n __m256 ymm_norm = _mm256_set1_ps(norm);\n\n if (((uintptr_t)arr & 0xf) == 0 && ((uintptr_t)out & 0x1f) == 0) {\n for (; arr != last_aligned; arr += 16, out += 16) {\n __m256 ymm_0 = _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 0)));\n __m256 ymm_1 = _mm256_cvtph_ps(_mm_load_si128((__m128i *)(arr + 8)));\n ymm_0 = _mm256_div_ps(ymm_0, ymm_norm);\n ymm_1 = _mm256_div_ps(ymm_1, ymm_norm);\n _mm256_store_ps(out + 0, ymm_0);\n _mm256_store_ps(out + 8, ymm_1);\n }\n\n if (last >= last_aligned + 8) {\n _mm256_store_ps(\n out, _mm256_div_ps(_mm256_cvtph_ps(_mm_load_si128((__m128i *)arr)),\n ymm_norm));\n arr += 8;\n out += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 16, out += 16) {\n __m256 ymm_0 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 0)));\n __m256 ymm_1 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i *)(arr + 8)));\n ymm_0 = _mm256_div_ps(ymm_0, ymm_norm);\n ymm_1 = _mm256_div_ps(ymm_1, ymm_norm);\n _mm256_storeu_ps(out + 0, ymm_0);\n _mm256_storeu_ps(out + 8, ymm_1);\n }\n\n if (last >= last_aligned + 8) {\n _mm256_storeu_ps(\n out, _mm256_div_ps(_mm256_cvtph_ps(_mm_loadu_si128((__m128i *)arr)),\n ymm_norm));\n arr += 8;\n out += 8;\n }\n }\n switch (last - arr) {\n case 7:\n out[6] = _cvtsh_ss(arr[6]) / norm;\n /* FALLTHRU */\n case 6:\n out[5] = _cvtsh_ss(arr[5]) / norm;\n /* FALLTHRU */\n case 5:\n out[4] = _cvtsh_ss(arr[4]) / norm;\n /* FALLTHRU */\n case 4:\n out[3] = _cvtsh_ss(arr[3]) / norm;\n /* FALLTHRU */\n case 3:\n out[2] = _cvtsh_ss(arr[2]) / norm;\n /* FALLTHRU */\n case 2:\n out[1] = _cvtsh_ss(arr[1]) / norm;\n /* FALLTHRU */\n case 1:\n out[0] = _cvtsh_ss(arr[0]) / norm;\n }\n}\n\nstatic inline void convert_fp32_to_fp16_avx(const float *arr, size_t size,\n uint16_t *out) {\n const float *last = arr + size;\n const float *last_aligned = arr + ((size >> 4) << 4);\n\n if (((uintptr_t)arr & 0x1f) == 0 && ((uintptr_t)out & 0xf) == 0) {\n for (; arr != last_aligned; arr += 16, out += 16) {\n _mm_store_si128(\n (__m128i *)(out + 0),\n _mm256_cvtps_ph(_mm256_load_ps(arr + 0), _MM_FROUND_NO_EXC));\n _mm_store_si128(\n (__m128i *)(out + 8),\n _mm256_cvtps_ph(_mm256_load_ps(arr + 8), _MM_FROUND_NO_EXC));\n }\n\n if (last >= last_aligned + 8) {\n _mm_store_si128(\n (__m128i *)(out + 0),\n _mm256_cvtps_ph(_mm256_load_ps(arr + 0), _MM_FROUND_NO_EXC));\n arr += 8;\n out += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 16, out += 16) {\n _mm_storeu_si128(\n (__m128i *)(out + 0),\n _mm256_cvtps_ph(_mm256_loadu_ps(arr + 0), _MM_FROUND_NO_EXC));\n _mm_storeu_si128(\n (__m128i *)(out + 8),\n _mm256_cvtps_ph(_mm256_loadu_ps(arr + 8), _MM_FROUND_NO_EXC));\n }\n\n if (last >= last_aligned + 8) {\n _mm_storeu_si128(\n (__m128i *)(out + 0),\n _mm256_cvtps_ph(_mm256_loadu_ps(arr + 0), _MM_FROUND_NO_EXC));\n arr += 8;\n out += 8;\n }\n }\n switch (last - arr) {\n case 7:\n out[6] = _cvtss_sh(arr[6], _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 6:\n out[5] = _cvtss_sh(arr[5], _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 5:\n out[4] = _cvtss_sh(arr[4], _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 4:\n out[3] = _cvtss_sh(arr[3], _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 3:\n out[2] = _cvtss_sh(arr[2], _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 2:\n out[1] = _cvtss_sh(arr[1], _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 1:\n out[0] = _cvtss_sh(arr[0], _MM_FROUND_NO_EXC);\n }\n}\n\nstatic inline void convert_fp32_to_fp16_avx(const float *arr, size_t size,\n float norm, uint16_t *out) {\n const float *last = arr + size;\n const float *last_aligned = arr + ((size >> 4) << 4);\n __m256 ymm_norm = _mm256_set1_ps(norm);\n\n if (((uintptr_t)arr & 0x1f) == 0 && ((uintptr_t)out & 0xf) == 0) {\n for (; arr != last_aligned; arr += 16, out += 16) {\n __m256 ymm_0 = _mm256_load_ps(arr + 0);\n __m256 ymm_1 = _mm256_load_ps(arr + 8);\n ymm_0 = _mm256_div_ps(ymm_0, ymm_norm);\n ymm_1 = _mm256_div_ps(ymm_1, ymm_norm);\n _mm_store_si128((__m128i *)(out + 0),\n _mm256_cvtps_ph(ymm_0, _MM_FROUND_NO_EXC));\n _mm_store_si128((__m128i *)(out + 8),\n _mm256_cvtps_ph(ymm_1, _MM_FROUND_NO_EXC));\n }\n\n if (last >= last_aligned + 8) {\n _mm_store_si128(\n (__m128i *)out,\n _mm256_cvtps_ph(_mm256_div_ps(_mm256_load_ps(arr), ymm_norm),\n _MM_FROUND_NO_EXC));\n arr += 8;\n out += 8;\n }\n } else {\n for (; arr != last_aligned; arr += 16, out += 16) {\n __m256 ymm_0 = _mm256_loadu_ps(arr + 0);\n __m256 ymm_1 = _mm256_loadu_ps(arr + 8);\n ymm_0 = _mm256_div_ps(ymm_0, ymm_norm);\n ymm_1 = _mm256_div_ps(ymm_1, ymm_norm);\n _mm_storeu_si128((__m128i *)(out + 0),\n _mm256_cvtps_ph(ymm_0, _MM_FROUND_NO_EXC));\n _mm_storeu_si128((__m128i *)(out + 8),\n _mm256_cvtps_ph(ymm_1, _MM_FROUND_NO_EXC));\n }\n\n if (last >= last_aligned + 8) {\n _mm_storeu_si128(\n (__m128i *)out,\n _mm256_cvtps_ph(_mm256_div_ps(_mm256_loadu_ps(arr), ymm_norm),\n _MM_FROUND_NO_EXC));\n arr += 8;\n out += 8;\n }\n }\n switch (last - arr) {\n case 7:\n out[6] = _cvtss_sh(arr[6] / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 6:\n out[5] = _cvtss_sh(arr[5] / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 5:\n out[4] = _cvtss_sh(arr[4] / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 4:\n out[3] = _cvtss_sh(arr[3] / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 3:\n out[2] = _cvtss_sh(arr[2] / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 2:\n out[1] = _cvtss_sh(arr[1] / norm, _MM_FROUND_NO_EXC);\n /* FALLTHRU */\n case 1:\n out[0] = _cvtss_sh(arr[0] / norm, _MM_FROUND_NO_EXC);\n }\n}\n#endif // __F16C__ && __AVX__\n\nstatic inline void convert_fp16_to_fp32_fallback(const uint16_t *arr,\n size_t size, float *out) {\n for (size_t i = 0; i != size; ++i) {\n out[i] = float32(arr[i]);\n }\n}\n\nstatic inline void convert_fp16_to_fp32_fallback(const uint16_t *arr,\n size_t size, float norm,\n float *out) {\n for (size_t i = 0; i != size; ++i) {\n out[i] = float32(arr[i]) / norm;\n }\n}\n\nstatic inline void convert_fp32_to_fp16_fallback(const float *arr, size_t size,\n uint16_t *out) {\n for (size_t i = 0; i != size; ++i) {\n out[i] = float16(arr[i]);\n }\n}\n\nstatic inline void convert_fp32_to_fp16_fallback(const float *arr, size_t size,\n float norm, uint16_t *out) {\n for (size_t i = 0; i != size; ++i) {\n out[i] = float16(arr[i] / norm);\n }\n}\n\nstatic inline void convert_fp16_to_fp32(const uint16_t *arr, size_t size,\n float *out) {\n#if defined(__F16C__) && defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C &&\n zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n return convert_fp16_to_fp32_avx512f(arr, size, out);\n }\n#endif\n\n#if defined(__F16C__) && defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C &&\n zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n return convert_fp16_to_fp32_avx(arr, size, out);\n }\n#endif\n\n return convert_fp16_to_fp32_fallback(arr, size, out);\n}\n\nstatic inline void convert_fp16_to_fp32(const uint16_t *arr, size_t size,\n float norm, float *out) {\n#if defined(__F16C__) && defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C &&\n zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n return convert_fp16_to_fp32_avx512f(arr, size, norm, out);\n }\n#endif\n\n#if defined(__F16C__) && defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C &&\n zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n return convert_fp16_to_fp32_avx(arr, size, norm, out);\n }\n#endif\n\n return convert_fp16_to_fp32_fallback(arr, size, norm, out);\n}\n\nstatic inline void convert_fp32_to_fp16(const float *arr, size_t size,\n uint16_t *out) {\n#if defined(__F16C__) && defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C &&\n zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n return convert_fp32_to_fp16_avx512f(arr, size, out);\n }\n#endif\n\n#if defined(__F16C__) && defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C &&\n zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n return convert_fp32_to_fp16_avx(arr, size, out);\n }\n#endif\n\n return convert_fp32_to_fp16_fallback(arr, size, out);\n}\n\nstatic inline void convert_fp32_to_fp16(const float *arr, size_t size,\n float norm, uint16_t *out) {\n#if defined(__F16C__) && defined(__AVX512F__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C &&\n zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {\n return convert_fp32_to_fp16_avx512f(arr, size, norm, out);\n }\n#endif\n\n#if defined(__F16C__) && defined(__AVX__)\n if (zvec::ailego::internal::CpuFeatures::static_flags_.F16C &&\n zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {\n return convert_fp32_to_fp16_avx(arr, size, norm, out);\n }\n#endif\n\n return convert_fp32_to_fp16_fallback(arr, size, norm, out);\n}\n\n#endif //\n\nnamespace zvec {\nnamespace ailego {\n\nfloat FloatHelper::ToFP32(uint16_t val) {\n return float32(val);\n}\n\nvoid FloatHelper::ToFP32(const uint16_t *arr, size_t size, float *out) {\n return convert_fp16_to_fp32(arr, size, out);\n}\n\nvoid FloatHelper::ToFP32(const uint16_t *arr, size_t size, float norm,\n float *out) {\n return convert_fp16_to_fp32(arr, size, norm, out);\n}\n\nuint16_t FloatHelper::ToFP16(float val) {\n return float16(val);\n}\n\nvoid FloatHelper::ToFP16(const float *arr, size_t size, uint16_t *out) {\n return convert_fp32_to_fp16(arr, size, out);\n}\n\nvoid FloatHelper::ToFP16(const float *arr, size_t size, float norm,\n uint16_t *out) {\n return convert_fp32_to_fp16(arr, size, norm, out);\n}\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/utility/math_helper.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \nnamespace zvec {\nnamespace ailego {\n\n/*! Math Helper\n */\nstruct MathHelper {\n //! Calculate the absolute value\n template \n static inline auto Absolute(const T &x) ->\n typename std::enable_if::value, R>::type {\n return static_cast(std::abs(x));\n }\n\n //! Calculate the absolute value\n template \n static inline R Absolute(const Float16 &x) {\n return static_cast(Float16::Absolute(x));\n }\n\n //! Calculate the absolute difference\n template \n static inline auto AbsoluteDifference(const T &x, const T &y) ->\n typename std::enable_if::value, R>::type {\n auto m = ((x ^ y) & -(x < y));\n auto d =\n static_cast::type>((x ^ m) - (y ^ m));\n return static_cast(d);\n }\n\n //! Calculate the absolute difference\n template \n static inline auto AbsoluteDifference(const T &x, const T &y) ->\n typename std::enable_if::value, R>::type {\n return static_cast(std::abs(x - y));\n }\n\n //! Calculate the absolute difference\n template \n static inline R AbsoluteDifference(const Float16 &x, const Float16 &y) {\n return static_cast(std::abs(x - y));\n }\n\n //! Calculate the squared difference\n template \n static inline auto SquaredDifference(const T &x, const T &y) ->\n typename std::enable_if::value, R>::type {\n auto m = ((x ^ y) & -(x < y));\n auto d =\n static_cast::type>((x ^ m) - (y ^ m));\n return static_cast(d * d);\n }\n\n //! Calculate the squared difference\n template \n static inline auto SquaredDifference(const T &x, const T &y) ->\n typename std::enable_if::value, R>::type {\n auto d = x - y;\n return static_cast(d * d);\n }\n\n //! Calculate the squared difference\n template \n static inline R SquaredDifference(const Float16 &x, const Float16 &y) {\n auto d = x - y;\n return static_cast(d * d);\n }\n\n //! Test whether two integral numbers are equal\n template \n static inline auto IsAlmostEqual(const T &x, const T &y, int) ->\n typename std::enable_if::value, bool>::type {\n return (x == y);\n }\n\n //! Test whether two floating point numbers are equal\n template \n static inline auto IsAlmostEqual(const T &x, const T &y, int ulp) ->\n typename std::enable_if::value, bool>::type {\n // the machine epsilon has to be scaled to the magnitude of the values used\n // and multiplied by the desired precision in ULPs (units in the last place)\n return ((std::fabs(x - y) <=\n std::numeric_limits::epsilon() * std::fabs(x + y) * ulp) ||\n (std::fabs(x - y) < std::numeric_limits::min()));\n }\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/utility/matrix_helper.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace ailego {\n\nstruct MatrixHelper {\n //! Transpose a matrix\n template \n static inline void Transpose(const void *src, size_t N, void *dst) {\n for (size_t i = 0; i < M; ++i) {\n for (size_t j = 0; j < N; ++j) {\n *(reinterpret_cast(dst) + (j * M + i)) =\n *(reinterpret_cast(src) + (i * N + j));\n }\n }\n }\n\n //! Reverse transpose a matrix\n template \n static inline void ReverseTranspose(const void *src, size_t N, void *dst) {\n for (size_t i = 0; i < N; ++i) {\n for (size_t j = 0; j < M; ++j) {\n *(reinterpret_cast(dst) + (j * N + i)) =\n *(reinterpret_cast(src) + (i * M + j));\n }\n }\n }\n\n //! Transpose a matrix\n template \n static inline void Transpose(const void *src, size_t M, size_t N, void *dst) {\n for (size_t i = 0; i < M; ++i) {\n for (size_t j = 0; j < N; ++j) {\n *(reinterpret_cast(dst) + (j * M + i)) =\n *(reinterpret_cast(src) + (i * N + j));\n }\n }\n }\n\n //! Reverse transpose a matrix\n template \n static inline void ReverseTranspose(const void *src, size_t M, size_t N,\n void *dst) {\n for (size_t i = 0; i < N; ++i) {\n for (size_t j = 0; j < M; ++j) {\n *(reinterpret_cast(dst) + (j * N + i)) =\n *(reinterpret_cast(src) + (i * M + j));\n }\n }\n }\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/utility/memory_helper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"memory_helper.h\"\n#include \n#include \n#include \n#include \n#include \n\n#if defined(_WIN64) || defined(_WIN32)\n#include \n#include \n#else\n#if defined(__linux__) || defined(__linux)\n#include \n#elif defined(__APPLE__) && defined(__MACH__)\n#include \n#include \n#endif\n#include \n#endif\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(__linux__) || defined(__linux)\nbool MemoryHelper::SelfUsage(size_t *vsz, size_t *rss) {\n FILE *fp = fopen(\"/proc/self/statm\", \"r\");\n if (!fp) {\n return false;\n }\n\n if (fscanf(fp, \"%zd %zd\", vsz, rss) == EOF) {\n fclose(fp);\n return false;\n }\n fclose(fp);\n\n long pagesz = sysconf(_SC_PAGESIZE);\n *vsz *= (size_t)pagesz;\n *rss *= (size_t)pagesz;\n return true;\n}\n\nsize_t MemoryHelper::SelfRSS(void) {\n FILE *fp = fopen(\"/proc/self/statm\", \"r\");\n if (!fp) {\n return 0;\n }\n\n size_t rss = 0;\n if (fscanf(fp, \"%*d %zd %*d\", &rss) == EOF) {\n fclose(fp);\n return 0;\n }\n fclose(fp);\n return (rss * sysconf(_SC_PAGESIZE));\n}\n\nsize_t MemoryHelper::SelfPeakRSS(void) {\n struct rusage rusage;\n getrusage(RUSAGE_SELF, &rusage);\n return (size_t)(rusage.ru_maxrss * 1024);\n}\n\nsize_t MemoryHelper::TotalRamSize(void) {\n return (sysconf(_SC_PHYS_PAGES) * sysconf(_SC_PAGESIZE));\n}\n\nsize_t MemoryHelper::AvailableRamSize(void) {\n FILE *fp = fopen(\"/proc/meminfo\", \"r\");\n if (!fp) {\n return 0;\n }\n\n size_t avail = 0;\n char buf[128];\n while (fgets(buf, sizeof(buf), fp)) {\n if (strncmp(buf, \"MemAvailable:\", 13) == 0) {\n avail = (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10);\n break;\n }\n }\n\n // No found 'MemAvailable'\n if (avail == 0) {\n fseek(fp, 0L, SEEK_SET);\n\n size_t count = 0;\n while (fgets(buf, sizeof(buf), fp)) {\n switch (buf[0]) {\n case 'M':\n if (strncmp(buf, \"MemFree:\", 8) == 0) {\n avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10);\n ++count;\n }\n break;\n\n case 'B':\n if (strncmp(buf, \"Buffers:\", 8) == 0) {\n avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10);\n ++count;\n }\n break;\n\n case 'C':\n if (strncmp(buf, \"Cached:\", 7) == 0) {\n avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10);\n ++count;\n }\n break;\n }\n // All read\n if (count == 3) {\n break;\n }\n }\n }\n fclose(fp);\n return (avail * 1024);\n}\n\nsize_t MemoryHelper::UsedRamSize(void) {\n FILE *fp = fopen(\"/proc/meminfo\", \"r\");\n if (!fp) {\n return 0;\n }\n\n size_t total = 0, avail = 0, count = 0;\n char buf[128];\n\n while (fgets(buf, sizeof(buf), fp)) {\n switch (buf[0]) {\n case 'M':\n if (strncmp(buf, \"MemTotal:\", 9) == 0) {\n total = (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10);\n ++count;\n } else if (strncmp(buf, \"MemFree:\", 8) == 0) {\n avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10);\n ++count;\n }\n break;\n\n case 'B':\n if (strncmp(buf, \"Buffers:\", 8) == 0) {\n avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10);\n ++count;\n }\n break;\n\n case 'C':\n if (strncmp(buf, \"Cached:\", 7) == 0) {\n avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10);\n ++count;\n }\n break;\n\n case 'S':\n if (strncmp(buf, \"Slab:\", 5) == 0) {\n avail += (size_t)strtoull(strchr(buf, ':') + 1, NULL, 10);\n ++count;\n }\n break;\n }\n // All read\n if (count == 5) {\n break;\n }\n }\n fclose(fp);\n\n if (total == 0) {\n total = (sysconf(_SC_PHYS_PAGES) * sysconf(_SC_PAGESIZE)) / 1024;\n }\n return ((total - avail) * 1024);\n}\n\nsize_t MemoryHelper::ContainerAwareTotalRamSize(void) {\n size_t total_ram_size = TotalRamSize();\n std::string limit_in_bytes = \"/sys/fs/cgroup/memory/memory.limit_in_bytes\";\n if (FileHelper::IsExist(limit_in_bytes.c_str())) {\n std::ifstream memory_limit_ifs;\n std::string memory_limit_str{\"\"};\n memory_limit_ifs.open(limit_in_bytes, std::ios::in);\n if (memory_limit_ifs.is_open()) {\n uint64_t limit = 0;\n memory_limit_ifs >> memory_limit_str;\n if (memory_limit_str != \"-1\") {\n // Refer to:\n // https://access.redhat.com/documentation/zh-cn/red_hat_enterprise_linux/7/html/resource_management_guide/sec-memory\n StringHelper::ToUint64(memory_limit_str, &limit);\n if (limit != 0x7FFFFFFFFFFFF000) {\n // Refer to:\n // https://stackoverflow.com/questions/70332396/why-cgroups-file-memory-limit-in-bytes-use-9223372036854771712-as-a-default-valu\n total_ram_size = static_cast(limit);\n }\n }\n memory_limit_ifs.close();\n }\n }\n return total_ram_size;\n}\n\n#elif defined(__APPLE__) && defined(__MACH__)\nbool MemoryHelper::SelfUsage(size_t *vsz, size_t *rss) {\n struct mach_task_basic_info info;\n mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT;\n\n if (task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&info,\n &count) != KERN_SUCCESS) {\n return false;\n }\n *vsz = info.virtual_size;\n *rss = info.resident_size;\n return true;\n}\n\nsize_t MemoryHelper::SelfRSS(void) {\n struct mach_task_basic_info info;\n mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT;\n\n if (task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&info,\n &count) != KERN_SUCCESS) {\n return 0;\n }\n return info.resident_size;\n}\n\nsize_t MemoryHelper::SelfPeakRSS(void) {\n struct mach_task_basic_info info;\n mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT;\n\n if (task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&info,\n &count) != KERN_SUCCESS) {\n return 0;\n }\n return info.resident_size_max;\n}\n\nsize_t MemoryHelper::TotalRamSize(void) {\n int mib[2] = {CTL_HW, HW_MEMSIZE};\n uint64_t size = 0;\n size_t len = sizeof(size);\n if (sysctl(mib, 2, &size, &len, 0, 0) != 0) {\n return 0;\n }\n return (size_t)size;\n}\n\nsize_t MemoryHelper::AvailableRamSize(void) {\n struct vm_statistics stat;\n mach_msg_type_number_t count = HOST_VM_INFO_COUNT;\n vm_size_t pagesize = 0;\n\n if (host_page_size(mach_host_self(), &pagesize) != KERN_SUCCESS) {\n return 0;\n }\n if (host_statistics(mach_host_self(), HOST_VM_INFO, (host_info_t)&stat,\n &count) != KERN_SUCCESS) {\n return 0;\n }\n return ((stat.free_count + stat.inactive_count) * pagesize);\n}\n\nsize_t MemoryHelper::UsedRamSize(void) {\n struct vm_statistics stat;\n mach_msg_type_number_t count = HOST_VM_INFO_COUNT;\n vm_size_t pagesize = 0;\n\n if (host_page_size(mach_host_self(), &pagesize) != KERN_SUCCESS) {\n return 0;\n }\n if (host_statistics(mach_host_self(), HOST_VM_INFO, (host_info_t)&stat,\n &count) != KERN_SUCCESS) {\n return 0;\n }\n return ((stat.active_count + stat.wire_count) * pagesize);\n}\n\nsize_t MemoryHelper::ContainerAwareTotalRamSize(void) {\n return 0u;\n}\n\n#elif defined(_WIN64) || defined(_WIN32)\nstatic inline int getpagesize(void) {\n SYSTEM_INFO info;\n GetSystemInfo(&info);\n return info.dwPageSize;\n}\n\nbool MemoryHelper::SelfUsage(size_t *vsz, size_t *rss) {\n PROCESS_MEMORY_COUNTERS info;\n if (!GetProcessMemoryInfo(GetCurrentProcess(), &info, sizeof(info))) {\n return false;\n }\n *vsz = (size_t)info.PagefileUsage;\n *rss = (size_t)info.WorkingSetSize;\n return true;\n}\n\nsize_t MemoryHelper::SelfRSS(void) {\n PROCESS_MEMORY_COUNTERS info;\n if (!GetProcessMemoryInfo(GetCurrentProcess(), &info, sizeof(info))) {\n return 0u;\n }\n return (size_t)info.WorkingSetSize;\n}\n\nsize_t MemoryHelper::SelfPeakRSS(void) {\n PROCESS_MEMORY_COUNTERS info;\n GetProcessMemoryInfo(GetCurrentProcess(), &info, sizeof(info));\n return (size_t)info.PeakWorkingSetSize;\n}\n\nsize_t MemoryHelper::TotalRamSize(void) {\n MEMORYSTATUSEX status;\n status.dwLength = sizeof(status);\n GlobalMemoryStatusEx(&status);\n return (size_t)status.ullTotalPhys;\n}\n\nsize_t MemoryHelper::AvailableRamSize(void) {\n MEMORYSTATUSEX status;\n status.dwLength = sizeof(status);\n GlobalMemoryStatusEx(&status);\n return (size_t)status.ullAvailPhys;\n}\n\nsize_t MemoryHelper::UsedRamSize(void) {\n MEMORYSTATUSEX status;\n status.dwLength = sizeof(status);\n GlobalMemoryStatusEx(&status);\n return (size_t)(status.ullTotalPhys - status.ullAvailPhys);\n}\n\nsize_t MemoryHelper::ContainerAwareTotalRamSize(void) {\n return 0u;\n}\n\n#else\nbool MemoryHelper::SelfUsage(size_t *vsz, size_t *rss) {\n *vsz = 0u;\n *rss = 0u;\n return false;\n}\n\nsize_t MemoryHelper::SelfRSS(void) {\n return 0u;\n}\n\nsize_t MemoryHelper::SelfPeakRSS(void) {\n return 0u;\n}\n\nsize_t MemoryHelper::TotalRamSize(void) {\n return 0u;\n}\n\nsize_t MemoryHelper::AvailableRamSize(void) {\n return 0u;\n}\n\nsize_t MemoryHelper::UsedRamSize(void) {\n return 0u;\n}\n\nsize_t MemoryHelper::ContainerAwareTotalRamSize(void) {\n return 0u;\n}\n#endif\n\nsize_t MemoryHelper::PageSize(void) {\n static size_t page_size = static_cast(getpagesize());\n return page_size;\n}\n\nsize_t MemoryHelper::HugePageSize(void) {\n static size_t page_size = static_cast(2 * 1024 * 1024);\n return page_size;\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/utility/memory_helper.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace ailego {\n\n/*! Memory Helper\n */\nstruct MemoryHelper {\n //! Retrieve the page size of memory\n static size_t PageSize(void);\n\n //! Retrieve the huge page size of memory\n static size_t HugePageSize(void);\n\n //! Retrieve the VSZ and RSS of self process in bytes\n static bool SelfUsage(size_t *vsz, size_t *rss);\n\n //! Retrieve the RSS of self process in bytes\n static size_t SelfRSS(void);\n\n //! Retrieve the peak RSS of self process in bytes\n static size_t SelfPeakRSS(void);\n\n //! Retrieve the total size of physical memory (RAM) in bytes\n static size_t TotalRamSize(void);\n\n //! Retrieve the available size of physical memory (RAM) in bytes\n static size_t AvailableRamSize(void);\n\n //! Retrieve the used size of physical memory (RAM) in bytes\n static size_t UsedRamSize(void);\n\n //! Retrieve the total size of physical memory (RAM) in bytes in container\n static size_t ContainerAwareTotalRamSize(void);\n};\n\n} // namespace ailego\n} // namespace zvec\n" }, { "path": "src/ailego/utility/string_helper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace ailego {\n\nbool StringHelper::StartsWith(const std::string &ref,\n const std::string &prefix) {\n return (ref.size() >= prefix.size()) &&\n (ref.compare(0, prefix.size(), prefix) == 0);\n}\n\nbool StringHelper::EndsWith(const std::string &ref, const std::string &suffix) {\n size_t s1 = ref.size();\n size_t s2 = suffix.size();\n return (s1 >= s2) && (ref.compare(s1 - s2, s2, suffix) == 0);\n}\n\nvoid StringHelper::LeftTrim(std::string &str) {\n str.erase(str.begin(), std::find_if(str.begin(), str.end(), [](int ch) {\n return !std::isspace(ch);\n }));\n}\n\nvoid StringHelper::RightTrim(std::string &str) {\n str.erase(std::find_if(str.rbegin(), str.rend(),\n [](int ch) { return !std::isspace(ch); })\n .base(),\n str.end());\n}\n\nvoid StringHelper::Trim(std::string &str) {\n StringHelper::RightTrim(str);\n StringHelper::LeftTrim(str);\n}\n\nstd::string StringHelper::CopyLeftTrim(std::string str) {\n StringHelper::LeftTrim(str);\n return str;\n}\n\nstd::string StringHelper::CopyRightTrim(std::string str) {\n StringHelper::RightTrim(str);\n return str;\n}\n\nstd::string StringHelper::CopyTrim(std::string str) {\n StringHelper::Trim(str);\n return str;\n}\n\n#if defined(_MSC_VER)\n#define strncasecmp _strnicmp\n#endif\n\nbool StringHelper::CompareIgnoreCase(const std::string &a,\n const std::string &b) {\n if (a.size() != b.size()) {\n return false;\n }\n return (strncasecmp(a.data(), b.data(), a.size()) == 0);\n}\n\nvoid StringHelper::Append(std::string *str, const internal::Alphameric &a) {\n str->reserve(str->size() + a.size());\n str->append(a.data(), a.size());\n}\n\nvoid StringHelper::Append(std::string *str, const internal::Alphameric &a,\n const internal::Alphameric &b) {\n str->reserve(str->size() + a.size() + b.size());\n str->append(a.data(), a.size());\n str->append(b.data(), b.size());\n}\n\nvoid StringHelper::Append(std::string *str, const internal::Alphameric &a,\n const internal::Alphameric &b,\n const internal::Alphameric &c) {\n str->reserve(str->size() + a.size() + b.size() + c.size());\n str->append(a.data(), a.size());\n str->append(b.data(), b.size());\n str->append(c.data(), c.size());\n}\n\nvoid StringHelper::Append(std::string *str, const internal::Alphameric &a,\n const internal::Alphameric &b,\n const internal::Alphameric &c,\n const internal::Alphameric &d) {\n str->reserve(str->size() + a.size() + b.size() + c.size() + d.size());\n str->append(a.data(), a.size());\n str->append(b.data(), b.size());\n str->append(c.data(), c.size());\n str->append(d.data(), d.size());\n}\n\nvoid StringHelper::AppendViews(std::string *str,\n std::initializer_list views) {\n size_t new_size = str->size();\n for (auto &v : views) {\n new_size += v.size();\n }\n str->reserve(new_size);\n for (auto &v : views) {\n str->append(v.data(), v.size());\n }\n}\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/utility/time_helper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n\n#if defined(_WIN64) || defined(_WIN32)\n#include \n#endif\n\nnamespace zvec {\nnamespace ailego {\n\n#if defined(_WIN64) || defined(_WIN32)\nuint64_t Monotime::NanoSeconds(void) {\n LARGE_INTEGER stamp, freq;\n QueryPerformanceFrequency(&freq);\n QueryPerformanceCounter(&stamp);\n return (uint64_t)((double)stamp.QuadPart *\n (1000000000.0 / (double)freq.QuadPart));\n}\n\nuint64_t Monotime::MicroSeconds(void) {\n LARGE_INTEGER stamp, freq;\n QueryPerformanceFrequency(&freq);\n QueryPerformanceCounter(&stamp);\n return (stamp.QuadPart * 1000000u / freq.QuadPart);\n}\n\nuint64_t Monotime::MilliSeconds(void) {\n LARGE_INTEGER stamp, freq;\n QueryPerformanceFrequency(&freq);\n QueryPerformanceCounter(&stamp);\n return (stamp.QuadPart * 1000u / freq.QuadPart);\n}\n\nuint64_t Monotime::Seconds(void) {\n LARGE_INTEGER stamp, freq;\n QueryPerformanceFrequency(&freq);\n QueryPerformanceCounter(&stamp);\n return (stamp.QuadPart / freq.QuadPart);\n}\n\n// January 1, 1970 (start of Unix epoch) in \"ticks\"\n#define UNIX_TIME_START 0x019DB1DED53E8000ull\n\nuint64_t Realtime::NanoSeconds(void) {\n LARGE_INTEGER stamp;\n FILETIME file;\n GetSystemTimeAsFileTime(&file);\n stamp.HighPart = file.dwHighDateTime;\n stamp.LowPart = file.dwLowDateTime;\n return (stamp.QuadPart - UNIX_TIME_START) * 100u;\n}\n\nuint64_t Realtime::MicroSeconds(void) {\n LARGE_INTEGER stamp;\n FILETIME file;\n GetSystemTimeAsFileTime(&file);\n stamp.HighPart = file.dwHighDateTime;\n stamp.LowPart = file.dwLowDateTime;\n return (stamp.QuadPart - UNIX_TIME_START) / 10u;\n}\n\nuint64_t Realtime::MilliSeconds(void) {\n LARGE_INTEGER stamp;\n FILETIME file;\n GetSystemTimeAsFileTime(&file);\n stamp.HighPart = file.dwHighDateTime;\n stamp.LowPart = file.dwLowDateTime;\n return (stamp.QuadPart - UNIX_TIME_START) / 10000u;\n}\n\nuint64_t Realtime::Seconds(void) {\n LARGE_INTEGER stamp;\n FILETIME file;\n GetSystemTimeAsFileTime(&file);\n stamp.HighPart = file.dwHighDateTime;\n stamp.LowPart = file.dwLowDateTime;\n return (stamp.QuadPart - UNIX_TIME_START) / 10000000u;\n}\n\nsize_t Realtime::Localtime(uint64_t stamp, const char *format, char *buf,\n size_t len) {\n time_t val = static_cast(stamp);\n return strftime(buf, len, format, localtime(&val));\n}\n\nsize_t Realtime::Gmtime(uint64_t stamp, const char *format, char *buf,\n size_t len) {\n time_t val = static_cast(stamp);\n return strftime(buf, len, format, gmtime(&val));\n}\n\nsize_t Realtime::Localtime(const char *format, char *buf, size_t len) {\n time_t now = time(0);\n return strftime(buf, len, format, localtime(&now));\n}\n\nsize_t Realtime::Gmtime(const char *format, char *buf, size_t len) {\n time_t now = time(0);\n return strftime(buf, len, format, gmtime(&now));\n}\n#else\nuint64_t Monotime::NanoSeconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_MONOTONIC, &tspec);\n return (tspec.tv_sec * 1000000000u + tspec.tv_nsec);\n}\n\nuint64_t Monotime::MicroSeconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_MONOTONIC, &tspec);\n return (tspec.tv_sec * 1000000u + tspec.tv_nsec / 1000u);\n}\n\nuint64_t Monotime::MilliSeconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_MONOTONIC, &tspec);\n return (tspec.tv_sec * 1000u + tspec.tv_nsec / 1000000u);\n}\n\nuint64_t Monotime::Seconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_MONOTONIC, &tspec);\n return (tspec.tv_sec);\n}\n\nuint64_t Realtime::NanoSeconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_REALTIME, &tspec);\n return (tspec.tv_sec * 1000000000u + tspec.tv_nsec);\n}\n\nuint64_t Realtime::MicroSeconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_REALTIME, &tspec);\n return (tspec.tv_sec * 1000000u + tspec.tv_nsec / 1000u);\n}\n\nuint64_t Realtime::MilliSeconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_REALTIME, &tspec);\n return (tspec.tv_sec * 1000u + tspec.tv_nsec / 1000000u);\n}\n\nuint64_t Realtime::Seconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_REALTIME, &tspec);\n return (tspec.tv_sec);\n}\n\nsize_t Realtime::Localtime(uint64_t stamp, const char *format, char *buf,\n size_t len) {\n struct tm tmbuf;\n time_t val = static_cast(stamp);\n return strftime(buf, len, format, localtime_r(&val, &tmbuf));\n}\n\nsize_t Realtime::Gmtime(uint64_t stamp, const char *format, char *buf,\n size_t len) {\n struct tm tmbuf;\n time_t val = static_cast(stamp);\n return strftime(buf, len, format, gmtime_r(&val, &tmbuf));\n}\n\nsize_t Realtime::Localtime(const char *format, char *buf, size_t len) {\n struct tm tmbuf;\n time_t now = time(0);\n return strftime(buf, len, format, localtime_r(&now, &tmbuf));\n}\n\nsize_t Realtime::Gmtime(const char *format, char *buf, size_t len) {\n struct tm tmbuf;\n time_t now = time(0);\n return strftime(buf, len, format, gmtime_r(&now, &tmbuf));\n}\n\nuint64_t CPUtime::NanoSeconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_THREAD_CPUTIME_ID, &tspec);\n return (tspec.tv_sec * 1000000000u + tspec.tv_nsec);\n}\n\nuint64_t CPUtime::MicroSeconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_THREAD_CPUTIME_ID, &tspec);\n return (tspec.tv_sec * 1000000u + tspec.tv_nsec / 1000u);\n}\n\nuint64_t CPUtime::MilliSeconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_THREAD_CPUTIME_ID, &tspec);\n return (tspec.tv_sec * 1000u + tspec.tv_nsec / 1000000u);\n}\n\nuint64_t CPUtime::Seconds(void) {\n struct timespec tspec;\n clock_gettime(CLOCK_THREAD_CPUTIME_ID, &tspec);\n return (tspec.tv_sec);\n}\n#endif // _WIN64 || _WIN32\n\n} // namespace ailego\n} // namespace zvec" }, { "path": "src/ailego/version.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"version.h\"\n#include \"version.i\"\n\n#ifdef ailego_VERSION\n#define AILEGO_VERSION_STRING ailego_VERSION\n#else\n#define AILEGO_VERSION_STRING \"unknown\"\n#endif\n\nnamespace zvec {\n\nnamespace ailego {\n\nstatic const char AILEGO_VERSION_DETAILS[] = AILEGO_VERSION_COMPILE_DETAILS(\n \"AiLego Library Version \" AILEGO_VERSION_STRING\n \".\\nCopyright (C) The Software Authors. All rights reserved.\\n\");\n\nconst char *Version::String(void) {\n return AILEGO_VERSION_STRING;\n}\n\nconst char *Version::Details(void) {\n return AILEGO_VERSION_DETAILS;\n}\n\n} // namespace ailego\n} // namespace zvec\n\n// extern \"C\" int __wrap_main(int, char *[]) {\n// fwrite(ailego::AILEGO_VERSION_DETAILS, 1,\n// strlen(ailego::AILEGO_VERSION_DETAILS), stdout);\n// fflush(stdout);\n// _Exit(0);\n// }\n" }, { "path": "src/ailego/version.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\nnamespace zvec {\n\nnamespace ailego {\n\n/*! AiLego Version\n */\nstruct Version {\n //! Retrieve the version number in string\n static const char *String(void);\n\n //! Retrieve the detailed version information\n static const char *Details(void);\n};\n\n} // namespace ailego\n\n} // namespace zvec" }, { "path": "src/ailego/version.i", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n\n#ifndef AILEGO_VERSION_TO_STRING_\n#define AILEGO_VERSION_TO_STRING_(x) #x\n#endif\n\n#ifndef AILEGO_VERSION_TO_STRING\n#define AILEGO_VERSION_TO_STRING(x) AILEGO_VERSION_TO_STRING_(x)\n#endif\n\n/*! http://nadeausoftware.com/articles/2012/01/\n * c_c_tip_how_use_compiler_predefined_macros_detect_operating_system\n */\n#if defined(__linux) || defined(__linux__)\n#define AILEGO_VERSION_PLATFORM \"Linux\"\n#elif defined(__FreeBSD__)\n#define AILEGO_VERSION_PLATFORM \"FreeBSD\"\n#elif defined(__NetBSD__)\n#define AILEGO_VERSION_PLATFORM \"NetBSD\"\n#elif defined(__OpenBSD__)\n#define AILEGO_VERSION_PLATFORM \"OpenBSD\"\n#elif defined(__APPLE__) || defined(__MACH__)\n#define AILEGO_VERSION_PLATFORM \"Darwin\"\n#elif defined(__CYGWIN__) && !defined(_WIN32)\n#define AILEGO_VERSION_PLATFORM \"Cygwin\"\n#elif defined(_WIN64)\n#define AILEGO_VERSION_PLATFORM \"Microsoft Windows (64-bit)\"\n#elif defined(_WIN32)\n#define AILEGO_VERSION_PLATFORM \"Microsoft Windows (32-bit)\"\n#elif defined(__sun) && defined(__SVR4)\n#define AILEGO_VERSION_PLATFORM \"Solaris\"\n#elif defined(_AIX)\n#define AILEGO_VERSION_PLATFORM \"AIX\"\n#elif defined(__hpux)\n#define AILEGO_VERSION_PLATFORM \"HP-UX\"\n#elif defined(__unix) || defined(__unix__)\n#define AILEGO_VERSION_PLATFORM \"Unix\"\n#else\n#define AILEGO_VERSION_PLATFORM \"Unknown Platform\"\n#endif\n\n/*! http://nadeausoftware.com/articles/2012/10/\n * c_c_tip_how_detect_compiler_name_and_version_using_compiler_predefined_macros\n */\n#if defined(__NVCC__)\n#define AILEGO_VERSION_COMPILER_NAME \"Nvidia CUDA Compiler\"\n#elif defined(__clang__)\n#define AILEGO_VERSION_COMPILER_NAME \"Clang/LLVM\"\n#elif defined(__ICC) || defined(__INTEL_COMPILER)\n#define AILEGO_VERSION_COMPILER_NAME \"Intel ICC/ICPC\"\n#elif defined(__GNUC__) || defined(__GNUG__)\n#define AILEGO_VERSION_COMPILER_NAME \"GNU GCC/G++\"\n#elif defined(__HP_cc) || defined(__HP_aCC)\n#define AILEGO_VERSION_COMPILER_NAME \"Hewlett-Packard C/aC++\"\n#elif defined(__IBMC__) || defined(__IBMCPP__)\n#define AILEGO_VERSION_COMPILER_NAME \"IBM XL C/C++\"\n#elif defined(_MSC_VER)\n#define AILEGO_VERSION_COMPILER_NAME \"Microsoft Visual C++\"\n#elif defined(__PGI)\n#define AILEGO_VERSION_COMPILER_NAME \"Portland Group PGCC/PGCPP\"\n#elif defined(__SUNPRO_C) || defined(__SUNPRO_CC)\n#define AILEGO_VERSION_COMPILER_NAME \"Oracle Solaris Studio\"\n#else\n#define AILEGO_VERSION_COMPILER_NAME \"Unknown Compiler\"\n#endif\n\n#if defined(__CUDACC_VER_MAJOR__)\n#define AILEGO_VERSION_COMPILER \\\n AILEGO_VERSION_COMPILER_NAME \\\n \" (\" AILEGO_VERSION_TO_STRING(__CUDACC_VER_MAJOR__) \\\n \".\" AILEGO_VERSION_TO_STRING(__CUDACC_VER_MINOR__) \\\n \".\" AILEGO_VERSION_TO_STRING(__CUDACC_VER_BUILD__) \")\"\n#elif defined(__VERSION__)\n#define AILEGO_VERSION_COMPILER \\\n AILEGO_VERSION_COMPILER_NAME \" (\" __VERSION__ \")\"\n#elif defined(_MSC_FULL_VER)\n#define AILEGO_VERSION_COMPILER \\\n AILEGO_VERSION_COMPILER_NAME \" (\" AILEGO_VERSION_TO_STRING(_MSC_FULL_VER) \")\"\n#elif defined(_MSC_VER)\n#define AILEGO_VERSION_COMPILER \\\n AILEGO_VERSION_COMPILER_NAME \" (\" AILEGO_VERSION_TO_STRING(_MSC_VER) \")\"\n#elif defined(__PGIC__)\n#define AILEGO_VERSION_COMPILER \\\n AILEGO_VERSION_COMPILER_NAME \\\n \" (\" AILEGO_VERSION_TO_STRING(__PGIC__) \".\" AILEGO_VERSION_TO_STRING( \\\n __PGIC_MINOR__) \".\" AILEGO_VERSION_TO_STRING(__PGIC_PATCHLEVEL__) \")\"\n#elif defined(__xlc__)\n#define AILEGO_VERSION_COMPILER AILEGO_VERSION_COMPILER_NAME \" (\" __xlc__ \")\"\n#elif defined(__SUNPRO_C)\n#define AILEGO_VERSION_COMPILER \\\n AILEGO_VERSION_COMPILER_NAME \" (\" AILEGO_VERSION_TO_STRING(__SUNPRO_C) \")\"\n#elif defined(__HP_cc)\n#define AILEGO_VERSION_COMPILER \\\n AILEGO_VERSION_COMPILER_NAME \" (\" AILEGO_VERSION_TO_STRING(__HP_cc) \")\"\n#else\n#define AILEGO_VERSION_COMPILER AILEGO_VERSION_COMPILER_NAME\n#endif\n\n#if defined(__x86_64__) || defined(_M_X64)\n#define AILEGO_VERSION_PROCESSOR \"x86 64-bit Processor\"\n#elif defined(__i386) || defined(_M_IX86)\n#define AILEGO_VERSION_PROCESSOR \"x86 32-bit Processor\"\n#elif defined(__ARM_ARCH)\n#if defined(__ARM_64BIT_STATE)\n#define AILEGO_VERSION_PROCESSOR \"ARM 64-bit Processor\"\n#else\n#define AILEGO_VERSION_PROCESSOR \"ARM 32-bit Processor\"\n#endif\n#elif defined(__ia64) || defined(__itanium__) || defined(_M_IA64)\n#define AILEGO_VERSION_PROCESSOR \"Itanium Processor\"\n#elif defined(__powerpc64__) || defined(__ppc64__) || defined(__PPC64__)\n#define AILEGO_VERSION_PROCESSOR \"PowerPC 64-bit Processor\"\n#elif defined(__powerpc__) || defined(__ppc__) || defined(__PPC__)\n#define AILEGO_VERSION_PROCESSOR \"PowerPC 32-bit Processor\"\n#elif defined(__sparc)\n#define AILEGO_VERSION_PROCESSOR \"SPARC Processor\"\n#else\n#define AILEGO_VERSION_PROCESSOR \"Unknown Processor\"\n#endif\n\n#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__\n#define AILEGO_VERSION_BYTE_ORDER \" Little-endian Byte Order\\n\"\n#elif __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__\n#define AILEGO_VERSION_BYTE_ORDER \" Big-endian Byte Order\\n\"\n#elif __BYTE_ORDER__ == __ORDER_PDP_ENDIAN__\n#define AILEGO_VERSION_BYTE_ORDER \" PDP-endian Byte Order\\n\"\n#else\n#define AILEGO_VERSION_BYTE_ORDER \"\"\n#endif\n\n#if defined(_DEBUG) || (!defined(__OPTIMIZE__) && !defined(NDEBUG))\n#define AILEGO_VERSION_DEBUG_INFO \" Debug Information\\n\"\n#else\n#define AILEGO_VERSION_DEBUG_INFO \"\"\n#endif\n\n#if defined(__SANITIZE_ADDRESS__)\n#define AILEGO_VERSION_ASAN \" Address Sanitizer\\n\"\n#else\n#define AILEGO_VERSION_ASAN \"\"\n#endif\n\n#if defined(__STDC_VERSION__)\n#define AILEGO_VERSION_STDC \\\n \" C Standard \" AILEGO_VERSION_TO_STRING(__STDC_VERSION__) \"\\n\"\n#else\n#define AILEGO_VERSION_STDC \"\"\n#endif\n\n#if defined(__cplusplus)\n#define AILEGO_VERSION_CPLUSPLUS \\\n \" C++ Standard \" AILEGO_VERSION_TO_STRING(__cplusplus) \"\\n\"\n#else\n#define AILEGO_VERSION_CPLUSPLUS \"\"\n#endif\n\n#if defined(__GXX_ABI_VERSION)\n#define AILEGO_VERSION_GXX_ABI \\\n \" GNU C++ ABI \" AILEGO_VERSION_TO_STRING(__GXX_ABI_VERSION) \"\\n\"\n#else\n#define AILEGO_VERSION_GXX_ABI \"\"\n#endif\n\n#if defined(__GLIBC__)\n#define AILEGO_VERSION_GLIBC \\\n \" GNU glibc \" AILEGO_VERSION_TO_STRING( \\\n __GLIBC__) \".\" AILEGO_VERSION_TO_STRING(__GLIBC_MINOR__) \"\\n\"\n#else\n#define AILEGO_VERSION_GLIBC \"\"\n#endif\n\n#if defined(WINVER)\n#define AILEGO_VERSION_WINSDK \\\n \" Microsoft Windows SDK \" AILEGO_VERSION_TO_STRING(WINVER) \"\\n\"\n#else\n#define AILEGO_VERSION_WINSDK \"\"\n#endif\n\n#if defined(__CLR_VER)\n#define AILEGO_VERSION_CLR \\\n \" Microsoft CLR \" AILEGO_VERSION_TO_STRING(__CLR_VER) \"\\n\"\n#else\n#define AILEGO_VERSION_CLR \"\"\n#endif\n\n#if defined(__LSB_VERSION__)\n#define AILEGO_VERSION_LSB \\\n \" Linux Standards Base \" AILEGO_VERSION_TO_STRING(__LSB_VERSION__) \"\\n\"\n#else\n#define AILEGO_VERSION_LSB \"\"\n#endif\n\n#if defined(_POSIX_VERSION)\n#define AILEGO_VERSION_POSIX \\\n \" POSIX Specification \" AILEGO_VERSION_TO_STRING(_POSIX_VERSION) \"\\n\"\n#else\n#define AILEGO_VERSION_POSIX \"\"\n#endif\n\n#if defined(_XOPEN_VERSION)\n#define AILEGO_VERSION_XOPEN \\\n \" X/Open Specification \" AILEGO_VERSION_TO_STRING(_XOPEN_VERSION) \"\\n\"\n#else\n#define AILEGO_VERSION_XOPEN \"\"\n#endif\n\n#if defined(_OPENMP)\n#define AILEGO_VERSION_OPENMP \\\n \" OpenMP API \" AILEGO_VERSION_TO_STRING(_OPENMP) \"\\n\"\n#else\n#define AILEGO_VERSION_OPENMP \"\"\n#endif\n\n#if defined(__ARM_NEON)\n#define AILEGO_VERSION_SIMD \" Arm Neon Instruction Set\\n\"\n#elif defined(__AVX512FP16__)\n#define AILEGO_VERSION_SIMD \" AVX-512FP16 Instruction Set\\n\"\n#elif defined(__AVX512F__)\n#define AILEGO_VERSION_SIMD \" AVX-512F Instruction Set\\n\"\n#elif defined(__AVX2__)\n#define AILEGO_VERSION_SIMD \" AVX-2 Instruction Set\\n\"\n#elif defined(__AVX__)\n#define AILEGO_VERSION_SIMD \" AVX Instruction Set\\n\"\n#elif defined(__SSE4_2__)\n#define AILEGO_VERSION_SIMD \" SSE-4.2 Instruction Set\\n\"\n#elif defined(__SSE4_1__)\n#define AILEGO_VERSION_SIMD \" SSE-4.1 Instruction Set\\n\"\n#elif defined(__SSSE3__)\n#define AILEGO_VERSION_SIMD \" SSSE-3 Instruction Set\\n\"\n#elif defined(__SSE3__)\n#define AILEGO_VERSION_SIMD \" SSE-3 Instruction Set\\n\"\n#elif defined(__SSE2__)\n#define AILEGO_VERSION_SIMD \" SSE-2 Instruction Set\\n\"\n#elif defined(__SSE__)\n#define AILEGO_VERSION_SIMD \" SSE Instruction Set\\n\"\n#elif defined(__MMX__)\n#define AILEGO_VERSION_SIMD \" MMX Instruction Set\\n\"\n#else\n#define AILEGO_VERSION_SIMD \"\"\n#endif\n\n#if defined(PY_VERSION)\n#if PY_RELEASE_LEVEL == PY_RELEASE_LEVEL_ALPHA\n#define AILEGO_VERSION_PYTHON \\\n \" Python API \" PY_VERSION \\\n \" Alpha \" AILEGO_VERSION_TO_STRING(PY_RELEASE_SERIAL) \"\\n\"\n#elif PY_RELEASE_LEVEL == PY_RELEASE_LEVEL_BETA\n#define AILEGO_VERSION_PYTHON \\\n \" Python API \" PY_VERSION \\\n \" Beta \" AILEGO_VERSION_TO_STRING(PY_RELEASE_SERIAL) \"\\n\"\n#elif PY_RELEASE_LEVEL == PY_RELEASE_LEVEL_GAMMA\n#define AILEGO_VERSION_PYTHON \\\n \" Python API \" PY_VERSION \\\n \" Release Candidate \" AILEGO_VERSION_TO_STRING(PY_RELEASE_SERIAL) \"\\n\"\n#elif PY_RELEASE_LEVEL == PY_RELEASE_LEVEL_FINAL\n#define AILEGO_VERSION_PYTHON \" Python API \" PY_VERSION \" Final\\n\"\n#else\n#define AILEGO_VERSION_PYTHON \" Python API \" PY_VERSION \"\\n\"\n#endif\n#else\n#define AILEGO_VERSION_PYTHON \"\"\n#endif\n\n//! Gather information of compiling\n#define AILEGO_VERSION_COMPILE_DETAILS(__PREFIX_INFO__) \\\n __PREFIX_INFO__ \\\n \"Compiled by \" AILEGO_VERSION_COMPILER \\\n \".\\n\" \\\n \"Compiled for \" AILEGO_VERSION_PROCESSOR \\\n \".\\n\" \\\n \"Compiled on \" AILEGO_VERSION_PLATFORM \" on \" __DATE__ \" \" __TIME__ \\\n \".\\n\" \\\n \"Compiled with: \\n\" \\\n \"\" AILEGO_VERSION_BYTE_ORDER \"\" AILEGO_VERSION_SIMD \\\n \"\" AILEGO_VERSION_DEBUG_INFO \"\" AILEGO_VERSION_ASAN \"\" AILEGO_VERSION_STDC \\\n \"\" AILEGO_VERSION_CPLUSPLUS \"\" AILEGO_VERSION_GXX_ABI \\\n \"\" AILEGO_VERSION_POSIX \"\" AILEGO_VERSION_XOPEN \"\" AILEGO_VERSION_LSB \\\n \"\" AILEGO_VERSION_GLIBC \"\" AILEGO_VERSION_WINSDK \"\" AILEGO_VERSION_CLR \\\n \"\" AILEGO_VERSION_OPENMP \"\" AILEGO_VERSION_PYTHON \"\\n\"\n" }, { "path": "src/binding/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\n# Retrieve version from git repository\ngit_version(ZVEC_VERSION ${CMAKE_CURRENT_SOURCE_DIR})\n\n# Add repository\ncc_directory(python)" }, { "path": "src/binding/python/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\nset(CMAKE_CXX_STANDARD 17)\nset(CMAKE_CXX_STANDARD_REQUIRED ON)\n\nfind_package(pybind11 REQUIRED)\n\nset(SRC_LISTS\n binding.cc\n model/python_collection.cc\n model/python_doc.cc\n model/param/python_param.cc\n model/schema/python_schema.cc\n model/common/python_config.cc\n typing/python_type.cc\n)\n\npybind11_add_module(_zvec ${SRC_LISTS})\n\nif (CMAKE_SYSTEM_NAME STREQUAL \"Linux\")\n target_link_libraries(_zvec PRIVATE\n -Wl,--whole-archive\n $\n $\n $\n $\n $\n $\n $\n $\n $\n $\n $\n -Wl,--no-whole-archive\n zvec_db\n )\n target_link_options(_zvec PRIVATE\n \"LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/exports.map\"\n )\nelseif (APPLE)\n target_link_libraries(_zvec PRIVATE\n -Wl,-force_load,$\n -Wl,-force_load,$\n -Wl,-force_load,$\n -Wl,-force_load,$\n -Wl,-force_load,$\n -Wl,-force_load,$\n -Wl,-force_load,$\n -Wl,-force_load,$\n -Wl,-force_load,$\n -Wl,-force_load,$\n -Wl,-force_load,$\n zvec_db\n )\n target_link_libraries(_zvec PRIVATE\n -Wl,-exported_symbols_list,${CMAKE_CURRENT_SOURCE_DIR}/exports.mac\n )\nendif ()\n\ntarget_include_directories(_zvec PRIVATE ${PYBIND11_INCLUDE_DIR} ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/binding/python/include)\n" }, { "path": "src/binding/python/binding.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"python_collection.h\"\n#include \"python_config.h\"\n#include \"python_doc.h\"\n#include \"python_param.h\"\n#include \"python_schema.h\"\n#include \"python_type.h\"\n\nnamespace zvec {\nPYBIND11_MODULE(_zvec, m) {\n m.doc() = \"Zvec core module\";\n\n ZVecPyTyping::Initialize(m);\n ZVecPyParams::Initialize(m);\n ZVecPySchemas::Initialize(m);\n ZVecPyConfig::Initialize(m);\n ZVecPyDoc::Initialize(m);\n ZVecPyCollection::Initialize(m);\n}\n} // namespace zvec\n" }, { "path": "src/binding/python/exports.mac", "content": "_PyInit__zvec\n" }, { "path": "src/binding/python/include/python_collection.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.#pragma once\n\n#include \n#include \n\nnamespace py = pybind11;\n\nnamespace zvec {\n\nclass ZVecPyCollection {\n public:\n ZVecPyCollection() = delete;\n\n public:\n static void Initialize(py::module_ &m);\n\n private:\n static void bind_db_methods(py::class_ &col);\n static void bind_ddl_methods(py::class_ &col);\n static void bind_dml_methods(py::class_ &col);\n static void bind_dql_methods(py::class_ &col);\n};\n\n} // namespace zvec\n" }, { "path": "src/binding/python/include/python_config.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.#pragma once\n\n#include \n#include \n\nnamespace py = pybind11;\n\nnamespace zvec {\n\nclass ZVecPyConfig {\n public:\n ZVecPyConfig() = delete;\n\n public:\n static void Initialize(py::module_ &m);\n};\n\n} // namespace zvec\n" }, { "path": "src/binding/python/include/python_doc.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.#pragma once\n\n#include \n#include \n\nnamespace py = pybind11;\n\nnamespace zvec {\n\nclass ZVecPyDoc {\n public:\n ZVecPyDoc() = delete;\n\n public:\n static void Initialize(py::module_ &m);\n\n private:\n static void bind_doc_operator(py::module_ &m);\n static void bind_doc(py::module_ &m);\n};\n\n} // namespace zvec\n" }, { "path": "src/binding/python/include/python_param.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.#pragma once\n\n#include \n#include \n#include \n\nnamespace py = pybind11;\n\nnamespace zvec {\n\nclass ZVecPyParams {\n public:\n ZVecPyParams() = delete;\n\n public:\n static void Initialize(py::module_ &m);\n\n private:\n static void bind_index_params(py::module_ &m);\n static void bind_query_params(py::module_ &m);\n static void bind_options(py::module_ &m);\n static void bind_vector_query(py::module_ &m);\n};\n\n} // namespace zvec\n" }, { "path": "src/binding/python/include/python_schema.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.#pragma once\n\n#include \n#include \n\nnamespace py = pybind11;\n\nnamespace zvec {\n\nclass ZVecPySchemas {\n public:\n ZVecPySchemas() = delete;\n\n public:\n static void Initialize(py::module_ &m);\n\n private:\n static void bind_field_schema(py::module_ &m);\n static void bind_collection_schema(py::module_ &m);\n static void bind_collection_stats(py::module_ &m);\n};\n\n} // namespace zvec\n" }, { "path": "src/binding/python/include/python_type.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.#pragma once\n\n#include \n#include \n#include \n\nnamespace py = pybind11;\n\nnamespace zvec {\n\nclass ZVecPyTyping {\n public:\n ZVecPyTyping() = delete;\n\n public:\n static void Initialize(py::module_ &m);\n\n private:\n static void bind_datatypes(py::module_ &m);\n static void bind_index_types(py::module_ &m);\n static void bind_metric_types(py::module_ &m);\n static void bind_quantize_types(py::module_ &m);\n static void bind_status(py::module_ &m);\n};\n\n} // namespace zvec\n" }, { "path": "src/binding/python/model/common/python_config.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"python_config.h\"\n#include \n\nnamespace zvec {\n\ninline bool has_key(py::dict d, const std::string &key) {\n return py::bool_(d.contains(key));\n}\n\ntemplate \nstd::optional get_if(py::dict d, const std::string &key) {\n if (has_key(d, key)) {\n try {\n py::object obj = d[py::str(key)];\n return obj.cast();\n } catch (const py::cast_error &) {\n throw py::type_error(\"Key '\" + key + \"' is not of expected type.\");\n }\n }\n return std::nullopt;\n}\n\ninline std::string to_lower(const std::string &s) {\n std::string lower;\n lower.reserve(s.size());\n std::transform(s.begin(), s.end(), std::back_inserter(lower), ::tolower);\n return lower;\n}\n\ninline bool iequals(const std::string &a, const std::string &b) {\n return to_lower(a) == to_lower(b);\n}\n\nGlobalConfig::LogLevel str_to_loglevel(const std::string &s) {\n if (iequals(s, \"debug\")) return GlobalConfig::LogLevel::DEBUG;\n if (iequals(s, \"info\")) return GlobalConfig::LogLevel::INFO;\n if (iequals(s, \"warn\") || iequals(s, \"warning\"))\n return GlobalConfig::LogLevel::WARN;\n if (iequals(s, \"error\")) return GlobalConfig::LogLevel::ERROR;\n if (iequals(s, \"fatal\")) return GlobalConfig::LogLevel::FATAL;\n throw py::value_error(\"Invalid log level: \");\n}\n\n\nvoid ZVecPyConfig::Initialize(pybind11::module_ &m) {\n m.def(\"Initialize\", [](py::args args, py::kwargs kwargs) -> py::none {\n py::dict config_dict;\n // parse args\n for (auto &arg : args) {\n if (py::isinstance(arg)) {\n for (auto item : arg.cast()) {\n config_dict[item.first] = item.second;\n }\n } else {\n throw py::type_error(\"Positional argument must be a dict if provided\");\n }\n }\n\n // parser kwargs\n if (kwargs) {\n for (auto item : kwargs) {\n config_dict[item.first] = item.second;\n }\n }\n\n if (config_dict.empty()) {\n return py::none();\n }\n\n GlobalConfig::ConfigData data;\n // config memory_limit_mb\n if (has_key(config_dict, \"memory_limit_mb\")) {\n auto mb = get_if(config_dict, \"memory_limit_mb\").value();\n if (mb <= 0) throw py::value_error(\"memory_limit_mb must be positive\");\n data.memory_limit_bytes = static_cast(mb) * 1024 * 1024;\n }\n\n // config log\n bool has_log_type = has_key(config_dict, \"log_type\");\n bool has_log_level = has_key(config_dict, \"log_level\");\n if (has_log_type || has_log_level) {\n std::string log_type = \"console\";\n std::string log_level_str = \"warn\";\n\n if (has_log_type) {\n log_type = config_dict[\"log_type\"].cast();\n }\n if (has_log_level) {\n log_level_str = config_dict[\"log_level\"].cast();\n }\n auto log_level = str_to_loglevel(log_level_str);\n if (iequals(log_type, \"file\")) {\n std::string dir = DEFAULT_LOG_DIR;\n std::string basename = DEFAULT_LOG_BASENAME;\n uint32_t file_size = DEFAULT_LOG_FILE_SIZE;\n uint32_t overdue_days = DEFAULT_LOG_OVERDUE_DAYS;\n\n if (has_key(config_dict, \"log_dir\")) {\n dir = get_if(config_dict, \"log_dir\").value();\n }\n if (has_key(config_dict, \"log_basename\")) {\n basename = get_if(config_dict, \"log_basename\").value();\n }\n if (has_key(config_dict, \"log_file_size\")) {\n auto s = get_if(config_dict, \"log_file_size\").value();\n if (s <= 0) {\n throw py::value_error(\"log_file_size must be positive\");\n }\n file_size = static_cast(s);\n }\n if (has_key(config_dict, \"log_overdue_days\")) {\n std::cout << \" ** log_overdue_days: \" << overdue_days << std::endl;\n auto d = get_if(config_dict, \"log_overdue_days\").value();\n if (d <= 0) {\n throw py::value_error(\"log_overdue_days must be positive\");\n }\n overdue_days = static_cast(d);\n }\n\n data.log_config = std::make_shared(\n log_level, dir, basename, file_size, overdue_days);\n\n } else if (iequals(log_type, \"console\")) {\n data.log_config =\n std::make_shared(log_level);\n } else {\n throw py::value_error(\"log_type must be 'console' or 'file'\");\n }\n }\n\n // set query thread count\n if (has_key(config_dict, \"query_threads\")) {\n auto q = get_if(config_dict, \"query_threads\").value();\n if (q <= 0) throw py::value_error(\"query_threads must be positive\");\n data.query_thread_count = static_cast(q);\n }\n\n // set optimize thread count\n if (has_key(config_dict, \"optimize_threads\")) {\n auto o = get_if(config_dict, \"optimize_threads\").value();\n if (o <= 0) throw py::value_error(\"optimize_threads must be positive\");\n data.optimize_thread_count = static_cast(o);\n }\n\n // set invert_to_forward_scan_ratio\n if (has_key(config_dict, \"invert_to_forward_scan_ratio\")) {\n auto v =\n get_if(config_dict, \"invert_to_forward_scan_ratio\").value();\n if (v < 0.0 || v > 1.0) {\n throw py::value_error(\n \"invert_to_forward_scan_ratio must be in [0.0, 1.0]\");\n }\n data.invert_to_forward_scan_ratio = static_cast(v);\n }\n\n // set brute_force_by_keys_ratio\n if (has_key(config_dict, \"brute_force_by_keys_ratio\")) {\n auto v = get_if(config_dict, \"brute_force_by_keys_ratio\").value();\n if (v < 0.0 || v > 1.0) {\n throw py::value_error(\n \"brute_force_by_keys_ratio must be in [0.0, 1.0]\");\n }\n data.brute_force_by_keys_ratio = static_cast(v);\n }\n\n // initialize (contains validate)\n Status status = GlobalConfig::Instance().Initialize(data);\n if (!status.ok()) {\n throw std::runtime_error(\"Initialization failed: \" + status.message());\n }\n return py::none();\n });\n}\n\n\n} // namespace zvec" }, { "path": "src/binding/python/model/param/python_param.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"python_param.h\"\n#include \n#include \n#include \n#include \n#include \"python_doc.h\"\n\nnamespace zvec {\n\nstatic std::string index_type_to_string(const IndexType type) {\n switch (type) {\n case IndexType::INVERT:\n return \"INVERT\";\n case IndexType::FLAT:\n return \"FLAT\";\n case IndexType::IVF:\n return \"IVF\";\n case IndexType::HNSW:\n return \"HNSW\";\n case IndexType::HNSW_RABITQ:\n return \"HNSW_RABITQ\";\n default:\n return \"UNDEFINED\";\n }\n}\n\nstatic std::string metric_type_to_string(const MetricType type) {\n switch (type) {\n case MetricType::COSINE:\n return \"COSINE\";\n case MetricType::IP:\n return \"IP\";\n case MetricType::L2:\n return \"L2\";\n default:\n return \"UNDEFINED\";\n }\n}\n\nstatic std::string quantize_type_to_string(const QuantizeType type) {\n switch (type) {\n case QuantizeType::UNDEFINED:\n return \"UNDEFINED\";\n case QuantizeType::INT8:\n return \"INT8\";\n case QuantizeType::INT4:\n return \"INT4\";\n case QuantizeType::FP16:\n return \"FP16\";\n case QuantizeType::RABITQ:\n return \"RABITQ\";\n default:\n return \"UNDEFINED\";\n }\n}\n\ntemplate \nT checked_cast(const py::handle &h, const std::string &vector_field,\n const std::string &expected_type) {\n try {\n return py::cast(h);\n } catch (const py::cast_error &e) {\n std::string actual_type = std::string(py::str(py::type::of(h)));\n std::string msg =\n vector_field + \": expected \" + expected_type + \", got \" + actual_type;\n throw py::type_error(msg);\n }\n}\n\ntemplate \nstd::string serialize_vector(const T *data, size_t n) {\n std::string buf;\n buf.resize(n * sizeof(T));\n std::memcpy(buf.data(), data, n * sizeof(T));\n return buf;\n}\n\ntemplate \nstd::pair serialize_sparse_vector(\n const py::dict &sparse_dict, ValueCastFn &&value_caster) {\n const size_t n = sparse_dict.size();\n if (n == 0) return {{}, {}};\n\n std::string indices_buf;\n indices_buf.resize(n * sizeof(uint32_t));\n auto *indices_ptr = reinterpret_cast(indices_buf.data());\n\n std::string values_buf;\n values_buf.resize(n * sizeof(ValueType));\n auto *values_ptr = reinterpret_cast(values_buf.data());\n\n size_t i = 0;\n for (const auto &[py_key, py_val] : sparse_dict) {\n indices_ptr[i] = checked_cast(py_key, \"Sparse indices\", \"UINT32\");\n values_ptr[i] = value_caster(py_val, i);\n ++i;\n }\n return {std::move(indices_buf), std::move(values_buf)};\n}\n\nvoid ZVecPyParams::Initialize(pybind11::module_ &parent) {\n auto m =\n parent.def_submodule(\"param\", \"This module contains the params of Zvec\");\n\n // binding index_params [invert/hnsw/flat/ivf]\n bind_index_params(m);\n\n // bind query_params [hnsw/ivf]\n bind_query_params(m);\n\n // bind options [collection/index/optimize/column]\n bind_options(m);\n\n // bind vector query\n bind_vector_query(m);\n}\n\nvoid ZVecPyParams::bind_index_params(pybind11::module_ &m) {\n // binding base index params\n py::class_> index_params(\n m, \"IndexParam\", R\"pbdoc(\nBase class for all index parameter configurations.\n\nThis abstract base class defines the common interface for index types.\nIt should not be instantiated directly; use derived classes instead.\n\nAttributes:\n type (IndexType): The type of the index (e.g., HNSW, FLAT, INVERT).\n)pbdoc\");\n index_params\n .def_property_readonly(\n \"type\",\n [](const IndexParams &self) -> IndexType { return self.type(); },\n \"IndexType: The type of the index.\")\n .def(\"clone\", &IndexParams::clone, py::return_value_policy::copy)\n .def(\n \"__eq__\",\n [](const IndexParams &self, const py::object &other) {\n if (!py::isinstance(other)) return false;\n return self == other.cast();\n },\n py::is_operator())\n .def(\n \"to_dict\",\n [](const IndexParams &self) -> py::dict {\n py::dict dict;\n dict[\"type\"] = index_type_to_string(self.type());\n return dict;\n },\n \"Convert to dictionary with all fields\")\n .def(py::pickle(\n [](const IndexParams &self) { // __getstate__\n return py::make_tuple(self.type());\n },\n [](py::tuple t) { // __setstate__\n if (t.size() != 1)\n throw std::runtime_error(\"Invalid state for IndexParams\");\n return std::shared_ptr();\n }));\n\n // binding invert index params\n py::class_>\n invert_params(m, \"InvertIndexParam\", R\"pbdoc(\nParameters for configuring an invert index.\n\nThis class controls whether range query\noptimization is enabled for invert index structures.\n\nAttributes:\n type (IndexType): Always `IndexType.INVERTED`.\n enable_range_optimization (bool): Whether range optimization is enabled.\n enable_extended_wildcard (bool): Whether extended wildcard (suffix and infix) search is enabled.\n\nExamples:\n >>> params = InvertIndexParam(enable_range_optimization=True, enable_extended_wildcard=False)\n >>> print(params.enable_range_optimization)\n True\n >>> print(params.enable_extended_wildcard)\n False\n >>> config = params.to_dict()\n >>> print(config)\n {'enable_range_optimization': True, 'enable_extended_wildcard': False}\n)pbdoc\");\n invert_params\n .def(py::init(), py::arg(\"enable_range_optimization\") = false,\n py::arg(\"enable_extended_wildcard\") = false,\n R\"pbdoc(\nConstructs an InvertIndexParam instance.\n\nArgs:\n enable_range_optimization (bool, optional): If True, enables range query\n optimization for the invert index. Defaults to False.\n enable_extended_wildcard (bool, optional): If True, enables extended wildcard\n search including suffix and infix patterns. Defaults to False.\n)pbdoc\")\n .def_property_readonly(\"enable_range_optimization\",\n &InvertIndexParams::enable_range_optimization,\n R\"pbdoc(\nbool: Whether range optimization is enabled for this inverted index.\n)pbdoc\")\n .def_property_readonly(\"enable_extended_wildcard\",\n &InvertIndexParams::enable_extended_wildcard,\n R\"pbdoc(\nbool: Whether extended wildcard (suffix and infix) search is enabled.\nNote: Prefix search is always enabled regardless of this setting.\n)pbdoc\")\n .def(\n \"to_dict\",\n [](const InvertIndexParams &self) -> py::dict {\n py::dict dict;\n dict[\"enable_range_optimization\"] =\n self.enable_range_optimization();\n dict[\"enable_extended_wildcard\"] = self.enable_extended_wildcard();\n return dict;\n },\n \"Convert to dictionary with all fields\")\n .def(\"__repr__\",\n [](const InvertIndexParams &self) -> std::string {\n return \"{\"\n \"\\\"enable_range_optimization\\\":\" +\n std::to_string(self.enable_range_optimization()) +\n \",\"\n \"\\\"enable_extended_wildcard\\\":\" +\n std::to_string(self.enable_extended_wildcard()) + \"}\";\n })\n .def(py::pickle(\n [](const InvertIndexParams &self) { // __getstate__\n return py::make_tuple(self.enable_range_optimization(),\n self.enable_extended_wildcard());\n },\n [](py::tuple t) { // __setstate__\n if (t.size() != 2)\n throw std::runtime_error(\"Invalid state for InvertIndexParams\");\n return std::make_shared(t[0].cast(),\n t[1].cast());\n }));\n\n // binding base vector index params\n py::class_>\n vector_params(m, \"VectorIndexParam\", R\"pbdoc(\nBase class for vector index parameter configurations.\n\nEncapsulates common settings for all vector index types.\n\nAttributes:\n type (IndexType): The specific vector index type (e.g., HNSW, FLAT).\n metric_type (MetricType): Distance metric used for similarity search.\n quantize_type (QuantizeType): Optional vector quantization type.\n)pbdoc\");\n vector_params\n .def_property_readonly(\n \"metric_type\",\n [](const VectorIndexParams &self) -> MetricType {\n return self.metric_type();\n },\n \"MetricType: Distance metric (e.g., IP, COSINE, L2).\")\n .def_property_readonly(\n \"quantize_type\",\n [](const VectorIndexParams &self) -> QuantizeType {\n return self.quantize_type();\n },\n \"QuantizeType: Vector quantization type (e.g., FP16, INT8).\")\n .def(\n \"to_dict\",\n [](const VectorIndexParams &self) -> py::dict {\n py::dict dict;\n dict[\"type\"] = index_type_to_string(self.type());\n dict[\"metric_type\"] = metric_type_to_string(self.metric_type());\n dict[\"quantize_type\"] =\n quantize_type_to_string(self.quantize_type());\n return dict;\n },\n \"Convert to dictionary with all fields\")\n .def(py::pickle(\n [](const VectorIndexParams &self) { // __getstate__\n return py::make_tuple(self.type(), self.metric_type(),\n self.quantize_type());\n },\n [](py::tuple t) { // __setstate__\n if (t.size() != 3)\n throw std::runtime_error(\"Invalid state for VectorIndexParams\");\n // 基类，不能直接实例化，用于子类\n return std::shared_ptr();\n }));\n\n // binding hnsw index params\n py::class_>\n hnsw_params(m, \"HnswIndexParam\", R\"pbdoc(\nParameters for configuring an HNSW (Hierarchical Navigable Small World) index.\n\nHNSW is a graph-based approximate nearest neighbor search index. This class\nencapsulates its construction hyperparameters.\n\nAttributes:\n metric_type (MetricType): Distance metric used for similarity computation.\n Default is ``MetricType.IP`` (inner product).\n m (int): Number of bi-directional links created for every new element\n during construction. Higher values improve accuracy but increase\n memory usage and construction time. Default is 50.\n ef_construction (int): Size of the dynamic candidate list for nearest\n neighbors during index construction. Larger values yield better\n graph quality at the cost of slower build time. Default is 500.\n quantize_type (QuantizeType): Optional quantization type for vector\n compression (e.g., FP16, INT8). Default is `QuantizeType.UNDEFINED` to\n disable quantization.\n\nExamples:\n >>> from zvec.typing import MetricType, QuantizeType\n >>> params = HnswIndexParam(\n ... metric_type=MetricType.COSINE,\n ... m=16,\n ... ef_construction=200,\n ... quantize_type=QuantizeType.INT8\n ... )\n >>> print(params)\n {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8'}\n)pbdoc\");\n hnsw_params\n .def(py::init(),\n py::arg(\"metric_type\") = MetricType::IP,\n py::arg(\"m\") = core_interface::kDefaultHnswNeighborCnt,\n py::arg(\"ef_construction\") =\n core_interface::kDefaultHnswEfConstruction,\n py::arg(\"quantize_type\") = QuantizeType::UNDEFINED)\n .def_property_readonly(\n \"m\", &HnswIndexParams::m,\n \"int: Maximum number of neighbors per node in upper layers.\")\n .def_property_readonly(\n \"ef_construction\", &HnswIndexParams::ef_construction,\n \"int: Candidate list size during index construction.\")\n .def(\n \"to_dict\",\n [](const HnswIndexParams &self) -> py::dict {\n py::dict dict;\n dict[\"type\"] = index_type_to_string(self.type());\n dict[\"metric_type\"] = metric_type_to_string(self.metric_type());\n dict[\"m\"] = self.m();\n dict[\"ef_construction\"] = self.ef_construction();\n dict[\"quantize_type\"] =\n quantize_type_to_string(self.quantize_type());\n return dict;\n },\n \"Convert to dictionary with all fields\")\n .def(\"__repr__\",\n [](const HnswIndexParams &self) -> std::string {\n return \"{\"\n \"\\\"metric_type\\\":\" +\n metric_type_to_string(self.metric_type()) +\n \", \\\"m\\\":\" + std::to_string(self.m()) +\n \", \\\"ef_construction\\\":\" +\n std::to_string(self.ef_construction()) +\n \", \\\"quantize_type\\\":\" +\n quantize_type_to_string(self.quantize_type()) + \"}\";\n })\n .def(py::pickle(\n [](const HnswIndexParams &self) {\n return py::make_tuple(self.metric_type(), self.m(),\n self.ef_construction(), self.quantize_type());\n },\n [](py::tuple t) {\n if (t.size() != 4)\n throw std::runtime_error(\"Invalid state for HnswIndexParams\");\n return std::make_shared(\n t[0].cast(), t[1].cast(), t[2].cast(),\n t[3].cast());\n }));\n\n // binding hnsw rabitq index params\n py::class_>\n hnsw_rabitq_params(m, \"HnswRabitqIndexParam\", R\"pbdoc(\nParameters for configuring an HNSW (Hierarchical Navigable Small World) index with RabitQ quantization.\n\nHNSW is a graph-based approximate nearest neighbor search index. RabitQ is a\nquantization method that provides high compression with minimal accuracy loss.\n\nAttributes:\n metric_type (MetricType): Distance metric used for similarity computation.\n Default is ``MetricType.IP`` (inner product).\n m (int): Number of bi-directional links created for every new element\n during construction. Higher values improve accuracy but increase\n memory usage and construction time. Default is 50.\n ef_construction (int): Size of the dynamic candidate list for nearest\n neighbors during index construction. Larger values yield better\n graph quality at the cost of slower build time. Default is 500.\n\nExamples:\n >>> from zvec.typing import MetricType\n >>> params = HnswRabitqIndexParam(\n ... metric_type=MetricType.COSINE,\n ... m=16,\n ... ef_construction=200\n ... )\n >>> print(params)\n {'metric_type': 'COSINE', 'm': 16, 'ef_construction': 200}\n)pbdoc\");\n hnsw_rabitq_params\n .def(py::init(),\n py::arg(\"metric_type\") = MetricType::IP,\n py::arg(\"total_bits\") = core_interface::kDefaultRabitqTotalBits,\n py::arg(\"num_clusters\") = core_interface::kDefaultRabitqNumClusters,\n py::arg(\"m\") = core_interface::kDefaultHnswNeighborCnt,\n py::arg(\"ef_construction\") =\n core_interface::kDefaultHnswEfConstruction,\n py::arg(\"sample_count\") = 0)\n .def_property_readonly(\"m\", &HnswRabitqIndexParams::m,\n \"int: Maximum number of neighbors per node.\")\n .def_property_readonly(\n \"ef_construction\", &HnswRabitqIndexParams::ef_construction,\n \"int: Candidate list size during index construction.\")\n .def_property_readonly(\"total_bits\", &HnswRabitqIndexParams::total_bits,\n \"int: Total bits for RabitQ quantization.\")\n .def_property_readonly(\"num_clusters\",\n &HnswRabitqIndexParams::num_clusters,\n \"int: Number of clusters for RabitQ.\")\n .def_property_readonly(\"sample_count\",\n &HnswRabitqIndexParams::sample_count,\n \"int: Sample count for RabitQ training.\")\n .def(\n \"to_dict\",\n [](const HnswRabitqIndexParams &self) -> py::dict {\n py::dict dict;\n dict[\"type\"] = index_type_to_string(self.type());\n dict[\"metric_type\"] = metric_type_to_string(self.metric_type());\n dict[\"quantize_type\"] =\n quantize_type_to_string(self.quantize_type());\n dict[\"total_bits\"] = self.total_bits();\n dict[\"num_clusters\"] = self.num_clusters();\n dict[\"sample_count\"] = self.sample_count();\n dict[\"m\"] = self.m();\n dict[\"ef_construction\"] = self.ef_construction();\n return dict;\n },\n \"Convert to dictionary with all fields\")\n .def(\n \"__repr__\",\n [](const HnswRabitqIndexParams &self) -> std::string {\n return \"{\"\n \"\\\"type\\\":\\\"\" +\n index_type_to_string(self.type()) +\n \"\\\", \\\"metric_type\\\":\\\"\" +\n metric_type_to_string(self.metric_type()) +\n \"\\\", \\\"total_bits\\\":\" + std::to_string(self.total_bits()) +\n \", \\\"num_clusters\\\":\" + std::to_string(self.num_clusters()) +\n \", \\\"sample_count\\\":\" + std::to_string(self.sample_count()) +\n \", \\\"m\\\":\" + std::to_string(self.m()) +\n \", \\\"ef_construction\\\":\" +\n std::to_string(self.ef_construction()) +\n \", \\\"quantize_type\\\":\\\"\" +\n quantize_type_to_string(self.quantize_type()) + \"\\\"}\";\n })\n .def(py::pickle(\n [](const HnswRabitqIndexParams &self) {\n return py::make_tuple(self.metric_type(), self.total_bits(),\n self.num_clusters(), self.m(),\n self.ef_construction(), self.sample_count());\n },\n [](py::tuple t) {\n if (t.size() != 6)\n throw std::runtime_error(\n \"Invalid state for HnswRabitqIndexParams\");\n return std::make_shared(\n t[0].cast(), t[1].cast(), t[2].cast(),\n t[3].cast(), t[4].cast(), t[5].cast());\n }));\n\n // FlatIndexParams\n py::class_>\n flat_params(m, \"FlatIndexParam\", R\"pbdoc(\nParameters for configuring a flat (brute-force) index.\n\nA flat index performs exact nearest neighbor search by comparing the query\nvector against all vectors in the collection. It is simple, accurate, and\nsuitable for small to medium datasets or as a baseline.\n\nAttributes:\n metric_type (MetricType): Distance metric used for similarity computation.\n Default is ``MetricType.IP`` (inner product).\n quantize_type (QuantizeType): Optional quantization type for vector\n compression (e.g., FP16, INT8). Use ``QuantizeType.UNDEFINED`` to\n disable quantization. Default is ``QuantizeType.UNDEFINED``.\n\nExamples:\n >>> from zvec.typing import MetricType, QuantizeType\n >>> params = FlatIndexParam(\n ... metric_type=MetricType.L2,\n ... quantize_type=QuantizeType.FP16\n ... )\n >>> print(params)\n {'metric_type': 'L2', 'quantize_type': 'FP16'}\n)pbdoc\");\n flat_params\n .def(py::init(),\n py::arg(\"metric_type\") = MetricType::IP,\n py::arg(\"quantize_type\") = QuantizeType::UNDEFINED,\n R\"pbdoc(\nConstructs a FlatIndexParam instance.\n\nArgs:\n metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.\n quantize_type (QuantizeType, optional): Vector quantization type.\n Defaults to QuantizeType.UNDEFINED (no quantization).\n)pbdoc\")\n .def(\n \"to_dict\",\n [](const FlatIndexParams &self) -> py::dict {\n py::dict dict;\n dict[\"metric_type\"] = metric_type_to_string(self.metric_type());\n dict[\"quantize_type\"] =\n quantize_type_to_string(self.quantize_type());\n return dict;\n },\n \"Convert to dictionary with all fields\")\n .def(\"__repr__\",\n [](const FlatIndexParams &self) -> std::string {\n return \"{\"\n \"\\\"metric_type\\\":\" +\n metric_type_to_string(self.metric_type()) +\n \", \\\"quantize_type\\\":\" +\n quantize_type_to_string(self.quantize_type()) + \"}\";\n })\n .def(py::pickle(\n [](const FlatIndexParams &self) {\n return py::make_tuple(self.metric_type(), self.quantize_type());\n },\n [](py::tuple t) {\n if (t.size() != 2)\n throw std::runtime_error(\"Invalid state for FlatIndexParams\");\n return std::make_shared(t[0].cast(),\n t[1].cast());\n }));\n\n // IVFIndexParams\n py::class_>\n ivf_params(m, \"IVFIndexParam\", R\"pbdoc(\nParameters for configuring an IVF (Inverted File Index) index.\n\nIVF partitions the vector space into clusters (inverted lists). At query time,\nonly a subset of clusters is searched, providing a trade-off between speed\nand accuracy.\n\nAttributes:\n metric_type (MetricType): Distance metric used for similarity computation.\n Default is ``MetricType.IP`` (inner product).\n n_list (int): Number of clusters (inverted lists) to partition the dataset into.\n If set to 0, the system will auto-select a reasonable value based on data size.\n Default is 0 (auto).\n n_iters (int): Number of iterations for k-means clustering during index training.\n Higher values yield more stable centroids. Default is 10.\n use_soar (bool): Whether to enable SOAR (Scalable Optimized Adaptive Routing)\n for improved IVF search performance. Default is False.\n quantize_type (QuantizeType): Optional quantization type for vector\n compression (e.g., FP16, INT8). Default is ``QuantizeType.UNDEFINED``.\n\nExamples:\n >>> from zvec.typing import MetricType, QuantizeType\n >>> params = IVFIndexParam(\n ... metric_type=MetricType.COSINE,\n ... n_list=100,\n ... n_iters=15,\n ... use_soar=True,\n ... quantize_type=QuantizeType.INT8\n ... )\n >>> print(params.n_list)\n 100\n)pbdoc\");\n ivf_params\n .def(py::init(),\n py::arg(\"metric_type\") = MetricType::IP, py::arg(\"n_list\") = 0,\n py::arg(\"n_iters\") = 10, py::arg(\"use_soar\") = false,\n py::arg(\"quantize_type\") = QuantizeType::UNDEFINED,\n R\"pbdoc(\nConstructs an IVFIndexParam instance.\n\nArgs:\n metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.\n n_list (int, optional): Number of inverted lists (clusters). Set to 0 for auto.\n Defaults to 0.\n n_iters (int, optional): Number of k-means iterations during training.\n Defaults to 10.\n use_soar (bool, optional): Enable SOAR optimization. Defaults to False.\n quantize_type (QuantizeType, optional): Vector quantization type.\n Defaults to QuantizeType.UNDEFINED.\n)pbdoc\")\n .def_property_readonly(\"n_list\", &IVFIndexParams::n_list,\n \"int: Number of inverted lists (0 = auto).\")\n .def_property_readonly(\n \"n_iters\", &IVFIndexParams::n_iters,\n \"int: Number of k-means iterations during training.\")\n .def_property_readonly(\"use_soar\", &IVFIndexParams::use_soar,\n \"bool: Whether SOAR optimization is enabled.\")\n .def(\n \"to_dict\",\n [](const IVFIndexParams &self) -> py::dict {\n py::dict dict;\n dict[\"type\"] = index_type_to_string(self.type());\n dict[\"metric_type\"] = metric_type_to_string(self.metric_type());\n dict[\"n_list\"] = self.n_list();\n dict[\"n_iters\"] = self.n_iters();\n dict[\"use_soar\"] = self.use_soar();\n dict[\"quantize_type\"] =\n quantize_type_to_string(self.quantize_type());\n return dict;\n },\n \"Convert to dictionary with all fields\")\n .def(\"__repr__\",\n [](const IVFIndexParams &self) {\n return \"{\"\n \"\\\"metric_type\\\":\" +\n metric_type_to_string(self.metric_type()) +\n \", \\\"n_list\\\":\" + std::to_string(self.n_list()) +\n \", \\\"n_iters\\\":\" + std::to_string(self.n_iters()) +\n \", \\\"use_soar\\\":\" + std::to_string(self.use_soar()) +\n \", \\\"quantize_type\\\":\" +\n quantize_type_to_string(self.quantize_type()) + \"}\";\n })\n .def(py::pickle(\n [](const IVFIndexParams &self) {\n return py::make_tuple(self.metric_type(), self.n_list(),\n self.n_iters(), self.use_soar(),\n self.quantize_type());\n },\n [](py::tuple t) {\n if (t.size() != 5)\n throw std::runtime_error(\"Invalid state for IVFIndexParams\");\n return std::make_shared(\n t[0].cast(), t[1].cast(), t[2].cast(),\n t[3].cast(), t[4].cast());\n }));\n}\n\nvoid ZVecPyParams::bind_query_params(py::module_ &m) {\n // binding base query params\n py::class_> query_params(\n m, \"QueryParam\", R\"pbdoc(\nBase class for all query parameter configurations.\n\nThis abstract base class defines common query settings such as search radius\nand whether to force linear (brute-force) search. It should not be instantiated\ndirectly; use derived classes like `HnswQueryParam` or `IVFQueryParam`.\n\nAttributes:\n type (IndexType): The index type this query is configured for.\n radius (float): Search radius for range queries. Used in combination with\n top-k to filter results. Default is 0.0 (disabled).\n is_linear (bool): If True, forces brute-force linear search instead of\n using the index. Useful for debugging or small datasets. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n)pbdoc\");\n query_params\n .def_property_readonly(\n \"type\",\n [](const QueryParams &self) -> IndexType { return self.type(); },\n \"IndexType: The type of index this query targets.\")\n .def_property_readonly(\n \"radius\",\n [](const QueryParams &self) -> float { return self.radius(); },\n \"IndexType: The type of index this query targets.\")\n .def_property_readonly(\n \"is_linear\",\n [](const QueryParams &self) -> bool { return self.is_linear(); },\n \"bool: Whether to bypass the index and use brute-force linear \"\n \"search.\")\n .def_property_readonly(\n \"is_using_refiner\",\n [](const QueryParams &self) -> bool {\n return self.is_using_refiner();\n },\n \"bool: Whether to use refiner for the query.\")\n .def(py::pickle(\n [](const QueryParams &self) { // __getstate__\n return py::make_tuple(self.type(), self.radius(), self.is_linear());\n },\n [](py::tuple t) { // __setstate__\n if (t.size() != 3)\n throw std::runtime_error(\"Invalid state for QueryParams\");\n return std::shared_ptr();\n }));\n\n // binding hnsw query params\n py::class_>\n hnsw_params(m, \"HnswQueryParam\", R\"pbdoc(\nQuery parameters for HNSW (Hierarchical Navigable Small World) index.\n\nControls the trade-off between search speed and accuracy via the `ef` parameter.\n\nAttributes:\n type (IndexType): Always ``IndexType.HNSW``.\n ef (int): Size of the dynamic candidate list during search.\n Larger values improve recall but slow down search.\n Default is 300.\n radius (float): Search radius for range queries. Default is 0.0.\n is_linear (bool): Force linear search. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n\nExamples:\n >>> params = HnswQueryParam(ef=300)\n >>> print(params.ef)\n 300\n >>> print(params.to_dict() if hasattr(params, 'to_dict') else params)\n {\"type\":\"HNSW\", \"ef\":300}\n)pbdoc\");\n hnsw_params\n .def(py::init(),\n py::arg(\"ef\") = core_interface::kDefaultHnswEfSearch,\n py::arg(\"radius\") = 0.0f, py::arg(\"is_linear\") = false,\n py::arg(\"is_using_refiner\") = false,\n R\"pbdoc(\nConstructs an HnswQueryParam instance.\n\nArgs:\n ef (int, optional): Search-time candidate list size.\n Higher values improve accuracy. Defaults to 100.\n radius (float, optional): Search radius for range queries. Default is 0.0.\n is_linear (bool, optional): Force linear search. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n)pbdoc\")\n .def_property_readonly(\n \"ef\", [](const HnswQueryParams &self) -> int { return self.ef(); },\n \"int: Size of the dynamic candidate list during HNSW search.\")\n .def(\"__repr__\",\n [](const HnswQueryParams &self) -> std::string {\n return \"{\"\n \"\\\"type\\\":\" +\n index_type_to_string(self.type()) +\n \", \\\"ef\\\":\" + std::to_string(self.ef()) +\n \", \\\"radius\\\":\" + std::to_string(self.radius()) +\n \", \\\"is_linear\\\":\" + std::to_string(self.is_linear()) +\n \", \\\"is_using_refiner\\\":\" +\n std::to_string(self.is_using_refiner()) + \"}\";\n })\n .def(py::pickle(\n [](const HnswQueryParams &self) {\n return py::make_tuple(self.ef(), self.radius(), self.is_linear(),\n self.is_using_refiner());\n },\n [](py::tuple t) {\n if (t.size() != 4)\n throw std::runtime_error(\"Invalid state for HnswQueryParams\");\n auto obj = std::make_shared(t[0].cast());\n obj->set_radius(t[1].cast());\n obj->set_is_linear(t[2].cast());\n obj->set_is_using_refiner(t[3].cast());\n return obj;\n }));\n\n // binding ivf query params\n py::class_>\n ivf_params(m, \"IVFQueryParam\", R\"pbdoc(\nQuery parameters for IVF (Inverted File Index) index.\n\nControls how many inverted lists (`nprobe`) to visit during search.\n\nAttributes:\n type (IndexType): Always ``IndexType.IVF``.\n nprobe (int): Number of closest clusters (inverted lists) to search.\n Higher values improve recall but increase latency.\n Default is 10.\n radius (float): Search radius for range queries. Default is 0.0.\n is_linear (bool): Force linear search. Default is False.\n\nExamples:\n >>> params = IVFQueryParam(nprobe=20)\n >>> print(params.nprobe)\n 20\n)pbdoc\");\n ivf_params\n .def(py::init(), py::arg(\"nprobe\") = 10, R\"pbdoc(\nConstructs an IVFQueryParam instance.\n\nArgs:\n nprobe (int, optional): Number of inverted lists to probe during search.\n Higher values improve accuracy. Defaults to 10.\n)pbdoc\")\n .def_property_readonly(\n \"nprobe\",\n [](const IVFQueryParams &self) -> int { return self.nprobe(); },\n \"int: Number of inverted lists to search during IVF query.\")\n .def(\"__repr__\",\n [](const IVFQueryParams &self) -> std::string {\n return \"{\"\n \"\\\"type\\\":\" +\n index_type_to_string(self.type()) +\n \", \\\"nprobe\\\":\" + std::to_string(self.nprobe()) + \"}\";\n })\n .def(py::pickle(\n [](const IVFQueryParams &self) {\n return py::make_tuple(self.nprobe(), self.radius(),\n self.is_linear());\n },\n [](py::tuple t) {\n if (t.size() != 3)\n throw std::runtime_error(\"Invalid state for IVFQueryParams\");\n auto obj = std::make_shared(t[0].cast());\n obj->set_radius(t[1].cast());\n obj->set_is_linear(t[2].cast());\n return obj;\n }));\n\n // binding hnsw rabitq query params\n py::class_>\n hnsw_rabitq_query_params(m, \"HnswRabitqQueryParam\", R\"pbdoc(\nQuery parameters for HNSW RaBitQ (Hierarchical Navigable Small World with RaBitQ quantization) index.\n\nControls the trade-off between search speed and accuracy via the `ef` parameter.\nRaBitQ provides efficient quantization while maintaining high search quality.\n\nAttributes:\n type (IndexType): Always ``IndexType.HNSW_RABITQ``.\n ef (int): Size of the dynamic candidate list during search.\n Larger values improve recall but slow down search.\n Default is 300.\n radius (float): Search radius for range queries. Default is 0.0.\n is_linear (bool): Force linear search. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n\nExamples:\n >>> params = HnswRabitqQueryParam(ef=300)\n >>> print(params.ef)\n 300\n >>> print(params.to_dict() if hasattr(params, 'to_dict') else params)\n {\"type\":\"HNSW_RABITQ\", \"ef\":300}\n)pbdoc\");\n hnsw_rabitq_query_params\n .def(py::init(),\n py::arg(\"ef\") = core_interface::kDefaultHnswEfSearch,\n py::arg(\"radius\") = 0.0f, py::arg(\"is_linear\") = false,\n py::arg(\"is_using_refiner\") = false,\n R\"pbdoc(\nConstructs an HnswRabitqQueryParam instance.\n\nArgs:\n ef (int, optional): Search-time candidate list size.\n Higher values improve accuracy. Defaults to 300.\n radius (float, optional): Search radius for range queries. Default is 0.0.\n is_linear (bool, optional): Force linear search. Default is False.\n is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.\n)pbdoc\")\n .def_property_readonly(\n \"ef\",\n [](const HnswRabitqQueryParams &self) -> int { return self.ef(); },\n \"int: Size of the dynamic candidate list during HNSW RaBitQ search.\")\n .def(\"__repr__\",\n [](const HnswRabitqQueryParams &self) -> std::string {\n return \"{\"\n \"\\\"type\\\":\\\"\" +\n index_type_to_string(self.type()) +\n \"\\\", \\\"ef\\\":\" + std::to_string(self.ef()) +\n \", \\\"radius\\\":\" + std::to_string(self.radius()) +\n \", \\\"is_linear\\\":\" + std::to_string(self.is_linear()) +\n \", \\\"is_using_refiner\\\":\" +\n std::to_string(self.is_using_refiner()) + \"}\";\n })\n .def(py::pickle(\n [](const HnswRabitqQueryParams &self) {\n return py::make_tuple(self.ef(), self.radius(), self.is_linear(),\n self.is_using_refiner());\n },\n [](py::tuple t) {\n if (t.size() != 4)\n throw std::runtime_error(\n \"Invalid state for HnswRabitqQueryParams\");\n auto obj =\n std::make_shared(t[0].cast());\n obj->set_radius(t[1].cast());\n obj->set_is_linear(t[2].cast());\n obj->set_is_using_refiner(t[3].cast());\n return obj;\n }));\n}\n\nvoid ZVecPyParams::bind_options(py::module_ &m) { // binding collection options\n py::class_(m, \"CollectionOption\", R\"pbdoc(\nOptions for opening or creating a collection.\n\nAttributes:\n read_only (bool): Whether the collection is opened in read-only mode.\n Default is False.\n enable_mmap (bool): Whether to use memory-mapped I/O for data files.\n Default is True.\n\nExamples:\n >>> opt = CollectionOption(read_only=True, enable_mmap=False)\n >>> print(opt.read_only)\n True\n)pbdoc\")\n .def(py::init(), py::arg(\"read_only\") = false,\n py::arg(\"enable_mmap\") = true,\n R\"pbdoc(\nConstructs a CollectionOption instance.\n\nArgs:\n read_only (bool, optional): Open collection in read-only mode.\n Defaults to False.\n enable_mmap (bool, optional): Enable memory-mapped I/O.\n Defaults to True.\n)pbdoc\")\n .def_property_readonly(\n \"enable_mmap\",\n [](const CollectionOptions &self) { return self.enable_mmap_; })\n .def_property_readonly(\n \"read_only\",\n [](const CollectionOptions &self) { return self.read_only_; })\n .def(\"__repr__\",\n [](const CollectionOptions &self) -> std::string {\n return \"{\"\n \"\\\"enable_mmap\\\":\" +\n std::to_string(self.enable_mmap_) +\n \", \\\"read_only\\\":\" + std::to_string(self.read_only_) + \"}\";\n })\n .def(py::pickle(\n [](const CollectionOptions &self) {\n return py::make_tuple(self.read_only_, self.enable_mmap_,\n self.max_buffer_size_);\n },\n [](py::tuple t) {\n if (t.size() != 3)\n throw std::runtime_error(\n \"Invalid pickle data for CollectionOptions\");\n CollectionOptions obj{};\n obj.read_only_ = t[0].cast();\n obj.enable_mmap_ = t[1].cast();\n obj.max_buffer_size_ = t[2].cast();\n return obj;\n }));\n\n // SegmentOptions\n py::class_(m, \"SegmentOption\", R\"pbdoc(\nOptions for segment-level operations.\n\nCurrently, this class mirrors CollectionOption and is used internally.\nIt supports read-only mode, memory mapping, and buffer configuration.\n\nNote:\n This class is primarily for internal use. Most users should use\n CollectionOption instead.\n\nExamples:\n >>> opt = SegmentOption()\n >>> print(opt.enable_mmap)\n True\n)pbdoc\")\n .def(py::init<>(), \"Constructs a SegmentOption with default settings.\")\n .def_property_readonly(\n \"enable_mmap\",\n [](const SegmentOptions &self) { return self.enable_mmap_; },\n \"bool: Whether memory-mapped I/O is enabled.\")\n .def_property_readonly(\n \"read_only\",\n [](const SegmentOptions &self) { return self.read_only_; },\n \"bool: Whether the segment is read-only.\")\n .def_property_readonly(\n \"max_buffer_size\",\n [](const SegmentOptions &self) { return self.max_buffer_size_; },\n \"int: Maximum buffer size in bytes (internal use).\")\n .def(\"__repr__\",\n [](const SegmentOptions &self) -> std::string {\n return \"{\"\n \"\\\"enable_mmap\\\":\" +\n std::to_string(self.enable_mmap_) +\n \", \\\"read_only\\\":\" + std::to_string(self.read_only_) +\n \", \\\"max_buffer_size\\\":\" +\n std::to_string(self.max_buffer_size_) + \"}\";\n })\n .def(py::pickle(\n [](const SegmentOptions &self) {\n return py::make_tuple(self.read_only_, self.enable_mmap_,\n self.max_buffer_size_);\n },\n [](py::tuple t) {\n if (t.size() != 3)\n throw std::runtime_error(\n \"Invalid pickle data for SegmentOptions\");\n SegmentOptions obj{};\n obj.read_only_ = t[0].cast();\n obj.enable_mmap_ = t[1].cast();\n obj.max_buffer_size_ = t[2].cast();\n return obj;\n }));\n\n // CreateIndexOptions\n py::class_(m, \"IndexOption\",\n R\"pbdoc(\nOptions for creating an index.\n\nAttributes:\n concurrency (int): Number of threads to use during index creation.\n If 0, the system will choose an optimal value automatically.\n Default is 0.\n\nExamples:\n >>> opt = IndexOption(concurrency=4)\n >>> print(opt.concurrency)\n 4\n)pbdoc\")\n .def(py::init(), py::arg(\"concurrency\") = 0,\n R\"pbdoc(\nConstructs an IndexOption instance.\n\nArgs:\n concurrency (int, optional): Number of concurrent threads.\n 0 means auto-detect. Defaults to 0.\n)pbdoc\")\n .def_property_readonly(\n \"concurrency\",\n [](const CreateIndexOptions &self) { return self.concurrency_; },\n \"int: Number of threads used for index creation (0 = auto).\")\n .def(py::pickle(\n [](const CreateIndexOptions &self) {\n return py::make_tuple(self.concurrency_);\n },\n [](py::tuple t) {\n if (t.size() != 1)\n throw std::runtime_error(\n \"Invalid pickle data for CreateIndexOptions\");\n CreateIndexOptions obj{};\n obj.concurrency_ = t[0].cast();\n return obj;\n }));\n\n // OptimizeOptions\n py::class_(m, \"OptimizeOption\", R\"pbdoc(\nOptions for optimizing a collection (e.g., merging segments).\n\nAttributes:\n concurrency (int): Number of threads to use during optimization.\n If 0, the system will choose an optimal value automatically.\n Default is 0.\n\nExamples:\n >>> opt = OptimizeOption(concurrency=2)\n >>> print(opt.concurrency)\n 2\n)pbdoc\")\n .def(py::init(), py::arg(\"concurrency\") = 0,\n R\"pbdoc(\nConstructs an OptimizeOption instance.\n\nArgs:\n concurrency (int, optional): Number of concurrent threads.\n 0 means auto-detect. Defaults to 0.\n)pbdoc\")\n .def_property_readonly(\n \"concurrency\",\n [](const OptimizeOptions &self) { return self.concurrency_; },\n \"int: Number of threads used for optimization (0 = auto).\")\n .def(py::pickle(\n [](const OptimizeOptions &self) {\n return py::make_tuple(self.concurrency_);\n },\n [](py::tuple t) {\n if (t.size() != 1)\n throw std::runtime_error(\n \"Invalid pickle data for OptimizeOptions\");\n OptimizeOptions obj{};\n obj.concurrency_ = t[0].cast();\n return obj;\n }));\n\n // AddColumnOptions\n py::class_(m, \"AddColumnOption\",\n R\"pbdoc(\nOptions for adding a new column to a collection.\n\nAttributes:\n concurrency (int): Number of threads to use when backfilling data\n for the new column. If 0, auto-detect is used. Default is 0.\n\nExamples:\n >>> opt = AddColumnOption(concurrency=1)\n >>> print(opt.concurrency)\n 1\n)pbdoc\")\n .def(py::init(), py::arg(\"concurrency\") = 0,\n R\"pbdoc(\nConstructs an AddColumnOption instance.\n\nArgs:\n concurrency (int, optional): Number of threads for data backfill.\n 0 means auto-detect. Defaults to 0.\n)pbdoc\")\n .def_property_readonly(\n \"concurrency\",\n [](const AddColumnOptions &self) { return self.concurrency_; },\n \"int: Number of threads used when adding a column (0 = auto).\")\n .def(py::pickle(\n [](const AddColumnOptions &self) {\n return py::make_tuple(self.concurrency_);\n },\n [](py::tuple t) {\n if (t.size() != 1)\n throw std::runtime_error(\n \"Invalid pickle data for AddColumnOptions\");\n AddColumnOptions obj{};\n obj.concurrency_ = t[0].cast();\n return obj;\n }));\n\n // AlterColumnOptions\n py::class_(m, \"AlterColumnOption\", R\"pbdoc(\nOptions for altering an existing column (e.g., changing index settings).\n\nAttributes:\n concurrency (int): Number of threads to use during the alteration process.\n If 0, the system will choose an optimal value automatically.\n Default is 0.\n\nExamples:\n >>> opt = AlterColumnOption(concurrency=1)\n >>> print(opt.concurrency)\n 1\n)pbdoc\")\n .def(py::init(), py::arg(\"concurrency\") = 0,\n R\"pbdoc(\nConstructs an AlterColumnOption instance.\n\nArgs:\n concurrency (int, optional): Number of threads for column alteration.\n 0 means auto-detect. Defaults to 0.\n)pbdoc\")\n .def_property_readonly(\n \"concurrency\",\n [](const AlterColumnOptions &self) { return self.concurrency_; },\n \"int: Number of threads used when altering a column (0 = auto).\")\n .def(py::pickle(\n [](const AlterColumnOptions &self) {\n return py::make_tuple(self.concurrency_);\n },\n [](py::tuple t) {\n if (t.size() != 1)\n throw std::runtime_error(\n \"Invalid pickle data for AlterColumnOptions\");\n AlterColumnOptions obj{};\n obj.concurrency_ = t[0].cast();\n return obj;\n }));\n}\n\nvoid ZVecPyParams::bind_vector_query(py::module_ &m) {\n py::class_(m, \"_VectorQuery\")\n .def(py::init<>())\n // properties\n .def_readwrite(\"topk\", &VectorQuery::topk_)\n .def_readwrite(\"field_name\", &VectorQuery::field_name_)\n .def_readwrite(\"filter\", &VectorQuery::filter_)\n .def_readwrite(\"include_vector\", &VectorQuery::include_vector_)\n .def_readwrite(\"query_params\", &VectorQuery::query_params_)\n .def_readwrite(\"output_fields\", &VectorQuery::output_fields_)\n // vector\n .def(\"set_vector\",\n [](VectorQuery &self, const FieldSchema &field_schema,\n const py::object &obj) {\n const DataType data_type = field_schema.data_type();\n\n // dense vector\n if (FieldSchema::is_dense_vector_field(data_type)) {\n if (!py::isinstance(obj)) {\n throw py::type_error(\"Dense vector[\" + field_schema.name() +\n \"] expects a ndarray, got \" +\n std::string(py::str(py::type::of(obj))));\n }\n const auto arr = obj.cast();\n if (arr.ndim() != 1) {\n throw py::type_error(\"Dense vector expects 1D array, got \" +\n std::to_string(arr.ndim()) + \"D\");\n }\n const auto buf = arr.request();\n switch (data_type) {\n case DataType::VECTOR_FP32: {\n self.query_vector_ = serialize_vector(\n static_cast(buf.ptr), buf.size);\n return;\n }\n case DataType::VECTOR_FP64: {\n self.query_vector_ = serialize_vector(\n static_cast(buf.ptr), buf.size);\n return;\n }\n case DataType::VECTOR_INT8: {\n self.query_vector_ = serialize_vector(\n static_cast(buf.ptr), buf.size);\n return;\n }\n case DataType::VECTOR_FP16: {\n self.query_vector_ = serialize_vector(\n static_cast(buf.ptr), buf.size);\n return;\n }\n default:\n throw py::type_error(\n \"Unsupported dense vector type for ndarray input: \" +\n std::to_string(static_cast(data_type)));\n }\n }\n // sparse vector\n if (FieldSchema::is_sparse_vector_field(data_type)) {\n if (!py::isinstance(obj)) {\n throw py::type_error(\"Sparse vector[\" + field_schema.name() +\n \"] expects a Python dict, got \" +\n std::string(py::str(py::type::of(obj))));\n }\n const auto sparse = obj.cast();\n\n switch (data_type) {\n case DataType::SPARSE_VECTOR_FP16: {\n auto [indices, values] =\n serialize_sparse_vector(\n sparse, [](const py::handle &h, size_t idx) {\n float f = checked_cast(\n h, \"Sparse value[\" + std::to_string(idx) + \"]\",\n \"FLOAT\");\n return ailego::Float16(f);\n });\n self.query_sparse_indices_ = std::move(indices);\n self.query_sparse_values_ = std::move(values);\n break;\n }\n case DataType::SPARSE_VECTOR_FP32: {\n auto [indices, values] = serialize_sparse_vector(\n sparse, [](const py::handle &h, size_t idx) {\n return checked_cast(\n h, \"Sparse value[\" + std::to_string(idx) + \"]\",\n \"FLOAT\");\n });\n self.query_sparse_indices_ = std::move(indices);\n self.query_sparse_values_ = std::move(values);\n break;\n }\n default:\n throw py::type_error(\n \"Unsupported sparse vector type: \" +\n std::to_string(static_cast(data_type)));\n }\n return;\n }\n\n throw py::type_error(\"Unsupported vector field type for field: \" +\n field_schema.name());\n })\n .def(\n \"get_vector\",\n [](const VectorQuery &self,\n const FieldSchema &field_schema) -> py::object {\n DataType data_type = field_schema.data_type();\n if (FieldSchema::is_dense_vector_field(data_type)) {\n if (self.query_vector_.empty()) {\n throw std::runtime_error(\"No dense vector has been set\");\n }\n\n size_t byte_size = self.query_vector_.size();\n const void *data = self.query_vector_.data();\n\n switch (data_type) {\n case DataType::VECTOR_FP32: {\n if (byte_size % sizeof(float) != 0) {\n throw std::runtime_error(\n \"Invalid buffer size for VECTOR_FP32\");\n }\n size_t dim = byte_size / sizeof(float);\n return py::array_t({dim}, {sizeof(float)},\n static_cast(data));\n }\n case DataType::VECTOR_FP64: {\n if (byte_size % sizeof(double) != 0) {\n throw std::runtime_error(\n \"Invalid buffer size for VECTOR_FP64\");\n }\n size_t dim = byte_size / sizeof(double);\n return py::array_t({dim}, {sizeof(double)},\n static_cast(data));\n }\n case DataType::VECTOR_INT8: {\n if (byte_size % sizeof(int8_t) != 0) {\n throw std::runtime_error(\n \"Invalid buffer size for VECTOR_INT8\");\n }\n size_t dim = byte_size / sizeof(int8_t);\n return py::array_t({dim}, {sizeof(int8_t)},\n static_cast(data));\n }\n case DataType::VECTOR_FP16: {\n if (byte_size % 2 != 0) {\n throw std::runtime_error(\n \"Invalid buffer size for VECTOR_FP16\");\n }\n size_t dim = byte_size / 2;\n return py::array(py::dtype(\"float16\"), {dim}, {2}, data);\n }\n\n default:\n throw py::type_error(\n \"Unsupported dense vector type for get_vector: \" +\n std::to_string(static_cast(data_type)));\n }\n }\n if (FieldSchema::is_sparse_vector_field(data_type)) {\n if (self.query_sparse_indices_.empty()) {\n return py::dict();\n }\n\n // Deserialize indices: stored as uint32_t[]\n size_t indices_byte_size = self.query_sparse_indices_.size();\n if (indices_byte_size % sizeof(uint32_t) != 0) {\n throw std::runtime_error(\n \"Sparse indices buffer size not aligned to uint32_t\");\n }\n size_t n = indices_byte_size / sizeof(uint32_t);\n const uint32_t *indices = reinterpret_cast(\n self.query_sparse_indices_.data());\n\n // Deserialize values\n switch (data_type) {\n case DataType::SPARSE_VECTOR_FP32: {\n if (self.query_sparse_values_.size() != n * sizeof(float)) {\n throw std::runtime_error(\n \"Sparse FP32 values buffer size mismatch\");\n }\n const float *values = reinterpret_cast(\n self.query_sparse_values_.data());\n py::dict result;\n for (size_t i = 0; i < n; ++i) {\n result[py::int_(indices[i])] = py::float_(values[i]);\n }\n return result;\n }\n case DataType::SPARSE_VECTOR_FP16: {\n if (self.query_sparse_values_.size() !=\n n * sizeof(uint16_t)) {\n throw std::runtime_error(\n \"Sparse FP16 values buffer size mismatch\");\n }\n const uint16_t *raw_bits = reinterpret_cast(\n self.query_sparse_values_.data());\n py::dict result;\n for (size_t i = 0; i < n; ++i) {\n float f = ailego::FloatHelper::ToFP32(raw_bits[i]);\n result[py::int_(indices[i])] = py::float_(f);\n }\n return result;\n }\n default:\n throw py::type_error(\"Unsupported sparse vector type...\");\n }\n }\n\n\n throw py::type_error(\"Unsupported vector field type: \" +\n field_schema.name());\n },\n py::arg(\"field_schema\"))\n .def(py::pickle(\n [](const VectorQuery &self) {\n return py::make_tuple(\n self.topk_, self.field_name_, self.query_vector_,\n self.query_sparse_indices_, self.query_sparse_values_,\n self.filter_, self.include_vector_, self.output_fields_,\n self.query_params_ ? py::cast(self.query_params_) : py::none());\n },\n [](py::tuple t) {\n if (t.size() != 9)\n throw std::runtime_error(\"Invalid pickle data for VectorQuery\");\n\n VectorQuery obj{};\n obj.topk_ = t[0].cast();\n obj.field_name_ = t[1].cast();\n obj.query_vector_ = t[2].cast();\n obj.query_sparse_indices_ = t[3].cast();\n obj.query_sparse_values_ = t[4].cast();\n obj.filter_ = t[5].cast();\n obj.include_vector_ = t[6].cast();\n obj.output_fields_ = t[7].cast>();\n\n if (!t[8].is_none()) {\n obj.query_params_ = t[8].cast();\n }\n return obj;\n }));\n}\n} // namespace zvec" }, { "path": "src/binding/python/model/python_collection.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"python_collection.h\"\n#include \n#include \n\nnamespace zvec {\n\ninline void throw_if_error(const Status &status) {\n switch (status.code()) {\n case StatusCode::OK:\n return;\n case StatusCode::NOT_FOUND:\n throw py::key_error(status.message());\n case StatusCode::INVALID_ARGUMENT:\n throw py::value_error(status.message());\n case StatusCode::INTERNAL_ERROR:\n case StatusCode::ALREADY_EXISTS:\n case StatusCode::NOT_SUPPORTED:\n case StatusCode::PERMISSION_DENIED:\n case StatusCode::FAILED_PRECONDITION:\n case StatusCode::UNKNOWN:\n default:\n throw std::runtime_error(status.message());\n }\n}\n\n\ntemplate \nT unwrap_expected(const tl::expected &exp) {\n if (exp.has_value()) {\n return exp.value();\n }\n throw_if_error(exp.error());\n return T{};\n}\n\nvoid ZVecPyCollection::Initialize(pybind11::module_ &m) {\n py::class_ collection(m, \"_Collection\");\n bind_db_methods(collection);\n bind_ddl_methods(collection);\n bind_dml_methods(collection);\n bind_dql_methods(collection);\n collection.def(py::pickle(\n [](const Collection &c) {\n return py::make_tuple(c.Path(), c.Schema(), c.Options());\n },\n [](py::tuple t) {\n if (t.size() != 3) {\n throw std::runtime_error(\"Invalid tuple size for Collection pickle\");\n }\n std::string path = t[0].cast();\n auto schema = t[1].cast();\n CollectionOptions options = t[2].cast();\n auto result = Collection::Open(path, options);\n // auto result = Collection::CreateAndOpen(path, schema, options);\n return unwrap_expected(result);\n }));\n}\n\nvoid ZVecPyCollection::bind_db_methods(\n py::class_ &col) {\n col.def_static(\"CreateAndOpen\",\n [](const std::string &path, const CollectionSchema &schema,\n const CollectionOptions &options) {\n auto result =\n Collection::CreateAndOpen(path, schema, options);\n return unwrap_expected(result);\n })\n .def_static(\"Open\", [](const std::string &path,\n const CollectionOptions &options) {\n auto result = Collection::Open(path, options);\n return unwrap_expected(result);\n });\n}\n\n\nvoid ZVecPyCollection::bind_ddl_methods(\n py::class_ &col) {\n // bind collection properties\n col.def(\"Path\",\n [](const Collection &self) {\n auto ret = self.Path();\n return unwrap_expected(ret);\n })\n .def(\"Options\",\n [](const Collection &self) {\n auto ret = self.Options();\n return unwrap_expected(ret);\n })\n .def(\"Schema\",\n [](const Collection &self) {\n auto ret = self.Schema();\n return unwrap_expected(ret);\n })\n .def(\"Stats\", [](const Collection &self) {\n auto ret = self.Stats();\n return unwrap_expected(ret);\n });\n\n // bind collection ddl methods\n col.def(\"Destroy\",\n [](Collection &self) {\n const auto status = self.Destroy();\n throw_if_error(status);\n })\n .def(\"Flush\", [](Collection &self) {\n auto status = self.Flush();\n throw_if_error(status);\n });\n\n // binding index ddl methods\n col.def(\"CreateIndex\",\n [](Collection &self, const std::string &column_name,\n const IndexParams::Ptr &index_options,\n const CreateIndexOptions &options) {\n const auto status =\n self.CreateIndex(column_name, index_options, options);\n throw_if_error(status);\n })\n .def(\"DropIndex\",\n [](Collection &self, const std::string &column_name) {\n const auto status = self.DropIndex(column_name);\n throw_if_error(status);\n })\n .def(\"Optimize\", [](Collection &self, const OptimizeOptions &options) {\n const auto status = self.Optimize(options);\n throw_if_error(status);\n });\n\n // binding column ddl methods\n col.def(\"AddColumn\",\n [](Collection &self, const FieldSchema::Ptr &column_schema,\n const std::string &expression, const AddColumnOptions &options) {\n const auto status =\n self.AddColumn(column_schema, expression, options);\n throw_if_error(status);\n })\n .def(\"DropColumn\",\n [](Collection &self, std::string &column_name) {\n auto status = self.DropColumn(column_name);\n throw_if_error(status);\n })\n .def(\"AlterColumn\", [](Collection &self, std::string &column_name,\n const std::string &rename,\n const FieldSchema::Ptr &new_column_schema,\n const AlterColumnOptions &options) {\n const auto status =\n self.AlterColumn(column_name, rename, new_column_schema, options);\n throw_if_error(status);\n });\n}\n\nvoid ZVecPyCollection::bind_dml_methods(\n py::class_ &col) {\n // bind collection upsert/insert/update/delete methods\n col.def(\"Insert\",\n [](Collection &self, std::vector &docs) {\n const auto result = self.Insert(docs);\n return unwrap_expected(result);\n })\n .def(\"Update\",\n [](Collection &self, std::vector &docs) {\n const auto result = self.Update(docs);\n return unwrap_expected(result);\n })\n .def(\"Upsert\",\n [](Collection &self, std::vector &docs) {\n const auto result = self.Upsert(docs);\n return unwrap_expected(result);\n })\n .def(\"Delete\",\n [](Collection &self, const std::vector &pks) {\n const auto result = self.Delete(pks);\n return unwrap_expected(result);\n })\n .def(\"DeleteByFilter\", [](Collection &self, const std::string &filter) {\n const auto status = self.DeleteByFilter(filter);\n throw_if_error(status);\n });\n}\n\nvoid ZVecPyCollection::bind_dql_methods(\n py::class_ &col) {\n col.def(\"Query\",\n [](const Collection &self, const VectorQuery &query) {\n const auto result = self.Query(query);\n // return DocPtrList\n return unwrap_expected(result);\n })\n .def(\"GroupByQuery\",\n [](const Collection &self, const GroupByVectorQuery &query) {\n const auto result = self.GroupByQuery(query);\n // return GroupResults\n return unwrap_expected(result);\n })\n .def(\"Fetch\",\n [](const Collection &self, const std::vector &pks) {\n const auto result = self.Fetch(pks);\n // return DocPtrMap\n return unwrap_expected(result);\n });\n}\n\n} // namespace zvec" }, { "path": "src/binding/python/model/python_doc.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"python_doc.h\"\n#include \n#include \n\nnamespace zvec {\n\ntemplate \nT checked_cast(const py::object &obj, const std::string &field,\n const std::string &expected_type) {\n try {\n return obj.cast();\n } catch (const py::cast_error &e) {\n std::string actual_type = std::string(py::str(py::type::of(obj)));\n std::string msg = \"Field '\" + field + \"': expected \" + expected_type +\n \", got \" + actual_type;\n throw py::type_error(msg);\n }\n}\n\nvoid ZVecPyDoc::Initialize(pybind11::module_ &m) {\n bind_doc_operator(m);\n bind_doc(m);\n}\n\n\nvoid ZVecPyDoc::bind_doc_operator(py::module_ &m) {\n py::enum_(m, \"_DocOp\")\n .value(\"INSERT\", Operator::INSERT)\n .value(\"UPDATE\", Operator::UPDATE)\n .value(\"DELETE\", Operator::DELETE)\n .value(\"UPSERT\", Operator::UPSERT);\n}\n\n\nvoid ZVecPyDoc::bind_doc(py::module_ &m) {\n // binding doc\n py::class_ doc(m, \"_Doc\");\n\n doc.def(py::init([]() { return std::make_shared(); }))\n .def(\"set_pk\", &Doc::set_pk)\n .def(\"pk\", &Doc::pk)\n .def(\"set_score\", &Doc::set_score)\n .def(\"score\", &Doc::score)\n .def(\"has_field\", &Doc::has)\n .def(\"field_names\", &Doc::field_names)\n .def(py::pickle(\n [](const Doc &d) {\n std::vector data = d.serialize();\n return py::bytes(reinterpret_cast(data.data()),\n data.size());\n },\n [](py::bytes b) {\n py::buffer_info info(py::buffer(b).request());\n const uint8_t *buf = reinterpret_cast(info.ptr);\n size_t size = static_cast(info.size);\n Doc::Ptr d = Doc::deserialize(buf, size);\n if (!d) throw std::runtime_error(\"Failed to unpickle Doc\");\n return d;\n }));\n\n\n // binding doc set field\n doc.def(\n \"set_any\",\n [](Doc &self, const std::string &field, const FieldSchema &field_schema,\n const py::object &obj) -> bool {\n if (obj.is_none()) {\n if (field_schema.nullable()) {\n self.set_null(field);\n return true;\n }\n throw py::value_error(\"Field '\" + field +\n \"': expected non-nullable type\");\n }\n switch (field_schema.data_type()) {\n // base datatypes\n case DataType::STRING:\n return self.set(field,\n checked_cast(obj, field, \"STRING\"));\n case DataType::BOOL:\n return self.set(field, checked_cast(obj, field, \"BOOL\"));\n case DataType::INT32:\n return self.set(field, checked_cast(obj, field, \"INT32\"));\n case DataType::INT64:\n return self.set(field, checked_cast(obj, field, \"INT64\"));\n case DataType::UINT32:\n return self.set(field,\n checked_cast(obj, field, \"UINT32\"));\n case DataType::UINT64:\n return self.set(field,\n checked_cast(obj, field, \"UINT64\"));\n case DataType::FLOAT:\n return self.set(field, checked_cast(obj, field, \"FLOAT\"));\n case DataType::DOUBLE:\n return self.set(field, checked_cast(obj, field, \"DOUBLE\"));\n\n // array datatypes\n case DataType::ARRAY_STRING:\n return self.set(field, checked_cast>(\n obj, field, \"ARRAY_STRING\"));\n case DataType::ARRAY_BOOL:\n return self.set(field, checked_cast>(\n obj, field, \"ARRAY_BOOL\"));\n case DataType::ARRAY_INT32:\n return self.set(field, checked_cast>(\n obj, field, \"ARRAY_INT32\"));\n case DataType::ARRAY_UINT32:\n return self.set(field, checked_cast>(\n obj, field, \"ARRAY_UINT32\"));\n case DataType::ARRAY_INT64:\n return self.set(field, checked_cast>(\n obj, field, \"ARRAY_INT64\"));\n case DataType::ARRAY_UINT64:\n return self.set(field, checked_cast>(\n obj, field, \"ARRAY_UINT64\"));\n case DataType::ARRAY_FLOAT:\n return self.set(field, checked_cast>(\n obj, field, \"ARRAY_FLOAT\"));\n case DataType::ARRAY_DOUBLE:\n return self.set(field, checked_cast>(\n obj, field, \"ARRAY_DOUBLE\"));\n\n // dense vector datatypes\n case DataType::VECTOR_FP16: {\n const auto value = checked_cast(\n obj, field, \"VECTOR_FP16 (list of numbers)\");\n std::vector new_value;\n new_value.reserve(value.size());\n for (const auto &item : value) {\n try {\n new_value.emplace_back(item.cast());\n } catch (const py::cast_error &e) {\n throw py::type_error(\"Vector '\" + field +\n \"': expected VECTOR_FP16, got \" +\n std::string(py::str(py::type::of(obj))));\n }\n }\n return self.set(field, new_value);\n }\n case DataType::VECTOR_FP32:\n return self.set(field, checked_cast>(\n obj, field, \"VECTOR_FP32\"));\n case DataType::VECTOR_FP64:\n return self.set(field, checked_cast>(\n obj, field, \"VECTOR_FP64\"));\n case DataType::VECTOR_INT8:\n return self.set(field, checked_cast>(\n obj, field, \"VECTOR_INT8\"));\n\n // sparse vector datatypes\n case DataType::SPARSE_VECTOR_FP32: {\n const auto sparse_dict =\n checked_cast(obj, field, \"SPARSE_VECTOR_FP32 (dict)\");\n std::vector indices;\n std::vector values;\n for (const auto &item : sparse_dict) {\n try {\n indices.push_back(item.first.cast());\n values.push_back(item.second.cast());\n } catch (const py::cast_error &e) {\n throw py::type_error(\n \"Vector '\" + field +\n \"': sparse vector key/value must be (uint32, float), \"\n \"got key=\" +\n std::string(py::str(py::type::of(item.first))) +\n \", value=\" +\n std::string(py::str(py::type::of(item.second))));\n }\n }\n const std::pair, std::vector>\n sparse_vector{std::move(indices), std::move(values)};\n return self.set(field, sparse_vector);\n }\n case DataType::SPARSE_VECTOR_FP16: {\n const auto sparse_dict =\n checked_cast(obj, field, \"SPARSE_VECTOR_FP16 (dict)\");\n std::vector indices;\n std::vector values;\n for (const auto &item : sparse_dict) {\n try {\n indices.push_back(item.first.cast());\n values.push_back(ailego::Float16(item.second.cast()));\n } catch (const py::cast_error &e) {\n throw py::type_error(\n \"Field '\" + field +\n \"': sparse vector key/value must be (uint32, float), \"\n \"got key=\" +\n std::string(py::str(py::type::of(item.first))) +\n \", value=\" +\n std::string(py::str(py::type::of(item.second))));\n }\n }\n const std::pair, std::vector>\n sparse_vector{std::move(indices), std::move(values)};\n return self.set(field, sparse_vector);\n }\n default:\n throw py::type_error(\"Unsupported type for field: \" + field);\n }\n });\n\n // binding doc get field\n doc.def(\n \"get_any\",\n [](Doc &self, const std::string &field,\n const DataType &type) -> py::object {\n switch (type) {\n // base datatypes\n case DataType::STRING:\n return py::cast(self.get(field));\n case DataType::BOOL:\n return py::cast(self.get(field));\n case DataType::INT32:\n return py::cast(self.get(field));\n case DataType::UINT32:\n return py::cast(self.get(field));\n case DataType::INT64:\n return py::cast(self.get(field));\n case DataType::UINT64:\n return py::cast(self.get(field));\n case DataType::FLOAT:\n return py::cast(self.get(field));\n case DataType::DOUBLE:\n return py::cast(self.get(field));\n\n // array datatypes\n case DataType::ARRAY_STRING:\n return py::cast(self.get>(field));\n case DataType::ARRAY_INT32:\n return py::cast(self.get>(field));\n case DataType::ARRAY_INT64:\n return py::cast(self.get>(field));\n case DataType::ARRAY_UINT32:\n return py::cast(self.get>(field));\n case DataType::ARRAY_UINT64:\n return py::cast(self.get>(field));\n case DataType::ARRAY_FLOAT:\n return py::cast(self.get>(field));\n case DataType::ARRAY_DOUBLE:\n return py::cast(self.get>(field));\n case DataType::ARRAY_BOOL:\n return py::cast(self.get>(field));\n\n // vector datatypes\n case DataType::VECTOR_INT8:\n return py::cast(self.get>(field));\n case DataType::VECTOR_FP16: {\n auto value = self.get>(field);\n if (value.has_value()) {\n std::vector new_value;\n new_value.reserve(value.value().size());\n for (auto &item : value.value()) {\n new_value.push_back(static_cast(item));\n }\n return py::cast(new_value);\n }\n return py::none();\n }\n case DataType::VECTOR_FP32:\n return py::cast(self.get>(field));\n case DataType::VECTOR_FP64:\n return py::cast(self.get>(field));\n case DataType::SPARSE_VECTOR_FP16: {\n auto vector = self.get<\n std::pair, std::vector>>(\n field);\n const auto &indices = vector->first;\n const auto &values = vector->second;\n py::dict d;\n for (size_t i = 0; i < indices.size(); ++i) {\n d[py::int_(indices[i])] =\n py::float_(static_cast(values[i]));\n }\n return std::move(d);\n }\n case DataType::SPARSE_VECTOR_FP32: {\n auto vector =\n self.get, std::vector>>(\n field);\n const auto &indices = vector->first;\n const auto &values = vector->second;\n py::dict d;\n for (size_t i = 0; i < indices.size(); ++i) {\n d[py::int_(indices[i])] = py::float_(values[i]);\n }\n return std::move(d);\n }\n default:\n throw py::type_error(\"Unsupported type for field: \" + field);\n }\n });\n doc.def(\n \"get_all\",\n [](Doc &self, const CollectionSchema &schema) -> py::tuple {\n py::tuple result(4);\n // 1. set doc id and score\n result[0] = py::str(self.pk());\n result[1] = py::float_(self.score());\n\n if (self.is_empty()) {\n result[2] = py::none();\n result[3] = py::none();\n return result;\n }\n // 2. set scalar fields\n py::dict fields;\n for (const auto &field_meta : schema.forward_fields()) {\n const std::string &field = field_meta->name();\n if (!self.has_value(field)) {\n continue;\n }\n\n try {\n auto val = [&]() -> py::object {\n switch (field_meta->data_type()) {\n // base datatypes\n case DataType::STRING:\n return py::str(self.get(field).value());\n case DataType::BOOL:\n return py::cast(self.get(field));\n case DataType::INT32:\n return py::cast(self.get(field));\n case DataType::UINT32:\n return py::cast(self.get(field));\n case DataType::INT64:\n return py::cast(self.get(field));\n case DataType::UINT64:\n return py::cast(self.get(field));\n case DataType::FLOAT:\n return py::cast(self.get(field));\n case DataType::DOUBLE:\n return py::cast(self.get(field));\n\n // array datatypes\n case DataType::ARRAY_STRING:\n return py::cast(self.get>(field));\n case DataType::ARRAY_INT32:\n return py::cast(self.get>(field));\n case DataType::ARRAY_INT64:\n return py::cast(self.get>(field));\n case DataType::ARRAY_UINT32:\n return py::cast(self.get>(field));\n case DataType::ARRAY_UINT64:\n return py::cast(self.get>(field));\n case DataType::ARRAY_FLOAT:\n return py::cast(self.get>(field));\n case DataType::ARRAY_DOUBLE:\n return py::cast(self.get>(field));\n case DataType::ARRAY_BOOL:\n return py::cast(self.get>(field));\n default:\n throw py::type_error(\"Unsupported type for field: \" + field);\n }\n }();\n fields[py::str(field)] = val;\n } catch (const std::exception &e) {\n fields[py::str(field)] = py::none();\n }\n }\n if (!fields.empty()) {\n result[2] = fields;\n } else {\n result[2] = py::none();\n }\n // 3. set vector fields\n py::dict vectors;\n for (const auto &vec_meta : schema.vector_fields()) {\n const std::string &vec = vec_meta->name();\n if (!self.has_value(vec)) continue;\n\n try {\n auto array = [&]() -> py::object {\n switch (vec_meta->data_type()) {\n case DataType::VECTOR_INT8:\n return py::cast(self.get>(vec));\n case DataType::VECTOR_FP16: {\n auto value = self.get>(vec);\n if (value.has_value()) {\n std::vector new_value;\n new_value.reserve(value.value().size());\n for (auto &item : value.value()) {\n new_value.push_back(static_cast(item));\n }\n return py::cast(new_value);\n }\n return py::none();\n }\n case DataType::VECTOR_FP32:\n return py::cast(self.get>(vec));\n case DataType::VECTOR_FP64:\n return py::cast(self.get>(vec));\n case DataType::SPARSE_VECTOR_FP16: {\n auto vector =\n self.get,\n std::vector>>(vec);\n const auto &indices = vector->first;\n const auto &values = vector->second;\n py::dict d;\n for (size_t i = 0; i < indices.size(); ++i) {\n d[py::int_(indices[i])] =\n py::float_(static_cast(values[i]));\n }\n return std::move(d);\n }\n case DataType::SPARSE_VECTOR_FP32: {\n auto vector = self.get<\n std::pair, std::vector>>(\n vec);\n const auto &indices = vector->first;\n const auto &values = vector->second;\n py::dict d;\n for (size_t i = 0; i < indices.size(); ++i) {\n d[py::int_(indices[i])] = py::float_(values[i]);\n }\n return std::move(d);\n }\n default:\n throw py::type_error(\"Unsupported type for field: \" + vec);\n }\n }();\n vectors[py::str(vec)] = array;\n } catch (const std::exception &e) {\n vectors[py::str(vec)] = py::none();\n }\n }\n if (!vectors.empty()) {\n result[3] = vectors;\n } else {\n result[3] = py::none();\n }\n return result;\n },\n py::arg(\"schema\"),\n \"Get all fields and vectors as a tuple: (id, score, fields, vectors). \"\n \"Vectors are zero-copy numpy arrays (dense: ndarray, sparse: (indices, \"\n \"values) tuple).\");\n}\n} // namespace zvec" }, { "path": "src/binding/python/model/schema/python_schema.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"python_schema.h\"\n#include \n#include \n#include \n\nnamespace zvec {\n\nvoid ZVecPySchemas::Initialize(pybind11::module_ &parent) {\n auto m =\n parent.def_submodule(\"schema\", \"This module contains the schema of Zvec\");\n\n bind_field_schema(m);\n bind_collection_schema(m);\n bind_collection_stats(m);\n}\n\nvoid ZVecPySchemas::bind_field_schema(py::module_ &m) {\n py::class_(m, \"_FieldSchema\")\n .def(py::init(),\n py::arg(\"name\"), py::arg(\"data_type\"), py::arg(\"dimension\") = 0,\n py::arg(\"nullable\") = false, py::arg(\"index_param\") = nullptr)\n .def_property_readonly(\"name\", &FieldSchema::name)\n .def_property_readonly(\"data_type\", &FieldSchema::data_type)\n .def_property_readonly(\"nullable\", &FieldSchema::nullable)\n .def_property_readonly(\"dimension\", &FieldSchema::dimension)\n .def_property_readonly(\"is_dense_vector\", &FieldSchema::is_dense_vector)\n .def_property_readonly(\"is_sparse_vector\", &FieldSchema::is_sparse_vector)\n .def_property_readonly(\"index_type\",\n [](const FieldSchema &self) {\n return self.index_params()\n ? self.index_type()\n : IndexType::UNDEFINED;\n })\n .def_property_readonly(\"index_param\",\n [](const FieldSchema &self) -> py::object {\n if (self.index_params()) {\n return py::cast(self.index_params());\n }\n return py::none();\n })\n .def(\"__eq__\", &FieldSchema::operator==)\n .def(\"__ne__\", &FieldSchema::operator!=)\n .def(py::pickle(\n [](const FieldSchema &self) {\n return py::make_tuple(self.name(), self.data_type(),\n self.dimension(), self.nullable(),\n self.index_params()\n ? py::cast(self.index_params())\n : py::none());\n },\n [](py::tuple t) {\n if (t.size() != 5) {\n throw std::runtime_error(\n \"Invalid tuple size for FieldSchema pickle\");\n }\n std::string name = t[0].cast();\n DataType dtype = t[1].cast();\n uint32_t dim = t[2].cast();\n bool nullable = t[3].cast();\n\n IndexParams::Ptr idx_params = nullptr;\n if (!t[4].is_none()) {\n idx_params = t[4].cast();\n }\n\n return std::make_shared(name, dtype, dim, nullable,\n idx_params);\n }));\n}\n\nvoid ZVecPySchemas::bind_collection_schema(py::module_ &m) {\n py::class_(m, \"_CollectionSchema\")\n .def(py::init(),\n py::arg(\"name\"), py::arg(\"fields\"),\n \"Construct with name and list of fields\")\n .def_property_readonly(\"name\", &CollectionSchema::name)\n .def(\"has_field\", &CollectionSchema::has_field, py::arg(\"field_name\"),\n \"Check if a field exists.\")\n .def(\n \"get_field\",\n [](const CollectionSchema &self, const std::string &name)\n -> const FieldSchema * { return self.get_field(name); },\n py::arg(\"field_name\"), py::return_value_policy::reference_internal,\n \"Get field by name (const pointer), returns None if not found.\")\n .def(\n \"get_forward_field\",\n [](const CollectionSchema &self, const std::string &name)\n -> const FieldSchema * { return self.get_forward_field(name); },\n py::arg(\"field_name\"), py::return_value_policy::reference_internal,\n \"Get forward field (used for filtering).\")\n .def(\n \"get_vector_field\",\n [](const CollectionSchema &self, const std::string &name)\n -> const FieldSchema * { return self.get_vector_field(name); },\n py::arg(\"field_name\"), py::return_value_policy::reference_internal,\n \"Get vector field by name.\")\n .def(\"fields\", &CollectionSchema::fields,\n \"Return list of all field schemas.\", py::return_value_policy::copy)\n .def(\"forward_fields\", &CollectionSchema::forward_fields,\n \"Return list of forward-indexed fields.\",\n py::return_value_policy::copy)\n .def(\"vector_fields\", &CollectionSchema::vector_fields,\n \"Return list of vector fields.\", py::return_value_policy::copy)\n .def(\"__eq__\", &CollectionSchema::operator==)\n .def(\"__ne__\", &CollectionSchema::operator!=)\n .def(py::pickle(\n [](const CollectionSchema &cs) {\n return py::make_tuple(cs.name(), cs.fields(),\n cs.max_doc_count_per_segment());\n },\n [](py::tuple t) {\n if (t.size() != 3)\n throw std::runtime_error(\"Invalid state for CollectionSchema!\");\n\n auto name = t[0].cast();\n auto fields = t[1].cast();\n auto max_docs = t[2].cast();\n\n auto cs = std::make_shared(name, fields);\n cs->set_max_doc_count_per_segment(max_docs);\n return cs;\n }));\n}\n\nvoid ZVecPySchemas::bind_collection_stats(py::module_ &m) {\n pybind11::class_(m, \"CollectionStats\")\n .def(pybind11::init<>())\n .def_property_readonly(\n \"doc_count\", [](const CollectionStats &c) { return c.doc_count; })\n .def_property_readonly(\n \"index_completeness\",\n [](const CollectionStats &c) { return c.index_completeness; })\n .def(\"__repr__\", [](const CollectionStats &c) {\n std::string map_str = \"{\";\n bool first = true;\n for (const auto &[k, v] : c.index_completeness) {\n if (!first) map_str += \", \";\n map_str += \"\\\"\" + k + \"\\\":\" + std::to_string(v);\n first = false;\n }\n map_str += \"}\";\n return \"{\\\"doc_count\\\":\" + std::to_string(c.doc_count) +\n \", \\\"index_completeness\\\":\" + map_str + \"}\";\n });\n}\n\n} // namespace zvec" }, { "path": "src/binding/python/typing/python_type.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"python_type.h\"\n\nnamespace zvec {\n\nvoid ZVecPyTyping::Initialize(pybind11::module_ &parent) {\n auto m = parent.def_submodule(\n \"typing\", \"This module contains the basic data types of Zvec\");\n // binding base types\n bind_datatypes(m);\n bind_index_types(m);\n bind_metric_types(m);\n bind_quantize_types(m);\n bind_status(m);\n}\n\nvoid ZVecPyTyping::bind_datatypes(pybind11::module_ &m) {\n py::enum_(m, \"DataType\", R\"pbdoc(\nEnumeration of supported data types in Zvec.\n\nIncludes scalar types, dense/sparse vector types, and array types.\n\nExamples:\n >>> from zvec.typing import DataType\n >>> print(DataType.FLOAT)\n DataType.FLOAT\n >>> print(DataType.VECTOR_FP32)\n DataType.VECTOR_FP32\n)pbdoc\")\n // field type\n .value(\"STRING\", DataType::STRING)\n .value(\"BOOL\", DataType::BOOL)\n .value(\"INT32\", DataType::INT32)\n .value(\"INT64\", DataType::INT64)\n .value(\"FLOAT\", DataType::FLOAT)\n .value(\"DOUBLE\", DataType::DOUBLE)\n .value(\"UINT32\", DataType::UINT32)\n .value(\"UINT64\", DataType::UINT64)\n\n\n // dense vector type\n .value(\"VECTOR_FP16\", DataType::VECTOR_FP16)\n .value(\"VECTOR_FP32\", DataType::VECTOR_FP32)\n .value(\"VECTOR_FP64\", DataType::VECTOR_FP64)\n .value(\"VECTOR_INT8\", DataType::VECTOR_INT8)\n\n\n // sparse vector type\n .value(\"SPARSE_VECTOR_FP32\", DataType::SPARSE_VECTOR_FP32)\n .value(\"SPARSE_VECTOR_FP16\", DataType::SPARSE_VECTOR_FP16)\n\n\n // array type [not support bool/bytes]\n .value(\"ARRAY_STRING\", DataType::ARRAY_STRING)\n .value(\"ARRAY_INT32\", DataType::ARRAY_INT32)\n .value(\"ARRAY_INT64\", DataType::ARRAY_INT64)\n .value(\"ARRAY_FLOAT\", DataType::ARRAY_FLOAT)\n .value(\"ARRAY_DOUBLE\", DataType::ARRAY_DOUBLE)\n .value(\"ARRAY_BOOL\", DataType::ARRAY_BOOL)\n .value(\"ARRAY_UINT32\", DataType::ARRAY_UINT32)\n .value(\"ARRAY_UINT64\", DataType::ARRAY_UINT64)\n\n\n // non support\n // .value(\"BINARY\", DataType::BINARY)\n // .value(\"ARRAY_BINARY\", DataType::ARRAY_BINARY)\n // .value(\"VECTOR_INT4\", DataType::VECTOR_INT4)\n // .value(\"VECTOR_INT16\", DataType::VECTOR_INT16)\n // .value(\"VECTOR_BINARY32\", DataType::VECTOR_BINARY32)\n // .value(\"VECTOR_BINARY64\", DataType::VECTOR_BINARY64)\n // .value(\"UNDEFINED\", DataType::UNDEFINED)\n ;\n}\n\nvoid ZVecPyTyping::bind_index_types(pybind11::module_ &m) {\n py::enum_(m, \"IndexType\", R\"pbdoc(\nEnumeration of supported index types in Zvec.\n\nExamples:\n >>> from zvec.typing import IndexType\n >>> print(IndexType.HNSW)\n IndexType.HNSW\n)pbdoc\")\n .value(\"UNDEFINED\", IndexType::UNDEFINED)\n .value(\"HNSW\", IndexType::HNSW)\n .value(\"HNSW_RABITQ\", IndexType::HNSW_RABITQ)\n .value(\"IVF\", IndexType::IVF)\n .value(\"FLAT\", IndexType::FLAT)\n .value(\"INVERT\", IndexType::INVERT);\n}\n\nvoid ZVecPyTyping::bind_metric_types(pybind11::module_ &m) {\n py::enum_(m, \"MetricType\", R\"pbdoc(\nEnumeration of supported distance/similarity metrics.\n\n- COSINE: Cosine similarity.\n- IP: Inner product (dot product).\n- L2: Euclidean distance (L2 norm).\n\nExamples:\n >>> from zvec.typing import MetricType\n >>> print(MetricType.COSINE)\n MetricType.COSINE\n)pbdoc\")\n .value(\"COSINE\", MetricType::COSINE)\n .value(\"IP\", MetricType::IP)\n .value(\"L2\", MetricType::L2);\n}\n\nvoid ZVecPyTyping::bind_quantize_types(py::module_ &m) {\n py::enum_(m, \"QuantizeType\", R\"pbdoc(\nEnumeration of supported quantization types for vector compression.\n\nExamples:\n >>> from zvec.typing import QuantizeType\n >>> print(QuantizeType.INT8)\n QuantizeType.INT8\n)pbdoc\")\n .value(\"UNDEFINED\", QuantizeType::UNDEFINED)\n .value(\"FP16\", QuantizeType::FP16)\n .value(\"INT8\", QuantizeType::INT8)\n .value(\"INT4\", QuantizeType::INT4)\n .value(\"RABITQ\", QuantizeType::RABITQ);\n}\n\nvoid ZVecPyTyping::bind_status(py::module_ &m) {\n // bind status code\n py::enum_(m, \"StatusCode\", R\"pbdoc(\nEnumeration of possible status codes for Zvec operations.\n\nUsed by the `Status` class to indicate success or failure reason.\n)pbdoc\")\n .value(\"OK\", StatusCode::OK)\n .value(\"NOT_FOUND\", StatusCode::NOT_FOUND)\n .value(\"ALREADY_EXISTS\", StatusCode::ALREADY_EXISTS)\n .value(\"INVALID_ARGUMENT\", StatusCode::INVALID_ARGUMENT)\n .value(\"PERMISSION_DENIED\", StatusCode::PERMISSION_DENIED)\n .value(\"FAILED_PRECONDITION\", StatusCode::FAILED_PRECONDITION)\n .value(\"RESOURCE_EXHAUSTED\", StatusCode::RESOURCE_EXHAUSTED)\n .value(\"UNAVAILABLE\", StatusCode::UNAVAILABLE)\n .value(\"INTERNAL_ERROR\", StatusCode::INTERNAL_ERROR)\n .value(\"NOT_SUPPORTED\", StatusCode::NOT_SUPPORTED)\n .value(\"UNKNOWN\", StatusCode::UNKNOWN);\n\n // bind status\n py::class_(m, \"Status\", R\"pbdoc(\nRepresents the outcome of a Zvec operation.\n\nA `Status` object is either OK (success) or carries an error code and message.\n\nExamples:\n >>> from zvec.typing import Status, StatusCode\n >>> s = Status()\n >>> print(s.ok())\n True\n >>> s = Status(StatusCode.INVALID_ARGUMENT, \"Field not found\")\n >>> print(s.code() == StatusCode.INVALID_ARGUMENT)\n True\n >>> print(s.message())\n Field not found\n)pbdoc\")\n .def(py::init<>())\n .def(py::init(), py::arg(\"code\"),\n py::arg(\"message\") = \"\", R\"pbdoc(\nConstruct a status with the given code and optional message.\n\nArgs:\n code (StatusCode): The status code.\n message (str, optional): Error message. Defaults to empty string.\n)pbdoc\")\n .def(\"ok\", &Status::ok, \"bool: Returns True if the status is OK.\")\n .def(\"code\", &Status::code, \"StatusCode: Returns the status code.\")\n .def(\"message\", &Status::message,\n \"str: Returns the error message (may be empty).\")\n .def_static(\"OK\", &Status::OK, \"Create an OK status.\")\n .def_static(\n \"InvalidArgument\",\n [](const std::string &msg) { return Status::InvalidArgument(msg); },\n py::arg(\"message\"))\n .def_static(\n \"NotFound\",\n [](const std::string &msg) { return Status::NotFound(msg); },\n py::arg(\"message\"))\n .def_static(\n \"AlreadyExists\",\n [](const std::string &msg) { return Status::AlreadyExists(msg); },\n py::arg(\"message\"))\n .def_static(\n \"InternalError\",\n [](const std::string &msg) { return Status::InternalError(msg); },\n py::arg(\"message\"))\n .def_static(\n \"PermissionDenied\",\n [](const std::string &msg) { return Status::PermissionDenied(msg); },\n py::arg(\"message\"))\n .def(\"__eq__\", [](const Status &self,\n const Status &other) { return self == other; })\n .def(\"__ne__\", [](const Status &self,\n const Status &other) { return self != other; })\n .def(\"__repr__\", [](const Status &self) {\n std::string result =\n \"{\"\n \"\\\"code\\\":\" +\n std::to_string(static_cast(self.code()));\n\n if (!self.message().empty()) {\n result += \", \\\"message\\\":\\\"\" + self.message() + \"\\\"\";\n }\n\n result += \"}\";\n return result;\n });\n}\n\n} // namespace zvec" }, { "path": "src/core/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\nif(RABITQ_SUPPORTED AND AUTO_DETECT_ARCH)\n set(HNSW_RABITQ_FILES\n hnsw_rabitq_query_algorithm.cc\n hnsw_rabitq_streamer.cc\n hnsw_rabitq_searcher.cc\n hnsw_rabitq_entity.cc\n rabitq_reformer.cc\n rabitq_converter.cc\n )\n set(HNSW_RABITQ_FILES_FULL ${HNSW_RABITQ_FILES})\n list(TRANSFORM HNSW_RABITQ_FILES_FULL PREPEND \"algorithm/hnsw_rabitq/\")\n foreach(FILE ${HNSW_RABITQ_FILES_FULL})\n set_source_files_properties(\n ${FILE}\n PROPERTIES\n COMPILE_FLAGS \"${RABITQ_ARCH_FLAG}\"\n )\n endforeach()\nendif()\n\ncc_directory(framework)\ncc_directory(algorithm)\ncc_directory(metric)\ncc_directory(quantizer)\ncc_directory(utility)\ncc_directory(interface)\ncc_directory(mixed_reducer)\n\ngit_version(GIT_SRCS_VER ${CMAKE_CURRENT_SOURCE_DIR})\nfile(GLOB_RECURSE ALL_CORE_SRCS *.cc *.c *.h)\n\n# Remove algorithm/hnsw_rabitq implementation files if not supported.\n# interface/indexes/hnsw_rabitq_index.cc is kept because it provides the vtable\n# for HNSWRabitqIndex and guards rabitqlib usage with #if RABITQ_SUPPORTED.\nif(NOT RABITQ_SUPPORTED)\n list(FILTER ALL_CORE_SRCS EXCLUDE REGEX \".*/algorithm/hnsw_rabitq/.*\")\nendif()\n\ncc_library(\n NAME zvec_core STATIC STRICT PACKED\n SRCS ${ALL_CORE_SRCS}\n LIBS zvec_ailego zvec_turbo sparsehash magic_enum rabitqlib\n INCS . ${PROJECT_ROOT_DIR}/src/core\n VERSION \"${GIT_SRCS_VER}\"\n)\n" }, { "path": "src/core/algorithm/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\ncc_directory(cluster)\ncc_directory(flat)\ncc_directory(flat_sparse)\ncc_directory(ivf)\ncc_directory(hnsw)\ncc_directory(hnsw_sparse)\nif(RABITQ_SUPPORTED)\n message(STATUS \"BUILD RABITQ\")\n cc_directory(hnsw_rabitq)\nelse()\n message(STATUS \"NOT BUILD RABITQ\")\n # Empty stub library for unsupported platforms\n file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/rabitq_stub.cc\n \"// Stub implementation for unsupported platforms\\n\"\n \"// RaBitQ only supports Linux x86_64\\n\"\n \"namespace zvec { namespace core { /* empty namespace for compatibility */ } }\\n\"\n )\n\n cc_library(\n NAME core_knn_hnsw_rabitq\n STATIC SHARED STRICT ALWAYS_LINK\n SRCS ${CMAKE_CURRENT_BINARY_DIR}/rabitq_stub.cc\n LIBS core_framework\n INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )\nendif()\n" }, { "path": "src/core/algorithm/cluster/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\ncc_library(\n NAME core_knn_cluster STATIC SHARED STRICT ALWAYS_LINK\n SRCS *.cc\n LIBS zvec_ailego core_framework \n INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/cluster\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )\n" }, { "path": "src/core/algorithm/cluster/cluster_params.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\n//! General\nstatic const std::string GENERAL_CLUSTER_COUNT =\n \"proxima.general.cluster.count\";\nstatic const std::string GENERAL_THREAD_COUNT =\n \"proxima.general.cluster.thread_count\";\n\n//! Optimize K-means\nstatic const std::string OPTKMEANS_CLUSTER_COUNT =\n \"proxima.optkmeans.cluster.count\";\nstatic const std::string OPTKMEANS_CLUSTER_MAX_ITERATIONS =\n \"proxima.optkmeans.cluster.max_iterations\";\nstatic const std::string OPTKMEANS_CLUSTER_EPSILON =\n \"proxima.optkmeans.cluster.epsilon\";\nstatic const std::string OPTKMEANS_CLUSTER_SHARD_FACTOR =\n \"proxima.optkmeans.cluster.shard_factor\";\nstatic const std::string OPTKMEANS_CLUSTER_PURGE_EMPTY =\n \"proxima.optkmeans.cluster.purge_empty\";\nstatic const std::string OPTKMEANS_CLUSTER_MARKOV_CHAIN_LENGTH =\n \"proxima.optkmeans.cluster.markov_chain_length\";\nstatic const std::string OPTKMEANS_CLUSTER_ASSUMPTION_FREE =\n \"proxima.optkmeans.cluster.assumption_free\";\n\n//! K-means\nstatic const std::string KMEANS_CLUSTER_COUNT = \"proxima.kmeans.cluster.count\";\nstatic const std::string KMEANS_CLUSTER_SHARD_FACTOR =\n \"proxima.kmeans.cluster.shard_factor\";\nstatic const std::string KMEANS_CLUSTER_EPSILON =\n \"proxima.kmeans.cluster.epsilon\";\nstatic const std::string KMEANS_CLUSTER_MAX_ITERATIONS =\n \"proxima.kmeans.cluster.max_iterations\";\nstatic const std::string KMEANS_CLUSTER_PURGE_EMPTY =\n \"proxima.kmeans.cluster.purge_empty\";\nstatic const std::string KMEANS_CLUSTER_BATCH = \"proxima.kmeans.cluster.batch\";\nstatic const std::string KMEANS_CLUSTER_SEEKER_CLASS =\n \"proxima.kmeans.cluster.seeker_class\";\nstatic const std::string KMEANS_CLUSTER_SEEKER_PARAMS =\n \"proxima.kmeans.cluster.seeker_params\";\n\n//! Mini Batch K-means\nstatic const std::string MINIBATCHKMEANS_CLUSTER_COUNT =\n \"proxima.minibatchkmeans.cluster.count\";\nstatic const std::string MINIBATCHKMEANS_CLUSTER_SHARD_FACTOR =\n \"proxima.minibatchkmeans.cluster.shard_factor\";\nstatic const std::string MINIBATCHKMEANS_CLUSTER_EPSILON =\n \"proxima.minibatchkmeans.cluster.epsilon\";\nstatic const std::string MINIBATCHKMEANS_CLUSTER_MAX_ITERATIONS =\n \"proxima.minibatchkmeans.cluster.max_iterations\";\nstatic const std::string MINIBATCHKMEANS_CLUSTER_PURGE_EMPTY =\n \"proxima.minibatchkmeans.cluster.purge_empty\";\nstatic const std::string MINIBATCHKMEANS_CLUSTER_TRY_COUNT =\n \"proxima.minibatchkmeans.cluster.try_count\";\nstatic const std::string MINIBATCHKMEANS_CLUSTER_BATCH_COUNT =\n \"proxima.minibatchkmeans.cluster.batch_count\";\nstatic const std::string MINIBATCHKMEANS_CLUSTER_SEEKER_CLASS =\n \"proxima.minibatchkmeans.cluster.seeker_class\";\nstatic const std::string MINIBATCHKMEANS_CLUSTER_SEEKER_PARAMS =\n \"proxima.minibatchkmeans.cluster.seeker_params\";\n\n//! K-means++\nstatic const std::string KMEANSPP_CLUSTER_COUNT =\n \"proxima.kmeanspp.cluster.count\";\nstatic const std::string KMEANSPP_CLUSTER_SHARD_FACTOR =\n \"proxima.kmeanspp.cluster.shard_factor\";\nstatic const std::string KMEANSPP_CLUSTER_CLASS =\n \"proxima.kmeanspp.cluster.class\";\nstatic const std::string KMEANSPP_CLUSTER_PARAMS =\n \"proxima.kmeanspp.cluster.params\";\n\n//! K-MC2\nstatic const std::string KMC2_CLUSTER_COUNT = \"proxima.kmc2.cluster.count\";\nstatic const std::string KMC2_CLUSTER_SHARD_FACTOR =\n \"proxima.kmc2.cluster.shard_factor\";\nstatic const std::string KMC2_CLUSTER_MARKOV_CHAIN_LENGTH =\n \"proxima.kmc2.cluster.markov_chain_length\";\nstatic const std::string KMC2_CLUSTER_ASSUMPTION_FREE =\n \"proxima.kmc2.cluster.assumption_free\";\nstatic const std::string KMC2_CLUSTER_CLASS = \"proxima.kmc2.cluster.class\";\nstatic const std::string KMC2_CLUSTER_PARAMS = \"proxima.kmc2.cluster.params\";\n\n//! Bisecting K-means\nstatic const std::string BIKMEANS_CLUSTER_COUNT =\n \"proxima.bikmeans.cluster.count\";\nstatic const std::string BIKMEANS_CLUSTER_INIT_COUNT =\n \"proxima.bikmeans.cluster.init_count\";\nstatic const std::string BIKMEANS_CLUSTER_PURGE_EMPTY =\n \"proxima.bikmeans.cluster.purge_empty\";\nstatic const std::string BIKMEANS_CLUSTER_FIRST_CLASS =\n \"proxima.bikmeans.cluster.first_class\";\nstatic const std::string BIKMEANS_CLUSTER_SECOND_CLASS =\n \"proxima.bikmeans.cluster.second_class\";\nstatic const std::string BIKMEANS_CLUSTER_FIRST_PARAMS =\n \"proxima.bikmeans.cluster.first_params\";\nstatic const std::string BIKMEANS_CLUSTER_SECOND_PARAMS =\n \"proxima.bikmeans.cluster.second_params\";\n\n//! K-medoids\nstatic const std::string KMEDOIDS_CLUSTER_COUNT =\n \"proxima.kmedoids.cluster.count\";\nstatic const std::string KMEDOIDS_CLUSTER_SHARD_FACTOR =\n \"proxima.kmedoids.cluster.shard_factor\";\nstatic const std::string KMEDOIDS_CLUSTER_EPSILON =\n \"proxima.kmedoids.cluster.epsilon\";\nstatic const std::string KMEDOIDS_CLUSTER_MAX_ITERATIONS =\n \"proxima.kmedoids.cluster.max_iterations\";\nstatic const std::string KMEDOIDS_CLUSTER_PURGE_EMPTY =\n \"proxima.kmedoids.cluster.purge_empty\";\nstatic const std::string KMEDOIDS_CLUSTER_BENCH_RATIO =\n \"proxima.kmedoids.cluster.bench_ratio\";\nstatic const std::string KMEDOIDS_CLUSTER_ONLY_MEANS =\n \"proxima.kmedoids.cluster.only_means\";\nstatic const std::string KMEDOIDS_CLUSTER_WITHOUT_MEANS =\n \"proxima.kmedoids.cluster.without_means\";\nstatic const std::string KMEDOIDS_CLUSTER_SEEKER_CLASS =\n \"proxima.kmedoids.cluster.seeker_class\";\nstatic const std::string KMEDOIDS_CLUSTER_SEEKER_PARAMS =\n \"proxima.kmedoids.cluster.seeker_params\";\n\n//! Stratified\nstatic const std::string STRATIFIED_CLUSTER_COUNT =\n \"proxima.stratified.cluster.count\";\nstatic const std::string STRATIFIED_CLUSTER_FIRST_CLASS =\n \"proxima.stratified.cluster.first_class\";\nstatic const std::string STRATIFIED_CLUSTER_SECOND_CLASS =\n \"proxima.stratified.cluster.second_class\";\nstatic const std::string STRATIFIED_CLUSTER_FIRST_COUNT =\n \"proxima.stratified.cluster.first_count\";\nstatic const std::string STRATIFIED_CLUSTER_SECOND_COUNT =\n \"proxima.stratified.cluster.second_count\";\nstatic const std::string STRATIFIED_CLUSTER_FIRST_PARAMS =\n \"proxima.stratified.cluster.first_params\";\nstatic const std::string STRATIFIED_CLUSTER_SECOND_PARAMS =\n \"proxima.stratified.cluster.second_params\";\nstatic const std::string STRATIFIED_CLUSTER_AUTO_TUNING =\n \"proxima.stratified.cluster.auto_tuning\";\nstatic const std::string STRATIFIED_CLUSTER_SECOND_POOL_COUNT =\n \"proxima.stratified.cluster.second_pool_count\";\n\n//! Gap Statistics\nstatic const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MIN =\n \"proxima.gapstats.cluster_estimater.k_min\";\nstatic const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MAX =\n \"proxima.gapstats.cluster_estimater.k_max\";\nstatic const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MIN_STEP =\n \"proxima.gapstats.cluster_estimater.k_min_step\";\nstatic const std::string GAPSTATS_CLUSTER_ESTIMATER_K_MAX_STEP =\n \"proxima.gapstats.cluster_estimater.k_max_step\";\nstatic const std::string GAPSTATS_CLUSTER_ESTIMATER_TRY_COUNT =\n \"proxima.gapstats.cluster_estimater.try_count\";\nstatic const std::string GAPSTATS_CLUSTER_ESTIMATER_SHARD_FACTOR =\n \"proxima.gapstats.cluster_estimater.shard_factor\";\nstatic const std::string GAPSTATS_CLUSTER_ESTIMATER_ENABLE_MC2 =\n \"proxima.gapstats.cluster_estimater.enable_mc2\";\nstatic const std::string GAPSTATS_CLUSTER_ESTIMATER_MARKOV_CHAIN_LENGTH =\n \"proxima.gapstats.cluster_estimater.markov_chain_length\";\nstatic const std::string GAPSTATS_CLUSTER_ESTIMATER_CLUSTER_CLASS =\n \"proxima.gapstats.cluster_estimater.cluster_class\";\n\nstatic const std::string CLUSTER_TRAINER_SAMPLE_COUNT =\n \"proxima.cluster.trainer.sample_count\";\nstatic const std::string CLUSTER_TRAINER_SAMPLE_RATIO =\n \"proxima.cluster.trainer.sample_ratio\";\nstatic const std::string CLUSTER_TRAINER_THREAD_COUNT =\n \"proxima.cluster.trainer.thread_count\";\nstatic const std::string CLUSTER_TRAINER_FILE_NAME =\n \"proxima.cluster.trainer.file_name\";\nstatic const std::string CLUSTER_TRAINER_CLASS_NAME =\n \"proxima.cluster.trainer.class_name\";\n\nstatic const std::string STRATIFIED_TRAINER_SAMPLE_COUNT =\n \"proxima.stratified.trainer.sample_count\";\nstatic const std::string STRATIFIED_TRAINER_SAMPLE_RATIO =\n \"proxima.stratified.trainer.sample_ratio\";\nstatic const std::string STRATIFIED_TRAINER_THREAD_COUNT =\n \"proxima.stratified.trainer.thread_count\";\nstatic const std::string STRATIFIED_TRAINER_FILE_NAME =\n \"proxima.stratified.trainer.file_name\";\nstatic const std::string STRATIFIED_TRAINER_CLASS_NAME =\n \"proxima.stratified.trainer.class_name\";\nstatic const std::string STRATIFIED_TRAINER_CLUSTER_COUNT =\n \"proxima.stratified.trainer.cluster_count\";\nstatic const std::string STRATIFIED_TRAINER_AUTOAUNE =\n \"proxima.stratified.trainer.autotune\";\nstatic const std::string STRATIFIED_TRAINER_PARAMS_IN_LEVEL_PREFIX =\n \"proxima.stratified.trainer.cluster_params_in_level_\";\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/cluster/kmeans_cluster.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"cluster_params.h\"\n#include \"linear_seeker.h\"\n#include \"vector_mean.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Kmeans Cluster\n */\nclass KmeansCluster : public IndexCluster {\n public:\n //! Constructor\n KmeansCluster(void) {}\n\n //! Constructor\n KmeansCluster(size_t iters, bool batch)\n : max_iterations_(iters), batch_(batch) {}\n\n //! Constructor\n KmeansCluster(bool batch) : batch_(batch) {}\n\n //! Destructor\n virtual ~KmeansCluster(void) {}\n\n //! Initialize Cluster\n virtual int init(const IndexMeta &meta, const ailego::Params ¶ms);\n\n //! Cleanup Cluster\n virtual int cleanup(void);\n\n //! Reset Cluster\n virtual int reset(void);\n\n //! Update Cluster\n virtual int update(const ailego::Params ¶ms);\n\n //! Suggest dividing to K clusters\n virtual void suggest(uint32_t k);\n\n //! Mount features\n virtual int mount(IndexFeatures::Pointer feats);\n\n //! Cluster\n virtual int cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n //! Classify\n virtual int classify(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n //! Label\n virtual int label(IndexThreads::Pointer threads,\n const IndexCluster::CentroidList ¢s,\n std::vector *out);\n\n protected:\n //! Test if it is valid\n bool is_valid(void) const;\n\n //! Cluster once\n int clustering(IndexThreads *threads, IndexCluster::CentroidList ¢s,\n double *cost);\n\n //! Update parameters\n void update_params(const ailego::Params ¶ms);\n\n //! Init seeker\n int init_seeker(void);\n\n //! Build seeker\n int build_seeker(const IndexCluster::CentroidList ¢s);\n\n //! Check Centroids\n bool check_centroids(const IndexCluster::CentroidList ¢s);\n\n //! Initialize Centroids\n void init_centroids(size_t count, IndexCluster::CentroidList *out);\n\n //! Initialize Shard Containers\n void init_containers(size_t shard_count);\n\n //! Initialize Shard Features Containers\n void init_features_containers(size_t shard_count);\n\n //! Split Clusters\n void split_clusters(IndexThreads *threads,\n const IndexCluster::CentroidList ¢s);\n\n //! Update Centroids\n void update_centroids(IndexThreads *threads,\n IndexCluster::CentroidList ¢s);\n\n //! Update Clusters\n void update_clusters(IndexThreads *threads,\n const IndexCluster::CentroidList ¢s);\n\n //! Update Clusters' Features\n void update_features(IndexThreads *threads,\n IndexCluster::CentroidList ¢s);\n\n //! Update Labels\n void update_labels(IndexThreads *threads, std::vector *labels);\n\n //! Split Clusters in Thread\n void split_clusters_thread(size_t index_begin, size_t index_end,\n const IndexThreads *threads);\n\n //! Update Centroid in Thread\n void update_centroid_thread(size_t column, IndexCluster::CentroidList *out);\n\n //! Update Cluster in Thread\n void update_cluster_thread(size_t index_begin, size_t index_end,\n const IndexThreads *threads);\n\n //! Update Cluster's Features in Thread\n void update_features_thread(size_t column, IndexCluster::CentroidList *out);\n\n //! Update Labels in Thread\n void update_labels_thread(size_t index_begin, size_t index_end,\n std::vector *labels);\n\n protected:\n //! Members\n IndexMeta meta_{};\n IndexFeatures::Pointer features_{};\n LinearSeeker::Pointer seeker_{};\n std::vector shard_cluster_scores_{};\n std::vector> shard_cluster_features_{};\n std::shared_ptr shard_cluster_means_{};\n std::shared_ptr batch_means_{};\n std::vector batch_scores_{};\n double epsilon_{std::numeric_limits::epsilon()};\n float shard_factor_{16.0f};\n uint32_t max_iterations_{20u};\n uint32_t cluster_count_{0u};\n uint32_t thread_count_{0u};\n bool batch_{false};\n bool purge_empty_{false};\n};\n\n/*! Centroid Features\n */\nclass KmeansCentroidFeatures : public IndexFeatures {\n public:\n //! Constructor\n KmeansCentroidFeatures(const IndexMeta &meta,\n const IndexCluster::CentroidList ¢s)\n : centroids_(cents),\n feature_size_(meta.element_size()),\n feature_dimension_(meta.dimension()),\n data_type_(meta.data_type()) {}\n\n virtual size_t count(void) const {\n return centroids_.size();\n }\n\n virtual size_t dimension(void) const {\n return feature_dimension_;\n }\n\n virtual const void *element(size_t i) const {\n return centroids_[i].feature();\n }\n\n virtual IndexMeta::DataType data_type(void) const {\n return data_type_;\n }\n\n virtual size_t element_size(void) const {\n return feature_size_;\n }\n\n private:\n const IndexCluster::CentroidList ¢roids_;\n size_t feature_size_;\n size_t feature_dimension_;\n IndexMeta::DataType data_type_;\n};\n\nstatic inline std::shared_ptr NewVectorMean(const IndexMeta &meta) {\n switch (meta.data_type()) {\n case IndexMeta::DataType::DT_FP16:\n return std::make_shared>(\n meta.dimension());\n\n case IndexMeta::DataType::DT_FP32:\n return std::make_shared>(meta.dimension());\n\n case IndexMeta::DataType::DT_FP64:\n return std::make_shared>(meta.dimension());\n\n case IndexMeta::DataType::DT_INT8:\n return std::make_shared>(meta.dimension());\n\n case IndexMeta::DataType::DT_INT4:\n return std::make_shared>(meta.dimension());\n\n case IndexMeta::DataType::DT_INT16:\n return std::make_shared>(meta.dimension());\n\n default:\n break;\n }\n // As binary default\n return std::make_shared(meta.dimension());\n}\n\nstatic inline std::shared_ptr NewVectorMeanArray(\n const IndexMeta &meta) {\n switch (meta.data_type()) {\n case IndexMeta::DataType::DT_FP16:\n return std::make_shared<\n GeneralVectorMeanArray>>(\n meta.dimension());\n\n case IndexMeta::DataType::DT_FP32:\n return std::make_shared<\n GeneralVectorMeanArray>>(meta.dimension());\n\n case IndexMeta::DataType::DT_FP64:\n return std::make_shared<\n GeneralVectorMeanArray>>(\n meta.dimension());\n\n case IndexMeta::DataType::DT_INT8:\n return std::make_shared<\n GeneralVectorMeanArray>>(\n meta.dimension());\n\n case IndexMeta::DataType::DT_INT4:\n return std::make_shared<\n GeneralVectorMeanArray>>(meta.dimension());\n\n case IndexMeta::DataType::DT_INT16:\n return std::make_shared<\n GeneralVectorMeanArray>>(\n meta.dimension());\n\n default:\n break;\n }\n // As binary default\n return std::make_shared>(\n meta.dimension());\n}\n\nstatic inline std::shared_ptr NewVectorMeanArray(\n const IndexMeta &meta, const IndexCluster::CentroidList ¢s) {\n switch (meta.data_type()) {\n case IndexMeta::DataType::DT_FP16: {\n auto ptr = std::make_shared<\n GeneralVectorMeanArray>>(\n meta.dimension());\n\n for (const auto &it : cents) {\n ptr->emplace(reinterpret_cast(it.feature()),\n meta.dimension(), it.follows());\n }\n return ptr;\n }\n\n case IndexMeta::DataType::DT_FP32: {\n auto ptr =\n std::make_shared>>(\n meta.dimension());\n\n for (const auto &it : cents) {\n ptr->emplace(reinterpret_cast(it.feature()),\n meta.dimension(), it.follows());\n }\n return ptr;\n }\n\n case IndexMeta::DataType::DT_FP64: {\n auto ptr =\n std::make_shared>>(\n meta.dimension());\n\n for (const auto &it : cents) {\n ptr->emplace(reinterpret_cast(it.feature()),\n meta.dimension(), it.follows());\n }\n return ptr;\n }\n\n case IndexMeta::DataType::DT_INT8: {\n auto ptr =\n std::make_shared>>(\n meta.dimension());\n\n for (const auto &it : cents) {\n ptr->emplace(reinterpret_cast(it.feature()),\n meta.dimension(), it.follows());\n }\n return ptr;\n }\n\n case IndexMeta::DataType::DT_INT4: {\n auto ptr =\n std::make_shared>>(\n meta.dimension());\n\n for (const auto &it : cents) {\n ptr->emplace(reinterpret_cast(it.feature()),\n meta.dimension(), it.follows());\n }\n return ptr;\n }\n\n case IndexMeta::DataType::DT_INT16: {\n auto ptr = std::make_shared<\n GeneralVectorMeanArray>>(\n meta.dimension());\n\n for (const auto &it : cents) {\n ptr->emplace(reinterpret_cast(it.feature()),\n meta.dimension(), it.follows());\n }\n return ptr;\n }\n\n default:\n break;\n }\n\n // As binary default\n auto ptr = std::make_shared>(\n meta.dimension());\n\n for (const auto &it : cents) {\n ptr->emplace(it.feature(), meta.dimension(), it.follows());\n }\n return ptr;\n}\n\nstatic inline double CalculateSSE(const IndexCluster::CentroidList ¢s) {\n double accum = 0.0;\n for (const auto &it : cents) {\n accum += it.score();\n }\n return accum;\n}\n\nstatic inline void PurgeCentroids(IndexCluster::CentroidList ¢s,\n bool cutting) {\n size_t index = 0;\n size_t tamp = cents.size();\n\n while (index < tamp) {\n if (cents[index].follows() == 0) {\n size_t last_index = tamp - 1;\n\n if (index != last_index) {\n std::swap(cents[index], cents[last_index]);\n }\n tamp = last_index;\n continue;\n }\n ++index;\n }\n if (cutting) {\n cents.resize(tamp);\n }\n}\n\nint KmeansCluster::init(const IndexMeta &meta, const ailego::Params ¶ms) {\n meta_ = meta;\n this->update_params(params);\n\n return this->init_seeker();\n}\n\nint KmeansCluster::cleanup(void) {\n features_.reset();\n shard_cluster_scores_.clear();\n shard_cluster_features_.clear();\n shard_cluster_means_.reset();\n batch_means_.reset();\n batch_scores_.clear();\n seeker_->cleanup();\n return 0;\n}\n\nint KmeansCluster::reset(void) {\n features_.reset();\n shard_cluster_scores_.clear();\n shard_cluster_features_.clear();\n shard_cluster_means_->clear();\n batch_means_->clear();\n batch_scores_.clear();\n seeker_->reset();\n return 0;\n}\n\nint KmeansCluster::update(const ailego::Params ¶ms) {\n this->update_params(params);\n return 0;\n}\n\nvoid KmeansCluster::suggest(uint32_t k) {\n cluster_count_ = k;\n}\n\nint KmeansCluster::mount(IndexFeatures::Pointer feats) {\n if (!feats) {\n return IndexError_InvalidArgument;\n }\n if (!feats->is_matched(meta_)) {\n return IndexError_Mismatch;\n }\n\n // Check dimension\n auto data_type = meta_.data_type();\n switch (data_type) {\n case IndexMeta::DataType::DT_INT4:\n if (feats->dimension() % 2 != 0) {\n LOG_ERROR(\n \"Unsupported feature dimension %zu (dimension of int4 \"\n \"must be an integer multiple of 2).\",\n feats->dimension());\n return IndexError_Mismatch;\n }\n break;\n case IndexMeta::DataType::DT_BINARY32:\n if (feats->dimension() % 32 != 0) {\n LOG_ERROR(\n \"Unsupported feature dimension %zu (dimension of binary32 \"\n \"must be an integer multiple of 32).\",\n feats->dimension());\n return IndexError_Mismatch;\n }\n break;\n case IndexMeta::DataType::DT_BINARY64:\n if (feats->dimension() % 64 != 0) {\n LOG_ERROR(\n \"Unsupported feature dimension %zu (dimension of binary64 \"\n \"must be an integer multiple of 64).\",\n feats->dimension());\n return IndexError_Mismatch;\n }\n break;\n default:\n break;\n }\n\n features_ = std::move(feats);\n return 0;\n}\n\nint KmeansCluster::cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n ailego::ElapsedTime stamp;\n\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n if (cents.empty()) {\n if (cluster_count_ == 0) {\n LOG_ERROR(\"The count of cluster is unknown.\");\n return IndexError_NoReady;\n }\n this->init_centroids(cluster_count_, ¢s);\n }\n\n if (batch_) {\n batch_means_ = NewVectorMeanArray(meta_, cents);\n batch_scores_.clear();\n for (const auto &it : cents) {\n batch_scores_.push_back(it.score());\n }\n }\n\n double cost = 0.0;\n\n // we need to do clustering and update the centroids' follows, even if\n // cents.size() == 1. Otherwise, the centroid with empty follows will be\n // removed if purge_empty enabled\n for (uint32_t i = 0; (i < max_iterations_) && (cents.size() > 0); ++i) {\n double new_cost, new_epsilon;\n\n int result = this->clustering(threads.get(), cents, &new_cost);\n if (result != 0) {\n LOG_ERROR(\"(%u) Failed to cluster.\", i + 1);\n return result;\n }\n\n new_epsilon = new_cost - cost;\n LOG_DEBUG(\"(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> %f = %f\",\n i, cents.size(), features_->count(),\n (size_t)stamp.milli_seconds(), cost, new_cost, new_epsilon);\n stamp.reset();\n\n new_epsilon = std::abs(new_epsilon);\n if (new_epsilon < epsilon_) {\n break;\n }\n cost = new_cost;\n }\n\n // Purge the empty centroids\n PurgeCentroids(cents, purge_empty_);\n return 0;\n}\n\nint KmeansCluster::classify(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (cents.empty()) {\n LOG_ERROR(\"The input centroid's list is empty.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n int result = this->build_seeker(cents);\n if (result != 0) {\n LOG_ERROR(\"Failed to build the seeker.\");\n return result;\n }\n\n this->update_clusters(threads.get(), cents);\n this->update_features(threads.get(), cents);\n return 0;\n}\n\nint KmeansCluster::label(IndexThreads::Pointer threads,\n const IndexCluster::CentroidList ¢s,\n std::vector *out) {\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (cents.empty()) {\n LOG_ERROR(\"The input centroid's list is empty.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n int result = this->build_seeker(cents);\n if (result != 0) {\n LOG_ERROR(\"Failed to build the seeker.\");\n return result;\n }\n\n this->update_labels(threads.get(), out);\n return 0;\n}\n\nbool KmeansCluster::is_valid(void) const {\n if (!seeker_ || !features_ || !features_->count()) {\n return false;\n }\n return true;\n}\n\nint KmeansCluster::clustering(IndexThreads *threads,\n IndexCluster::CentroidList ¢s, double *cost) {\n int result = this->build_seeker(cents);\n if (result != 0) {\n LOG_ERROR(\"Failed to build the seeker.\");\n return result;\n }\n\n this->split_clusters(threads, cents);\n this->update_centroids(threads, cents);\n *cost = CalculateSSE(cents);\n return 0;\n}\n\nvoid KmeansCluster::update_params(const ailego::Params ¶ms) {\n params.get(GENERAL_THREAD_COUNT, &thread_count_);\n params.get(GENERAL_CLUSTER_COUNT, &cluster_count_);\n params.get(KMEANS_CLUSTER_COUNT, &cluster_count_);\n params.get(KMEANS_CLUSTER_SHARD_FACTOR, &shard_factor_);\n params.get(KMEANS_CLUSTER_EPSILON, &epsilon_);\n params.get(KMEANS_CLUSTER_MAX_ITERATIONS, &max_iterations_);\n params.get(KMEANS_CLUSTER_BATCH, &batch_);\n params.get(KMEANS_CLUSTER_PURGE_EMPTY, &purge_empty_);\n}\n\nint KmeansCluster::init_seeker(void) {\n seeker_.reset(new (std::nothrow) LinearSeeker);\n if (!seeker_) {\n LOG_ERROR(\"Failed to create linear seeker.\");\n return IndexError_NoMemory;\n }\n\n int result = seeker_->init(meta_);\n if (result != 0) {\n LOG_ERROR(\"Failed to initialize linear seeker.\");\n return result;\n }\n\n return 0;\n}\n\nint KmeansCluster::build_seeker(const IndexCluster::CentroidList ¢s) {\n int result =\n seeker_->mount(std::make_shared(meta_, cents));\n if (result != 0) {\n LOG_ERROR(\"Failed to mount features for linear seeker.\");\n return result;\n }\n\n return 0;\n}\n\nbool KmeansCluster::check_centroids(const IndexCluster::CentroidList ¢s) {\n for (const auto &it : cents) {\n if (it.size() != meta_.element_size()) {\n return false;\n }\n }\n return true;\n}\n\nvoid KmeansCluster::init_centroids(size_t count,\n IndexCluster::CentroidList *out) {\n size_t feature_size = features_->element_size();\n size_t features_count = features_->count();\n size_t sample_count = std::min(count, features_count);\n\n ailego::Reservoir sampler(sample_count);\n for (size_t i = 0; i < features_count; ++i) {\n sampler.fill(i);\n }\n\n // Save centroids\n out->reserve(sampler.pool().size());\n for (auto i : sampler.pool()) {\n out->emplace_back(features_->element(i), feature_size);\n }\n}\n\nvoid KmeansCluster::init_containers(size_t shard_count) {\n if (!shard_cluster_means_) {\n shard_cluster_means_ = NewVectorMeanArray(meta_);\n }\n shard_cluster_means_->clear();\n shard_cluster_means_->resize(shard_count);\n shard_cluster_scores_.clear();\n shard_cluster_scores_.resize(shard_count);\n}\n\nvoid KmeansCluster::init_features_containers(size_t shard_count) {\n shard_cluster_features_.resize(shard_count);\n for (auto &features : shard_cluster_features_) {\n features.clear();\n }\n}\n\nvoid KmeansCluster::split_clusters(IndexThreads *threads,\n const IndexCluster::CentroidList ¢s) {\n // Initilize containers\n this->init_containers(threads->count() * cents.size());\n auto task_group = threads->make_group();\n\n // Initilize base information\n size_t features_count = features_->count();\n size_t shard_count = std::max(\n static_cast(std::ceil(threads->count() * shard_factor_)), 1u);\n size_t fregment_count = (features_count + shard_count - 1) / shard_count;\n\n for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count);\n ++i) {\n size_t next_index = index + fregment_count;\n if (next_index > features_count) {\n next_index = features_count;\n }\n\n // Process in work thread\n task_group->submit(\n ailego::Closure::New(this, &KmeansCluster::split_clusters_thread, index,\n next_index, threads));\n\n // Next index\n index = next_index;\n }\n task_group->wait_finish();\n}\n\nvoid KmeansCluster::update_centroids(IndexThreads *threads,\n IndexCluster::CentroidList ¢s) {\n auto task_group = threads->make_group();\n for (size_t i = 0; i < cents.size(); ++i) {\n task_group->submit(ailego::Closure::New(\n this, &KmeansCluster::update_centroid_thread, i, ¢s));\n }\n task_group->wait_finish();\n}\n\nvoid KmeansCluster::update_clusters(IndexThreads *threads,\n const IndexCluster::CentroidList ¢s) {\n // Initilize containers\n this->init_features_containers(threads->count() * cents.size());\n auto task_group = threads->make_group();\n\n size_t features_count = features_->count();\n size_t shard_count = std::max(\n static_cast(std::ceil(threads->count() * shard_factor_)), 1u);\n size_t fregment_count = (features_count + shard_count - 1) / shard_count;\n\n for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count);\n ++i) {\n size_t next_index = index + fregment_count;\n if (next_index > features_count) {\n next_index = features_count;\n }\n // Process in work thread\n task_group->submit(\n ailego::Closure::New(this, &KmeansCluster::update_cluster_thread, index,\n next_index, threads));\n\n // Next index\n index = next_index;\n }\n task_group->wait_finish();\n}\n\nvoid KmeansCluster::update_features(IndexThreads *threads,\n IndexCluster::CentroidList ¢s) {\n auto task_group = threads->make_group();\n for (size_t i = 0; i < cents.size(); ++i) {\n // Process in work thread\n task_group->submit(ailego::Closure::New(\n this, &KmeansCluster::update_features_thread, i, ¢s));\n }\n task_group->wait_finish();\n}\n\nvoid KmeansCluster::update_labels(IndexThreads *threads,\n std::vector *labels) {\n size_t features_count = features_->count();\n size_t shard_count = std::max(\n static_cast(std::ceil(threads->count() * shard_factor_)), 1u);\n size_t fregment_count = (features_count + shard_count - 1) / shard_count;\n auto task_group = threads->make_group();\n\n // Prepare buffer\n labels->resize(features_count);\n\n for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count);\n ++i) {\n size_t next_index = index + fregment_count;\n if (next_index > features_count) {\n next_index = features_count;\n }\n\n // Process in work thread\n task_group->submit(ailego::Closure::New(\n this, &KmeansCluster::update_labels_thread, index, next_index, labels));\n\n // Next index\n index = next_index;\n }\n task_group->wait_finish();\n}\n\nvoid KmeansCluster::split_clusters_thread(size_t index_begin, size_t index_end,\n const IndexThreads *threads) {\n size_t feature_size = features_->element_size();\n size_t thread_offset = threads->indexof_this() * seeker_->original()->count();\n\n for (size_t i = index_begin; i != index_end; ++i) {\n const void *feat = features_->element(i);\n LinearSeeker::Document result(0, std::numeric_limits::max());\n\n // ignore error\n seeker_->seek(feat, meta_.element_size(), &result);\n\n size_t sel_column = thread_offset + result.index;\n shard_cluster_scores_[sel_column] += result.score;\n shard_cluster_means_->at(sel_column).plus(feat, feature_size);\n }\n}\n\nvoid KmeansCluster::update_centroid_thread(size_t column,\n IndexCluster::CentroidList *out) {\n size_t cluster_count = out->size();\n double cluster_score = 0.0;\n\n // Create Accumulator\n std::shared_ptr accum = NewVectorMean(meta_);\n if (batch_) {\n cluster_score += batch_scores_[column];\n accum->merge(batch_means_->at(column));\n }\n\n // Compute the score of centroid\n for (size_t i = column; i < shard_cluster_scores_.size();\n i += cluster_count) {\n cluster_score += shard_cluster_scores_[i];\n accum->merge(shard_cluster_means_->at(i));\n }\n\n // Update centroid\n IndexCluster::Centroid *centroid = &(out->at(column));\n centroid->set_score(cluster_score);\n centroid->set_follows(accum->count());\n accum->mean(centroid->mutable_buffer());\n}\n\nvoid KmeansCluster::update_cluster_thread(size_t index_begin, size_t index_end,\n const IndexThreads *threads) {\n size_t thread_offset = threads->indexof_this() * seeker_->original()->count();\n\n for (size_t i = index_begin; i != index_end; ++i) {\n const void *feat = features_->element(i);\n LinearSeeker::Document result(0, std::numeric_limits::max());\n\n // ignore error\n seeker_->seek(feat, meta_.element_size(), &result);\n\n size_t sel_column = thread_offset + result.index;\n shard_cluster_features_[sel_column].emplace_back(feat);\n }\n}\n\nvoid KmeansCluster::update_features_thread(size_t column,\n IndexCluster::CentroidList *out) {\n size_t cluster_count = out->size();\n size_t cluster_follows = 0u;\n\n // Compute the follows of cluster\n for (size_t i = column; i < shard_cluster_features_.size();\n i += cluster_count) {\n cluster_follows += shard_cluster_features_[i].size();\n }\n\n // Merge all features in cluster\n std::vector &cluster_features =\n *(out->at(column).mutable_similars());\n cluster_features.resize(cluster_follows);\n\n for (size_t i = column, j = 0; i < shard_cluster_features_.size();\n i += cluster_count) {\n const std::vector &it = shard_cluster_features_[i];\n std::memcpy(&cluster_features[j], it.data(), it.size() * sizeof(void *));\n j += it.size();\n }\n}\n\nvoid KmeansCluster::update_labels_thread(size_t index_begin, size_t index_end,\n std::vector *labels) {\n for (size_t i = index_begin; i != index_end; ++i) {\n const void *feat = features_->element(i);\n LinearSeeker::Document result(0, std::numeric_limits::max());\n\n // ignore error\n seeker_->seek(feat, meta_.element_size(), &result);\n (*labels)[i] = static_cast(result.index);\n }\n}\n\nINDEX_FACTORY_REGISTER_CLUSTER_ALIAS(KmeansCluster, KmeansCluster, false);\nINDEX_FACTORY_REGISTER_CLUSTER_ALIAS(BatchKmeansCluster, KmeansCluster, true);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/cluster/linear_seeker.cc", "content": "\n// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"linear_seeker.h\"\n\nnamespace zvec {\nnamespace core {\n\nint LinearSeeker::seek(const void *query, size_t len, Document *out) {\n if (ailego_unlikely(!query || !out || meta_.element_size() != len)) {\n return IndexError_InvalidArgument;\n }\n\n float sel_score = std::numeric_limits::max();\n uint32_t sel_column = 0;\n uint32_t total = static_cast(features_->count());\n\n for (uint32_t i = 0; i < total; ++i) {\n float score = 0.0f;\n\n distance_func_(features_->element(i), query, meta_.dimension(), &score);\n if (score < sel_score) {\n sel_score = score;\n sel_column = i;\n }\n }\n\n out->index = sel_column;\n out->score = sel_score;\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/cluster/linear_seeker.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"seeker.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Linear Seeker\n */\nclass LinearSeeker : public Seeker {\n public:\n typedef std::shared_ptr Pointer;\n\n //! Constructor\n LinearSeeker(void) : meta_(), metric_(), features_() {}\n\n //! Destructor\n ~LinearSeeker(void) {}\n\n //! Initialize Seeker\n int init(const IndexMeta &meta) override {\n meta_ = meta;\n\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"Create Metric %s failed.\", meta_.metric_name().c_str());\n\n return IndexError_Unsupported;\n }\n int ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"IndexMetric init failed wit ret %d.\", ret);\n\n return ret;\n }\n distance_func_ = metric_->distance_matrix(1, 1);\n if (!distance_func_) {\n LOG_ERROR(\"DistanceMatrix function is nullptr.\");\n\n return IndexError_Unsupported;\n }\n return 0;\n }\n\n //! Cleanup Seeker\n int cleanup(void) override {\n features_.reset();\n return 0;\n }\n\n //! Reset Seeker\n int reset(void) override {\n features_.reset();\n return 0;\n }\n\n //! Mount features\n int mount(IndexFeatures::Pointer feats) override {\n if (!feats) {\n return IndexError_InvalidArgument;\n }\n if (!feats->is_matched(meta_)) {\n return IndexError_Mismatch;\n }\n features_ = std::move(feats);\n return 0;\n }\n\n //! Seek (TOP 1 Document)\n int seek(const void *query, size_t len, Document *out) override;\n\n //! Retrieve the original features\n IndexFeatures::Pointer original(void) const override {\n return features_;\n }\n\n private:\n IndexMeta meta_{};\n IndexMetric::Pointer metric_{};\n IndexFeatures::Pointer features_{};\n IndexMetric::MatrixDistance distance_func_{nullptr};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/cluster/opt_kmeans_cluster.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n#include \n#include \n#include \n#include \n#include \"cluster_params.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Optimize K-Means cluster algorithm\n */\nclass OptKmeansAlgorithm : public IndexCluster {\n public:\n //! Constructor\n OptKmeansAlgorithm(void) {}\n\n //! Destructor\n virtual ~OptKmeansAlgorithm(void) {}\n\n //! Initialize Cluster\n int init(const IndexMeta &meta, const ailego::Params ¶ms);\n\n //! Mount features\n virtual int mount(IndexFeatures::Pointer feats);\n\n //! Suggest dividing to K clusters\n virtual void suggest(uint32_t k);\n\n //! Classify\n virtual int classify(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n //! Label\n virtual int label(IndexThreads::Pointer threads,\n const IndexCluster::CentroidList ¢s,\n std::vector *out);\n\n //! Cluster\n virtual int cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) = 0;\n\n //! Cleanup Cluster\n virtual int cleanup(void);\n\n //! Reset Cluster\n virtual int reset(void);\n\n //! Update Cluster\n virtual int update(const ailego::Params ¶ms);\n\n protected:\n //! Update parameters\n void update_params(const ailego::Params ¶ms);\n\n //! Init Kmeans Algorithm\n int init_algorithm();\n\n //! Init Distance function\n int init_distance_func();\n\n //! Check Centroids\n bool check_centroids(const IndexCluster::CentroidList ¢s);\n\n //! Test if it is valid\n bool is_valid(void) const;\n\n //! Update Clusters\n void update_clusters(IndexThreads *threads,\n const IndexCluster::CentroidList ¢s);\n\n //! Update Cluster in Thread\n void update_cluster_thread(size_t index_begin, size_t index_end,\n const IndexThreads *threads,\n const IndexCluster::CentroidList ¢s);\n\n //! Initialize Shard Features Containers\n void init_features_containers(size_t shard_count);\n\n //! Update Clusters' Features\n void update_features(IndexThreads *threads,\n IndexCluster::CentroidList ¢s);\n\n //! Update Cluster's Features in Thread\n void update_features_thread(size_t column, IndexCluster::CentroidList *out);\n\n //! Update Labels\n void update_labels(IndexThreads *threads, std::vector *labels,\n const IndexCluster::CentroidList ¢s);\n\n //! Update Labels in Thread\n void update_labels_thread(size_t index_begin, size_t index_end,\n std::vector *labels,\n const IndexCluster::CentroidList ¢s);\n\n //! Initialize Centroids\n void init_centroids(size_t count, IndexCluster::CentroidList *out);\n\n protected:\n uint32_t cluster_count_{0u};\n uint32_t thread_count_{0u};\n uint32_t max_iterations_{20u};\n double epsilon_{std::numeric_limits::epsilon()};\n float shard_factor_{16.0f};\n bool purge_empty_{false};\n bool assumption_free_{false};\n uint32_t markov_chain_length_{32};\n IndexMeta meta_{};\n IndexFeatures::Pointer features_{};\n std::vector> shard_cluster_features_{};\n IndexMetric::MatrixDistance distance_func_{nullptr};\n};\n\nbool OptKmeansAlgorithm::is_valid(void) const {\n if (!features_ || !features_->count()) {\n return false;\n }\n return true;\n}\n\nbool OptKmeansAlgorithm::check_centroids(\n const IndexCluster::CentroidList ¢s) {\n for (const auto &it : cents) {\n if (it.size() != meta_.element_size()) {\n return false;\n }\n }\n return true;\n}\n\nvoid OptKmeansAlgorithm::update_params(const ailego::Params ¶ms) {\n params.get(GENERAL_THREAD_COUNT, &thread_count_);\n params.get(GENERAL_CLUSTER_COUNT, &cluster_count_);\n params.get(OPTKMEANS_CLUSTER_COUNT, &cluster_count_);\n params.get(OPTKMEANS_CLUSTER_SHARD_FACTOR, &shard_factor_);\n params.get(OPTKMEANS_CLUSTER_EPSILON, &epsilon_);\n params.get(OPTKMEANS_CLUSTER_MAX_ITERATIONS, &max_iterations_);\n params.get(OPTKMEANS_CLUSTER_PURGE_EMPTY, &purge_empty_);\n params.get(OPTKMEANS_CLUSTER_MARKOV_CHAIN_LENGTH, &markov_chain_length_);\n params.get(OPTKMEANS_CLUSTER_ASSUMPTION_FREE, &assumption_free_);\n}\n\nint OptKmeansAlgorithm::init_distance_func() {\n IndexMetric::Pointer metric_{};\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"Create Metric %s failed.\", meta_.metric_name().c_str());\n return IndexError_Unsupported;\n }\n int ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"IndexMetric init failed wit ret %d.\", ret);\n return ret;\n }\n distance_func_ = metric_->distance_matrix(1, 1);\n if (!distance_func_) {\n LOG_ERROR(\"DistanceMatrix function is nullptr.\");\n return IndexError_Unsupported;\n }\n return 0;\n}\n\nvoid OptKmeansAlgorithm::update_clusters(\n IndexThreads *threads, const IndexCluster::CentroidList ¢s) {\n // Initilize containers\n this->init_features_containers(threads->count() * cents.size());\n auto task_group = threads->make_group();\n\n size_t features_count = features_->count();\n size_t shard_count = std::max(\n static_cast(std::ceil(threads->count() * shard_factor_)), 1u);\n size_t fregment_count = (features_count + shard_count - 1) / shard_count;\n\n for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count);\n ++i) {\n size_t next_index = index + fregment_count;\n if (next_index > features_count) {\n next_index = features_count;\n }\n\n // Process in work thread·\n task_group->submit(\n ailego::Closure::New(this, &OptKmeansAlgorithm::update_cluster_thread,\n index, next_index, threads, cents));\n\n // Next index\n index = next_index;\n }\n task_group->wait_finish();\n}\n\nvoid OptKmeansAlgorithm::update_cluster_thread(\n size_t index_begin, size_t index_end, const IndexThreads *threads,\n const IndexCluster::CentroidList ¢s) {\n size_t thread_offset = threads->indexof_this() * cents.size();\n\n for (size_t i = index_begin; i != index_end; ++i) {\n const void *feat = features_->element(i);\n uint32_t sel_index = 0;\n float sel_score = std::numeric_limits::max();\n\n // todo: get min distance\n uint32_t total = static_cast(cents.size());\n for (uint32_t j = 0; j < total; ++j) {\n float score = 0.0f;\n\n distance_func_(cents[j].feature(), feat, meta_.dimension(), &score);\n if (score < sel_score) {\n sel_score = score;\n sel_index = j;\n }\n }\n\n size_t sel_column = thread_offset + sel_index;\n shard_cluster_features_[sel_column].emplace_back(feat);\n }\n}\n\nvoid OptKmeansAlgorithm::init_features_containers(size_t shard_count) {\n shard_cluster_features_.resize(shard_count);\n for (auto &features : shard_cluster_features_) {\n features.clear();\n }\n}\n\nvoid OptKmeansAlgorithm::update_features(IndexThreads *threads,\n IndexCluster::CentroidList ¢s) {\n auto task_group = threads->make_group();\n for (size_t i = 0; i < cents.size(); ++i) {\n // Process in work thread\n task_group->submit(ailego::Closure::New(\n this, &OptKmeansAlgorithm::update_features_thread, i, ¢s));\n }\n task_group->wait_finish();\n}\n\nvoid OptKmeansAlgorithm::update_labels(\n IndexThreads *threads, std::vector *labels,\n const IndexCluster::CentroidList ¢s) {\n size_t features_count = features_->count();\n size_t shard_count = std::max(\n static_cast(std::ceil(threads->count() * shard_factor_)), 1u);\n size_t fregment_count = (features_count + shard_count - 1) / shard_count;\n auto task_group = threads->make_group();\n\n // Prepare buffer\n labels->resize(features_count);\n\n for (size_t i = 0, index = 0; (i != shard_count) && (index < features_count);\n ++i) {\n size_t next_index = index + fregment_count;\n if (next_index > features_count) {\n next_index = features_count;\n }\n\n // Process in work thread\n task_group->submit(\n ailego::Closure::New(this, &OptKmeansAlgorithm::update_labels_thread,\n index, next_index, labels, cents));\n\n // Next index\n index = next_index;\n }\n task_group->wait_finish();\n}\n\nvoid OptKmeansAlgorithm::update_labels_thread(\n size_t index_begin, size_t index_end, std::vector *labels,\n const IndexCluster::CentroidList ¢s) {\n for (size_t i = index_begin; i != index_end; ++i) {\n const void *feat = features_->element(i);\n\n uint32_t sel_index = 0;\n float sel_score = std::numeric_limits::max();\n\n // todo: get min distance\n uint32_t total = static_cast(cents.size());\n for (uint32_t j = 0; j < total; ++j) {\n float score = 0.0f;\n\n distance_func_(cents[j].feature(), feat, meta_.dimension(), &score);\n if (score < sel_score) {\n sel_score = score;\n sel_index = j;\n }\n }\n\n (*labels)[i] = static_cast(sel_index);\n }\n}\n\nvoid OptKmeansAlgorithm::init_centroids(size_t count,\n IndexCluster::CentroidList *out) {\n // Just resize, because the get random centroid step is done by cluster_once\n out->resize(count);\n}\n\nvoid OptKmeansAlgorithm::update_features_thread(\n size_t column, IndexCluster::CentroidList *out) {\n size_t cluster_count = out->size();\n size_t cluster_follows = 0u;\n\n // Compute the follows of cluster\n for (size_t i = column; i < shard_cluster_features_.size();\n i += cluster_count) {\n cluster_follows += shard_cluster_features_[i].size();\n }\n\n // Merge all features in cluster\n std::vector &cluster_features =\n *(out->at(column).mutable_similars());\n cluster_features.resize(cluster_follows);\n\n for (size_t i = column, j = 0; i < shard_cluster_features_.size();\n i += cluster_count) {\n const std::vector &it = shard_cluster_features_[i];\n std::memcpy(&cluster_features[j], it.data(), it.size() * sizeof(void *));\n j += it.size();\n }\n}\n\nstatic inline void PurgeCentroids(IndexCluster::CentroidList ¢s,\n bool cutting) {\n size_t index = 0;\n size_t tamp = cents.size();\n\n while (index < tamp) {\n if (cents[index].follows() == 0) {\n size_t last_index = tamp - 1;\n\n if (index != last_index) {\n std::swap(cents[index], cents[last_index]);\n }\n tamp = last_index;\n continue;\n }\n ++index;\n }\n if (cutting) {\n cents.resize(tamp);\n }\n}\n\nint OptKmeansAlgorithm::init(const IndexMeta &meta,\n const ailego::Params ¶ms) {\n meta_ = meta;\n this->update_params(params);\n\n return init_distance_func();\n}\n\nint OptKmeansAlgorithm::mount(IndexFeatures::Pointer feats) {\n if (!feats) {\n return IndexError_InvalidArgument;\n }\n if (!feats->is_matched(meta_)) {\n return IndexError_Mismatch;\n }\n\n // Check dimension\n auto type_ = meta_.data_type();\n switch (type_) {\n case IndexMeta::DataType::DT_INT4:\n if (feats->dimension() % 8 != 0) {\n LOG_ERROR(\n \"Unsupported feature dimension %zu (dimension of int4 \"\n \"must be an integer multiple of 8).\",\n feats->dimension());\n return IndexError_Mismatch;\n }\n break;\n case IndexMeta::DataType::DT_INT8:\n if (feats->dimension() % 4 != 0) {\n LOG_ERROR(\n \"Unsupported feature dimension %zu (dimension of int8 \"\n \"must be an integer multiple of 4).\",\n feats->dimension());\n return IndexError_Mismatch;\n }\n break;\n case IndexMeta::DataType::DT_BINARY32:\n case IndexMeta::DataType::DT_BINARY64:\n if (feats->dimension() % 32 != 0) {\n LOG_ERROR(\n \"Unsupported feature dimension %zu (dimension of binary \"\n \"must be an integer multiple of 32).\",\n feats->dimension());\n return IndexError_Mismatch;\n }\n break;\n default:\n break;\n }\n\n features_ = std::move(feats);\n return 0;\n}\n\nvoid OptKmeansAlgorithm::suggest(uint32_t k) {\n cluster_count_ = k;\n}\n\nint OptKmeansAlgorithm::classify(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (cents.empty()) {\n LOG_ERROR(\"The input centroid's list is empty.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n this->update_clusters(threads.get(), cents);\n this->update_features(threads.get(), cents);\n return 0;\n}\n\nint OptKmeansAlgorithm::label(IndexThreads::Pointer threads,\n const IndexCluster::CentroidList ¢s,\n std::vector *out) {\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (cents.empty()) {\n LOG_ERROR(\"The input centroid's list is empty.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n this->update_labels(threads.get(), out, cents);\n return 0;\n}\n\nint OptKmeansAlgorithm::update(const ailego::Params ¶ms) {\n this->update_params(params);\n // algorithm_->reset(cluster_count_);\n return 0;\n}\n\nint OptKmeansAlgorithm::reset(void) {\n features_.reset();\n shard_cluster_features_.clear();\n\n return 0;\n}\n\nint OptKmeansAlgorithm::cleanup(void) {\n features_.reset();\n shard_cluster_features_.clear();\n\n return 0;\n}\n\n\n/*! Numerical K-Means cluster algorithm\n */\ntemplate \nclass NumericalKmeansAlgorithm : public OptKmeansAlgorithm {\n public:\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n // Check supporting type\n static_assert(ailego::IsArithmetic::value,\n \"ValueType must be arithmetic\");\n\n //! Constructor\n NumericalKmeansAlgorithm(void) {}\n\n //! Destructor\n virtual ~NumericalKmeansAlgorithm(void) {}\n\n //! Cluster\n virtual int cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n protected:\n void update_centroids(\n IndexCluster::CentroidList ¢s,\n const ailego::NumericalKmeans &algorithm);\n};\n\ntemplate \nvoid NumericalKmeansAlgorithm::update_centroids(\n IndexCluster::CentroidList ¢s,\n const ailego::NumericalKmeans &algorithm) {\n this->init_centroids(algorithm.centroids().count(), ¢s);\n for (size_t i = 0; i < cents.size(); ++i) {\n IndexCluster::Centroid *centroid = &(cents.at(i));\n centroid->set_score(algorithm.context().clusters()[i].cost());\n centroid->set_follows(algorithm.context().clusters()[i].count());\n centroid->set_feature(algorithm.centroids()[i],\n meta_.dimension() * sizeof(T));\n }\n}\n\ntemplate \nint NumericalKmeansAlgorithm::cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n ailego::ElapsedTime stamp;\n\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n // get cluster algorithm\n size_t centroid_count =\n cents.empty()\n ? std::min(cluster_count_, static_cast(features_->count()))\n : cents.size();\n if (centroid_count == 0) {\n LOG_ERROR(\"The count of cluster is unknown.\");\n return IndexError_NoReady;\n }\n ailego::NumericalKmeans algorithm(centroid_count,\n meta_.dimension());\n\n // mount features into algorithm\n auto features_count = features_->count();\n auto dim = meta_.dimension();\n\n algorithm.feature_matrix_reserve(features_count);\n\n for (size_t i = 0; i < features_count; ++i) {\n auto vec = reinterpret_cast(features_->element(i));\n algorithm.append(vec, dim);\n }\n\n if (!cents.empty()) {\n auto centroids = algorithm.mutable_centroids();\n centroids->reserve(cents.size());\n for (const auto &it : cents) {\n centroids->append(reinterpret_cast(it.feature()),\n meta_.dimension());\n }\n } else {\n ailego::Kmc2CentroidsGenerator<\n ailego::NumericalKmeans, IndexThreads>\n g;\n g.set_chain_length(markov_chain_length_);\n g.set_assumption_free(assumption_free_);\n algorithm.init_centroids(*threads, g);\n }\n\n double cost = 0.0;\n\n for (uint32_t i = 0; i < max_iterations_; ++i) {\n double old_cost, new_epsilon;\n old_cost = cost;\n\n bool result = algorithm.cluster_once(*threads, &cost);\n if (result != true) {\n LOG_ERROR(\"(%u) Failed to cluster.\", i + 1);\n return -1;\n }\n\n new_epsilon = std::abs(cost - old_cost);\n LOG_DEBUG(\"(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> %f = %f\",\n i, algorithm.centroids().count(), features_->count(),\n (size_t)stamp.milli_seconds(), old_cost, cost, new_epsilon);\n stamp.reset();\n\n if (new_epsilon < epsilon_) {\n break;\n }\n }\n\n // update_centroids(cents);\n update_centroids(cents, algorithm);\n\n // Purge the empty centroids\n PurgeCentroids(cents, purge_empty_);\n return 0;\n}\n\n/*! Nibble K-Means cluster algorithm\n */\ntemplate \nclass NibbleKmeansAlgorithm : public OptKmeansAlgorithm {\n public:\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n // Check supporting type\n static_assert(ailego::IsArithmetic::value,\n \"ValueType must be arithmetic\");\n\n //! Constructor\n NibbleKmeansAlgorithm(void) {}\n\n //! Destructor\n virtual ~NibbleKmeansAlgorithm(void) {}\n\n //! Cluster\n virtual int cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n protected:\n //! update centroids\n void update_centroids(IndexCluster::CentroidList ¢s,\n const ailego::NibbleKmeans &algorithm);\n};\n\ntemplate \nvoid NibbleKmeansAlgorithm::update_centroids(\n IndexCluster::CentroidList ¢s,\n const ailego::NibbleKmeans &algorithm) {\n this->init_centroids(algorithm.centroids().count(), ¢s);\n for (size_t i = 0; i < cents.size(); ++i) {\n IndexCluster::Centroid *centroid = &(cents.at(i));\n centroid->set_score(algorithm.context().clusters()[i].cost());\n centroid->set_follows(algorithm.context().clusters()[i].count());\n centroid->set_feature(algorithm.centroids()[i], (meta_.dimension() >> 1));\n }\n}\n\ntemplate \nint NibbleKmeansAlgorithm::cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n ailego::ElapsedTime stamp;\n\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n // get cluster algorithm\n size_t centroid_count =\n cents.empty()\n ? std::min(cluster_count_, static_cast(features_->count()))\n : cents.size();\n if (centroid_count == 0) {\n LOG_ERROR(\"The count of cluster is unknown.\");\n return IndexError_NoReady;\n }\n ailego::NibbleKmeans algorithm(centroid_count,\n meta_.dimension());\n\n // mount features into algorithm\n auto features_count = features_->count();\n auto dim = meta_.dimension();\n for (size_t i = 0; i < features_count; ++i) {\n auto vec = reinterpret_cast::type *>(\n features_->element(i));\n algorithm.append(vec, dim);\n }\n\n if (!cents.empty()) {\n auto centroids = algorithm.mutable_centroids();\n centroids->reserve(cents.size());\n for (const auto &it : cents) {\n centroids->append(\n reinterpret_cast::type *>(\n it.feature()),\n size_t(meta_.dimension()));\n }\n } else {\n ailego::Kmc2CentroidsGenerator<\n ailego::NibbleKmeans, IndexThreads>\n g;\n g.set_chain_length(markov_chain_length_);\n g.set_assumption_free(assumption_free_);\n algorithm.init_centroids(*threads, g);\n }\n\n double cost = 0.0;\n\n for (uint32_t i = 0; i < max_iterations_; ++i) {\n double old_cost, new_epsilon;\n old_cost = cost;\n\n bool result = algorithm.cluster_once(*threads, &cost);\n if (result != true) {\n LOG_ERROR(\"(%u) Failed to cluster.\", i + 1);\n return -1;\n }\n\n new_epsilon = std::abs(cost - old_cost);\n LOG_DEBUG(\n \"(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> \"\n \"%f = %f\",\n i, algorithm.centroids().count(), features_->count(),\n (size_t)stamp.milli_seconds(), old_cost, cost, new_epsilon);\n stamp.reset();\n\n if (new_epsilon < epsilon_) {\n break;\n }\n }\n\n // update centroids\n update_centroids(cents, algorithm);\n\n // Purge the empty centroids\n PurgeCentroids(cents, purge_empty_);\n return 0;\n}\n\n/*! Binary K-Means cluster algorithm\n */\ntemplate \nclass BinaryKmeansAlgorithm : public OptKmeansAlgorithm {\n public:\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n // Check supporting type\n static_assert(ailego::IsArithmetic::value,\n \"ValueType must be arithmetic\");\n\n //! Constructor\n BinaryKmeansAlgorithm(void) {}\n\n //! Destructor\n virtual ~BinaryKmeansAlgorithm(void) {}\n\n //! Cluster\n virtual int cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n protected:\n //! update centroids\n void update_centroids(IndexCluster::CentroidList ¢s,\n const ailego::BinaryKmeans &algorithm);\n};\n\ntemplate \nvoid BinaryKmeansAlgorithm::update_centroids(\n IndexCluster::CentroidList ¢s,\n const ailego::BinaryKmeans &algorithm) {\n this->init_centroids(algorithm.centroids().count(), ¢s);\n for (size_t i = 0; i < cents.size(); ++i) {\n IndexCluster::Centroid *centroid = &(cents.at(i));\n centroid->set_score(algorithm.context().clusters()[i].cost());\n centroid->set_follows(algorithm.context().clusters()[i].count());\n centroid->set_feature(algorithm.centroids()[i], (meta_.dimension() >> 3));\n }\n}\n\ntemplate \nint BinaryKmeansAlgorithm::cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n ailego::ElapsedTime stamp;\n\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n // get cluster algorithm\n size_t centroid_count =\n cents.empty()\n ? std::min(cluster_count_, static_cast(features_->count()))\n : cents.size();\n if (centroid_count == 0) {\n LOG_ERROR(\"The count of cluster is unknown.\");\n return IndexError_NoReady;\n }\n ailego::BinaryKmeans algorithm(centroid_count,\n meta_.dimension());\n\n // mount features into algorithm\n auto features_count = features_->count();\n auto dim = meta_.dimension();\n for (size_t i = 0; i < features_count; ++i) {\n auto vec = reinterpret_cast(features_->element(i));\n algorithm.append(vec, dim);\n }\n\n if (!cents.empty()) {\n auto centroids = algorithm.mutable_centroids();\n centroids->reserve(cents.size());\n for (const auto &it : cents) {\n centroids->append(reinterpret_cast(it.feature()),\n meta_.dimension());\n }\n } else {\n ailego::Kmc2CentroidsGenerator<\n ailego::BinaryKmeans, IndexThreads>\n g;\n g.set_chain_length(markov_chain_length_);\n g.set_assumption_free(assumption_free_);\n algorithm.init_centroids(*threads, g);\n }\n\n double cost = 0.0;\n\n for (uint32_t i = 0; i < max_iterations_; ++i) {\n double old_cost, new_epsilon;\n old_cost = cost;\n\n bool result = algorithm.cluster_once(*threads, &cost);\n if (result != true) {\n LOG_ERROR(\"(%u) Failed to cluster.\", i + 1);\n return -1;\n }\n\n new_epsilon = std::abs(cost - old_cost);\n LOG_DEBUG(\n \"(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> \"\n \"%f = %f\",\n i, algorithm.centroids().count(), features_->count(),\n (size_t)stamp.milli_seconds(), old_cost, cost, new_epsilon);\n stamp.reset();\n\n if (new_epsilon < epsilon_) {\n break;\n }\n }\n\n // update centroids\n update_centroids(cents, algorithm);\n\n // Purge the empty centroids\n PurgeCentroids(cents, purge_empty_);\n return 0;\n}\n\n/*! Numerical K-Means cluster algorithm\n */\ntemplate \nclass NumericalInnerProductKmeansAlgorithm : public OptKmeansAlgorithm {\n public:\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n // Check supporting type\n static_assert(ailego::IsArithmetic::value,\n \"ValueType must be arithmetic\");\n\n //! Constructor\n NumericalInnerProductKmeansAlgorithm(void) {}\n\n //! Destructor\n virtual ~NumericalInnerProductKmeansAlgorithm(void) {}\n\n //! Cluster\n virtual int cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n protected:\n void update_centroids(\n IndexCluster::CentroidList ¢s,\n const ailego::NumericalInnerProductKmeans &algorithm);\n};\n\ntemplate \nvoid NumericalInnerProductKmeansAlgorithm::update_centroids(\n IndexCluster::CentroidList ¢s,\n const ailego::NumericalInnerProductKmeans &algorithm) {\n this->init_centroids(algorithm.centroids().count(), ¢s);\n for (size_t i = 0; i < cents.size(); ++i) {\n IndexCluster::Centroid *centroid = &(cents.at(i));\n centroid->set_score(algorithm.context().clusters()[i].cost());\n centroid->set_follows(algorithm.context().clusters()[i].count());\n centroid->set_feature(algorithm.centroids()[i],\n meta_.dimension() * sizeof(T));\n }\n}\n\ntemplate \nint NumericalInnerProductKmeansAlgorithm::cluster(\n IndexThreads::Pointer threads, IndexCluster::CentroidList ¢s) {\n ailego::ElapsedTime stamp;\n\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n // get cluster algorithm\n size_t centroid_count =\n cents.empty()\n ? std::min(cluster_count_, static_cast(features_->count()))\n : cents.size();\n if (centroid_count == 0) {\n LOG_ERROR(\"The count of cluster is unknown.\");\n return IndexError_NoReady;\n }\n ailego::NumericalInnerProductKmeans algorithm(\n centroid_count, meta_.dimension(), true);\n\n // mount features into algorithm\n auto features_count = features_->count();\n auto dim = meta_.dimension();\n\n algorithm.feature_matrix_reserve(features_count);\n\n for (size_t i = 0; i < features_count; ++i) {\n auto vec = reinterpret_cast(features_->element(i));\n algorithm.append(vec, dim);\n }\n\n if (!cents.empty()) {\n auto centroids = algorithm.mutable_centroids();\n centroids->reserve(cents.size());\n for (const auto &it : cents) {\n centroids->append(reinterpret_cast(it.feature()),\n meta_.dimension());\n }\n } else {\n ailego::Kmc2CentroidsGenerator<\n ailego::NumericalInnerProductKmeans,\n IndexThreads>\n g;\n g.set_chain_length(markov_chain_length_);\n g.set_assumption_free(assumption_free_);\n algorithm.init_centroids(*threads, g);\n }\n\n double cost = 0.0;\n\n for (uint32_t i = 0; i < max_iterations_; ++i) {\n double old_cost, new_epsilon;\n old_cost = cost;\n\n bool result = algorithm.cluster_once(*threads, &cost);\n if (result != true) {\n LOG_ERROR(\"(%u) Failed to cluster.\", i + 1);\n return -1;\n }\n\n new_epsilon = std::abs(cost - old_cost);\n LOG_DEBUG(\"(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> %f = %f\",\n i, algorithm.centroids().count(), features_->count(),\n (size_t)stamp.milli_seconds(), old_cost, cost, new_epsilon);\n stamp.reset();\n\n if (new_epsilon < epsilon_) {\n break;\n }\n }\n\n // update_centroids(cents);\n update_centroids(cents, algorithm);\n\n // Purge the empty centroids\n PurgeCentroids(cents, purge_empty_);\n return 0;\n}\n\n/*! Nibble Inner Product K-Means cluster algorithm\n */\ntemplate \nclass NibbleInnerProductKmeansAlgorithm : public OptKmeansAlgorithm {\n public:\n //! Type of value\n using ValueType = typename std::remove_cv::type;\n\n // Check supporting type\n static_assert(ailego::IsArithmetic::value,\n \"ValueType must be arithmetic\");\n\n //! Constructor\n NibbleInnerProductKmeansAlgorithm(void) {}\n\n //! Destructor\n virtual ~NibbleInnerProductKmeansAlgorithm(void) {}\n\n //! Cluster\n virtual int cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n protected:\n //! update centroids\n void update_centroids(\n IndexCluster::CentroidList ¢s,\n const ailego::NibbleInnerProductKmeans &algorithm);\n};\n\ntemplate \nvoid NibbleInnerProductKmeansAlgorithm::update_centroids(\n IndexCluster::CentroidList ¢s,\n const ailego::NibbleInnerProductKmeans &algorithm) {\n this->init_centroids(algorithm.centroids().count(), ¢s);\n for (size_t i = 0; i < cents.size(); ++i) {\n IndexCluster::Centroid *centroid = &(cents.at(i));\n centroid->set_score(algorithm.context().clusters()[i].cost());\n centroid->set_follows(algorithm.context().clusters()[i].count());\n centroid->set_feature(algorithm.centroids()[i], (meta_.dimension() >> 1));\n }\n}\n\ntemplate \nint NibbleInnerProductKmeansAlgorithm::cluster(\n IndexThreads::Pointer threads, IndexCluster::CentroidList ¢s) {\n ailego::ElapsedTime stamp;\n\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n // get cluster algorithm\n size_t centroid_count =\n cents.empty()\n ? std::min(cluster_count_, static_cast(features_->count()))\n : cents.size();\n if (centroid_count == 0) {\n LOG_ERROR(\"The count of cluster is unknown.\");\n return IndexError_NoReady;\n }\n ailego::NibbleInnerProductKmeans algorithm(\n centroid_count, meta_.dimension());\n\n // mount features into algorithm\n auto features_count = features_->count();\n auto dim = meta_.dimension();\n for (size_t i = 0; i < features_count; ++i) {\n auto vec = reinterpret_cast::type *>(\n features_->element(i));\n algorithm.append(vec, dim);\n }\n\n if (!cents.empty()) {\n auto centroids = algorithm.mutable_centroids();\n centroids->reserve(cents.size());\n for (const auto &it : cents) {\n centroids->append(\n reinterpret_cast::type *>(\n it.feature()),\n size_t(meta_.dimension()));\n }\n } else {\n ailego::Kmc2CentroidsGenerator<\n ailego::NibbleInnerProductKmeans, IndexThreads>\n g;\n g.set_chain_length(markov_chain_length_);\n g.set_assumption_free(assumption_free_);\n algorithm.init_centroids(*threads, g);\n }\n\n double cost = 0.0;\n\n for (uint32_t i = 0; i < max_iterations_; ++i) {\n double old_cost, new_epsilon;\n old_cost = cost;\n\n bool result = algorithm.cluster_once(*threads, &cost);\n if (result != true) {\n LOG_ERROR(\"(%u) Failed to cluster.\", i + 1);\n return -1;\n }\n\n new_epsilon = std::abs(cost - old_cost);\n LOG_DEBUG(\n \"(%u) Updated %zu Clusters, %zu Features: %zu ms, %f -> \"\n \"%f = %f\",\n i, algorithm.centroids().count(), features_->count(),\n (size_t)stamp.milli_seconds(), old_cost, cost, new_epsilon);\n stamp.reset();\n\n if (new_epsilon < epsilon_) {\n break;\n }\n }\n\n // update centroids\n update_centroids(cents, algorithm);\n\n // Purge the empty centroids\n PurgeCentroids(cents, purge_empty_);\n return 0;\n}\n\n/*! Kmeans Cluster\n */\nclass OptKmeansCluster : public IndexCluster {\n public:\n //! Constructor\n OptKmeansCluster(void) {}\n\n //! Destructor\n virtual ~OptKmeansCluster(void) {}\n\n //! Initialize Cluster\n virtual int init(const IndexMeta &meta, const ailego::Params ¶ms);\n\n //! Cleanup Cluster\n virtual int cleanup(void);\n\n //! Reset Cluster\n virtual int reset(void);\n\n //! Update Cluster\n virtual int update(const ailego::Params ¶ms);\n\n //! Suggest dividing to K clusters\n virtual void suggest(uint32_t k);\n\n //! Mount features\n virtual int mount(IndexFeatures::Pointer feats);\n\n //! Cluster\n virtual int cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n //! Classify\n virtual int classify(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n //! Label\n virtual int label(IndexThreads::Pointer threads,\n const IndexCluster::CentroidList ¢s,\n std::vector *out);\n\n protected:\n //! Members\n IndexCluster::Pointer algorithm_{};\n};\n\n//! Cluster\nint OptKmeansCluster::cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n return algorithm_->cluster(std::move(threads), cents);\n}\n\n//! Classify\nint OptKmeansCluster::classify(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n return algorithm_->classify(std::move(threads), cents);\n}\n\n//! Label\nint OptKmeansCluster::label(IndexThreads::Pointer threads,\n const IndexCluster::CentroidList ¢s,\n std::vector *out) {\n return algorithm_->label(std::move(threads), cents, out);\n}\n\n//! Update Cluster\nint OptKmeansCluster::update(const ailego::Params ¶ms) {\n return algorithm_->update(params);\n}\n\n//! Reset Cluster\nint OptKmeansCluster::reset(void) {\n return algorithm_->reset();\n}\n\n//! Cleanup Cluster\nint OptKmeansCluster::cleanup(void) {\n return algorithm_->cleanup();\n}\n\n//! Suggest dividing to K clusters\nvoid OptKmeansCluster::suggest(uint32_t k) {\n algorithm_->suggest(k);\n}\n\nint OptKmeansCluster::mount(IndexFeatures::Pointer feats) {\n return algorithm_->mount(feats);\n}\n\nint OptKmeansCluster::init(const IndexMeta &meta,\n const ailego::Params ¶ms) {\n auto type_ = meta.data_type();\n\n if (meta.metric_name() == \"InnerProduct\" || meta.metric_name() == \"Cosine\") {\n switch (type_) {\n case IndexMeta::DataType::DT_FP16: {\n algorithm_.reset(\n new (std::nothrow)\n NumericalInnerProductKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_FP32: {\n algorithm_.reset(new (std::nothrow)\n NumericalInnerProductKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_FP64: {\n algorithm_.reset(new (std::nothrow)\n NumericalInnerProductKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_INT8: {\n algorithm_.reset(new (std::nothrow)\n NumericalInnerProductKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_INT16: {\n algorithm_.reset(new (std::nothrow)\n NumericalInnerProductKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_INT4: {\n algorithm_.reset(new (std::nothrow)\n NibbleInnerProductKmeansAlgorithm);\n break;\n }\n default: {\n LOG_ERROR(\"Unsupported feature types %d.\", type_);\n return IndexError_Mismatch;\n }\n }\n } else {\n switch (type_) {\n case IndexMeta::DataType::DT_FP16: {\n algorithm_.reset(new (std::nothrow)\n NumericalKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_FP32: {\n algorithm_.reset(new (std::nothrow) NumericalKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_FP64: {\n algorithm_.reset(new (std::nothrow) NumericalKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_INT8: {\n algorithm_.reset(new (std::nothrow) NumericalKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_INT16: {\n algorithm_.reset(new (std::nothrow) NumericalKmeansAlgorithm);\n break;\n }\n case IndexMeta::DataType::DT_INT4: {\n algorithm_.reset(new (std::nothrow) NibbleKmeansAlgorithm);\n break;\n }\n // TODO\n case IndexMeta::DataType::DT_BINARY32: {\n algorithm_.reset(new (std::nothrow) BinaryKmeansAlgorithm);\n break;\n }\n\n#if defined(AILEGO_M64)\n case IndexMeta::DataType::DT_BINARY64: {\n algorithm_.reset(new (std::nothrow) BinaryKmeansAlgorithm);\n break;\n }\n#endif // AILEGO_M64\n\n default: {\n LOG_ERROR(\"Unsupported feature types %d.\", type_);\n return IndexError_Mismatch;\n }\n }\n }\n\n algorithm_->init(meta, params);\n\n return 0;\n}\n\nINDEX_FACTORY_REGISTER_CLUSTER(OptKmeansCluster);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/cluster/seeker.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\nclass Seeker {\n public:\n struct Document {\n uint32_t index;\n float score;\n\n //! Constructor\n Document(void) : index(0), score(0.0f) {}\n\n //! Constructor\n Document(uint32_t i, float v) : index(i), score(v) {}\n\n //! Constructor\n Document(const Document &rhs) : index(rhs.index), score(rhs.score) {}\n\n //! Assignment\n Document &operator=(const Document &rhs) {\n index = rhs.index;\n score = rhs.score;\n return *this;\n }\n\n //! Less than\n bool operator<(const Document &rhs) const {\n return (this->score < rhs.score);\n }\n\n //! Greater than\n bool operator>(const Document &rhs) const {\n return (this->score > rhs.score);\n }\n };\n\n public:\n //! Destructor\n virtual ~Seeker(void) {}\n\n virtual int init(const IndexMeta &meta) = 0;\n\n virtual int cleanup(void) = 0;\n\n virtual int reset(void) = 0;\n\n virtual int mount(IndexFeatures::Pointer feats) = 0;\n\n virtual int seek(const void *query, size_t len, Document *out) = 0;\n\n virtual IndexFeatures::Pointer original(void) const = 0;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/cluster/stratified_cluster.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n#include \n#include \n#include \n#include \"cluster_params.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Stratified Cluster\n */\nclass StratifiedCluster : public IndexCluster {\n public:\n //! Constructor\n StratifiedCluster(void) {}\n\n //! Destructor\n virtual ~StratifiedCluster(void) {}\n\n //! Initialize Cluster\n virtual int init(const IndexMeta &meta, const ailego::Params ¶ms) {\n meta_ = meta;\n this->update_params(params);\n return 0;\n }\n\n //! Cleanup Cluster\n virtual int cleanup(void) {\n features_.reset();\n return 0;\n }\n\n //! Reset Cluster\n virtual int reset(void) {\n features_.reset();\n return 0;\n }\n\n //! Update Cluster\n virtual int update(const ailego::Params ¶ms) {\n this->update_params(params);\n return 0;\n }\n\n //! Suggest dividing to K clusters\n virtual void suggest(uint32_t k) {\n cluster_count_ = k;\n }\n\n //! Mount features\n virtual int mount(IndexFeatures::Pointer feats) {\n if (!feats) {\n return IndexError_InvalidArgument;\n }\n if (!feats->is_matched(meta_)) {\n return IndexError_Mismatch;\n }\n features_ = std::move(feats);\n return 0;\n }\n\n //! Cluster\n virtual int cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n //! Classify\n virtual int classify(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s);\n\n //! Label\n virtual int label(IndexThreads::Pointer threads,\n const IndexCluster::CentroidList ¢s,\n std::vector *out);\n\n protected:\n //! Test if it is valid\n bool is_valid(void) const {\n if (!features_ || !features_->count()) {\n return false;\n }\n return true;\n }\n\n //! Update parameters\n void update_params(const ailego::Params ¶ms);\n\n //! Check Centroids\n bool check_centroids(const IndexCluster::CentroidList ¢s);\n\n //! Initialize Sub Clusters\n int init_sub_clusters(IndexCluster::Pointer *first,\n IndexCluster::Pointer *second);\n\n //! Initialize First Cluster\n int init_first_cluster(IndexCluster::Pointer *first);\n\n //! Initialize Second Cluster\n int init_second_cluster(IndexCluster::Pointer *second,\n IndexFeatures::Pointer features);\n\n private:\n //! Members\n IndexMeta meta_{};\n IndexFeatures::Pointer features_{};\n uint32_t cluster_count_{0u};\n uint32_t thread_count_{0u};\n uint32_t first_cluster_count_{0u};\n uint32_t second_cluster_count_{0u};\n bool auto_tuning_{false};\n std::string first_cluster_class_{\"OptKmeansCluster\"};\n std::string second_cluster_class_{\"OptKmeansCluster\"};\n ailego::Params first_cluster_params_{};\n ailego::Params second_cluster_params_{};\n\n // TODO: Maybe optimize later\n uint32_t second_threads_count_{10u}; // todo\n};\n\nint StratifiedCluster::cluster(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n IndexCluster::Pointer first_cluster;\n int result = init_first_cluster(&first_cluster);\n if (result != 0) {\n LOG_ERROR(\"Failed to initialize the first cluster.\");\n return result;\n }\n\n if (first_cluster_count_) {\n first_cluster->suggest(first_cluster_count_);\n }\n\n // The first clustering\n LOG_DEBUG(\"Clustering with first cluster: %s.\", first_cluster_class_.c_str());\n result = first_cluster->cluster(threads, cents);\n if (result != 0) {\n LOG_ERROR(\"Failed to cluster in first cluster: %s.\",\n first_cluster_class_.c_str());\n return result;\n }\n\n result = first_cluster->classify(threads, cents);\n if (result != 0) {\n LOG_ERROR(\"Failed to classify in first cluster: %s.\",\n first_cluster_class_.c_str());\n return result;\n }\n\n // Cleanup for saving memory\n first_cluster.reset();\n\n // Calculate the total cluster count\n uint32_t total_cluster_count = cents.size() * second_cluster_count_;\n if (cluster_count_) {\n total_cluster_count = cluster_count_;\n }\n\n // Use thread_threads cluster instead\n uint32_t tail_threads = threads->count() % second_threads_count_;\n std::vector> threads_cluster;\n\n // TODO: reusing threads pool?\n // Incase the threads count less than second threads count\n if (threads->count() / second_threads_count_ == 0) {\n for (size_t threads_idx = 0; threads_idx < tail_threads; threads_idx++) {\n std::shared_ptr curr_threads =\n std::make_shared(1, false);\n threads_cluster.push_back(curr_threads);\n }\n } else {\n for (size_t threads_idx = 0; threads_idx < second_threads_count_;\n threads_idx++) {\n uint32_t curr_threads_count = threads->count() / second_threads_count_;\n if (threads_idx >= second_threads_count_ - tail_threads) {\n curr_threads_count++;\n }\n std::shared_ptr curr_threads =\n std::make_shared(curr_threads_count, false);\n threads_cluster.push_back(curr_threads);\n }\n }\n\n auto task_group = threads->make_group();\n // The second clustering\n for (size_t i = 0; i < cents.size(); ++i) {\n if (cents[i].similars().empty()) {\n continue;\n }\n\n IndexThreads::Pointer &curr_threads =\n threads_cluster[i % (threads_cluster.size())];\n\n task_group->submit(ailego::Closure::New(\n [this, &curr_threads, &total_cluster_count, ¢s](size_t index) {\n auto &it = cents[index];\n IndexCluster::Pointer second_cluster;\n std::shared_ptr features =\n std::make_shared(\n meta_, it.similars().data(), it.similars().size());\n\n int ret = this->init_second_cluster(&second_cluster, features);\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize the second cluster.\");\n return;\n }\n\n if (auto_tuning_) {\n if (total_cluster_count) {\n double factor = static_cast(it.similars().size()) /\n static_cast(this->features_->count());\n second_cluster->suggest(\n std::max(static_cast(\n std::floor(total_cluster_count * factor)),\n 1u));\n }\n } else if (second_cluster_count_) {\n second_cluster->suggest(second_cluster_count_);\n }\n\n LOG_DEBUG(\"Clustering with second cluster: %s.\",\n second_cluster_class_.c_str());\n ret = second_cluster->cluster(curr_threads, *(it.mutable_subitems()));\n if (ret != 0) {\n LOG_ERROR(\"Failed to cluster in second cluster: %s.\",\n second_cluster_class_.c_str());\n }\n },\n i));\n }\n task_group->wait_finish();\n return 0;\n}\n\nint StratifiedCluster::classify(IndexThreads::Pointer threads,\n IndexCluster::CentroidList ¢s) {\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (cents.empty()) {\n LOG_ERROR(\"The input centroid's list is empty.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n IndexCluster::Pointer first_cluster, second_cluster;\n int result = init_sub_clusters(&first_cluster, &second_cluster);\n if (result != 0) {\n LOG_ERROR(\"Failed to initialize the subclusters.\");\n return result;\n }\n\n // The first classifying\n result = first_cluster->classify(threads, cents);\n if (result != 0) {\n LOG_ERROR(\"Failed to classify in first cluster: %s.\",\n first_cluster_class_.c_str());\n return result;\n }\n\n // Cleanup for saving memory\n first_cluster.reset();\n\n std::shared_ptr shell =\n std::make_shared(meta_);\n\n // The second classifying\n for (IndexCluster::Centroid &it : cents) {\n const auto &feats = it.similars();\n\n if (feats.empty()) {\n continue;\n }\n\n shell->mount(feats.data(), feats.size());\n result = second_cluster->mount(shell);\n if (result != 0) {\n LOG_ERROR(\"Failed to mount features for second cluster: %s.\",\n second_cluster_class_.c_str());\n return result;\n }\n\n result = second_cluster->classify(threads, *it.mutable_subitems());\n if (result != 0) {\n LOG_ERROR(\"Failed to classify in second cluster: %s.\",\n second_cluster_class_.c_str());\n return result;\n }\n }\n return 0;\n}\n\nint StratifiedCluster::label(IndexThreads::Pointer threads,\n const IndexCluster::CentroidList ¢s,\n std::vector *out) {\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n if (cents.empty()) {\n LOG_ERROR(\"The input centroid's list is empty.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->check_centroids(cents)) {\n LOG_ERROR(\"The input centroid's list includes some invalid centroids.\");\n return IndexError_InvalidArgument;\n }\n\n if (!this->is_valid()) {\n LOG_ERROR(\"The cluster is not ready.\");\n return IndexError_NoReady;\n }\n\n IndexCluster::Pointer first_cluster;\n int result = init_first_cluster(&first_cluster);\n if (result != 0) {\n LOG_ERROR(\"Failed to initialize the subclusters.\");\n return result;\n }\n\n result = first_cluster->label(threads, cents, out);\n if (result != 0) {\n LOG_ERROR(\"Failed to label in first cluster: %s.\",\n first_cluster_class_.c_str());\n return result;\n }\n return 0;\n}\n\nvoid StratifiedCluster::update_params(const ailego::Params ¶ms) {\n params.get(GENERAL_THREAD_COUNT, &thread_count_);\n params.get(GENERAL_CLUSTER_COUNT, &cluster_count_);\n params.get(STRATIFIED_CLUSTER_COUNT, &cluster_count_);\n params.get(STRATIFIED_CLUSTER_FIRST_COUNT, &first_cluster_count_);\n params.get(STRATIFIED_CLUSTER_SECOND_COUNT, &second_cluster_count_);\n params.get(STRATIFIED_CLUSTER_FIRST_CLASS, &first_cluster_class_);\n params.get(STRATIFIED_CLUSTER_SECOND_CLASS, &second_cluster_class_);\n params.get(STRATIFIED_CLUSTER_FIRST_PARAMS, &first_cluster_params_);\n params.get(STRATIFIED_CLUSTER_SECOND_PARAMS, &second_cluster_params_);\n params.get(STRATIFIED_CLUSTER_AUTO_TUNING, &auto_tuning_);\n params.get(STRATIFIED_CLUSTER_SECOND_POOL_COUNT, &second_threads_count_);\n}\n\nbool StratifiedCluster::check_centroids(\n const IndexCluster::CentroidList ¢s) {\n for (const auto &it : cents) {\n if (it.size() != meta_.element_size()) {\n return false;\n }\n }\n return true;\n}\n\nint StratifiedCluster::init_sub_clusters(IndexCluster::Pointer *first,\n IndexCluster::Pointer *second) {\n IndexCluster::Pointer first_cluster =\n IndexFactory::CreateCluster(first_cluster_class_);\n\n if (!first_cluster) {\n LOG_ERROR(\"Failed to create first cluster: %s.\",\n first_cluster_class_.c_str());\n return IndexError_NoExist;\n }\n\n IndexCluster::Pointer second_cluster =\n IndexFactory::CreateCluster(second_cluster_class_);\n\n if (!second_cluster) {\n LOG_ERROR(\"Failed to create second cluster: %s.\",\n first_cluster_class_.c_str());\n return IndexError_NoExist;\n }\n\n int result = first_cluster->init(meta_, first_cluster_params_);\n if (result != 0) {\n LOG_ERROR(\"Failed to initialize first cluster: %s.\",\n first_cluster_class_.c_str());\n return result;\n }\n\n result = second_cluster->init(meta_, second_cluster_params_);\n if (result != 0) {\n LOG_ERROR(\"Failed to initialize second cluster: %s.\",\n second_cluster_class_.c_str());\n return result;\n }\n\n result = first_cluster->mount(features_);\n if (result != 0) {\n LOG_ERROR(\"Failed to mount features for first cluster: %s.\",\n first_cluster_class_.c_str());\n return result;\n }\n\n *first = std::move(first_cluster);\n *second = std::move(second_cluster);\n return 0;\n}\n\nint StratifiedCluster::init_first_cluster(IndexCluster::Pointer *first) {\n IndexCluster::Pointer first_cluster =\n IndexFactory::CreateCluster(first_cluster_class_);\n\n if (!first_cluster) {\n LOG_ERROR(\"Failed to create first cluster: %s.\",\n first_cluster_class_.c_str());\n return IndexError_NoExist;\n }\n\n int result = first_cluster->init(meta_, first_cluster_params_);\n if (result != 0) {\n LOG_ERROR(\"Failed to initialize first cluster: %s.\",\n first_cluster_class_.c_str());\n return result;\n }\n\n result = first_cluster->mount(features_);\n if (result != 0) {\n LOG_ERROR(\"Failed to mount features for first cluster: %s.\",\n first_cluster_class_.c_str());\n return result;\n }\n\n *first = std::move(first_cluster);\n return 0;\n}\n\nint StratifiedCluster::init_second_cluster(IndexCluster::Pointer *second,\n IndexFeatures::Pointer features) {\n IndexCluster::Pointer second_cluster =\n IndexFactory::CreateCluster(second_cluster_class_);\n\n if (!second_cluster) {\n LOG_ERROR(\"Failed to create second cluster: %s.\",\n second_cluster_class_.c_str());\n return IndexError_NoExist;\n }\n\n int result = second_cluster->init(meta_, second_cluster_params_);\n if (result != 0) {\n LOG_ERROR(\"Failed to initialize second cluster: %s.\",\n second_cluster_class_.c_str());\n return result;\n }\n\n result = second_cluster->mount(features);\n if (result != 0) {\n LOG_ERROR(\"Failed to mount features for second cluster: %s.\",\n second_cluster_class_.c_str());\n return result;\n }\n\n *second = std::move(second_cluster);\n return 0;\n}\n\nINDEX_FACTORY_REGISTER_CLUSTER(StratifiedCluster);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/cluster/stratified_cluster_trainer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"stratified_cluster_trainer.h\"\n#include \n#include \n#include \n#include \n#include \n#include \"cluster_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nconst std::string StratifiedClusterTrainer::SEP_TOKEN = \"*\";\nconst std::string StratifiedClusterTrainer::DEFAULT_CLUSTER_CLASS =\n \"OptKmeansCluster\";\n\nint StratifiedClusterTrainer::init_params(const ailego::Params ¶ms) {\n params.get(STRATIFIED_TRAINER_SAMPLE_COUNT, &sample_count_);\n params.get(STRATIFIED_TRAINER_SAMPLE_RATIO, &sample_ratio_);\n params.get(STRATIFIED_TRAINER_THREAD_COUNT, &thread_count_);\n cluster_auto_tuning_ = params.get_as_bool(STRATIFIED_TRAINER_AUTOAUNE);\n\n std::string centroids_num =\n params.get_as_string(STRATIFIED_TRAINER_CLUSTER_COUNT);\n if (!centroids_num.empty()) {\n ailego::StringHelper::Split(centroids_num, SEP_TOKEN, ¢roid_num_vec_);\n for (size_t i = 0; i < centroid_num_vec_.size(); ++i) {\n if (centroid_num_vec_[i] == 0) {\n LOG_ERROR(\"Invalid centroid num %s\", centroids_num.c_str());\n return IndexError_InvalidArgument;\n }\n }\n } else {\n LOG_ERROR(\"No centroids_num configed.\");\n return IndexError_InvalidArgument;\n }\n\n size_t level_cnt = centroid_num_vec_.size();\n for (size_t i = 1; i <= level_cnt; ++i) {\n std::string level_params_key =\n STRATIFIED_TRAINER_PARAMS_IN_LEVEL_PREFIX + std::to_string(i);\n ailego::Params level_params;\n params.get(level_params_key, &level_params);\n cluster_params_.push_back(level_params);\n }\n\n std::string cluster_class(DEFAULT_CLUSTER_CLASS);\n params.get(STRATIFIED_TRAINER_CLASS_NAME, &cluster_class);\n ailego::StringHelper::Split(cluster_class, SEP_TOKEN, &cluster_class_);\n if (cluster_class_.size() == 1) {\n // repeat the cluster class to level_cnt\n for (size_t i = 1; i < level_cnt; ++i) {\n cluster_class_.push_back(cluster_class_[0]);\n }\n } else if (cluster_class_.size() != level_cnt) {\n LOG_ERROR(\"Cluster class should be equal to level count\");\n return IndexError_InvalidArgument;\n }\n return 0;\n}\n\nint StratifiedClusterTrainer::init(const IndexMeta &index_meta,\n const ailego::Params ¶ms) {\n int err = init_params(params);\n if (err != 0) {\n LOG_ERROR(\"init params failed, errno:%d,%s\", err, IndexError::What(err));\n return err;\n }\n\n meta_ = index_meta;\n ailego::Params cluster_params;\n if (centroid_num_vec_.size() == 0) {\n LOG_ERROR(\"invalid centroid num\");\n return IndexError_InvalidArgument;\n } else if (centroid_num_vec_.size() == 1) {\n // one level clustering\n class_name_ = cluster_class_[0];\n cluster_params = cluster_params_[0];\n suggest_centriod_cnt_ = centroid_num_vec_[0];\n } else if (centroid_num_vec_.size() == 2) {\n // cluster level > 1\n class_name_ = \"StratifiedCluster\";\n int level_cnt = centroid_num_vec_.size();\n cluster_params.set(STRATIFIED_CLUSTER_FIRST_CLASS,\n cluster_class_[level_cnt - 2]);\n cluster_params.set(STRATIFIED_CLUSTER_SECOND_CLASS,\n cluster_class_[level_cnt - 1]);\n cluster_params.set(STRATIFIED_CLUSTER_FIRST_COUNT,\n centroid_num_vec_[level_cnt - 2]);\n cluster_params.set(STRATIFIED_CLUSTER_SECOND_COUNT,\n centroid_num_vec_[level_cnt - 1]);\n cluster_params.set(STRATIFIED_CLUSTER_FIRST_PARAMS,\n cluster_params_[level_cnt - 2]);\n cluster_params.set(STRATIFIED_CLUSTER_SECOND_PARAMS,\n cluster_params_[level_cnt - 1]);\n cluster_params.set(STRATIFIED_CLUSTER_AUTO_TUNING, cluster_auto_tuning_);\n suggest_centriod_cnt_ =\n centroid_num_vec_[level_cnt - 1] * centroid_num_vec_[level_cnt - 2];\n } else {\n LOG_ERROR(\"Unsupported more than 2 level clustering.\");\n return IndexError_Unsupported;\n }\n\n cluster_ = IndexFactory::CreateCluster(class_name_);\n if (!cluster_) {\n LOG_ERROR(\"Failed to create cluster[%s]\", class_name_.c_str());\n return IndexError_InvalidArgument;\n }\n int result = cluster_->init(meta_, cluster_params);\n if (result != 0) {\n LOG_ERROR(\"Failed to initialize of cluster[%s], error: %d, %s\",\n class_name_.c_str(), result, IndexError::What(result));\n return result;\n }\n if (suggest_centriod_cnt_ > 0) {\n cluster_->suggest(suggest_centriod_cnt_);\n }\n\n return 0;\n}\n\nint StratifiedClusterTrainer::cleanup(void) {\n cluster_ = nullptr;\n centroids_.clear();\n return 0;\n}\n\nint StratifiedClusterTrainer::train(IndexThreads::Pointer threads,\n IndexHolder::Pointer holder) {\n ailego::ElapsedTime timer;\n if (!holder) {\n return IndexError_InvalidArgument;\n }\n if (!holder->is_matched(meta_)) {\n return IndexError_Mismatch;\n }\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n\n size_t train_sample_count = std::max(\n sample_count_, static_cast(sample_ratio_ * holder->count()));\n\n IndexFeatures::Pointer features;\n if (train_sample_count > 0) {\n LOG_INFO(\n \"Train sampling, SampleCount=%u, SampleRatio=%f, HolderCount=%lu, \"\n \"TrainCount=%lu\",\n sample_count_, sample_ratio_, holder->count(), train_sample_count);\n\n auto sampler = std::make_shared>(\n meta_, train_sample_count);\n size_t pre_reserve = train_sample_count < holder->count()\n ? train_sample_count\n : holder->count();\n sampler->reserve(pre_reserve);\n for (auto iter = holder->create_iterator(); iter && iter->is_valid();\n iter->next()) {\n sampler->emplace(iter->data());\n }\n features = sampler;\n stats_.set_trained_count(train_sample_count);\n } else {\n LOG_INFO(\n \"Do no sampling, SampleCount=%u, SampleRatio=%f, \"\n \"HolderCount=%lu, TrainCount=%lu\",\n sample_count_, sample_ratio_, holder->count(), holder->count());\n\n auto no_sampler = std::make_shared(meta_);\n for (auto iter = holder->create_iterator(); iter && iter->is_valid();\n iter->next()) {\n no_sampler->emplace(iter->data());\n }\n\n features = no_sampler;\n stats_.set_trained_count(holder->count());\n }\n stats_.set_discarded_count(0);\n\n // Holder is not needed, cleanup it.\n holder.reset();\n\n int result = cluster_->mount(features);\n if (result != 0) {\n LOG_ERROR(\"Failed to mount features of cluster[%s], error: %d, %s\",\n class_name_.c_str(), result, IndexError::What(result));\n return result;\n }\n\n centroids_.clear();\n result = cluster_->cluster(std::move(threads), centroids_);\n if (result != 0) {\n LOG_ERROR(\"Failed to cluster features of cluster[%s], error: %d, %s\",\n class_name_.c_str(), result, IndexError::What(result));\n return result;\n }\n\n // check build result\n std::vector level_size;\n std::function\n cal_centroid_cnt =\n [&cal_centroid_cnt, &level_size](\n const IndexCluster::CentroidList ¢s, size_t level) {\n if (level > level_size.size()) {\n level_size.resize(level);\n }\n level_size[level - 1] += cents.size();\n for (const auto &it : cents) {\n if (!it.subitems().empty()) {\n cal_centroid_cnt(it.subitems(), level + 1);\n }\n }\n };\n cal_centroid_cnt(centroids_, 1);\n\n size_t centroids_num = level_size[level_size.size() - 1];\n if (centroids_num > suggest_centriod_cnt_) {\n LOG_WARN(\n \"Built centroid(%zd level) count[%zd] bigger than expected \"\n \"count[%d]\",\n level_size.size(), centroids_num, suggest_centriod_cnt_);\n } else {\n LOG_INFO(\"Built centroid(%zd level) count[%zd], expected count[%d]\",\n level_size.size(), centroids_num, suggest_centriod_cnt_);\n }\n\n stats_.set_trained_costtime(timer.milli_seconds());\n\n return 0;\n}\n\nint StratifiedClusterTrainer::load(IndexStorage::Pointer cntr) {\n if (!cntr) {\n LOG_ERROR(\"IndexStorage is nullptr.\");\n return IndexError_InvalidArgument;\n }\n std::shared_ptr bundle =\n std::make_shared();\n if (!bundle) {\n LOG_ERROR(\"New MemoryInndexBundle failed.\");\n return IndexError_NoMemory;\n }\n\n auto results = cntr->get_all();\n for (auto &it : results) {\n IndexStorage::Segment::Pointer &seg = it.second;\n if (!seg) {\n LOG_ERROR(\"Get Segment %s failed.\", it.first.c_str());\n return IndexError_InvalidArgument;\n }\n size_t data_size = seg->data_size();\n const void *data = nullptr;\n size_t actual_size = seg->read(0, &data, data_size);\n if (actual_size != data_size) {\n LOG_ERROR(\"Read data failed expect %zu, actual %zu.\", data_size,\n actual_size);\n return IndexError_ReadData;\n }\n bundle->set(it.first, data, data_size);\n }\n\n int result = IndexHelper::DeserializeFromStorage(cntr.get(), &meta_);\n if (result != 0) {\n LOG_ERROR(\"Failed to deserialize meta from container\");\n return result;\n }\n\n result = IndexCluster::Deserialize(meta_, std::move(bundle), ¢roids_);\n if (result != 0) {\n LOG_ERROR(\"Failed to deserialize index: %d\", result);\n return result;\n }\n return 0;\n}\n\nint StratifiedClusterTrainer::dump(const IndexDumper::Pointer &dumper) {\n IndexBundle::Pointer bundle;\n int result = IndexCluster::Serialize(meta_, centroids_, &bundle);\n if (result != 0) {\n LOG_ERROR(\"IndexCluster Serialize failed with ret %d.\", result);\n return result;\n }\n\n result = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (result != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return result;\n }\n\n for (const auto &it : bundle->all()) {\n size_t data_size = it.second.size();\n result = dumper->append(it.first, data_size, 0, 0);\n if (result != 0) {\n LOG_ERROR(\"Dumper append meta %s %zu failed.\", it.first.c_str(),\n data_size);\n return IndexError_PackIndex;\n }\n size_t actual_size = dumper->write(it.second.buffer(), data_size);\n if (actual_size != data_size) {\n LOG_ERROR(\"Dumper segment %s expect %zu, actual %zu.\", it.first.c_str(),\n data_size, actual_size);\n return IndexError_PackIndex;\n }\n }\n return 0;\n}\n\nconst IndexMeta &StratifiedClusterTrainer::meta(void) const {\n return meta_;\n}\n\nconst IndexTrainer::Stats &StratifiedClusterTrainer::stats(void) const {\n return stats_;\n}\n\nIndexBundle::Pointer StratifiedClusterTrainer::indexes(void) const {\n IndexBundle::Pointer bundle;\n IndexCluster::Serialize(meta_, centroids_, &bundle);\n return bundle;\n}\n\n//! Register Cluster Trainer in Factory\nINDEX_FACTORY_REGISTER_TRAINER(StratifiedClusterTrainer);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/cluster/stratified_cluster_trainer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\n/*! Cluster Trainer\n */\nclass StratifiedClusterTrainer : public IndexTrainer {\n public:\n typedef std::shared_ptr Pointer;\n\n //! Constructor\n StratifiedClusterTrainer(void) {}\n\n //! Destructor\n ~StratifiedClusterTrainer(void) {}\n\n protected:\n //! Initialize Trainer\n virtual int init(const IndexMeta &meta, const ailego::Params ¶ms);\n\n //! Cleanup Trainer\n virtual int cleanup(void);\n\n //! Train the data\n virtual int train(IndexThreads::Pointer threads, IndexHolder::Pointer holder);\n\n //! Load index from file path or dir\n virtual int load(IndexStorage::Pointer cntr);\n\n //! Dump index into file path or dir\n virtual int dump(const IndexDumper::Pointer &dumper);\n\n //! Retrieve Index Meta\n virtual const IndexMeta &meta(void) const;\n\n //! Retrieve statistics\n virtual const IndexTrainer::Stats &stats(void) const;\n\n //! Retrieve the output indexes\n virtual IndexBundle::Pointer indexes(void) const;\n\n private:\n int init_params(const ailego::Params ¶ms);\n\n private:\n IndexMeta meta_{};\n uint32_t sample_count_{0u};\n float sample_ratio_{0.0};\n uint32_t thread_count_{0u};\n bool cluster_auto_tuning_{false};\n IndexCluster::Pointer cluster_{};\n IndexCluster::CentroidList centroids_{};\n\n uint32_t suggest_centriod_cnt_{0u};\n std::string class_name_;\n std::vector cluster_class_;\n std::vector centroid_num_vec_;\n std::vector cluster_params_;\n IndexTrainer::Stats stats_{};\n\n private:\n static const std::string SEP_TOKEN;\n static const std::string DEFAULT_CLUSTER_CLASS;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/cluster/vector_mean.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\n/*! Vector Mean\n */\nstruct VectorMean {\n //! Destructor\n virtual ~VectorMean(void) {}\n\n //! Reset accumulator\n virtual void reset(void) = 0;\n\n //! Plus a vector\n virtual bool plus(const void *vec, size_t len) = 0;\n\n //! Retrieve the mean of vectors\n virtual bool mean(void *out, size_t len) const = 0;\n\n //! Retrieve the mean of vectors\n virtual void mean(std::string *out) const = 0;\n\n //! Merge another vector mean\n virtual bool merge(const VectorMean &rhs) = 0;\n\n //! Retrieve the count of vectors\n virtual size_t count(void) const = 0;\n\n //! Retrieve the dimension of vectors\n virtual size_t dimension(void) const = 0;\n};\n\n/*! Vector Mean Array\n */\nstruct VectorMeanArray {\n //! Destructor\n virtual ~VectorMeanArray(void) {}\n\n //! Operator []\n VectorMean &operator[](size_t i) {\n return this->at(i);\n }\n\n //! Operator []\n const VectorMean &operator[](size_t i) const {\n return this->at(i);\n }\n\n //! Resize accumulators\n virtual void resize(size_t cnt) = 0;\n\n //! Clear accumulators\n virtual void clear(void) = 0;\n\n //! Retrieve an accumulator\n virtual VectorMean &at(size_t i) = 0;\n\n //! Retrieve an accumulator\n virtual const VectorMean &at(size_t i) const = 0;\n\n //! Retrieve the count of accumulators\n virtual size_t count(void) const = 0;\n\n //! Retrieve the dimension of accumulators\n virtual size_t dimension(void) const = 0;\n};\n\n/*! General Vector Mean Array\n */\ntemplate ::type>\nclass GeneralVectorMeanArray : public VectorMeanArray {\n public:\n //! Constructor\n GeneralVectorMeanArray(size_t dim) : dimension_(dim), array_() {}\n\n //! Constructor\n GeneralVectorMeanArray(const GeneralVectorMeanArray &rhs)\n : dimension_(rhs.dimension_), array_(rhs.array_) {}\n\n //! Constructor\n GeneralVectorMeanArray(GeneralVectorMeanArray &&rhs)\n : dimension_(rhs.dimension_), array_(std::move(rhs.array_)) {}\n\n //! Emplace an accumulator\n template \n bool emplace(TArgs &&...args) {\n T accum(std::forward(args)...);\n if (accum.dimension() != dimension_) {\n return false;\n }\n array_.push_back(std::move(accum));\n return true;\n }\n\n //! Resize accumulators\n virtual void resize(size_t cnt) {\n if (array_.size() < cnt) {\n for (size_t i = array_.size(); i < cnt; ++i) {\n array_.emplace_back(dimension_);\n }\n } else {\n array_.resize(cnt);\n }\n }\n\n //! Clear accumulators\n virtual void clear(void) {\n array_.clear();\n }\n\n //! Retrieve an accumulator\n virtual VectorMean &at(size_t i) {\n return array_[i];\n }\n\n //! Retrieve an accumulator\n virtual const VectorMean &at(size_t i) const {\n return array_[i];\n }\n\n //! Retrieve the count of accumulators\n virtual size_t count(void) const {\n return array_.size();\n }\n\n //! Retrieve the dimension of accumulators\n virtual size_t dimension(void) const {\n return dimension_;\n }\n\n private:\n //! Disable them\n GeneralVectorMeanArray(void) = delete;\n\n //! Members\n size_t dimension_;\n std::vector array_;\n};\n\n/*! Numerical Vector Mean\n */\ntemplate ::value>::type>\nclass NumericalVectorMean : public VectorMean {\n public:\n //! Constructor\n NumericalVectorMean(void) : count_(0), accums_() {}\n\n //! Constructor\n NumericalVectorMean(const NumericalVectorMean &rhs)\n : count_(rhs.count_), accums_(rhs.accums_) {}\n\n //! Constructor\n NumericalVectorMean(NumericalVectorMean &&rhs)\n : count_(rhs.count_), accums_(std::move(rhs.accums_)) {}\n\n //! Constructor\n NumericalVectorMean(size_t dim) : count_(0), accums_(dim) {}\n\n //! Constructor\n NumericalVectorMean(const T *means, size_t dim, size_t cnt)\n : count_(cnt), accums_(dim) {\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] = static_cast(means[i]) * count_;\n }\n }\n\n //! Reset accumulator\n void reset(size_t dim) {\n count_ = 0u;\n accums_.clear();\n accums_.resize(dim, 0.0);\n }\n\n //! Reset accumulator\n virtual void reset(void) {\n this->reset(accums_.size());\n }\n\n //! Plus a vector\n virtual bool plus(const void *vec, size_t len) {\n size_t dim = accums_.size();\n if (dim * sizeof(T) != len) {\n return false;\n }\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] += *(static_cast(vec) + i);\n }\n ++count_;\n return true;\n }\n\n //! Retrieve the mean of vectors\n virtual bool mean(void *out, size_t len) const {\n size_t dim = accums_.size();\n if (dim * sizeof(T) != len) {\n return false;\n }\n for (size_t i = 0; i < dim; ++i) {\n *(static_cast(out) + i) = FloatCast(accums_[i] / count_);\n }\n return true;\n }\n\n //! Retrieve the mean of vectors\n virtual void mean(std::string *out) const {\n ailego::NumericalVector &vec =\n *static_cast *>(out);\n\n size_t dim = accums_.size();\n vec.resize(dim);\n for (size_t i = 0; i < dim; ++i) {\n vec[i] = FloatCast(accums_[i] / count_);\n }\n }\n\n //! Merge another vector mean\n virtual bool merge(const VectorMean &rhs) {\n const NumericalVectorMean &src =\n dynamic_cast &>(rhs);\n\n size_t dim = accums_.size();\n if (dim != src.accums_.size()) {\n return false;\n }\n count_ += src.count_;\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] += src.accums_[i];\n }\n return true;\n }\n\n //! Retrieve the count of vectors\n virtual size_t count(void) const {\n return count_;\n }\n\n //! Retrieve dimension of accumulator\n virtual size_t dimension(void) const {\n return accums_.size();\n }\n\n protected:\n //! Convert float type to another type\n template \n static auto FloatCast(const double &val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(val);\n }\n\n //! Convert float type to another type\n template \n static auto FloatCast(const double &val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(std::round(val));\n }\n\n private:\n //! Members\n size_t count_;\n std::vector accums_;\n};\n\n/*! Numerical Vector Harmonic Mean\n */\ntemplate ::value>::type>\nclass NumericalVectorHarmonicMean : public VectorMean {\n public:\n //! Constructor\n NumericalVectorHarmonicMean(void) : count_(0), accums_() {}\n\n //! Constructor\n NumericalVectorHarmonicMean(const NumericalVectorHarmonicMean &rhs)\n : count_(rhs.count_), accums_(rhs.accums_) {}\n\n //! Constructor\n NumericalVectorHarmonicMean(NumericalVectorHarmonicMean &&rhs)\n : count_(rhs.count_), accums_(std::move(rhs.accums_)) {}\n\n //! Constructor\n NumericalVectorHarmonicMean(size_t dim) : count_(0), accums_(dim) {}\n\n //! Constructor\n NumericalVectorHarmonicMean(const T *means, size_t dim, size_t cnt)\n : count_(cnt), accums_(dim) {\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] = static_cast(count_) / static_cast(means[i]);\n }\n }\n\n //! Reset accumulator\n void reset(size_t dim) {\n count_ = 0u;\n accums_.clear();\n accums_.resize(dim, 0.0);\n }\n\n //! Reset accumulator\n virtual void reset(void) {\n this->reset(accums_.size());\n }\n\n //! Plus a vector (harmonic)\n virtual bool plus(const void *vec, size_t len) {\n size_t dim = accums_.size();\n if (dim * sizeof(T) != len) {\n return false;\n }\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] += 1.0 / *(static_cast(vec) + i);\n }\n ++count_;\n return true;\n }\n\n //! Retrieve the mean of vectors (harmonic)\n virtual bool mean(void *out, size_t len) const {\n size_t dim = accums_.size();\n if (dim * sizeof(T) != len) {\n return false;\n }\n for (size_t i = 0; i < dim; ++i) {\n *(static_cast(out) + i) = FloatCast(count_ / accums_[i]);\n }\n return true;\n }\n\n //! Retrieve the mean of vectors\n virtual void mean(std::string *out) const {\n ailego::NumericalVector &vec =\n *static_cast *>(out);\n\n size_t dim = accums_.size();\n vec.resize(dim);\n for (size_t i = 0; i < dim; ++i) {\n vec[i] = FloatCast(count_ / accums_[i]);\n }\n }\n\n //! Merge another vector mean\n virtual bool merge(const VectorMean &rhs) {\n const NumericalVectorHarmonicMean &src =\n dynamic_cast &>(rhs);\n\n size_t dim = accums_.size();\n if (dim != src.accums_.size()) {\n return false;\n }\n count_ += src.count_;\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] += src.accums_[i];\n }\n return true;\n }\n\n //! Retrieve the count of vectors\n virtual size_t count(void) const {\n return count_;\n }\n\n //! Retrieve dimension of accumulator\n virtual size_t dimension(void) const {\n return accums_.size();\n }\n\n protected:\n //! Convert float type to another type\n template \n static auto FloatCast(const double &val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(val);\n }\n\n //! Convert float type to another type\n template \n static auto FloatCast(const double &val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(std::round(val));\n }\n\n private:\n //! Members\n size_t count_;\n std::vector accums_;\n};\n\n/*! Numerical Vector Geometric Mean\n */\ntemplate ::value>::type>\nclass NumericalVectorGeometricMean : public VectorMean {\n public:\n //! Constructor\n NumericalVectorGeometricMean(void) : count_(0), accums_() {}\n\n //! Constructor\n NumericalVectorGeometricMean(const NumericalVectorGeometricMean &rhs)\n : count_(rhs.count_), accums_(rhs.accums_) {}\n\n //! Constructor\n NumericalVectorGeometricMean(NumericalVectorGeometricMean &&rhs)\n : count_(rhs.count_), accums_(std::move(rhs.accums_)) {}\n\n //! Constructor\n NumericalVectorGeometricMean(size_t dim) : count_(0), accums_(dim, 1.0) {}\n\n //! Constructor\n NumericalVectorGeometricMean(const T *means, size_t dim, size_t cnt)\n : count_(cnt), accums_(dim, 1.0) {\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] = std::pow(static_cast(means[i]), count_);\n }\n }\n\n //! Reset accumulator\n void reset(size_t dim) {\n count_ = 0u;\n accums_.clear();\n accums_.resize(dim, 1.0);\n }\n\n //! Reset accumulator\n virtual void reset(void) {\n this->reset(accums_.size());\n }\n\n //! Plus a vector (geometric)\n virtual bool plus(const void *vec, size_t len) {\n size_t dim = accums_.size();\n if (dim * sizeof(T) != len) {\n return false;\n }\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] *= *(static_cast(vec) + i);\n }\n ++count_;\n return true;\n }\n\n //! Retrieve the mean of vectors (geometric)\n virtual bool mean(void *out, size_t len) const {\n size_t dim = accums_.size();\n if (dim * sizeof(T) != len) {\n return false;\n }\n for (size_t i = 0; i < dim; ++i) {\n *(static_cast(out) + i) =\n FloatCast(std::pow(accums_[i], 1.0 / count_));\n }\n return true;\n }\n\n //! Retrieve the mean of vectors\n virtual void mean(std::string *out) const {\n ailego::NumericalVector &vec =\n *static_cast *>(out);\n\n size_t dim = accums_.size();\n vec.resize(dim);\n for (size_t i = 0; i < dim; ++i) {\n vec[i] = FloatCast(std::pow(accums_[i], 1.0 / count_));\n }\n }\n\n //! Merge another vector mean\n virtual bool merge(const VectorMean &rhs) {\n const NumericalVectorGeometricMean &src =\n dynamic_cast &>(rhs);\n\n size_t dim = accums_.size();\n if (dim != src.accums_.size()) {\n return false;\n }\n count_ += src.count_;\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] *= src.accums_[i];\n }\n return true;\n }\n\n //! Retrieve the count of vectors\n virtual size_t count(void) const {\n return count_;\n }\n\n //! Retrieve dimension of accumulator\n virtual size_t dimension(void) const {\n return accums_.size();\n }\n\n protected:\n //! Convert float type to another type\n template \n static auto FloatCast(const double &val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(val);\n }\n\n //! Convert float type to another type\n template \n static auto FloatCast(const double &val) ->\n typename std::enable_if::value, U>::type {\n return static_cast(std::round(val));\n }\n\n private:\n //! Members\n size_t count_;\n std::vector accums_;\n};\n\n/*! Binary Vector Mean\n */\nclass BinaryVectorMean : public VectorMean {\n public:\n //! Constructor\n BinaryVectorMean(void) : count_(0), accums_() {}\n\n //! Constructor\n BinaryVectorMean(const BinaryVectorMean &rhs)\n : count_(rhs.count_), accums_(rhs.accums_) {}\n\n //! Constructor\n BinaryVectorMean(BinaryVectorMean &&rhs)\n : count_(rhs.count_), accums_(std::move(rhs.accums_)) {}\n\n //! Constructor\n BinaryVectorMean(size_t dim) : count_(0), accums_(((dim + 7) >> 3) << 3) {}\n\n //! Constructor\n BinaryVectorMean(const void *means, size_t dim, size_t cnt)\n : count_(cnt), accums_(((dim + 7) >> 3) << 3) {\n const uint8_t *bits = reinterpret_cast(means);\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] = (count_ >> 1);\n\n if (bits[i >> 3] & static_cast(1 << (i & 0x7))) {\n accums_[i] += 1;\n }\n }\n }\n\n //! Reset accumulator\n void reset(size_t dim) {\n count_ = 0u;\n accums_.clear();\n accums_.resize(dim);\n }\n\n //! Reset accumulator\n virtual void reset(void) {\n this->reset(accums_.size());\n }\n\n //! Plus a vector\n virtual bool plus(const void *vec, size_t len) {\n size_t dim = accums_.size();\n if (dim != (len << 3)) {\n return false;\n }\n\n const uint8_t *bits = reinterpret_cast(vec);\n for (size_t i = 0; i < dim; ++i) {\n if (bits[i >> 3] & static_cast(1 << (i & 0x7))) {\n accums_[i] += 1;\n }\n }\n ++count_;\n return true;\n }\n\n //! Retrieve the mean of vectors\n virtual bool mean(void *out, size_t len) const {\n size_t dim = accums_.size();\n if (dim != (len << 3)) {\n return false;\n }\n memset(out, 0, len);\n\n uint8_t *bits = reinterpret_cast(out);\n size_t half_count = count_ >> 1;\n for (size_t i = 0; i < dim; ++i) {\n if (accums_[i] > half_count) {\n bits[i >> 3] |= static_cast(1 << (i & 0x7));\n }\n }\n return true;\n }\n\n //! Retrieve the mean of vectors\n virtual void mean(std::string *out) const {\n size_t dim = accums_.size();\n out->clear();\n out->resize((dim + 7) / 8);\n\n uint8_t *bits =\n reinterpret_cast(const_cast(out->data()));\n size_t half_count = count_ >> 1;\n for (size_t i = 0; i < dim; ++i) {\n if (accums_[i] > half_count) {\n bits[i >> 3] |= static_cast(1 << (i & 0x7));\n }\n }\n }\n\n //! Merge another vector mean\n virtual bool merge(const VectorMean &rhs) {\n const BinaryVectorMean &src = dynamic_cast(rhs);\n\n size_t dim = accums_.size();\n if (dim != src.accums_.size()) {\n return false;\n }\n count_ += src.count_;\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] += src.accums_[i];\n }\n return true;\n }\n\n //! Retrieve the count of vectors\n virtual size_t count(void) const {\n return count_;\n }\n\n //! Retrieve dimension of accumulator\n virtual size_t dimension(void) const {\n return accums_.size();\n }\n\n private:\n //! Members\n size_t count_;\n std::vector accums_;\n};\n\n/*! Numerical Vector Mean\n */\ntemplate ::value>::type>\nclass NibbleVectorMean : public VectorMean {\n public:\n //! Constructor\n NibbleVectorMean(void) : count_(0), accums_() {}\n\n //! Constructor\n NibbleVectorMean(const NibbleVectorMean &rhs)\n : count_(rhs.count_), accums_(rhs.accums_) {}\n\n //! Constructor\n NibbleVectorMean(NibbleVectorMean &&rhs)\n : count_(rhs.count_), accums_(std::move(rhs.accums_)) {}\n\n //! Constructor\n NibbleVectorMean(size_t dim) : count_(0), accums_(dim) {}\n\n //! Constructor\n NibbleVectorMean(const void *means, size_t dim, size_t cnt)\n : count_(cnt), accums_(dim) {\n const uint8_t *arr = reinterpret_cast(means);\n for (size_t i = 0; i != dim; i += 2) {\n uint8_t val = arr[i >> 1];\n int lo = ((int8_t)(val << 4) >> 4);\n int hi = ((int8_t)(val) >> 4);\n accums_[i] = static_cast(lo) * count_;\n accums_[i + 1] = static_cast(hi) * count_;\n }\n }\n\n //! Reset accumulator\n void reset(size_t dim) {\n count_ = 0u;\n accums_.clear();\n accums_.resize(dim, 0.0);\n }\n\n //! Reset accumulator\n virtual void reset(void) {\n this->reset(accums_.size());\n }\n\n //! Plus a vector\n virtual bool plus(const void *vec, size_t len) {\n size_t dim = accums_.size();\n if (dim != (len << 1)) {\n return false;\n }\n\n const uint8_t *arr = reinterpret_cast(vec);\n for (size_t i = 0; i != dim; i += 2) {\n uint8_t val = arr[i >> 1];\n accums_[i] += ((int8_t)(val << 4) >> 4);\n accums_[i + 1] += ((int8_t)(val) >> 4);\n }\n ++count_;\n return true;\n }\n\n //! Retrieve the mean of vectors\n virtual bool mean(void *out, size_t len) const {\n size_t dim = accums_.size();\n if (dim != (len << 1)) {\n return false;\n }\n memset(out, 0, len);\n\n uint8_t *arr = reinterpret_cast(out);\n\n for (size_t i = 0; i != dim; i += 2) {\n int lo = static_cast(std::round(accums_[i] / count_));\n int hi = static_cast(std::round(accums_[i + 1] / count_));\n arr[i >> 1] = (uint8_t)((hi << 4) & 0xf0) | (uint8_t)(lo & 0xf);\n }\n\n return true;\n }\n\n //! Retrieve the mean of vectors\n virtual void mean(std::string *out) const {\n size_t dim = accums_.size();\n out->clear();\n out->resize(dim >> 1);\n\n uint8_t *arr = reinterpret_cast(const_cast(out->data()));\n\n for (size_t i = 0; i != dim; i += 2) {\n int lo = static_cast(std::round(accums_[i] / count_));\n int hi = static_cast(std::round(accums_[i + 1] / count_));\n arr[i >> 1] = (uint8_t)((hi << 4) & 0xf0) | (uint8_t)(lo & 0xf);\n }\n }\n\n //! Merge another vector mean\n virtual bool merge(const VectorMean &rhs) {\n const NibbleVectorMean &src = dynamic_cast(rhs);\n\n size_t dim = accums_.size();\n if (dim != src.accums_.size()) {\n return false;\n }\n count_ += src.count_;\n for (size_t i = 0; i < dim; ++i) {\n accums_[i] += src.accums_[i];\n }\n return true;\n }\n\n //! Retrieve the count of vectors\n virtual size_t count(void) const {\n return count_;\n }\n\n //! Retrieve dimension of accumulator\n virtual size_t dimension(void) const {\n return accums_.size();\n }\n\n private:\n //! Members\n size_t count_;\n std::vector accums_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n#message(STATUS \"PROJECT_ROOT_DIR = ${PROJECT_ROOT_DIR}\")\ncc_library(\n NAME core_knn_flat \n STATIC SHARED STRICT ALWAYS_LINK\n SRCS *.cc\n LIBS core_framework \n INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm ${PROJECT_ROOT_DIR}/src/core/framework\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )\n" }, { "path": "src/core/algorithm/flat/flat_builder.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"flat_builder.h\"\n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\ntemplate \nint FlatBuilder::init(const IndexMeta &meta,\n const ailego::Params ¶ms) {\n meta_ = meta;\n\n // Set the major order\n bool column_major_order = false;\n if (params.get(PARAM_FLAT_COLUMN_MAJOR_ORDER, &column_major_order)) {\n meta_.set_major_order(column_major_order ? IndexMeta::MO_COLUMN\n : IndexMeta::MO_ROW);\n }\n\n // Verify column major order\n if (meta_.major_order() != IndexMeta::MO_ROW) {\n IndexMeta::DataType dt = meta_.data_type();\n\n bool support_column_major = false;\n if ((dt != IndexMeta::DataType::DT_FP32 &&\n dt != IndexMeta::DataType::DT_FP16 &&\n dt != IndexMeta::DataType::DT_INT8 && dt != IndexMeta::DT_INT4 &&\n dt != IndexMeta::DT_BINARY32 && dt != IndexMeta::DT_BINARY64) ||\n (meta_.unit_size() != IndexMeta::UnitSizeof(dt))) {\n if (meta_.major_order() == IndexMeta::MO_COLUMN) {\n LOG_ERROR(\"Unsupported type %d with unit size %u.\", dt,\n meta_.unit_size());\n return IndexError_Unsupported;\n } else {\n support_column_major = false;\n }\n }\n if (meta_.element_size() % IndexMeta::AlignSizeof(dt) != 0) {\n if (meta_.major_order() == IndexMeta::MO_COLUMN) {\n LOG_ERROR(\"Unsupported type %d with dimension %u.\", dt,\n meta_.dimension());\n return IndexError_Unsupported;\n } else {\n support_column_major = false;\n }\n }\n\n if (meta_.major_order() == IndexMeta::MO_UNDEFINED &&\n support_column_major) {\n meta_.set_major_order(IndexMeta::MO_COLUMN);\n }\n }\n\n if (!VerifyMetric(meta_)) {\n LOG_ERROR(\"Invalid index measure %s.\", meta_.metric_name().c_str());\n return IndexError_InvalidArgument;\n }\n\n std::string tag = std::to_string(BATCH_SIZE);\n ailego::Params searcher_params;\n searcher_params.set(PARAM_FLAT_BATCH_SIZE, BATCH_SIZE);\n meta_.set_searcher(\"FlatSearcher\" + tag, 0, searcher_params);\n meta_.set_builder(\"FlatBuilder\" + tag, 0, params);\n return 0;\n}\n\ntemplate \nint FlatBuilder::build(IndexThreads::Pointer,\n IndexHolder::Pointer holder) {\n ailego::ElapsedTime stamp;\n if (!holder->is_matched(meta_)) {\n LOG_ERROR(\"The holder is unmatched with initialized meta.\");\n return IndexError_Mismatch;\n }\n\n holder_ = std::move(holder);\n stats_.set_built_count(holder_->count());\n stats_.set_built_costtime(stamp.milli_seconds());\n return 0;\n}\n\ntemplate \nint FlatBuilder::dump(const IndexDumper::Pointer &dumper) {\n ailego::ElapsedTime stamp;\n if (!holder_) {\n return IndexError_NoReady;\n }\n\n std::vector keys;\n if (meta_.major_order() == IndexMeta::MO_COLUMN) {\n int error_code = this->write_column_index(dumper.get(), &keys);\n if (error_code != 0) {\n return error_code;\n }\n } else {\n int error_code = this->write_row_index(dumper.get(), &keys);\n if (error_code != 0) {\n return error_code;\n }\n }\n\n int error_code = this->write_keys(keys, dumper.get());\n if (error_code != 0) {\n return error_code;\n }\n\n error_code = this->write_mapping(keys, dumper.get());\n if (error_code != 0) {\n return error_code;\n }\n\n error_code = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (error_code != 0) {\n return error_code;\n }\n\n stats_.set_dumped_count(keys.size());\n stats_.set_dumped_costtime(stamp.milli_seconds());\n return 0;\n}\n\ntemplate \nint FlatBuilder::write_keys(const std::vector &keys,\n IndexDumper *dumper) {\n size_t keys_size = keys.size() * sizeof(uint64_t);\n size_t keys_padding_size = ailego_align(keys_size, 32) - keys_size;\n if (dumper->write(keys.data(), keys_size) != keys_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n\n // Write the padding if need\n if (keys_padding_size) {\n std::string padding(keys_padding_size, '\\0');\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n }\n return dumper->append(FLAT_SEGMENT_KEYS_SEG_ID, keys_size, keys_padding_size,\n 0);\n}\n\ntemplate \nint FlatBuilder::write_mapping(const std::vector &keys,\n IndexDumper *dumper) {\n std::vector mapping(keys.size());\n std::iota(mapping.begin(), mapping.end(), 0);\n std::sort(\n mapping.begin(), mapping.end(),\n [&keys](uint32_t lhs, uint32_t rhs) { return (keys[lhs] < keys[rhs]); });\n\n size_t mapping_size = mapping.size() * sizeof(uint32_t);\n size_t mapping_padding_size = ailego_align(mapping_size, 32) - mapping_size;\n if (dumper->write(mapping.data(), mapping_size) != mapping_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n\n // Write the padding if need\n if (mapping_padding_size) {\n std::string padding(mapping_padding_size, '\\0');\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n }\n return dumper->append(FLAT_SEGMENT_MAPPING_SEG_ID, mapping_size,\n mapping_padding_size, 0);\n}\n\ntemplate \ntemplate \nint FlatBuilder::write_column_index(IndexDumper *dumper,\n std::vector *keys) {\n auto iter = holder_->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Failed to create iterator of holder\");\n return IndexError_Runtime;\n }\n\n // Write features\n size_t element_size = holder_->element_size();\n size_t block_size = element_size * BATCH_SIZE;\n std::string block1, block2;\n block1.reserve(block_size);\n block2.reserve(block_size);\n\n for (; iter->is_valid(); iter->next()) {\n block1.append(reinterpret_cast(iter->data()), element_size);\n keys->emplace_back(iter->key());\n\n if (block1.size() == block_size) {\n ailego::MatrixHelper::Transpose(\n block1.data(), element_size / sizeof(T), (void *)block2.data());\n\n if (dumper->write(block2.data(), block_size) != block_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\",\n dumper->name().c_str());\n return IndexError_WriteData;\n }\n block1.clear();\n }\n }\n\n if (!block1.empty()) {\n if (dumper->write(block1.data(), block1.size()) != block1.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n }\n\n // Write the padding if need\n size_t features_size = keys->size() * element_size;\n size_t features_padding_size =\n ailego_align(features_size, 32) - features_size;\n if (features_padding_size) {\n std::string padding(features_padding_size, '\\0');\n\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n }\n return dumper->append(FLAT_SEGMENT_FEATURES_SEG_ID, features_size,\n features_padding_size, 0);\n}\n\ntemplate \nint FlatBuilder::write_row_index(IndexDumper *dumper,\n std::vector *keys) {\n auto iter = holder_->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Failed to create iterator of holder\");\n return IndexError_Runtime;\n }\n\n // Write features\n size_t element_size = holder_->element_size();\n for (; iter->is_valid(); iter->next()) {\n if (dumper->write(iter->data(), element_size) != element_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n keys->emplace_back(iter->key());\n }\n\n // Write the padding if need\n size_t features_size = keys->size() * element_size;\n size_t features_padding_size =\n ailego_align(features_size, 32) - features_size;\n if (features_padding_size) {\n std::string padding(features_padding_size, '\\0');\n\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n }\n return dumper->append(FLAT_SEGMENT_FEATURES_SEG_ID, features_size,\n features_padding_size, 0);\n}\n\nINDEX_FACTORY_REGISTER_BUILDER_ALIAS(LinearBuilder, FlatBuilder<32>);\nINDEX_FACTORY_REGISTER_BUILDER_ALIAS(FlatBuilder, FlatBuilder<32>);\nINDEX_FACTORY_REGISTER_BUILDER_ALIAS(FlatBuilder16, FlatBuilder<16>);\nINDEX_FACTORY_REGISTER_BUILDER_ALIAS(FlatBuilder32, FlatBuilder<32>);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_builder.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Flat Builder\n */\ntemplate \nclass FlatBuilder : public IndexBuilder {\n public:\n //! Destructor\n virtual ~FlatBuilder(void) {}\n\n //! Initialize the builder\n int init(const IndexMeta &meta, const ailego::Params ¶ms) override;\n\n //! Cleanup the builder\n int cleanup(void) override {\n holder_ = nullptr;\n return 0;\n }\n\n //! Train the data\n int train(IndexThreads::Pointer, IndexHolder::Pointer) override {\n stats_.set_trained_count(0u);\n stats_.set_trained_costtime(0u);\n return 0;\n }\n\n //! Train the data\n int train(const IndexTrainer::Pointer &) override {\n stats_.set_trained_count(0u);\n stats_.set_trained_costtime(0u);\n return 0;\n }\n\n //! Build the index\n int build(IndexThreads::Pointer, IndexHolder::Pointer holder) override;\n\n //! Dump index into storage\n int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n const IndexBuilder::Stats &stats(void) const override {\n return stats_;\n }\n\n protected:\n //! Dump index keys\n int write_keys(const std::vector &keys, IndexDumper *dumper);\n\n //! Dump index keys mapping\n int write_mapping(const std::vector &keys, IndexDumper *dumper);\n\n //! Dump index using column-major-order format\n template \n int write_column_index(IndexDumper *dumper, std::vector *keys);\n\n //! Dump index using column-major-order format\n int write_column_index(IndexDumper *dumper, std::vector *keys) {\n switch (IndexMeta::AlignSizeof(meta_.data_type())) {\n case 2:\n return this->write_column_index(dumper, keys);\n case 4:\n return this->write_column_index(dumper, keys);\n case 8:\n return this->write_column_index(dumper, keys);\n default:\n ailego_check_with(0, \"BAD CASE\");\n }\n return IndexError_Runtime;\n }\n\n //! Dump index using row-major-order format\n int write_row_index(IndexDumper *dumper, std::vector *keys);\n\n private:\n IndexMeta meta_{};\n IndexBuilder::Stats stats_{};\n IndexHolder::Pointer holder_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_distance_matrix.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Brute Force Distance Tuple\n */\ntemplate \nclass FlatDistanceTuple;\n\n/*! Brute Force Distance Tuple\n */\ntemplate <>\nclass FlatDistanceTuple<1> {\n public:\n //! Retrieve non-zero if all distances are valid.\n bool is_valid(void) const {\n return !!distance_;\n }\n\n //! Retrieve non-zero if a distance is valid.\n bool is_valid(size_t m) const {\n return m == 1 && !!distance_;\n }\n\n //! Initialize the distance tuple\n void initialize(const IndexMetric &measure) {\n distance_ = measure.distance_matrix(1, 1);\n }\n\n //! Initialize the distance tuple\n void initialize(const IndexMetric &measure, size_t m) {\n distance_ = measure.distance_matrix(m, 1);\n }\n\n //! Compute the distance between matrix and query\n template \n auto distance(const void *m, const void *q, size_t dim, float *out) const ->\n typename std::enable_if::type {\n distance_(m, q, dim, out);\n }\n\n private:\n IndexMetric::MatrixDistance distance_{};\n};\n\n/*! Brute Force Distance Tuple\n */\ntemplate \nclass FlatDistanceTuple<\n K, typename std::enable_if::value>::type> {\n public:\n //! Retrieve non-zero if all distances are valid.\n bool is_valid(void) const {\n return (distance_tuple_.is_valid() && !!distance_);\n }\n\n //! Retrieve non-zero if a distance is valid.\n bool is_valid(size_t m) const {\n return (m == K ? (!!distance_)\n : (m < K ? distance_tuple_.is_valid(m) : false));\n }\n\n //! Initialize the distance tuple\n void initialize(const IndexMetric &measure) {\n distance_tuple_.initialize(measure);\n distance_ = measure.distance_matrix(K, 1);\n }\n\n //! Initialize the distance tuple\n void initialize(const IndexMetric &measure, size_t m) {\n distance_tuple_.initialize(measure, m);\n distance_ = measure.distance_matrix(m, K);\n }\n\n //! Compute the distance between matrix and query\n template \n auto distance(const void *m, const void *q, size_t dim, float *out) const ->\n typename std::enable_if::type {\n distance_(m, q, dim, out);\n }\n\n //! Compute the distance between matrix and query\n template \n auto distance(const void *m, const void *q, size_t dim, float *out) const ->\n typename std::enable_if<(K > M) && IsEqualPowerofTwo::value>::type {\n distance_tuple_.template distance(m, q, dim, out);\n }\n\n private:\n FlatDistanceTuple<(K >> 1)> distance_tuple_{};\n IndexMetric::MatrixDistance distance_{};\n};\n\n/*! Brute Force Distance Matrix\n */\ntemplate \nclass FlatDistanceMatrix;\n\n/*! Brute Force Distance Matrix\n */\ntemplate <>\nclass FlatDistanceMatrix<1> {\n public:\n //! Retrieve non-zero if all distances are valid.\n bool is_valid(void) const {\n return (!!distance_);\n }\n\n //! Initialize the distance tuple\n void initialize(const IndexMetric &measure) {\n distance_ = measure.distance_matrix(1, 1);\n }\n\n //! Compute the distance between matrix and query\n template \n auto distance(const void *m, const void *q, size_t dim, float *out) const ->\n typename std::enable_if::type {\n distance_(m, q, dim, out);\n }\n\n private:\n IndexMetric::MatrixDistance distance_{};\n};\n\n/*! Brute Force Distance Matrix\n */\ntemplate \nclass FlatDistanceMatrix<\n K, typename std::enable_if::value>::type> {\n public:\n //! Retrieve non-zero if all distances are valid.\n bool is_valid(void) const {\n return (tuple_h_.is_valid() && tuple_v_.is_valid());\n }\n\n //! Retrieve non-zero if a distance is valid.\n bool is_valid(size_t m, size_t n) const {\n return (m == K ? tuple_h_.is_valid(n)\n : (m < K && n == 1 ? tuple_v_.is_valid(m) : false));\n }\n\n //! Initialize the distance tuple\n void initialize(const IndexMetric &measure) {\n tuple_h_.initialize(measure, K);\n tuple_v_.initialize(measure);\n }\n\n //! Compute the distance between matrix and query\n template \n auto distance(const void *m, const void *q, size_t dim, float *out) const ->\n typename std::enable_if<(K == M) && (K >= N)>::type {\n tuple_h_.template distance(m, q, dim, out);\n }\n\n //! Compute the distance between matrix and query\n template \n auto distance(const void *m, const void *q, size_t dim, float *out) const ->\n typename std::enable_if<(K > M) && (N == 1u)>::type {\n tuple_v_.template distance(m, q, dim, out);\n }\n\n private:\n FlatDistanceTuple tuple_h_{};\n FlatDistanceTuple<(K >> 1)> tuple_v_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_index_format.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\nusing node_id_t = uint32_t;\nusing key_t = uint64_t;\nusing level_t = int32_t;\nusing dist_t = float;\nusing TopkHeap = ailego::KeyValueHeap;\nusing CandidateHeap =\n ailego::KeyValueHeap>;\nconstexpr node_id_t kInvalidNodeId = static_cast(-1);\nconstexpr key_t kInvalidKey = static_cast(-1);\n\n/*! Index Format of Linear Index Header\n */\nstruct LinearIndexHeader {\n LinearIndexHeader()\n : header_size(0),\n total_vector_count(0),\n linear_body_size(0),\n linear_list_count(0),\n block_vector_count(0),\n block_size(0),\n block_count(0),\n index_meta_size(0) {\n memset(reserved_, 0, sizeof(reserved_));\n memset(index_meta, 0, sizeof(index_meta));\n }\n uint32_t header_size{0};\n uint32_t total_vector_count{0};\n uint64_t linear_body_size{0};\n uint32_t linear_list_count{0};\n uint32_t block_vector_count{0};\n uint32_t block_size{0};\n uint32_t block_count{0};\n uint32_t index_meta_size{0};\n char reserved_[28] = {0};\n char index_meta[0];\n};\n\n/*! Index Format of Linear Index Meta for each Linear list\n */\nstruct LinearListMeta {\n LinearListMeta() : offset(0), block_count(0), vector_count(0), id_offset(0) {\n memset(reserved_, 0, sizeof(reserved_));\n }\n uint64_t offset{0};\n uint32_t block_count{0};\n uint32_t vector_count{0};\n uint32_t id_offset{0};\n char reserved_[16] = {0};\n};\n\n/*! Index Format of Location in Linear Index for each vector\n */\nstruct LinearVecLocation {\n LinearVecLocation(size_t off, bool col)\n : offset(off), column_major(col), reserved(0u) {}\n\n uint64_t offset : 48; // feature offset in posting block segment\n uint64_t column_major : 1; // coloum major if true\n uint64_t reserved : 15;\n};\n\n/*! Index Format of Integer Quantizer params for each linear list\n */\nstruct LinearIntegerQuantizerParams {\n float scale{1.0};\n float bias{0.0};\n};\n\n/*! Location of Vectors Block in Storage Segment\n */\nstruct BlockLocation {\n uint32_t segment_id{0};\n uint32_t block_index{0};\n};\n\n/*! The Header of a Block in Storage Segment\n */\nstruct BlockHeader {\n BlockHeader() : vector_count(0u), column_major(0u), reserved(0u) {}\n BlockLocation next;\n uint16_t vector_count{0};\n uint16_t column_major : 1;\n uint16_t reserved : 15;\n};\n\nstruct DeletionMap {\n void set(uint32_t index) {\n bitset.set(index);\n }\n\n void reset(uint32_t index) {\n bitset.reset(index);\n }\n\n bool test(uint32_t index) const {\n return bitset.test(index);\n }\n\n bool is_dirty() const {\n return bitset.test_any();\n }\n\n ailego::FixedBitset<32> bitset{};\n};\n\nstatic_assert(sizeof(DeletionMap) == 4, \"DeletionMap must be 4 bytes\");\n\n/*! Meta Information of Streamer Entity\n */\nstruct StreamerLinearMeta {\n StreamerLinearMeta()\n : create_time(0),\n update_time(0),\n revision_id(0),\n segment_count(0),\n segment_size(0) {\n memset(reserved_, 0, sizeof(reserved_));\n }\n uint64_t create_time{0};\n uint64_t update_time{0};\n uint64_t revision_id{0};\n uint32_t segment_count{0};\n uint32_t segment_size{0};\n uint8_t reserved_[32] = {0};\n LinearIndexHeader header;\n};\n\n/*! Location of Vector in Storage Segment\n */\nstruct VectorLocation {\n //! Constructor\n VectorLocation(void)\n : segment_id(0u), column_major(0u), reserved(0u), offset(0u) {}\n\n //! Constructor\n VectorLocation(uint32_t id, bool col, uint32_t off)\n : segment_id(id), column_major(col), reserved(0u), offset(off) {}\n\n uint32_t segment_id{0};\n uint16_t column_major : 1;\n uint16_t reserved : 15;\n uint32_t offset{0};\n\n public:\n bool operator==(const VectorLocation &other) const {\n return segment_id == other.segment_id &&\n column_major == other.column_major && offset == other.offset;\n }\n};\n\n// static_assert(sizeof(VectorLocation) == sizeof(uint64_t),\n// \"VectorLocation must be size of 8 bytes\");\n\nstruct KeyInfo {\n KeyInfo(void) : centroid_idx(0u) {}\n KeyInfo(uint32_t idx, const VectorLocation &loc)\n : centroid_idx(idx), location(loc) {}\n KeyInfo(VectorLocation loc) : location(loc) {}\n uint32_t centroid_idx{0};\n VectorLocation location;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_searcher.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"flat_searcher.h\"\n#include \n#include \n#include \"flat_distance_matrix.h\"\n#include \"flat_searcher_context.h\"\n#include \"flat_searcher_provider.h\"\n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\ntemplate \nIndexProvider::Pointer FlatSearcher::create_provider(void) const {\n std::lock_guard lock(mapping_mutex_);\n\n if (mapping_.empty()) {\n auto mapping_segment = container_->get(FLAT_SEGMENT_MAPPING_SEG_ID);\n if (!mapping_segment) {\n LOG_ERROR(\"Failed to fetch segment %s\",\n FLAT_SEGMENT_MAPPING_SEG_ID.c_str());\n return nullptr;\n }\n\n if (mapping_segment->data_size() % sizeof(uint32_t) != 0) {\n LOG_ERROR(\"Invalid data size %zu of mapping segment\",\n mapping_segment->data_size());\n return nullptr;\n }\n\n size_t mapping_count = mapping_segment->data_size() / sizeof(uint32_t);\n if (mapping_count * meta_.element_size() !=\n features_segment_->data_size()) {\n LOG_ERROR(\"Invalid data size %zd of mapping segment\",\n features_segment_->data_size());\n return nullptr;\n }\n\n const uint32_t *mapping = nullptr;\n if (mapping_segment->read(0, reinterpret_cast(&mapping),\n mapping_segment->data_size()) !=\n mapping_segment->data_size()) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from mapping segment\",\n mapping_segment->data_size());\n return nullptr;\n }\n mapping_.clear();\n mapping_.reserve(mapping_count);\n std::copy(mapping, mapping + mapping_count, std::back_inserter(mapping_));\n }\n return IndexProvider::Pointer(new (std::nothrow)\n FlatSearcherProvider(this));\n}\n\ntemplate \nint FlatSearcher::load(IndexStorage::Pointer cntr,\n IndexMetric::Pointer measure) {\n ailego::ElapsedTime stamp;\n if (!cntr) {\n return IndexError_InvalidArgument;\n }\n\n int error_code = IndexHelper::DeserializeFromStorage(cntr.get(), &meta_);\n if (error_code != 0) {\n LOG_ERROR(\n \"Failed to deserialize index meta from container %s, error=%d, %s\",\n cntr->name().c_str(), error_code, IndexError::What(error_code));\n return error_code;\n }\n\n if (!measure) {\n error_code = InitializeMetric(meta_, &measure_);\n if (error_code != 0) {\n LOG_ERROR(\"Failed to initialize index measure %s, error=%d, %s\",\n meta_.metric_name().c_str(), error_code,\n IndexError::What(error_code));\n return error_code;\n }\n if (measure_->query_metric()) {\n measure_ = measure_->query_metric();\n }\n } else {\n if (!measure->is_matched(meta_)) {\n LOG_ERROR(\n \"The index measure is unmatched with index meta from container.\");\n return IndexError_Mismatch;\n }\n measure_ = std::move(measure);\n }\n\n column_major_order_ = (meta_.major_order() == IndexMeta::MO_COLUMN);\n distance_matrix_.initialize(*measure_);\n\n if (column_major_order_) {\n if (!distance_matrix_.is_valid()) {\n LOG_ERROR(\"Lack of distance functions to support column index.\");\n return IndexError_Unsupported;\n }\n } else {\n if (!distance_matrix_.is_valid(1, 1)) {\n LOG_ERROR(\"Lack of distance functions to support row index.\");\n return IndexError_Unsupported;\n }\n }\n\n auto keys_segment = cntr->get(FLAT_SEGMENT_KEYS_SEG_ID);\n if (!keys_segment) {\n LOG_ERROR(\"Failed to fetch segment %s\", FLAT_SEGMENT_KEYS_SEG_ID.c_str());\n return IndexError_NoExist;\n }\n features_segment_ = cntr->get(FLAT_SEGMENT_FEATURES_SEG_ID);\n if (!features_segment_) {\n LOG_ERROR(\"Failed to fetch segment %s\", FLAT_SEGMENT_KEYS_SEG_ID.c_str());\n return IndexError_NoExist;\n }\n\n if (keys_segment->data_size() % sizeof(uint64_t) != 0) {\n LOG_ERROR(\"Invalid data size %zu of keys segment\",\n keys_segment->data_size());\n return IndexError_InvalidLength;\n }\n\n size_t keys_count = keys_segment->data_size() / sizeof(uint64_t);\n if (keys_count * meta_.element_size() != features_segment_->data_size()) {\n LOG_ERROR(\"Invalid data size %zd of features segment\",\n features_segment_->data_size());\n return IndexError_Mismatch;\n }\n\n if (keys_segment->read(0, reinterpret_cast(&keys_),\n keys_segment->data_size()) !=\n keys_segment->data_size()) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from keys segment\",\n keys_segment->data_size());\n return IndexError_ReadData;\n }\n\n for (size_t i = 0; i < keys_count; i++) {\n key_id_mapping_[keys_[i]] = i;\n }\n\n container_ = cntr;\n magic_ = IndexContext::GenerateMagic();\n stats_.set_loaded_count(keys_count);\n stats_.set_loaded_costtime(stamp.milli_seconds());\n return 0;\n}\n\ntemplate \nint FlatSearcher::search_impl(const void *query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n ailego_assert(query && !!context);\n ailego_assert(measure_->is_matched(meta_, qmeta));\n\n FlatSearcherContext *bf_context =\n dynamic_cast *>(context.get());\n if (!bf_context) {\n LOG_ERROR(\"Invalid brute-force searcher context\");\n return IndexError_InvalidArgument;\n }\n\n if (bf_context->magic() != magic_) {\n bf_context->reset(this);\n }\n if (bf_context->group_by_search()) {\n return bf_context->group_by_search_impl(query, qmeta, 1);\n } else {\n return (column_major_order_ ? bf_context->search_column(query, qmeta)\n : bf_context->search_row(query, qmeta));\n }\n}\n\ntemplate \nint FlatSearcher::search_impl(const void *query,\n const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const {\n ailego_assert(query && count && !!context);\n ailego_assert(measure_->is_matched(meta_, qmeta));\n\n FlatSearcherContext *bf_context =\n dynamic_cast *>(context.get());\n if (!bf_context) {\n LOG_ERROR(\"Invalid brute-force searcher context\");\n return IndexError_InvalidArgument;\n }\n\n if (bf_context->magic() != magic_) {\n bf_context->reset(this);\n }\n\n if (bf_context->group_by_search()) {\n return bf_context->group_by_search_impl(query, qmeta, count);\n } else {\n return (column_major_order_ ? bf_context->search_column(query, qmeta, count)\n : bf_context->search_row(query, qmeta, count));\n }\n}\n\ntemplate \nint FlatSearcher::search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n ailego_assert(query && count && !!context);\n ailego_assert(measure_->is_matched(meta_, qmeta));\n\n if (ailego_unlikely(p_keys.size() != count)) {\n LOG_ERROR(\"The size of p_keys is not equal to count\");\n return IndexError_InvalidArgument;\n }\n\n FlatSearcherContext *bf_context =\n dynamic_cast *>(context.get());\n if (!bf_context) {\n LOG_ERROR(\"Invalid brute-force searcher context\");\n return IndexError_InvalidArgument;\n }\n\n if (bf_context->magic() != magic_) {\n bf_context->reset(this);\n }\n\n return bf_context->search_bf_by_p_keys_impl(query, p_keys, qmeta, count);\n}\n\ntemplate \nIndexSearcher::Context::Pointer FlatSearcher::create_context(\n void) const {\n return IndexSearcher::Context::Pointer(\n new FlatSearcherContext(this));\n}\n\nINDEX_FACTORY_REGISTER_SEARCHER_ALIAS(LinearSearcher, FlatSearcher<32>);\nINDEX_FACTORY_REGISTER_SEARCHER_ALIAS(FlatSearcher, FlatSearcher<32>);\nINDEX_FACTORY_REGISTER_SEARCHER_ALIAS(FlatSearcher16, FlatSearcher<16>);\nINDEX_FACTORY_REGISTER_SEARCHER_ALIAS(FlatSearcher32, FlatSearcher<32>);\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_searcher.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n#include \n#include \n#include \n#include \"flat_distance_matrix.h\"\n#include \"flat_index_format.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Flat Searcher\n */\ntemplate \nclass FlatSearcher : public IndexSearcher {\n public:\n //! Destructor\n virtual ~FlatSearcher(void) = default;\n\n //! Initialize Searcher\n int init(const ailego::Params &index_params) override {\n params_ = index_params;\n read_block_size_ = FLAT_DEFAULT_READ_BLOCK_SIZE;\n index_params.get(PARAM_FLAT_READ_BLOCK_SIZE, &read_block_size_);\n return 0;\n }\n\n //! Cleanup Searcher\n int cleanup(void) override {\n return this->unload();\n }\n\n //! Load index from container\n int load(IndexStorage::Pointer cntr, IndexMetric::Pointer measure) override;\n\n //! Unload index\n int unload(void) override {\n container_ = nullptr;\n measure_ = nullptr;\n features_segment_ = nullptr;\n keys_ = nullptr;\n key_id_mapping_.clear();\n return 0;\n }\n\n //! Similarity brute force search\n int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override {\n return this->search_impl(query, qmeta, context);\n }\n\n //! Similarity brute force search\n int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const override {\n return this->search_impl(query, qmeta, count, context);\n }\n\n //! Similarity search\n int search_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity search\n int search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const override;\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const void *query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override {\n return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context);\n }\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const void *query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const override;\n\n //! Retrieve statistics\n const IndexSearcher::Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve meta of index\n const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n //! Retrieve params of index\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n //! Create a searcher context\n IndexSearcher::Context::Pointer create_context(void) const override;\n\n //! Create a searcher provider\n IndexProvider::Pointer create_provider(void) const override;\n\n //! Retrieve magic number\n uint32_t magic(void) const {\n return magic_;\n }\n\n //! Retrieve block size of data read\n uint32_t read_block_size(void) const {\n return read_block_size_;\n }\n\n //! Retrieve primary key via index id\n uint64_t key(size_t i) const {\n return keys_[i];\n }\n\n // Retrieve index id via primary key\n node_id_t get_id(key_t key) const {\n auto it = key_id_mapping_.find(key);\n if (it != key_id_mapping_.end()) {\n return it->second;\n } else {\n return kInvalidNodeId;\n }\n }\n\n //! Retrieve primary key via index id\n uint32_t local_index(size_t i) const {\n return mapping_[i];\n }\n\n //! Retrieve primary key via index id\n inline bool column_major_order(void) const {\n return column_major_order_;\n }\n\n //! Retrieve the distance matrix\n const FlatDistanceMatrix &distance_matrix(void) const {\n return distance_matrix_;\n }\n\n //! Clone a features segment\n IndexStorage::Segment::Pointer clone_features_segment(void) const {\n return features_segment_->clone();\n }\n\n const void *get_vector(key_t key) const override {\n auto provider = this->create_provider();\n return provider->get_vector(key);\n }\n\n private:\n //! Members\n const uint64_t *keys_{nullptr};\n std::unordered_map key_id_mapping_;\n uint32_t magic_{IndexContext::GenerateMagic()};\n uint32_t read_block_size_{FLAT_DEFAULT_READ_BLOCK_SIZE};\n bool column_major_order_{false};\n IndexMeta meta_{};\n IndexStorage::Pointer container_{};\n IndexMetric::Pointer measure_{};\n ailego::Params params_{};\n IndexStorage::Segment::Pointer features_segment_{};\n mutable std::vector mapping_{};\n mutable std::mutex mapping_mutex_{};\n FlatDistanceMatrix distance_matrix_{};\n IndexSearcher::Stats stats_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_searcher_context.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \"flat_index_format.h\"\n#include \"flat_searcher.h\"\n#include \"flat_utility.h\"\n\n\nnamespace zvec {\nnamespace core {\n\n/*! Brute Force Searcher Context\n */\ntemplate \nclass FlatSearcherContext : public IndexSearcher::Context {\n public:\n //! Constructor\n FlatSearcherContext(const FlatSearcher *owner) {\n this->reset(owner);\n }\n\n //! Destructor\n virtual ~FlatSearcherContext(void) {}\n\n //! Set topk of search result\n void set_topk(uint32_t topk) override {\n topk_ = topk;\n }\n\n //! Retrieve search result\n const IndexDocumentList &result(void) const override {\n return result_heaps_.at(0);\n }\n\n //! Retrieve search result with index\n const IndexDocumentList &result(size_t index) const override {\n return result_heaps_.at(index);\n }\n\n //! Retrieve result object for output\n IndexDocumentList *mutable_result(size_t idx) override {\n return &result_heaps_.at(idx);\n }\n\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(void) const override {\n return group_results_[0];\n }\n\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(\n size_t idx) const override {\n return group_results_[idx];\n }\n\n //! Update the parameters of context\n int update(const ailego::Params & /*params*/) override {\n return 0;\n }\n\n //! Retrieve magic number\n uint32_t magic(void) const override {\n return magic_;\n }\n\n //! Get group topk\n inline uint32_t group_topk() const {\n return group_topk_;\n }\n\n //! Get group num\n inline uint32_t group_num() const {\n return group_num_;\n }\n\n inline std::map &group_topk_heaps() {\n return group_topk_heaps_;\n }\n\n void set_fetch_vector(bool v) override {\n fetch_vector_ = v;\n }\n\n bool fetch_vector() const override {\n return fetch_vector_;\n }\n\n inline void resize_group_results(size_t size) {\n if (group_by_search()) {\n group_results_.resize(size);\n }\n }\n\n void topk_to_group_result(uint32_t idx) {\n ailego_assert_with(idx < group_results_.size(), \"invalid idx\");\n\n group_results_[idx].clear();\n\n std::vector> group_topk_list;\n std::vector> best_score_in_groups;\n for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end();\n itr++) {\n const std::string &group_id = (*itr).first;\n auto &heap = (*itr).second;\n heap.sort();\n\n if (heap.size() > 0) {\n float best_score = heap[0].second;\n best_score_in_groups.push_back(std::make_pair(group_id, best_score));\n }\n }\n\n std::sort(best_score_in_groups.begin(), best_score_in_groups.end(),\n [](const std::pair &a,\n const std::pair &b) -> int {\n return a.second < b.second;\n });\n\n // truncate to group num\n for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size();\n ++i) {\n const std::string &group_id = best_score_in_groups[i].first;\n\n group_topk_list.emplace_back(\n std::make_pair(group_id, group_topk_heaps_[group_id]));\n }\n\n group_results_[idx].resize(group_topk_list.size());\n\n for (uint32_t i = 0; i < group_topk_list.size(); ++i) {\n const std::string &group_id = group_topk_list[i].first;\n group_results_[idx][i].set_group_id(group_id);\n\n uint32_t size = std::min(\n group_topk_, static_cast(group_topk_list[i].second.size()));\n\n for (uint32_t j = 0; j < size; ++j) {\n auto score = group_topk_list[i].second[j].second;\n if (score > this->threshold()) {\n break;\n }\n\n node_id_t id = group_topk_list[i].second[j].first;\n\n auto provider = owner_->create_provider();\n\n if (fetch_vector_) {\n group_results_[idx][i].mutable_docs()->emplace_back(\n id, score, id, provider->get_vector(id));\n } else {\n group_results_[idx][i].mutable_docs()->emplace_back(id, score, id);\n }\n }\n }\n }\n\n //! Get if group by search\n bool group_by_search() {\n return group_num_ > 0;\n }\n\n //! Set group params\n void set_group_params(uint32_t group_num, uint32_t group_topk) override {\n group_num_ = group_num;\n group_topk_ = group_topk;\n group_topk_heaps_.clear();\n }\n\n void reset() override {}\n\n //! Reset the context\n void reset(const FlatSearcher *owner) {\n magic_ = owner->magic();\n feature_size_ = owner->meta().element_size();\n\n uint32_t block_size = feature_size_ * BATCH_SIZE;\n actual_read_size_ =\n (owner->read_block_size() + block_size - 1) / block_size * block_size;\n features_segment_ = owner->clone_features_segment();\n owner_ = owner;\n }\n\n //! Similarity search\n int search_row(const void *query, const IndexQueryMeta &qmeta) {\n return (this->filter().is_valid()\n ? this->search_row_filter(query, qmeta)\n : this->search_row_nofilter(query, qmeta));\n }\n\n //! Similarity search\n int search_row(const void *query, const IndexQueryMeta &qmeta, size_t count) {\n return (this->filter().is_valid()\n ? this->batch_search_row_filter(query, qmeta, count)\n : this->batch_search_row_nofilter(query, qmeta, count));\n }\n\n //! Similarity search\n int search_column(const void *query, const IndexQueryMeta &qmeta) {\n return (this->filter().is_valid()\n ? this->search_column_filter(query, qmeta)\n : this->search_column_nofilter(query, qmeta));\n }\n\n //! Similarity search\n int search_column(const void *query, const IndexQueryMeta &qmeta,\n size_t count) {\n return (this->filter().is_valid()\n ? this->batch_search_column_filter(query, qmeta, count)\n : this->batch_search_column_nofilter(query, qmeta, count));\n }\n\n int group_by_search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count);\n\n int search_bf_by_p_keys_impl(const void *query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count);\n\n protected:\n //! Enqueue items into the search heaps (without filter)\n template \n auto batch_enqueue_nofilter(const void *block, size_t block_index,\n size_t query_index, const IndexQueryMeta &qmeta,\n size_t query_count) ->\n typename std::enable_if::value>::type {\n size_t query_batch_count = query_count / K;\n\n for (size_t i = 0; i != query_batch_count; ++i) {\n owner_->distance_matrix().template distance(\n block, &batch_queries_[query_index * qmeta.element_size()],\n qmeta.dimension(), scores_);\n\n for (size_t k = 0; k != K; ++k) {\n IndexDocumentHeap *heap = &result_heaps_[query_index++];\n for (size_t j = 0; j != BATCH_SIZE; ++j) {\n heap->emplace(0, scores_[k * BATCH_SIZE + j], block_index + j);\n }\n } // end of for\n } // end of for\n\n size_t query_left_count = query_count % K;\n if (query_left_count != 0) {\n this->batch_enqueue_nofilter<(K >> 1)>(block, block_index, query_index,\n qmeta, query_left_count);\n }\n }\n\n //! Enqueue items into the search heaps (without filter)\n template \n auto batch_enqueue_nofilter(const void *block, size_t block_index,\n size_t query_index, const IndexQueryMeta &qmeta,\n size_t query_count) ->\n typename std::enable_if::type {\n ailego_assert(query_count == 1);\n (void)query_count;\n\n owner_->distance_matrix().template distance(\n block, &batch_queries_[query_index * qmeta.element_size()],\n qmeta.dimension(), scores_);\n\n IndexDocumentHeap *heap = &result_heaps_[query_index];\n for (size_t i = 0; i != BATCH_SIZE; ++i) {\n heap->emplace(0, scores_[i], block_index + i);\n }\n }\n\n //! Enqueue items into the search heaps (with filter)\n template \n auto batch_enqueue_filter(const void *block, size_t block_index,\n size_t block_mask, size_t query_index,\n const IndexQueryMeta &qmeta, size_t query_count) ->\n typename std::enable_if::value>::type {\n size_t query_batch_count = query_count / K;\n\n for (size_t i = 0; i != query_batch_count; ++i) {\n owner_->distance_matrix().template distance(\n block, &batch_queries_[query_index * qmeta.element_size()],\n qmeta.dimension(), scores_);\n\n for (size_t k = 0; k != K; ++k) {\n IndexDocumentHeap *heap = &result_heaps_[query_index++];\n for (size_t j = 0; j != BATCH_SIZE; ++j) {\n if ((block_mask & (1 << j)) != 0) {\n heap->emplace(0, scores_[k * BATCH_SIZE + j], block_index + j);\n }\n }\n } // end of for\n } // end of for\n\n size_t query_left_count = query_count % K;\n if (query_left_count != 0) {\n this->batch_enqueue_filter<(K >> 1)>(\n block, block_index, block_mask, query_index, qmeta, query_left_count);\n }\n }\n\n //! Enqueue items into the search heaps (with filter)\n template \n auto batch_enqueue_filter(const void *block, size_t block_index,\n size_t block_mask, size_t query_index,\n const IndexQueryMeta &qmeta, size_t query_count) ->\n typename std::enable_if::type {\n ailego_assert(query_count == 1);\n (void)query_count;\n\n owner_->distance_matrix().template distance(\n block, &batch_queries_[query_index * qmeta.element_size()],\n qmeta.dimension(), scores_);\n\n IndexDocumentHeap *heap = &result_heaps_[query_index];\n for (size_t i = 0; i != BATCH_SIZE; ++i) {\n if ((block_mask & (1 << i)) != 0) {\n heap->emplace(0, scores_[i], block_index + i);\n }\n }\n }\n\n //! Enqueue items into the search heaps (without filter)\n template \n auto single_enqueue_nofilter(const void *feature, size_t feature_index,\n size_t query_index, const IndexQueryMeta &qmeta,\n size_t query_count) ->\n typename std::enable_if::value>::type {\n size_t query_batch_count = query_count / K;\n\n for (size_t i = 0; i != query_batch_count; ++i) {\n owner_->distance_matrix().template distance(\n &batch_queries_[query_index * qmeta.element_size()], feature,\n qmeta.dimension(), scores_);\n\n for (size_t k = 0; k != K; ++k) {\n result_heaps_[query_index++].emplace(0, scores_[k], feature_index);\n }\n }\n size_t query_left_count = query_count % K;\n if (query_left_count != 0) {\n this->single_enqueue_nofilter<(K >> 1)>(\n feature, feature_index, query_index, qmeta, query_left_count);\n }\n }\n\n //! Enqueue items into the search heaps (without filter)\n template \n auto single_enqueue_nofilter(const void *feature, size_t feature_index,\n size_t query_index, const IndexQueryMeta &qmeta,\n size_t query_count) ->\n typename std::enable_if::type {\n ailego_assert(query_count == 1);\n (void)query_count;\n\n owner_->distance_matrix().template distance<1>(\n feature, &batch_queries_[query_index * qmeta.element_size()],\n qmeta.dimension(), scores_);\n result_heaps_[query_index].emplace(0, scores_[0], feature_index);\n }\n\n protected:\n //! Similarity search (1 column without filter)\n int search_column_nofilter(const void *query, const IndexQueryMeta &qmeta);\n\n //! Similarity search (1 column with filter)\n int search_column_filter(const void *query, const IndexQueryMeta &qmeta);\n\n //! Similarity search (1 row without filter)\n int search_row_nofilter(const void *query, const IndexQueryMeta &qmeta);\n\n //! Similarity search (1 row with filter)\n int search_row_filter(const void *query, const IndexQueryMeta &qmeta);\n\n //! Similarity search (n columns without filter)\n int batch_search_column_nofilter(const void *query,\n const IndexQueryMeta &qmeta,\n size_t query_count);\n\n //! Similarity search (n columns with filter)\n int batch_search_column_filter(const void *query, const IndexQueryMeta &qmeta,\n size_t query_count);\n\n //! Similarity search (n rows without filter)\n int batch_search_row_nofilter(const void *query, const IndexQueryMeta &qmeta,\n size_t query_count);\n\n //! Similarity search (n rows with filter)\n int batch_search_row_filter(const void *query, const IndexQueryMeta &qmeta,\n size_t query_count);\n\n private:\n const FlatSearcher *owner_{nullptr};\n uint32_t magic_{0};\n uint32_t topk_{0};\n uint32_t feature_size_{0};\n uint32_t actual_read_size_{0};\n IndexStorage::Segment::Pointer features_segment_{};\n std::vector result_heaps_{1};\n std::string batch_queries_{};\n float scores_[BATCH_SIZE * BATCH_SIZE];\n bool fetch_vector_{false};\n\n // group\n uint32_t group_num_{0}, group_topk_{0};\n std::map group_topk_heaps_{};\n std::vector group_results_{};\n};\n\ntemplate \nint FlatSearcherContext::search_column_nofilter(\n const void *query, const IndexQueryMeta &qmeta) {\n IndexDocumentHeap *heap = &result_heaps_[0];\n heap->clear();\n heap->limit(topk_);\n heap->set_threshold(this->threshold());\n\n size_t left_size = features_segment_->data_size();\n size_t block_size = feature_size_ * BATCH_SIZE;\n size_t read_offset = 0;\n size_t feature_index = 0;\n auto matrix = this->owner_->distance_matrix();\n\n while (left_size >= actual_read_size_) {\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, actual_read_size_) !=\n actual_read_size_) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n actual_read_size_);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < actual_read_size_; offset += block_size) {\n matrix.template distance(\n (const char *)data + offset, query, qmeta.dimension(), scores_);\n\n for (size_t i = 0; i != BATCH_SIZE; ++i) {\n heap->emplace(0, scores_[i], feature_index++);\n }\n }\n read_offset += actual_read_size_;\n left_size -= actual_read_size_;\n }\n\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, left_size) != left_size) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from features segment\",\n left_size);\n return IndexError_ReadData;\n }\n\n // Process left block features\n size_t left_size_aligned = left_size / block_size * block_size;\n for (size_t offset = 0; offset != left_size_aligned; offset += block_size) {\n matrix.template distance((const char *)data + offset, query,\n qmeta.dimension(), scores_);\n\n for (size_t i = 0; i != BATCH_SIZE; ++i) {\n heap->emplace(0, scores_[i], feature_index++);\n }\n }\n\n // Process left single features\n for (size_t offset = left_size_aligned; offset < left_size;\n offset += feature_size_) {\n float score;\n matrix.template distance<1>((const char *)data + offset, query,\n qmeta.dimension(), &score);\n heap->emplace(0, score, feature_index++);\n }\n\n for (auto &it : *heap) {\n it.set_key(owner_->key(it.index()));\n }\n heap->sort();\n return 0;\n}\n\ntemplate \nint FlatSearcherContext::search_column_filter(\n const void *query, const IndexQueryMeta &qmeta) {\n IndexDocumentHeap *heap = &result_heaps_[0];\n heap->clear();\n heap->limit(topk_);\n heap->set_threshold(this->threshold());\n\n size_t left_size = features_segment_->data_size();\n size_t block_size = feature_size_ * BATCH_SIZE;\n size_t read_offset = 0;\n size_t feature_index = 0;\n auto matrix = owner_->distance_matrix();\n\n while (left_size >= actual_read_size_) {\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, actual_read_size_) !=\n actual_read_size_) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n actual_read_size_);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < actual_read_size_; offset += block_size) {\n matrix.template distance(\n (const char *)data + offset, query, qmeta.dimension(), scores_);\n\n for (size_t i = 0; i != BATCH_SIZE; ++i) {\n uint64_t feature_key = owner_->key(feature_index);\n\n if (!this->filter()(feature_key)) {\n if (group_by_search()) {\n }\n heap->emplace(feature_key, scores_[i], feature_index);\n }\n feature_index += 1;\n }\n }\n read_offset += actual_read_size_;\n left_size -= actual_read_size_;\n }\n\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, left_size) != left_size) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from features segment\",\n left_size);\n return IndexError_ReadData;\n }\n\n // Process left block features\n size_t left_size_aligned = left_size / block_size * block_size;\n for (size_t offset = 0; offset != left_size_aligned; offset += block_size) {\n matrix.template distance((const char *)data + offset, query,\n qmeta.dimension(), scores_);\n\n for (size_t i = 0; i != BATCH_SIZE; ++i) {\n uint64_t feature_key = owner_->key(feature_index);\n\n if (!this->filter()(feature_key)) {\n heap->emplace(feature_key, scores_[i], feature_index);\n }\n feature_index += 1;\n }\n }\n\n // Process left single features\n for (size_t offset = left_size_aligned; offset < left_size;\n offset += feature_size_) {\n uint64_t feature_key = owner_->key(feature_index);\n if (!this->filter()(feature_key)) {\n float score;\n matrix.template distance<1>((const char *)data + offset, query,\n qmeta.dimension(), &score);\n heap->emplace(feature_key, score, feature_index);\n }\n feature_index += 1;\n }\n heap->sort();\n return 0;\n}\n\ntemplate \nint FlatSearcherContext::search_row_nofilter(\n const void *query, const IndexQueryMeta &qmeta) {\n IndexDocumentHeap *heap = &result_heaps_[0];\n heap->clear();\n heap->limit(topk_);\n heap->set_threshold(this->threshold());\n\n size_t left_size = features_segment_->data_size();\n size_t read_offset = 0;\n size_t feature_index = 0;\n auto matrix = owner_->distance_matrix();\n\n while (left_size >= actual_read_size_) {\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, actual_read_size_) !=\n actual_read_size_) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n actual_read_size_);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < actual_read_size_;\n offset += feature_size_) {\n float score;\n matrix.template distance<1>((const char *)data + offset, query,\n qmeta.dimension(), &score);\n heap->emplace(0, score, feature_index++);\n }\n read_offset += actual_read_size_;\n left_size -= actual_read_size_;\n }\n\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, left_size) != left_size) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from features segment\",\n left_size);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < left_size; offset += feature_size_) {\n float score;\n matrix.template distance<1>((const char *)data + offset, query,\n qmeta.dimension(), &score);\n heap->emplace(0, score, feature_index++);\n }\n for (auto &it : *heap) {\n it.set_key(owner_->key(it.index()));\n }\n heap->sort();\n return 0;\n}\n\ntemplate \nint FlatSearcherContext::search_row_filter(\n const void *query, const IndexQueryMeta &qmeta) {\n IndexDocumentHeap *heap = &result_heaps_[0];\n heap->clear();\n heap->limit(topk_);\n heap->set_threshold(this->threshold());\n\n size_t left_size = features_segment_->data_size();\n size_t read_offset = 0;\n size_t feature_index = 0;\n auto matrix = owner_->distance_matrix();\n\n while (left_size >= actual_read_size_) {\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, actual_read_size_) !=\n actual_read_size_) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n actual_read_size_);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < actual_read_size_;\n offset += feature_size_) {\n uint64_t feature_key = owner_->key(feature_index);\n if (!this->filter()(feature_key)) {\n float score;\n matrix.template distance<1>((const char *)data + offset, query,\n qmeta.dimension(), &score);\n heap->emplace(feature_key, score, feature_index);\n }\n feature_index += 1;\n }\n read_offset += actual_read_size_;\n left_size -= actual_read_size_;\n }\n\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, left_size) != left_size) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from features segment\",\n left_size);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < left_size; offset += feature_size_) {\n uint64_t feature_key = owner_->key(feature_index);\n if (!this->filter()(feature_key)) {\n float score;\n matrix.template distance<1>((const char *)data + offset, query,\n qmeta.dimension(), &score);\n heap->emplace(feature_key, score, feature_index);\n }\n feature_index += 1;\n }\n heap->sort();\n return 0;\n}\n\ntemplate \nint FlatSearcherContext::batch_search_column_nofilter(\n const void *query, const IndexQueryMeta &qmeta, size_t query_count) {\n // Initialize resources\n result_heaps_.resize(query_count);\n for (auto &heap : result_heaps_) {\n heap.clear();\n heap.limit(topk_);\n heap.set_threshold(this->threshold());\n }\n\n // Transpose queries\n batch_queries_.clear();\n batch_queries_.reserve(query_count * qmeta.element_size());\n TransposeQueries(query, qmeta, query_count, &batch_queries_);\n\n size_t left_size = features_segment_->data_size();\n size_t block_size = feature_size_ * BATCH_SIZE;\n size_t read_offset = 0;\n size_t block_index = 0;\n\n // Process feature blocks\n while (left_size >= actual_read_size_) {\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, actual_read_size_) !=\n actual_read_size_) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n actual_read_size_);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < actual_read_size_; offset += block_size) {\n this->batch_enqueue_nofilter(\n (const char *)data + offset, block_index, 0, qmeta, query_count);\n block_index += BATCH_SIZE;\n }\n read_offset += actual_read_size_;\n left_size -= actual_read_size_;\n }\n\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, left_size) != left_size) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from features segment\",\n left_size);\n return IndexError_ReadData;\n }\n\n // Process left block features\n size_t left_size_aligned = left_size / block_size * block_size;\n for (size_t offset = 0; offset != left_size_aligned; offset += block_size) {\n this->batch_enqueue_nofilter(\n (const char *)data + offset, block_index, 0, qmeta, query_count);\n block_index += BATCH_SIZE;\n }\n\n // Process left single features\n for (size_t offset = left_size_aligned; offset < left_size;\n offset += feature_size_) {\n this->single_enqueue_nofilter(\n (const char *)data + offset, block_index, 0, qmeta, query_count);\n block_index += 1;\n }\n\n // Normalize results\n for (auto &heap : result_heaps_) {\n for (auto &it : heap) {\n it.set_key(owner_->key(it.index()));\n }\n heap.sort();\n }\n return 0;\n}\n\ntemplate \nint FlatSearcherContext::batch_search_column_filter(\n const void *query, const IndexQueryMeta &qmeta, size_t query_count) {\n // Initialize resources\n result_heaps_.resize(query_count);\n for (auto &heap : result_heaps_) {\n heap.clear();\n heap.limit(topk_);\n heap.set_threshold(this->threshold());\n }\n\n // Transpose queries\n batch_queries_.clear();\n batch_queries_.reserve(query_count * qmeta.element_size());\n TransposeQueries(query, qmeta, query_count, &batch_queries_);\n\n size_t left_size = features_segment_->data_size();\n size_t block_size = feature_size_ * BATCH_SIZE;\n size_t read_offset = 0;\n size_t block_index = 0;\n\n // Process feature blocks\n while (left_size >= actual_read_size_) {\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, actual_read_size_) !=\n actual_read_size_) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n actual_read_size_);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < actual_read_size_; offset += block_size) {\n size_t block_mask = 0;\n for (size_t i = 0; i != BATCH_SIZE; ++i) {\n if (!this->filter()(this->owner_->key(block_index + i))) {\n block_mask |= (1 << i);\n }\n }\n if (block_mask != 0) {\n this->batch_enqueue_filter((const char *)data + offset,\n block_index, block_mask, 0,\n qmeta, query_count);\n }\n block_index += BATCH_SIZE;\n }\n read_offset += actual_read_size_;\n left_size -= actual_read_size_;\n }\n\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, left_size) != left_size) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from features segment\",\n left_size);\n return IndexError_ReadData;\n }\n\n // Process left block features\n size_t left_size_aligned = left_size / block_size * block_size;\n for (size_t offset = 0; offset != left_size_aligned; offset += block_size) {\n size_t block_mask = 0;\n for (size_t i = 0; i != BATCH_SIZE; ++i) {\n if (!this->filter()(this->owner_->key(block_index + i))) {\n block_mask |= (1 << i);\n }\n }\n if (block_mask != 0) {\n this->batch_enqueue_filter((const char *)data + offset,\n block_index, block_mask, 0, qmeta,\n query_count);\n }\n block_index += BATCH_SIZE;\n }\n\n // Process left single features\n for (size_t offset = left_size_aligned; offset < left_size;\n offset += feature_size_) {\n if (!this->filter()(owner_->key(block_index))) {\n this->single_enqueue_nofilter(\n (const char *)data + offset, block_index, 0, qmeta, query_count);\n }\n block_index += 1;\n }\n\n // Normalize results\n for (auto &heap : result_heaps_) {\n for (auto &it : heap) {\n it.set_key(owner_->key(it.index()));\n }\n heap.sort();\n }\n return 0;\n}\n\ntemplate \nint FlatSearcherContext::batch_search_row_nofilter(\n const void *query, const IndexQueryMeta &qmeta, size_t query_count) {\n // Initialize resources\n result_heaps_.resize(query_count);\n for (auto &heap : result_heaps_) {\n heap.clear();\n heap.limit(topk_);\n heap.set_threshold(this->threshold());\n }\n\n size_t left_size = features_segment_->data_size();\n size_t read_offset = 0;\n size_t feature_index = 0;\n auto matrix = owner_->distance_matrix();\n\n // Process feature blocks\n while (left_size >= actual_read_size_) {\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, actual_read_size_) !=\n actual_read_size_) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n actual_read_size_);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < actual_read_size_;\n offset += feature_size_) {\n size_t query_offset = 0;\n const void *feature = (const char *)data + offset;\n\n for (auto &heap : result_heaps_) {\n float score;\n matrix.template distance<1>(feature, (const char *)query + query_offset,\n qmeta.dimension(), &score);\n heap.emplace(0, score, feature_index);\n query_offset += qmeta.element_size();\n }\n feature_index += 1;\n }\n read_offset += actual_read_size_;\n left_size -= actual_read_size_;\n }\n\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, left_size) != left_size) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from features segment\",\n left_size);\n return IndexError_ReadData;\n }\n\n // Process left features\n for (size_t offset = 0; offset < left_size; offset += feature_size_) {\n size_t query_offset = 0;\n const void *feature = (const char *)data + offset;\n\n for (auto &heap : result_heaps_) {\n float score;\n matrix.template distance<1>(feature, (const char *)query + query_offset,\n qmeta.dimension(), &score);\n heap.emplace(0, score, feature_index);\n query_offset += qmeta.element_size();\n }\n feature_index += 1;\n }\n\n // Normalize results\n for (auto &heap : result_heaps_) {\n for (auto &it : heap) {\n it.set_key(owner_->key(it.index()));\n }\n heap.sort();\n }\n return 0;\n}\n\ntemplate \nint FlatSearcherContext::batch_search_row_filter(\n const void *query, const IndexQueryMeta &qmeta, size_t query_count) {\n // Initialize resources\n result_heaps_.resize(query_count);\n for (auto &heap : result_heaps_) {\n heap.clear();\n heap.limit(topk_);\n heap.set_threshold(this->threshold());\n }\n\n size_t left_size = features_segment_->data_size();\n size_t read_offset = 0;\n size_t feature_index = 0;\n auto matrix = owner_->distance_matrix();\n\n // Process feature blocks\n while (left_size >= actual_read_size_) {\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, actual_read_size_) !=\n actual_read_size_) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n actual_read_size_);\n return IndexError_ReadData;\n }\n\n for (size_t offset = 0; offset < actual_read_size_;\n offset += feature_size_) {\n uint64_t feature_key = owner_->key(feature_index);\n\n if (!this->filter()(feature_key)) {\n size_t query_offset = 0;\n const void *feature = (const char *)data + offset;\n\n for (auto &heap : result_heaps_) {\n float score;\n matrix.template distance<1>(feature,\n (const char *)query + query_offset,\n qmeta.dimension(), &score);\n heap.emplace(feature_key, score, feature_index);\n query_offset += qmeta.element_size();\n }\n }\n feature_index += 1;\n }\n read_offset += actual_read_size_;\n left_size -= actual_read_size_;\n }\n\n const void *data = nullptr;\n if (features_segment_->read(read_offset, &data, left_size) != left_size) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from features segment\",\n left_size);\n return IndexError_ReadData;\n }\n\n // Process left features\n for (size_t offset = 0; offset < left_size; offset += feature_size_) {\n uint64_t feature_key = owner_->key(feature_index);\n\n if (!this->filter()(feature_key)) {\n size_t query_offset = 0;\n const void *feature = (const char *)data + offset;\n\n for (auto &heap : result_heaps_) {\n float score;\n matrix.template distance<1>(feature, (const char *)query + query_offset,\n qmeta.dimension(), &score);\n heap.emplace(feature_key, score, feature_index);\n query_offset += qmeta.element_size();\n }\n }\n feature_index += 1;\n }\n\n // Normalize results\n for (auto &heap : result_heaps_) {\n heap.sort();\n }\n return 0;\n}\n\ntemplate \nint FlatSearcherContext::group_by_search_impl(\n const void *query, const IndexQueryMeta &qmeta, uint32_t count) {\n this->resize_group_results(count);\n if (!this->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](uint64_t key) {\n return this->group_by()(key);\n };\n\n auto provider = owner_->create_provider();\n\n for (size_t q = 0; q < count; ++q) {\n this->group_topk_heaps().clear();\n\n for (node_id_t id = 0; id < provider->count(); ++id) {\n if (!this->filter().is_valid() || !this->filter()(owner_->key(id))) {\n dist_t dist = 0;\n owner_->distance_matrix().template distance<1>(\n query, provider->get_vector(owner_->key(id)), provider->dimension(),\n &dist);\n\n std::string group_id = group_by(owner_->key(id));\n auto &topk_heap = this->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(this->group_topk());\n }\n topk_heap.emplace(id, dist);\n }\n }\n this->topk_to_group_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n return 0;\n}\n\ntemplate \nint FlatSearcherContext::search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count) {\n auto provider = owner_->create_provider();\n if (this->group_by_search()) {\n this->resize_group_results(count);\n if (!this->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](uint64_t key) {\n return this->group_by()(key);\n };\n\n for (size_t q = 0; q < count; ++q) {\n this->group_topk_heaps().clear();\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!this->filter().is_valid() || !this->filter()(pk)) {\n dist_t dist = 0;\n owner_->distance_matrix().template distance<1>(\n query, provider->get_vector(pk), provider->dimension(), &dist);\n\n std::string group_id = group_by(pk);\n auto &topk_heap = this->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(this->group_topk());\n }\n topk_heap.emplace(owner_->get_id(pk), dist);\n }\n }\n this->topk_to_group_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n } else {\n result_heaps_.resize(count);\n for (auto &heap : result_heaps_) {\n heap.clear();\n heap.limit(topk_);\n heap.set_threshold(this->threshold());\n }\n for (size_t q = 0; q < count; ++q) {\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!this->filter().is_valid() || !this->filter()(pk)) {\n dist_t dist = 0;\n owner_->distance_matrix().template distance<1>(\n query, provider->get_vector(pk), provider->dimension(), &dist);\n result_heaps_[q].emplace(pk, dist, owner_->get_id(pk));\n }\n }\n query = static_cast(query) + qmeta.element_size();\n }\n for (auto &heap : result_heaps_) {\n heap.sort();\n }\n }\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_searcher_provider.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \"flat_distance_matrix.h\"\n#include \"flat_searcher.h\"\n// #include \"flat_streamer.h\"\n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Brute Force Searcher Provider\n */\ntemplate \nclass FlatSearcherProvider : public IndexProvider {\n public:\n //! Constructor\n FlatSearcherProvider(const FlatSearcher *owner) {\n feature_size_ = owner->meta().element_size();\n features_segment_ = owner->clone_features_segment();\n total_vector_count_ =\n features_segment_->data_size() / owner->meta().element_size();\n owner_ = owner;\n block_buffer_.resize(BATCH_SIZE * feature_size_);\n }\n\n //! Create a new iterator\n IndexProvider::Iterator::Pointer create_iterator(void) override {\n return IndexProvider::Iterator::Pointer(\n new (std::nothrow) FlatSearcherProvider::Iterator(owner_));\n }\n\n //! Retrieve count of vectors\n size_t count(void) const override {\n return total_vector_count_;\n }\n\n //! Retrieve dimension of vector\n size_t dimension(void) const override {\n return owner_->meta().dimension();\n }\n\n //! Retrieve type of vector\n IndexMeta::DataType data_type(void) const override {\n return owner_->meta().data_type();\n }\n\n //! Retrieve vector size in bytes\n size_t element_size(void) const override {\n return owner_->meta().element_size();\n }\n\n //! Retrieve a vector using a primary key\n const void *get_vector(uint64_t key) const override {\n return this->get_vector_by_index(owner_->get_id(key));\n }\n\n //! Retrieve the owner class\n const std::string &owner_class(void) const override {\n return owner_->name();\n }\n\n protected:\n /*! Brute Force Provider Iterator\n */\n class Iterator : public IndexProvider::Iterator {\n public:\n //! Constructor\n Iterator(const FlatSearcher *owner) {\n block_buffer_.resize(BATCH_SIZE * owner->meta().element_size());\n feature_size_ = owner->meta().element_size();\n features_segment_ = owner->clone_features_segment();\n total_vector_count_ =\n features_segment_->data_size() / owner->meta().element_size();\n owner_ = owner;\n cursor_index_ = 0;\n offset_ = 0;\n this->next_block();\n }\n\n //! Retrieve pointer of data\n //! NOTICE: the vec feature will be changed after iterating to next, so\n //! the caller need to keep a copy of it before iterator to next vector\n const void *data(void) const override {\n return data_;\n }\n\n //! Test if the iterator is valid\n bool is_valid(void) const override {\n return (!invalid_ && cursor_index_ < total_vector_count_);\n }\n\n //! Retrieve primary key\n uint64_t key(void) const override {\n return owner_->key(cursor_index_);\n }\n\n //! Next iterator\n void next(void) override {\n ++cursor_index_;\n\n if ((cursor_index_ % BATCH_SIZE) != 0) {\n data_ += feature_size_;\n } else {\n this->next_block();\n }\n }\n\n protected:\n //! Read a block of data\n void next_block(void) {\n const void *read_data = nullptr;\n size_t read_size = 0;\n\n if (cursor_index_ >= total_vector_count_) {\n invalid_ = true;\n return;\n }\n\n if (cursor_index_ + BATCH_SIZE < total_vector_count_) {\n read_size = BATCH_SIZE * feature_size_;\n } else {\n read_size = (total_vector_count_ - cursor_index_) * feature_size_;\n }\n if (features_segment_->read(offset_, &read_data, read_size) !=\n read_size) {\n LOG_ERROR(\"Failed to read data (%zu bytes) from features segment\",\n read_size);\n invalid_ = true;\n return;\n }\n offset_ += read_size;\n\n // The order of data may be a column format, convert it to the row format.\n if (owner_->column_major_order() &&\n read_size == BATCH_SIZE * feature_size_) {\n uint32_t align_size =\n IndexMeta::AlignSizeof(owner_->meta().data_type());\n ReverseTranspose(align_size, read_data,\n feature_size_ / align_size,\n &block_buffer_[0]);\n data_ = block_buffer_.data();\n } else {\n data_ = reinterpret_cast(read_data);\n }\n }\n\n private:\n const FlatSearcher *owner_{nullptr};\n IndexStorage::Segment::Pointer features_segment_{};\n uint32_t total_vector_count_{0};\n uint32_t feature_size_{0};\n std::vector block_buffer_{};\n const uint8_t *data_{nullptr};\n uint64_t offset_{0};\n uint32_t cursor_index_{0};\n bool invalid_{false};\n };\n\n //! Retrieve a vector via local index\n const void *get_vector_by_index(uint32_t index) const {\n const void *read_data = nullptr;\n if (index == kInvalidNodeId) {\n LOG_ERROR(\"Failed to get vector by Invalid Id.\");\n return nullptr;\n }\n\n if (owner_->column_major_order() &&\n index < (total_vector_count_ - (total_vector_count_ % BATCH_SIZE))) {\n uint32_t block_size = feature_size_ * BATCH_SIZE;\n uint64_t offset = (index - (index % BATCH_SIZE)) * feature_size_;\n\n if (features_segment_->read(offset, &read_data, block_size) !=\n block_size) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n block_size);\n return nullptr;\n }\n\n uint32_t align_size = IndexMeta::AlignSizeof(owner_->meta().data_type());\n ReverseTranspose(\n align_size, read_data, feature_size_ / align_size, &block_buffer_[0]);\n read_data = block_buffer_.data() + ((index % BATCH_SIZE) * feature_size_);\n\n } else {\n if (features_segment_->read(index * feature_size_, &read_data,\n feature_size_) != feature_size_) {\n LOG_ERROR(\"Failed to read data (%u bytes) from features segment\",\n feature_size_);\n return nullptr;\n }\n }\n return read_data;\n }\n\n private:\n //! Members\n const FlatSearcher *owner_{nullptr};\n IndexStorage::Segment::Pointer features_segment_{};\n uint32_t feature_size_{0};\n uint32_t total_vector_count_{0};\n mutable std::vector block_buffer_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_streamer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"flat_streamer.h\"\n#include \n#include \"flat_streamer_context.h\"\n#include \"flat_streamer_dumper.h\"\n#include \"flat_streamer_provider.h\"\n\nnamespace zvec {\nnamespace core {\n\n#define WRITE_LOCK_GUARD(MUTEX, LOCK_NAME) \\\n ailego::WriteLock write_lock(MUTEX); \\\n std::unique_lock LOCK_NAME(write_lock);\n\n#define READ_LOCK_GUARD_DEFER(MUTEX, LOCK_NAME) \\\n ailego::ReadLock read_lock(MUTEX); \\\n std::unique_lock LOCK_NAME(read_lock, std::defer_lock);\n\ntemplate \nFlatStreamer::FlatStreamer() : entity_(stats_) {}\n\ntemplate \nFlatStreamer::~FlatStreamer() {\n if (state_ == STATE_INITED) {\n this->cleanup();\n }\n}\n\ntemplate \nint FlatStreamer::init(const IndexMeta &imeta,\n const ailego::Params ¶ms) {\n meta_ = imeta;\n meta_.set_streamer(\"FlatStreamer\", 0U, params);\n\n int error_code = InitializeMetric(meta_, &metric_);\n if (error_code != 0) {\n LOG_ERROR(\"Failed to initialize index metric %s, error=%d, %s\",\n meta_.metric_name().c_str(), error_code,\n IndexError::What(error_code));\n return error_code;\n }\n if (metric_->query_metric()) {\n metric_ = metric_->query_metric();\n }\n\n // 参数设置\n if (params.get(PARAM_FLAT_COLUMN_MAJOR_ORDER, &column_major_order_)) {\n meta_.set_major_order(column_major_order_ ? IndexMeta::MO_COLUMN\n : IndexMeta::MO_ROW);\n }\n // Verify column major order\n if (meta_.major_order() != IndexMeta::MO_ROW) {\n IndexMeta::DataType ft = meta_.data_type();\n\n bool support_column_major = true;\n if ((ft != IndexMeta::DT_FP32 && ft != IndexMeta::DT_FP16 &&\n ft != IndexMeta::DT_INT8 && ft != IndexMeta::DT_INT4 &&\n ft != IndexMeta::DT_BINARY32 && ft != IndexMeta::DT_BINARY64) ||\n (meta_.unit_size() != IndexMeta::UnitSizeof(ft))) {\n if (meta_.major_order() == IndexMeta::MO_COLUMN) {\n LOG_ERROR(\"Unsupported type %d with unit size %u.\", ft,\n meta_.unit_size());\n return IndexError_Unsupported;\n } else {\n support_column_major = false;\n }\n }\n if (meta_.element_size() % IndexMeta::AlignSizeof(ft) != 0) {\n if (meta_.major_order() == IndexMeta::MO_COLUMN) {\n LOG_ERROR(\"Unsupported type %d with dimension %u.\", ft,\n meta_.dimension());\n return IndexError_Unsupported;\n } else {\n support_column_major = false;\n }\n }\n\n if (meta_.major_order() == IndexMeta::MO_UNDEFINED &&\n support_column_major) {\n meta_.set_major_order(IndexMeta::MO_ROW);\n }\n }\n\n if (!VerifyMetric(meta_)) {\n LOG_ERROR(\"Invalid index metric %s.\", meta_.metric_name().c_str());\n return IndexError_InvalidArgument;\n }\n\n read_block_size_ = FLAT_DEFAULT_READ_BLOCK_SIZE;\n params.get(PARAM_FLAT_READ_BLOCK_SIZE, &read_block_size_);\n params.get(PARAM_FLAT_USE_ID_MAP, &use_key_info_map_);\n\n // entity init\n uint32_t block_vector_count = kDefaultBlockVecCount;\n uint32_t segment_size = kDefaultSegmentSize;\n bool filter_same_key = true;\n entity_.set_block_vector_count(block_vector_count);\n entity_.set_segment_size(segment_size);\n entity_.enable_filter_same_key(filter_same_key);\n entity_.set_linear_list_count(1);\n entity_.set_use_key_info_map(use_key_info_map_);\n *entity_.mutable_meta() = meta_;\n\n state_ = STATE_INITED;\n\n return 0;\n}\n\ntemplate \nint FlatStreamer::cleanup() {\n if (state_ == STATE_OPENED) {\n this->close();\n }\n\n LOG_DEBUG(\"FlatStreamer cleanup\");\n state_ = STATE_INIT;\n return 0;\n}\n\ntemplate \nint FlatStreamer::open(IndexStorage::Pointer stg) {\n if (!stg) {\n LOG_ERROR(\"Failed to open for invalid storage\");\n return IndexError_InvalidArgument;\n }\n if (ailego_unlikely(state_ != STATE_INITED)) {\n LOG_ERROR(\"Open storage failed, init streamer first!\");\n return IndexError_NoReady;\n }\n\n LOG_DEBUG(\"FlatStreamer open with %s\", stg->name().c_str());\n\n int ret = entity_.open(std::move(stg), meta_);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to open storage\");\n return ret;\n }\n magic_ = IndexContext::GenerateMagic();\n\n state_ = STATE_OPENED;\n\n return 0;\n}\n\ntemplate \nint FlatStreamer::close(void) {\n LOG_DEBUG(\"FlatStreamer close\");\n\n entity_.flush_linear_meta();\n\n stats_.clear();\n\n int ret = entity_.close();\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n state_ = STATE_INITED;\n return 0;\n}\n\ntemplate \nint FlatStreamer::flush(uint64_t checkpoint) {\n LOG_INFO(\"FlatStreamer flush with checkpoint %zu\", (size_t)checkpoint);\n return entity_.flush(checkpoint);\n}\n\ntemplate \nint FlatStreamer::dump(const IndexDumper::Pointer &dumper) {\n std::string searcher_name = \"FlatSearcher\";\n if (BATCH_SIZE == 16) {\n searcher_name = \"FlatSearcher16\";\n }\n meta_.set_searcher(searcher_name, 0U, ailego::Params());\n WRITE_LOCK_GUARD(dump_mutex_, dump_lock);\n std::shared_ptr> bf_dumper =\n std::make_shared>(this);\n int ret = bf_dumper->dump(dumper);\n *(stats_.mutable_dumped_size()) += bf_dumper->dump_size();\n return ret;\n}\n\ntemplate \nIndexStreamer::Context::UPointer FlatStreamer::create_context(\n void) const {\n if (state_ != STATE_OPENED) {\n LOG_ERROR(\"Failed to create Context, open storage first!\");\n return Context::UPointer();\n }\n return IndexStreamer::Context::Pointer(\n new FlatStreamerContext(this));\n}\n\ntemplate \nIndexProvider::Pointer FlatStreamer::create_provider(void) const {\n return IndexProvider::Pointer(new (std::nothrow)\n FlatStreamerProvider(this));\n}\n\ntemplate \nint FlatStreamer::add_impl(uint64_t pkey, const void *query,\n const IndexQueryMeta &qmeta,\n Context::UPointer &context) {\n if (!query || qmeta.dimension() != meta_.dimension() ||\n qmeta.data_type() != meta_.data_type() ||\n qmeta.element_size() != meta_.element_size()) {\n LOG_ERROR(\n \"Failed to add for invalid arguments, query=%p, qmeta(type=%u \"\n \"dim=%u size=%u) vs meta(type=%u dim=%u size=%u)\",\n query, qmeta.data_type(), qmeta.dimension(), qmeta.element_size(),\n meta_.data_type(), meta_.dimension(), meta_.element_size());\n (*stats_.mutable_discarded_count())++;\n return IndexError_InvalidArgument;\n }\n\n auto *ctx = dynamic_cast *>(context.get());\n if (!ctx) {\n LOG_ERROR(\"Failed to cast FlatStreamerContext\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Cast;\n }\n\n READ_LOCK_GUARD_DEFER(dump_mutex_, dump_lock);\n\n if (!dump_lock.try_lock()) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n\n // IndexQueryMeta iv_qmeta;\n // int ret = entity_.convert(query, qmeta, &query, &iv_qmeta);\n // if (ret != 0) {\n // LOG_ERROR(\"Failed to convert record for %s\",\n // IndexError::What(ret));\n // (*stats_.mutable_discarded_count())++;\n // return ret;\n // }\n\n int ret = entity_.add(pkey, query, qmeta.element_size());\n if (ret != 0) {\n LOG_ERROR(\"Failed to add record for %s\", IndexError::What(ret));\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n return 0;\n}\n\ntemplate \nint FlatStreamer::add_with_id_impl(uint32_t id, const void *query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) {\n if (!query || qmeta.dimension() != meta_.dimension() ||\n qmeta.data_type() != meta_.data_type() ||\n qmeta.element_size() != meta_.element_size()) {\n LOG_ERROR(\n \"Failed to add for invalid arguments, query=%p, qmeta(type=%u \"\n \"dim=%u size=%u) vs meta(type=%u dim=%u size=%u)\",\n query, qmeta.data_type(), qmeta.dimension(), qmeta.element_size(),\n meta_.data_type(), meta_.dimension(), meta_.element_size());\n (*stats_.mutable_discarded_count())++;\n return IndexError_InvalidArgument;\n }\n\n auto *ctx = dynamic_cast *>(context.get());\n if (!ctx) {\n LOG_ERROR(\"Failed to cast FlatStreamerContext\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Cast;\n }\n\n READ_LOCK_GUARD_DEFER(dump_mutex_, dump_lock);\n\n if (!dump_lock.try_lock()) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n\n int ret = entity_.add_vector_with_id(id, query, qmeta.element_size());\n if (ret != 0) {\n LOG_ERROR(\"Failed to add record for %s\", IndexError::What(ret));\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n return 0;\n}\n\ntemplate \nint FlatStreamer::search_bf_impl(const void *query,\n const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const {\n ailego_assert(query && count && !!context);\n ailego_assert(metric_->is_matched(meta_, qmeta));\n\n FlatStreamerContext *bf_context =\n dynamic_cast *>(context.get());\n if (!bf_context) {\n LOG_ERROR(\"Invalid brute-force streamer context\");\n return IndexError_InvalidArgument;\n }\n\n if (bf_context->magic() != magic_) {\n bf_context->reset(this);\n }\n\n if (bf_context->group_by_search()) {\n return group_by_search_impl(query, qmeta, count, context);\n }\n\n bf_context->reset_results(count);\n auto &filter = bf_context->filter();\n\n for (size_t q = 0; q < count; ++q) {\n auto *heap = bf_context->result_heap();\n auto *context_stats = bf_context->mutable_stats(q);\n uint32_t scan_count = 0;\n int ret = entity_.search(query, filter, &scan_count, heap, context_stats);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to search for %s\", IndexError::What(ret));\n return ret;\n }\n heap->sort();\n bf_context->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n return 0;\n}\n\ntemplate \nint FlatStreamer::search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n ailego_assert(query && count && !!context);\n ailego_assert(metric_->is_matched(meta_, qmeta));\n\n FlatStreamerContext *bf_context =\n dynamic_cast *>(context.get());\n if (!bf_context) {\n LOG_ERROR(\"Invalid brute-force streamer context\");\n return IndexError_InvalidArgument;\n }\n\n if (bf_context->magic() != magic_) {\n bf_context->reset(this);\n }\n\n if (bf_context->group_by_search()) {\n return group_by_search_p_keys_impl(query, p_keys, qmeta, count, context);\n }\n\n bf_context->reset_results(count);\n auto &filter = bf_context->filter();\n\n for (size_t q = 0; q < count; ++q) {\n auto *heap = bf_context->result_heap();\n for (node_id_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t key = p_keys[q][idx];\n if (!filter.is_valid() || !filter(key)) {\n dist_t dist = 0;\n IndexStorage::MemoryBlock block;\n if (entity_.get_vector_by_key(key, block) != 0) continue;\n entity_.row_major_distance(query, block.data(), 1, &dist);\n heap->emplace(key, dist);\n }\n }\n heap->sort();\n bf_context->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n return 0;\n}\n\ntemplate \nint FlatStreamer::group_by_search_impl(\n const void *query, const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n FlatStreamerContext *bf_context =\n dynamic_cast *>(context.get());\n if (!bf_context) {\n LOG_ERROR(\"Invalid brute-force streamer context\");\n return IndexError_InvalidArgument;\n }\n\n bf_context->resize_group_results(count);\n if (!bf_context->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](uint64_t key) {\n return bf_context->group_by()(key);\n };\n\n auto iterator = entity_.creater_iterator();\n\n for (size_t q = 0; q < count; ++q) {\n bf_context->group_topk_heaps().clear();\n for (node_id_t id = 0; id < entity_.vector_count(); ++id) {\n uint64_t key = entity_.key(id);\n if (!bf_context->filter().is_valid() || !bf_context->filter()(key)) {\n dist_t dist = 0;\n IndexStorage::MemoryBlock block;\n if (entity_.get_vector_by_key(key, block) != 0) continue;\n entity_.row_major_distance(query, block.data(), 1, &dist);\n\n std::string group_id = group_by(key);\n auto &topk_heap = bf_context->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(bf_context->group_topk());\n }\n topk_heap.emplace(key, dist);\n }\n }\n bf_context->topk_to_group_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n return 0;\n}\n\ntemplate \nint FlatStreamer::group_by_search_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n FlatStreamerContext *bf_context =\n dynamic_cast *>(context.get());\n if (!bf_context) {\n LOG_ERROR(\"Invalid brute-force streamer context\");\n return IndexError_InvalidArgument;\n }\n\n bf_context->resize_group_results(count);\n if (!bf_context->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](uint64_t key) {\n return bf_context->group_by()(key);\n };\n\n auto iterator = entity_.creater_iterator();\n\n for (size_t q = 0; q < count; ++q) {\n bf_context->group_topk_heaps().clear();\n for (node_id_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t key = p_keys[q][idx];\n if (!bf_context->filter().is_valid() || !bf_context->filter()(key)) {\n dist_t dist = 0;\n IndexStorage::MemoryBlock block;\n if (entity_.get_vector_by_key(key, block) != 0) continue;\n entity_.row_major_distance(query, block.data(), 1, &dist);\n\n std::string group_id = group_by(key);\n auto &topk_heap = bf_context->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(bf_context->group_topk());\n }\n topk_heap.emplace(key, dist);\n }\n }\n bf_context->topk_to_group_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n return 0;\n}\n\nINDEX_FACTORY_REGISTER_STREAMER_ALIAS(LinearStreamer, FlatStreamer<32>);\nINDEX_FACTORY_REGISTER_STREAMER_ALIAS(FlatStreamer, FlatStreamer<32>);\nINDEX_FACTORY_REGISTER_STREAMER_ALIAS(FlatStreamer16, FlatStreamer<16>);\nINDEX_FACTORY_REGISTER_STREAMER_ALIAS(FlatStreamer32, FlatStreamer<32>);\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_streamer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \"flat_streamer_entity.h\"\n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Flat Streamer\n */\ntemplate \nclass FlatStreamer : public IndexStreamer {\n public:\n using ContextPointer = IndexStreamer::Context::UPointer;\n\n FlatStreamer(void);\n virtual ~FlatStreamer(void);\n\n FlatStreamer(const FlatStreamer &streamer) = delete;\n FlatStreamer &operator=(const FlatStreamer &streamer) = delete;\n\n public:\n //! Initialize Streamer\n int init(const IndexMeta &, const ailego::Params &) override;\n\n //! Cleanup Streamer\n int cleanup(void) override;\n\n //! Create a context\n IndexStreamer::Context::UPointer create_context(void) const override;\n\n //! Create a new iterator\n IndexProvider::Pointer create_provider(void) const override;\n\n //! Add a vector into index\n int add_impl(uint64_t pkey, const void *query, const IndexQueryMeta &qmeta,\n Context::UPointer &context) override;\n\n int add_with_id_impl(uint32_t id, const void *query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) override;\n\n //! Similarity search\n int search_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::UPointer &context) const override {\n return search_bf_impl(query, qmeta, 1, context);\n }\n\n //! Similarity search\n int search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::UPointer &context) const override {\n return search_bf_impl(query, qmeta, count, context);\n }\n\n //! Similarity brute force search\n int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::UPointer &context) const override {\n return search_bf_impl(query, qmeta, 1, context);\n }\n\n //! Similarity brute force search\n int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::UPointer &context) const override;\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const void *query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta,\n Context::UPointer &context) const override {\n return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context);\n }\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const void *query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::UPointer &context) const override;\n\n int group_by_search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::UPointer &context) const;\n\n int group_by_search_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const;\n\n //! Open index from file path\n int open(IndexStorage::Pointer stg) override;\n\n //! Close file\n int close(void) override;\n\n //! flush file\n int flush(uint64_t checkpoint) override;\n\n //! Dump index into storage\n int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve meta of index\n const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n const FlatStreamerEntity &entity(void) const {\n return entity_;\n }\n\n virtual const void *get_vector(uint64_t key) const override {\n return this->get_vector_by_key(key);\n }\n\n virtual int get_vector(const uint64_t key,\n IndexStorage::MemoryBlock &block) const override {\n return this->get_vector_by_key(key, block);\n }\n\n const void *get_vector_by_key(uint64_t key) const {\n return entity_.get_vector_by_key(key);\n }\n\n int get_vector_by_key(const uint64_t key,\n IndexStorage::MemoryBlock &block) const override {\n return entity_.get_vector_by_key(key, block);\n }\n const void *get_vector_by_id(uint32_t id) const override {\n return get_vector_by_key(id);\n }\n\n int get_vector_by_id(const uint32_t id,\n IndexStorage::MemoryBlock &block) const override {\n return get_vector_by_key(id, block);\n }\n\n uint32_t magic(void) const {\n return magic_;\n }\n\n //! Retrieve block size of data read\n uint32_t read_block_size(void) const {\n return read_block_size_;\n }\n\n private:\n //! Constants\n static constexpr uint32_t kDefaultBlockVecCount = 32u;\n static constexpr uint32_t kDefaultSegmentSize = 4 * 1024 * 1024u;\n static constexpr float kDefaultDocsSoftLimitRatio = 0.9f;\n\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 };\n\n //! Members\n uint32_t magic_{0};\n uint32_t docs_hard_limit_{std::numeric_limits::max()};\n uint32_t docs_soft_limit_{0};\n IndexMeta meta_{};\n std::vector> data_;\n IndexStreamer::Stats stats_{};\n IndexMetric::Pointer metric_{};\n State state_{STATE_INIT};\n mutable std::mutex mapping_mutex_{};\n ailego::SharedMutex dump_mutex_{};\n FlatStreamerEntity entity_;\n bool column_major_order_{false};\n bool use_key_info_map_{true};\n uint32_t read_block_size_{0};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_streamer_context.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \"flat_streamer.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Brute Force Streamer Context\n */\ntemplate \nclass FlatStreamerContext : public IndexStreamer::Context {\n public:\n //! Constructor\n FlatStreamerContext(const FlatStreamer *owner) {\n this->reset(owner);\n }\n\n //! Destructor\n virtual ~FlatStreamerContext(void) = default;\n\n //! Set topk of search result\n void set_topk(uint32_t topk) override {\n topk_ = topk;\n result_heap_.limit(topk);\n }\n\n //! Retrieve search result\n const IndexDocumentList &result(void) const override {\n return results_[0];\n }\n\n //! Retrieve search result with index\n const IndexDocumentList &result(size_t idx) const override {\n return results_[idx];\n }\n\n //! Retrieve result object for output\n IndexDocumentList *mutable_result(size_t idx) override {\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n return &results_[idx];\n }\n\n inline IndexDocumentHeap *result_heap() {\n return &result_heap_;\n }\n\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(void) const override {\n return group_results_[0];\n }\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(\n size_t idx) const override {\n return group_results_[idx];\n }\n\n //! Update the parameters of context\n int update(const ailego::Params & /*params*/) override {\n return 0;\n }\n\n //! Retrieve magic number\n uint32_t magic(void) const override {\n return magic_;\n }\n\n //! Get group topk\n inline uint32_t group_topk() const {\n return group_topk_;\n }\n //! Get group num\n inline uint32_t group_num() const {\n return group_num_;\n }\n inline std::map &group_topk_heaps() {\n return group_topk_heaps_;\n }\n void set_fetch_vector(bool v) override {\n fetch_vector_ = v;\n }\n bool fetch_vector() const override {\n return fetch_vector_;\n }\n inline void resize_group_results(size_t size) {\n if (group_by_search()) {\n group_results_.resize(size);\n }\n }\n\n void topk_to_result(uint32_t idx) {\n if (ailego_unlikely(result_heap_.size() == 0)) {\n return;\n }\n\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n int size = std::min(topk_, static_cast(result_heap_.size()));\n result_heap_.sort();\n results_[idx].clear();\n for (int i = 0; i < size; ++i) {\n auto score = result_heap_[i].score();\n if (score > this->threshold()) {\n break;\n }\n\n key_t key = result_heap_[i].key();\n if (fetch_vector_) {\n IndexStorage::MemoryBlock block;\n owner_->entity().get_vector_by_key(key, block);\n results_[idx].emplace_back(key, score, key, block);\n } else {\n results_[idx].emplace_back(key, score, key);\n }\n }\n }\n\n void topk_to_group_result(uint32_t idx) {\n ailego_assert_with(idx < group_results_.size(), \"invalid idx\");\n group_results_[idx].clear();\n std::vector> group_topk_list;\n std::vector> best_score_in_groups;\n for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end();\n itr++) {\n const std::string &group_id = (*itr).first;\n auto &heap = (*itr).second;\n heap.sort();\n if (heap.size() > 0) {\n float best_score = heap[0].second;\n best_score_in_groups.push_back(std::make_pair(group_id, best_score));\n }\n }\n std::sort(best_score_in_groups.begin(), best_score_in_groups.end(),\n [](const std::pair &a,\n const std::pair &b) -> int {\n return a.second < b.second;\n });\n // truncate to group num\n for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size();\n ++i) {\n const std::string &group_id = best_score_in_groups[i].first;\n group_topk_list.emplace_back(\n std::make_pair(group_id, group_topk_heaps_[group_id]));\n }\n group_results_[idx].resize(group_topk_list.size());\n for (uint32_t i = 0; i < group_topk_list.size(); ++i) {\n const std::string &group_id = group_topk_list[i].first;\n group_results_[idx][i].set_group_id(group_id);\n uint32_t size = std::min(\n group_topk_, static_cast(group_topk_list[i].second.size()));\n for (uint32_t j = 0; j < size; ++j) {\n auto score = group_topk_list[i].second[j].second;\n if (score > this->threshold()) {\n break;\n }\n node_id_t id = group_topk_list[i].second[j].first;\n auto provider = owner_->create_provider();\n if (fetch_vector_) {\n IndexStorage::MemoryBlock block;\n provider->get_vector(id, block);\n group_results_[idx][i].mutable_docs()->emplace_back(id, score, id,\n block);\n } else {\n group_results_[idx][i].mutable_docs()->emplace_back(id, score, id);\n }\n }\n }\n }\n\n //! Get if group by search\n bool group_by_search() {\n return group_num_ > 0;\n }\n //! Set group params\n void set_group_params(uint32_t group_num, uint32_t group_topk) override {\n group_num_ = group_num;\n group_topk_ = group_topk;\n group_topk_heaps_.clear();\n }\n\n void reset() override {\n for (auto &it : results_) {\n it.clear();\n }\n for (auto &it : group_results_) {\n it.clear();\n }\n }\n\n //! Reset the context\n void reset(const FlatStreamer *owner) {\n this->reset();\n magic_ = owner->magic();\n feature_size_ = owner->meta().element_size();\n\n uint32_t block_size = feature_size_ * BATCH_SIZE;\n actual_read_size_ =\n (owner->read_block_size() + block_size - 1) / block_size * block_size;\n owner_ = owner;\n }\n\n //! Reset all the query results\n void reset_results(size_t qnum) {\n results_.resize(qnum);\n stats_vec_.resize(qnum);\n for (size_t i = 0; i < qnum; ++i) {\n results_[i].clear();\n stats_vec_[i].clear();\n }\n result_heap_.clear();\n result_heap_.limit(topk_);\n result_heap_.set_threshold(this->threshold());\n }\n\n Stats *mutable_stats(size_t idx = 0) {\n ailego_assert_with(stats_vec_.size() > idx, \"invalid index\");\n return &stats_vec_[idx];\n }\n\n private:\n const FlatStreamer *owner_{nullptr};\n std::vector stats_vec_{};\n uint32_t magic_{0};\n uint32_t topk_{0};\n uint32_t feature_size_{0};\n uint32_t actual_read_size_{0};\n IndexDocumentHeap result_heap_;\n std::vector results_{};\n std::string batch_queries_{};\n float scores_[BATCH_SIZE * BATCH_SIZE];\n bool fetch_vector_{false};\n // group\n uint32_t group_num_{0};\n uint32_t group_topk_{0};\n std::map group_topk_heaps_{};\n std::vector group_results_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_streamer_dumper.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \"flat_streamer.h\"\n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\ntemplate \nclass FlatStreamerDumper {\n public:\n typedef std::unique_ptr Pointer;\n\n FlatStreamerDumper(const FlatStreamer *owner) {\n owner_ = owner;\n dump_size_ = 0;\n }\n\n int dump(const IndexDumper::Pointer &dumper) {\n ailego::ElapsedTime stamp;\n\n std::vector keys;\n if (owner_->meta().major_order() == IndexMeta::MO_COLUMN) {\n int error_code = this->write_column_index(dumper.get(), &keys);\n if (error_code != 0) {\n return error_code;\n }\n } else {\n int error_code = this->write_row_index(dumper.get(), &keys);\n if (error_code != 0) {\n return error_code;\n }\n }\n\n int error_code = this->write_keys(keys, dumper.get());\n if (error_code != 0) {\n return error_code;\n }\n\n error_code = this->write_mapping(keys, dumper.get());\n if (error_code != 0) {\n return error_code;\n }\n\n error_code = IndexHelper::SerializeToDumper(owner_->meta(), dumper.get());\n if (error_code != 0) {\n return error_code;\n }\n LOG_DEBUG(\"dumped_count: %zu, costtime: %zu\", keys.size(),\n (size_t)stamp.milli_seconds());\n return 0;\n }\n\n size_t dump_size() {\n return dump_size_;\n }\n\n private:\n int write_column_index(IndexDumper *dumper, std::vector *keys) {\n switch (IndexMeta::AlignSizeof(owner_->meta().data_type())) {\n case 2:\n return this->write_column_index(dumper, keys);\n case 4:\n return this->write_column_index(dumper, keys);\n case 8:\n return this->write_column_index(dumper, keys);\n default:\n ailego_check_with(0, \"BAD CASE\");\n }\n return IndexError_Runtime;\n }\n\n template \n int write_column_index(IndexDumper *dumper, std::vector *keys) {\n auto iter = owner_->entity().creater_iterator();\n if (!iter) {\n LOG_ERROR(\"Failed to create iterator\");\n return IndexError_Runtime;\n }\n\n // Write features\n size_t element_size = owner_->meta().element_size();\n size_t block_size = element_size * BATCH_SIZE;\n std::string block1, block2;\n block1.reserve(block_size);\n block2.reserve(block_size);\n\n for (; iter->is_valid(); iter->next()) {\n block1.append(reinterpret_cast(iter->data()), element_size);\n keys->emplace_back(iter->key());\n\n if (block1.size() == block_size) {\n ailego::MatrixHelper::Transpose(\n block1.data(), element_size / sizeof(T), (void *)block2.data());\n\n if (dumper->write(block2.data(), block_size) != block_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\",\n dumper->name().c_str());\n return IndexError_WriteData;\n }\n block1.clear();\n dump_size_ += block_size;\n }\n }\n\n if (!block1.empty()) {\n if (dumper->write(block1.data(), block1.size()) != block1.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\",\n dumper->name().c_str());\n return IndexError_WriteData;\n }\n dump_size_ += block1.size();\n }\n\n // Write the padding if need\n size_t features_size = keys->size() * element_size;\n size_t features_padding_size =\n ailego_align(features_size, 32) - features_size;\n if (features_padding_size) {\n std::string padding(features_padding_size, '\\0');\n\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\",\n dumper->name().c_str());\n return IndexError_WriteData;\n }\n dump_size_ += padding.size();\n }\n return dumper->append(FLAT_SEGMENT_FEATURES_SEG_ID, features_size,\n features_padding_size, 0);\n }\n\n int write_row_index(IndexDumper *dumper, std::vector *keys) {\n auto iter = owner_->entity().creater_iterator();\n if (!iter) {\n LOG_ERROR(\"Failed to create iterator\");\n return IndexError_Runtime;\n }\n\n // Write features\n size_t element_size = owner_->meta().element_size();\n for (; iter->is_valid(); iter->next()) {\n if (dumper->write(iter->data(), element_size) != element_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\",\n dumper->name().c_str());\n return IndexError_WriteData;\n }\n dump_size_ += element_size;\n keys->emplace_back(iter->key());\n }\n\n // Write the padding if need\n size_t features_size = keys->size() * element_size;\n size_t features_padding_size =\n ailego_align(features_size, 32) - features_size;\n if (features_padding_size) {\n std::string padding(features_padding_size, '\\0');\n\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\",\n dumper->name().c_str());\n return IndexError_WriteData;\n }\n dump_size_ += padding.size();\n }\n return dumper->append(FLAT_SEGMENT_FEATURES_SEG_ID, features_size,\n features_padding_size, 0);\n }\n\n int write_keys(const std::vector &keys, IndexDumper *dumper) {\n size_t keys_size = keys.size() * sizeof(uint64_t);\n size_t keys_padding_size = ailego_align(keys_size, 32) - keys_size;\n if (dumper->write(keys.data(), keys_size) != keys_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n dump_size_ += keys_size;\n\n // Write the padding if need\n if (keys_padding_size) {\n std::string padding(keys_padding_size, '\\0');\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\",\n dumper->name().c_str());\n return IndexError_WriteData;\n }\n dump_size_ += padding.size();\n }\n return dumper->append(FLAT_SEGMENT_KEYS_SEG_ID, keys_size,\n keys_padding_size, 0);\n }\n\n int write_mapping(const std::vector &keys, IndexDumper *dumper) {\n std::vector mapping(keys.size());\n std::iota(mapping.begin(), mapping.end(), 0);\n std::sort(mapping.begin(), mapping.end(),\n [&keys](uint32_t lhs, uint32_t rhs) {\n return (keys[lhs] < keys[rhs]);\n });\n\n size_t mapping_size = mapping.size() * sizeof(uint32_t);\n size_t mapping_padding_size = ailego_align(mapping_size, 32) - mapping_size;\n if (dumper->write(mapping.data(), mapping_size) != mapping_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n dump_size_ += mapping_size;\n\n // Write the padding if need\n if (mapping_padding_size) {\n std::string padding(mapping_padding_size, '\\0');\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\",\n dumper->name().c_str());\n return IndexError_WriteData;\n }\n dump_size_ += padding.size();\n }\n return dumper->append(FLAT_SEGMENT_MAPPING_SEG_ID, mapping_size,\n mapping_padding_size, 0);\n }\n\n private:\n const FlatStreamer *owner_{nullptr};\n size_t dump_size_{0};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_streamer_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"flat_streamer_entity.h\"\n#include \n#include \n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nFlatStreamerEntity::FlatStreamerEntity(IndexStreamer::Stats &stats)\n : stats_(stats) {}\n\nint FlatStreamerEntity::open(IndexStorage::Pointer storage,\n const IndexMeta & /*mt*/) {\n if (storage_) {\n LOG_ERROR(\"An storage instance is already opened\");\n return IndexError_Duplicate;\n }\n // segments_[0] store the meta information of the linear list\n ailego_assert_with(segments_.size() == 0, \"Invalid Size\");\n\n key_info_map_lock_ = std::make_shared();\n key_info_map_.clear();\n id_key_vector_.clear();\n withid_key_info_map_.clear();\n withid_key_map_.clear();\n\n vec_unit_size_ = IndexMeta::AlignSizeof(index_meta_.data_type());\n vec_cols_ = index_meta_.element_size() / vec_unit_size_;\n meta_.header.block_size =\n ailego_align(sizeof(BlockHeader) + sizeof(DeletionMap) +\n (index_meta_.element_size() + sizeof(uint64_t)) *\n meta_.header.block_vector_count,\n 32);\n\n if (storage->get(FLAT_LINEAR_LIST_HEAD_SEG_ID) ||\n storage->get(FLAT_LINEAR_META_SEG_ID)) {\n int ret = this->load_storage(storage);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to load storage index\");\n return ret;\n }\n } else {\n int ret = this->init_storage(storage);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to init storage\");\n return ret;\n }\n }\n\n storage_ = storage;\n\n //! Create the distance calculator\n auto metric = IndexFactory::CreateMetric(index_meta_.metric_name());\n if (!metric) {\n LOG_ERROR(\"Failed to create metric %s\", index_meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n int ret = metric->init(index_meta_, index_meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize metric %s\",\n index_meta_.metric_name().c_str());\n return ret;\n }\n row_distance_ = metric->distance();\n column_distance_ =\n metric->distance_matrix(meta_.header.block_vector_count, 1);\n\n LOG_DEBUG(\"Open storage %s done, metric=%s\", storage_->name().c_str(),\n index_meta_.metric_name().c_str());\n\n return 0;\n}\n\nint FlatStreamerEntity::close(void) {\n segments_.clear();\n storage_.reset();\n key_info_map_lock_.reset();\n key_info_map_.clear();\n withid_key_info_map_.clear();\n withid_key_map_.clear();\n id_key_vector_.clear();\n meta_.create_time = 0;\n meta_.update_time = 0;\n meta_.segment_count = 0;\n meta_.header.total_vector_count = 0;\n meta_.header.block_count = 0;\n meta_.header.block_size = 0;\n meta_.header.linear_body_size = 0;\n\n return 0;\n}\n\nint FlatStreamerEntity::flush_linear_meta(void) {\n if (!storage_) {\n return 0;\n }\n\n meta_.update_time = ailego::Realtime::Seconds();\n meta_.revision_id = stats_.revision_id();\n stats_.set_update_time(meta_.update_time);\n auto segment = storage_->get(FLAT_LINEAR_META_SEG_ID);\n if (ailego_unlikely(!segment)) {\n LOG_ERROR(\"Failed to get segment %s\", FLAT_LINEAR_META_SEG_ID.c_str());\n return IndexError_Runtime;\n }\n if (segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) {\n LOG_ERROR(\"Failed to write segment %s\", FLAT_LINEAR_META_SEG_ID.c_str());\n return IndexError_WriteData;\n }\n\n return 0;\n}\n\nint FlatStreamerEntity::flush(uint64_t checkpoint) {\n int ret = this->flush_linear_meta();\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n if (checkpoint != 0) {\n storage_->refresh(checkpoint);\n }\n ret = storage_->flush();\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to refresh storage for %s\", IndexError::What(ret));\n return ret;\n }\n if (checkpoint != 0) {\n stats_.set_check_point(checkpoint);\n }\n\n return 0;\n}\n\nint FlatStreamerEntity::add(uint64_t key, const void *vec, size_t size) {\n std::lock_guard lock(mutex_);\n if (filter_same_key_) {\n key_info_map_lock_->lock_shared();\n if (key_info_map_.find(key) != key_info_map_.end()) {\n key_info_map_lock_->unlock_shared();\n LOG_WARN(\"Try to add duplicate key, drop it\");\n return IndexError_Duplicate;\n }\n key_info_map_lock_->unlock_shared();\n }\n if (size != static_cast(index_meta_.element_size())) {\n LOG_ERROR(\"Failed to add, mismatch size %zu vs elemsize %u\", size,\n index_meta_.element_size());\n return IndexError_Mismatch;\n }\n\n IndexStorage::MemoryBlock head_block;\n this->get_head_block(head_block);\n const BlockLocation *bl =\n reinterpret_cast(head_block.data());\n if (ailego_unlikely(bl == nullptr)) {\n LOG_ERROR(\"Failed to get block loc\");\n return IndexError_ReadData;\n }\n BlockLocation block = *bl;\n\n if (!this->is_valid_block(block)) {\n int ret = this->alloc_block(block, &block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n ret = this->update_head_block(block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n }\n\n int ret = this->add_to_block(block, key, vec, size);\n if (ret == IndexError_IndexFull) {\n ret = this->alloc_block(block, &block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n ret = this->update_head_block(block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n ret = this->add_to_block(block, key, vec, size);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n }\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n (*stats_.mutable_added_count())++;\n stats_.set_revision_id(meta_.revision_id + 1);\n\n return 0;\n}\n\nint FlatStreamerEntity::search(const void *query, const IndexFilter &filter,\n uint32_t *scan_count, IndexDocumentHeap *heap,\n IndexContext::Stats *context_stats) const {\n IndexStorage::MemoryBlock head_block;\n this->get_head_block(head_block);\n const BlockLocation *bl =\n reinterpret_cast(head_block.data());\n if (ailego_unlikely(bl == nullptr)) {\n LOG_ERROR(\"Failed to get block loc\");\n return IndexError_ReadData;\n }\n\n BlockLocation block = *bl;\n\n while (this->is_valid_block(block)) {\n IndexStorage::MemoryBlock block_header_block;\n this->get_block_header(block, block_header_block);\n const BlockHeader *hd =\n reinterpret_cast(block_header_block.data());\n if (ailego_unlikely(hd == nullptr)) {\n LOG_ERROR(\"Failed to get block header\");\n return IndexError_ReadData;\n }\n\n if (hd->vector_count > 0) {\n *scan_count += hd->vector_count;\n IndexStorage::MemoryBlock deletion_map_block;\n this->get_block_deletion_map(block, deletion_map_block);\n const DeletionMap *deletion_map =\n reinterpret_cast(deletion_map_block.data());\n if (filter.is_valid() || deletion_map->is_dirty()) {\n this->search_block(query, block, hd, 1.0, filter, deletion_map, heap,\n context_stats);\n } else {\n *(context_stats->mutable_dist_calced_count()) += hd->vector_count;\n this->search_block(query, block, hd, 1.0, heap);\n }\n }\n block = hd->next;\n }\n return 0;\n}\n\n//! Search in a block\nvoid FlatStreamerEntity::search_block(const void *query,\n const BlockLocation &bl,\n const BlockHeader *hd, float norm_val,\n IndexDocumentHeap *heap) const {\n std::vector distances(block_vector_count());\n IndexStorage::MemoryBlock vecs_block;\n this->get_block_vectors(bl, vecs_block);\n const char *vecs = reinterpret_cast(vecs_block.data());\n IndexStorage::MemoryBlock keys_block;\n this->get_block_keys(bl, keys_block);\n const uint64_t *keys = reinterpret_cast(keys_block.data());\n row_major_distance(query, vecs, hd->vector_count, distances.data());\n for (size_t k = 0; k < hd->vector_count; ++k) {\n if (keys[k] != kInvalidKey) {\n heap->emplace(keys[k], distances[k] * norm_val);\n }\n }\n}\n\n//! Search in a block with filter\nvoid FlatStreamerEntity::search_block(\n const void *query, const BlockLocation &bl, const BlockHeader *hd,\n float norm_val, const IndexFilter &filter, const DeletionMap *deletion_map,\n IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const {\n std::vector distances(block_vector_count());\n\n IndexStorage::MemoryBlock vecs_block;\n this->get_block_vectors(bl, vecs_block);\n const char *vecs = reinterpret_cast(vecs_block.data());\n IndexStorage::MemoryBlock keys_block;\n this->get_block_keys(bl, keys_block);\n const uint64_t *keys = reinterpret_cast(keys_block.data());\n\n DeletionMap keeps;\n for (size_t k = 0; k < hd->vector_count; ++k) {\n const bool condition1 = !deletion_map->test(k);\n const bool condition2 = filter.is_valid() ? !filter(keys[k]) : true;\n const bool condition3 = keys[k] != kInvalidKey;\n if (condition1 && condition2 && condition3) {\n keeps.set(k);\n }\n }\n if (!keeps.is_dirty()) {\n (*context_stats->mutable_filtered_count()) += hd->vector_count;\n return;\n }\n for (size_t k = 0; k < hd->vector_count; ++k) {\n if (keeps.test(k)) {\n auto cur_vec = vecs + index_meta_.element_size() * k;\n row_major_distance(query, cur_vec, 1, distances.data() + k);\n ++(*context_stats->mutable_dist_calced_count());\n }\n }\n for (size_t k = 0; k < hd->vector_count; ++k) {\n if (keeps.test(k)) {\n heap->emplace(keys[k], distances[k] * norm_val);\n } else {\n ++(*context_stats->mutable_filtered_count());\n }\n }\n}\n\nint FlatStreamerEntity::search_bf(const void *query, const IndexFilter &filter,\n IndexDocumentHeap *heap,\n IndexContext::Stats *context_stats) const {\n uint32_t scan_count;\n return this->search(query, filter, &scan_count, heap, context_stats);\n}\n\nFlatStreamerEntity::Pointer FlatStreamerEntity::clone(void) const {\n std::vector segments;\n segments.reserve(segments_.size());\n for (size_t i = 0; i < segments_.size(); ++i) {\n segments.emplace_back(segments_[i]->clone());\n if (!segments[i]) {\n LOG_ERROR(\"Failed to clone segment, index=%zu\", i);\n return nullptr;\n }\n }\n auto entity = new (std::nothrow) FlatStreamerEntity(stats_);\n if (!entity) {\n LOG_ERROR(\"Failed to New FlatStreamerEntity object\");\n return nullptr;\n }\n entity->index_meta_ = this->index_meta_;\n entity->storage_ = this->storage_;\n // entity->reformer_ = this->reformer_;\n entity->segments_ = segments;\n entity->meta_ = this->meta_;\n entity->key_info_map_lock_ = this->key_info_map_lock_;\n entity->key_info_map_ = this->key_info_map_;\n entity->id_key_vector_ = this->id_key_vector_;\n entity->withid_key_info_map_ = this->withid_key_info_map_;\n entity->withid_key_map_ = this->withid_key_map_;\n entity->filter_same_key_ = this->filter_same_key_;\n entity->vec_unit_size_ = this->vec_unit_size_;\n entity->vec_cols_ = this->vec_cols_;\n return FlatStreamerEntity::Pointer(entity);\n}\n\nconst void *FlatStreamerEntity::get_vector_by_key(uint64_t key) const {\n VectorLocation loc{};\n key_info_map_lock_->lock_shared();\n if (use_key_info_map_) {\n auto iterator = key_info_map_.find(key);\n if (iterator == key_info_map_.end()) {\n key_info_map_lock_->unlock_shared();\n return nullptr;\n }\n loc = iterator->second;\n } else {\n if (key < withid_key_info_map_.size()) {\n loc = withid_key_info_map_[key];\n } else {\n key_info_map_lock_->unlock_shared();\n return nullptr;\n }\n }\n key_info_map_lock_->unlock_shared();\n\n auto segment = this->get_segment(loc.segment_id);\n const void *data = nullptr;\n if (segment->read(loc.offset, &data, index_meta_.element_size()) !=\n index_meta_.element_size()) {\n LOG_ERROR(\"Failed to read segment, size=%u\", index_meta_.element_size());\n return nullptr;\n }\n return data;\n}\n\nint FlatStreamerEntity::get_vector_by_key(\n const uint64_t key, IndexStorage::MemoryBlock &block) const {\n VectorLocation loc{};\n key_info_map_lock_->lock_shared();\n if (use_key_info_map_) {\n auto iterator = key_info_map_.find(key);\n if (iterator == key_info_map_.end()) {\n key_info_map_lock_->unlock_shared();\n return -1;\n }\n loc = iterator->second;\n } else {\n if (key < withid_key_info_map_.size()) {\n loc = withid_key_info_map_[key];\n } else {\n key_info_map_lock_->unlock_shared();\n return -1;\n }\n }\n key_info_map_lock_->unlock_shared();\n\n auto segment = this->get_segment(loc.segment_id);\n if (segment->read(loc.offset, block, index_meta_.element_size()) !=\n index_meta_.element_size()) {\n LOG_ERROR(\"Failed to read segment, size=%u\", index_meta_.element_size());\n return -1;\n }\n return 0;\n}\n\nIndexProvider::Iterator::Pointer FlatStreamerEntity::creater_iterator(\n void) const {\n auto entity = this->clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone entity\");\n return nullptr;\n }\n\n return Iterator::Pointer(new (std::nothrow)\n FlatStreamerEntity::Iterator(std::move(entity)));\n}\n\nvoid FlatStreamerEntity::Iterator::read_next_block(void) {\n auto block_size = entity_->linear_block_size();\n while (segment_id_ < entity_->segments_.size()) {\n auto &segment = entity_->segments_[segment_id_];\n size_t off = block_index_ * block_size;\n if (off + block_size > segment->data_size()) {\n ++segment_id_;\n block_index_ = 0;\n continue;\n }\n if (segment->read(off, block_, block_size) != block_size) {\n LOG_ERROR(\"Failed to read block, off=%zu\", off);\n break;\n }\n data_ = block_.data();\n auto hd = reinterpret_cast(\n static_cast(data_) + block_size - sizeof(BlockHeader));\n if (hd->vector_count == 0) {\n ++block_index_;\n continue;\n }\n\n block_vector_count_ = hd->vector_count;\n block_vector_index_ = 0;\n size_t elemsize = entity_->index_meta_.element_size();\n keys_ = reinterpret_cast(\n reinterpret_cast(data_) +\n elemsize * entity_->block_vector_count());\n return;\n }\n\n is_valid_ = false;\n}\n\nint FlatStreamerEntity::init_storage(IndexStorage::Pointer storage) {\n // Init Linear Meta Segment\n meta_.create_time = ailego::Realtime::Seconds();\n stats_.set_create_time(meta_.create_time);\n meta_.update_time = ailego::Realtime::Seconds();\n stats_.set_update_time(meta_.update_time);\n meta_.segment_count = 0;\n meta_.revision_id = 0;\n\n std::string str;\n index_meta_.serialize(&str);\n const size_t page = ailego::MemoryHelper::PageSize();\n\n meta_.header.header_size = sizeof(LinearIndexHeader) + str.size();\n meta_.header.total_vector_count = 0;\n meta_.header.linear_body_size = 0;\n meta_.header.block_count = 0;\n meta_.header.index_meta_size = str.size();\n meta_.header.linear_list_count = 1;\n\n AdjustSegmentSize(&meta_);\n\n LOG_DEBUG(\n \"Create Streamer Index, VecSize=%u, BlockSize=%u SegmentSize=%u \"\n \"LinearListCount=%u\",\n index_meta_.element_size(), meta_.header.block_size, meta_.segment_size,\n meta_.header.linear_list_count);\n\n size_t size = ailego_align(sizeof(meta_) + str.size(), page);\n int ret = storage->append(FLAT_LINEAR_META_SEG_ID, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to append segment %s\", FLAT_LINEAR_META_SEG_ID.c_str());\n return ret;\n }\n auto segment = storage->get(FLAT_LINEAR_META_SEG_ID);\n if (ailego_unlikely(!segment)) {\n LOG_ERROR(\"Failed to get segment %s\", FLAT_LINEAR_META_SEG_ID.c_str());\n return IndexError_Runtime;\n }\n if (segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) {\n LOG_ERROR(\"Failed to write segment data\");\n return IndexError_WriteData;\n }\n if (segment->write(sizeof(meta_), str.data(), str.size()) != str.size()) {\n LOG_ERROR(\"Failed to write segment data, size=%zu\", str.size());\n return IndexError_WriteData;\n }\n\n ret = storage->append(\"IndexMeta\", str.size());\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to append segment IndexMeta, code: %d\", ret);\n return ret;\n }\n auto index_meta_segment = storage->get(\"IndexMeta\");\n if (index_meta_segment->write(0, str.data(), str.size()) != str.size()) {\n LOG_ERROR(\"Failed to write segment data, size=%zu\", str.size());\n return IndexError_WriteData;\n }\n *stats_.mutable_index_size() += size;\n\n // Init Linear List Head Segment\n size = ailego_align(sizeof(BlockLocation) * linear_list_count(), page);\n ret = storage->append(FLAT_LINEAR_LIST_HEAD_SEG_ID, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to append segment %s for %s, size=%zu\",\n FLAT_LINEAR_LIST_HEAD_SEG_ID.c_str(), IndexError::What(ret),\n size);\n return ret;\n }\n segment = storage->get(FLAT_LINEAR_LIST_HEAD_SEG_ID);\n if (ailego_unlikely(!segment)) {\n LOG_ERROR(\"Failed to get segment %s\", FLAT_LINEAR_LIST_HEAD_SEG_ID.c_str());\n return IndexError_Runtime;\n }\n if (segment->resize(size) != size) {\n LOG_ERROR(\"Failed to resize segment, size=%zu\", size);\n return IndexError_WriteData;\n }\n segments_.emplace_back(std::move(segment));\n\n *stats_.mutable_index_size() += size;\n\n return 0;\n}\n\nint FlatStreamerEntity::load_linear_meta(IndexStorage::Pointer storage) {\n AdjustSegmentSize(&meta_);\n\n // Load Meta Segment\n auto segment = storage->get(FLAT_LINEAR_META_SEG_ID);\n if (!segment || segment->data_size() < sizeof(meta_)) {\n LOG_ERROR(\"Missing segment %s, or invalid segment size\",\n FLAT_LINEAR_META_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n IndexStorage::MemoryBlock data_block;\n if (segment->read(0, data_block, segment->data_size()) !=\n segment->data_size()) {\n LOG_ERROR(\"Failed to read storage, size=%zu\", segment->data_size());\n return IndexError_InvalidFormat;\n }\n auto *mt = reinterpret_cast(data_block.data());\n if (mt->header.block_vector_count != meta_.header.block_vector_count) {\n LOG_ERROR(\"Unmatched BlockVecCount Setting, Index %u vs Setting %u\",\n mt->header.block_vector_count, meta_.header.block_vector_count);\n return IndexError_Mismatch;\n }\n if (mt->header.block_size != meta_.header.block_size) {\n LOG_ERROR(\"Unmatched BlockSize Setting, Index %u vs Setting %u\",\n mt->header.block_size, meta_.header.block_size);\n return IndexError_Mismatch;\n }\n if (mt->header.index_meta_size + sizeof(meta_) > segment->data_size()) {\n LOG_ERROR(\"Invalid format, IndexMetaSize %u, SegmentSize %zu\",\n mt->header.index_meta_size, segment->data_size());\n return IndexError_InvalidFormat;\n }\n if (mt->header.linear_list_count != meta_.header.linear_list_count) {\n LOG_ERROR(\"Unmatch LinearListCount, Index size %u vs Setting %u\",\n mt->header.linear_list_count, meta_.header.linear_list_count);\n return IndexError_InvalidFormat;\n }\n IndexMeta index_meta;\n if (!index_meta.deserialize(mt->header.index_meta,\n mt->header.index_meta_size)) {\n LOG_ERROR(\"Failed to deserialize IndexMeta, size=%u\",\n mt->header.index_meta_size);\n return IndexError_InvalidFormat;\n }\n if (index_meta.data_type() != index_meta_.data_type() ||\n index_meta.dimension() != index_meta_.dimension() ||\n index_meta.element_size() != index_meta_.element_size() ||\n index_meta.metric_name() != index_meta_.metric_name()) {\n LOG_ERROR(\n \"Unmatch IndexMeta, Index(type=%u dim=%u elemsize=%u \"\n \"metric=%s) Setting(type=%u dim=%u elemsize=%u metric=%s)\",\n index_meta.data_type(), index_meta.dimension(),\n index_meta.element_size(), index_meta.metric_name().c_str(),\n index_meta_.data_type(), index_meta_.dimension(),\n index_meta_.element_size(), index_meta_.metric_name().c_str());\n return IndexError_Mismatch;\n }\n // Segment Size can be reconfigurable\n auto segment_size = meta_.segment_size;\n std::memcpy(&meta_, mt, sizeof(meta_));\n meta_.segment_size = segment_size;\n return 0;\n}\n\nint FlatStreamerEntity::load_segment_keys_to_map(BlockLocation block) {\n while (this->is_valid_block(block)) {\n auto segment = this->get_segment(block.segment_id);\n\n IndexStorage::MemoryBlock block_header_block;\n this->get_block_header(block, block_header_block);\n const BlockHeader *hd =\n reinterpret_cast(block_header_block.data());\n if (ailego_unlikely(hd == nullptr)) {\n LOG_ERROR(\"Failed to get block header\");\n return IndexError_ReadData;\n }\n IndexStorage::MemoryBlock keys_block;\n this->get_block_keys(block, keys_block);\n const uint64_t *keys =\n reinterpret_cast(keys_block.data());\n IndexStorage::MemoryBlock deletion_map_block;\n this->get_block_deletion_map(block, deletion_map_block);\n const DeletionMap *deletion_map =\n reinterpret_cast(deletion_map_block.data());\n\n for (uint32_t vector_index = 0; vector_index < hd->vector_count;\n ++vector_index) {\n if (deletion_map->test(vector_index)) {\n continue;\n }\n size_t vector_off =\n this->get_block_vector_offset(block.block_index, vector_index);\n key_info_map_[keys[vector_index]] =\n VectorLocation(block.segment_id, false, vector_off);\n id_key_vector_.push_back(keys[vector_index]);\n }\n block = hd->next;\n }\n return 0;\n}\n\nint FlatStreamerEntity::load_segment_keys_to_vector() {\n for (uint32_t i = 0; i < meta_.header.total_vector_count; i++) {\n size_t block_id = i / block_vector_count();\n uint32_t vector_index = i % block_vector_count();\n\n ailego_assert(segments_.size() > 1);\n size_t segment_block_count =\n segments_[1]->data_size() / linear_block_size();\n size_t segment_id = block_id / segment_block_count + 1;\n size_t real_block_id = block_id % segment_block_count;\n size_t vector_off =\n this->get_block_vector_offset(real_block_id, vector_index);\n\n withid_key_info_map_.push_back(\n VectorLocation(segment_id, false, vector_off));\n size_t key_off = get_block_key_offset(real_block_id, vector_index);\n withid_key_map_.push_back(key_off);\n }\n return 0;\n}\n\nint FlatStreamerEntity::load_storage(IndexStorage::Pointer storage) {\n int ret = this->load_linear_meta(storage);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n // Load Linear List\n auto hd_segment = storage->get(FLAT_LINEAR_LIST_HEAD_SEG_ID);\n if (ailego_unlikely(!hd_segment)) {\n LOG_ERROR(\"Failed to get segment %s\", FLAT_LINEAR_LIST_HEAD_SEG_ID.c_str());\n return IndexError_Runtime;\n }\n if (hd_segment->data_size() < linear_list_count() * sizeof(BlockLocation)) {\n LOG_ERROR(\"Invalid segment size, LinearListCount=%zu, size=%zu\",\n linear_list_count(), hd_segment->data_size());\n return IndexError_InvalidFormat;\n }\n segments_.emplace_back(hd_segment);\n\n size_t index_size = hd_segment->capacity();\n for (size_t i = 1; i <= meta_.segment_count; ++i) {\n std::string segment_id =\n ailego::StringHelper::Concat(FLAT_SEGMENT_FEATURES_SEG_ID, i);\n auto seg = storage->get(segment_id);\n if (!seg || seg->data_size() < meta_.header.block_size) {\n LOG_ERROR(\"Failed to get segment %s, or invalid segment size\",\n segment_id.c_str());\n return IndexError_InvalidFormat;\n }\n index_size += seg->capacity();\n segments_.emplace_back(std::move(seg));\n }\n\n for (size_t i = 0; i < linear_list_count(); i++) {\n IndexStorage::MemoryBlock head_block;\n this->get_head_block(head_block);\n const BlockLocation *bl =\n reinterpret_cast(head_block.data());\n if (ailego_unlikely(bl == nullptr)) {\n LOG_ERROR(\"Failed to get block loc\");\n return IndexError_ReadData;\n }\n BlockLocation block = *bl;\n if (use_key_info_map_) {\n ret = this->load_segment_keys_to_map(block);\n } else {\n ret = this->load_segment_keys_to_vector();\n }\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n }\n\n char create_time[32];\n char update_time[32];\n ailego::Realtime::Gmtime(meta_.create_time, \"%Y-%m-%d %H:%M:%S\", create_time,\n sizeof(create_time));\n ailego::Realtime::Gmtime(meta_.update_time, \"%Y-%m-%d %H:%M:%S\", update_time,\n sizeof(update_time));\n LOG_DEBUG(\n \"Load Index, IndexSize=%zu SegmentCount=%u SegmentSize=%u \"\n \"RevisionId=%zu BlockCount=%u BlockSize=%u \"\n \"BlockVectorCount=%u LinearListCount=%u TotalVecCount=%zu \"\n \"CreateTime=%s UpdateTime=%s\",\n index_size, meta_.segment_count, meta_.segment_size,\n static_cast(meta_.revision_id), meta_.header.block_count,\n meta_.header.block_size, meta_.header.block_vector_count,\n meta_.header.linear_list_count,\n static_cast(meta_.header.total_vector_count), create_time,\n update_time);\n\n stats_.set_index_size(index_size);\n stats_.set_check_point(storage->check_point());\n stats_.set_create_time(meta_.create_time);\n stats_.set_revision_id(meta_.revision_id);\n stats_.set_update_time(meta_.update_time);\n stats_.set_loaded_count(meta_.header.total_vector_count);\n\n return 0;\n}\n\nint FlatStreamerEntity::alloc_segment(void) {\n size_t index = segments_.size();\n if (index == kMaxSegmentId) {\n LOG_ERROR(\"Failed to alloc new segment, exceed max count %zu\",\n kMaxSegmentId);\n return IndexError_IndexFull;\n }\n\n std::string segment_id =\n ailego::StringHelper::Concat(FLAT_SEGMENT_FEATURES_SEG_ID, index);\n size_t size =\n ailego_align(meta_.segment_size, ailego::MemoryHelper::PageSize());\n auto segment = storage_->get(segment_id);\n if (segment) {\n if (segment->padding_size() < linear_block_size()) {\n LOG_ERROR(\n \"Unexpect segment, index=%zu, data_size=%zu \"\n \"padding_size=%zu block_size=%zu\",\n index, segment->data_size(), segment->padding_size(),\n linear_block_size());\n return IndexError_Runtime;\n }\n LOG_WARN(\"Alloc an existing segment=%s capacity=%zu\", segment_id.c_str(),\n segment->capacity());\n } else {\n int ret = storage_->append(segment_id, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to alloc segment from storage\");\n return ret;\n }\n segment = storage_->get(segment_id);\n if (ailego_unlikely(!segment)) {\n LOG_ERROR(\"Failed to get segment %s\", segment_id.c_str());\n return IndexError_Runtime;\n }\n }\n meta_.segment_count += 1;\n meta_.header.linear_body_size += size;\n segments_.emplace_back(std::move(segment));\n *stats_.mutable_index_size() += size;\n\n // Update meta information\n auto meta_segment = storage_->get(FLAT_LINEAR_META_SEG_ID);\n if (ailego_unlikely(!meta_segment)) {\n LOG_ERROR(\"Failed to get segment %s\", FLAT_LINEAR_META_SEG_ID.c_str());\n return IndexError_Runtime;\n }\n if (meta_segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) {\n LOG_ERROR(\"Failed to write meta segment\");\n return IndexError_WriteData;\n }\n\n return 0;\n}\n\nint FlatStreamerEntity::alloc_block(const BlockLocation &next,\n BlockLocation *block) {\n if (segments_.size() <= 1 ||\n segments_.back()->padding_size() < linear_block_size()) {\n int ret = this->alloc_segment();\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n }\n\n auto &segment = segments_.back();\n size_t block_index = segment->data_size() / linear_block_size();\n if (block_index == kMaxBlockId) {\n LOG_ERROR(\"Failed to alloc block, exceed max count %zu per segment\",\n kMaxBlockId);\n return IndexError_IndexFull;\n }\n\n BlockHeader header;\n header.next = next;\n header.vector_count = 0;\n header.column_major = false;\n\n size_t hd_off = segment->data_size() + linear_block_size() - sizeof(header);\n if (segment->write(hd_off, &header, sizeof(header)) != sizeof(header)) {\n LOG_ERROR(\"Failed to write block header\");\n return IndexError_WriteData;\n }\n\n size_t del_off = hd_off - sizeof(DeletionMap);\n DeletionMap reset_del_map{};\n if (segment->write(del_off, &reset_del_map, sizeof(reset_del_map)) !=\n sizeof(reset_del_map)) {\n LOG_ERROR(\"Failed to write block deletion map\");\n return IndexError_WriteData;\n }\n\n ++meta_.header.block_count;\n block->segment_id = segments_.size() - 1;\n block->block_index = (segment->data_size() / linear_block_size()) - 1;\n\n return 0;\n}\n\nint FlatStreamerEntity::add_to_block(const BlockLocation &block, uint64_t key,\n const void *data, size_t size) {\n IndexStorage::MemoryBlock block_header_block;\n this->get_block_header(block, block_header_block);\n const BlockHeader *header =\n reinterpret_cast(block_header_block.data());\n if (ailego_unlikely(header == nullptr)) {\n LOG_ERROR(\"Failed to get header\");\n return IndexError_ReadData;\n }\n\n if (header->vector_count == block_vector_count()) {\n return IndexError_IndexFull;\n }\n\n auto &segment = segments_[block.segment_id];\n\n size_t vector_off =\n get_block_vector_offset(block.block_index, header->vector_count);\n if (segment->write(vector_off, data, size) != size) {\n LOG_ERROR(\"Failed to write vector, off=%zu size=%zu\", vector_off, size);\n return IndexError_WriteData;\n }\n\n size_t key_off =\n get_block_key_offset(block.block_index, header->vector_count);\n if (segment->write(key_off, &key, sizeof(key)) != sizeof(key)) {\n LOG_ERROR(\"Failed to write key, off=%zu\", key_off);\n return IndexError_WriteData;\n }\n\n BlockHeader hd = *header;\n hd.vector_count += 1;\n size_t hd_off = get_block_header_offset(block.block_index);\n if (segment->write(hd_off, &hd, sizeof(hd)) != sizeof(hd)) {\n LOG_ERROR(\"Failed to write block header, off=%zu\", hd_off);\n return IndexError_WriteData;\n }\n\n VectorLocation loc(block.segment_id, false, vector_off);\n key_info_map_lock_->lock();\n key_info_map_[key] = loc;\n id_key_vector_.push_back(key);\n withid_key_info_map_.push_back(loc);\n withid_key_map_.push_back(key_off);\n key_info_map_lock_->unlock();\n\n ++meta_.header.total_vector_count;\n return 0;\n}\n\nint FlatStreamerEntity::add_vector_with_id(const uint32_t id, const void *query,\n const uint32_t size) {\n std::lock_guard lock(mutex_);\n // if (filter_same_key_) {\n // key_info_map_lock_->lock_shared();\n // if (key_info_map_.find(id) != key_info_map_.end()) {\n // key_info_map_lock_->unlock_shared();\n // LOG_WARN(\"Try to add duplicate key, drop it\");\n // return IndexError_Duplicate;\n // }\n // key_info_map_lock_->unlock_shared();\n // }\n\n if (size != static_cast(index_meta_.element_size())) {\n LOG_ERROR(\"Failed to add, mismatch size %u vs elemsize %u\", size,\n index_meta_.element_size());\n return IndexError_Mismatch;\n }\n\n\n if (id >= vector_count()) {\n IndexStorage::MemoryBlock head_block;\n this->get_head_block(head_block);\n BlockLocation block =\n *reinterpret_cast(head_block.data());\n if (!this->is_valid_block(block)) {\n int ret = this->alloc_block(block, &block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n ret = this->update_head_block(block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n }\n for (size_t start_id = vector_count(); start_id < id; ++start_id) {\n std::vector vec(size);\n int ret = this->add_to_block(block, kInvalidKey, vec.data(), size);\n if (ret == IndexError_IndexFull) {\n ret = this->alloc_block(block, &block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n ret = this->update_head_block(block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n ret = this->add_to_block(block, kInvalidKey, vec.data(), size);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n }\n }\n\n int ret = this->add_to_block(block, id, query, size);\n if (ret == IndexError_IndexFull) {\n ret = this->alloc_block(block, &block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n ret = this->update_head_block(block);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n ret = this->add_to_block(block, id, query, size);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n }\n } else {\n VectorLocation vector_loc = withid_key_info_map_[id];\n auto segment = this->get_segment(vector_loc.segment_id);\n size_t vector_off = vector_loc.offset;\n if (segment->write(vector_off, query, size) != size) {\n LOG_ERROR(\"Failed to write vector, off=%zu size=%u\", vector_off, size);\n return IndexError_WriteData;\n }\n size_t key_off = withid_key_map_[id];\n uint64_t key = id;\n if (segment->write(key_off, &key, sizeof(key)) != sizeof(key)) {\n LOG_ERROR(\"Failed to write key, off=%zu\", key_off);\n return IndexError_WriteData;\n }\n key_info_map_lock_->lock();\n key_info_map_[key] = vector_loc;\n key_info_map_lock_->unlock();\n }\n (*stats_.mutable_added_count())++;\n stats_.set_revision_id(meta_.revision_id + 1);\n\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_streamer_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"flat_index_format.h\"\n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Flat Streamer Entity\n */\nclass FlatStreamerEntity {\n public:\n typedef std::shared_ptr Pointer;\n\n //! Constructor\n explicit FlatStreamerEntity(IndexStreamer::Stats &stats);\n\n //! Destructor\n virtual ~FlatStreamerEntity(void) = default;\n\n //! Open the entity with storage\n int open(IndexStorage::Pointer storage, const IndexMeta &mt);\n\n //! Close the entity\n int close(void);\n\n //! Flush Linear Meta information to storage\n int flush_linear_meta(void);\n\n //! Flush linear index to storage\n int flush(uint64_t checkpoint);\n\n //! Add vector to linear index\n int add(uint64_t key, const void *vec, size_t size);\n\n //! Search in linear list with filter\n int search(const void *query, const IndexFilter &filter, uint32_t *scan_count,\n IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const;\n\n //! Search in a block\n void search_block(const void *query, const BlockLocation &bl,\n const BlockHeader *hd, float norm_val,\n IndexDocumentHeap *heap) const;\n\n //! Search in a block with filter\n void search_block(const void *query, const BlockLocation &bl,\n const BlockHeader *hd, float norm_val,\n const IndexFilter &filter, const DeletionMap *deletion_map,\n IndexDocumentHeap *heap,\n IndexContext::Stats *context_stats) const;\n\n //! Flat Search with filter\n int search_bf(const void *query, const IndexFilter &filter,\n IndexDocumentHeap *heap,\n IndexContext::Stats *context_stats) const;\n\n //! Clone the entity\n virtual FlatStreamerEntity::Pointer clone(void) const;\n\n //! Retrieve the total vectors in the index\n size_t vector_count(void) const {\n return meta_.header.total_vector_count;\n }\n\n //! Retrieve the linear list count\n size_t linear_list_count(void) const {\n return meta_.header.linear_list_count;\n }\n\n //! Retrieve block size of the linear vector\n size_t linear_block_size(void) const {\n return meta_.header.block_size;\n }\n\n //! Retrieve the vectors count in one block\n size_t block_vector_count(void) const {\n // assert(meta_.header.block_vector_count == 32);\n return meta_.header.block_vector_count;\n }\n\n //! Retrieve IndexMeta of the linear index\n const IndexMeta &meta(void) const {\n return index_meta_;\n }\n\n //! Retrieve mutable IndexMeta of the linear index\n IndexMeta *mutable_meta(void) {\n return &index_meta_;\n }\n\n //! Retrieve vector by local id\n const void *get_vector_by_key(uint64_t key) const;\n\n int get_vector_by_key(const uint64_t key,\n IndexStorage::MemoryBlock &block) const;\n\n //! Create a new iterator\n IndexProvider::Iterator::Pointer creater_iterator(void) const;\n\n\n //! Set params\n void set_block_vector_count(uint32_t count) {\n meta_.header.block_vector_count = count;\n }\n\n void set_use_key_info_map(bool use_id_map) {\n use_key_info_map_ = use_id_map;\n LOG_DEBUG(\"use_key_info_map_: %d\", (int)use_key_info_map_);\n }\n\n //! Set params\n void set_segment_size(uint32_t size) {\n meta_.segment_size = size;\n }\n\n //! Set params\n void set_linear_list_count(uint32_t count) {\n meta_.header.linear_list_count = count;\n }\n\n //! Set params\n void enable_filter_same_key(bool enabled) {\n filter_same_key_ = enabled;\n }\n\n inline uint64_t key(uint32_t id) const {\n if (id < id_key_vector_.size()) {\n return id_key_vector_[id];\n } else {\n return kInvalidKey;\n }\n }\n\n inline void row_major_distance(const void *query, const void *feature,\n size_t fnum, float *out) const {\n const uint8_t *cur_feature = reinterpret_cast(feature);\n for (size_t f = 0; f < fnum; ++f) {\n row_distance_(query, cur_feature, index_meta_.dimension(), out + f);\n cur_feature += index_meta_.element_size();\n }\n }\n\n int add_vector_with_id(const uint32_t id, const void *query,\n const uint32_t element_size);\n\n private:\n //! Disable them\n FlatStreamerEntity(const FlatStreamerEntity &) = delete;\n FlatStreamerEntity &operator=(const FlatStreamerEntity &) = delete;\n\n /*! Iterator of all the linear list\n */\n class Iterator : public IndexProvider::Iterator {\n public:\n //! Constructor\n Iterator(const FlatStreamerEntity::Pointer &entity) : entity_(entity) {\n this->read_next_block();\n }\n //! Retrieve pointer of data\n const void *data(void) const override {\n return reinterpret_cast(data_) +\n block_vector_index_ * entity_->index_meta_.element_size();\n }\n //! Test if the iterator is valid\n bool is_valid(void) const override {\n return is_valid_;\n }\n //! Retrieve primary key\n uint64_t key(void) const override {\n return keys_[block_vector_index_];\n }\n //! Next iterator\n void next(void) override {\n if (++block_vector_index_ == block_vector_count_) {\n ++block_index_;\n this->read_next_block();\n }\n }\n\n private:\n //! Read next non-empty block\n void read_next_block(void);\n\n //! Members\n std::string buffer_{};\n const FlatStreamerEntity::Pointer entity_;\n IndexStorage::MemoryBlock block_;\n const void *data_{nullptr};\n const uint64_t *keys_{nullptr};\n uint32_t segment_id_{1u}; // The first segment is header info\n uint32_t block_index_{0u};\n uint32_t block_vector_index_{0u};\n uint32_t block_vector_count_{0u};\n bool is_valid_{true};\n };\n\n //! Retrive storage segment by index\n const IndexStorage::Segment::Pointer get_segment(size_t index) const {\n for (size_t i = segments_.size(); i <= index; ++i) {\n auto segment_id =\n ailego::StringHelper::Concat(FLAT_SEGMENT_FEATURES_SEG_ID, i);\n auto segment = storage_->get(segment_id);\n if (!segment) {\n LOG_ERROR(\"Failed to get segment %s\", segment_id.c_str());\n return IndexStorage::Segment::Pointer();\n }\n segments_.emplace_back(std::move(segment));\n }\n return segments_[index];\n }\n\n //! Rejust the segment size as to aligned by page size\n void AdjustSegmentSize(StreamerLinearMeta *mt) {\n if (mt->segment_size < mt->header.block_size) {\n mt->segment_size = mt->header.block_size;\n }\n mt->segment_size = ailego_align(\n mt->segment_size / mt->header.block_size * mt->header.block_size,\n ailego::MemoryHelper::PageSize());\n }\n\n //! Init with an empty storage\n int init_storage(IndexStorage::Pointer storage);\n\n //! Load linear meta information from storage\n int load_linear_meta(IndexStorage::Pointer storage);\n\n //! Load keys to keys map\n int load_segment_keys_to_map(BlockLocation block);\n\n //! Load keys to keys map\n int load_segment_keys_to_vector(void);\n\n //! Load index from storage\n int load_storage(IndexStorage::Pointer storage);\n\n //! Check whether the block is empty\n bool is_valid_block(const BlockLocation &block) const {\n return block.segment_id != 0;\n }\n\n //! Update header block of an linear list\n int update_head_block(const BlockLocation &block) {\n ailego_assert_with(segments_.size() != 0, \"Invalid Segments\");\n\n auto &hd_segment = segments_[0];\n if (hd_segment->write(0, &block, sizeof(block)) != sizeof(block)) {\n LOG_ERROR(\"Failed to write head block location\");\n return IndexError_WriteData;\n }\n\n return 0;\n }\n\n //! Alloc a new segment\n int alloc_segment(void);\n\n //! Alloc a new block\n int alloc_block(const BlockLocation &next, BlockLocation *block);\n\n //! Add a record to a block\n int add_to_block(const BlockLocation &block, uint64_t key, const void *data,\n size_t size);\n\n private:\n size_t get_block_offset(uint32_t block_index) const {\n return block_index * linear_block_size();\n }\n\n size_t get_block_header_offset(uint32_t block_index) const {\n return get_block_offset(block_index) + linear_block_size() -\n sizeof(BlockHeader);\n }\n\n size_t get_block_deletion_map_offset(uint32_t block_index) const {\n return get_block_header_offset(block_index) - sizeof(DeletionMap);\n }\n\n size_t get_block_key_offset(uint32_t block_index,\n uint32_t vector_index) const {\n return get_block_offset(block_index) +\n block_vector_count() * index_meta_.element_size() +\n sizeof(uint64_t) * vector_index;\n }\n\n size_t get_block_vector_offset(uint32_t block_index,\n uint32_t vector_index) const {\n return this->get_block_offset(block_index) +\n vector_index * index_meta_.element_size();\n }\n\n //! Get header block of an linear list\n int get_head_block(IndexStorage::MemoryBlock &header_block) const {\n ailego_assert_with(segments_.size() != 0, \"Invalid Segments\");\n auto &hd_segment = segments_[0];\n if (hd_segment->read(0, header_block, sizeof(BlockLocation)) !=\n sizeof(BlockLocation)) {\n LOG_ERROR(\"Failed to read head block location\");\n return -1;\n }\n return 0;\n }\n\n //! Get BlockHeader of the block\n int get_block_header(const BlockLocation &block,\n IndexStorage::MemoryBlock &header_block) const {\n // The header is located in the end of a block to align features\n auto &segment = this->get_segment(block.segment_id);\n ailego_assert_with(segment != nullptr, \"Index Overflow\");\n size_t off = this->get_block_header_offset(block.block_index);\n if (segment->read(off, header_block, sizeof(BlockHeader)) !=\n sizeof(BlockHeader)) {\n LOG_ERROR(\"Failed to read block header, off=%zu\", off);\n return -1;\n }\n return 0;\n }\n int get_block_deletion_map(\n const BlockLocation &block,\n IndexStorage::MemoryBlock &deletion_map_block) const {\n auto &segment = this->get_segment(block.segment_id);\n ailego_assert_with(segment != nullptr, \"Index Overflow\");\n size_t off = this->get_block_deletion_map_offset(block.block_index);\n if (segment->read(off, deletion_map_block, sizeof(DeletionMap)) !=\n sizeof(DeletionMap)) {\n LOG_ERROR(\"Failed to read deletion map, off=%zu\", off);\n return -1;\n }\n return 0;\n }\n\n int get_block_keys(const BlockLocation &block,\n IndexStorage::MemoryBlock &keys_block) const {\n auto &segment = this->get_segment(block.segment_id);\n ailego_assert_with(segment != nullptr, \"Index Overflow\");\n size_t off = this->get_block_key_offset(block.block_index, 0);\n if (segment->read(off, keys_block,\n block_vector_count() * sizeof(uint64_t)) !=\n block_vector_count() * sizeof(uint64_t)) {\n LOG_ERROR(\"Failed to read block header, off=%zu\", off);\n return -1;\n }\n return 0;\n }\n\n int get_block_vectors(const BlockLocation &block,\n IndexStorage::MemoryBlock &vector_block) const {\n auto &segment = this->get_segment(block.segment_id);\n ailego_assert_with(segment != nullptr, \"Index Overflow\");\n size_t off = this->get_block_vector_offset(block.block_index, 0);\n if (segment->read(off, vector_block,\n block_vector_count() * index_meta_.element_size()) !=\n block_vector_count() * index_meta_.element_size()) {\n LOG_ERROR(\"Failed to read block header, off=%zu\", off);\n return -1;\n }\n return 0;\n }\n\n private:\n //! Constants\n static constexpr size_t kMaxSegmentId = std::numeric_limits::max();\n static constexpr size_t kMaxBlockId = std::numeric_limits::max();\n\n //! Members\n std::mutex mutex_{};\n IndexMeta index_meta_{};\n IndexStorage::Pointer storage_{};\n IndexMetric::MatrixDistance row_distance_{}, column_distance_{};\n mutable std::vector segments_{};\n StreamerLinearMeta meta_{};\n IndexStreamer::Stats &stats_;\n mutable std::shared_ptr key_info_map_lock_{};\n std::unordered_map key_info_map_{};\n std::vector withid_key_info_map_{};\n std::vector withid_key_map_{};\n std::vector id_key_vector_{};\n bool filter_same_key_{false};\n bool use_key_info_map_{true};\n uint32_t vec_unit_size_{0};\n uint32_t vec_cols_{0};\n mutable std::string vec_buf_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_streamer_provider.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \"flat_distance_matrix.h\"\n#include \"flat_searcher.h\"\n#include \"flat_streamer.h\"\n#include \"flat_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Brute Force Streamer Provider\n */\n\ntemplate \nclass FlatStreamerProvider : public IndexProvider {\n public:\n //! Constructor\n FlatStreamerProvider(const FlatStreamer *owner) {\n feature_size_ = owner->meta().element_size();\n total_vector_count_ = owner->entity().vector_count();\n owner_ = owner;\n block_buffer_.resize(BATCH_SIZE * feature_size_);\n }\n\n //! Create a new iterator\n IndexProvider::Iterator::Pointer create_iterator(void) override {\n return owner_->entity().creater_iterator();\n }\n\n //! Retrieve count of vectors\n size_t count(void) const override {\n return total_vector_count_;\n }\n\n //! Retrieve dimension of vector\n size_t dimension(void) const override {\n return owner_->meta().dimension();\n }\n\n //! Retrieve type of vector\n IndexMeta::DataType data_type(void) const override {\n return owner_->meta().data_type();\n }\n\n //! Retrieve vector size in bytes\n size_t element_size(void) const override {\n return owner_->meta().element_size();\n }\n\n //! Retrieve a vector using a primary key\n const void *get_vector(uint64_t key) const override {\n return this->get_vector_by_key(key);\n }\n\n int get_vector(const uint64_t key,\n IndexStorage::MemoryBlock &block) const override {\n return this->get_vector_by_key(key, block);\n }\n\n //! Retrieve the owner class\n const std::string &owner_class(void) const override {\n return owner_->name();\n }\n\n protected:\n //! Retrieve a vector via primary key\n const void *get_vector_by_key(uint64_t key) const {\n return owner_->get_vector_by_key(key);\n }\n\n int get_vector_by_key(const uint64_t key,\n IndexStorage::MemoryBlock &block) const {\n return owner_->get_vector_by_key(key, block);\n }\n\n private:\n //! Members\n const FlatStreamer *owner_{nullptr};\n IndexStorage::Segment::Pointer features_segment_{};\n uint32_t feature_size_{0};\n uint32_t total_vector_count_{0};\n mutable std::vector block_buffer_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat/flat_utility.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\n//! The default size of reading a block\nstatic constexpr uint32_t FLAT_DEFAULT_READ_BLOCK_SIZE = 4 * 1024 * 1024;\nstatic const std::string FLAT_LINEAR_META_SEG_ID = \"flat.linear_meta\";\nstatic const std::string FLAT_LINEAR_LIST_HEAD_SEG_ID = \"flat.linear_list_head\";\n\nstatic const std::string FLAT_SEGMENT_KEYS_SEG_ID(\"flat.keys\");\nstatic const std::string FLAT_SEGMENT_FEATURES_SEG_ID(\"flat.features\");\nstatic const std::string FLAT_SEGMENT_MAPPING_SEG_ID(\"flat.mapping\");\n\n// index params\nstatic const std::string PARAM_FLAT_COLUMN_MAJOR_ORDER(\n \"proxima.flat.column_major_order\");\nstatic const std::string PARAM_FLAT_BATCH_SIZE(\"proxima.flat.batch_size\");\nstatic const std::string PARAM_FLAT_READ_BLOCK_SIZE(\n \"proxima.flat.read_block_size\");\nstatic const std::string PARAM_FLAT_USE_ID_MAP(\"proxima.flat.use_id_map\");\n\n//! Determines if a number is equal to two to the power of n.\ntemplate \nstruct IsEqualPowerofTwo\n : std::integral_constant {};\n\n//! Transpose a block\ntemplate \nstatic inline void ReverseTranspose(size_t align_size, const void *src,\n size_t dim, void *dst) {\n switch (align_size) {\n case 2:\n ailego::MatrixHelper::ReverseTranspose(src, dim, dst);\n break;\n case 4:\n ailego::MatrixHelper::ReverseTranspose(src, dim, dst);\n break;\n case 8:\n ailego::MatrixHelper::ReverseTranspose(src, dim, dst);\n break;\n }\n}\n\nstatic inline void ReverseTranspose(size_t align_size, const void *src,\n size_t m, size_t dim, void *dst) {\n switch (align_size) {\n case 2:\n ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst);\n break;\n case 4:\n ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst);\n break;\n case 8:\n ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst);\n break;\n }\n}\n\ntemplate \nstatic inline void TransposeOne(const void *src, size_t M, size_t N,\n void *dst) {\n for (size_t i = 0; i < N; ++i) {\n reinterpret_cast(dst)[i] = reinterpret_cast(src)[i * M];\n }\n}\n\nstatic inline void Transpose(size_t align_size, const void *src, size_t m,\n size_t dim, void *dst) {\n switch (align_size) {\n case 2:\n ailego::MatrixHelper::Transpose(src, m, dim, dst);\n break;\n case 4:\n ailego::MatrixHelper::Transpose(src, m, dim, dst);\n break;\n case 8:\n ailego::MatrixHelper::Transpose(src, m, dim, dst);\n break;\n }\n}\n\n//! Transpose queries\ntemplate \nvoid TransposeQueries(const void *query, const IndexQueryMeta &qmeta,\n size_t query_count, std::string *out) {\n if (K <= 1) {\n ailego_assert(query_count == 1);\n (void)query_count;\n out->append(reinterpret_cast(query) + out->size(),\n qmeta.element_size());\n } else {\n ailego_assert_with(IsEqualPowerofTwo::value,\n \"K must be equal to two to the power of n.\");\n\n size_t query_batch_count = query_count / K;\n size_t query_offset = out->size();\n out->resize(query_offset + query_batch_count * K * qmeta.element_size());\n\n switch (IndexMeta::AlignSizeof(qmeta.data_type())) {\n case 2:\n for (size_t i = 0; i != query_batch_count; ++i) {\n ailego::MatrixHelper::Transpose(\n (const char *)query + query_offset,\n qmeta.element_size() / sizeof(uint16_t), &((*out)[query_offset]));\n query_offset += qmeta.element_size() * K;\n }\n break;\n\n case 4:\n for (size_t i = 0; i != query_batch_count; ++i) {\n ailego::MatrixHelper::Transpose(\n (const char *)query + query_offset,\n qmeta.element_size() / sizeof(uint32_t), &((*out)[query_offset]));\n\n query_offset += qmeta.element_size() * K;\n }\n break;\n\n case 8:\n for (size_t i = 0; i != query_batch_count; ++i) {\n ailego::MatrixHelper::Transpose(\n (const char *)query + query_offset,\n qmeta.element_size() / sizeof(uint64_t), &((*out)[query_offset]));\n query_offset += qmeta.element_size() * K;\n }\n break;\n\n default:\n ailego_check_with(0, \"BAD CASE\");\n }\n size_t query_left_count = query_count % K;\n if (query_left_count != 0) {\n TransposeQueries<(K >> 1)>(query, qmeta, query_left_count, out);\n }\n }\n}\n\n//! Create and initialize measure\nstatic inline int InitializeMetric(const IndexMeta &mt,\n IndexMetric::Pointer *out) {\n IndexMetric::Pointer measure = IndexFactory::CreateMetric(mt.metric_name());\n if (!measure) {\n return IndexError_NoExist;\n }\n\n int error_code = measure->init(mt, mt.metric_params());\n if (error_code != 0) {\n return error_code;\n }\n *out = measure;\n return 0;\n}\n\n//! Verify measure\nstatic inline bool VerifyMetric(const IndexMeta &meta) {\n IndexMetric::Pointer measure = IndexFactory::CreateMetric(meta.metric_name());\n if (!measure) {\n return false;\n }\n int error_code = measure->init(meta, meta.metric_params());\n if (error_code != 0) {\n return false;\n }\n return true;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\ncc_library(\n NAME core_knn_flat_sparse \n STATIC SHARED STRICT ALWAYS_LINK\n SRCS *.cc\n LIBS core_framework \n INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_builder.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"flat_sparse_builder.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"flat_sparse_index_format.h\"\n#include \"flat_sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nFlatSparseBuilder::FlatSparseBuilder() {}\n\nint FlatSparseBuilder::init(const IndexMeta &meta,\n const ailego::Params & /*params*/) {\n LOG_INFO(\"Begin FlatSparseBuilder::init\");\n\n meta_ = meta;\n\n state_ = BUILD_STATE_INITED;\n LOG_INFO(\"End FlatSparseBuilder::init\");\n return 0;\n}\n\nint FlatSparseBuilder::cleanup(void) {\n LOG_INFO(\"Begin FlatSparseBuilder::cleanup\");\n\n stats_.clear_attributes();\n stats_.set_trained_count(0UL);\n stats_.set_built_count(0UL);\n stats_.set_dumped_count(0UL);\n stats_.set_discarded_count(0UL);\n stats_.set_trained_costtime(0UL);\n stats_.set_built_costtime(0UL);\n stats_.set_dumped_costtime(0UL);\n state_ = BUILD_STATE_INIT;\n\n LOG_INFO(\"End FlatSparseBuilder::cleanup\");\n\n return 0;\n}\n\nint FlatSparseBuilder::train(IndexThreads::Pointer,\n IndexSparseHolder::Pointer /*holder*/) {\n if (state_ != BUILD_STATE_INITED) {\n LOG_ERROR(\"Init the builder before FlatSparseBuilder::train\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"Begin FlatSparseBuilder::train\");\n\n stats_.set_trained_count(0UL);\n stats_.set_trained_costtime(0UL);\n state_ = BUILD_STATE_TRAINED;\n\n LOG_INFO(\"End FlatSparseBuilder::train\");\n\n return 0;\n}\n\nint FlatSparseBuilder::train(const IndexTrainer::Pointer & /*trainer*/) {\n if (state_ != BUILD_STATE_INITED) {\n LOG_ERROR(\"Init the builder before FlatSparseBuilder::train\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"Begin FlatSparseBuilder::train by trainer\");\n\n stats_.set_trained_count(0UL);\n stats_.set_trained_costtime(0UL);\n state_ = BUILD_STATE_TRAINED;\n\n LOG_INFO(\"End FlatSparseBuilder::train by trainer\");\n\n return 0;\n}\n\nint FlatSparseBuilder::build(IndexThreads::Pointer,\n IndexSparseHolder::Pointer holder) {\n LOG_INFO(\"Begin FlatSparseBuilder::build\");\n\n ailego::ElapsedTime stamp;\n if (!holder) {\n LOG_ERROR(\"Input holder is nullptr while building index\");\n return IndexError_InvalidArgument;\n }\n\n if (!holder->is_matched(meta_)) {\n LOG_ERROR(\"Input holder doesn't match index meta while building index\");\n return IndexError_Mismatch;\n }\n\n holder_ = std::move(holder);\n\n stats_.set_built_count(holder_->count());\n stats_.set_built_costtime(stamp.milli_seconds());\n state_ = BUILD_STATE_BUILT;\n\n LOG_INFO(\"End FlatSparseBuilder::build\");\n return 0;\n}\n\nint FlatSparseBuilder::dump(const IndexDumper::Pointer &dumper) {\n if (state_ != BUILD_STATE_BUILT || !holder_) {\n LOG_INFO(\"Build the index before FlatSparseBuilder::dump\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"Begin FlatSparseBuilder::dump\");\n\n auto start_time = ailego::Monotime::MilliSeconds();\n\n int ret = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return ret;\n }\n\n uint32_t dump_count;\n ret = do_dump(dumper, &dump_count);\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump index\");\n return ret;\n }\n\n holder_ = nullptr;\n stats_.set_dumped_count(dump_count);\n stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time);\n\n LOG_INFO(\"End FlatSparseBuilder::dump\");\n return 0;\n}\n\nint FlatSparseBuilder::do_dump(const IndexDumper::Pointer &dumper,\n uint32_t *dump_count) {\n // bf meta\n int ret = dump_meta(dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump meta\");\n return ret;\n }\n\n std::vector keys;\n ret = dump_vector_and_offset(dumper.get(), &keys);\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump offset data\");\n return ret;\n }\n\n ret = dump_keys(keys, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump keys\");\n return ret;\n }\n\n ret = dump_mapping(keys, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump mapping\");\n return ret;\n }\n\n *dump_count = keys.size();\n\n return 0;\n}\n\nint FlatSparseBuilder::dump_meta(IndexDumper *dumper) {\n FlatSparseMeta meta;\n meta.create_time = ailego::Realtime::Seconds();\n meta.update_time = ailego::Realtime::Seconds();\n meta.doc_cnt = holder_->count();\n\n if (dumper->write(&meta, sizeof(meta)) != sizeof(meta)) {\n LOG_ERROR(\"Failed to write meta\");\n return IndexError_WriteData;\n }\n\n size_t meta_padding_size = ailego_align(sizeof(meta), 32) - sizeof(meta);\n if (meta_padding_size) {\n std::string padding(meta_padding_size, '\\0');\n if (dumper->write(padding.data(), meta_padding_size) != meta_padding_size) {\n LOG_ERROR(\"Failed to write meta padding\");\n return IndexError_WriteData;\n }\n }\n return dumper->append(PARAM_FLAT_SPARSE_META_SEG_ID, sizeof(meta),\n meta_padding_size, 0);\n}\n\nint FlatSparseBuilder::dump_vector_and_offset(IndexDumper *dumper,\n std::vector *keys) {\n // iterate the holder\n auto iter = holder_->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Failed to create iterator\");\n return IndexError_Runtime;\n }\n\n uint64_t written_length{0U};\n\n std::vector> offset_lens;\n while (iter->is_valid()) {\n keys->push_back(iter->key());\n\n uint32_t length;\n if (write_vector_data(iter->sparse_count(), iter->sparse_indices(),\n iter->sparse_data(), dumper, &length) != 0) {\n return IndexError_WriteData;\n }\n\n offset_lens.push_back({written_length, length});\n written_length += length;\n iter->next();\n }\n\n if (dumper->append(PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID, written_length, 0,\n 0) != 0) {\n LOG_ERROR(\"Failed to append offset data\");\n return IndexError_WriteData;\n }\n\n LOG_DEBUG(\"Data total written: %zu\", (size_t)written_length);\n\n for (auto &offset_len : offset_lens) {\n if (dumper->write(&offset_len.first, sizeof(offset_len.first)) !=\n sizeof(offset_len.first)) {\n LOG_ERROR(\"Failed to write offset\");\n return IndexError_WriteData;\n }\n\n if (dumper->write(&offset_len.second, sizeof(offset_len.second)) !=\n sizeof(offset_len.second)) {\n LOG_ERROR(\"Failed to write length\");\n return IndexError_WriteData;\n }\n }\n\n if (dumper->append(PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID,\n offset_lens.size() * (sizeof(uint64_t) + sizeof(uint32_t)),\n 0, 0) != 0) {\n LOG_ERROR(\"Failed to append offset data\");\n return IndexError_WriteData;\n }\n\n LOG_DEBUG(\"Offset total written: %zu\",\n offset_lens.size() * (sizeof(uint64_t) + sizeof(uint32_t)));\n\n return 0;\n}\n\nint FlatSparseBuilder::write_vector_data(const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_vec,\n IndexDumper *dumper,\n uint32_t *length) {\n std::string sparse_buffer;\n\n SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_vec,\n meta_.unit_size(), sparse_buffer);\n\n if (dumper->write(sparse_buffer.data(), sparse_buffer.size()) !=\n sparse_buffer.size()) {\n LOG_ERROR(\"Failed to write sparse data\");\n return IndexError_WriteData;\n }\n\n *length = sparse_buffer.size();\n\n return 0;\n}\n\nint FlatSparseBuilder::dump_keys(const std::vector &keys,\n IndexDumper *dumper) {\n size_t keys_size = keys.size() * sizeof(uint64_t);\n if (dumper->write(keys.data(), keys_size) != keys_size) {\n LOG_ERROR(\"Failed to write keys to dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n size_t keys_padding_size = ailego_align(keys_size, 32) - keys_size;\n if (keys_padding_size) {\n std::string padding(keys_padding_size, '\\0');\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write padding to dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n }\n return dumper->append(PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID, keys_size,\n keys_padding_size, 0);\n}\n\nint FlatSparseBuilder::dump_mapping(const std::vector &keys,\n IndexDumper *dumper) {\n std::vector mapping(keys.size());\n std::iota(mapping.begin(), mapping.end(), 0);\n std::sort(\n mapping.begin(), mapping.end(),\n [&keys](uint32_t lhs, uint32_t rhs) { return (keys[lhs] < keys[rhs]); });\n\n size_t mapping_size = mapping.size() * sizeof(uint32_t);\n size_t mapping_padding_size = ailego_align(mapping_size, 32) - mapping_size;\n if (dumper->write(mapping.data(), mapping_size) != mapping_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n\n // Write the padding if need\n if (mapping_padding_size) {\n std::string padding(mapping_padding_size, '\\0');\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n }\n return dumper->append(PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID, mapping_size,\n mapping_padding_size, 0);\n}\n\nINDEX_FACTORY_REGISTER_BUILDER(FlatSparseBuilder);\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_builder.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\n/*! Brute Force Sparse Builder\n */\nclass FlatSparseBuilder : public IndexBuilder {\n public:\n //! Constructor\n FlatSparseBuilder();\n\n //! Initialize the builder\n int init(const IndexMeta &meta, const ailego::Params ¶ms) override;\n\n //! Cleanup the builder\n int cleanup(void) override;\n\n //! Train the data\n int train(IndexThreads::Pointer, IndexSparseHolder::Pointer holder) override;\n\n //! Train the data\n int train(const IndexTrainer::Pointer &trainer) override;\n\n int train(IndexThreads::Pointer /*threads*/,\n IndexHolder::Pointer /*holder*/) override {\n return IndexError_NotImplemented;\n }\n\n int build(IndexThreads::Pointer /*threads*/,\n IndexHolder::Pointer /*holder*/) override {\n return IndexError_NotImplemented;\n }\n\n //! Build the index\n int build(IndexThreads::Pointer threads,\n IndexSparseHolder::Pointer holder) override;\n\n //! Dump index into storage\n int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n const Stats &stats(void) const override {\n return stats_;\n }\n\n private:\n int do_dump(const IndexDumper::Pointer &dumper, uint32_t *dump_count);\n\n int dump_meta(IndexDumper *dumper);\n\n int dump_keys(const std::vector &keys, IndexDumper *dumper);\n\n int dump_mapping(const std::vector &keys, IndexDumper *dumper);\n\n int dump_vector_and_offset(IndexDumper *dumper, std::vector *keys);\n\n int write_vector_data(const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_vec,\n IndexDumper *dumper, uint32_t *length);\n\n private:\n enum BUILD_STATE {\n BUILD_STATE_INIT = 0,\n BUILD_STATE_INITED = 1,\n BUILD_STATE_TRAINED = 2,\n BUILD_STATE_BUILT = 3\n };\n\n IndexSparseHolder::Pointer holder_{};\n\n std::atomic_bool error_{false};\n IndexMeta meta_{};\n IndexMetric::Pointer measure_{};\n std::mutex mutex_{};\n std::condition_variable cond_{};\n Stats stats_{};\n\n BUILD_STATE state_{BUILD_STATE_INIT};\n};\n\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_context.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"flat_sparse_context.h\"\n\nnamespace zvec {\nnamespace core {\n\nconst FlatSparseEntity *FlatSparseContext::entity() const {\n if (context_type_ == kStreamerContext) {\n return &streamer_owner_->entity();\n } else if (context_type_ == kSearcherContext) {\n return &searcher_owner_->entity();\n }\n return nullptr;\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_context.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \"flat_sparse_entity.h\"\n#include \"flat_sparse_searcher.h\"\n#include \"flat_sparse_streamer.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass FlatSparseStreamer;\nclass FlatSparseSearcher;\n\n/*! Brute Force Sparse Streamer Context\n */\nclass FlatSparseContext : public IndexContext {\n public:\n //! Constructor\n enum ContextType {\n kUnknownContext = 0,\n kSearcherContext = 1,\n kStreamerContext = 3\n };\n FlatSparseContext(const FlatSparseStreamer *streamer_ptr)\n : streamer_owner_(streamer_ptr), context_type_(kStreamerContext) {}\n\n FlatSparseContext(const FlatSparseSearcher *searcher_ptr)\n : searcher_owner_(searcher_ptr), context_type_(kSearcherContext) {}\n\n //! Destructor\n virtual ~FlatSparseContext(void) = default;\n\n //! Set topk of search result\n void set_topk(uint32_t topk) override {\n topk_ = topk;\n result_heap_.limit(topk_);\n result_heap_.set_threshold(this->threshold());\n }\n\n //! Retrieve search result\n const IndexDocumentList &result(void) const override {\n return results_.at(0);\n }\n\n //! Retrieve search result with index\n const IndexDocumentList &result(size_t index) const override {\n return results_.at(index);\n }\n\n //! Retrieve result object for output\n IndexDocumentList *mutable_result(size_t idx) override {\n return &results_.at(idx);\n }\n\n inline IndexDocumentHeap *result_heap() {\n return &result_heap_;\n }\n\n //! Update the parameters of context\n int update(const ailego::Params & /*params*/) override {\n return 0;\n }\n\n //! Retrieve magic number\n uint32_t magic(void) const override {\n return magic_;\n }\n\n void set_fetch_vector(bool v) override {\n fetch_vector_ = v;\n }\n\n bool fetch_vector() const override {\n return fetch_vector_;\n }\n\n //! Retrieve search group result with index\n const IndexGroupDocumentList &group_result(void) const override {\n return group_results_[0];\n }\n\n //! Retrieve search group result with index\n const IndexGroupDocumentList &group_result(size_t idx) const override {\n return group_results_[idx];\n }\n\n IndexGroupDocumentList *mutable_group_result(size_t idx) {\n return &group_results_[idx];\n }\n\n //! Set group params\n void set_group_params(uint32_t group_num, uint32_t group_topk) override {\n group_num_ = group_num;\n group_topk_ = group_topk;\n result_group_heap_.clear();\n }\n\n //! Get if group by search\n inline bool group_by_search() {\n return group_num_ > 0;\n }\n\n inline uint32_t group_topk() const {\n return group_topk_;\n }\n\n inline uint32_t group_num() const {\n return group_num_;\n }\n\n void reset() override {}\n\n //! Reset the context\n void reset(const FlatSparseStreamer *streamer_ptr) {\n magic_ = streamer_ptr->magic();\n streamer_owner_ = streamer_ptr;\n context_type_ = kStreamerContext;\n }\n\n void reset(const FlatSparseSearcher *searcher_ptr) {\n magic_ = searcher_ptr->magic();\n searcher_owner_ = searcher_ptr;\n context_type_ = kSearcherContext;\n }\n\n //! Reset all the query results\n void reset_results(size_t qnum) {\n if (group_by_search()) {\n group_results_.resize(qnum);\n } else {\n result_heap_.clear();\n result_heap_.limit(topk_);\n result_heap_.set_threshold(this->threshold());\n results_.resize(qnum);\n stats_vec_.resize(qnum);\n for (size_t i = 0; i < results_.size(); ++i) {\n results_[i].clear();\n stats_vec_[i].clear();\n }\n }\n }\n\n Stats *mutable_stats(size_t idx = 0) {\n ailego_assert_with(stats_vec_.size() > idx, \"invalid index\");\n return &stats_vec_[idx];\n }\n\n inline void topk_to_result(uint32_t idx) {\n if (ailego_unlikely(result_heap_.size() == 0)) {\n return;\n }\n\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n int size = std::min(topk_, static_cast(result_heap_.size()));\n result_heap_.sort();\n results_[idx].clear();\n for (int i = 0; i < size; ++i) {\n auto score = result_heap_[i].score();\n if (score > this->threshold()) {\n break;\n }\n\n key_t key = result_heap_[i].key();\n if (fetch_vector_) {\n node_id_t id = entity()->get_id(key);\n IndexStorage::MemoryBlock vec_block;\n entity()->get_sparse_vector(id, vec_block);\n const void *sparse_data = vec_block.data();\n IndexSparseDocument sparse_doc;\n if (sparse_data != nullptr) {\n SparseUtility::ReverseSparseFormat(sparse_data, sparse_doc,\n entity()->sparse_unit_size());\n }\n results_[idx].emplace_back(key, score, id, nullptr, sparse_doc);\n } else {\n results_[idx].emplace_back(key, score);\n }\n }\n }\n\n private:\n const FlatSparseEntity *entity() const;\n\n private:\n const FlatSparseStreamer *streamer_owner_{nullptr};\n const FlatSparseSearcher *searcher_owner_{nullptr};\n ContextType context_type_{kUnknownContext};\n std::vector stats_vec_{};\n uint32_t magic_{0};\n uint32_t topk_{0};\n IndexDocumentHeap result_heap_;\n // std::string batch_queries_{};\n bool fetch_vector_{false};\n\n // group\n uint32_t group_num_{0};\n uint32_t group_topk_{0};\n std::map result_group_heap_{};\n std::vector results_{};\n std::vector group_results_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \"flat_sparse_index_format.h\"\n\nnamespace zvec {\nnamespace core {\n\nusing node_id_t = uint32_t;\nconstexpr node_id_t kInvalidNodeId = static_cast(-1);\n\n/*! Flat Sparse Entity\n */\nclass FlatSparseEntity {\n public:\n typedef std::shared_ptr Pointer;\n\n //! Constructor\n explicit FlatSparseEntity() {}\n\n //! Destructor\n ~FlatSparseEntity() = default;\n\n //! Disable them\n FlatSparseEntity(const FlatSparseEntity &) = delete;\n FlatSparseEntity &operator=(const FlatSparseEntity &) = delete;\n\n //! Search in linear list with filter\n int search(const std::string &sparse_vector, const IndexFilter &filter,\n IndexDocumentHeap *heap) const {\n for (node_id_t i = 0; i < doc_cnt(); i++) {\n uint64_t key = get_key(i);\n if (ailego_unlikely(key == kInvalidKey)) {\n // LOG_ERROR(\"The key of node_id[%u] not found in keys map\", i);\n // return IndexError_Runtime;\n continue;\n }\n if (!filter.is_valid() || !filter(key)) {\n float dist = get_search_distance(sparse_vector, i);\n heap->emplace(key, dist);\n }\n }\n\n return 0;\n }\n\n //! Search in linear list with filter and target pkeys\n int search_p_keys(const std::string &sparse_vector,\n const std::vector &p_keys,\n const IndexFilter &filter, IndexDocumentHeap *heap) const {\n for (auto p_key : p_keys) {\n if (!filter.is_valid() || !filter(p_key)) {\n auto node_id = get_id(p_key);\n if (node_id != kInvalidNodeId) {\n float dist = get_search_distance(sparse_vector, node_id);\n heap->emplace(p_key, dist);\n }\n }\n }\n\n return 0;\n }\n\n //! Group search in linear list with filter\n int search_group(\n const std::string &sparse_vector, const IndexFilter &filter,\n const std::function &group_by_func, uint32_t topk,\n std::unordered_map *heap) const {\n for (node_id_t i = 0; i < doc_cnt(); i++) {\n uint64_t key = get_key(i);\n if (ailego_unlikely(key == kInvalidKey)) {\n LOG_ERROR(\"The key of node_id[%u] not found in keys map\", i);\n return IndexError_Runtime;\n }\n if (!filter.is_valid() || !filter(key)) {\n float dist = get_search_distance(sparse_vector, i);\n\n std::string group_id = group_by_func(key);\n\n auto &group_heap = (*heap)[group_id];\n if (group_heap.empty()) {\n group_heap.limit(topk);\n }\n group_heap.emplace(key, dist);\n }\n }\n\n return 0;\n }\n\n //! Group search in linear list with filter and target pkeys\n int search_group_p_keys(\n const std::string &sparse_vector, const std::vector &p_keys,\n const IndexFilter &filter,\n const std::function &group_by_func, uint32_t topk,\n std::unordered_map *heap) const {\n for (auto p_key : p_keys) {\n if (!filter.is_valid() || !filter(p_key)) {\n auto node_id = get_id(p_key);\n if (node_id != kInvalidNodeId) {\n float dist = get_search_distance(sparse_vector, node_id);\n\n std::string group_id = group_by_func(p_key);\n\n auto &group_heap = (*heap)[group_id];\n if (group_heap.empty()) {\n group_heap.limit(topk);\n }\n group_heap.emplace(p_key, dist);\n }\n }\n }\n\n return 0;\n }\n\n //! Get sparse vector by key\n int get_sparse_vector(uint64_t key, std::string *sparse_vector) const {\n const void *sparse_vector_ptr;\n uint32_t sparse_vector_len;\n int ret = get_sparse_vector_ptr_by_key(key, &sparse_vector_ptr,\n &sparse_vector_len);\n if (ret != 0) {\n return ret;\n }\n *sparse_vector = std::string(static_cast(sparse_vector_ptr),\n sparse_vector_len);\n return 0;\n }\n\n //! Get sparse vector by node id\n const void *get_sparse_vector(node_id_t id) const {\n const void *sparse_vector_ptr;\n uint32_t sparse_vector_len;\n int ret =\n get_sparse_vector_ptr_by_id(id, &sparse_vector_ptr, &sparse_vector_len);\n if (ret != 0) {\n return nullptr;\n }\n return sparse_vector_ptr;\n }\n\n int get_sparse_vector_by_key(const uint64_t key,\n std::string *sparse_vector) const {\n uint32_t sparse_vector_len;\n IndexStorage::MemoryBlock sparse_vector_block;\n int ret = get_sparse_vector_ptr_by_key(key, sparse_vector_block,\n &sparse_vector_len);\n if (ret != 0) {\n return ret;\n }\n *sparse_vector =\n std::string(static_cast(sparse_vector_block.data()),\n sparse_vector_len);\n return 0;\n }\n\n int get_sparse_vector(node_id_t id,\n IndexStorage::MemoryBlock &sparse_vector_block) const {\n uint32_t sparse_vector_len;\n return get_sparse_vector_ptr_by_id(id, sparse_vector_block,\n &sparse_vector_len);\n }\n\n int get_sparse_vector_ptr_by_key(uint64_t key, const void **sparse_vector_ptr,\n uint32_t *sparse_vector_len_ptr) const {\n auto node_id = get_id(key);\n if (node_id == kInvalidNodeId) {\n return IndexError_NoExist;\n }\n\n return get_sparse_vector_ptr_by_id(node_id, sparse_vector_ptr,\n sparse_vector_len_ptr);\n }\n\n int get_sparse_vector_ptr_by_key(\n const uint64_t key, IndexStorage::MemoryBlock &sparse_vector_block,\n uint32_t *sparse_vector_len_ptr) const {\n auto node_id = get_id(key);\n if (node_id == kInvalidNodeId) {\n return IndexError_NoExist;\n }\n\n return get_sparse_vector_ptr_by_id(node_id, sparse_vector_block,\n sparse_vector_len_ptr);\n }\n\n std::vector get_keys() const {\n std::vector keys;\n node_id_t doc_total_cnt = doc_cnt();\n for (node_id_t node_id = 0; node_id < doc_total_cnt; ++node_id) {\n uint64_t key = get_key(node_id);\n if (key == kInvalidKey) {\n return {kInvalidKey};\n } else {\n keys.push_back(key);\n }\n }\n\n return keys;\n }\n\n\n public:\n virtual uint32_t doc_cnt() const = 0;\n\n virtual uint32_t total_sparse_count() const = 0;\n\n virtual node_id_t get_id(uint64_t key) const = 0;\n\n virtual uint64_t get_key(node_id_t id) const = 0;\n\n virtual int get_sparse_vector_ptr_by_id(\n node_id_t id, const void **sparse_vector,\n uint32_t *sparse_vector_len) const = 0;\n\n virtual int get_sparse_vector_ptr_by_id(\n const node_id_t /*id*/,\n IndexStorage::MemoryBlock & /*sparse_vector_block*/,\n uint32_t * /*sparse_vector_len*/) const {\n return IndexError_NotImplemented;\n }\n\n\n virtual float get_search_distance(const std::string &vector,\n node_id_t target_node_id) const = 0;\n virtual size_t sparse_unit_size() const = 0;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_index_format.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\nstatic constexpr uint64_t kInvalidKey = std::numeric_limits::max();\nstatic constexpr uint32_t kDefaultOffsetChunkSize = 1024 * 1024; // 1MB\nstatic constexpr uint32_t kDefaultDataChunkSize = 8 * 1024 * 1024; // 8MB\n\nstruct FlatSparseMeta {\n uint64_t create_time{0};\n uint64_t update_time{0};\n uint32_t doc_cnt{0};\n uint32_t total_sparse_count{0};\n uint8_t reserved[8] = {0};\n};\n\nstatic_assert(sizeof(FlatSparseMeta) % 32 == 0,\n \"FlatSparseMeta must be aligned with 32 bytes\");\n\nstruct FlatSparseStreamerMeta {\n uint32_t offset_chunk_count{0};\n uint32_t offset_chunk_size{kDefaultOffsetChunkSize};\n uint32_t data_chunk_count{0};\n uint32_t data_chunk_size{kDefaultDataChunkSize};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_provider.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \"flat_sparse_streamer_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Brute Force Sparse Streamer Provider\n */\n// FlatSparseStreamerEntity or FlatSparseSearcherEntity\ntemplate \nclass FlatSparseIndexProvider : public IndexSparseProvider {\n public:\n //! Constructor\n FlatSparseIndexProvider(const std::shared_ptr entity,\n const IndexMeta &meta, const std::string &owner)\n : entity_(entity), meta_(meta), owner_class_(owner) {}\n\n //! Create a new iterator\n IndexSparseProvider::Iterator::Pointer create_iterator(void) override {\n return IndexSparseProvider::Iterator::Pointer(new (std::nothrow)\n Iterator(entity_, meta_));\n }\n\n //! Retrieve count of vectors\n size_t count(void) const override {\n return entity_->doc_cnt();\n }\n\n //! Retrieve type of vector\n IndexMeta::DataType data_type(void) const override {\n return meta_.data_type();\n }\n\n //! Retrieve a vector using a primary key\n int get_sparse_vector(uint64_t key, uint32_t *sparse_count,\n std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override {\n std::string sparse_data;\n\n int ret = entity_->get_sparse_vector_by_key(key, &sparse_data);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to get sparse vector, key=%zu, ret=%s\", (size_t)key,\n IndexError::What(ret));\n return ret;\n }\n\n SparseUtility::ReverseSparseFormat(sparse_data, sparse_count,\n sparse_indices_buffer,\n sparse_values_buffer, meta_.unit_size());\n return 0;\n }\n\n //! Retrieve the owner class\n const std::string &owner_class(void) const override {\n return owner_class_;\n }\n\n size_t total_sparse_count() const override {\n return entity_->total_sparse_count();\n }\n\n private:\n class Iterator : public IndexSparseProvider::Iterator {\n public:\n Iterator(const std::shared_ptr &entity,\n const IndexMeta &meta)\n : entity_(entity), meta_(meta), cur_id_(0U), valid_(false) {\n IndexStorage::MemoryBlock sparse_data_block;\n entity_->get_sparse_vector(cur_id_, sparse_data_block);\n const void *sparse_data = sparse_data_block.data();\n if (sparse_data != nullptr) {\n valid_ = true;\n\n sparse_indices_buffer_.clear();\n sparse_data_buffer_.clear();\n\n SparseUtility::ReverseSparseFormat(\n sparse_data, &sparse_count_, &sparse_indices_buffer_,\n &sparse_data_buffer_, meta.unit_size());\n }\n }\n\n //! Retrieve sparse count\n virtual uint32_t sparse_count() const override {\n return sparse_count_;\n }\n\n //! Retrieve sparse indices\n virtual const uint32_t *sparse_indices() const override {\n return reinterpret_cast(sparse_indices_buffer_.data());\n }\n\n //! Retrieve sparse data\n virtual const void *sparse_data() const override {\n return reinterpret_cast(sparse_data_buffer_.data());\n }\n\n //! Test if the iterator is valid\n virtual bool is_valid(void) const override {\n return cur_id_ < entity_->doc_cnt() && valid_;\n }\n\n //! Retrieve primary key\n virtual uint64_t key(void) const override {\n // std::cout << \"iter key=\" << cur_id_ << std::endl;\n return entity_->get_key(cur_id_);\n }\n\n //! Next iterator\n virtual void next(void) override {\n cur_id_ = get_next_valid_id(cur_id_ + 1);\n\n if (cur_id_ < entity_->doc_cnt()) {\n IndexStorage::MemoryBlock sparse_data_block;\n entity_->get_sparse_vector(cur_id_, sparse_data_block);\n const void *sparse_data = sparse_data_block.data();\n if (sparse_data != nullptr) {\n valid_ = true;\n\n sparse_indices_buffer_.clear();\n sparse_data_buffer_.clear();\n\n SparseUtility::ReverseSparseFormat(\n sparse_data, &sparse_count_, &sparse_indices_buffer_,\n &sparse_data_buffer_, meta_.unit_size());\n } else {\n valid_ = false;\n }\n }\n }\n\n //! Reset the iterator\n void reset(void) {\n cur_id_ = get_next_valid_id(0);\n IndexStorage::MemoryBlock sparse_data_block;\n entity_->get_sparse_vector(cur_id_, sparse_data_block);\n const void *sparse_data = sparse_data_block.data();\n if (sparse_data != nullptr) {\n valid_ = true;\n\n SparseUtility::ReverseSparseFormat(\n sparse_data, &sparse_count_, &sparse_indices_buffer_,\n &sparse_data_buffer_, meta_.unit_size());\n }\n }\n\n private:\n node_id_t get_next_valid_id(node_id_t start_id) {\n for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) {\n if (entity_->get_key(i) != kInvalidNodeId) {\n return i;\n }\n }\n return kInvalidNodeId;\n }\n\n private:\n const std::shared_ptr entity_{nullptr};\n const IndexMeta &meta_;\n node_id_t cur_id_;\n uint32_t sparse_count_;\n std::string sparse_indices_buffer_;\n std::string sparse_data_buffer_;\n bool valid_{false};\n };\n\n private:\n const std::shared_ptr entity_{nullptr};\n const IndexMeta &meta_;\n const std::string owner_class_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_search.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \"flat_sparse_context.h\"\n\nnamespace zvec {\nnamespace core {\n\nstatic inline IndexGroupDocumentList ConvertGroupMapToResult(\n std::unordered_map group_map,\n uint32_t group_num) {\n IndexGroupDocumentList result;\n\n std::vector> best_score_in_groups;\n for (auto itr = group_map.begin(); itr != group_map.end(); itr++) {\n const std::string &group_id = (*itr).first;\n auto &heap = (*itr).second;\n\n if (heap.size() > 0) {\n float best_score = heap[0].score();\n best_score_in_groups.push_back(std::make_pair(group_id, best_score));\n }\n }\n\n std::sort(best_score_in_groups.begin(), best_score_in_groups.end(),\n [](const std::pair &a,\n const std::pair &b) -> int {\n return a.second < b.second;\n });\n\n // truncate to group num\n for (uint32_t i = 0; i < group_num && i < best_score_in_groups.size(); ++i) {\n const std::string &group_id = best_score_in_groups[i].first;\n\n result.emplace_back(\n GroupIndexDocument(group_id, std::move(group_map[group_id])));\n }\n\n return result;\n}\n\nstatic inline int FlatSearch(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query, bool with_p_keys,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n const IndexMeta, IndexContext::Pointer &context,\n FlatSparseEntity *entity) {\n int ret;\n\n FlatSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to FlatSparseContext failed\");\n return IndexError_Cast;\n }\n\n // reset context results\n ctx->reset_results(count);\n\n const uint32_t *sparse_indices_tmp = sparse_indices;\n const void *sparse_query_tmp = sparse_query;\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](uint64_t key) {\n return ctx->group_by()(key);\n };\n\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n ailego::MinusInnerProductSparseMatrix::transform_sparse_format(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n sparse_query_buffer);\n\n std::unordered_map group_heap{};\n\n if (with_p_keys) {\n ret = entity->search_group_p_keys(sparse_query_buffer, p_keys[q],\n ctx->filter(), group_by,\n ctx->group_topk(), &group_heap);\n } else {\n ret = entity->search_group(sparse_query_buffer, ctx->filter(), group_by,\n ctx->group_topk(), &group_heap);\n }\n\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to search group, ret=%s\", IndexError::What(ret));\n return ret;\n }\n\n // sort group heap\n for (auto &group : group_heap) {\n group.second.sort();\n }\n\n auto group_result =\n ConvertGroupMapToResult(std::move(group_heap), ctx->group_num());\n ctx->mutable_group_result(q)->swap(group_result);\n }\n } else {\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n ailego::MinusInnerProductSparseMatrix::transform_sparse_format(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n sparse_query_buffer);\n\n auto heap = ctx->result_heap();\n\n if (with_p_keys) {\n ret = entity->search_p_keys(sparse_query_buffer, p_keys[q],\n ctx->filter(), heap);\n } else {\n ret = entity->search(sparse_query_buffer, ctx->filter(), heap);\n }\n\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to search, ret=%s\", IndexError::What(ret));\n return ret;\n }\n\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n }\n\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_searcher.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"flat_sparse_searcher.h\"\n#include \n#include \n#include \"flat_sparse_context.h\"\n#include \"flat_sparse_provider.h\"\n#include \"flat_sparse_search.h\"\n\nnamespace zvec {\nnamespace core {\n\nconst uint32_t FlatSparseSearcher::VERSION = 0U;\n\nFlatSparseSearcher::FlatSparseSearcher(void) {}\n\nFlatSparseSearcher::~FlatSparseSearcher(void) {}\n\nint FlatSparseSearcher::init(const ailego::Params & /*params*/) {\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint FlatSparseSearcher::cleanup(void) {\n this->unload();\n return 0;\n}\n\nint FlatSparseSearcher::load(IndexStorage::Pointer container,\n IndexMetric::Pointer /*measure*/) {\n if (state_ != STATE_INITED) {\n LOG_ERROR(\"Init the searcher first before load index\");\n return IndexError_Runtime;\n }\n\n LOG_INFO(\"Begin FlatSparseSearcher::load\");\n\n int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to deserialize meta from container\");\n return ret;\n }\n\n if (meta_.searcher_revision() != VERSION) {\n LOG_ERROR(\"Unsupported searcher revision %u\", meta_.searcher_revision());\n return IndexError_Unsupported;\n }\n\n ret = entity_.load(container, meta_);\n if (ret != 0) {\n LOG_ERROR(\"FlatSparseSearcher load index failed\");\n return ret;\n }\n\n state_ = STATE_LOADED;\n magic_ = IndexContext::GenerateMagic();\n\n LOG_INFO(\"End FlatSparseSearcher::load\");\n\n return 0;\n}\n\nint FlatSparseSearcher::unload(void) {\n LOG_INFO(\"Begin FlatSparseSearcher::unload\");\n\n meta_.clear();\n entity_.unload();\n state_ = STATE_INITED;\n\n LOG_INFO(\"End FlatSparseSearcher::unload\");\n\n return 0;\n}\n\nint FlatSparseSearcher::search_bf_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const {\n return do_search(sparse_count, sparse_indices, sparse_query, false, {}, qmeta,\n count, context);\n}\n\nint FlatSparseSearcher::search_bf_by_p_keys_impl(\n const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const {\n return do_search(sparse_count, sparse_indices, sparse_query, true, p_keys,\n qmeta, count, context);\n}\n\nint FlatSparseSearcher::get_sparse_vector(\n uint64_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const {\n if (state_ != STATE_LOADED) {\n LOG_ERROR(\"Failed to get sparse vector, load container first!\");\n return IndexError_NoIndexLoaded;\n }\n\n std::string sparse_data;\n\n int ret = entity_.get_sparse_vector(key, &sparse_data);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to get sparse vector, key=%zu, ret=%s\", (size_t)key,\n IndexError::What(ret));\n return ret;\n }\n\n SparseUtility::ReverseSparseFormat(sparse_data, sparse_count,\n sparse_indices_buffer,\n sparse_values_buffer, meta_.unit_size());\n\n return 0;\n}\n\nFlatSparseSearcher::ContextPointer FlatSparseSearcher::create_context() const {\n if (state_ != STATE_LOADED) {\n LOG_ERROR(\"Failed to create Context, load container first!\");\n return Context::UPointer();\n }\n FlatSparseSearcherEntity::Pointer entity = entity_.clone();\n return FlatSparseSearcher::ContextPointer(new FlatSparseContext(this));\n}\n\n//! Create a new iterator\nIndexSearcher::SparseProvider::Pointer\nFlatSparseSearcher::create_sparse_provider(void) const {\n if (state_ != STATE_LOADED) {\n LOG_ERROR(\"Failed to create provider, load container first!\");\n return SparseProvider::Pointer();\n }\n\n auto entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"Clone entity failed\");\n return SparseProvider::Pointer();\n }\n return SparseProvider::Pointer(\n new FlatSparseIndexProvider(\n entity, meta_, \"FlatSparseSearcher\"));\n}\n\nint FlatSparseSearcher::do_search(\n const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, bool with_p_keys,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const {\n if (state_ != STATE_LOADED) {\n LOG_ERROR(\"Failed to do search, load container first!\");\n return IndexError_NoIndexLoaded;\n }\n\n int ret = check_params(qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n return FlatSearch(sparse_count, sparse_indices, sparse_query, with_p_keys,\n p_keys, qmeta, count, meta_, context,\n (FlatSparseEntity *)&entity_);\n}\n\nINDEX_FACTORY_REGISTER_SEARCHER(FlatSparseSearcher);\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_searcher.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \"flat_sparse_searcher_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass FlatSparseSearcher : public IndexSearcher {\n public:\n static const uint32_t VERSION;\n\n public:\n using ContextPointer = IndexSearcher::Context::Pointer;\n\n public:\n FlatSparseSearcher(void);\n virtual ~FlatSparseSearcher(void);\n\n FlatSparseSearcher(const FlatSparseSearcher &) = delete;\n FlatSparseSearcher &operator=(const FlatSparseSearcher &) = delete;\n\n public:\n //! Initialize Searcher\n int init(const ailego::Params ¶ms) override;\n\n //! Cleanup Searcher\n int cleanup(void) override;\n\n //! Load Index from storage\n int load(IndexStorage::Pointer container,\n IndexMetric::Pointer /*measure*/) override;\n\n //! Unload index from storage\n int unload(void) override;\n\n int search_impl(const void * /*query*/, const IndexQueryMeta & /*qmeta*/,\n Context::Pointer & /*context*/) const override {\n return IndexError_NotImplemented;\n }\n\n int search_impl(const void * /*query*/, const IndexQueryMeta & /*qmeta*/,\n uint32_t /*count*/,\n Context::Pointer & /*context*/) const override {\n return IndexError_NotImplemented;\n }\n\n int search_bf_impl(const void * /*query*/, const IndexQueryMeta & /*qmeta*/,\n Context::Pointer & /*context*/) const override {\n return IndexError_NotImplemented;\n }\n\n int search_bf_impl(const void * /*query*/, const IndexQueryMeta & /*qmeta*/,\n uint32_t /*count*/,\n Context::Pointer & /*context*/) const override {\n return IndexError_NotImplemented;\n }\n\n //! Similarity search with sparse inputs\n int search_impl(const uint32_t sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override {\n return search_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1,\n context);\n }\n\n //! Similarity search with sparse inputs\n int search_impl(const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const override {\n return search_bf_impl(sparse_count, sparse_indices, sparse_query, qmeta,\n count, context);\n }\n\n //! Similarity brute force search with sparse inputs\n int search_bf_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override {\n return search_bf_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1,\n context);\n }\n\n //! Similarity brute force search with sparse inputs\n int search_bf_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const override;\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta,\n ContextPointer &context) const override {\n return search_bf_by_p_keys_impl(&sparse_count, sparse_indices, sparse_query,\n p_keys, qmeta, 1, context);\n }\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const override;\n\n //! Fetch sparser vector by key\n int get_sparse_vector(uint64_t key, uint32_t *sparse_count,\n std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override;\n\n //! Create a searcher context\n ContextPointer create_context() const override;\n\n //! Create a new iterator\n IndexSearcher::SparseProvider::Pointer create_sparse_provider(\n void) const override;\n\n //! Retrieve statistics\n const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve meta of index\n const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n //! Retrieve params of index\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n const FlatSparseSearcherEntity &entity(void) const {\n return entity_;\n }\n\n uint32_t magic(void) const {\n return magic_;\n }\n\n private:\n inline int check_params(const IndexQueryMeta &qmeta) const {\n if (ailego_unlikely(qmeta.data_type() != meta_.data_type())) {\n LOG_ERROR(\"Unsupported query meta\");\n return IndexError_Mismatch;\n }\n return 0;\n }\n\n int do_search(const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, bool with_p_keys,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const;\n\n private:\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 };\n\n FlatSparseSearcherEntity entity_{};\n IndexMeta meta_{};\n ailego::Params params_{};\n uint32_t magic_{0U};\n\n Stats stats_;\n State state_{STATE_INIT};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"flat_sparse_searcher_entity.h\"\n#include \n#include \n#include \"flat_sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nFlatSparseSearcherEntity::FlatSparseSearcherEntity() {}\n\nint FlatSparseSearcherEntity::load(const IndexStorage::Pointer &container,\n const IndexMeta &index_meta) {\n if (container_) {\n LOG_ERROR(\"An storage instance is already opened\");\n return IndexError_Duplicate;\n }\n\n int ret = this->load_container(container);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to load storage index\");\n return ret;\n }\n\n if (init_measure(index_meta) != 0) {\n LOG_ERROR(\"Failed to init measure\");\n return IndexError_InvalidFormat;\n }\n\n container_ = container;\n return 0;\n}\n\nint FlatSparseSearcherEntity::init_measure(const IndexMeta &meta) {\n measure_ = IndexFactory::CreateMetric(meta.metric_name());\n if (!measure_) {\n LOG_ERROR(\"Failed to create measure %s\", meta.metric_name().c_str());\n return IndexError_NoExist;\n }\n int ret = measure_->init(meta, meta.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"Failled to init measure, ret=%d\", ret);\n return ret;\n }\n\n if (!measure_->sparse_distance()) {\n LOG_ERROR(\"Invalid measure distance\");\n return IndexError_InvalidArgument;\n }\n\n search_sparse_distance_ = measure_->sparse_distance();\n\n if (measure_->query_metric() && measure_->query_metric()->distance()) {\n search_sparse_distance_ = measure_->query_metric()->sparse_distance();\n }\n sparse_unit_size_ = meta.unit_size();\n\n return 0;\n}\n\nint FlatSparseSearcherEntity::load_container(\n const IndexStorage::Pointer &container) {\n // meta\n auto segment = container->get(PARAM_FLAT_SPARSE_META_SEG_ID);\n if (!segment || segment->data_size() < sizeof(meta_)) {\n LOG_ERROR(\"Missing segment %s, or invalid segment size\",\n PARAM_FLAT_SPARSE_META_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n const void *data;\n if (ailego_unlikely(segment->read(0, &data, sizeof(meta_)) !=\n sizeof(meta_))) {\n LOG_ERROR(\"Failed to read meta segment %s\",\n PARAM_FLAT_SPARSE_META_SEG_ID.c_str());\n return IndexError_ReadData;\n }\n meta_ = *(reinterpret_cast(data));\n\n // keys segment\n keys_chunk_ = container->get(PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID);\n if (!keys_chunk_) {\n LOG_ERROR(\"Missing segment %s\", PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n\n // mapping segment\n mapping_chunk_ = container->get(PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID);\n if (!mapping_chunk_) {\n LOG_ERROR(\"Missing segment %s\",\n PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n\n // offset segment\n sparse_offset_chunk_ = container->get(PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID);\n if (!sparse_offset_chunk_) {\n LOG_ERROR(\"Missing segment %s\",\n PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n\n // data segment\n sparse_data_chunk_ = container->get(PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID);\n if (!sparse_data_chunk_) {\n LOG_ERROR(\"Missing segment %s\", PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n\n return 0;\n}\n\nint FlatSparseSearcherEntity::unload() {\n container_.reset();\n sparse_data_chunk_.reset();\n sparse_offset_chunk_.reset();\n keys_chunk_.reset();\n mapping_chunk_.reset();\n\n return 0;\n}\n\nFlatSparseSearcherEntity::Pointer FlatSparseSearcherEntity::clone() const {\n auto entity = new (std::nothrow)\n FlatSparseSearcherEntity(meta_, sparse_data_chunk_, sparse_offset_chunk_,\n keys_chunk_, mapping_chunk_);\n return FlatSparseSearcherEntity::Pointer(entity);\n}\n\nint FlatSparseSearcherEntity::get_sparse_vector_ptr_by_id(\n node_id_t id, const void **sparse_vector_ptr,\n uint32_t *sparse_vector_len_ptr) const {\n uint32_t offset_chunk_offset = id * offset_size_per_node();\n\n const void *offset_info = nullptr;\n if (ailego_unlikely(sparse_offset_chunk_->read(\n offset_chunk_offset, &offset_info,\n offset_size_per_node()) != offset_size_per_node())) {\n LOG_ERROR(\"Read offset info failed, offset=%u\", offset_chunk_offset);\n return IndexError_ReadData;\n };\n\n // sparse offset\n uint64_t sparse_offset = *(uint64_t *)offset_info;\n uint32_t sparse_vector_len =\n *(uint32_t *)((uint8_t *)offset_info + sizeof(uint64_t));\n\n if (sparse_vector_len > 0) {\n const void *sparse_data =\n get_sparse_vector_data(sparse_offset, sparse_vector_len);\n if (ailego_unlikely(sparse_data == nullptr)) {\n LOG_ERROR(\"Get nullptr sparse, offset=%zu, len=%u\", (size_t)sparse_offset,\n sparse_vector_len);\n\n return IndexError_ReadData;\n }\n *sparse_vector_ptr = sparse_data;\n *sparse_vector_len_ptr = sparse_vector_len;\n }\n\n return 0;\n}\n\nconst void *FlatSparseSearcherEntity::get_sparse_vector_data(\n uint64_t offset, uint32_t length) const {\n const void *data;\n auto size = sparse_data_chunk_->read(offset, &data, length);\n if (size != length) {\n LOG_ERROR(\n \"read sparse vector data failed: offset=%zu, \"\n \"length=%u, size=%zu\",\n (size_t)offset, length, size);\n return nullptr;\n }\n return data;\n}\n\n\nnode_id_t FlatSparseSearcherEntity::get_id(uint64_t key) const {\n if (ailego_unlikely(!mapping_chunk_)) {\n LOG_ERROR(\"Index missing mapping segment\");\n return kInvalidNodeId;\n }\n\n //! Do binary search\n node_id_t start = 0UL;\n node_id_t end = doc_cnt();\n const void *data;\n node_id_t idx = 0u;\n while (start < end) {\n idx = start + (end - start) / 2;\n if (ailego_unlikely(mapping_chunk_->read(idx * sizeof(node_id_t), &data,\n sizeof(node_id_t)) !=\n sizeof(node_id_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidNodeId;\n }\n const uint64_t *mkey;\n node_id_t local_id = *reinterpret_cast(data);\n if (ailego_unlikely(keys_chunk_->read(\n local_id * sizeof(uint64_t), (const void **)(&mkey),\n sizeof(uint64_t)) != sizeof(uint64_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidNodeId;\n }\n if (*mkey < key) {\n start = idx + 1;\n } else if (*mkey > key) {\n end = idx;\n } else {\n return local_id;\n }\n }\n return kInvalidNodeId;\n}\n\nuint64_t FlatSparseSearcherEntity::get_key(node_id_t id) const {\n const void *key;\n if (ailego_unlikely(\n keys_chunk_->read(id * sizeof(uint64_t), &key, sizeof(uint64_t)) !=\n sizeof(uint64_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidKey;\n }\n return *(reinterpret_cast(key));\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_searcher_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \"flat_sparse_entity.h\"\n#include \"flat_sparse_index_format.h\"\n\nnamespace zvec {\nnamespace core {\n\n\n/*! Flat Sparse Searcher Entity\n */\nclass FlatSparseSearcherEntity : public FlatSparseEntity {\n public:\n typedef std::shared_ptr Pointer;\n\n using Chunk = IndexStorage::Segment;\n\n //! Constructor\n explicit FlatSparseSearcherEntity();\n\n //! Destructor\n virtual ~FlatSparseSearcherEntity() = default;\n\n //! Disable them\n FlatSparseSearcherEntity(const FlatSparseSearcherEntity &) = delete;\n FlatSparseSearcherEntity &operator=(const FlatSparseSearcherEntity &) =\n delete;\n\n //! Load the entity with container\n int load(const IndexStorage::Pointer &container, const IndexMeta &index_meta);\n\n //! Unload the entity\n int unload();\n\n public:\n inline uint32_t doc_cnt() const override {\n return meta_.doc_cnt;\n }\n\n inline uint32_t total_sparse_count() const override {\n return meta_.total_sparse_count;\n }\n\n size_t sparse_unit_size() const override {\n return sparse_unit_size_;\n }\n\n float get_search_distance(const std::string &vector,\n node_id_t target_node_id) const override {\n float dist;\n const void *target_vector;\n uint32_t target_vector_len;\n get_sparse_vector_ptr_by_id(target_node_id, &target_vector,\n &target_vector_len);\n search_sparse_distance_(vector.c_str(), target_vector, &dist);\n return dist;\n }\n\n FlatSparseSearcherEntity::Pointer clone() const;\n\n node_id_t get_id(uint64_t key) const override;\n\n uint64_t get_key(node_id_t id) const override;\n\n int get_sparse_vector_ptr_by_id(node_id_t id, const void **sparse_vector,\n uint32_t *sparse_vector_len) const override;\n\n private:\n int load_container(const IndexStorage::Pointer &container);\n\n int init_measure(const IndexMeta &meta);\n\n inline uint32_t offset_size_per_node() const {\n return sizeof(uint64_t) + sizeof(uint32_t);\n }\n\n const void *get_sparse_vector_data(uint64_t offset, uint32_t length) const;\n\n private:\n FlatSparseSearcherEntity(const FlatSparseMeta &meta,\n Chunk::Pointer sparse_data_chunk,\n Chunk::Pointer sparse_offset_chunk,\n Chunk::Pointer keys_chunk,\n Chunk::Pointer mapping_chunk)\n : meta_(meta),\n sparse_data_chunk_(sparse_data_chunk),\n sparse_offset_chunk_(sparse_offset_chunk),\n keys_chunk_(keys_chunk),\n mapping_chunk_(mapping_chunk) {}\n\n private:\n IndexStorage::Pointer container_{};\n\n // meta\n FlatSparseMeta meta_;\n\n // measure\n IndexMetric::Pointer measure_{};\n IndexMetric::MatrixSparseDistance search_sparse_distance_{};\n\n // chunk\n Chunk::Pointer sparse_data_chunk_;\n Chunk::Pointer sparse_offset_chunk_;\n Chunk::Pointer keys_chunk_;\n Chunk::Pointer mapping_chunk_;\n\n size_t sparse_unit_size_{0U};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_streamer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"flat_sparse_streamer.h\"\n#include \n#include \n#include \n#include \n#include \n#include \"flat_sparse_context.h\"\n#include \"flat_sparse_provider.h\"\n#include \"flat_sparse_search.h\"\n\nnamespace zvec {\nnamespace core {\n\nconst uint32_t FlatSparseStreamer::VERSION = 0U;\n\nFlatSparseStreamer::FlatSparseStreamer() : entity_(stats_) {}\n\nFlatSparseStreamer::~FlatSparseStreamer() {\n this->close();\n}\n\nint FlatSparseStreamer::init(const IndexMeta &imeta,\n const ailego::Params ¶ms) {\n LOG_DEBUG(\"FlatSparseStreamer init\");\n\n meta_ = imeta;\n meta_.set_streamer(\"FlatSparseStreamer\", VERSION, params);\n\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint FlatSparseStreamer::cleanup() {\n LOG_DEBUG(\"FlatSparseStreamer cleanup\");\n\n this->close();\n\n meta_.clear();\n\n return 0;\n}\n\nint FlatSparseStreamer::open(IndexStorage::Pointer stg) {\n LOG_DEBUG(\"FlatSparseStreamer open\");\n\n if (ailego_unlikely(state_ != STATE_INITED)) {\n LOG_ERROR(\"Open storage failed, init streamer first!\");\n return IndexError_NoReady;\n }\n\n int ret = entity_.open(std::move(stg), meta_);\n if (ret != 0) {\n LOG_ERROR(\"FlatSparseStreamer entity failed to open storage\");\n return ret;\n }\n\n IndexMeta index_meta;\n ret = entity_.get_index_sparse_meta(&index_meta);\n if (ret == IndexError_NoExist) {\n // Set IndexMeta for the new index\n ret = entity_.set_index_sparse_meta(meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to set index meta for %s\", IndexError::What(ret));\n return ret;\n }\n } else {\n if (index_meta.streamer_revision() != meta_.streamer_revision()) {\n LOG_ERROR(\"Streamer revision mismatch, expect=%u, actual=%u\",\n meta_.streamer_revision(), index_meta.streamer_revision());\n return IndexError_Mismatch;\n }\n if (index_meta.metric_name() != meta_.metric_name() ||\n index_meta.data_type() != meta_.data_type()) {\n LOG_ERROR(\"IndexMeta mismatch from the previous in index\");\n return IndexError_Mismatch;\n }\n // The IndexMeasure Params may be updated like MipsSquaredEuclidean\n auto metric_params = index_meta.metric_params();\n metric_params.merge(meta_.metric_params());\n meta_.set_metric(index_meta.metric_name(), 0, metric_params);\n }\n\n state_ = STATE_OPENED;\n magic_ = IndexContext::GenerateMagic();\n\n return 0;\n}\n\nint FlatSparseStreamer::close() {\n if (state_ != STATE_OPENED) {\n return 0;\n }\n\n LOG_DEBUG(\"FlatSparseStreamer close\");\n\n stats_.clear();\n int ret = entity_.close();\n if (ret != 0) {\n LOG_ERROR(\"Failed to close entity %s\", IndexError::What(ret));\n return ret;\n }\n state_ = STATE_INITED;\n return 0;\n}\n\nint FlatSparseStreamer::flush(uint64_t checkpoint) {\n if (state_ != STATE_OPENED) {\n LOG_ERROR(\"Failed to flush, open streamer first!\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"FlatSparseStreamer flush, checkpoint=%zu\", (size_t)checkpoint);\n\n return entity_.flush(checkpoint);\n}\n\nint FlatSparseStreamer::dump(const IndexDumper::Pointer &dumper) {\n if (state_ != STATE_OPENED) {\n LOG_ERROR(\"Failed to dump, open streamer first!\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"FlatSparseStreamer dump\");\n\n shared_mutex_.lock();\n AILEGO_DEFER([&]() { shared_mutex_.unlock(); });\n\n meta_.set_searcher(\"FlatSparseSearcher\", VERSION, ailego::Params());\n\n int ret = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return ret;\n }\n\n return entity_.dump(dumper);\n}\n\nFlatSparseStreamer::ContextPointer FlatSparseStreamer::create_context() const {\n if (state_ != STATE_OPENED) {\n LOG_ERROR(\"Failed to create Context, open streamer first!\");\n return Context::UPointer();\n }\n FlatSparseStreamerEntity::Pointer entity = entity_.clone();\n return FlatSparseStreamer::ContextPointer(new FlatSparseContext(this));\n}\n\nIndexStreamer::SparseProvider::Pointer\nFlatSparseStreamer::create_sparse_provider(void) const {\n if (state_ != STATE_OPENED) {\n LOG_ERROR(\"Failed to create provider, open streamer first!\");\n return SparseProvider::Pointer();\n }\n\n auto entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"Clone entity failed\");\n return SparseProvider::Pointer();\n }\n return SparseProvider::Pointer(\n new FlatSparseIndexProvider(\n entity, meta_, \"FlatSparseStreamerProvider\"));\n}\n\nint FlatSparseStreamer::add_impl(uint64_t pkey, const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) {\n if (state_ != STATE_OPENED) {\n LOG_ERROR(\"Failed to add_impl, open streamer first!\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_NoReady;\n }\n\n int ret = check_params(qmeta);\n if (ailego_unlikely(ret != 0)) {\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n if (ailego_unlikely(sparse_count > PARAM_FLAT_SPARSE_MAX_DIM_SIZE)) {\n LOG_ERROR(\n \"Failed to add sparse vector: number of non-zero elements (%u) exceeds \"\n \"maximum allowed (%u), key=%zu\",\n sparse_count, PARAM_FLAT_SPARSE_MAX_DIM_SIZE, (size_t)pkey);\n (*stats_.mutable_discarded_count())++;\n return IndexError_InvalidValue;\n }\n\n // context is trivial here\n FlatSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to FlatSparseContext failed\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Cast;\n }\n\n if (ailego_unlikely(!shared_mutex_.try_lock_shared())) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); });\n\n // convert to sparse format and add to entity\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_query,\n meta_.unit_size(), sparse_query_buffer);\n\n ret = entity_.add(pkey, sparse_query_buffer, sparse_count);\n if (ret != 0) {\n LOG_ERROR(\"Failed to add sparse vector, key=%zu, ret=%s\", (size_t)pkey,\n IndexError::What(ret));\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n (*stats_.mutable_added_count())++;\n return 0;\n}\n\nint FlatSparseStreamer::add_with_id_impl(uint32_t pkey,\n const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) {\n if (state_ != STATE_OPENED) {\n LOG_ERROR(\"Failed to add_with_id_impl, open streamer first!\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_NoReady;\n }\n\n int ret = check_params(qmeta);\n if (ailego_unlikely(ret != 0)) {\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n if (ailego_unlikely(sparse_count > PARAM_FLAT_SPARSE_MAX_DIM_SIZE)) {\n LOG_ERROR(\n \"Failed to add sparse vector: number of non-zero elements (%u) exceeds \"\n \"maximum allowed (%u), key=%zu\",\n sparse_count, PARAM_FLAT_SPARSE_MAX_DIM_SIZE, (size_t)pkey);\n (*stats_.mutable_discarded_count())++;\n return IndexError_InvalidValue;\n }\n\n // context is trivial here\n FlatSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to FlatSparseContext failed\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Cast;\n }\n\n if (ailego_unlikely(!shared_mutex_.try_lock_shared())) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); });\n\n // convert to sparse format and add to entity\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_query,\n meta_.unit_size(), sparse_query_buffer);\n\n ret = entity_.add_vector_with_id(pkey, sparse_query_buffer, sparse_count);\n if (ret != 0) {\n LOG_ERROR(\"Failed to add sparse vector, key=%zu, ret=%s\", (size_t)pkey,\n IndexError::What(ret));\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n (*stats_.mutable_added_count())++;\n return 0;\n}\n\n//! Similarity search with sparse inputs\nint FlatSparseStreamer::search_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n return search_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1,\n context);\n}\n\n//! Similarity search with sparse inputs\nint FlatSparseStreamer::search_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n return search_bf_impl(sparse_count, sparse_indices, sparse_query, qmeta,\n count, context);\n}\n\n//! Similarity brute force search with sparse inputs\nint FlatSparseStreamer::search_bf_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n return search_bf_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1,\n context);\n}\n\n//! Linear search by primary keys\nint FlatSparseStreamer::search_bf_by_p_keys_impl(\n const uint32_t sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, ContextPointer &context) const {\n return search_bf_by_p_keys_impl(&sparse_count, sparse_indices, sparse_query,\n p_keys, qmeta, 1, context);\n}\n\n//! Similarity brute force search with sparse inputs\nint FlatSparseStreamer::search_bf_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const {\n return do_search(sparse_count, sparse_indices, sparse_query, false, {}, qmeta,\n count, context);\n}\n\n//! Linear search by primary keys with sparse inputs\nint FlatSparseStreamer::search_bf_by_p_keys_impl(\n const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const {\n return do_search(sparse_count, sparse_indices, sparse_query, true, p_keys,\n qmeta, count, context);\n}\n\n//! Fetch sparse vector by key\nint FlatSparseStreamer::get_sparse_vector(\n uint64_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const {\n if (state_ != STATE_OPENED) {\n LOG_ERROR(\"Failed to get_sparse_vector, open streamer first!\");\n return IndexError_NoReady;\n }\n\n std::string sparse_data;\n\n int ret = entity_.get_sparse_vector_by_key(key, &sparse_data);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to get sparse vector, key=%zu, ret=%s\", (size_t)key,\n IndexError::What(ret));\n return ret;\n }\n\n SparseUtility::ReverseSparseFormat(sparse_data, sparse_count,\n sparse_indices_buffer,\n sparse_values_buffer, meta_.unit_size());\n\n return 0;\n}\n\nint FlatSparseStreamer::do_search(\n const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, bool with_p_keys,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const {\n if (state_ != STATE_OPENED) {\n LOG_ERROR(\"Failed to do_search, open streamer first!\");\n return IndexError_NoReady;\n }\n\n int ret = check_params(qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n FlatSparseContext *ctx = dynamic_cast(context.get());\n if (ctx->magic() != magic_) {\n ctx->reset(this);\n }\n\n return FlatSearch(sparse_count, sparse_indices, sparse_query, with_p_keys,\n p_keys, qmeta, count, meta_, context,\n (FlatSparseEntity *)&entity_);\n}\n\nINDEX_FACTORY_REGISTER_STREAMER(FlatSparseStreamer);\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_streamer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \"flat_sparse_streamer_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Flat Sparse Streamer\n */\nclass FlatSparseStreamer : public IndexStreamer {\n public:\n static const uint32_t VERSION;\n\n public:\n using ContextPointer = IndexStreamer::Context::Pointer;\n\n FlatSparseStreamer(void);\n virtual ~FlatSparseStreamer(void);\n\n FlatSparseStreamer(const FlatSparseStreamer &streamer) = delete;\n FlatSparseStreamer &operator=(const FlatSparseStreamer &streamer) = delete;\n\n public:\n //! Initialize Streamer\n int init(const IndexMeta &, const ailego::Params &) override;\n\n //! Cleanup Streamer\n int cleanup(void) override;\n\n //! Open index from file path\n int open(IndexStorage::Pointer stg) override;\n\n //! Close file\n int close(void) override;\n\n //! flush file\n int flush(uint64_t checkpoint) override;\n\n //! Dump index into storage\n int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Create a context\n ContextPointer create_context(void) const override;\n\n //! Create a new iterator\n IndexStreamer::SparseProvider::Pointer create_sparse_provider(\n void) const override;\n\n int add_impl(uint64_t pkey, const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta, Context::Pointer &context) override;\n\n int add_with_id_impl(uint32_t pkey, const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) override;\n\n //! Similarity search with sparse inputs\n int search_impl(const uint32_t sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity search with sparse inputs\n int search_impl(const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const override;\n\n //! Similarity brute force search with sparse inputs\n int search_bf_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity brute force search with sparse inputs\n int search_bf_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const override;\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta,\n ContextPointer &context) const override;\n\n //! Linear search by primary keys with sparse inputs\n int search_bf_by_p_keys_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const override;\n\n //! Fetch sparse vector by key\n int get_sparse_vector(uint64_t key, uint32_t *sparse_count,\n std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override;\n\n int get_sparse_vector_by_id(\n uint32_t id, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override {\n return get_sparse_vector(id, sparse_count, sparse_indices_buffer,\n sparse_values_buffer);\n }\n\n //! Retrieve statistics\n const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve meta of index\n const IndexMeta &meta(void) const override {\n return meta_;\n }\n const FlatSparseStreamerEntity &entity(void) const {\n return entity_;\n }\n\n uint32_t magic(void) const {\n return magic_;\n }\n\n private:\n inline int check_params(const IndexQueryMeta &qmeta) const {\n if (ailego_unlikely(qmeta.data_type() != meta_.data_type())) {\n LOG_ERROR(\"Unsupported query meta, type=%d, expected=%d\",\n qmeta.data_type(), meta_.data_type());\n return IndexError_Mismatch;\n }\n return 0;\n }\n\n int do_search(const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, bool with_p_keys,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const;\n\n private:\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 };\n\n IndexMeta meta_{};\n FlatSparseStreamerEntity entity_;\n\n uint32_t magic_{0U};\n Stats stats_{};\n State state_{STATE_INIT};\n\n //! avoid add vector while dumping index\n ailego::SharedMutex shared_mutex_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"flat_sparse_streamer_entity.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"flat_sparse_index_format.h\"\n#include \"flat_sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nFlatSparseStreamerEntity::FlatSparseStreamerEntity(IndexStreamer::Stats &stats)\n : stats_(stats) {}\n\nint FlatSparseStreamerEntity::open(IndexStorage::Pointer storage,\n const IndexMeta &meta) {\n if (storage_) {\n LOG_ERROR(\"An storage instance is already opened\");\n return IndexError_Duplicate;\n }\n\n keys_map_lock_ = std::make_shared();\n if (!keys_map_lock_) {\n LOG_ERROR(\"FlatSparseStreamerEntity new object failed\");\n return IndexError_NoMemory;\n }\n keys_map_ = std::make_shared>();\n\n if (storage->get(PARAM_FLAT_SPARSE_META_SEG_ID) ||\n storage->get(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID)) {\n int ret = this->load_storage(storage, meta);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to load storage index\");\n return ret;\n }\n } else {\n int ret = this->init_storage(storage, meta);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to load storage index\");\n return ret;\n }\n }\n\n if (init_metric(meta) != 0) {\n LOG_ERROR(\"Failed to init metric\");\n return IndexError_InvalidFormat;\n }\n\n // reserve data chunk\n meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_MAX_DATA_CHUNK_CNT,\n &max_data_chunk_cnt_);\n sparse_data_chunks_.reserve(max_data_chunk_cnt_);\n\n // reserve offset chunk\n meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_MAX_DOC_CNT,\n &max_doc_cnt_);\n sparse_offset_chunks_.reserve(max_doc_cnt_ / doc_cnt_per_offset_chunk() + 1);\n sparse_unit_size_ = meta.unit_size();\n\n LOG_DEBUG(\n \"FlatSparseStreamerEntity open success, doc_count[%u], \"\n \"data_chunk_size[%u], offset_chunk_size[%u], data_chunk_count[%zu], \"\n \"offset_chunk_count[%zu]\",\n meta_.doc_cnt, streamer_meta_.data_chunk_size,\n streamer_meta_.offset_chunk_size, sparse_data_chunks_.size(),\n sparse_offset_chunks_.size());\n\n storage_ = storage;\n return 0;\n}\n\nint FlatSparseStreamerEntity::init_metric(const IndexMeta &meta) {\n metric_ = IndexFactory::CreateMetric(meta.metric_name());\n if (!metric_) {\n LOG_ERROR(\"Failed to create metric %s\", meta.metric_name().c_str());\n return IndexError_NoExist;\n }\n int ret = metric_->init(meta, meta.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"Failled to init metric, ret=%d\", ret);\n return ret;\n }\n\n if (!metric_->sparse_distance()) {\n LOG_ERROR(\"Invalid metric distance\");\n return IndexError_InvalidArgument;\n }\n\n search_sparse_distance_ = metric_->sparse_distance();\n\n if (metric_->query_metric() && metric_->query_metric()->distance()) {\n search_sparse_distance_ = metric_->query_metric()->sparse_distance();\n }\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::load_storage(IndexStorage::Pointer storage,\n const IndexMeta &meta) {\n size_t index_size{0};\n\n // load meta\n auto segment = storage->get(PARAM_FLAT_SPARSE_META_SEG_ID);\n\n if (!segment || segment->data_size() < sizeof(meta_)) {\n LOG_ERROR(\"Missing segment %s, or invalid segment size\",\n PARAM_FLAT_SPARSE_META_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n IndexStorage::MemoryBlock data_block;\n if (ailego_unlikely(segment->read(0, data_block, sizeof(meta_)) !=\n sizeof(meta_))) {\n LOG_ERROR(\"Failed to read meta segment %s\",\n PARAM_FLAT_SPARSE_META_SEG_ID.c_str());\n return IndexError_ReadData;\n }\n meta_ = *(reinterpret_cast(data_block.data()));\n index_size += segment->capacity();\n\n // load streamer meta\n segment = storage->get(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID);\n if (!segment || segment->data_size() < sizeof(streamer_meta_)) {\n LOG_ERROR(\"Missing segment %s, or invalid segment size\",\n PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n if (ailego_unlikely(segment->read(0, data_block, sizeof(streamer_meta_)) !=\n sizeof(streamer_meta_))) {\n LOG_ERROR(\"Failed to read streamer meta segment %s\",\n PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str());\n return IndexError_ReadData;\n }\n streamer_meta_ =\n *(reinterpret_cast(data_block.data()));\n index_size += segment->capacity();\n\n uint32_t meta_data_chunk_size{streamer_meta_.data_chunk_size};\n uint32_t meta_offset_chunk_size{streamer_meta_.offset_chunk_size};\n meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_DATA_CHUNK_SIZE,\n &meta_data_chunk_size);\n meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_OFFSET_CHUNK_SIZE,\n &meta_offset_chunk_size);\n if (streamer_meta_.data_chunk_size != meta_data_chunk_size ||\n streamer_meta_.offset_chunk_size != meta_offset_chunk_size) {\n LOG_ERROR(\n \"Invalid streamer meta chunk size data[%u] offset[%u], expect data[%u] \"\n \"offset[%u]\",\n streamer_meta_.data_chunk_size, streamer_meta_.offset_chunk_size,\n meta_data_chunk_size, meta_offset_chunk_size);\n return IndexError_InvalidFormat;\n }\n\n // check chunk cnt\n if (streamer_meta_.data_chunk_count > max_data_chunk_cnt_ ||\n meta_.doc_cnt > max_doc_cnt_) {\n LOG_ERROR(\n \"Invalid data chunk count[%u] doc count[%u], expect less than \"\n \"chunk count[%u] doc count[%u]\",\n streamer_meta_.data_chunk_count, meta_.doc_cnt, max_data_chunk_cnt_,\n max_doc_cnt_);\n return IndexError_InvalidFormat;\n }\n\n // load offset chunks\n for (size_t i = 0; i < streamer_meta_.offset_chunk_count; ++i) {\n std::string segment_id =\n ailego::StringHelper::Concat(PARAM_FLAT_SPARSE_OFFSET_SEG_ID_PREFIX, i);\n segment = storage->get(segment_id);\n if (!segment) {\n LOG_ERROR(\"Missing segment %s\", segment_id.c_str());\n return IndexError_InvalidFormat;\n }\n sparse_offset_chunks_.emplace_back(segment);\n index_size += segment->capacity();\n }\n // load data chunks\n for (size_t i = 0; i < streamer_meta_.data_chunk_count; ++i) {\n std::string segment_id =\n ailego::StringHelper::Concat(PARAM_FLAT_SPARSE_DATA_SEG_ID_PREFIX, i);\n segment = storage->get(segment_id);\n if (!segment) {\n LOG_ERROR(\"Missing segment %s\", segment_id.c_str());\n }\n sparse_data_chunks_.emplace_back(segment);\n index_size += segment->capacity();\n }\n\n // load keys\n for (node_id_t i = 0; i < meta_.doc_cnt; ++i) {\n (*keys_map_)[get_key(i)] = i;\n }\n\n stats_.set_index_size(index_size);\n stats_.set_check_point(storage->check_point());\n stats_.set_create_time(meta_.create_time);\n stats_.set_update_time(meta_.update_time);\n stats_.set_loaded_count(keys_map_->size());\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::init_storage(IndexStorage::Pointer storage,\n const IndexMeta &meta) {\n meta_.create_time = ailego::Realtime::Seconds();\n stats_.set_create_time(meta_.create_time);\n meta_.update_time = ailego::Realtime::Seconds();\n stats_.set_update_time(meta_.update_time);\n meta_.doc_cnt = 0;\n\n meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_DATA_CHUNK_SIZE,\n &streamer_meta_.data_chunk_size);\n meta.streamer_params().get(PARAM_FLAT_SPARSE_STREAMER_OFFSET_CHUNK_SIZE,\n &streamer_meta_.offset_chunk_size);\n\n // append meta segment\n size_t size = ailego_align(sizeof(meta_), ailego::MemoryHelper::PageSize());\n int ret = storage->append(PARAM_FLAT_SPARSE_META_SEG_ID, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to append meta segment %s\",\n PARAM_FLAT_SPARSE_META_SEG_ID.c_str());\n return ret;\n }\n auto segment = storage->get(PARAM_FLAT_SPARSE_META_SEG_ID);\n if (ailego_unlikely(!segment)) {\n LOG_ERROR(\"Failed to get meta segment %s\",\n PARAM_FLAT_SPARSE_META_SEG_ID.c_str());\n return IndexError_Runtime;\n }\n if (segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) {\n LOG_ERROR(\"Failed to write meta segment %s\",\n PARAM_FLAT_SPARSE_META_SEG_ID.c_str());\n return IndexError_WriteData;\n }\n\n *stats_.mutable_index_size() += size;\n\n // append streamer meta segment\n size = ailego_align(sizeof(streamer_meta_), ailego::MemoryHelper::PageSize());\n ret = storage->append(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to append streamer meta segment %s\",\n PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str());\n return ret;\n }\n segment = storage->get(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID);\n if (ailego_unlikely(!segment)) {\n LOG_ERROR(\"Failed to get streamer meta segment %s\",\n PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str());\n return IndexError_Runtime;\n }\n if (segment->write(0, &streamer_meta_, sizeof(streamer_meta_)) !=\n sizeof(streamer_meta_)) {\n LOG_ERROR(\"Failed to write streamer meta segment %s\",\n PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str());\n return IndexError_WriteData;\n }\n\n *stats_.mutable_index_size() += size;\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::close() {\n storage_.reset();\n sparse_data_chunks_.clear();\n sparse_offset_chunks_.clear();\n\n keys_map_lock_.reset();\n keys_map_.reset();\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::flush(uint64_t checkpoint) {\n // flush meta\n meta_.update_time = ailego::Realtime::Seconds();\n stats_.set_update_time(meta_.update_time);\n auto segment = storage_->get(PARAM_FLAT_SPARSE_META_SEG_ID);\n if (ailego_unlikely(!segment)) {\n LOG_ERROR(\"Failed to get meta segment %s\",\n PARAM_FLAT_SPARSE_META_SEG_ID.c_str());\n return IndexError_Runtime;\n }\n if (segment->write(0, &meta_, sizeof(meta_)) != sizeof(meta_)) {\n LOG_ERROR(\"Failed to write meta segment %s\",\n PARAM_FLAT_SPARSE_META_SEG_ID.c_str());\n return IndexError_WriteData;\n }\n\n // flush streamer meta\n streamer_meta_.data_chunk_count = sparse_data_chunks_.size();\n streamer_meta_.offset_chunk_count = sparse_offset_chunks_.size();\n segment = storage_->get(PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID);\n if (ailego_unlikely(!segment)) {\n LOG_ERROR(\"Failed to get streamer meta segment %s\",\n PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str());\n return IndexError_Runtime;\n }\n if (segment->write(0, &streamer_meta_, sizeof(streamer_meta_)) !=\n sizeof(streamer_meta_)) {\n LOG_ERROR(\"Failed to write streamer meta segment %s\",\n PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID.c_str());\n return IndexError_WriteData;\n }\n\n if (checkpoint != 0) {\n storage_->refresh(checkpoint);\n }\n int ret = storage_->flush();\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Failed to flush storage for %s\", IndexError::What(ret));\n return ret;\n }\n if (checkpoint != 0) {\n stats_.set_check_point(checkpoint);\n }\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::dump(const IndexDumper::Pointer &dumper) {\n ailego::ElapsedTime stamp;\n\n int ret;\n // meta\n ret = dump_meta(dumper.get());\n if (ret != 0) {\n return ret;\n }\n\n auto duration_dump_meta = stamp.milli_seconds();\n\n // offset & data\n ret = dump_offset_data(dumper.get());\n if (ret != 0) {\n return ret;\n }\n\n auto duration_dump_offset_data = stamp.milli_seconds() - duration_dump_meta;\n\n // keys\n std::vector keys = get_keys();\n ret = dump_keys(keys, dumper.get());\n if (ret != 0) {\n return ret;\n }\n\n auto duration_dump_keys =\n stamp.milli_seconds() - duration_dump_offset_data - duration_dump_meta;\n\n // mapping\n ret = dump_mapping(keys, dumper.get());\n if (ret != 0) {\n return ret;\n }\n\n auto duration_dump_mapping = stamp.milli_seconds() -\n duration_dump_offset_data - duration_dump_meta -\n duration_dump_keys;\n\n LOG_INFO(\n \"Dump index meta: %zu ms, offset & data: %zu ms, keys: %zu ms, \"\n \"mapping: %zu ms\",\n (size_t)duration_dump_meta, (size_t)duration_dump_offset_data,\n (size_t)duration_dump_keys, (size_t)duration_dump_mapping);\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::dump_offset_data(IndexDumper *dumper) {\n ailego::ElapsedTime stamp;\n\n uint64_t init_offset = dump_size_;\n std::vector> offset_length;\n\n // write data\n int ret;\n node_id_t total_doc_cnt = doc_cnt();\n for (node_id_t node_id = 0; node_id < total_doc_cnt; node_id++) {\n uint32_t target_vector_len;\n IndexStorage::MemoryBlock target_vector_block;\n ret = get_sparse_vector_ptr_by_id(node_id, target_vector_block,\n &target_vector_len);\n if (ret != 0) {\n LOG_ERROR(\"Failed to get vector, node_id=%u, error: %s\", node_id,\n IndexError::What(ret));\n return ret;\n }\n const void *target_vector = target_vector_block.data();\n ret = dump_sparse_vector_data(target_vector, target_vector_len, dumper);\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump sparse vector data, node_id=%u, error: %s\",\n node_id, IndexError::What(ret));\n return ret;\n }\n\n offset_length.push_back({dump_size_ - init_offset, target_vector_len});\n dump_size_ += target_vector_len;\n }\n\n // append data segment\n if (dumper->append(PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID,\n dump_size_ - init_offset, 0, 0) != 0) {\n LOG_ERROR(\"append data segment failed\");\n return IndexError_WriteData;\n }\n\n auto duration_dump_data = stamp.milli_seconds();\n\n // write offset\n for (auto &offset_length_pair : offset_length) {\n if (dumper->write(&offset_length_pair.first,\n sizeof(offset_length_pair.first)) !=\n sizeof(offset_length_pair.first)) {\n return IndexError_WriteData;\n }\n if (dumper->write(&offset_length_pair.second,\n sizeof(offset_length_pair.second)) !=\n sizeof(offset_length_pair.second)) {\n return IndexError_WriteData;\n }\n dump_size_ +=\n sizeof(offset_length_pair.first) + sizeof(offset_length_pair.second);\n }\n\n // append offset segment\n if (dumper->append(\n PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID,\n offset_length.size() * (sizeof(uint64_t) + sizeof(uint32_t)), 0,\n 0) != 0) {\n LOG_ERROR(\"append offset segment failed\");\n return IndexError_WriteData;\n }\n\n auto duration_dump_offset = stamp.milli_seconds() - duration_dump_data;\n\n LOG_INFO(\"Dump offset: %zu ms, data: %zu ms\", (size_t)duration_dump_offset,\n (size_t)duration_dump_data);\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::dump_sparse_vector_data(const void *data,\n uint32_t length,\n IndexDumper *dumper) {\n if (dumper->write(data, length) != length) {\n return IndexError_WriteData;\n }\n return 0;\n}\n\nint FlatSparseStreamerEntity::dump_meta(IndexDumper *dumper) {\n if (dumper->write(&meta_, sizeof(meta_)) != sizeof(meta_)) {\n LOG_ERROR(\"write meta failed\");\n return IndexError_WriteData;\n }\n\n size_t meta_padding_size = ailego_align(sizeof(meta_), 32) - sizeof(meta_);\n if (meta_padding_size) {\n std::string padding(meta_padding_size, '\\0');\n if (dumper->write(padding.data(), meta_padding_size) != meta_padding_size) {\n LOG_ERROR(\"write meta padding failed\");\n return IndexError_WriteData;\n }\n }\n return dumper->append(PARAM_FLAT_SPARSE_META_SEG_ID, sizeof(meta_),\n meta_padding_size, 0);\n}\n\nint FlatSparseStreamerEntity::dump_keys(const std::vector &keys,\n IndexDumper *dumper) {\n if (keys.size() == 1 && keys.back() == kInvalidKey) {\n return IndexError_Runtime;\n }\n\n size_t keys_size = keys.size() * sizeof(uint64_t);\n if (dumper->write(keys.data(), keys_size) != keys_size) {\n LOG_ERROR(\"Failed to write keys to dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n size_t keys_padding_size = ailego_align(keys_size, 32) - keys_size;\n if (keys_padding_size) {\n std::string padding(keys_padding_size, '\\0');\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write padding to dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n }\n return dumper->append(PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID, keys_size,\n keys_padding_size, 0);\n}\n\nint FlatSparseStreamerEntity::dump_mapping(const std::vector &keys,\n IndexDumper *dumper) {\n std::vector mapping(keys.size());\n std::iota(mapping.begin(), mapping.end(), 0);\n std::sort(\n mapping.begin(), mapping.end(),\n [&keys](uint32_t lhs, uint32_t rhs) { return (keys[lhs] < keys[rhs]); });\n\n size_t mapping_size = mapping.size() * sizeof(uint32_t);\n size_t mapping_padding_size = ailego_align(mapping_size, 32) - mapping_size;\n if (dumper->write(mapping.data(), mapping_size) != mapping_size) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n\n // Write the padding if need\n if (mapping_padding_size) {\n std::string padding(mapping_padding_size, '\\0');\n if (dumper->write(padding.data(), padding.size()) != padding.size()) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper->name().c_str());\n return IndexError_WriteData;\n }\n }\n return dumper->append(PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID, mapping_size,\n mapping_padding_size, 0);\n}\n\nFlatSparseStreamerEntity::Pointer FlatSparseStreamerEntity::clone() const {\n auto entity = new (std::nothrow) FlatSparseStreamerEntity(\n stats_, meta_, streamer_meta_, keys_map_lock_, keys_map_,\n sparse_data_chunks_, sparse_offset_chunks_);\n return FlatSparseStreamerEntity::Pointer(entity);\n}\n\nint FlatSparseStreamerEntity::add(uint64_t key,\n const std::string &sparse_vector,\n const uint32_t sparse_count) {\n uint32_t sparse_vector_len = sparse_vector.size();\n\n sparse_vector_len = AlignSize(sparse_vector_len);\n\n if (sparse_vector_len > streamer_meta_.data_chunk_size) {\n LOG_ERROR(\n \"Sparse Vector Length exceed the chunk size, sparse vec len: %u, chunk \"\n \"size: %u\",\n sparse_vector_len, streamer_meta_.data_chunk_size);\n return IndexError_InvalidArgument;\n }\n\n std::lock_guard lock(mutex_);\n node_id_t local_id = doc_cnt();\n\n if (ailego_unlikely(local_id >= max_doc_cnt_)) {\n LOG_ERROR(\"Add vector failed for exceed max doc count: %u\", max_doc_cnt_);\n return IndexError_IndexFull;\n }\n\n // duplicate check\n if (ailego_unlikely(get_id(key) != kInvalidNodeId)) {\n LOG_WARN(\"Try to add duplicate key, ignore it\");\n return IndexError_Duplicate;\n }\n\n // get sparse data chunk and offset for write sparse vector\n Chunk::Pointer sparse_data_chunk;\n uint32_t sparse_data_chunk_offset = -1U;\n uint32_t sparse_data_chunk_index = sparse_data_chunks_.size() - 1U;\n if (sparse_data_chunk_index == -1U ||\n sparse_data_chunks_[sparse_data_chunk_index]->data_size() +\n sparse_vector_len >\n streamer_meta_.data_chunk_size) {\n if (ailego_unlikely(sparse_data_chunks_.capacity() ==\n sparse_data_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n if (sparse_data_chunk_index != -1U) {\n LOG_ERROR(\n \"capacity: %zu, chunk used size: %zu, chunk size: %u, \"\n \"sparse_vector_len: %u\",\n sparse_data_chunks_.capacity(),\n sparse_data_chunks_[sparse_data_chunk_index]->data_size(),\n streamer_meta_.data_chunk_size, sparse_vector_len);\n }\n return IndexError_IndexFull;\n }\n\n sparse_data_chunk = alloc_new_data_chunk(sparse_data_chunks_.size());\n if (ailego_unlikely(!sparse_data_chunk)) {\n LOG_ERROR(\"allocate data chunk failed\");\n return IndexError_NoMemory;\n }\n sparse_data_chunks_.emplace_back(sparse_data_chunk);\n sparse_data_chunk_index = sparse_data_chunks_.size() - 1U;\n sparse_data_chunk_offset = 0UL;\n } else {\n sparse_data_chunk = sparse_data_chunks_[sparse_data_chunk_index];\n sparse_data_chunk_offset = sparse_data_chunk->data_size();\n }\n\n // write sparse vector\n if (sparse_vector.size() > 0) {\n if (ailego_unlikely(write_sparse_vector_data(\n sparse_data_chunk_index, sparse_data_chunk_offset,\n sparse_vector.data(), sparse_vector.size()) != 0)) {\n LOG_ERROR(\"write sparse vector failed\");\n return IndexError_NoMemory;\n }\n }\n\n uint64_t sparse_offset = sparse_data_chunk_index;\n sparse_offset = (sparse_offset << 32U) + sparse_data_chunk_offset;\n\n // get sparse offset chunk and offset for write new info\n Chunk::Pointer sparse_offset_chunk;\n uint32_t sparse_offset_chunk_offset = -1U;\n uint32_t sparse_offset_chunk_index = sparse_offset_chunks_.size() - 1U;\n if (sparse_offset_chunk_index == -1U ||\n sparse_offset_chunks_[sparse_offset_chunk_index]->data_size() +\n offset_size_per_node() >\n streamer_meta_.offset_chunk_size) {\n // no space left and need to allocate new offset chunk\n if (ailego_unlikely(sparse_offset_chunks_.capacity() ==\n sparse_offset_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n\n sparse_offset_chunk = alloc_new_offset_chunk(sparse_offset_chunks_.size());\n if (ailego_unlikely(!sparse_offset_chunk)) {\n LOG_ERROR(\"allocate offset chunk failed\");\n return IndexError_NoMemory;\n }\n sparse_offset_chunks_.emplace_back(sparse_offset_chunk);\n sparse_offset_chunk_index = sparse_offset_chunks_.size() - 1U;\n sparse_offset_chunk_offset = 0UL;\n } else {\n sparse_offset_chunk = sparse_offset_chunks_[sparse_offset_chunk_index];\n sparse_offset_chunk_offset = sparse_offset_chunk->data_size();\n }\n\n // write offset\n size_t size = sparse_offset_chunk->write(sparse_offset_chunk_offset,\n &sparse_offset, sizeof(uint64_t));\n if (ailego_unlikely(size != sizeof(uint64_t))) {\n LOG_ERROR(\"Chunk write sparse vec offset failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n // write length\n size =\n sparse_offset_chunk->write(sparse_offset_chunk_offset + sizeof(uint64_t),\n &sparse_vector_len, sizeof(uint32_t));\n if (ailego_unlikely(size != sizeof(uint32_t))) {\n LOG_ERROR(\"Chunk write sparse vec len failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n // write key\n size = sparse_offset_chunk->write(\n sparse_offset_chunk_offset + 2 * sizeof(uint64_t), &key,\n sizeof(uint64_t));\n if (ailego_unlikely(size != sizeof(uint64_t))) {\n LOG_ERROR(\"Chunk write key failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n // LOG_INFO(\"Write sparse vector, key=%lu, offset chunk=%u, offset=%u,\n // len=%u\",\n // key, sparse_offset_chunk_index, sparse_offset_chunk_offset,\n // offset_size_per_node());\n\n // LOG_INFO(\"Write sparse vector, key=%lu, data chunk=%u, offset=%u, len=%u\",\n // key, sparse_data_chunk_index, sparse_data_chunk_offset,\n // sparse_vector_len);\n\n // resize chunk\n if (sparse_vector_len > 0) {\n sparse_data_chunk_offset += sparse_vector_len;\n if (ailego_unlikely(sparse_data_chunk->resize(sparse_data_chunk_offset) !=\n sparse_data_chunk_offset)) {\n LOG_ERROR(\"Sparse Chunk resize to %u failed\", sparse_data_chunk_offset);\n return IndexError_Runtime;\n }\n }\n\n // persist in keys_map\n {\n keys_map_lock_->lock();\n (*keys_map_)[key] = local_id;\n keys_map_lock_->unlock();\n }\n\n inc_doc_count();\n inc_total_sparse_count(sparse_count);\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::add_vector_with_id(\n uint32_t id, const std::string &sparse_vector,\n const uint32_t sparse_count) {\n uint32_t sparse_vector_len = sparse_vector.size();\n\n sparse_vector_len = AlignSize(sparse_vector_len);\n\n if (sparse_vector_len > streamer_meta_.data_chunk_size) {\n LOG_ERROR(\n \"Sparse Vector Length exceed the chunk size, sparse vec len: %u, chunk \"\n \"size: %u\",\n sparse_vector_len, streamer_meta_.data_chunk_size);\n return IndexError_InvalidArgument;\n }\n\n std::lock_guard lock(mutex_);\n\n if (id >= doc_cnt()) {\n for (auto i = doc_cnt(); i <= id; i++) {\n node_id_t local_id = doc_cnt();\n if (ailego_unlikely(local_id >= max_doc_cnt_)) {\n LOG_ERROR(\"Add vector failed for exceed max doc count: %u\",\n max_doc_cnt_);\n return IndexError_IndexFull;\n }\n uint32_t sparse_data_chunk_index, sparse_data_chunk_offset,\n sparse_offset_chunk_index, sparse_offset_chunk_offset;\n if (i < id) {\n write_sparse_vector_to_chunk(\"\", 0, sparse_data_chunk_index,\n sparse_data_chunk_offset);\n } else {\n write_sparse_vector_to_chunk(sparse_vector, sparse_vector_len,\n sparse_data_chunk_index,\n sparse_data_chunk_offset);\n }\n uint64_t sparse_offset =\n ((uint64_t)sparse_data_chunk_index << 32U) + sparse_data_chunk_offset;\n get_new_sparse_offset_chunk(sparse_offset_chunk_index,\n sparse_offset_chunk_offset);\n uint64_t written_key = kInvalidKey;\n if (i == id) {\n written_key = i;\n }\n write_sparse_offset_to_chunk(sparse_offset_chunk_index,\n sparse_offset_chunk_offset, sparse_offset,\n sparse_vector_len, written_key);\n {\n keys_map_lock_->lock();\n (*keys_map_)[i] = written_key;\n keys_map_lock_->unlock();\n }\n inc_doc_count();\n }\n } else {\n uint32_t sparse_data_chunk_index, sparse_data_chunk_offset;\n write_sparse_vector_to_chunk(sparse_vector, sparse_vector_len,\n sparse_data_chunk_index,\n sparse_data_chunk_offset);\n uint64_t sparse_offset =\n ((uint64_t)sparse_data_chunk_index << 32U) + sparse_data_chunk_offset;\n uint32_t sparse_offset_chunk_index =\n id / get_offset_info_number_per_chunk();\n uint32_t sparse_offset_chunk_offset =\n id % get_offset_info_number_per_chunk() * offset_size_per_node();\n write_sparse_offset_to_chunk(sparse_offset_chunk_index,\n sparse_offset_chunk_offset, sparse_offset,\n sparse_vector_len, id);\n {\n keys_map_lock_->lock();\n (*keys_map_)[id] = id;\n keys_map_lock_->unlock();\n }\n }\n inc_total_sparse_count(sparse_count);\n return 0;\n}\n\nint FlatSparseStreamerEntity::write_sparse_vector_to_chunk(\n const std::string &sparse_vector, const uint32_t sparse_vector_len,\n uint32_t &sparse_data_chunk_index, uint32_t &sparse_data_chunk_offset) {\n // get sparse data chunk and offset for write sparse vector\n Chunk::Pointer sparse_data_chunk;\n sparse_data_chunk_offset = -1U;\n sparse_data_chunk_index = sparse_data_chunks_.size() - 1U;\n if (sparse_data_chunk_index == -1U ||\n sparse_data_chunks_[sparse_data_chunk_index]->data_size() +\n sparse_vector_len >\n streamer_meta_.data_chunk_size) {\n if (ailego_unlikely(sparse_data_chunks_.capacity() ==\n sparse_data_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n if (sparse_data_chunk_index != -1U) {\n LOG_ERROR(\n \"capacity: %zu, chunk used size: %zu, chunk size: %u, \"\n \"sparse_vector_len: %u\",\n sparse_data_chunks_.capacity(),\n sparse_data_chunks_[sparse_data_chunk_index]->data_size(),\n streamer_meta_.data_chunk_size, sparse_vector_len);\n }\n return IndexError_IndexFull;\n }\n\n sparse_data_chunk = alloc_new_data_chunk(sparse_data_chunks_.size());\n if (ailego_unlikely(!sparse_data_chunk)) {\n LOG_ERROR(\"allocate data chunk failed\");\n return IndexError_NoMemory;\n }\n sparse_data_chunks_.emplace_back(sparse_data_chunk);\n sparse_data_chunk_index = sparse_data_chunks_.size() - 1U;\n sparse_data_chunk_offset = 0UL;\n } else {\n sparse_data_chunk = sparse_data_chunks_[sparse_data_chunk_index];\n sparse_data_chunk_offset = sparse_data_chunk->data_size();\n }\n\n // write sparse vector\n if (sparse_vector.size() > 0) {\n if (ailego_unlikely(write_sparse_vector_data(\n sparse_data_chunk_index, sparse_data_chunk_offset,\n sparse_vector.data(), sparse_vector.size()) != 0)) {\n LOG_ERROR(\"write sparse vector failed\");\n return IndexError_NoMemory;\n }\n }\n\n // resize chunk\n if (sparse_vector_len > 0) {\n uint32_t sparse_data_chunk_size =\n sparse_data_chunk_offset + sparse_vector_len;\n if (ailego_unlikely(sparse_data_chunk->resize(sparse_data_chunk_size) !=\n sparse_data_chunk_size)) {\n LOG_ERROR(\"Sparse Chunk resize to %u failed\", sparse_data_chunk_size);\n return IndexError_Runtime;\n }\n }\n return 0;\n}\n\nint FlatSparseStreamerEntity::get_new_sparse_offset_chunk(\n uint32_t &sparse_offset_chunk_index, uint32_t &sparse_offset_chunk_offset) {\n // get sparse offset chunk and offset for write new info\n Chunk::Pointer sparse_offset_chunk;\n sparse_offset_chunk_offset = -1U;\n sparse_offset_chunk_index = sparse_offset_chunks_.size() - 1U;\n if (sparse_offset_chunk_index == -1U ||\n sparse_offset_chunks_[sparse_offset_chunk_index]->data_size() +\n offset_size_per_node() >\n streamer_meta_.offset_chunk_size) {\n // no space left and need to allocate new offset chunk\n if (ailego_unlikely(sparse_offset_chunks_.capacity() ==\n sparse_offset_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n\n sparse_offset_chunk = alloc_new_offset_chunk(sparse_offset_chunks_.size());\n if (ailego_unlikely(!sparse_offset_chunk)) {\n LOG_ERROR(\"allocate offset chunk failed\");\n return IndexError_NoMemory;\n }\n sparse_offset_chunks_.emplace_back(sparse_offset_chunk);\n sparse_offset_chunk_index = sparse_offset_chunks_.size() - 1U;\n sparse_offset_chunk_offset = 0UL;\n } else {\n sparse_offset_chunk = sparse_offset_chunks_[sparse_offset_chunk_index];\n sparse_offset_chunk_offset = sparse_offset_chunk->data_size();\n }\n return 0;\n}\n\nint FlatSparseStreamerEntity::write_sparse_offset_to_chunk(\n const uint32_t sparse_offset_chunk_index,\n const uint32_t sparse_offset_chunk_offset, const uint64_t sparse_offset,\n const uint32_t sparse_vector_len, const uint64_t node_id) {\n // write offset\n Chunk::Pointer sparse_offset_chunk =\n sparse_offset_chunks_[sparse_offset_chunk_index];\n size_t size = sparse_offset_chunk->write(sparse_offset_chunk_offset,\n &sparse_offset, sizeof(uint64_t));\n if (ailego_unlikely(size != sizeof(uint64_t))) {\n LOG_ERROR(\"Chunk write sparse vec offset failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n // write length\n size =\n sparse_offset_chunk->write(sparse_offset_chunk_offset + sizeof(uint64_t),\n &sparse_vector_len, sizeof(uint32_t));\n if (ailego_unlikely(size != sizeof(uint32_t))) {\n LOG_ERROR(\"Chunk write sparse vec len failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n // write key\n size = sparse_offset_chunk->write(\n sparse_offset_chunk_offset + 2 * sizeof(uint64_t), &node_id,\n sizeof(uint64_t));\n if (ailego_unlikely(size != sizeof(uint64_t))) {\n LOG_ERROR(\"Chunk write key failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n return 0;\n}\n\nuint64_t FlatSparseStreamerEntity::get_key(node_id_t node_id) const {\n uint32_t offset_chunk_index = node_id / get_offset_info_number_per_chunk();\n uint32_t offset_chunk_key_offset =\n node_id % get_offset_info_number_per_chunk() * offset_size_per_node() +\n 2 * sizeof(uint64_t);\n\n IndexStorage::MemoryBlock block;\n if (ailego_unlikely(sparse_offset_chunks_[offset_chunk_index]->read(\n offset_chunk_key_offset, block, sizeof(uint64_t)) !=\n sizeof(uint64_t))) {\n LOG_ERROR(\"Read key failed, offset=%u, node_id=%u\", offset_chunk_key_offset,\n node_id);\n return kInvalidKey;\n };\n\n return *reinterpret_cast(block.data());\n}\n\nint FlatSparseStreamerEntity::get_sparse_vector_ptr_by_id(\n node_id_t node_id, const void **sparse_vector_ptr,\n uint32_t *sparse_vector_len_ptr) const {\n uint32_t offset_chunk_index = node_id / get_offset_info_number_per_chunk();\n uint32_t offset_chunk_offset =\n node_id % get_offset_info_number_per_chunk() * offset_size_per_node();\n\n // LOG_DEBUG(\"Read sparse vector, offset chunk=%u, offset=%u, len=%u\",\n // offset_chunk_index, offset_chunk_offset, offset_size_per_node());\n\n auto offset_chunk = sparse_offset_chunks_[offset_chunk_index];\n\n const void *offset_info = nullptr;\n size_t read_len = offset_chunk->read(offset_chunk_offset, &offset_info,\n offset_size_per_node());\n if (ailego_unlikely(read_len != offset_size_per_node())) {\n LOG_ERROR(\"Read offset info failed, offset=%u, read_len=%zu, expect=%u\",\n offset_chunk_offset, read_len, offset_size_per_node());\n return IndexError_ReadData;\n };\n\n // sparse offset\n uint64_t sparse_offset = *(uint64_t *)offset_info;\n uint32_t sparse_vector_len =\n *(uint32_t *)((uint8_t *)offset_info + sizeof(uint64_t));\n\n uint32_t sparse_data_chunk_index =\n static_cast((sparse_offset >> 32) & 0xFFFFFFFF);\n uint32_t sparse_data_chunk_offset =\n static_cast(sparse_offset & 0xFFFFFFFF);\n\n if (sparse_vector_len > 0) {\n const void *sparse_data = get_sparse_vector_data(\n sparse_data_chunk_index, sparse_data_chunk_offset, sparse_vector_len);\n if (ailego_unlikely(sparse_data == nullptr)) {\n LOG_ERROR(\"Get nullptr sparse, offset=%zu, len=%u\", (size_t)sparse_offset,\n sparse_vector_len);\n\n return IndexError_ReadData;\n }\n *sparse_vector_ptr = sparse_data;\n *sparse_vector_len_ptr = sparse_vector_len;\n }\n\n // LOG_DEBUG(\"Read sparse vector, data chunk=%u, offset=%u, len=%u\",\n // sparse_data_chunk_index, sparse_data_chunk_offset,\n // sparse_vector_len);\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::get_sparse_vector_ptr_by_id(\n node_id_t node_id, IndexStorage::MemoryBlock &sparse_vector_block,\n uint32_t *sparse_vector_len_ptr) const {\n uint32_t offset_chunk_index = node_id / get_offset_info_number_per_chunk();\n uint32_t offset_chunk_offset =\n node_id % get_offset_info_number_per_chunk() * offset_size_per_node();\n\n // LOG_DEBUG(\"Read sparse vector, offset chunk=%u, offset=%u, len=%u\",\n // offset_chunk_index, offset_chunk_offset, offset_size_per_node());\n\n auto offset_chunk = sparse_offset_chunks_[offset_chunk_index];\n\n const void *offset_info = nullptr;\n IndexStorage::MemoryBlock offset_info_block;\n size_t read_len = offset_chunk->read(offset_chunk_offset, offset_info_block,\n offset_size_per_node());\n if (ailego_unlikely(read_len != offset_size_per_node())) {\n LOG_ERROR(\"Read offset info failed, offset=%u, read_len=%zu, expect=%u\",\n offset_chunk_offset, read_len, offset_size_per_node());\n return IndexError_ReadData;\n };\n offset_info = offset_info_block.data();\n\n // sparse offset\n uint64_t sparse_offset = *(uint64_t *)offset_info;\n uint32_t sparse_vector_len =\n *(uint32_t *)((uint8_t *)offset_info + sizeof(uint64_t));\n\n uint32_t sparse_data_chunk_index =\n static_cast((sparse_offset >> 32) & 0xFFFFFFFF);\n uint32_t sparse_data_chunk_offset =\n static_cast(sparse_offset & 0xFFFFFFFF);\n\n if (sparse_vector_len > 0) {\n get_sparse_vector_data(sparse_data_chunk_index, sparse_data_chunk_offset,\n sparse_vector_len, sparse_vector_block);\n if (ailego_unlikely(sparse_vector_block.data() == nullptr)) {\n LOG_ERROR(\"Get nullptr sparse, offset=%zu, len=%u\", (size_t)sparse_offset,\n sparse_vector_len);\n\n return IndexError_ReadData;\n }\n *sparse_vector_len_ptr = sparse_vector_len;\n }\n\n return 0;\n}\n\nint FlatSparseStreamerEntity::write_sparse_vector_data(uint32_t chunk_index,\n uint64_t offset,\n const void *data,\n uint32_t length) {\n auto size = sparse_data_chunks_[chunk_index]->write(offset, data, length);\n if (size != length) {\n LOG_ERROR(\n \"write sparse vector data failed: chunk_index=%u, offset=%zu, \"\n \"length=%u, size=%zu, chunk_data_size=%zu\",\n chunk_index, (size_t)offset, length, size,\n sparse_data_chunks_[chunk_index]->data_size());\n return IndexError_WriteData;\n }\n // LOG_DEBUG(\n // \"write_sparse_vector_data: chunk_index=%u, offset=%lu, length=%u, \"\n // \"data=%p\",\n // chunk_index, offset, length, data);\n return 0;\n}\n\nconst void *FlatSparseStreamerEntity::get_sparse_vector_data(\n uint32_t chunk_index, uint64_t offset, uint32_t length) const {\n const void *data;\n auto size = sparse_data_chunks_[chunk_index]->read(offset, &data, length);\n if (size != length) {\n LOG_ERROR(\n \"read sparse vector data failed: chunk_index=%u, offset=%zu, \"\n \"length=%u, size=%zu\",\n chunk_index, (size_t)offset, length, size);\n return nullptr;\n }\n // LOG_DEBUG(\n // \"get_sparse_vector_data: chunk_index=%u, offset=%lu, length=%u, \"\n // \"data=%p\",\n // chunk_index, offset, length, data);\n return data;\n}\n\nint FlatSparseStreamerEntity::get_sparse_vector_data(\n uint32_t chunk_index, uint64_t offset, uint32_t length,\n IndexStorage::MemoryBlock &block) const {\n auto size = sparse_data_chunks_[chunk_index]->read(offset, block, length);\n if (size != length) {\n LOG_ERROR(\n \"read sparse vector data failed: chunk_index=%u, offset=%zu, \"\n \"length=%u, size=%zu\",\n chunk_index, (size_t)offset, length, size);\n return IndexError_ReadData;\n }\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_streamer_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"flat_sparse_entity.h\"\n#include \"flat_sparse_index_format.h\"\n#include \"flat_sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n\n/*! Flat Sparse Streamer Entity\n */\nclass FlatSparseStreamerEntity : public FlatSparseEntity {\n public:\n typedef std::shared_ptr Pointer;\n\n using Chunk = IndexStorage::Segment;\n\n //! Constructor\n explicit FlatSparseStreamerEntity(IndexStreamer::Stats &stats);\n\n //! Destructor\n virtual ~FlatSparseStreamerEntity() = default;\n\n //! Disable them\n FlatSparseStreamerEntity(const FlatSparseStreamerEntity &) = delete;\n FlatSparseStreamerEntity &operator=(const FlatSparseStreamerEntity &) =\n delete;\n\n //! Open the entity with storage\n int open(IndexStorage::Pointer storage, const IndexMeta &meta);\n\n //! Close the entity\n int close();\n\n //! Flush linear index to storage\n int flush(uint64_t checkpoint);\n\n //! Dump index by dumper\n int dump(const IndexDumper::Pointer &dumper);\n\n //! Add sparse vector to linear index\n int add(uint64_t key, const std::string &sparse_vector,\n const uint32_t sparse_count);\n\n //! Add sparse vector to linear index with id\n int add_vector_with_id(uint32_t id, const std::string &sparse_vector,\n uint32_t sparse_count);\n\n //! Clone entity\n FlatSparseStreamerEntity::Pointer clone() const;\n\n int get_index_sparse_meta(IndexMeta *meta) const {\n return IndexHelper::DeserializeFromStorage(storage_.get(), meta);\n }\n\n int set_index_sparse_meta(const IndexMeta &meta) const {\n return IndexHelper::SerializeToStorage(meta, storage_.get());\n }\n\n public:\n inline uint32_t doc_cnt() const override {\n return meta_.doc_cnt;\n }\n\n inline uint32_t total_sparse_count() const override {\n return meta_.total_sparse_count;\n }\n\n size_t sparse_unit_size() const override {\n return sparse_unit_size_;\n }\n\n inline node_id_t get_id(uint64_t key) const override {\n keys_map_lock_->lock_shared();\n auto it = keys_map_->find(key);\n keys_map_lock_->unlock_shared();\n return it == keys_map_->end() ? kInvalidNodeId : it->second;\n }\n\n uint64_t get_key(node_id_t node_id) const override;\n\n int get_sparse_vector_ptr_by_id(node_id_t id, const void **sparse_vector,\n uint32_t *sparse_vector_len) const override;\n int get_sparse_vector_ptr_by_id(\n const node_id_t id, IndexStorage::MemoryBlock &sparse_vector_block,\n uint32_t *sparse_vector_len) const override;\n\n float get_search_distance(const std::string &vector,\n node_id_t target_node_id) const override {\n float dist;\n const void *target_vector;\n uint32_t target_vector_len;\n get_sparse_vector_ptr_by_id(target_node_id, &target_vector,\n &target_vector_len);\n search_sparse_distance_(vector.c_str(), target_vector, &dist);\n return dist;\n }\n\n private:\n void inc_doc_count() {\n meta_.doc_cnt++;\n }\n void inc_total_sparse_count(uint32_t count) {\n meta_.total_sparse_count += count;\n }\n\n int init_metric(const IndexMeta &meta);\n\n int init_storage(IndexStorage::Pointer storage, const IndexMeta &meta);\n\n int load_storage(IndexStorage::Pointer storage, const IndexMeta &meta);\n\n static inline size_t AlignSize(size_t size) {\n return (size + 0x1F) & (~0x1F);\n }\n\n inline uint32_t offset_size_per_node() const {\n return 3 * sizeof(uint64_t);\n }\n\n inline uint32_t doc_cnt_per_offset_chunk() const {\n return streamer_meta_.offset_chunk_size / offset_size_per_node();\n }\n\n Chunk::Pointer alloc_new_offset_chunk(uint32_t chunk_id) {\n std::string segment_id = ailego::StringHelper::Concat(\n PARAM_FLAT_SPARSE_OFFSET_SEG_ID_PREFIX, chunk_id);\n // LOG_INFO(\"Alloc new offset chunk %s\", segment_id.c_str());\n return alloc_new_chunk(segment_id, streamer_meta_.offset_chunk_size);\n }\n\n Chunk::Pointer alloc_new_data_chunk(uint32_t chunk_id) {\n std::string segment_id = ailego::StringHelper::Concat(\n PARAM_FLAT_SPARSE_DATA_SEG_ID_PREFIX, chunk_id);\n // LOG_INFO(\"Alloc new data chunk %s\", segment_id.c_str());\n return alloc_new_chunk(segment_id, streamer_meta_.data_chunk_size);\n }\n\n Chunk::Pointer alloc_new_chunk(const std::string &segment_id, uint32_t size) {\n int ret = storage_->append(segment_id, size);\n if (ailego_unlikely(ret != 0)) {\n return nullptr;\n }\n *stats_.mutable_index_size() += size;\n return storage_->get(segment_id);\n }\n\n inline uint32_t get_offset_info_number_per_chunk() const {\n return streamer_meta_.offset_chunk_size / offset_size_per_node();\n }\n\n int write_sparse_vector_to_chunk(const std::string &sparse_vector,\n const uint32_t sparse_vector_len,\n uint32_t &sparse_data_chunk_index,\n uint32_t &sparse_data_chunk_offset);\n\n int get_new_sparse_offset_chunk(uint32_t &sparse_offset_chunk_index,\n uint32_t &sparse_offset_chunk_offset);\n\n int write_sparse_offset_to_chunk(const uint32_t sparse_offset_chunk_index,\n const uint32_t sparse_offset_chunk_offset,\n const uint64_t sparse_offset,\n const uint32_t sparse_vector_len,\n const uint64_t node_id);\n\n int write_sparse_vector_data(uint32_t chunk_index, uint64_t offset,\n const void *data, uint32_t length);\n\n const void *get_sparse_vector_data(uint32_t chunk_index, uint64_t offset,\n uint32_t length) const;\n\n int get_sparse_vector_data(uint32_t chunk_index, uint64_t offset,\n uint32_t length,\n IndexStorage::MemoryBlock &block) const;\n\n int dump_sparse_vector_data(const void *data, uint32_t length,\n IndexDumper *dumper);\n\n int dump_meta(IndexDumper *dumper);\n\n int dump_index_meta(IndexDumper *dumper);\n\n int dump_keys(const std::vector &keys, IndexDumper *dumper);\n\n int dump_mapping(const std::vector &keys, IndexDumper *dumper);\n\n int dump_offset_data(IndexDumper *dumper);\n\n\n private:\n FlatSparseStreamerEntity(\n IndexStreamer::Stats &stats, const FlatSparseMeta &meta,\n const FlatSparseStreamerMeta &streamer_meta,\n std::shared_ptr keys_map_lock,\n std::shared_ptr> keys_map,\n std::vector sparse_data_chunks,\n std::vector sparse_offset_chunks)\n : stats_(stats),\n meta_(meta),\n streamer_meta_(streamer_meta),\n keys_map_lock_(keys_map_lock),\n keys_map_(keys_map),\n sparse_data_chunks_(std::move(sparse_data_chunks)),\n sparse_offset_chunks_(std::move(sparse_offset_chunks)) {}\n\n private:\n IndexStorage::Pointer storage_{};\n IndexStreamer::Stats &stats_;\n\n // meta\n FlatSparseMeta meta_;\n FlatSparseStreamerMeta streamer_meta_;\n\n // metric\n IndexMetric::Pointer metric_{};\n IndexMetric::MatrixSparseDistance search_sparse_distance_{};\n\n std::mutex mutex_{};\n\n // keys map\n mutable std::shared_ptr keys_map_lock_{};\n std::shared_ptr> keys_map_{};\n\n // chunks\n mutable std::vector sparse_data_chunks_{};\n mutable std::vector sparse_offset_chunks_{};\n\n // config\n uint32_t max_doc_cnt_{1 << 24U}; // 16 million\n uint32_t max_data_chunk_cnt_{\n 1 << 10U}; // 1024, default single_data_chunk_size = 8M,\n // default_total_max = 1024 * 8M = 8G\n\n uint64_t dump_size_{0U};\n size_t sparse_unit_size_{0U};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/flat_sparse/flat_sparse_utility.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\nstatic constexpr uint32_t PARAM_FLAT_SPARSE_MAX_DIM_SIZE = 16384;\nstatic const std::string PARAM_FLAT_SPARSE_META_SEG_ID =\n \"bruteforce_sparse_meta\";\n\n// streamer\nstatic const std::string PARAM_FLAT_SPARSE_STREAMER_META_SEG_ID =\n \"bruteforce_sparse_streamer_meta\";\nstatic const std::string PARAM_FLAT_SPARSE_OFFSET_SEG_ID_PREFIX =\n \"bruteforce_sparse_streamer_offset_\";\nstatic const std::string PARAM_FLAT_SPARSE_DATA_SEG_ID_PREFIX =\n \"bruteforce_sparse_streamer_data_\";\n\n// searcher\nstatic const std::string PARAM_FLAT_SPARSE_DUMP_OFFSET_SEG_ID =\n \"bruteforce_sparse_searcher_offset_segment\";\nstatic const std::string PARAM_FLAT_SPARSE_DUMP_DATA_SEG_ID =\n \"bruteforce_sparse_searcher_data_segment\";\nstatic const std::string PARAM_FLAT_SPARSE_DUMP_KEYS_SEG_ID =\n \"bruteforce_sparse_searcher_keys_segment\";\nstatic const std::string PARAM_FLAT_SPARSE_DUMP_MAPPING_SEG_ID =\n \"bruteforce_sparse_searcher_mapping_segment\";\n\n// streamer\nstatic const std::string PARAM_FLAT_SPARSE_STREAMER_OFFSET_CHUNK_SIZE(\n \"proxima.bruteforce.sparse_streamer.offset_chunk_size\");\n\nstatic const std::string PARAM_FLAT_SPARSE_STREAMER_DATA_CHUNK_SIZE(\n \"proxima.bruteforce.sparse_streamer.data_chunk_size\");\n\nstatic const std::string PARAM_FLAT_SPARSE_STREAMER_MAX_DOC_CNT(\n \"proxima.bruteforce.sparse_streamer.max_doc_cnt\");\n\nstatic const std::string PARAM_FLAT_SPARSE_STREAMER_MAX_DATA_CHUNK_CNT(\n \"proxima.bruteforce.sparse_streamer.max_data_chunk_cnt\");\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\ncc_library(\n NAME core_knn_hnsw \n STATIC SHARED STRICT ALWAYS_LINK\n SRCS *.cc\n LIBS core_framework sparsehash\n INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )\n" }, { "path": "src/core/algorithm/hnsw/hnsw_algorithm.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_algorithm.h\"\n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nHnswAlgorithm::HnswAlgorithm(HnswEntity &entity)\n : entity_(entity),\n mt_(std::chrono::system_clock::now().time_since_epoch().count()),\n lock_pool_(kLockCnt) {}\n\nint HnswAlgorithm::cleanup() {\n return 0;\n}\n\nint HnswAlgorithm::add_node(node_id_t id, level_t level, HnswContext *ctx) {\n spin_lock_.lock();\n\n // std::cout << \"id: \" << id << \", level: \" << level << std::endl;\n\n auto cur_max_level = entity_.cur_max_level();\n auto entry_point = entity_.entry_point();\n if (ailego_unlikely(entry_point == kInvalidNodeId)) {\n entity_.update_ep_and_level(id, level);\n spin_lock_.unlock();\n return 0;\n }\n spin_lock_.unlock();\n\n if (ailego_unlikely(level > cur_max_level)) {\n mutex_.lock();\n // re-check max level\n cur_max_level = entity_.cur_max_level();\n entry_point = entity_.entry_point();\n if (level <= cur_max_level) {\n mutex_.unlock();\n }\n }\n\n level_t cur_level = cur_max_level;\n dist_t dist = ctx->dist_calculator()(entry_point);\n for (; cur_level > level; --cur_level) {\n select_entry_point(cur_level, &entry_point, &dist, ctx);\n }\n\n for (; cur_level >= 0; --cur_level) {\n search_neighbors(cur_level, &entry_point, &dist, ctx->level_topk(cur_level),\n ctx);\n }\n\n // add neighbors from down level to top level, to avoid upper level visible\n // to knn_search but the under layer level not ready\n for (cur_level = 0; cur_level <= level; ++cur_level) {\n add_neighbors(id, cur_level, ctx->level_topk(cur_level), ctx);\n ctx->level_topk(cur_level).clear();\n }\n\n if (ailego_unlikely(level > cur_max_level)) {\n spin_lock_.lock();\n entity_.update_ep_and_level(id, level);\n spin_lock_.unlock();\n mutex_.unlock();\n }\n\n return 0;\n}\n\nint HnswAlgorithm::search(HnswContext *ctx) const {\n spin_lock_.lock();\n auto maxLevel = entity_.cur_max_level();\n auto entry_point = entity_.entry_point();\n spin_lock_.unlock();\n\n if (ailego_unlikely(entry_point == kInvalidNodeId)) {\n return 0;\n }\n\n dist_t dist = ctx->dist_calculator().dist(entry_point);\n for (level_t cur_level = maxLevel; cur_level >= 1; --cur_level) {\n select_entry_point(cur_level, &entry_point, &dist, ctx);\n }\n\n auto &topk_heap = ctx->topk_heap();\n topk_heap.clear();\n search_neighbors(0, &entry_point, &dist, topk_heap, ctx);\n\n if (ctx->group_by_search()) {\n expand_neighbors_by_group(topk_heap, ctx);\n }\n\n return 0;\n}\n\n//! select_entry_point on hnsw level, ef = 1\nvoid HnswAlgorithm::select_entry_point(level_t level, node_id_t *entry_point,\n dist_t *dist, HnswContext *ctx) const {\n auto &entity = ctx->get_entity();\n HnswDistCalculator &dc = ctx->dist_calculator();\n while (true) {\n const Neighbors neighbors = entity.get_neighbors(level, *entry_point);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n uint32_t size = neighbors.size();\n if (size == 0) {\n break;\n }\n\n std::vector neighbor_vec_blocks;\n int ret = entity.get_vector(&neighbors[0], size, neighbor_vec_blocks);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_vector())++;\n }\n if (ailego_unlikely(ret != 0)) {\n break;\n }\n\n bool find_closer = false;\n\n std::vector dists(size);\n std::vector neighbor_vecs(size);\n for (uint32_t i = 0; i < size; ++i) {\n neighbor_vecs[i] = neighbor_vec_blocks[i].data();\n }\n\n dc.batch_dist(neighbor_vecs.data(), size, dists.data());\n\n for (uint32_t i = 0; i < size; ++i) {\n dist_t cur_dist = dists[i];\n\n if (cur_dist < *dist) {\n *entry_point = neighbors[i];\n *dist = cur_dist;\n find_closer = true;\n }\n }\n\n if (!find_closer) {\n break;\n }\n }\n\n return;\n}\n\nvoid HnswAlgorithm::add_neighbors(node_id_t id, level_t level,\n TopkHeap &topk_heap, HnswContext *ctx) {\n if (ailego_unlikely(topk_heap.size() == 0)) {\n return;\n }\n\n HnswDistCalculator &dc = ctx->dist_calculator();\n\n update_neighbors(dc, id, level, topk_heap);\n\n // reverse update neighbors\n for (size_t i = 0; i < topk_heap.size(); ++i) {\n reverse_update_neighbors(dc, topk_heap[i].first, level, id,\n topk_heap[i].second, ctx->update_heap());\n }\n\n return;\n}\n\nvoid HnswAlgorithm::search_neighbors(level_t level, node_id_t *entry_point,\n dist_t *dist, TopkHeap &topk,\n HnswContext *ctx) const {\n const auto &entity = ctx->get_entity();\n HnswDistCalculator &dc = ctx->dist_calculator();\n VisitFilter &visit = ctx->visit_filter();\n CandidateHeap &candidates = ctx->candidates();\n std::function filter = [](node_id_t) { return false; };\n if (ctx->filter().is_valid()) {\n filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); };\n }\n\n candidates.clear();\n visit.clear();\n visit.set_visited(*entry_point);\n if (!filter(*entry_point)) {\n topk.emplace(*entry_point, *dist);\n }\n\n candidates.emplace(*entry_point, *dist);\n while (!candidates.empty() && !ctx->reach_scan_limit()) {\n auto top = candidates.begin();\n node_id_t main_node = top->first;\n dist_t main_dist = top->second;\n\n if (topk.full() && main_dist > topk[0].second) {\n break;\n }\n\n candidates.pop();\n const Neighbors neighbors = entity.get_neighbors(level, main_node);\n ailego_prefetch(neighbors.data);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n\n std::vector neighbor_ids(neighbors.size());\n uint32_t size = 0;\n for (uint32_t i = 0; i < neighbors.size(); ++i) {\n node_id_t node = neighbors[i];\n if (visit.visited(node)) {\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_visit_dup_cnt())++;\n }\n continue;\n }\n visit.set_visited(node);\n neighbor_ids[size++] = node;\n }\n if (size == 0) {\n continue;\n }\n\n std::vector neighbor_vec_blocks;\n int ret = entity.get_vector(neighbor_ids.data(), size, neighbor_vec_blocks);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_vector())++;\n }\n if (ailego_unlikely(ret != 0)) {\n break;\n }\n\n // do prefetch\n static constexpr node_id_t BATCH_SIZE = 12;\n static constexpr node_id_t PREFETCH_STEP = 2;\n for (uint32_t i = 0; i < std::min(BATCH_SIZE * PREFETCH_STEP, size); ++i) {\n ailego_prefetch(neighbor_vec_blocks[i].data());\n }\n // done\n\n std::vector dists(size);\n std::vector neighbor_vecs(size);\n\n for (uint32_t i = 0; i < size; ++i) {\n neighbor_vecs[i] = neighbor_vec_blocks[i].data();\n }\n\n dc.batch_dist(neighbor_vecs.data(), size, dists.data());\n\n for (uint32_t i = 0; i < size; ++i) {\n node_id_t node = neighbor_ids[i];\n dist_t cur_dist = dists[i];\n\n if ((!topk.full()) || cur_dist < topk[0].second) {\n candidates.emplace(node, cur_dist);\n // update entry_point for next level scan\n if (cur_dist < *dist) {\n *entry_point = node;\n *dist = cur_dist;\n }\n if (!filter(node)) {\n topk.emplace(node, cur_dist);\n }\n } // end if\n } // end for\n } // while\n\n return;\n}\n\nvoid HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk,\n HnswContext *ctx) const {\n if (!ctx->group_by().is_valid()) {\n return;\n }\n\n const auto &entity = ctx->get_entity();\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity.get_key(id));\n };\n\n // devide into groups\n std::map &group_topk_heaps = ctx->group_topk_heaps();\n for (uint32_t i = 0; i < topk.size(); ++i) {\n node_id_t id = topk[i].first;\n auto score = topk[i].second;\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = group_topk_heaps[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, score);\n }\n\n // stage 2, expand to reach group num as possible\n if (group_topk_heaps.size() < ctx->group_num()) {\n VisitFilter &visit = ctx->visit_filter();\n CandidateHeap &candidates = ctx->candidates();\n HnswDistCalculator &dc = ctx->dist_calculator();\n\n std::function filter = [](node_id_t) { return false; };\n if (ctx->filter().is_valid()) {\n filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); };\n }\n\n // refill to get enough groups\n candidates.clear();\n visit.clear();\n for (uint32_t i = 0; i < topk.size(); ++i) {\n node_id_t id = topk[i].first;\n float score = topk[i].second;\n\n visit.set_visited(id);\n candidates.emplace_back(id, score);\n }\n\n // do expand\n while (!candidates.empty() && !ctx->reach_scan_limit()) {\n auto top = candidates.begin();\n node_id_t main_node = top->first;\n\n candidates.pop();\n const Neighbors neighbors = entity.get_neighbors(0, main_node);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n\n std::vector neighbor_ids(neighbors.size());\n uint32_t size = 0;\n for (uint32_t i = 0; i < neighbors.size(); ++i) {\n node_id_t node = neighbors[i];\n if (visit.visited(node)) {\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_visit_dup_cnt())++;\n }\n continue;\n }\n visit.set_visited(node);\n neighbor_ids[size++] = node;\n }\n if (size == 0) {\n continue;\n }\n\n std::vector neighbor_vec_blocks;\n int ret =\n entity.get_vector(neighbor_ids.data(), size, neighbor_vec_blocks);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_vector())++;\n }\n if (ailego_unlikely(ret != 0)) {\n break;\n }\n\n static constexpr node_id_t PREFETCH_STEP = 2;\n for (uint32_t i = 0; i < size; ++i) {\n node_id_t node = neighbor_ids[i];\n node_id_t prefetch_id = i + PREFETCH_STEP;\n if (prefetch_id < size) {\n ailego_prefetch(neighbor_vec_blocks[prefetch_id].data());\n }\n dist_t cur_dist = dc.dist(neighbor_vec_blocks[i].data());\n\n if (!filter(node)) {\n std::string group_id = group_by(node);\n\n auto &topk_heap = group_topk_heaps[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(node, cur_dist);\n\n if (group_topk_heaps.size() >= ctx->group_num()) {\n break;\n }\n }\n\n candidates.emplace(node, cur_dist);\n } // end for\n } // end while\n } // end if\n}\n\nvoid HnswAlgorithm::update_neighbors(HnswDistCalculator &dc, node_id_t id,\n level_t level, TopkHeap &topk_heap) {\n topk_heap.sort();\n\n uint32_t max_neighbor_cnt = entity_.neighbor_cnt(level);\n if (topk_heap.size() <= static_cast(entity_.prune_cnt())) {\n if (topk_heap.size() <= static_cast(max_neighbor_cnt)) {\n entity_.update_neighbors(level, id, topk_heap);\n return;\n }\n }\n\n uint32_t cur_size = 0;\n for (size_t i = 0; i < topk_heap.size(); ++i) {\n node_id_t cur_node = topk_heap[i].first;\n dist_t cur_node_dist = topk_heap[i].second;\n bool good = true;\n for (uint32_t j = 0; j < cur_size; ++j) {\n dist_t tmp_dist = dc.dist(cur_node, topk_heap[j].first);\n if (tmp_dist <= cur_node_dist) {\n good = false;\n break;\n }\n }\n\n if (good) {\n topk_heap[cur_size].first = cur_node;\n topk_heap[cur_size].second = cur_node_dist;\n cur_size++;\n if (cur_size >= max_neighbor_cnt) {\n break;\n }\n }\n }\n\n // when after-prune neighbor count is too seldom,\n // we use this strategy to make-up enough edges\n // not only just make-up out-degrees\n // we also make-up enough in-degrees\n uint32_t min_neighbors = entity_.min_neighbor_cnt();\n for (size_t k = cur_size; cur_size < min_neighbors && k < topk_heap.size();\n ++k) {\n bool exist = false;\n for (size_t j = 0; j < cur_size; ++j) {\n if (topk_heap[j].first == topk_heap[k].first) {\n exist = true;\n break;\n }\n }\n if (!exist) {\n topk_heap[cur_size].first = topk_heap[k].first;\n topk_heap[cur_size].second = topk_heap[k].second;\n cur_size++;\n }\n }\n\n topk_heap.resize(cur_size);\n entity_.update_neighbors(level, id, topk_heap);\n\n return;\n}\n\nvoid HnswAlgorithm::reverse_update_neighbors(HnswDistCalculator &dc,\n node_id_t id, level_t level,\n node_id_t link_id, dist_t dist,\n TopkHeap &update_heap) {\n const size_t max_neighbor_cnt = entity_.neighbor_cnt(level);\n\n uint32_t lock_idx = id & kLockMask;\n lock_pool_[lock_idx].lock();\n const Neighbors neighbors = entity_.get_neighbors(level, id);\n size_t size = neighbors.size();\n ailego_assert_with(size <= max_neighbor_cnt, \"invalid neighbor size\");\n if (size < max_neighbor_cnt) {\n entity_.add_neighbor(level, id, size, link_id);\n lock_pool_[lock_idx].unlock();\n return;\n }\n\n update_heap.emplace(link_id, dist);\n\n for (size_t i = 0; i < size; ++i) {\n node_id_t node = neighbors[i];\n dist_t cur_dist = dc.dist(id, node);\n update_heap.emplace(node, cur_dist);\n }\n\n //! TODO: optimize prune\n //! prune edges\n update_heap.sort();\n size_t cur_size = 0;\n for (size_t i = 0; i < update_heap.size(); ++i) {\n node_id_t cur_node = update_heap[i].first;\n dist_t cur_node_dist = update_heap[i].second;\n bool good = true;\n for (size_t j = 0; j < cur_size; ++j) {\n dist_t tmp_dist = dc.dist(cur_node, update_heap[j].first);\n if (tmp_dist <= cur_node_dist) {\n good = false;\n break;\n }\n }\n\n if (good) {\n update_heap[cur_size].first = cur_node;\n update_heap[cur_size].second = cur_node_dist;\n cur_size++;\n if (cur_size >= max_neighbor_cnt) {\n break;\n }\n }\n }\n\n update_heap.resize(cur_size);\n entity_.update_neighbors(level, id, update_heap);\n\n lock_pool_[lock_idx].unlock();\n\n update_heap.clear();\n\n return;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_algorithm.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"hnsw_context.h\"\n#include \"hnsw_dist_calculator.h\"\n#include \"hnsw_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! hnsw graph algorithm implement\nclass HnswAlgorithm {\n public:\n typedef std::unique_ptr UPointer;\n\n public:\n //! Constructor\n explicit HnswAlgorithm(HnswEntity &entity);\n\n //! Destructor\n ~HnswAlgorithm() = default;\n\n //! Cleanup HnswAlgorithm\n int cleanup();\n\n //! Add a node to hnsw graph\n //! @id: the node unique id\n //! @level: a node will be add to graph in each level [0, level]\n //! return 0 on success, or errCode in failure\n int add_node(node_id_t id, level_t level, HnswContext *ctx);\n\n //! do knn search in graph\n //! return 0 on success, or errCode in failure. results saved in ctx\n int search(HnswContext *ctx) const;\n\n //! Initiate HnswAlgorithm\n int init() {\n level_probas_.clear();\n double level_mult =\n 1 / std::log(static_cast(entity_.scaling_factor()));\n for (int level = 0;; level++) {\n // refers faiss get_random_level alg\n double proba =\n std::exp(-level / level_mult) * (1 - std::exp(-1 / level_mult));\n if (proba < 1e-9) {\n break;\n }\n level_probas_.push_back(proba);\n }\n\n return 0;\n }\n\n //! Generate a random level\n //! return graph level\n uint32_t get_random_level() const {\n // gen rand float (0, 1)\n double f = mt_() / static_cast(mt_.max());\n for (size_t level = 0; level < level_probas_.size(); level++) {\n if (f < level_probas_[level]) {\n return level;\n }\n f -= level_probas_[level];\n }\n return level_probas_.size() - 1;\n }\n\n private:\n //! Select in upper layer to get entry point for next layer search\n void select_entry_point(level_t level, node_id_t *entry_point, dist_t *dist,\n HnswContext *ctx) const;\n\n //! update node id neighbors from topkHeap, and reverse link is also updated\n void add_neighbors(node_id_t id, level_t level, TopkHeap &topk_heap,\n HnswContext *ctx);\n\n //! Given a node id and level, search the nearest neighbors in graph\n //! Note: the nearest neighbors result keeps in topk, and entry_point and\n //! dist will be updated to current level nearest node id and distance\n void search_neighbors(level_t level, node_id_t *entry_point, dist_t *dist,\n TopkHeap &topk, HnswContext *ctx) const;\n\n //! Update the node's neighbors\n void update_neighbors(HnswDistCalculator &dc, node_id_t id, level_t level,\n TopkHeap &topk_heap);\n\n //! Checking linkId could be id's new neighbor, and add as neighbor if true\n //! @dc distance calculator\n //! @updateHeap temporary heap in updating neighbors\n void reverse_update_neighbors(HnswDistCalculator &dc, node_id_t id,\n level_t level, node_id_t link_id, dist_t dist,\n TopkHeap &update_heap);\n\n //! expand neighbors until group nums are reached\n void expand_neighbors_by_group(TopkHeap &topk, HnswContext *ctx) const;\n\n private:\n HnswAlgorithm(const HnswAlgorithm &) = delete;\n HnswAlgorithm &operator=(const HnswAlgorithm &) = delete;\n\n private:\n static constexpr uint32_t kLockCnt{1U << 8};\n static constexpr uint32_t kLockMask{kLockCnt - 1U};\n\n HnswEntity &entity_;\n mutable std::mt19937 mt_{};\n std::vector level_probas_{};\n\n mutable ailego::SpinMutex spin_lock_{}; // global spin lock\n std::mutex mutex_{}; // global mutex\n // TODO: spin lock?\n std::vector lock_pool_{};\n};\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw/hnsw_builder.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_builder.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"hnsw_algorithm.h\"\n#include \"hnsw_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswBuilder::HnswBuilder() = default;\n\nint HnswBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) {\n LOG_INFO(\"Begin HnswBuilder::init\");\n\n meta_ = meta;\n auto params_copy = params;\n meta_.set_builder(\"HnswBuilder\", HnswEntity::kRevision,\n std::move(params_copy));\n\n size_t memory_quota = 0UL;\n params.get(PARAM_HNSW_BUILDER_MEMORY_QUOTA, &memory_quota);\n params.get(PARAM_HNSW_BUILDER_THREAD_COUNT, &thread_cnt_);\n params.get(PARAM_HNSW_BUILDER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_);\n params.get(PARAM_HNSW_BUILDER_EFCONSTRUCTION, &ef_construction_);\n params.get(PARAM_HNSW_BUILDER_CHECK_INTERVAL_SECS, &check_interval_secs_);\n\n params.get(PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT, &upper_max_neighbor_cnt_);\n float multiplier = HnswEntity::kDefaultL0MaxNeighborCntMultiplier;\n params.get(PARAM_HNSW_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER, &multiplier);\n l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_;\n scaling_factor_ = upper_max_neighbor_cnt_;\n params.get(PARAM_HNSW_BUILDER_SCALING_FACTOR, &scaling_factor_);\n\n multiplier = HnswEntity::kDefaultNeighborPruneMultiplier;\n params.get(PARAM_HNSW_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier);\n size_t prune_cnt = multiplier * upper_max_neighbor_cnt_;\n\n if (ef_construction_ == 0) {\n ef_construction_ = HnswEntity::kDefaultEfConstruction;\n }\n if (upper_max_neighbor_cnt_ == 0) {\n upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt;\n }\n if (upper_max_neighbor_cnt_ > kMaxNeighborCnt) {\n LOG_ERROR(\"[%s] must be in range (0,%d]\",\n PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT.c_str(), kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) {\n LOG_ERROR(\"[%s]-[%d] must be <= [%s]-[%d]\",\n PARAM_HNSW_BUILDER_MIN_NEIGHBOR_COUNT.c_str(), min_neighbor_cnt_,\n PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT.c_str(),\n upper_max_neighbor_cnt_);\n return IndexError_InvalidArgument;\n }\n if (l0_max_neighbor_cnt_ == 0) {\n l0_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt;\n }\n if (l0_max_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) {\n LOG_ERROR(\"L0MaxNeighborCnt must be in range (0,%d)\",\n HnswEntity::kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (scaling_factor_ == 0U) {\n scaling_factor_ = HnswEntity::kDefaultScalingFactor;\n }\n if (scaling_factor_ < 5 || scaling_factor_ > 1000) {\n LOG_ERROR(\"[%s] must be in range [5,1000]\",\n PARAM_HNSW_BUILDER_SCALING_FACTOR.c_str());\n return IndexError_InvalidArgument;\n }\n if (thread_cnt_ == 0) {\n thread_cnt_ = std::thread::hardware_concurrency();\n }\n if (thread_cnt_ > std::thread::hardware_concurrency()) {\n LOG_WARN(\"[%s] greater than cpu cores %u\",\n PARAM_HNSW_BUILDER_THREAD_COUNT.c_str(),\n std::thread::hardware_concurrency());\n }\n if (prune_cnt == 0UL) {\n prune_cnt = upper_max_neighbor_cnt_;\n }\n\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"CreateMetric failed, name: %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n int ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"IndexMetric init failed, ret=%d\", ret);\n return ret;\n }\n\n entity_.set_vector_size(meta_.element_size());\n\n entity_.set_ef_construction(ef_construction_);\n entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_);\n entity_.set_min_neighbor_cnt(min_neighbor_cnt_);\n entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_);\n entity_.set_scaling_factor(scaling_factor_);\n entity_.set_memory_quota(memory_quota);\n entity_.set_prune_cnt(prune_cnt);\n\n ret = entity_.init();\n if (ret != 0) {\n return ret;\n }\n\n alg_ = HnswAlgorithm::UPointer(new HnswAlgorithm(entity_));\n\n ret = alg_->init();\n if (ret != 0) {\n return ret;\n }\n\n state_ = BUILD_STATE_INITED;\n LOG_INFO(\n \"End HnswBuilder::init, params: vectorSize=%u efConstruction=%u \"\n \"l0NeighborCnt=%u upperNeighborCnt=%u scalingFactor=%u \"\n \"memoryQuota=%zu neighborPruneCnt=%zu metricName=%s \",\n meta_.element_size(), ef_construction_, l0_max_neighbor_cnt_,\n upper_max_neighbor_cnt_, scaling_factor_, memory_quota, prune_cnt,\n meta_.metric_name().c_str());\n\n return 0;\n}\n\nint HnswBuilder::cleanup(void) {\n LOG_INFO(\"Begin HnswBuilder::cleanup\");\n\n l0_max_neighbor_cnt_ = HnswEntity::kDefaultL0MaxNeighborCnt;\n min_neighbor_cnt_ = 0;\n upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt;\n ef_construction_ = HnswEntity::kDefaultEfConstruction;\n scaling_factor_ = HnswEntity::kDefaultScalingFactor;\n check_interval_secs_ = kDefaultLogIntervalSecs;\n errcode_ = 0;\n error_ = false;\n entity_.cleanup();\n alg_->cleanup();\n meta_.clear();\n metric_.reset();\n stats_.clear_attributes();\n stats_.set_trained_count(0UL);\n stats_.set_built_count(0UL);\n stats_.set_dumped_count(0UL);\n stats_.set_discarded_count(0UL);\n stats_.set_trained_costtime(0UL);\n stats_.set_built_costtime(0UL);\n stats_.set_dumped_costtime(0UL);\n state_ = BUILD_STATE_INIT;\n\n LOG_INFO(\"End HnswBuilder::cleanup\");\n\n return 0;\n}\n\nint HnswBuilder::train(IndexThreads::Pointer, IndexHolder::Pointer holder) {\n if (state_ != BUILD_STATE_INITED) {\n LOG_ERROR(\"Init the builder before HnswBuilder::train\");\n return IndexError_NoReady;\n }\n\n if (!holder) {\n LOG_ERROR(\"Input holder is nullptr while training index\");\n return IndexError_InvalidArgument;\n }\n if (!holder->is_matched(meta_)) {\n LOG_ERROR(\"Input holder doesn't match index meta while training index\");\n return IndexError_Mismatch;\n }\n LOG_INFO(\"Begin HnswBuilder::train\");\n size_t trained_cost_time = 0;\n size_t trained_count = 0;\n\n if (metric_->support_train()) {\n auto start_time = ailego::Monotime::MilliSeconds();\n auto iter = holder->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Create iterator for holder failed\");\n return IndexError_Runtime;\n }\n while (iter->is_valid()) {\n int ret = metric_->train(iter->data(), meta_.dimension());\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw build measure train failed, ret=%d\", ret);\n return ret;\n }\n iter->next();\n ++trained_count;\n }\n trained_cost_time = ailego::Monotime::MilliSeconds() - start_time;\n }\n stats_.set_trained_count(trained_count);\n stats_.set_trained_costtime(trained_cost_time);\n state_ = BUILD_STATE_TRAINED;\n\n LOG_INFO(\"End HnswBuilder::train\");\n\n return 0;\n}\n\nint HnswBuilder::train(const IndexTrainer::Pointer & /*trainer*/) {\n if (state_ != BUILD_STATE_INITED) {\n LOG_ERROR(\"Init the builder before HnswBuilder::train\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"Begin HnswBuilder::train by trainer\");\n\n stats_.set_trained_count(0UL);\n stats_.set_trained_costtime(0UL);\n state_ = BUILD_STATE_TRAINED;\n\n LOG_INFO(\"End HnswBuilder::train by trainer\");\n\n return 0;\n}\n\nint HnswBuilder::build(IndexThreads::Pointer threads,\n IndexHolder::Pointer holder) {\n if (state_ != BUILD_STATE_TRAINED) {\n LOG_ERROR(\"Train the index before HnswBuilder::build\");\n return IndexError_NoReady;\n }\n\n if (!holder) {\n LOG_ERROR(\"Input holder is nullptr while building index\");\n return IndexError_InvalidArgument;\n }\n if (!holder->is_matched(meta_)) {\n LOG_ERROR(\"Input holder doesn't match index meta while building index\");\n return IndexError_Mismatch;\n }\n if (!threads) {\n threads = std::make_shared(thread_cnt_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n\n auto start_time = ailego::Monotime::MilliSeconds();\n LOG_INFO(\"Begin HnswBuilder::build\");\n\n if (holder->count() != static_cast(-1)) {\n LOG_DEBUG(\"HnswBuilder holder documents count %lu\", holder->count());\n int ret = entity_.reserve_space(holder->count());\n if (ret != 0) {\n LOG_ERROR(\"HnswBuilde reserver space failed\");\n return ret;\n }\n }\n auto iter = holder->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Create iterator for holder failed\");\n return IndexError_Runtime;\n }\n int ret;\n error_ = false;\n while (iter->is_valid()) {\n level_t level = alg_->get_random_level();\n node_id_t id;\n\n const void *vec = iter->data();\n ret = entity_.add_vector(level, iter->key(), vec, &id);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n iter->next();\n }\n // Holder is not needed, cleanup it.\n holder.reset();\n\n LOG_INFO(\"Finished save vector, start build graph...\");\n\n auto task_group = threads->make_group();\n if (!task_group) {\n LOG_ERROR(\"Failed to create task group\");\n return IndexError_Runtime;\n }\n\n std::atomic finished{0};\n for (size_t i = 0; i < threads->count(); ++i) {\n task_group->submit(ailego::Closure ::New(this, &HnswBuilder::do_build, i,\n threads->count(), &finished));\n }\n\n while (!task_group->is_finished()) {\n std::unique_lock lk(mutex_);\n cond_.wait_until(lk, std::chrono::system_clock::now() +\n std::chrono::seconds(check_interval_secs_));\n if (error_.load(std::memory_order_acquire)) {\n LOG_ERROR(\"Failed to build index while waiting finish\");\n return errcode_;\n }\n LOG_INFO(\"Built cnt %u, finished percent %.3f%%\", finished.load(),\n finished.load() * 100.0f / entity_.doc_cnt());\n }\n if (error_.load(std::memory_order_acquire)) {\n LOG_ERROR(\"Failed to build index while waiting finish\");\n return errcode_;\n }\n task_group->wait_finish();\n\n stats_.set_built_count(finished.load());\n stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time);\n state_ = BUILD_STATE_BUILT;\n\n LOG_INFO(\"End HnswBuilder::build\");\n return 0;\n}\n\nvoid HnswBuilder::do_build(node_id_t idx, size_t step_size,\n std::atomic *finished) {\n AILEGO_DEFER([&]() {\n std::lock_guard latch(mutex_);\n cond_.notify_one();\n });\n HnswContext *ctx = new (std::nothrow)\n HnswContext(meta_.dimension(), metric_,\n std::shared_ptr(&entity_, [](HnswEntity *) {}));\n if (ailego_unlikely(ctx == nullptr)) {\n if (!error_.exchange(true)) {\n LOG_ERROR(\"Failed to create context\");\n errcode_ = IndexError_NoMemory;\n }\n return;\n }\n HnswContext::Pointer auto_ptr(ctx);\n ctx->set_max_scan_num(entity_.doc_cnt());\n int ret = ctx->init(HnswContext::kBuilderContext);\n if (ret != 0) {\n if (!error_.exchange(true)) {\n LOG_ERROR(\"Failed to init context\");\n errcode_ = IndexError_Runtime;\n }\n return;\n }\n\n IndexQueryMeta qmeta(meta_.data_type(), meta_.dimension());\n for (node_id_t id = idx; id < entity_.doc_cnt(); id += step_size) {\n ctx->reset_query(entity_.get_vector(id));\n ret = alg_->add_node(id, entity_.get_level(id), ctx);\n if (ailego_unlikely(ret != 0)) {\n if (!error_.exchange(true)) {\n LOG_ERROR(\"Hnsw graph add node failed\");\n errcode_ = ret;\n }\n return;\n }\n ctx->clear();\n (*finished)++;\n }\n}\n\nint HnswBuilder::dump(const IndexDumper::Pointer &dumper) {\n if (state_ != BUILD_STATE_BUILT) {\n LOG_INFO(\"Build the index before HnswBuilder::dump\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"Begin HnswBuilder::dump\");\n\n meta_.set_searcher(\"HnswSearcher\", HnswEntity::kRevision, ailego::Params());\n auto start_time = ailego::Monotime::MilliSeconds();\n\n int ret = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return ret;\n }\n\n ret = entity_.dump(dumper);\n if (ret != 0) {\n LOG_ERROR(\"HnswBuilder dump index failed\");\n return ret;\n }\n\n stats_.set_dumped_count(entity_.doc_cnt());\n stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time);\n\n LOG_INFO(\"EndHnswBuilder::dump\");\n return 0;\n}\n\nINDEX_FACTORY_REGISTER_BUILDER(HnswBuilder);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_builder.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"hnsw_algorithm.h\"\n#include \"hnsw_builder_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswBuilder : public IndexBuilder {\n public:\n //! Constructor\n HnswBuilder();\n\n //! Initialize the builder\n virtual int init(const IndexMeta &meta,\n const ailego::Params ¶ms) override;\n\n //! Cleanup the builder\n virtual int cleanup(void) override;\n\n //! Train the data\n virtual int train(IndexThreads::Pointer,\n IndexHolder::Pointer holder) override;\n\n //! Train the data\n virtual int train(const IndexTrainer::Pointer &trainer) override;\n\n\n //! Build the index\n virtual int build(IndexThreads::Pointer threads,\n IndexHolder::Pointer holder) override;\n\n //! Dump index into storage\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n virtual const Stats &stats(void) const override {\n return stats_;\n }\n\n private:\n void do_build(node_id_t idx, size_t step_size,\n std::atomic *finished);\n\n constexpr static uint32_t kDefaultLogIntervalSecs = 15U;\n constexpr static uint32_t kMaxNeighborCnt = 65535;\n\n private:\n enum BUILD_STATE {\n BUILD_STATE_INIT = 0,\n BUILD_STATE_INITED = 1,\n BUILD_STATE_TRAINED = 2,\n BUILD_STATE_BUILT = 3\n };\n\n HnswBuilderEntity entity_{};\n HnswAlgorithm::UPointer alg_; // impl graph algorithm\n uint32_t thread_cnt_{0};\n uint32_t min_neighbor_cnt_{0};\n uint32_t upper_max_neighbor_cnt_{HnswEntity::kDefaultUpperMaxNeighborCnt};\n uint32_t l0_max_neighbor_cnt_{HnswEntity::kDefaultL0MaxNeighborCnt};\n uint32_t ef_construction_{HnswEntity::kDefaultEfConstruction};\n uint32_t scaling_factor_{HnswEntity::kDefaultScalingFactor};\n uint32_t check_interval_secs_{kDefaultLogIntervalSecs};\n\n int errcode_{0};\n std::atomic_bool error_{false};\n IndexMeta meta_{};\n IndexMetric::Pointer metric_{};\n std::mutex mutex_{};\n std::condition_variable cond_{};\n Stats stats_{};\n\n BUILD_STATE state_{BUILD_STATE_INIT};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_builder_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_builder_entity.h\"\n#include \n#include \n#include \"utility/sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswBuilderEntity::HnswBuilderEntity() {\n update_ep_and_level(kInvalidNodeId, 0U);\n}\n\nint HnswBuilderEntity::cleanup() {\n memory_quota_ = 0UL;\n neighbors_size_ = 0U;\n upper_neighbors_size_ = 0U;\n padding_size_ = 0U;\n vectors_buffer_.clear();\n keys_buffer_.clear();\n neighbors_buffer_.clear();\n upper_neighbors_buffer_.clear();\n neighbors_index_.clear();\n\n vectors_buffer_.shrink_to_fit();\n keys_buffer_.shrink_to_fit();\n neighbors_buffer_.shrink_to_fit();\n upper_neighbors_buffer_.shrink_to_fit();\n neighbors_index_.shrink_to_fit();\n\n this->HnswEntity::cleanup();\n\n return 0;\n}\n\nint HnswBuilderEntity::init() {\n size_t size = vector_size();\n\n //! aligned size to 32\n set_node_size(AlignSize(size));\n //! if node size is aligned to 1k, the build performance will downgrade\n if (node_size() % 1024 == 0) {\n set_node_size(AlignSize(node_size() + 1));\n }\n\n padding_size_ = node_size() - size;\n\n neighbors_size_ = neighbors_size();\n upper_neighbors_size_ = upper_neighbors_size();\n\n return 0;\n}\n\nint HnswBuilderEntity::reserve_space(size_t docs) {\n if (memory_quota_ > 0 && (node_size() * docs + neighbors_size_ * docs +\n sizeof(NeighborIndex) * docs >\n memory_quota_)) {\n return IndexError_NoMemory;\n }\n\n vectors_buffer_.reserve(node_size() * docs);\n keys_buffer_.reserve(sizeof(key_t) * docs);\n neighbors_buffer_.reserve(neighbors_size_ * docs);\n neighbors_index_.reserve(docs);\n\n return 0;\n}\n\nint HnswBuilderEntity::add_vector(level_t level, key_t key, const void *vec,\n node_id_t *id) {\n if (memory_quota_ > 0 &&\n (vectors_buffer_.capacity() + keys_buffer_.capacity() +\n neighbors_buffer_.capacity() + upper_neighbors_buffer_.capacity() +\n neighbors_index_.capacity() * sizeof(NeighborIndex)) > memory_quota_) {\n LOG_ERROR(\"Add vector failed, used memory exceed quota, cur_doc=%u\",\n doc_cnt());\n return IndexError_NoMemory;\n }\n\n vectors_buffer_.append(reinterpret_cast(vec), vector_size());\n vectors_buffer_.append(padding_size_, '\\0');\n keys_buffer_.append(reinterpret_cast(&key), sizeof(key));\n\n // init level 0 neighbors\n neighbors_buffer_.append(neighbors_size_, '\\0');\n\n neighbors_index_.emplace_back(upper_neighbors_buffer_.size(), level);\n\n // init upper layer neighbors\n for (level_t cur_level = 1; cur_level <= level; ++cur_level) {\n upper_neighbors_buffer_.append(upper_neighbors_size_, '\\0');\n }\n\n *id = (*mutable_doc_cnt())++;\n\n return 0;\n}\n\nkey_t HnswBuilderEntity::get_key(node_id_t id) const {\n return *(reinterpret_cast(keys_buffer_.data() +\n id * sizeof(key_t)));\n}\n\nconst void *HnswBuilderEntity::get_vector(node_id_t id) const {\n return vectors_buffer_.data() + id * node_size();\n}\n\nint HnswBuilderEntity::get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const {\n const void *vec = get_vector(id);\n block.reset((void *)vec);\n return 0;\n}\n\nint HnswBuilderEntity::get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const {\n for (uint32_t i = 0; i < count; ++i) {\n vecs[i] = vectors_buffer_.data() + ids[i] * node_size();\n }\n\n return 0;\n}\n\nint HnswBuilderEntity::get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const {\n const void *vecs[count];\n get_vector(ids, count, vecs);\n for (uint32_t i = 0; i < count; ++i) {\n vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i]));\n }\n return 0;\n}\n\nconst Neighbors HnswBuilderEntity::get_neighbors(level_t level,\n node_id_t id) const {\n const NeighborsHeader *hd = get_neighbor_header(level, id);\n return {hd->neighbor_cnt, hd->neighbors};\n}\n\nint HnswBuilderEntity::update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) {\n NeighborsHeader *hd =\n const_cast(get_neighbor_header(level, id));\n for (size_t i = 0; i < neighbors.size(); ++i) {\n hd->neighbors[i] = neighbors[i].first;\n }\n hd->neighbor_cnt = neighbors.size();\n\n // std::cout << \"id: \" << id << \", neighbour, id: \";\n // for (size_t i = 0; i < neighbors.size(); ++i) {\n // if (i == neighbors.size()-1)\n // std::cout << neighbors[i].first << \", score:\" << neighbors[i].second <<\n // std::endl;\n // else\n // std::cout << neighbors[i].first << \", score:\" << neighbors[i].second <<\n // \", id: \";\n // }\n\n return 0;\n}\n\nvoid HnswBuilderEntity::add_neighbor(level_t level, node_id_t id,\n uint32_t /*size*/, node_id_t neighbor_id) {\n NeighborsHeader *hd =\n const_cast(get_neighbor_header(level, id));\n hd->neighbors[hd->neighbor_cnt++] = neighbor_id;\n\n return;\n}\n\nint HnswBuilderEntity::dump(const IndexDumper::Pointer &dumper) {\n key_t *keys =\n reinterpret_cast(const_cast(keys_buffer_.data()));\n auto ret =\n dump_segments(dumper, keys, [&](node_id_t id) { return get_level(id); });\n if (ailego_unlikely(ret < 0)) {\n return ret;\n }\n\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_builder_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"hnsw_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswBuilderEntity : public HnswEntity {\n public:\n //! Add vector and key to hnsw entity, and local id will be saved to id\n virtual int add_vector(level_t level, key_t key, const void *vec,\n node_id_t *id) override;\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const override;\n\n //! Get vector feature data by key\n virtual const void *get_vector(node_id_t id) const override;\n\n //! Batch get vectors feature data by keys\n virtual int get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const override;\n\n virtual int get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n virtual int get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const override;\n\n //! Get the node id's neighbors on graph level\n const NeighborsHeader *get_neighbor_header(level_t level,\n node_id_t id) const {\n if (level == 0) {\n return reinterpret_cast(\n neighbors_buffer_.data() + neighbors_size_ * id);\n } else {\n size_t offset = neighbors_index_[id].offset;\n return reinterpret_cast(\n upper_neighbors_buffer_.data() + offset +\n (level - 1) * upper_neighbors_size_);\n }\n }\n\n //! Get the node id's neighbors on graph level\n virtual const Neighbors get_neighbors(level_t level,\n node_id_t id) const override;\n\n //! Replace node id in level's neighbors\n virtual int update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) override;\n\n //! add a neighbor to id in graph level\n virtual void add_neighbor(level_t level, node_id_t id, uint32_t size,\n node_id_t neighbor_id) override;\n\n //! Dump the hnsw graph to dumper\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Cleanup the entity\n virtual int cleanup(void) override;\n\n public:\n //! Constructor\n HnswBuilderEntity();\n\n //! Get the node graph level by id\n level_t get_level(node_id_t id) const {\n return neighbors_index_[id].level;\n }\n\n //! Init builerEntity\n int init();\n\n //! reserve buffer space for documents\n //! @param docs number of documents\n int reserve_space(size_t docs);\n\n //! Set memory quota params\n inline void set_memory_quota(size_t memory_quota) {\n memory_quota_ = memory_quota;\n }\n\n //! Get neighbors size\n inline size_t neighbors_size() const {\n return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t);\n }\n\n //! Get upper neighbors size\n inline size_t upper_neighbors_size() const {\n return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t);\n }\n\n public:\n HnswBuilderEntity(const HnswBuilderEntity &) = delete;\n HnswBuilderEntity &operator=(const HnswBuilderEntity &) = delete;\n\n private:\n friend class HnswSearcherEntity;\n //! class internal used only\n struct NeighborIndex {\n NeighborIndex(size_t off, level_t l) : offset(off), level(l) {}\n uint64_t offset : 48;\n uint64_t level : 16;\n };\n\n std::string vectors_buffer_{}; // aligned vectors\n std::string keys_buffer_{}; // aligned vectors\n std::string neighbors_buffer_{}; // level 0 neighbors buffer\n std::string upper_neighbors_buffer_{}; // upper layer neighbors buffer\n\n std::string sparse_data_buffer_{}; // aligned spase data buffer\n size_t sparse_data_offset_{0}; //\n\n // upper layer offset + level in upper_neighbors_buffer_\n std::vector neighbors_index_{};\n size_t memory_quota_{0UL};\n size_t neighbors_size_{0U}; // level 0 neighbors size\n size_t upper_neighbors_size_{0U}; // level 0 neighbors size\n size_t padding_size_{}; // padding size for each vector element\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_chunk.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_chunk.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nint ChunkBroker::init_storage(size_t chunk_size) {\n chunk_meta_.clear();\n chunk_meta_.chunk_size = chunk_size;\n chunk_meta_.create_time = ailego::Realtime::Seconds();\n stats_.set_create_time(chunk_meta_.create_time);\n chunk_meta_.update_time = ailego::Realtime::Seconds();\n stats_.set_update_time(chunk_meta_.update_time);\n\n //! alloc meta chunk\n size_t size = sizeof(HnswChunkMeta);\n size = (size + page_mask_) & (~page_mask_);\n const std::string segment_id =\n make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId);\n int ret = stg_->append(segment_id, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Storage append segment failed for %s\", IndexError::What(ret));\n return ret;\n }\n chunk_meta_segment_ = get_chunk(CHUNK_TYPE_META, kDefaultChunkSeqId);\n if (ailego_unlikely(!chunk_meta_segment_)) {\n LOG_ERROR(\"Get meta segment failed\");\n return IndexError_Runtime;\n }\n\n //! update meta info and write to storage\n chunk_meta_.chunk_cnts[CHUNK_TYPE_META] += 1;\n chunk_meta_.total_size += size;\n (*stats_.mutable_index_size()) += size;\n size = chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta));\n if (ailego_unlikely(size != sizeof(HnswChunkMeta))) {\n LOG_ERROR(\"Storage write data failed, wsize=%zu\", size);\n return IndexError_WriteData;\n }\n\n return 0;\n}\n\nint ChunkBroker::load_storage(size_t chunk_size) {\n IndexStorage::MemoryBlock data_block;\n size_t size = chunk_meta_segment_->read(0UL, data_block,\n chunk_meta_segment_->data_size());\n if (size != sizeof(HnswChunkMeta)) {\n LOG_ERROR(\"Invalid hnsw meta chunk, read size=%zu chunk size=%zu\", size,\n chunk_meta_segment_->data_size());\n return IndexError_InvalidFormat;\n }\n std::memcpy(&chunk_meta_, data_block.data(), size);\n if (chunk_meta_.chunk_size != chunk_size) {\n LOG_ERROR(\n \"Params hnsw chunk size=%zu mismatch from previous %zu \"\n \"in index\",\n chunk_size, (size_t)chunk_meta_.chunk_size);\n return IndexError_Mismatch;\n }\n\n *stats_.mutable_check_point() = stg_->check_point();\n stats_.set_revision_id(chunk_meta_.revision_id);\n stats_.set_update_time(chunk_meta_.update_time);\n stats_.set_create_time(chunk_meta_.create_time);\n\n char create_time[32];\n char update_time[32];\n ailego::Realtime::Gmtime(chunk_meta_.create_time, \"%Y-%m-%d %H:%M:%S\",\n create_time, sizeof(create_time));\n ailego::Realtime::Gmtime(chunk_meta_.update_time, \"%Y-%m-%d %H:%M:%S\",\n update_time, sizeof(update_time));\n LOG_DEBUG(\n \"Load index, indexSize=%zu chunkSize=%zu nodeChunks=%zu \"\n \"upperNeighborChunks=%zu revisionId=%zu \"\n \"createTime=%s updateTime=%s\",\n (size_t)chunk_meta_.total_size, (size_t)chunk_meta_.chunk_size,\n (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_NODE],\n (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_UPPER_NEIGHBOR],\n (size_t)chunk_meta_.revision_id, create_time, update_time);\n\n return 0;\n}\n\nint ChunkBroker::open(IndexStorage::Pointer stg, size_t max_index_size,\n size_t chunk_size, bool check_crc) {\n if (ailego_unlikely(stg_)) {\n LOG_ERROR(\"An storage instance is already opened\");\n return IndexError_Duplicate;\n }\n stg_ = std::move(stg);\n if (stg_->isHugePage()) {\n page_mask_ = ailego::MemoryHelper::HugePageSize() - 1;\n } else {\n page_mask_ = ailego::MemoryHelper::PageSize() - 1;\n }\n check_crc_ = check_crc;\n max_chunks_size_ = max_index_size;\n dirty_ = false;\n\n const std::string segment_id =\n make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId);\n chunk_meta_segment_ = stg_->get(segment_id);\n if (!chunk_meta_segment_) {\n LOG_DEBUG(\"Create new index\");\n return init_storage(chunk_size);\n }\n\n return load_storage(chunk_size);\n}\n\nint ChunkBroker::close(void) {\n flush(0UL);\n\n stg_.reset();\n check_crc_ = false;\n dirty_ = false;\n\n return 0;\n}\n\nint ChunkBroker::flush(uint64_t checkpoint) {\n ailego_assert_with(chunk_meta_segment_, \"invalid meta segment\");\n\n chunk_meta_.update_time = ailego::Realtime::Seconds();\n stats_.set_update_time(chunk_meta_.update_time);\n\n size_t size =\n chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta));\n if (ailego_unlikely(size != sizeof(HnswChunkMeta))) {\n LOG_ERROR(\"Storage write data failed, wsize=%zu\", size);\n }\n\n stg_->refresh(checkpoint);\n int ret = stg_->flush();\n if (ret == 0) {\n (*stats_.mutable_check_point()) = checkpoint;\n } else {\n LOG_ERROR(\"Storage flush failed for %s\", IndexError::What(ret));\n }\n return ret;\n}\n\nstd::pair ChunkBroker::alloc_chunk(int type,\n uint64_t seq_id,\n size_t size) {\n ailego_assert_with(type < CHUNK_TYPE_MAX, \"chunk type overflow\");\n\n Chunk::Pointer chunk;\n if (ailego_unlikely(!stg_)) {\n LOG_ERROR(\"Init storage first\");\n return std::make_pair(IndexError_Uninitialized, chunk);\n }\n\n //! check exist a empty chunk with the same name\n chunk = get_chunk(type, seq_id);\n if (chunk) {\n if (ailego_unlikely(chunk->capacity() == size &&\n chunk->data_size() == 0UL)) {\n LOG_ERROR(\"Exist invalid chunk size %zu, expect size %zu\",\n chunk->capacity(), size);\n chunk.reset();\n return std::make_pair(IndexError_Runtime, chunk);\n }\n return std::make_pair(0, chunk);\n }\n //! align to page size\n size = (size + page_mask_) & (~page_mask_);\n if (ailego_unlikely(chunk_meta_.total_size + size >= max_chunks_size_)) {\n LOG_ERROR(\"No space to new a chunk, curIndexSize=%zu allocSize=%zu\",\n (size_t)chunk_meta_.total_size, size);\n return std::make_pair(IndexError_IndexFull, chunk);\n }\n\n std::string segment_id = make_segment_id(type, seq_id);\n int ret = stg_->append(segment_id, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Storage append segment failed for %s\", IndexError::What(ret));\n return std::make_pair(ret, chunk);\n }\n chunk_meta_.chunk_cnts[type] += 1;\n chunk_meta_.total_size += size;\n (*stats_.mutable_index_size()) += size;\n\n size = chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta));\n if (ailego_unlikely(size != sizeof(HnswChunkMeta))) {\n LOG_ERROR(\"Storage append segment failed, wsize=%zu\", size);\n }\n\n chunk = get_chunk(type, seq_id);\n return std::make_pair(chunk ? 0 : IndexError_NoMemory, chunk);\n}\n\nChunk::Pointer ChunkBroker::get_chunk(int type, uint64_t seq_id) const {\n std::string segment_id = make_segment_id(type, seq_id);\n return stg_->get(segment_id);\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_chunk.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nusing Chunk = IndexStorage::Segment;\n\nclass ChunkBroker {\n public:\n typedef std::shared_ptr Pointer;\n\n enum CHUNK_TYPE {\n CHUNK_TYPE_HEADER = 1,\n CHUNK_TYPE_META = 2,\n CHUNK_TYPE_NODE = 3,\n CHUNK_TYPE_UPPER_NEIGHBOR = 4,\n CHUNK_TYPE_NEIGHBOR_INDEX = 5,\n CHUNK_TYPE_SPARSE_NODE = 6,\n CHUNK_TYPE_MAX = 8\n };\n static constexpr size_t kDefaultChunkSeqId = 0UL;\n\n ChunkBroker(IndexStreamer::Stats &stats) : stats_(stats) {}\n\n //! Open storage\n int open(IndexStorage::Pointer stg, size_t max_index_size, size_t chunk_size,\n bool check_crc);\n\n int close(void);\n\n int flush(uint64_t checkpoint);\n\n //! alloc a new chunk with size, not thread-safe\n std::pair alloc_chunk(int type, uint64_t seq_id,\n size_t size);\n\n //! alloc a new chunk with chunk size\n inline std::pair alloc_chunk(int type, uint64_t seq_id) {\n return alloc_chunk(type, seq_id, chunk_meta_.chunk_size);\n }\n\n Chunk::Pointer get_chunk(int type, uint64_t seq_id) const;\n\n inline size_t get_chunk_cnt(int type) const {\n ailego_assert_with(type < CHUNK_TYPE_MAX, \"chunk type overflow\");\n return chunk_meta_.chunk_cnts[type];\n }\n\n inline bool dirty(void) const {\n return dirty_;\n }\n\n inline void mark_dirty(void) {\n if (!dirty_) {\n dirty_ = true;\n chunk_meta_.revision_id += 1;\n stats_.set_revision_id(chunk_meta_.revision_id);\n }\n }\n\n const IndexStorage::Pointer storage(void) const {\n return stg_;\n }\n\n private:\n ChunkBroker(const ChunkBroker &) = delete;\n ChunkBroker &operator=(const ChunkBroker &) = delete;\n\n struct HnswChunkMeta {\n HnswChunkMeta(void) {\n memset(this, 0, sizeof(HnswChunkMeta));\n }\n void clear() {\n memset(this, 0, sizeof(HnswChunkMeta));\n }\n\n uint64_t chunk_cnts[CHUNK_TYPE_MAX];\n uint64_t chunk_size; // size of per chunk\n uint64_t total_size; // total size of allocated chunk\n uint64_t revision_id; // index revision\n uint64_t create_time;\n uint64_t update_time;\n uint64_t reserved[3];\n };\n\n static_assert(sizeof(HnswChunkMeta) % 32 == 0,\n \"HnswChunkMeta must be aligned with 32 bytes\");\n\n //! Init the storage after open an empty index\n int init_storage(size_t chunk_size);\n\n //! Load index from storage\n int load_storage(size_t chunk_size);\n\n static inline const std::string make_segment_id(int type, uint64_t seq_id) {\n return \"HnswT\" + ailego::StringHelper::ToString(type) + \"S\" +\n ailego::StringHelper::ToString(seq_id);\n }\n\n private:\n IndexStreamer::Stats &stats_;\n HnswChunkMeta chunk_meta_{};\n size_t page_mask_{0UL};\n size_t max_chunks_size_{0UL};\n IndexStorage::Pointer stg_{};\n IndexStorage::Segment::Pointer chunk_meta_segment_{};\n bool check_crc_{false};\n bool dirty_{false}; // set as true if index is modified , the flag\n // will not be cleared even if flushed\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_context.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_context.h\"\n#include \n#include \"hnsw_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswContext::HnswContext(size_t dimension, const IndexMetric::Pointer &metric,\n const HnswEntity::Pointer &entity)\n : IndexContext(metric),\n entity_(entity),\n dc_(entity_.get(), metric, dimension) {}\n\nHnswContext::HnswContext(const IndexMetric::Pointer &metric,\n const HnswEntity::Pointer &entity)\n : IndexContext(metric), entity_(entity), dc_(entity_.get(), metric) {}\n\nHnswContext::~HnswContext() {\n visit_filter_.destroy();\n}\n\nint HnswContext::init(ContextType type) {\n int ret;\n uint32_t doc_cnt;\n\n type_ = type;\n\n switch (type) {\n case kBuilderContext:\n ret = visit_filter_.init(VisitFilter::ByteMap, entity_->doc_cnt(),\n max_scan_num_, negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n candidates_.limit(max_scan_num_);\n update_heap_.limit(entity_->l0_neighbor_cnt() + 1);\n break;\n\n case kSearcherContext:\n ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_,\n negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n candidates_.limit(max_scan_num_);\n break;\n\n case kStreamerContext:\n // maxScanNum is unknown if inited from streamer, so the docCnt may\n // change. we need to compute maxScanNum by scan ratio, and preserve\n // max_doc_cnt space from visit filter\n doc_cnt = entity_->doc_cnt();\n max_scan_num_ = compute_max_scan_num(doc_cnt);\n reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt);\n ret = visit_filter_.init(filter_mode_, reserve_max_doc_cnt_,\n max_scan_num_, negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n\n update_heap_.limit(entity_->l0_neighbor_cnt() + 1);\n candidates_.limit(max_scan_num_);\n\n check_need_adjuct_ctx();\n break;\n\n default:\n LOG_ERROR(\"Init context failed\");\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nint HnswContext::update(const ailego::Params ¶ms) {\n auto update_visit_filter_param = [&]() {\n bool need_update = false;\n std::string p;\n switch (type_) {\n case kSearcherContext:\n p = PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_ENABLE;\n break;\n case kStreamerContext:\n p = PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_ENABLE;\n break;\n }\n\n if (params.has(p)) {\n bool bf_enabled;\n params.get(p, &bf_enabled);\n if (bf_enabled ^ (filter_mode_ == VisitFilter::BloomFilter)) {\n need_update = true;\n filter_mode_ =\n bf_enabled ? VisitFilter::BloomFilter : VisitFilter::ByteMap;\n }\n }\n\n float prob = negative_probability_;\n p.clear();\n switch (type_) {\n case kSearcherContext:\n p = PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB;\n break;\n case kStreamerContext:\n p = PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB;\n break;\n }\n params.get(p, &prob);\n if (filter_mode_ == VisitFilter::BloomFilter &&\n std::abs(prob - negative_probability_) > 1e-6) {\n need_update = true;\n }\n if (need_update) {\n visit_filter_.destroy();\n int max_doc_cnt = 0;\n if (type_ == kSearcherContext) {\n max_doc_cnt = entity_->doc_cnt();\n } else {\n max_doc_cnt = reserve_max_doc_cnt_;\n }\n int ret = visit_filter_.init(filter_mode_, max_doc_cnt, max_scan_num_,\n negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n }\n return 0;\n };\n\n switch (type_) {\n case kSearcherContext:\n if (params.has(PARAM_HNSW_SEARCHER_EF)) {\n params.get(PARAM_HNSW_SEARCHER_EF, &ef_);\n topk_heap_.limit(std::max(topk_, ef_));\n }\n\n if (params.has(PARAM_HNSW_SEARCHER_MAX_SCAN_RATIO)) {\n params.get(PARAM_HNSW_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_);\n max_scan_num_ =\n static_cast(max_scan_ratio_ * entity_->doc_cnt());\n max_scan_num_ = std::max(10000U, max_scan_num_);\n }\n\n if (params.has(PARAM_HNSW_SEARCHER_BRUTE_FORCE_THRESHOLD)) {\n params.get(PARAM_HNSW_SEARCHER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n }\n\n return update_visit_filter_param();\n\n case kStreamerContext:\n if (params.has(PARAM_HNSW_STREAMER_EF)) {\n params.get(PARAM_HNSW_STREAMER_EF, &ef_);\n topk_heap_.limit(std::max(topk_, ef_));\n }\n params.get(PARAM_HNSW_STREAMER_EF, &ef_);\n params.get(PARAM_HNSW_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_);\n params.get(PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_);\n params.get(PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_);\n if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) {\n LOG_ERROR(\"[%s] must be in range (0.0f,1.0f]\",\n PARAM_HNSW_STREAMER_MAX_SCAN_RATIO.c_str());\n return IndexError_InvalidArgument;\n }\n if (max_scan_limit_ < min_scan_limit_) {\n LOG_ERROR(\"[%s] must be >= [%s]\",\n PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT.c_str(),\n PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (params.has(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD)) {\n params.get(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n }\n\n return update_visit_filter_param();\n\n default:\n LOG_ERROR(\"update context failed, type=%u\", type_);\n return IndexError_Runtime;\n }\n}\n\nint HnswContext::update_context(ContextType type, const IndexMeta &meta,\n const IndexMetric::Pointer &metric,\n const HnswEntity::Pointer &entity,\n uint32_t magic_num) {\n uint32_t doc_cnt;\n\n if (ailego_unlikely(type != type_)) {\n LOG_ERROR(\n \"HnswContext doesn't support shared by different type, \"\n \"src=%u dst=%u\",\n type_, type);\n return IndexError_Unsupported;\n }\n\n magic_ = kInvalidMgic;\n\n // TODO: support change filter mode?\n switch (type) {\n case kBuilderContext:\n LOG_ERROR(\"BuildContext doesn't support update\");\n return IndexError_NotImplemented;\n\n case kSearcherContext:\n if (!visit_filter_.reset(entity->doc_cnt(), max_scan_num_)) {\n LOG_ERROR(\"Reset filter failed, mode %d\", visit_filter_.get_mode());\n return IndexError_Runtime;\n }\n\n candidates_.limit(max_scan_num_);\n topk_heap_.limit(std::max(topk_, ef_));\n break;\n\n case kStreamerContext:\n doc_cnt = entity->doc_cnt();\n max_scan_num_ = compute_max_scan_num(doc_cnt);\n reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt);\n if (!visit_filter_.reset(reserve_max_doc_cnt_, max_scan_num_)) {\n LOG_ERROR(\"Reset filter failed, mode %d\", visit_filter_.get_mode());\n return IndexError_Runtime;\n }\n\n update_heap_.limit(entity->l0_neighbor_cnt() + 1);\n candidates_.limit(max_scan_num_);\n topk_heap_.limit(std::max(topk_, ef_));\n break;\n\n default:\n LOG_ERROR(\"update context failed\");\n return IndexError_Runtime;\n }\n\n entity_ = entity;\n dc_.update(entity_.get(), metric, meta.dimension());\n magic_ = magic_num;\n level_topks_.clear();\n\n return 0;\n}\n\nvoid HnswContext::fill_random_to_topk_full(void) {\n static std::mt19937 mt(\n std::chrono::system_clock::now().time_since_epoch().count());\n std::uniform_int_distribution dt(0, entity_->doc_cnt() - 1);\n std::function gen;\n node_id_t seqid;\n std::function myfilter = [](node_id_t) { return false; };\n if (this->filter().is_valid()) {\n myfilter = [&](node_id_t id) {\n return this->filter()(entity_->get_key(id));\n };\n }\n\n if (topk_heap_.limit() < entity_->doc_cnt() / 2) {\n gen = [&](void) { return dt(mt); };\n } else {\n // If topk limit is big value, gen sequential id from an random initial\n seqid = dt(mt);\n gen = [&](void) {\n seqid = seqid == (entity_->doc_cnt() - 1) ? 0 : (seqid + 1);\n return seqid;\n };\n }\n\n for (size_t i = 0; !topk_heap_.full() && i < entity_->doc_cnt(); ++i) {\n const auto id = gen();\n if (!visit_filter_.visited(id) && !myfilter(id)) {\n visit_filter_.set_visited(id);\n topk_heap_.emplace(id, dc_.dist(id));\n }\n }\n return;\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw/hnsw_context.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"utility/sparse_utility.h\"\n#include \"utility/visit_filter.h\"\n#include \"hnsw_dist_calculator.h\"\n#include \"hnsw_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswContext : public IndexContext {\n public:\n //! Index Context Pointer\n typedef std::unique_ptr Pointer;\n\n enum ContextType {\n kUnknownContext = 0,\n kSearcherContext = 1,\n kBuilderContext = 2,\n kStreamerContext = 3\n };\n\n //! Construct\n HnswContext(size_t dimension, const IndexMetric::Pointer &metric,\n const HnswEntity::Pointer &entity);\n\n //! Construct\n HnswContext(const IndexMetric::Pointer &metric,\n const HnswEntity::Pointer &entity);\n\n //! Destructor\n virtual ~HnswContext();\n\n public:\n //! Set topk of search result\n virtual void set_topk(uint32_t val) override {\n topk_ = val;\n topk_heap_.limit(std::max(val, ef_));\n }\n\n //! Retrieve search result\n virtual const IndexDocumentList &result(void) const override {\n return results_[0];\n }\n\n //! Retrieve search result\n virtual const IndexDocumentList &result(size_t idx) const override {\n return results_[idx];\n }\n\n //! Retrieve result object for output\n virtual IndexDocumentList *mutable_result(size_t idx) override {\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n return &results_[idx];\n }\n\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(void) const override {\n return group_results_[0];\n }\n\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(\n size_t idx) const override {\n return group_results_[idx];\n }\n\n virtual uint32_t magic(void) const override {\n return magic_;\n }\n\n //! Set mode of debug\n virtual void set_debug_mode(bool enable) override {\n debug_mode_ = enable;\n }\n\n //! Retrieve mode of debug\n virtual bool debug_mode(void) const override {\n return this->debugging();\n }\n\n //! Retrieve string of debug\n virtual std::string debug_string(void) const override {\n char buf[4096];\n size_t size = snprintf(\n buf, sizeof(buf),\n \"scan_cnt=%zu,get_vector_cnt=%u,get_neighbors_cnt=%u,dup_node=%u\",\n get_scan_num(), stats_get_vector_cnt_, stats_get_neighbors_cnt_,\n stats_visit_dup_cnt_);\n return std::string(buf, size);\n }\n\n //! Update the parameters of context\n virtual int update(const ailego::Params ¶ms) override;\n\n public:\n //! Init context\n int init(ContextType type);\n\n //! Update context, the context may be shared by different searcher/streamer\n int update_context(ContextType type, const IndexMeta &meta,\n const IndexMetric::Pointer &metric,\n const HnswEntity::Pointer &entity, uint32_t magic_num);\n\n inline const HnswEntity &get_entity() const {\n return *entity_;\n }\n\n inline void resize_results(size_t size) {\n if (group_by_search()) {\n group_results_.resize(size);\n } else {\n results_.resize(size);\n }\n }\n\n inline void topk_to_result() {\n return topk_to_result(0);\n }\n\n //! Construct result from topk heap, result will be normalized\n inline void topk_to_result(uint32_t idx) {\n if (group_by_search()) {\n topk_to_group_result(idx);\n } else {\n topk_to_single_result(idx);\n }\n }\n\n inline void recal_topk_dist() {\n TopkHeap heap(topk_heap_);\n topk_heap_.clear();\n\n for (size_t i = 0; i < heap.size(); ++i) {\n node_id_t id = heap[i].first;\n dist_t dist = dc_.dist(id);\n topk_heap_.emplace_back(id, dist);\n }\n }\n\n inline void topk_to_single_result(uint32_t idx) {\n if (force_padding_topk_ && !topk_heap_.full() &&\n topk_heap_.size() < entity_->doc_cnt()) {\n this->fill_random_to_topk_full();\n }\n if (ailego_unlikely(topk_heap_.size() == 0)) {\n return;\n }\n\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n int size = std::min(topk_, static_cast(topk_heap_.size()));\n topk_heap_.sort();\n results_[idx].clear();\n\n for (int i = 0; i < size; ++i) {\n auto score = topk_heap_[i].second;\n if (score > this->threshold()) {\n break;\n }\n\n node_id_t id = topk_heap_[i].first;\n if (fetch_vector_) {\n results_[idx].emplace_back(entity_->get_key(id), score, id,\n entity_->get_vector(id));\n } else {\n results_[idx].emplace_back(entity_->get_key(id), score, id);\n }\n }\n\n return;\n }\n\n //! Construct result from topk heap, result will be normalized\n inline void topk_to_group_result(uint32_t idx) {\n ailego_assert_with(idx < group_results_.size(), \"invalid idx\");\n\n group_results_[idx].clear();\n\n std::vector> group_topk_list;\n std::vector> best_score_in_groups;\n for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end();\n itr++) {\n const std::string &group_id = (*itr).first;\n auto &heap = (*itr).second;\n heap.sort();\n\n if (heap.size() > 0) {\n float best_score = heap[0].second;\n best_score_in_groups.push_back(std::make_pair(group_id, best_score));\n }\n }\n\n std::sort(best_score_in_groups.begin(), best_score_in_groups.end(),\n [](const std::pair &a,\n const std::pair &b) -> int {\n return a.second < b.second;\n });\n\n // truncate to group num\n for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size();\n ++i) {\n const std::string &group_id = best_score_in_groups[i].first;\n\n group_topk_list.emplace_back(\n std::make_pair(group_id, group_topk_heaps_[group_id]));\n }\n\n group_results_[idx].resize(group_topk_list.size());\n\n for (uint32_t i = 0; i < group_topk_list.size(); ++i) {\n const std::string &group_id = group_topk_list[i].first;\n group_results_[idx][i].set_group_id(group_id);\n\n uint32_t size = std::min(\n group_topk_, static_cast(group_topk_list[i].second.size()));\n\n for (uint32_t j = 0; j < size; ++j) {\n auto score = group_topk_list[i].second[j].second;\n if (score > this->threshold()) {\n break;\n }\n\n node_id_t id = group_topk_list[i].second[j].first;\n\n if (fetch_vector_) {\n group_results_[idx][i].mutable_docs()->emplace_back(\n entity_->get_key(id), score, id, entity_->get_vector(id));\n } else {\n group_results_[idx][i].mutable_docs()->emplace_back(\n entity_->get_key(id), score, id);\n }\n }\n }\n }\n\n inline void reset_query(const void *query) {\n if (auto query_preprocess_func = index_metric_->get_query_preprocess_func();\n query_preprocess_func != nullptr) {\n size_t dim = dc_.dimension();\n preprocess_buffer_.resize(dim);\n memcpy(preprocess_buffer_.data(), query, dim);\n query_preprocess_func(preprocess_buffer_.data(), dim);\n query = preprocess_buffer_.data();\n }\n\n dc_.reset_query(query);\n dc_.clear_compare_cnt();\n }\n\n inline HnswDistCalculator &dist_calculator() {\n return dc_;\n }\n\n inline TopkHeap &topk_heap() {\n return topk_heap_;\n }\n\n inline TopkHeap &update_heap() {\n return update_heap_;\n }\n\n inline VisitFilter &visit_filter() {\n return visit_filter_;\n }\n\n inline CandidateHeap &candidates() {\n return candidates_;\n }\n\n inline void set_max_scan_num(uint32_t max_scan_num) {\n max_scan_num_ = max_scan_num;\n }\n\n inline void set_max_scan_limit(uint32_t max_scan_limit) {\n max_scan_limit_ = max_scan_limit;\n }\n\n inline void set_min_scan_limit(uint32_t min_scan_limit) {\n min_scan_limit_ = min_scan_limit;\n }\n\n inline void set_ef(uint32_t v) {\n ef_ = v;\n }\n\n inline void set_filter_mode(uint32_t v) {\n filter_mode_ = v;\n }\n\n inline void set_filter_negative_probability(float v) {\n negative_probability_ = v;\n }\n\n inline void set_max_scan_ratio(float v) {\n max_scan_ratio_ = v;\n }\n\n virtual void set_magic(uint32_t v) {\n magic_ = v;\n }\n\n virtual void set_force_padding_topk(bool v) {\n force_padding_topk_ = v;\n }\n\n void set_bruteforce_threshold(uint32_t v) override {\n bruteforce_threshold_ = v;\n }\n\n inline uint32_t get_bruteforce_threshold() const {\n return bruteforce_threshold_;\n }\n\n void set_fetch_vector(bool v) override {\n fetch_vector_ = v;\n }\n\n bool fetch_vector() const override {\n return fetch_vector_;\n }\n\n //! Reset context\n void reset(void) override {\n this->clear();\n set_filter(nullptr);\n reset_threshold();\n set_fetch_vector(false);\n set_group_params(0, 0);\n reset_group_by();\n }\n\n inline std::map &group_topk_heaps() {\n return group_topk_heaps_;\n }\n\n inline TopkHeap &level_topk(int level) {\n if (ailego_unlikely(level_topks_.size() <= static_cast(level))) {\n int cur_level = level_topks_.size();\n level_topks_.resize(level + 1);\n for (; cur_level <= level; ++cur_level) {\n size_t heap_size = std::max(entity_->neighbor_cnt(cur_level),\n entity_->ef_construction());\n level_topks_[cur_level].clear();\n level_topks_[cur_level].limit(heap_size);\n }\n }\n\n return level_topks_[level];\n }\n\n inline void check_need_adjuct_ctx(void) {\n check_need_adjuct_ctx(entity_->doc_cnt());\n }\n\n inline size_t compute_reserve_cnt(uint32_t cur_doc) const {\n if (cur_doc > kMaxReserveDocCnt) {\n return kMaxReserveDocCnt;\n } else if (cur_doc < kMinReserveDocCnt) {\n return kMinReserveDocCnt;\n }\n return cur_doc;\n }\n\n //! candidates heap and visitfilter need to resize as doc cnt growing up\n inline void check_need_adjuct_ctx(uint32_t doc_cnt) {\n if (ailego_unlikely(doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_)) {\n while (doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_) {\n reserve_max_doc_cnt_ =\n reserve_max_doc_cnt_ + compute_reserve_cnt(reserve_max_doc_cnt_);\n }\n uint32_t max_scan_cnt = compute_max_scan_num(reserve_max_doc_cnt_);\n max_scan_num_ = max_scan_cnt;\n visit_filter_.reset(reserve_max_doc_cnt_, max_scan_cnt);\n candidates_.clear();\n candidates_.limit(max_scan_num_);\n }\n }\n\n inline uint32_t compute_max_scan_num(uint32_t max_doc_cnt) const {\n uint32_t max_scan = max_doc_cnt * max_scan_ratio_;\n if (max_scan < min_scan_limit_) {\n max_scan = min_scan_limit_;\n } else if (max_scan > max_scan_limit_) {\n max_scan = max_scan_limit_;\n }\n return max_scan;\n }\n\n inline size_t get_scan_num() const {\n return dc_.compare_cnt();\n }\n\n inline uint64_t reach_scan_limit() const {\n return dc_.compare_cnt() >= max_scan_num_;\n }\n\n inline bool error() const {\n return dc_.error();\n }\n\n inline void clear() {\n dc_.clear();\n if (ailego_unlikely(this->debugging())) {\n stats_get_neighbors_cnt_ = 0u;\n stats_get_vector_cnt_ = 0u;\n stats_visit_dup_cnt_ = 0u;\n }\n // do not clear results_ for the next query will need it\n for (auto &it : results_) {\n it.clear();\n }\n for (auto &it : group_results_) {\n it.clear();\n }\n }\n\n uint32_t *mutable_stats_get_neighbors() {\n return &stats_get_neighbors_cnt_;\n }\n\n uint32_t *mutable_stats_get_vector() {\n return &stats_get_vector_cnt_;\n }\n\n uint32_t *mutable_stats_visit_dup_cnt() {\n return &stats_visit_dup_cnt_;\n }\n\n inline bool debugging(void) const {\n return debug_mode_;\n }\n\n inline void update_dist_caculator_distance(\n const IndexMetric::MatrixDistance &distance,\n const IndexMetric::MatrixBatchDistance &batch_distance) {\n dc_.update_distance(distance, batch_distance);\n }\n\n //! Get topk\n inline uint32_t topk() const override {\n return topk_;\n }\n\n //! Get group topk\n inline uint32_t group_topk() const {\n return group_topk_;\n }\n\n //! Get group num\n inline uint32_t group_num() const {\n return group_num_;\n }\n\n //! Get if group by search\n inline bool group_by_search() {\n return group_num_ > 0;\n }\n\n //! Set group params\n void set_group_params(uint32_t group_num, uint32_t group_topk) override {\n group_num_ = group_num;\n group_topk_ = group_topk;\n\n topk_ = group_topk_ * group_num_;\n\n topk_heap_.limit(std::max(topk_, ef_));\n\n group_topk_heaps_.clear();\n }\n\n private:\n // Filling random nodes if topk not full\n void fill_random_to_topk_full(void);\n\n constexpr static uint32_t kTriggerReserveCnt = 4096UL;\n constexpr static uint32_t kMinReserveDocCnt = 4096UL;\n constexpr static uint32_t kMaxReserveDocCnt = 128 * 1024UL;\n constexpr static uint32_t kInvalidMgic = -1U;\n\n private:\n HnswEntity::Pointer entity_;\n HnswDistCalculator dc_;\n IndexMetric::Pointer metric_;\n\n bool debug_mode_{false};\n bool force_padding_topk_{false};\n uint32_t max_scan_num_{0};\n uint32_t max_scan_limit_{0};\n uint32_t min_scan_limit_{0};\n uint32_t reserve_max_doc_cnt_{kMinReserveDocCnt};\n uint32_t topk_{0};\n uint32_t group_topk_{0};\n uint32_t filter_mode_{VisitFilter::ByteMap};\n float negative_probability_{HnswEntity::kDefaultBFNegativeProbability};\n uint32_t ef_{HnswEntity::kDefaultEf};\n float max_scan_ratio_{HnswEntity::kDefaultScanRatio};\n uint32_t magic_{0U};\n std::vector results_{};\n std::vector group_results_{};\n TopkHeap topk_heap_{};\n TopkHeap update_heap_{};\n std::vector level_topks_{};\n CandidateHeap candidates_{};\n VisitFilter visit_filter_{};\n uint32_t bruteforce_threshold_{};\n bool fetch_vector_{false};\n\n uint32_t group_num_{0};\n std::map group_topk_heaps_{};\n\n uint32_t type_{kUnknownContext};\n //! debug stats info\n uint32_t stats_get_neighbors_cnt_{0u};\n uint32_t stats_get_vector_cnt_{0u};\n uint32_t stats_visit_dup_cnt_{0u};\n std::string preprocess_buffer_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_dist_calculator.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"hnsw_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswDistCalculator {\n public:\n typedef std::shared_ptr Pointer;\n\n public:\n enum DistType {\n DIST_NONE = 0,\n DIST_DENSE = 1,\n DIST_HYBRID = 2,\n DIST_SPARSE = 3\n };\n\n public:\n //! Constructor\n HnswDistCalculator(const HnswEntity *entity,\n const IndexMetric::Pointer &metric, uint32_t dim)\n : entity_(entity),\n distance_(metric->distance()),\n batch_distance_(metric->batch_distance()),\n query_(nullptr),\n dim_(dim),\n compare_cnt_(0) {}\n\n //! Constructor\n HnswDistCalculator(const HnswEntity *entity,\n const IndexMetric::Pointer &metric, uint32_t dim,\n const void *query)\n : entity_(entity),\n distance_(metric->distance()),\n batch_distance_(metric->batch_distance()),\n query_(query),\n dim_(dim),\n compare_cnt_(0) {}\n\n //! Constructor\n HnswDistCalculator(const HnswEntity *entity,\n const IndexMetric::Pointer &metric)\n : entity_(entity),\n distance_(metric->distance()),\n batch_distance_(metric->batch_distance()),\n query_(nullptr),\n dim_(0),\n compare_cnt_(0) {}\n\n void update(const HnswEntity *entity, const IndexMetric::Pointer &metric) {\n entity_ = entity;\n distance_ = metric->distance();\n batch_distance_ = metric->batch_distance();\n }\n\n void update(const HnswEntity *entity, const IndexMetric::Pointer &metric,\n uint32_t dim) {\n entity_ = entity;\n distance_ = metric->distance();\n batch_distance_ = metric->batch_distance();\n dim_ = dim;\n }\n\n inline void update_distance(\n const IndexMetric::MatrixDistance &distance,\n const IndexMetric::MatrixBatchDistance &batch_distance) {\n distance_ = distance;\n batch_distance_ = batch_distance;\n }\n\n //! Reset query vector data\n inline void reset_query(const void *query) {\n error_ = false;\n query_ = query;\n }\n\n //! Returns distance\n inline dist_t dist(const void *vec_lhs, const void *vec_rhs) {\n if (ailego_unlikely(vec_lhs == nullptr || vec_rhs == nullptr)) {\n LOG_ERROR(\"Nullptr of dense vector\");\n error_ = true;\n return 0.0f;\n }\n\n float score{0.0f};\n\n distance_(vec_lhs, vec_rhs, dim_, &score);\n\n return score;\n }\n\n //! Returns distance between query and vec.\n inline dist_t dist(const void *vec) {\n compare_cnt_++;\n\n return dist(vec, query_);\n }\n\n //! Return distance between query and node id.\n inline dist_t dist(node_id_t id) {\n compare_cnt_++;\n\n const void *feat = entity_->get_vector(id);\n if (ailego_unlikely(feat == nullptr)) {\n LOG_ERROR(\"Get nullptr vector, id=%u\", id);\n error_ = true;\n return 0.0f;\n }\n\n return dist(feat, query_);\n }\n\n //! Return dist node lhs between node rhs\n inline dist_t dist(node_id_t lhs, node_id_t rhs) {\n compare_cnt_++;\n\n const void *feat = entity_->get_vector(lhs);\n const void *query = entity_->get_vector(rhs);\n if (ailego_unlikely(feat == nullptr || query == nullptr)) {\n LOG_ERROR(\"Get nullptr vector\");\n error_ = true;\n return 0.0f;\n }\n\n return dist(feat, query);\n }\n\n dist_t operator()(const void *vec) {\n return dist(vec);\n }\n\n dist_t operator()(id_t i) {\n return dist(i);\n }\n\n dist_t operator()(id_t lhs, id_t rhs) {\n return dist(lhs, rhs);\n }\n\n void batch_dist(const void **vecs, size_t num, dist_t *distances) {\n compare_cnt_++;\n\n batch_distance_(vecs, query_, num, dim_, distances);\n }\n\n inline dist_t batch_dist(node_id_t id) {\n compare_cnt_++;\n\n const void *feat = entity_->get_vector(id);\n if (ailego_unlikely(feat == nullptr)) {\n LOG_ERROR(\"Get nullptr vector, id=%u\", id);\n error_ = true;\n return 0.0f;\n }\n dist_t score = 0;\n batch_distance_(&feat, query_, 1, dim_, &score);\n\n return score;\n }\n\n inline void clear() {\n compare_cnt_ = 0;\n error_ = false;\n }\n\n inline void clear_compare_cnt() {\n compare_cnt_ = 0;\n }\n\n inline bool error() const {\n return error_;\n }\n\n //! Get distances compute times\n inline uint32_t compare_cnt() const {\n return compare_cnt_;\n }\n\n inline uint32_t dimension() const {\n return dim_;\n }\n\n private:\n HnswDistCalculator(const HnswDistCalculator &) = delete;\n HnswDistCalculator &operator=(const HnswDistCalculator &) = delete;\n\n private:\n const HnswEntity *entity_;\n\n IndexMetric::MatrixDistance distance_;\n IndexMetric::MatrixBatchDistance batch_distance_;\n\n const void *query_;\n uint32_t dim_;\n\n uint32_t compare_cnt_; // record distance compute times\n // uint32_t compare_cnt_batch_; // record batch distance compute time\n bool error_{false};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_entity.h\"\n#include \n#include \"utility/sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nconst std::string HnswEntity::kGraphHeaderSegmentId = \"graph.header\";\nconst std::string HnswEntity::kGraphFeaturesSegmentId = \"graph.features\";\nconst std::string HnswEntity::kGraphKeysSegmentId = \"graph.keys\";\nconst std::string HnswEntity::kGraphNeighborsSegmentId = \"graph.neighbors\";\nconst std::string HnswEntity::kGraphOffsetsSegmentId = \"graph.offsets\";\nconst std::string HnswEntity::kGraphMappingSegmentId = \"graph.mapping\";\nconst std::string HnswEntity::kHnswHeaderSegmentId = \"hnsw.header\";\nconst std::string HnswEntity::kHnswNeighborsSegmentId = \"hnsw.neighbors\";\nconst std::string HnswEntity::kHnswOffsetsSegmentId = \"hnsw.offsets\";\n\nint HnswEntity::CalcAndAddPadding(const IndexDumper::Pointer &dumper,\n size_t data_size, size_t *padding_size) {\n *padding_size = AlignSize(data_size) - data_size;\n if (*padding_size == 0) {\n return 0;\n }\n\n std::string padding(*padding_size, '\\0');\n if (dumper->write(padding.data(), *padding_size) != *padding_size) {\n LOG_ERROR(\"Append padding failed, size %lu\", *padding_size);\n return IndexError_WriteData;\n }\n return 0;\n}\n\nint64_t HnswEntity::dump_segment(const IndexDumper::Pointer &dumper,\n const std::string &segment_id,\n const void *data, size_t size) const {\n size_t len = dumper->write(data, size);\n if (len != size) {\n LOG_ERROR(\"Dump segment %s data failed, expect: %lu, actual: %lu\",\n segment_id.c_str(), size, len);\n return IndexError_WriteData;\n }\n\n size_t padding_size = AlignSize(size) - size;\n if (padding_size > 0) {\n std::string padding(padding_size, '\\0');\n if (dumper->write(padding.data(), padding_size) != padding_size) {\n LOG_ERROR(\"Append padding failed, size %lu\", padding_size);\n return IndexError_WriteData;\n }\n }\n\n uint32_t crc = ailego::Crc32c::Hash(data, size);\n int ret = dumper->append(segment_id, size, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s meta failed, ret=%d\", segment_id.c_str(), ret);\n return ret;\n }\n\n return len + padding_size;\n}\n\nint64_t HnswEntity::dump_header(const IndexDumper::Pointer &dumper,\n const HNSWHeader &hd) const {\n //! dump basic graph header. header is aligned and does not need padding\n int64_t graph_hd_size =\n dump_segment(dumper, kGraphHeaderSegmentId, &hd.graph, hd.graph.size);\n if (graph_hd_size < 0) {\n return graph_hd_size;\n }\n\n //! dump basic graph header. header is aligned and does not need padding\n int64_t hnsw_hd_size =\n dump_segment(dumper, kHnswHeaderSegmentId, &hd.hnsw, hd.hnsw.size);\n if (hnsw_hd_size < 0) {\n return hnsw_hd_size;\n }\n\n return graph_hd_size + hnsw_hd_size;\n}\n\nvoid HnswEntity::reshuffle_vectors(\n const std::function & /*get_level*/,\n std::vector * /*n2o_mapping*/,\n std::vector * /*o2n_mapping*/, key_t * /*keys*/) const {\n // TODO\n return;\n}\n\nint64_t HnswEntity::dump_mapping_segment(const IndexDumper::Pointer &dumper,\n const key_t *keys) const {\n std::vector mapping(doc_cnt());\n\n std::iota(mapping.begin(), mapping.end(), 0U);\n std::sort(mapping.begin(), mapping.end(),\n [&](node_id_t i, node_id_t j) { return keys[i] < keys[j]; });\n\n size_t size = mapping.size() * sizeof(node_id_t);\n\n return dump_segment(dumper, kGraphMappingSegmentId, mapping.data(), size);\n}\n\nint64_t HnswEntity::dump_segments(\n const IndexDumper::Pointer &dumper, key_t *keys,\n const std::function &get_level) const {\n HNSWHeader dump_hd(header());\n\n dump_hd.graph.node_size = AlignSize(vector_size());\n\n std::vector n2o_mapping; // map new id to origin id\n std::vector o2n_mapping; // map origin id to new id\n reshuffle_vectors(get_level, &n2o_mapping, &o2n_mapping, keys);\n if (!o2n_mapping.empty()) {\n dump_hd.hnsw.entry_point = o2n_mapping[entry_point()];\n }\n\n //! Dump header\n int64_t hd_size = dump_header(dumper, dump_hd);\n if (hd_size < 0) {\n return hd_size;\n }\n\n //! Dump vectors\n int64_t vecs_size = dump_vectors(dumper, n2o_mapping);\n if (vecs_size < 0) {\n return vecs_size;\n }\n\n //! Dump neighbors\n auto neighbors_size =\n dump_neighbors(dumper, get_level, n2o_mapping, o2n_mapping);\n if (neighbors_size < 0) {\n return neighbors_size;\n }\n //! free memory\n n2o_mapping = std::vector();\n o2n_mapping = std::vector();\n\n //! Dump keys\n size_t key_segment_size = doc_cnt() * sizeof(key_t);\n int64_t keys_size =\n dump_segment(dumper, kGraphKeysSegmentId, keys, key_segment_size);\n if (keys_size < 0) {\n return keys_size;\n }\n\n //! Dump mapping\n int64_t mapping_size = dump_mapping_segment(dumper, keys);\n if (mapping_size < 0) {\n return mapping_size;\n }\n\n return hd_size + keys_size + vecs_size + neighbors_size + mapping_size;\n}\n\nint64_t HnswEntity::dump_vectors(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping) const {\n size_t vector_dump_size = vector_size();\n\n size_t padding_size = AlignSize(vector_dump_size) - vector_dump_size;\n\n std::vector padding(padding_size);\n memset(padding.data(), 0, sizeof(char) * padding_size);\n const void *data = nullptr;\n uint32_t crc = 0U;\n size_t vecs_size = 0UL;\n\n //! dump vectors\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n data = get_vector(reorder_mapping.empty() ? id : reorder_mapping[id]);\n if (ailego_unlikely(!data)) {\n return IndexError_ReadData;\n }\n size_t len = dumper->write(data, vector_size());\n if (len != vector_size()) {\n LOG_ERROR(\"Dump vectors failed, write=%zu expect=%zu\", len,\n vector_size());\n return IndexError_WriteData;\n }\n\n crc = ailego::Crc32c::Hash(data, vector_size(), crc);\n vecs_size += vector_size();\n\n if (padding_size == 0) {\n continue;\n }\n\n len = dumper->write(padding.data(), padding_size);\n if (len != padding_size) {\n LOG_ERROR(\"Dump vectors failed, write=%zu expect=%zu\", len, padding_size);\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(padding.data(), padding_size, crc);\n vecs_size += padding_size;\n }\n\n int ret = dumper->append(kGraphFeaturesSegmentId, vecs_size, 0UL, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump vectors segment meta failed, ret %d\", ret);\n return ret;\n }\n\n return vecs_size;\n}\n\nint64_t HnswEntity::dump_graph_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const {\n std::vector graph_meta;\n graph_meta.reserve(doc_cnt());\n size_t offset = 0;\n uint32_t crc = 0;\n std::vector mapping(l0_neighbor_cnt());\n\n uint32_t min_neighbor_count = 10000;\n uint32_t max_neighbor_count = 0;\n size_t sum_neighbor_count = 0;\n\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n const Neighbors neighbors =\n get_neighbors(0, reorder_mapping.empty() ? id : reorder_mapping[id]);\n ailego_assert_with(!!neighbors.data, \"invalid neighbors\");\n ailego_assert_with(neighbors.size() <= l0_neighbor_cnt(),\n \"invalid neighbors\");\n\n uint32_t neighbor_count = neighbors.size();\n if (neighbor_count < min_neighbor_count) {\n min_neighbor_count = neighbor_count;\n }\n if (neighbor_count > max_neighbor_count) {\n max_neighbor_count = neighbor_count;\n }\n sum_neighbor_count += neighbor_count;\n\n graph_meta.emplace_back(offset, neighbor_count);\n size_t size = neighbors.size() * sizeof(node_id_t);\n const node_id_t *data = &neighbors[0];\n if (!neighbor_mapping.empty()) {\n for (node_id_t i = 0; i < neighbors.size(); ++i) {\n mapping[i] = neighbor_mapping[neighbors[i]];\n }\n data = mapping.data();\n }\n if (dumper->write(data, size) != size) {\n LOG_ERROR(\"Dump graph neighbor id=%u failed, size %lu\", id, size);\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(data, size, crc);\n offset += size;\n }\n\n uint32_t average_neighbor_count = 0;\n if (doc_cnt() > 0) {\n average_neighbor_count = sum_neighbor_count / doc_cnt();\n }\n LOG_INFO(\n \"Dump hnsw graph: min_neighbor_count[%u] max_neighbor_count[%u] \"\n \"average_neighbor_count[%u]\",\n min_neighbor_count, max_neighbor_count, average_neighbor_count);\n\n size_t padding_size = 0;\n int ret = CalcAndAddPadding(dumper, offset, &padding_size);\n if (ret != 0) {\n return ret;\n }\n ret = dumper->append(kGraphNeighborsSegmentId, offset, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s failed, ret %d\",\n kGraphNeighborsSegmentId.c_str(), ret);\n return ret;\n }\n\n //! dump level 0 neighbors meta\n auto len = dump_segment(dumper, kGraphOffsetsSegmentId, graph_meta.data(),\n graph_meta.size() * sizeof(GraphNeighborMeta));\n if (len < 0) {\n return len;\n }\n\n return len + offset + padding_size;\n}\n\nint64_t HnswEntity::dump_upper_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::function &get_level,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const {\n std::vector hnsw_meta;\n hnsw_meta.reserve(doc_cnt());\n size_t offset = 0;\n uint32_t crc = 0;\n std::vector buffer(upper_neighbor_cnt() + 1);\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n node_id_t new_id = reorder_mapping.empty() ? id : reorder_mapping[id];\n auto level = get_level(new_id);\n if (level == 0) {\n hnsw_meta.emplace_back(0U, 0U);\n continue;\n }\n hnsw_meta.emplace_back(offset, level);\n ailego_assert_with((size_t)level < kMaxGraphLayers, \"invalid level\");\n for (level_t cur_level = 1; cur_level <= level; ++cur_level) {\n const Neighbors neighbors = get_neighbors(cur_level, new_id);\n ailego_assert_with(!!neighbors.data, \"invalid neighbors\");\n ailego_assert_with(neighbors.size() <= neighbor_cnt(cur_level),\n \"invalid neighbors\");\n memset(buffer.data(), 0, sizeof(node_id_t) * buffer.size());\n buffer[0] = neighbors.size();\n if (neighbor_mapping.empty()) {\n memcpy(&buffer[1], &neighbors[0], neighbors.size() * sizeof(node_id_t));\n } else {\n for (node_id_t i = 0; i < neighbors.size(); ++i) {\n buffer[i + 1] = neighbor_mapping[neighbors[i]];\n }\n }\n if (dumper->write(buffer.data(), sizeof(node_id_t) * buffer.size()) !=\n sizeof(node_id_t) * buffer.size()) {\n LOG_ERROR(\"Dump graph neighbor id=%u failed, size %lu\", id,\n sizeof(node_id_t) * buffer.size());\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(buffer.data(),\n sizeof(node_id_t) * buffer.size(), crc);\n offset += sizeof(node_id_t) * buffer.size();\n }\n }\n size_t padding_size = 0;\n int ret = CalcAndAddPadding(dumper, offset, &padding_size);\n if (ret != 0) {\n return ret;\n }\n\n ret = dumper->append(kHnswNeighborsSegmentId, offset, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s failed, ret %d\", kHnswNeighborsSegmentId.c_str(),\n ret);\n return ret;\n }\n\n //! dump level 0 neighbors meta\n auto len = dump_segment(dumper, kHnswOffsetsSegmentId, hnsw_meta.data(),\n hnsw_meta.size() * sizeof(HnswNeighborMeta));\n if (len < 0) {\n return len;\n }\n\n return len + offset + padding_size;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nusing node_id_t = uint32_t;\nusing key_t = uint64_t;\nusing level_t = int32_t;\nusing dist_t = float;\nusing TopkHeap = ailego::KeyValueHeap;\nusing CandidateHeap =\n ailego::KeyValueHeap>;\nconstexpr node_id_t kInvalidNodeId = static_cast(-1);\nconstexpr key_t kInvalidKey = static_cast(-1);\nclass DistCalculator;\n\nstruct GraphHeader {\n uint32_t size;\n uint32_t version;\n uint32_t graph_type;\n uint32_t doc_count;\n uint32_t vector_size;\n uint32_t node_size;\n uint32_t l0_neighbor_count;\n uint32_t prune_type;\n uint32_t prune_neighbor_count;\n uint32_t ef_construction;\n uint32_t options;\n uint32_t min_neighbor_count;\n uint8_t reserved_[4080];\n};\n\nstatic_assert(sizeof(GraphHeader) % 32 == 0,\n \"GraphHeader must be aligned with 32 bytes\");\n\n//! Hnsw upper neighbor header\nstruct HnswHeader {\n uint32_t size; // header size\n uint32_t revision; // current total docs of the graph\n uint32_t upper_neighbor_count;\n uint32_t ef_construction;\n uint32_t scaling_factor;\n uint32_t max_level;\n uint32_t entry_point;\n uint32_t options;\n uint8_t reserved_[30];\n};\n\nstatic_assert(sizeof(HnswHeader) % 32 == 0,\n \"GraphHeader must be aligned with 32 bytes\");\n\n//! Hnsw common header and upper neighbor header\nstruct HNSWHeader {\n HNSWHeader() {\n clear();\n }\n\n HNSWHeader(const HNSWHeader &header) {\n memcpy(this, &header, sizeof(header));\n }\n\n HNSWHeader &operator=(const HNSWHeader &header) {\n memcpy(this, &header, sizeof(header));\n return *this;\n }\n\n //! Reset state to zero, and the params is untouched\n void inline reset() {\n graph.doc_count = 0U;\n hnsw.entry_point = kInvalidNodeId;\n hnsw.max_level = 0;\n }\n\n //! Clear all fields to init value\n void inline clear() {\n memset(this, 0, sizeof(HNSWHeader));\n hnsw.entry_point = kInvalidNodeId;\n graph.size = sizeof(GraphHeader);\n hnsw.size = sizeof(HnswHeader);\n }\n\n size_t l0_neighbor_cnt() const {\n return graph.l0_neighbor_count;\n }\n\n size_t upper_neighbor_cnt() const {\n return hnsw.upper_neighbor_count;\n }\n\n size_t vector_size() const {\n return graph.vector_size;\n }\n\n size_t ef_construction() const {\n return graph.ef_construction;\n }\n\n size_t scaling_factor() const {\n return hnsw.scaling_factor;\n }\n\n size_t neighbor_prune_cnt() const {\n return graph.prune_neighbor_count;\n }\n\n node_id_t entry_point() const {\n return hnsw.entry_point;\n }\n\n node_id_t doc_cnt() const {\n return graph.doc_count;\n }\n\n GraphHeader graph;\n HnswHeader hnsw;\n};\n\nstruct NeighborsHeader {\n uint32_t neighbor_cnt;\n node_id_t neighbors[0];\n};\n\nstruct Neighbors {\n Neighbors() : cnt{0}, data{nullptr} {}\n\n Neighbors(uint32_t cnt_in, const node_id_t *data_in)\n : cnt{cnt_in}, data{data_in} {}\n\n Neighbors(const IndexStorage::MemoryBlock &mem_block)\n : neighbor_block{mem_block} {\n auto hd = reinterpret_cast(neighbor_block.data());\n cnt = hd->neighbor_cnt;\n data = hd->neighbors;\n }\n\n size_t size(void) const {\n return cnt;\n }\n\n const node_id_t &operator[](size_t idx) const {\n return data[idx];\n }\n\n uint32_t cnt;\n const node_id_t *data;\n IndexStorage::MemoryBlock neighbor_block;\n};\n\n//! level 0 neighbors offset\nstruct GraphNeighborMeta {\n GraphNeighborMeta(size_t o, size_t cnt) : offset(o), neighbor_cnt(cnt) {}\n\n uint64_t offset : 48;\n uint64_t neighbor_cnt : 16;\n};\n\n//! hnsw upper neighbors meta\nstruct HnswNeighborMeta {\n HnswNeighborMeta(size_t o, size_t l) : offset(o), level(l) {}\n\n uint64_t offset : 48; // offset = idx * upper neighors size\n uint64_t level : 16;\n};\n\nclass HnswEntity {\n public:\n //! Constructor\n HnswEntity() {}\n\n //! Constructor\n HnswEntity(const HNSWHeader &hd) {\n header_ = hd;\n }\n\n //! Destructor\n virtual ~HnswEntity() {}\n\n //! HnswEntity Pointerd;\n typedef std::shared_ptr Pointer;\n\n //! Get max neighbor size of graph level\n inline size_t neighbor_cnt(level_t level) const {\n return level == 0 ? header_.graph.l0_neighbor_count\n : header_.hnsw.upper_neighbor_count;\n }\n\n //! get max neighbor size of graph level 0\n inline size_t l0_neighbor_cnt() const {\n return header_.graph.l0_neighbor_count;\n }\n\n //! get min neighbor size of graph\n inline size_t min_neighbor_cnt() const {\n return header_.graph.min_neighbor_count;\n }\n\n //! get upper neighbor size of graph level other than 0\n inline size_t upper_neighbor_cnt() const {\n return header_.hnsw.upper_neighbor_count;\n }\n\n //! Get current total doc of the hnsw graph\n inline node_id_t *mutable_doc_cnt() {\n return &header_.graph.doc_count;\n }\n\n inline node_id_t doc_cnt() const {\n return header_.graph.doc_count;\n }\n\n //! Get hnsw graph scaling params\n inline size_t scaling_factor() const {\n return header_.hnsw.scaling_factor;\n }\n\n //! Get prune_size\n inline size_t prune_cnt() const {\n return header_.graph.prune_neighbor_count;\n }\n\n //! Current entity of top level graph\n inline node_id_t entry_point() const {\n return header_.hnsw.entry_point;\n }\n\n //! Current max graph level\n inline level_t cur_max_level() const {\n return header_.hnsw.max_level;\n }\n\n //! Retrieve index vector size\n size_t vector_size() const {\n return header_.graph.vector_size;\n }\n\n //! Retrieve node size\n size_t node_size() const {\n return header_.graph.node_size;\n }\n\n //! Retrieve ef constuction\n size_t ef_construction() const {\n return header_.graph.ef_construction;\n }\n\n void set_vector_size(size_t size) {\n header_.graph.vector_size = size;\n }\n\n void set_prune_cnt(size_t v) {\n header_.graph.prune_neighbor_count = v;\n }\n\n void set_scaling_factor(size_t val) {\n header_.hnsw.scaling_factor = val;\n }\n\n void set_l0_neighbor_cnt(size_t cnt) {\n header_.graph.l0_neighbor_count = cnt;\n }\n\n void set_min_neighbor_cnt(size_t cnt) {\n header_.graph.min_neighbor_count = cnt;\n }\n\n void set_upper_neighbor_cnt(size_t cnt) {\n header_.hnsw.upper_neighbor_count = cnt;\n }\n\n void set_ef_construction(size_t ef) {\n header_.graph.ef_construction = ef;\n }\n\n protected:\n inline const HNSWHeader &header() const {\n return header_;\n }\n\n inline HNSWHeader *mutable_header() {\n return &header_;\n }\n\n inline size_t header_size() const {\n return sizeof(header_);\n }\n\n void set_node_size(size_t size) {\n header_.graph.node_size = size;\n }\n\n //! Dump all segment by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_segments(\n const IndexDumper::Pointer &dumper, key_t *keys,\n const std::function &get_level) const;\n\n private:\n //! dump mapping segment, for get_vector_by_key in provider\n int64_t dump_mapping_segment(const IndexDumper::Pointer &dumper,\n const key_t *keys) const;\n\n //! dump hnsw head by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_header(const IndexDumper::Pointer &dumper,\n const HNSWHeader &hd) const;\n\n //! dump vectors by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_vectors(const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping) const;\n\n //! dump hnsw neighbors by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_neighbors(const IndexDumper::Pointer &dumper,\n const std::function &get_level,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const {\n auto len1 = dump_graph_neighbors(dumper, reorder_mapping, neighbor_mapping);\n if (len1 < 0) {\n return len1;\n }\n auto len2 = dump_upper_neighbors(dumper, get_level, reorder_mapping,\n neighbor_mapping);\n if (len2 < 0) {\n return len2;\n }\n\n return len1 + len2;\n }\n\n //! dump segment by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_segment(const IndexDumper::Pointer &dumper,\n const std::string &segment_id, const void *data,\n size_t size) const;\n\n //! Dump level 0 neighbors\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_graph_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const;\n\n //! Dump upper level neighbors\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_upper_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::function &get_level,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const;\n\n public:\n //! Cleanup the entity\n virtual int cleanup(void) {\n header_.clear();\n return 0;\n }\n\n //! Make a copy of searcher entity, to support thread-safe operation.\n //! The segment in container cannot be read concurrenly\n virtual const HnswEntity::Pointer clone() const {\n LOG_ERROR(\"Update neighbors not implemented\");\n return HnswEntity::Pointer();\n }\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const = 0;\n\n //! Get vector feature data by key\n virtual const void *get_vector(node_id_t id) const = 0;\n\n //! Get vectors feature data by keys\n virtual int get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const = 0;\n\n virtual int get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const = 0;\n virtual int get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const = 0;\n\n //! Retrieve a vector using a primary key\n virtual const void *get_vector_by_key(uint64_t /*key*/) const {\n LOG_ERROR(\"get vector not implemented\");\n return nullptr;\n }\n\n virtual int get_vector_by_key(const key_t /*key*/,\n IndexStorage::MemoryBlock & /*block*/) const {\n return IndexError_NotImplemented;\n }\n\n //! Get the node id's neighbors on graph level\n //! Note: the neighbors cannot be modified, using the following\n //! method to get WritableNeighbors if want to\n virtual const Neighbors get_neighbors(level_t level, node_id_t id) const = 0;\n\n //! Add vector and key to hnsw entity, and local id will be saved in id\n virtual int add_vector(level_t /*level*/, key_t /*key*/, const void * /*vec*/,\n node_id_t * /*id*/) {\n return IndexError_NotImplemented;\n }\n\n //! Add vector and id to hnsw entity\n virtual int add_vector_with_id(level_t /*level*/, node_id_t /*id*/,\n const void * /*vec*/) {\n return IndexError_NotImplemented;\n }\n\n virtual int update_neighbors(\n level_t /*level*/, node_id_t /*id*/,\n const std::vector> & /*neighbors*/) {\n LOG_ERROR(\"Update neighbors dense not implemented\");\n\n return 0;\n }\n\n //! Append neighbor_id to node id neighbors on level, size is the current\n //! neighbors size. Notice: the caller must be ensure the neighbors not full\n virtual void add_neighbor(level_t /*level*/, node_id_t /*id*/,\n uint32_t /*size*/, node_id_t /*neighbor_id*/) {\n LOG_ERROR(\"Add neighbor not implemented\");\n }\n\n //! Update entry point and max level\n virtual void update_ep_and_level(node_id_t ep, level_t level) {\n header_.hnsw.entry_point = ep;\n header_.hnsw.max_level = level;\n }\n\n virtual int load(const IndexStorage::Pointer & /*container*/,\n bool /*check_crc*/) {\n LOG_ERROR(\"Load not implemented\");\n return IndexError_NotImplemented;\n }\n\n virtual int dump(const IndexDumper::Pointer & /*dumper*/) {\n LOG_ERROR(\"Dump not implemented\");\n return IndexError_NotImplemented;\n }\n\n static int CalcAndAddPadding(const IndexDumper::Pointer &dumper,\n size_t data_size, size_t *padding_size);\n\n protected:\n static inline size_t AlignSize(size_t size) {\n return (size + 0x1F) & (~0x1F);\n }\n\n static inline size_t AlignPageSize(size_t size) {\n size_t page_mask = ailego::MemoryHelper::PageSize() - 1;\n return (size + page_mask) & (~page_mask);\n }\n\n static inline size_t AlignHugePageSize(size_t size) {\n size_t page_mask = ailego::MemoryHelper::HugePageSize() - 1;\n return (size + page_mask) & (~page_mask);\n }\n\n //! rearrange vectors to improve cache locality\n void reshuffle_vectors(const std::function &get_level,\n std::vector *n2o_mapping,\n std::vector *o2n_mapping,\n key_t *keys) const;\n\n public:\n const static std::string kGraphHeaderSegmentId;\n const static std::string kGraphFeaturesSegmentId;\n const static std::string kGraphKeysSegmentId;\n const static std::string kGraphNeighborsSegmentId;\n const static std::string kGraphOffsetsSegmentId;\n const static std::string kGraphMappingSegmentId;\n const static std::string kHnswHeaderSegmentId;\n const static std::string kHnswNeighborsSegmentId;\n const static std::string kHnswOffsetsSegmentId;\n\n constexpr static uint32_t kRevision = 0U;\n constexpr static size_t kMaxGraphLayers = 15;\n constexpr static uint32_t kDefaultEfConstruction = 500;\n constexpr static uint32_t kDefaultEf = 500;\n constexpr static uint32_t kDefaultUpperMaxNeighborCnt = 50; // M of HNSW\n constexpr static uint32_t kDefaultL0MaxNeighborCnt = 100;\n constexpr static uint32_t kMaxNeighborCnt = 65535;\n constexpr static float kDefaultScanRatio = 0.1f;\n constexpr static uint32_t kDefaultMinScanLimit = 10000;\n constexpr static uint32_t kDefaultMaxScanLimit =\n std::numeric_limits::max();\n constexpr static float kDefaultBFNegativeProbability = 0.001f;\n constexpr static uint32_t kDefaultScalingFactor = 50U;\n constexpr static uint32_t kDefaultBruteForceThreshold = 1000U;\n constexpr static uint32_t kDefaultDocsHardLimit = 1 << 30U; // 1 billion\n constexpr static float kDefaultDocsSoftLimitRatio = 0.9f;\n constexpr static size_t kMaxChunkSize = 0xFFFFFFFF;\n constexpr static size_t kDefaultChunkSize = 2UL * 1024UL * 1024UL;\n constexpr static size_t kDefaultMaxChunkCnt = 50000UL;\n constexpr static float kDefaultNeighborPruneMultiplier =\n 1.0f; // prune_cnt = upper_max_neighbor_cnt * multiplier\n constexpr static float kDefaultL0MaxNeighborCntMultiplier =\n 2.0f; // l0_max_neighbor_cnt = upper_max_neighbor_cnt * multiplier\n\n protected:\n HNSWHeader header_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_index_hash.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"hnsw_chunk.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! Persistent hashmap implement through open addressing algorithm\ntemplate ::value>::type>\nclass HnswIndexHashMap {\n using key_type = Key;\n using val_type = Val;\n\n struct Iterator {\n key_type first;\n val_type second;\n };\n typedef Iterator *iterator;\n typedef Iterator Item;\n typedef const Iterator *const_iterator;\n\n class Slot {\n public:\n Slot(Chunk::Pointer &&chunk, const void *data)\n : chunk_(std::move(chunk)),\n items_(reinterpret_cast(data)) {}\n //! Return a empty loc or the key item loc\n\n Slot(Chunk::Pointer &&chunk, IndexStorage::MemoryBlock &&mem_block)\n : chunk_(std::move(chunk)), items_block_(std::move(mem_block)) {\n items_ = reinterpret_cast(items_block_.data());\n }\n const_iterator find(key_type key, uint32_t max_items, uint32_t mask) const {\n auto it = &items_[key & mask];\n for (auto i = 0U; i < max_items; ++i) {\n if (it->first == key || it->second == EmptyVal) {\n // LOG_DEBUG(\"i=%u\", i);\n return it;\n }\n ++it;\n if (it == &items_[max_items]) {\n it = &items_[0];\n }\n }\n return nullptr;\n }\n\n bool update(const_iterator it) {\n uint32_t offset = reinterpret_cast(it) -\n reinterpret_cast(&items_[0]);\n if (ailego_unlikely(chunk_->write(offset, it, sizeof(Item)) !=\n sizeof(Item))) {\n LOG_ERROR(\"Chunk write failed\");\n return false;\n }\n return true;\n }\n\n private:\n Chunk::Pointer chunk_{};\n const Item *items_{nullptr}; // point to chunk data\n IndexStorage::MemoryBlock items_block_{};\n };\n\n public:\n //! Init the hash\n //! broker the index allocator\n //! chunk_size the size of per chunk allocated, actual size may greater\n //! factor factor = 1/ratio, ratio is the probability of a squence\n //! number inserted to this container\n //! max the max number key can be inserted\n //! expansion_ratio memory expansion ratio\n int init(ChunkBroker::Pointer &broker, uint32_t chunk_size, uint32_t factor,\n size_t max, float expansion_ratio) {\n ailego_assert_with(expansion_ratio > 1.0f, \"ratio must > 1.0f\");\n broker_ = broker;\n\n size_t items = std::ceil(chunk_size * 1.0f / sizeof(Item));\n slot_items_ = 1UL << static_cast((std::ceil(std::log2(items))));\n size_t range = slot_items_ * factor / expansion_ratio;\n mask_bits_ = std::floor(std::log2(range));\n range = 1UL << mask_bits_;\n size_t max_slots = std::ceil(max * 1.0f / range);\n slots_.reserve(max_slots);\n slot_loc_mask_ = slot_items_ - 1U;\n int ret = load();\n if (ret != 0) {\n return ret;\n }\n\n LOG_DEBUG(\n \"HnswIndexHash init, chunkSize=%u factor=%u max=%zu \"\n \"ratio=%f slotItems=%u maxSlots=%zu maskBits=%u \"\n \"range=%zu\",\n chunk_size, factor, max, expansion_ratio, slot_items_, max_slots,\n mask_bits_, range);\n\n return 0;\n }\n\n int cleanup(void) {\n broker_.reset();\n slots_.clear();\n slots_.shrink_to_fit();\n mask_bits_ = 0U;\n slot_items_ = 0U;\n slot_loc_mask_ = 0U;\n\n return 0;\n }\n\n const_iterator end(void) const {\n return nullptr;\n }\n\n const_iterator find(const key_type key) const {\n auto idx = key >> mask_bits_;\n if (idx >= slots_.size()) {\n return end();\n }\n auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_);\n return it && it->second != EmptyVal ? it : nullptr;\n }\n\n bool insert(key_type key, val_type val) {\n auto idx = key >> mask_bits_;\n if (idx >= slots_.size()) {\n if (ailego_unlikely(idx >= slots_.capacity())) {\n LOG_ERROR(\"no space to insert\");\n return false;\n }\n for (auto i = slots_.size(); i <= idx; ++i) {\n if (ailego_unlikely(!alloc_slot(i))) {\n return false;\n }\n }\n }\n auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_);\n if (ailego_unlikely(it == nullptr)) {\n LOG_ERROR(\"no space to insert\");\n return false;\n }\n\n //! TODO: write memory is ok?\n const_cast(it)->first = key;\n const_cast(it)->second = val;\n\n return slots_[idx].update(it);\n }\n\n private:\n bool alloc_slot(size_t idx) {\n ailego_assert_with(idx == slots_.size(), \"invalid idx\");\n\n size_t size = slot_items_ * sizeof(Item);\n auto p =\n broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, idx, size);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return false;\n }\n Chunk::Pointer chunk = p.second;\n if (ailego_unlikely(chunk->resize(size) != size)) {\n LOG_ERROR(\"Chunk resize failed, size=%zu\", size);\n return false;\n }\n //! Read the whole data to memory\n IndexStorage::MemoryBlock data_block;\n if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {\n LOG_ERROR(\"Chunk read failed, size=%zu\", size);\n return false;\n }\n\n slots_.emplace_back(std::move(chunk), std::move(data_block));\n return true;\n }\n\n int load(void) {\n size_t slots_cnt =\n broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX);\n for (size_t i = 0UL; i < slots_cnt; ++i) {\n auto chunk =\n broker_->get_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, i);\n if (!chunk) {\n LOG_ERROR(\"Get chunk failed, seq=%zu\", i);\n return IndexError_InvalidFormat;\n }\n size_t size = sizeof(Item) * slot_items_;\n if (chunk->data_size() < size) {\n LOG_ERROR(\n \"Hash params may be mismatch, seq=%zu, data_size=%zu \"\n \"expect=%zu\",\n i, chunk->data_size(), size);\n return IndexError_InvalidFormat;\n }\n //! Read the whole data to memory\n IndexStorage::MemoryBlock data_block;\n if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {\n LOG_ERROR(\"Chunk read failed, size=%zu\", size);\n return false;\n }\n slots_.emplace_back(std::move(chunk), std::move(data_block));\n }\n return 0;\n }\n\n private:\n ChunkBroker::Pointer broker_{}; // chunk broker\n std::vector slots_{};\n uint32_t mask_bits_{0U};\n uint32_t slot_items_{}; // must be a power of 2\n uint32_t slot_loc_mask_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_index_provider.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \"hnsw_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswIndexProvider : public IndexProvider {\n public:\n HnswIndexProvider(const IndexMeta &meta, const HnswEntity::Pointer &entity,\n const std::string &owner)\n : meta_(meta), entity_(entity), owner_class_(owner) {}\n\n HnswIndexProvider(const HnswIndexProvider &) = delete;\n HnswIndexProvider &operator=(const HnswIndexProvider &) = delete;\n\n public: // holder interface\n //! Create a new iterator\n IndexProvider::Iterator::Pointer create_iterator() override {\n return HnswIndexProvider::Iterator::Pointer(new (std::nothrow)\n Iterator(entity_));\n }\n\n //! Retrieve count of vectors\n size_t count(void) const override {\n return entity_->doc_cnt();\n }\n\n //! Retrieve dimension of vector\n size_t dimension(void) const override {\n return meta_.dimension();\n }\n\n //! Retrieve type of vector\n IndexMeta::DataType data_type(void) const override {\n return meta_.data_type();\n }\n\n //! Retrieve vector size in bytes\n size_t element_size(void) const override {\n return meta_.element_size();\n }\n\n public: // provider's unique interface\n //! Retrieve a vector using a primary key\n const void *get_vector(uint64_t key) const override {\n return entity_->get_vector_by_key(key);\n }\n\n int get_vector(const uint64_t key,\n IndexStorage::MemoryBlock &block) const override {\n return entity_->get_vector_by_key(key, block);\n }\n\n //! Retrieve the owner class\n const std::string &owner_class(void) const override {\n return owner_class_;\n }\n\n private:\n class Iterator : public IndexProvider::Iterator {\n public:\n Iterator(const HnswEntity::Pointer &entity)\n : entity_(entity), cur_id_(0U) {}\n\n //! Retrieve pointer of data\n //! NOTICE: the vec feature will be changed after iterating to next, so\n //! the caller need to keep a copy of it before iterator to next vector\n virtual const void *data(void) const override {\n return entity_->get_vector(cur_id_);\n }\n\n //! Test if the iterator is valid\n virtual bool is_valid(void) const override {\n return cur_id_ < entity_->doc_cnt();\n }\n\n //! Retrieve primary key\n virtual uint64_t key(void) const override {\n return entity_->get_key(cur_id_);\n }\n\n //! Next iterator\n virtual void next(void) override {\n // cur_id_ += 1;\n cur_id_ = get_next_valid_id(cur_id_ + 1);\n }\n\n //! Reset the iterator\n void reset(void) {\n cur_id_ = get_next_valid_id(0);\n }\n\n private:\n node_id_t get_next_valid_id(node_id_t start_id) {\n for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) {\n if (entity_->get_key(i) != kInvalidNodeId) {\n cur_id_ = i;\n return i;\n }\n }\n return kInvalidNodeId;\n }\n\n private:\n const HnswEntity::Pointer entity_;\n node_id_t cur_id_;\n };\n\n private:\n const IndexMeta &meta_;\n const HnswEntity::Pointer entity_;\n const std::string owner_class_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_params.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\nstatic const std::string PARAM_HNSW_BUILDER_THREAD_COUNT(\n \"proxima.hnsw.builder.thread_count\");\nstatic const std::string PARAM_HNSW_BUILDER_MEMORY_QUOTA(\n \"proxima.hnsw.builder.memory_quota\");\nstatic const std::string PARAM_HNSW_BUILDER_EFCONSTRUCTION(\n \"proxima.hnsw.builder.efconstruction\");\nstatic const std::string PARAM_HNSW_BUILDER_SCALING_FACTOR(\n \"proxima.hnsw.builder.scaling_factor\");\nstatic const std::string PARAM_HNSW_BUILDER_CHECK_INTERVAL_SECS(\n \"proxima.hnsw.builder.check_interval_secs\");\nstatic const std::string PARAM_HNSW_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER(\n \"proxima.hnsw.builder.neighbor_prune_multiplier\");\nstatic const std::string PARAM_HNSW_BUILDER_MIN_NEIGHBOR_COUNT(\n \"proxima.hnsw.builder.min_neighbor_count\");\nstatic const std::string PARAM_HNSW_BUILDER_MAX_NEIGHBOR_COUNT(\n \"proxima.hnsw.builder.max_neighbor_count\");\nstatic const std::string PARAM_HNSW_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER(\n \"proxima.hnsw.builder.l0_max_neighbor_count_multiplier\");\n\nstatic const std::string PARAM_HNSW_SEARCHER_EF(\"proxima.hnsw.searcher.ef\");\nstatic const std::string PARAM_HNSW_SEARCHER_BRUTE_FORCE_THRESHOLD(\n \"proxima.hnsw.searcher.brute_force_threshold\");\nstatic const std::string PARAM_HNSW_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE(\n \"proxima.hnsw.searcher.neighbors_in_memory_enable\");\nstatic const std::string PARAM_HNSW_SEARCHER_MAX_SCAN_RATIO(\n \"proxima.hnsw.searcher.max_scan_ratio\");\nstatic const std::string PARAM_HNSW_SEARCHER_CHECK_CRC_ENABLE(\n \"proxima.hnsw.searcher.check_crc_enable\");\nstatic const std::string PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_ENABLE(\n \"proxima.hnsw.searcher.visit_bloomfilter_enable\");\nstatic const std::string PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB(\n \"proxima.hnsw.searcher.visit_bloomfilter_negative_prob\");\nstatic const std::string PARAM_HNSW_SEARCHER_FORCE_PADDING_RESULT_ENABLE(\n \"proxima.hnsw.searcher.force_padding_result_enable\");\n\nstatic const std::string PARAM_HNSW_STREAMER_MAX_SCAN_RATIO(\n \"proxima.hnsw.streamer.max_scan_ratio\");\nstatic const std::string PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT(\n \"proxima.hnsw.streamer.min_scan_limit\");\nstatic const std::string PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT(\n \"proxima.hnsw.streamer.max_scan_limit\");\nstatic const std::string PARAM_HNSW_STREAMER_EF(\"proxima.hnsw.streamer.ef\");\nstatic const std::string PARAM_HNSW_STREAMER_EFCONSTRUCTION(\n \"proxima.hnsw.streamer.efconstruction\");\nstatic const std::string PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT(\n \"proxima.hnsw.streamer.max_neighbor_count\");\nstatic const std::string PARAM_HNSW_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER(\n \"proxima.hnsw.streamer.l0_max_neighbor_count_multiplier\");\nstatic const std::string PARAM_HNSW_STREAMER_SCALING_FACTOR(\n \"proxima.hnsw.streamer.scaling_factor\");\nstatic const std::string PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD(\n \"proxima.hnsw.streamer.brute_force_threshold\");\nstatic const std::string PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT(\n \"proxima.hnsw.streamer.docs_hard_limit\");\nstatic const std::string PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT(\n \"proxima.hnsw.streamer.docs_soft_limit\");\nstatic const std::string PARAM_HNSW_STREAMER_MAX_INDEX_SIZE(\n \"proxima.hnsw.streamer.max_index_size\");\nstatic const std::string PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_ENABLE(\n \"proxima.hnsw.streamer.visit_bloomfilter_enable\");\nstatic const std::string PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB(\n \"proxima.hnsw.streamer.visit_bloomfilter_negative_prob\");\nstatic const std::string PARAM_HNSW_STREAMER_CHECK_CRC_ENABLE(\n \"proxima.hnsw.streamer.check_crc_enable\");\nstatic const std::string PARAM_HNSW_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER(\n \"proxima.hnsw.streamer.neighbor_prune_multiplier\");\nstatic const std::string PARAM_HNSW_STREAMER_CHUNK_SIZE(\n \"proxima.hnsw.streamer.chunk_size\");\nstatic const std::string PARAM_HNSW_STREAMER_FILTER_SAME_KEY(\n \"proxima.hnsw.streamer.filter_same_key\");\nstatic const std::string PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE(\n \"proxima.hnsw.streamer.get_vector_enable\");\nstatic const std::string PARAM_HNSW_STREAMER_MIN_NEIGHBOR_COUNT(\n \"proxima.hnsw.streamer.min_neighbor_count\");\nstatic const std::string PARAM_HNSW_STREAMER_FORCE_PADDING_RESULT_ENABLE(\n \"proxima.hnsw.streamer.force_padding_result_enable\");\nstatic const std::string PARAM_HNSW_STREAMER_ESTIMATE_DOC_COUNT(\n \"proxima.hnsw.streamer.estimate_doc_count\");\nstatic const std::string PARAM_HNSW_STREAMER_USE_ID_MAP(\n \"proxima.hnsw.streamer.use_id_map\");\n\nstatic const std::string PARAM_HNSW_REDUCER_WORKING_PATH(\n \"proxima.hnsw.reducer.working_path\");\nstatic const std::string PARAM_HNSW_REDUCER_NUM_OF_ADD_THREADS(\n \"proxima.hnsw.reducer.num_of_add_threads\");\nstatic const std::string PARAM_HNSW_REDUCER_INDEX_NAME(\n \"proxima.hnsw.reducer.index_name\");\nstatic const std::string PARAM_HNSW_REDUCER_EFCONSTRUCTION(\n \"proxima.hnsw.reducer.efconstruction\");\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_searcher.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_searcher.h\"\n#include \"hnsw_algorithm.h\"\n#include \"hnsw_index_provider.h\"\n#include \"hnsw_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswSearcher::HnswSearcher() = default;\n\nHnswSearcher::~HnswSearcher() = default;\n\nint HnswSearcher::init(const ailego::Params &search_params) {\n params_ = search_params;\n params_.get(PARAM_HNSW_SEARCHER_EF, &ef_);\n params_.get(PARAM_HNSW_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_);\n params_.get(PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_ENABLE, &bf_enabled_);\n params_.get(PARAM_HNSW_SEARCHER_CHECK_CRC_ENABLE, &check_crc_enabled_);\n params_.get(PARAM_HNSW_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE,\n &neighbors_in_memory_enabled_);\n params_.get(PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB,\n &bf_negative_probability_);\n params_.get(PARAM_HNSW_SEARCHER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n params_.get(PARAM_HNSW_SEARCHER_FORCE_PADDING_RESULT_ENABLE,\n &force_padding_topk_enabled_);\n\n if (ef_ == 0) {\n ef_ = HnswEntity::kDefaultEf;\n }\n if (bf_negative_probability_ <= 0.0f || bf_negative_probability_ >= 1.0f) {\n LOG_ERROR(\"[%s] must be in range (0,1)\",\n PARAM_HNSW_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str());\n return IndexError_InvalidArgument;\n }\n\n entity_.set_neighbors_in_memory(neighbors_in_memory_enabled_);\n\n state_ = STATE_INITED;\n\n LOG_DEBUG(\n \"Init params: ef=%u maxScanRatio=%f bfEnabled=%u checkCrcEnabled=%u \"\n \"neighborsInMemoryEnabled=%u bfNagtiveProb=%f bruteForceThreshold=%u \"\n \"forcePadding=%u\",\n ef_, max_scan_ratio_, bf_enabled_, check_crc_enabled_,\n neighbors_in_memory_enabled_, bf_negative_probability_,\n bruteforce_threshold_, force_padding_topk_enabled_);\n\n return 0;\n}\n\nvoid HnswSearcher::print_debug_info() {\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n Neighbors neighbours = entity_.get_neighbors(0, id);\n std::cout << \"node: \" << id << \"; \";\n for (uint32_t i = 0; i < neighbours.size(); ++i) {\n std::cout << neighbours[i];\n\n if (i == neighbours.size() - 1) {\n std::cout << std::endl;\n } else {\n std::cout << \", \";\n }\n }\n }\n}\n\nint HnswSearcher::cleanup() {\n LOG_INFO(\"Begin HnswSearcher:cleanup\");\n\n metric_.reset();\n meta_.clear();\n stats_.clear_attributes();\n stats_.set_loaded_count(0UL);\n stats_.set_loaded_costtime(0UL);\n max_scan_ratio_ = HnswEntity::kDefaultScanRatio;\n max_scan_num_ = 0U;\n ef_ = HnswEntity::kDefaultEf;\n bf_enabled_ = false;\n bf_negative_probability_ = HnswEntity::kDefaultBFNegativeProbability;\n bruteforce_threshold_ = HnswEntity::kDefaultBruteForceThreshold;\n check_crc_enabled_ = false;\n neighbors_in_memory_enabled_ = false;\n entity_.cleanup();\n state_ = STATE_INIT;\n\n LOG_INFO(\"End HnswSearcher:cleanup\");\n\n return 0;\n}\n\nint HnswSearcher::load(IndexStorage::Pointer container,\n IndexMetric::Pointer metric) {\n if (state_ != STATE_INITED) {\n LOG_ERROR(\"Init the searcher first before load index\");\n return IndexError_Runtime;\n }\n\n LOG_INFO(\"Begin HnswSearcher:load\");\n\n auto start_time = ailego::Monotime::MilliSeconds();\n\n int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to deserialize meta from container\");\n return ret;\n }\n\n ret = entity_.load(container, check_crc_enabled_);\n if (ret != 0) {\n LOG_ERROR(\"HnswSearcher load index failed\");\n return ret;\n }\n\n alg_ = HnswAlgorithm::UPointer(new HnswAlgorithm(entity_));\n\n if (metric) {\n metric_ = metric;\n } else {\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"CreateMetric failed, name: %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"IndexMetric init failed, ret=%d\", ret);\n return ret;\n }\n if (metric_->query_metric()) {\n metric_ = metric_->query_metric();\n }\n }\n\n if (!metric_->is_matched(meta_)) {\n LOG_ERROR(\"IndexMetric not match index meta\");\n return IndexError_Mismatch;\n }\n\n max_scan_num_ = static_cast(max_scan_ratio_ * entity_.doc_cnt());\n max_scan_num_ = std::max(4096U, max_scan_num_);\n\n stats_.set_loaded_count(entity_.doc_cnt());\n stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time);\n state_ = STATE_LOADED;\n magic_ = IndexContext::GenerateMagic();\n\n LOG_INFO(\"End HnswSearcher::load\");\n\n return 0;\n}\n\nint HnswSearcher::unload() {\n LOG_INFO(\"HnswSearcher unload index\");\n\n meta_.clear();\n entity_.cleanup();\n metric_.reset();\n max_scan_num_ = 0;\n stats_.set_loaded_count(0UL);\n stats_.set_loaded_costtime(0UL);\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint HnswSearcher::update_context(HnswContext *ctx) const {\n const HnswEntity::Pointer entity = entity_.clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone search context entity\");\n return IndexError_Runtime;\n }\n ctx->set_max_scan_num(max_scan_num_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n\n return ctx->update_context(HnswContext::kSearcherContext, meta_, metric_,\n entity, magic_);\n}\n\nint HnswSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const {\n if (ailego_unlikely(!query || !context)) {\n LOG_ERROR(\"The context is not created by this searcher\");\n return IndexError_Mismatch;\n }\n HnswContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswContext failed\");\n return IndexError_Cast;\n }\n\n if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) {\n return search_bf_impl(query, qmeta, count, context);\n }\n\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n int ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->resize_results(count);\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n int ret = alg_->search(ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher fast search failed\");\n return ret;\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nint HnswSearcher::search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const {\n if (ailego_unlikely(!query || !context)) {\n LOG_ERROR(\"The context is not created by this searcher\");\n return IndexError_Mismatch;\n }\n HnswContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n int ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->resize_results(count);\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n ctx->group_topk_heaps().clear();\n\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n dist_t dist = ctx->dist_calculator().batch_dist(id);\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n } else {\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n ctx->topk_heap().clear();\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n dist_t dist = ctx->dist_calculator().batch_dist(id);\n ctx->topk_heap().emplace(id, dist);\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nint HnswSearcher::search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n if (ailego_unlikely(!query || !context)) {\n LOG_ERROR(\"The context is not created by this searcher\");\n return IndexError_Mismatch;\n }\n\n if (ailego_unlikely(p_keys.size() != count)) {\n LOG_ERROR(\"The size of p_keys is not equal to count\");\n return IndexError_InvalidArgument;\n }\n\n HnswContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n int ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->resize_results(count);\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n ctx->group_topk_heaps().clear();\n\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n dist_t dist = ctx->dist_calculator().batch_dist(id);\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n } else {\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n ctx->topk_heap().clear();\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n dist_t dist = ctx->dist_calculator().batch_dist(id);\n ctx->topk_heap().emplace(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nIndexSearcher::Context::Pointer HnswSearcher::create_context() const {\n if (ailego_unlikely(state_ != STATE_LOADED)) {\n LOG_ERROR(\"Load the index first before create context\");\n return Context::Pointer();\n }\n const HnswEntity::Pointer search_ctx_entity = entity_.clone();\n if (!search_ctx_entity) {\n LOG_ERROR(\"Failed to create search context entity\");\n return Context::Pointer();\n }\n HnswContext *ctx = new (std::nothrow)\n HnswContext(meta_.dimension(), metric_, search_ctx_entity);\n if (ailego_unlikely(ctx == nullptr)) {\n LOG_ERROR(\"Failed to new HnswContext\");\n return Context::Pointer();\n }\n ctx->set_ef(ef_);\n ctx->set_max_scan_num(max_scan_num_);\n uint32_t filter_mode =\n bf_enabled_ ? VisitFilter::BloomFilter : VisitFilter::ByteMap;\n ctx->set_filter_mode(filter_mode);\n ctx->set_filter_negative_probability(bf_negative_probability_);\n ctx->set_magic(magic_);\n ctx->set_force_padding_topk(force_padding_topk_enabled_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n if (ailego_unlikely(ctx->init(HnswContext::kSearcherContext)) != 0) {\n LOG_ERROR(\"Init HnswContext failed\");\n delete ctx;\n return Context::Pointer();\n }\n\n return Context::Pointer(ctx);\n}\n\nIndexProvider::Pointer HnswSearcher::create_provider(void) const {\n LOG_DEBUG(\"HnswSearcher create provider\");\n\n auto entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"Clone HnswEntity failed\");\n return Provider::Pointer();\n }\n return Provider::Pointer(\n new (std::nothrow) HnswIndexProvider(meta_, entity, \"HnswSearcher\"));\n}\n\nconst void *HnswSearcher::get_vector(uint64_t key) const {\n return entity_.get_vector_by_key(key);\n}\n\nINDEX_FACTORY_REGISTER_SEARCHER(HnswSearcher);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_searcher.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"hnsw_searcher_entity.h\"\n#include \"hnsw_streamer.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSearcher : public IndexSearcher {\n public:\n using ContextPointer = IndexSearcher::Context::Pointer;\n\n public:\n HnswSearcher(void);\n ~HnswSearcher(void);\n\n HnswSearcher(const HnswSearcher &) = delete;\n HnswSearcher &operator=(const HnswSearcher &) = delete;\n\n protected:\n //! Initialize Searcher\n virtual int init(const ailego::Params ¶ms) override;\n\n //! Cleanup Searcher\n virtual int cleanup(void) override;\n\n //! Load Index from storage\n virtual int load(IndexStorage::Pointer container,\n IndexMetric::Pointer metric) override;\n\n //! Unload index from storage\n virtual int unload(void) override;\n\n //! KNN Search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n ContextPointer &context) const override {\n return search_impl(query, qmeta, 1, context);\n }\n\n //! KNN Search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n ContextPointer &context) const override;\n\n //! Linear Search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n ContextPointer &context) const override {\n return search_bf_impl(query, qmeta, 1, context);\n }\n\n //! Linear Search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n ContextPointer &context) const override;\n\n //! Linear search by primary keys\n virtual int search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, ContextPointer &context) const override {\n return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context);\n }\n\n //! Linear search by primary keys\n virtual int search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const override;\n\n //! Fetch vector by key\n virtual const void *get_vector(uint64_t key) const override;\n\n //! Create a searcher context\n virtual ContextPointer create_context() const override;\n\n //! Create a new iterator\n virtual IndexProvider::Pointer create_provider(void) const override;\n\n //! Retrieve statistics\n virtual const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve meta of index\n virtual const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n //! Retrieve params of index\n virtual const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n virtual void print_debug_info() override;\n\n private:\n //! To share ctx across streamer/searcher, we need to update the context for\n //! current streamer/searcher\n int update_context(HnswContext *ctx) const;\n\n private:\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 };\n\n HnswSearcherEntity entity_{};\n HnswAlgorithm::UPointer alg_; // impl graph algorithm\n\n IndexMetric::Pointer metric_{};\n IndexMeta meta_{};\n ailego::Params params_{};\n Stats stats_;\n uint32_t ef_{HnswEntity::kDefaultEf};\n uint32_t max_scan_num_{0U};\n uint32_t bruteforce_threshold_{HnswEntity::kDefaultBruteForceThreshold};\n float max_scan_ratio_{HnswEntity::kDefaultScanRatio};\n bool bf_enabled_{false};\n bool check_crc_enabled_{false};\n bool neighbors_in_memory_enabled_{false};\n bool force_padding_topk_enabled_{false};\n float bf_negative_probability_{HnswEntity::kDefaultBFNegativeProbability};\n uint32_t magic_{0U};\n\n State state_{STATE_INIT};\n};\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw/hnsw_searcher_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_searcher_entity.h\"\n#include \n#include \"utility/sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswSearcherEntity::HnswSearcherEntity() {}\n\nint HnswSearcherEntity::cleanup(void) {\n storage_.reset();\n vectors_.reset();\n keys_.reset();\n neighbors_.reset();\n neighbors_meta_.reset();\n neighbors_in_memory_enabled_ = false;\n loaded_ = false;\n\n this->HnswEntity::cleanup();\n\n return 0;\n}\n\nkey_t HnswSearcherEntity::get_key(node_id_t id) const {\n const void *key;\n if (ailego_unlikely(keys_->read(id * sizeof(key_t), &key, sizeof(key_t)) !=\n sizeof(key_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidKey;\n }\n return *(reinterpret_cast(key));\n}\n\n//! Get vector local id by key\nnode_id_t HnswSearcherEntity::get_id(key_t key) const {\n if (ailego_unlikely(!mapping_)) {\n LOG_ERROR(\"Index missing mapping segment\");\n return kInvalidNodeId;\n }\n\n //! Do binary search\n node_id_t start = 0UL;\n node_id_t end = doc_cnt();\n const void *data;\n node_id_t idx = 0u;\n while (start < end) {\n idx = start + (end - start) / 2;\n if (ailego_unlikely(\n mapping_->read(idx * sizeof(node_id_t), &data, sizeof(node_id_t)) !=\n sizeof(node_id_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidNodeId;\n }\n const key_t *mkey;\n node_id_t local_id = *reinterpret_cast(data);\n if (ailego_unlikely(keys_->read(local_id * sizeof(key_t),\n (const void **)(&mkey),\n sizeof(key_t)) != sizeof(key_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidNodeId;\n }\n if (*mkey < key) {\n start = idx + 1;\n } else if (*mkey > key) {\n end = idx;\n } else {\n return local_id;\n }\n }\n return kInvalidNodeId;\n}\n\nconst void *HnswSearcherEntity::get_vector_by_key(key_t key) const {\n node_id_t local_id = get_id(key);\n if (ailego_unlikely(local_id == kInvalidNodeId)) {\n return nullptr;\n }\n\n return get_vector(local_id);\n}\n\nconst void *HnswSearcherEntity::get_vector(node_id_t id) const {\n size_t read_size = vector_size();\n size_t offset = node_size() * id;\n\n const void *vec;\n if (ailego_unlikely(vectors_->read(offset, &vec, read_size) != read_size)) {\n LOG_ERROR(\"Read vector from segment failed\");\n return nullptr;\n }\n return vec;\n}\n\nint HnswSearcherEntity::get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const {\n const void *vec = get_vector(id);\n block.reset((void *)vec);\n return 0;\n}\n\nconst void *HnswSearcherEntity::get_vectors() const {\n const void *vec;\n size_t len = node_size() * doc_cnt();\n if (vectors_->read(0, &vec, len) != len) {\n LOG_ERROR(\"Read vectors from segment failed\");\n return nullptr;\n }\n return vec;\n}\n\nint HnswSearcherEntity::get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const {\n ailego_assert_with(count <= segment_datas_.size(), \"invalid count\");\n\n size_t read_size = vector_size();\n\n for (uint32_t i = 0; i < count; ++i) {\n segment_datas_[i].offset = node_size() * ids[i];\n segment_datas_[i].length = read_size;\n\n ailego_assert_with(segment_datas_[i].offset < vectors_->data_size(),\n \"invalid offset\");\n }\n if (ailego_unlikely(!vectors_->read(&segment_datas_[0], count))) {\n LOG_ERROR(\"Read vectors from segment failed\");\n return IndexError_ReadData;\n }\n for (uint32_t i = 0; i < count; ++i) {\n vecs[i] = segment_datas_[i].data;\n }\n\n return 0;\n}\n\nint HnswSearcherEntity::get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const {\n const void *vecs[count];\n get_vector(ids, count, vecs);\n for (uint32_t i = 0; i < count; ++i) {\n vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i]));\n }\n return 0;\n}\n\nconst Neighbors HnswSearcherEntity::get_neighbors(level_t level,\n node_id_t id) const {\n if (level == 0) {\n if (neighbors_in_memory_enabled_) {\n auto hd = reinterpret_cast(\n fixed_neighbors_.get() + neighbors_size() * id);\n return {hd->neighbor_cnt, hd->neighbors};\n }\n\n const GraphNeighborMeta *m;\n if (ailego_unlikely(neighbors_meta_->read(id * sizeof(GraphNeighborMeta),\n (const void **)(&m),\n sizeof(GraphNeighborMeta)) !=\n sizeof(GraphNeighborMeta))) {\n LOG_ERROR(\"Read neighbors meta from segment failed\");\n return {0, nullptr};\n }\n\n const void *data;\n if (ailego_unlikely(neighbors_->read(m->offset, &data,\n m->neighbor_cnt * sizeof(node_id_t)) !=\n m->neighbor_cnt * sizeof(node_id_t))) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return {0, nullptr};\n }\n return {static_cast(m->neighbor_cnt),\n reinterpret_cast(data)};\n }\n\n //! Read level > 0 neighbors\n const HnswNeighborMeta *m;\n if (ailego_unlikely(upper_neighbors_meta_->read(id * sizeof(HnswNeighborMeta),\n (const void **)(&m),\n sizeof(HnswNeighborMeta)) !=\n sizeof(HnswNeighborMeta))) {\n LOG_ERROR(\"Read neighbors meta from segment failed\");\n return {0, nullptr};\n }\n\n ailego_assert_with(level <= m->level, \"invalid level\");\n size_t offset = m->offset + (level - 1) * upper_neighbors_size();\n ailego_assert_with(offset <= upper_neighbors_->data_size(), \"invalid offset\");\n const void *data;\n if (ailego_unlikely(\n upper_neighbors_->read(offset, &data, upper_neighbors_size()) !=\n upper_neighbors_size())) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return {0, nullptr};\n }\n\n auto hd = reinterpret_cast(data);\n return {hd->neighbor_cnt, hd->neighbors};\n}\n\nint HnswSearcherEntity::load(const IndexStorage::Pointer &container,\n bool check_crc) {\n storage_ = container;\n\n int ret = load_segments(check_crc);\n if (ret != 0) {\n return ret;\n }\n\n loaded_ = true;\n\n LOG_INFO(\n \"Index info: docCnt=%u entryPoint=%u maxLevel=%d efConstruct=%zu \"\n \"l0NeighborCnt=%zu upperNeighborCnt=%zu scalingFactor=%zu \"\n \"vectorSize=%zu nodeSize=%zu vectorSegmentSize=%zu keySegmentSize=%zu \"\n \"neighborsSegmentSize=%zu neighborsMetaSegmentSize=%zu \",\n doc_cnt(), entry_point(), cur_max_level(), ef_construction(),\n l0_neighbor_cnt(), upper_neighbor_cnt(), scaling_factor(), vector_size(),\n node_size(), vectors_->data_size(), keys_->data_size(),\n neighbors_ == nullptr ? 0 : neighbors_->data_size(),\n neighbors_meta_ == nullptr ? 0 : neighbors_meta_->data_size());\n\n return 0;\n}\n\nint HnswSearcherEntity::load_segments(bool check_crc) {\n //! load header\n const void *data = nullptr;\n HNSWHeader hd;\n auto graph_hd_segment = storage_->get(kGraphHeaderSegmentId);\n if (!graph_hd_segment || graph_hd_segment->data_size() < sizeof(hd.graph)) {\n LOG_ERROR(\"Miss or invalid segment %s\", kGraphHeaderSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n if (graph_hd_segment->read(0, reinterpret_cast(&data),\n sizeof(hd.graph)) != sizeof(hd.graph)) {\n LOG_ERROR(\"Read segment %s failed\", kGraphHeaderSegmentId.c_str());\n return IndexError_ReadData;\n }\n memcpy(&hd.graph, data, sizeof(hd.graph));\n\n auto hnsw_hd_segment = storage_->get(kHnswHeaderSegmentId);\n if (!hnsw_hd_segment || hnsw_hd_segment->data_size() < sizeof(hd.hnsw)) {\n LOG_ERROR(\"Miss or invalid segment %s\", kHnswHeaderSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n if (hnsw_hd_segment->read(0, reinterpret_cast(&data),\n sizeof(hd.hnsw)) != sizeof(hd.hnsw)) {\n LOG_ERROR(\"Read segment %s failed\", kHnswHeaderSegmentId.c_str());\n return IndexError_ReadData;\n }\n memcpy(&hd.hnsw, data, sizeof(hd.hnsw));\n *mutable_header() = hd;\n segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt()));\n\n vectors_ = storage_->get(kGraphFeaturesSegmentId);\n if (!vectors_) {\n LOG_ERROR(\"IndexStorage get segment %s failed\",\n kGraphFeaturesSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n keys_ = storage_->get(kGraphKeysSegmentId);\n if (!keys_) {\n LOG_ERROR(\"IndexStorage get segment %s failed\",\n kGraphKeysSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n\n neighbors_ = storage_->get(kGraphNeighborsSegmentId);\n if (!neighbors_ || (neighbors_->data_size() == 0 && doc_cnt() > 1)) {\n LOG_ERROR(\"IndexStorage get segment %s failed or empty\",\n kGraphNeighborsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n neighbors_meta_ = storage_->get(kGraphOffsetsSegmentId);\n if (!neighbors_meta_ ||\n neighbors_meta_->data_size() < sizeof(GraphNeighborMeta) * doc_cnt()) {\n LOG_ERROR(\"IndexStorage get segment %s failed or invalid size\",\n kGraphOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n upper_neighbors_ = storage_->get(kHnswNeighborsSegmentId);\n if (!upper_neighbors_ ||\n (upper_neighbors_->data_size() == 0 && cur_max_level() > 0)) {\n LOG_ERROR(\"IndexStorage get segment %s failed or empty\",\n kHnswNeighborsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n upper_neighbors_meta_ = storage_->get(kHnswOffsetsSegmentId);\n if (!upper_neighbors_meta_ || upper_neighbors_meta_->data_size() <\n sizeof(HnswNeighborMeta) * doc_cnt()) {\n LOG_ERROR(\"IndexStorage get segment %s failed or invalid size\",\n kHnswOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n mapping_ = storage_->get(kGraphMappingSegmentId);\n if (!mapping_ || mapping_->data_size() < sizeof(node_id_t) * doc_cnt()) {\n LOG_ERROR(\"IndexStorage get segment %s failed or invalid size\",\n kGraphMappingSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (check_crc) {\n std::vector segments;\n segments.emplace_back(graph_hd_segment);\n segments.emplace_back(hnsw_hd_segment);\n segments.emplace_back(vectors_);\n segments.emplace_back(keys_);\n\n segments.emplace_back(neighbors_);\n segments.emplace_back(neighbors_meta_);\n segments.emplace_back(upper_neighbors_);\n segments.emplace_back(upper_neighbors_meta_);\n\n if (!do_crc_check(segments)) {\n LOG_ERROR(\"Check index crc failed, the index may broken\");\n return IndexError_Runtime;\n }\n }\n\n if (neighbors_in_memory_enabled_) {\n int ret = load_and_flat_neighbors();\n if (ret != 0) {\n return ret;\n }\n }\n\n return 0;\n}\n\nint HnswSearcherEntity::load_and_flat_neighbors() {\n fixed_neighbors_.reset(\n new (std::nothrow) char[neighbors_size() * doc_cnt()]{},\n std::default_delete());\n if (!fixed_neighbors_) {\n LOG_ERROR(\"Malloc memory failed\");\n return IndexError_NoMemory;\n }\n\n //! Get a new segemnt to release the buffer after loading neighbors\n auto neighbors_meta = storage_->get(kGraphOffsetsSegmentId);\n if (!neighbors_meta) {\n LOG_ERROR(\"IndexStorage get segment graph.offsets failed\");\n return IndexError_InvalidArgument;\n }\n\n const GraphNeighborMeta *neighbors_index = nullptr;\n if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index),\n neighbors_meta->data_size()) !=\n neighbors_meta->data_size()) {\n LOG_ERROR(\"Read segment %s data failed\", kGraphOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n const char *neighbor_data;\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n size_t rd_size = neighbors_index[id].neighbor_cnt * sizeof(node_id_t);\n if (ailego_unlikely(\n neighbors_->read(neighbors_index[id].offset,\n reinterpret_cast(&neighbor_data),\n rd_size) != rd_size)) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return IndexError_ReadData;\n }\n // copy level 0 neighbors to fixed size neighbors memory\n char *dst = fixed_neighbors_.get() + neighbors_size() * id;\n *reinterpret_cast(dst) = neighbors_index[id].neighbor_cnt;\n memcpy(dst + sizeof(uint32_t), neighbor_data, rd_size);\n }\n\n return 0;\n}\n\nint HnswSearcherEntity::get_fixed_neighbors(\n std::vector *fixed_neighbors) const {\n //! Get a new segemnt to release the buffer after loading neighbors\n auto neighbors_meta = storage_->get(kGraphOffsetsSegmentId);\n if (!neighbors_meta) {\n LOG_ERROR(\"IndexStorage get segment graph.offsets failed\");\n return IndexError_InvalidArgument;\n }\n\n const GraphNeighborMeta *neighbors_index = nullptr;\n size_t meta_size = neighbors_meta->data_size();\n if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index),\n meta_size) != meta_size) {\n LOG_ERROR(\"Read segment %s data failed\", kGraphOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n size_t fixed_neighbor_cnt = l0_neighbor_cnt();\n fixed_neighbors->resize((fixed_neighbor_cnt + 1) * doc_cnt(), kInvalidNodeId);\n\n size_t neighbors_cnt_offset = fixed_neighbor_cnt * doc_cnt();\n size_t total_neighbor_cnt = 0;\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n size_t cur_neighbor_cnt = neighbors_index[id].neighbor_cnt;\n if (cur_neighbor_cnt == 0) {\n (*fixed_neighbors)[neighbors_cnt_offset + id] = 0;\n continue;\n }\n size_t rd_size = cur_neighbor_cnt * sizeof(node_id_t);\n const uint32_t *neighbors;\n if (neighbors_->read(neighbors_index[id].offset,\n reinterpret_cast(&neighbors),\n rd_size) != rd_size) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return IndexError_ReadData;\n }\n\n // copy level 0 neighbors to fixed size neighbors memory\n auto it = fixed_neighbors->begin() + id * fixed_neighbor_cnt;\n std::copy(neighbors, neighbors + cur_neighbor_cnt, it);\n\n (*fixed_neighbors)[neighbors_cnt_offset + id] = cur_neighbor_cnt;\n total_neighbor_cnt += cur_neighbor_cnt;\n }\n LOG_INFO(\"total neighbor cnt: %zu, average neighbor cnt: %zu\",\n total_neighbor_cnt, total_neighbor_cnt / doc_cnt());\n\n return 0;\n}\n\nbool HnswSearcherEntity::do_crc_check(\n std::vector &segments) const {\n constexpr size_t blk_size = 4096;\n const void *data;\n for (auto &segment : segments) {\n size_t offset = 0;\n size_t rd_size;\n uint32_t crc = 0;\n while (offset < segment->data_size()) {\n size_t size = std::min(blk_size, segment->data_size() - offset);\n if ((rd_size = segment->read(offset, &data, size)) <= 0) {\n break;\n }\n offset += rd_size;\n crc = ailego::Crc32c::Hash(data, rd_size, crc);\n }\n if (crc != segment->data_crc()) {\n return false;\n }\n }\n return true;\n}\n\nconst HnswEntity::Pointer HnswSearcherEntity::clone() const {\n auto vectors = vectors_->clone();\n if (ailego_unlikely(!vectors)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphFeaturesSegmentId.c_str());\n return HnswEntity::Pointer();\n }\n auto keys = keys_->clone();\n if (ailego_unlikely(!keys)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphKeysSegmentId.c_str());\n return HnswEntity::Pointer();\n }\n\n auto mapping = mapping_->clone();\n if (ailego_unlikely(!mapping)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphMappingSegmentId.c_str());\n return HnswEntity::Pointer();\n }\n\n auto neighbors = neighbors_->clone();\n if (ailego_unlikely(!neighbors)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphNeighborsSegmentId.c_str());\n return HnswEntity::Pointer();\n }\n auto upper_neighbors = upper_neighbors_->clone();\n if (ailego_unlikely(!neighbors)) {\n LOG_ERROR(\"clone segment %s failed\", kHnswNeighborsSegmentId.c_str());\n return HnswEntity::Pointer();\n }\n auto neighbors_meta = neighbors_meta_->clone();\n if (ailego_unlikely(!neighbors_meta)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphOffsetsSegmentId.c_str());\n return HnswEntity::Pointer();\n }\n auto upper_neighbors_meta = upper_neighbors_meta_->clone();\n if (ailego_unlikely(!upper_neighbors_meta)) {\n LOG_ERROR(\"clone segment %s failed\", kHnswOffsetsSegmentId.c_str());\n return HnswEntity::Pointer();\n }\n\n SegmentGroupParam neighbor_group{neighbors, neighbors_meta, upper_neighbors,\n upper_neighbors_meta};\n\n HnswSearcherEntity *entity = new (std::nothrow)\n HnswSearcherEntity(header(), vectors, keys, mapping, neighbor_group,\n fixed_neighbors_, neighbors_in_memory_enabled_);\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"HnswSearcherEntity new failed\");\n }\n\n return HnswEntity::Pointer(entity);\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw/hnsw_searcher_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"hnsw_builder_entity.h\"\n#include \"hnsw_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSearcherEntity : public HnswEntity {\n public:\n using Pointer = std::shared_ptr;\n using SegmentPointer = IndexStorage::Segment::Pointer;\n\n public:\n struct SegmentGroupParam {\n SegmentGroupParam(SegmentPointer neighbors_in,\n SegmentPointer neighbors_meta_in,\n SegmentPointer upper_neighbors_in,\n SegmentPointer upper_neighbors_meta_in)\n : neighbors{neighbors_in},\n neighbors_meta{neighbors_meta_in},\n upper_neighbors{upper_neighbors_in},\n upper_neighbors_meta{upper_neighbors_meta_in} {}\n\n SegmentPointer neighbors{nullptr};\n SegmentPointer neighbors_meta{nullptr};\n SegmentPointer upper_neighbors{nullptr};\n SegmentPointer upper_neighbors_meta{nullptr};\n };\n\n //! Constructor\n HnswSearcherEntity();\n\n //! Make a copy of searcher entity, to support thread-safe operation.\n //! The segment in container cannot be read concurrenly\n virtual const HnswEntity::Pointer clone() const override;\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const override;\n\n //! Get vector local id by key\n node_id_t get_id(key_t key) const;\n\n //! Get vector feature data by key\n virtual const void *get_vector_by_key(key_t key) const override;\n\n //! Get vector feature data by id\n virtual const void *get_vector(node_id_t id) const override;\n\n //! Get vector feature data by id\n virtual int get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const override;\n\n virtual int get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n virtual int get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const override;\n\n //! Get all vectors\n const void *get_vectors() const;\n\n //! Get the node id's neighbors on graph level\n virtual const Neighbors get_neighbors(level_t level,\n node_id_t id) const override;\n\n virtual int load(const IndexStorage::Pointer &container,\n bool check_crc) override;\n\n int load_segments(bool check_crc);\n\n virtual int cleanup(void) override;\n\n public:\n bool is_loaded() const {\n return loaded_;\n }\n\n void set_neighbors_in_memory(bool enabled) {\n neighbors_in_memory_enabled_ = enabled;\n }\n\n //! get fixed length neighbors data\n int get_fixed_neighbors(std::vector *fixed_neighbors) const;\n\n private:\n //! Constructor\n HnswSearcherEntity(const HNSWHeader &hd, const SegmentPointer &vectors,\n const SegmentPointer &keys, const SegmentPointer &mapping,\n const SegmentGroupParam &neighbor_group,\n const std::shared_ptr &fixed_neighbors,\n bool neighbors_in_memory_enabled)\n : HnswEntity(hd),\n vectors_(vectors),\n keys_(keys),\n mapping_(mapping),\n neighbors_(neighbor_group.neighbors),\n neighbors_meta_(neighbor_group.neighbors_meta),\n upper_neighbors_(neighbor_group.upper_neighbors),\n upper_neighbors_meta_(neighbor_group.upper_neighbors_meta),\n neighbors_in_memory_enabled_(neighbors_in_memory_enabled) {\n segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt()),\n IndexStorage::SegmentData(0U, 0U));\n fixed_neighbors_ = fixed_neighbors;\n }\n\n bool do_crc_check(std::vector &segments) const;\n\n inline size_t neighbors_size() const {\n return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t);\n }\n\n inline size_t upper_neighbors_size() const {\n return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t);\n }\n\n //! If neighbors_in_memory_enabled, load the level0 neighbors to memory\n int load_and_flat_neighbors(void);\n\n public:\n HnswSearcherEntity(const HnswSearcherEntity &) = delete;\n HnswSearcherEntity &operator=(const HnswSearcherEntity &) = delete;\n\n private:\n IndexStorage::Pointer storage_{};\n\n SegmentPointer vectors_{};\n SegmentPointer keys_{};\n SegmentPointer mapping_{};\n\n SegmentPointer neighbors_{};\n SegmentPointer neighbors_meta_{};\n SegmentPointer upper_neighbors_{};\n SegmentPointer upper_neighbors_meta_{};\n\n mutable std::vector segment_datas_{};\n std::shared_ptr fixed_neighbors_{}; // level 0 fixed size neighbors\n bool neighbors_in_memory_enabled_{false};\n bool loaded_{false};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_streamer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_streamer.h\"\n#include \n#include \n#include \n#include \n#include \"utility/sparse_utility.h\"\n#include \"hnsw_algorithm.h\"\n#include \"hnsw_context.h\"\n#include \"hnsw_dist_calculator.h\"\n#include \"hnsw_index_provider.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswStreamer::HnswStreamer() : entity_(stats_) {}\n\nHnswStreamer::~HnswStreamer() {\n if (state_ == STATE_INITED) {\n this->cleanup();\n }\n}\n\nint HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) {\n meta_ = imeta;\n meta_.set_streamer(\"HnswStreamer\", HnswEntity::kRevision, params);\n\n params.get(PARAM_HNSW_STREAMER_MAX_INDEX_SIZE, &max_index_size_);\n\n params.get(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, &upper_max_neighbor_cnt_);\n float multiplier = HnswEntity::kDefaultL0MaxNeighborCntMultiplier;\n params.get(PARAM_HNSW_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER, &multiplier);\n l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_;\n\n multiplier = HnswEntity::kDefaultNeighborPruneMultiplier;\n params.get(PARAM_HNSW_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier);\n size_t prune_cnt = multiplier * upper_max_neighbor_cnt_;\n scaling_factor_ = upper_max_neighbor_cnt_;\n params.get(PARAM_HNSW_STREAMER_SCALING_FACTOR, &scaling_factor_);\n\n params.get(PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT, &docs_hard_limit_);\n params.get(PARAM_HNSW_STREAMER_EF, &ef_);\n params.get(PARAM_HNSW_STREAMER_EFCONSTRUCTION, &ef_construction_);\n params.get(PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_ENABLE, &bf_enabled_);\n params.get(PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB,\n &bf_negative_prob_);\n params.get(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, &bruteforce_threshold_);\n params.get(PARAM_HNSW_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_);\n params.get(PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_);\n params.get(PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_);\n params.get(PARAM_HNSW_STREAMER_CHECK_CRC_ENABLE, &check_crc_enabled_);\n params.get(PARAM_HNSW_STREAMER_CHUNK_SIZE, &chunk_size_);\n params.get(PARAM_HNSW_STREAMER_FILTER_SAME_KEY, &filter_same_key_);\n params.get(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, &get_vector_enabled_);\n params.get(PARAM_HNSW_STREAMER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_);\n params.get(PARAM_HNSW_STREAMER_FORCE_PADDING_RESULT_ENABLE,\n &force_padding_topk_enabled_);\n params.get(PARAM_HNSW_STREAMER_USE_ID_MAP, &use_id_map_);\n entity_.set_use_key_info_map(use_id_map_);\n\n params.get(PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit_);\n if (docs_soft_limit_ > 0 && docs_soft_limit_ > docs_hard_limit_) {\n LOG_ERROR(\"[%s] must be >= [%s]\",\n PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT.c_str(),\n PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT.c_str());\n return IndexError_InvalidArgument;\n } else if (docs_soft_limit_ == 0UL) {\n docs_soft_limit_ =\n docs_hard_limit_ * HnswEntity::kDefaultDocsSoftLimitRatio;\n }\n\n if (ef_ == 0U) {\n ef_ = HnswEntity::kDefaultEf;\n }\n if (ef_construction_ == 0U) {\n ef_construction_ = HnswEntity::kDefaultEfConstruction;\n }\n if (upper_max_neighbor_cnt_ == 0U) {\n upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt;\n }\n if (upper_max_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) {\n LOG_ERROR(\"[%s] must be in range (0,%d)\",\n PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT.c_str(),\n HnswEntity::kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (l0_max_neighbor_cnt_ == 0U) {\n l0_max_neighbor_cnt_ = HnswEntity::kDefaultL0MaxNeighborCnt;\n }\n if (l0_max_neighbor_cnt_ > HnswEntity::kMaxNeighborCnt) {\n LOG_ERROR(\"MaxL0NeighborCnt must be in range (0,%d)\",\n HnswEntity::kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) {\n LOG_ERROR(\"[%s]-[%u] must be <= [%s]-[%u]\",\n PARAM_HNSW_STREAMER_MIN_NEIGHBOR_COUNT.c_str(), min_neighbor_cnt_,\n PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT.c_str(),\n upper_max_neighbor_cnt_);\n return IndexError_InvalidArgument;\n }\n\n if (bf_negative_prob_ <= 0.0f || bf_negative_prob_ >= 1.0f) {\n LOG_ERROR(\"[%s] must be in range (0,1)\",\n PARAM_HNSW_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (scaling_factor_ == 0U) {\n scaling_factor_ = HnswEntity::kDefaultScalingFactor;\n }\n if (scaling_factor_ < 5 || scaling_factor_ > 1000) {\n LOG_ERROR(\"[%s] must be in range [5,1000]\",\n PARAM_HNSW_STREAMER_SCALING_FACTOR.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) {\n LOG_ERROR(\"[%s] must be in range (0.0f,1.0f]\",\n PARAM_HNSW_STREAMER_MAX_SCAN_RATIO.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (max_scan_limit_ < min_scan_limit_) {\n LOG_ERROR(\"[%s] must be >= [%s]\",\n PARAM_HNSW_STREAMER_MAX_SCAN_LIMIT.c_str(),\n PARAM_HNSW_STREAMER_MIN_SCAN_LIMIT.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (prune_cnt == 0UL) {\n prune_cnt = upper_max_neighbor_cnt_;\n }\n if (chunk_size_ == 0UL) {\n chunk_size_ = HnswEntity::kDefaultChunkSize;\n }\n if (chunk_size_ > HnswEntity::kMaxChunkSize) {\n LOG_ERROR(\"[%s] must be < %zu\", PARAM_HNSW_STREAMER_CHUNK_SIZE.c_str(),\n HnswEntity::kMaxChunkSize);\n return IndexError_InvalidArgument;\n }\n\n entity_.set_ef_construction(ef_construction_);\n entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_);\n entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_);\n entity_.set_scaling_factor(scaling_factor_);\n entity_.set_prune_cnt(prune_cnt);\n\n entity_.set_vector_size(meta_.element_size());\n\n entity_.set_chunk_size(chunk_size_);\n entity_.set_filter_same_key(filter_same_key_);\n entity_.set_get_vector(get_vector_enabled_);\n entity_.set_min_neighbor_cnt(min_neighbor_cnt_);\n\n int ret = entity_.init(docs_hard_limit_);\n if (ret != 0) {\n LOG_ERROR(\"Hnsw entity init failed for %s\", IndexError::What(ret));\n return ret;\n }\n\n LOG_DEBUG(\n \"Init params: maxIndexSize=%zu docsHardLimit=%zu docsSoftLimit=%zu \"\n \"efConstruction=%u ef=%u upperMaxNeighborCnt=%u l0MaxNeighborCnt=%u \"\n \"scalingFactor=%u maxScanRatio=%.3f minScanLimit=%zu maxScanLimit=%zu \"\n \"bfEnabled=%d bruteFoceThreshold=%zu bfNegativeProbability=%.5f \"\n \"checkCrcEnabled=%d pruneSize=%zu vectorSize=%u chunkSize=%zu \"\n \"filterSameKey=%u getVectorEnabled=%u minNeighborCount=%u \"\n \"forcePadding=%u \",\n max_index_size_, docs_hard_limit_, docs_soft_limit_, ef_construction_,\n ef_, upper_max_neighbor_cnt_, l0_max_neighbor_cnt_, scaling_factor_,\n max_scan_ratio_, min_scan_limit_, max_scan_limit_, bf_enabled_,\n bruteforce_threshold_, bf_negative_prob_, check_crc_enabled_, prune_cnt,\n meta_.element_size(), chunk_size_, filter_same_key_, get_vector_enabled_,\n min_neighbor_cnt_, force_padding_topk_enabled_);\n\n alg_ = HnswAlgorithm::UPointer(new HnswAlgorithm(entity_));\n\n ret = alg_->init();\n if (ret != 0) {\n return ret;\n }\n\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint HnswStreamer::cleanup(void) {\n if (state_ == STATE_OPENED) {\n this->close();\n }\n\n LOG_INFO(\"HnswStreamer cleanup\");\n\n meta_.clear();\n metric_.reset();\n stats_.clear();\n entity_.cleanup();\n\n if (alg_) {\n alg_->cleanup();\n }\n\n max_index_size_ = 0UL;\n docs_hard_limit_ = HnswEntity::kDefaultDocsHardLimit;\n docs_soft_limit_ = 0UL;\n upper_max_neighbor_cnt_ = HnswEntity::kDefaultUpperMaxNeighborCnt;\n l0_max_neighbor_cnt_ = HnswEntity::kDefaultL0MaxNeighborCnt;\n ef_ = HnswEntity::kDefaultEf;\n ef_construction_ = HnswEntity::kDefaultEfConstruction;\n bf_enabled_ = false;\n scaling_factor_ = HnswEntity::kDefaultScalingFactor;\n bruteforce_threshold_ = HnswEntity::kDefaultBruteForceThreshold;\n max_scan_limit_ = HnswEntity::kDefaultMaxScanLimit;\n min_scan_limit_ = HnswEntity::kDefaultMinScanLimit;\n chunk_size_ = HnswEntity::kDefaultChunkSize;\n bf_negative_prob_ = HnswEntity::kDefaultBFNegativeProbability;\n max_scan_ratio_ = HnswEntity::kDefaultScanRatio;\n state_ = STATE_INIT;\n check_crc_enabled_ = false;\n filter_same_key_ = false;\n get_vector_enabled_ = false;\n\n return 0;\n}\n\nint HnswStreamer::open(IndexStorage::Pointer stg) {\n LOG_INFO(\"HnswStreamer open\");\n\n if (ailego_unlikely(state_ != STATE_INITED)) {\n LOG_ERROR(\"Open storage failed, init streamer first!\");\n return IndexError_NoReady;\n }\n int ret = entity_.open(std::move(stg), max_index_size_, check_crc_enabled_);\n if (ret != 0) {\n return ret;\n }\n IndexMeta index_meta;\n ret = entity_.get_index_meta(&index_meta);\n if (ret == IndexError_NoExist) {\n // Set IndexMeta for the new index\n ret = entity_.set_index_meta(meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to set index meta for %s\", IndexError::What(ret));\n return ret;\n }\n } else if (ret != 0) {\n LOG_ERROR(\"Failed to get index meta for %s\", IndexError::What(ret));\n return ret;\n } else {\n if (index_meta.dimension() != meta_.dimension() ||\n index_meta.element_size() != meta_.element_size() ||\n index_meta.metric_name() != meta_.metric_name() ||\n index_meta.data_type() != meta_.data_type()) {\n LOG_ERROR(\"IndexMeta mismatch from the previous in index\");\n return IndexError_Mismatch;\n }\n // The IndexMetric Params may be updated like MipsSquaredEuclidean\n auto metric_params = index_meta.metric_params();\n metric_params.merge(meta_.metric_params());\n meta_.set_metric(index_meta.metric_name(), 0, metric_params);\n }\n\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"Failed to create metric %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"Failed to init metric, ret=%d\", ret);\n return ret;\n }\n\n if (!metric_->distance()) {\n LOG_ERROR(\"Invalid metric distance\");\n return IndexError_InvalidArgument;\n }\n\n if (!metric_->batch_distance()) {\n LOG_ERROR(\"Invalid metric batch distance\");\n return IndexError_InvalidArgument;\n }\n\n add_distance_ = metric_->distance();\n add_batch_distance_ = metric_->batch_distance();\n\n search_distance_ = add_distance_;\n search_batch_distance_ = add_batch_distance_;\n\n if (metric_->query_metric() && metric_->query_metric()->distance() &&\n metric_->query_metric()->batch_distance()) {\n search_distance_ = metric_->query_metric()->distance();\n search_batch_distance_ = metric_->query_metric()->batch_distance();\n }\n\n state_ = STATE_OPENED;\n magic_ = IndexContext::GenerateMagic();\n\n return 0;\n}\n\nint HnswStreamer::close(void) {\n LOG_INFO(\"HnswStreamer close\");\n\n stats_.clear();\n meta_.set_metric(metric_->name(), 0, metric_->params());\n entity_.set_index_meta(meta_);\n int ret = entity_.close();\n if (ret != 0) {\n return ret;\n }\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint HnswStreamer::flush(uint64_t checkpoint) {\n LOG_INFO(\"HnswStreamer flush checkpoint=%zu\", (size_t)checkpoint);\n\n meta_.set_metric(metric_->name(), 0, metric_->params());\n entity_.set_index_meta(meta_);\n return entity_.flush(checkpoint);\n}\n\nint HnswStreamer::dump(const IndexDumper::Pointer &dumper) {\n LOG_INFO(\"HnswStreamer dump\");\n\n shared_mutex_.lock();\n AILEGO_DEFER([&]() { shared_mutex_.unlock(); });\n\n meta_.set_searcher(\"HnswSearcher\", HnswEntity::kRevision, ailego::Params());\n\n int ret = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return ret;\n }\n return entity_.dump(dumper);\n}\n\nIndexStreamer::Context::Pointer HnswStreamer::create_context(void) const {\n if (ailego_unlikely(state_ != STATE_OPENED)) {\n LOG_ERROR(\"Create context failed, open storage first!\");\n return Context::Pointer();\n }\n\n HnswEntity::Pointer entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"CreateContext clone init failed\");\n return Context::Pointer();\n }\n HnswContext *ctx =\n new (std::nothrow) HnswContext(meta_.dimension(), metric_, entity);\n if (ailego_unlikely(ctx == nullptr)) {\n LOG_ERROR(\"Failed to new HnswContext\");\n return Context::Pointer();\n }\n ctx->set_ef(ef_);\n ctx->set_max_scan_limit(max_scan_limit_);\n ctx->set_min_scan_limit(min_scan_limit_);\n ctx->set_max_scan_ratio(max_scan_ratio_);\n ctx->set_filter_mode(bf_enabled_ ? VisitFilter::BloomFilter\n : VisitFilter::ByteMap);\n ctx->set_filter_negative_probability(bf_negative_prob_);\n ctx->set_magic(magic_);\n ctx->set_force_padding_topk(force_padding_topk_enabled_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n\n if (ailego_unlikely(ctx->init(HnswContext::kStreamerContext)) != 0) {\n LOG_ERROR(\"Init HnswContext failed\");\n delete ctx;\n return Context::Pointer();\n }\n uint32_t estimate_doc_count = 0;\n if (meta_.streamer_params().get(PARAM_HNSW_STREAMER_ESTIMATE_DOC_COUNT,\n &estimate_doc_count)) {\n LOG_DEBUG(\"HnswStreamer doc_count[%zu] estimate[%zu]\",\n (size_t)entity_.doc_cnt(), (size_t)estimate_doc_count);\n }\n ctx->check_need_adjuct_ctx(std::max(entity_.doc_cnt(), estimate_doc_count));\n\n return Context::Pointer(ctx);\n}\n\nIndexProvider::Pointer HnswStreamer::create_provider(void) const {\n LOG_DEBUG(\"HnswStreamer create provider\");\n\n auto entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"Clone HnswEntity failed\");\n return nullptr;\n }\n return Provider::Pointer(\n new HnswIndexProvider(meta_, entity, \"HnswStreamer\"));\n}\n\nint HnswStreamer::update_context(HnswContext *ctx) const {\n const HnswEntity::Pointer entity = entity_.clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone search context entity\");\n return IndexError_Runtime;\n }\n ctx->set_max_scan_limit(max_scan_limit_);\n ctx->set_min_scan_limit(min_scan_limit_);\n ctx->set_max_scan_ratio(max_scan_ratio_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n return ctx->update_context(HnswContext::kStreamerContext, meta_, metric_,\n entity, magic_);\n}\n\n//! Add a vector with id into index\nint HnswStreamer::add_with_id_impl(uint32_t id, const void *query,\n const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) {\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n HnswContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) {\n if (entity_.doc_cnt() >= docs_hard_limit_) {\n LOG_ERROR(\"Current docs %u exceed [%s]\", entity_.doc_cnt(),\n PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT.c_str());\n const std::lock_guard lk(mutex_);\n (*stats_.mutable_discarded_count())++;\n return IndexError_IndexFull;\n } else {\n LOG_WARN(\"Current docs %u exceed [%s]\", entity_.doc_cnt(),\n PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT.c_str());\n }\n }\n if (ailego_unlikely(!shared_mutex_.try_lock_shared())) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); });\n\n ctx->clear();\n ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_);\n ctx->reset_query(query);\n ctx->check_need_adjuct_ctx(entity_.doc_cnt());\n\n if (metric_->support_train()) {\n const std::lock_guard lk(mutex_);\n ret = metric_->train(query, meta_.dimension());\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer metric train failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n }\n\n level_t level = alg_->get_random_level();\n ret = entity_.add_vector_with_id(level, id, query);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer add vector failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n ret = alg_->add_node(id, level, ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw steamer add node failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n if (ailego_unlikely(ctx->error())) {\n (*stats_.mutable_discarded_count())++;\n return IndexError_Runtime;\n }\n (*stats_.mutable_added_count())++;\n\n return 0;\n}\n\n//! Add a vector into index\nint HnswStreamer::add_impl(uint64_t pkey, const void *query,\n const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) {\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n HnswContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) {\n if (entity_.doc_cnt() >= docs_hard_limit_) {\n LOG_ERROR(\"Current docs %u exceed [%s]\", entity_.doc_cnt(),\n PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT.c_str());\n const std::lock_guard lk(mutex_);\n (*stats_.mutable_discarded_count())++;\n return IndexError_IndexFull;\n } else {\n LOG_WARN(\"Current docs %u exceed [%s]\", entity_.doc_cnt(),\n PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT.c_str());\n }\n }\n if (ailego_unlikely(!shared_mutex_.try_lock_shared())) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); });\n\n ctx->clear();\n ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_);\n ctx->reset_query(query);\n ctx->check_need_adjuct_ctx(entity_.doc_cnt());\n\n if (metric_->support_train()) {\n const std::lock_guard lk(mutex_);\n ret = metric_->train(query, meta_.dimension());\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer metric train failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n }\n\n level_t level = alg_->get_random_level();\n node_id_t id;\n ret = entity_.add_vector(level, pkey, query, &id);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer add vector failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n ret = alg_->add_node(id, level, ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw steamer add node failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n if (ailego_unlikely(ctx->error())) {\n (*stats_.mutable_discarded_count())++;\n return IndexError_Runtime;\n }\n (*stats_.mutable_added_count())++;\n\n return 0;\n}\n\n\nint HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) const {\n return search_impl(query, qmeta, 1, context);\n}\n\n//! Similarity search\nint HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n IndexStreamer::Context::Pointer &context) const {\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n HnswContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswContext failed\");\n return IndexError_Cast;\n }\n\n if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) {\n return search_bf_impl(query, qmeta, count, context);\n }\n\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_);\n ctx->resize_results(count);\n ctx->check_need_adjuct_ctx(entity_.doc_cnt());\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n ret = alg_->search(ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher fast search failed\");\n return ret;\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nvoid HnswStreamer::print_debug_info() {\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n Neighbors neighbours = entity_.get_neighbors(0, id);\n std::cout << \"node: \" << id << \"; \";\n if (neighbours.size() == 0) std::cout << std::endl;\n for (uint32_t i = 0; i < neighbours.size(); ++i) {\n std::cout << neighbours[i];\n\n if (i == neighbours.size() - 1) {\n std::cout << std::endl;\n } else {\n std::cout << \", \";\n }\n }\n }\n\n // entity_.print_key_map();\n}\n\nint HnswStreamer::search_bf_impl(\n const void *query, const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) const {\n return search_bf_impl(query, qmeta, 1, context);\n}\n\nint HnswStreamer::search_bf_impl(\n const void *query, const IndexQueryMeta &qmeta, uint32_t count,\n IndexStreamer::Context::Pointer &context) const {\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n HnswContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_);\n ctx->resize_results(count);\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n ctx->group_topk_heaps().clear();\n\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n dist_t dist = ctx->dist_calculator().batch_dist(id);\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n } else {\n auto &filter = ctx->filter();\n auto &topk = ctx->topk_heap();\n\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n topk.clear();\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n\n if (!filter.is_valid() || !filter(entity_.get_key(id))) {\n dist_t dist = ctx->dist_calculator().batch_dist(id);\n topk.emplace(id, dist);\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nint HnswStreamer::search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n if (ailego_unlikely(p_keys.size() != count)) {\n LOG_ERROR(\"The size of p_keys is not equal to count\");\n return IndexError_InvalidArgument;\n }\n\n HnswContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_);\n ctx->resize_results(count);\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n ctx->group_topk_heaps().clear();\n\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n dist_t dist = ctx->dist_calculator().batch_dist(id);\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n } else {\n auto &filter = ctx->filter();\n auto &topk = ctx->topk_heap();\n\n for (size_t q = 0; q < count; ++q) {\n ctx->reset_query(query);\n topk.clear();\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n key_t pk = p_keys[q][idx];\n if (!filter.is_valid() || !filter(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n dist_t dist = ctx->dist_calculator().batch_dist(id);\n topk.emplace(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\n\nINDEX_FACTORY_REGISTER_STREAMER(HnswStreamer);\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw/hnsw_streamer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"hnsw_algorithm.h\"\n#include \"hnsw_streamer_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswStreamer : public IndexStreamer {\n public:\n using ContextPointer = IndexStreamer::Context::Pointer;\n\n HnswStreamer(void);\n virtual ~HnswStreamer(void);\n\n HnswStreamer(const HnswStreamer &streamer) = delete;\n HnswStreamer &operator=(const HnswStreamer &streamer) = delete;\n\n protected:\n //! Initialize Streamer\n virtual int init(const IndexMeta &imeta,\n const ailego::Params ¶ms) override;\n\n //! Cleanup Streamer\n virtual int cleanup(void) override;\n\n //! Create a context\n virtual Context::Pointer create_context(void) const override;\n\n //! Create a new iterator\n virtual IndexProvider::Pointer create_provider(void) const override;\n\n //! Add a vector into index\n virtual int add_impl(uint64_t pkey, const void *query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) override;\n\n //! Add a vector with id into index\n virtual int add_with_id_impl(uint32_t id, const void *query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) override;\n\n //! Similarity search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const override;\n\n //! Similarity brute force search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity brute force search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const override;\n\n //! Linear search by primary keys\n virtual int search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, ContextPointer &context) const override {\n return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context);\n }\n\n //! Linear search by primary keys\n virtual int search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const override;\n\n //! Fetch vector by key\n virtual const void *get_vector(uint64_t key) const override {\n return entity_.get_vector_by_key(key);\n }\n\n virtual int get_vector(const uint64_t key,\n IndexStorage::MemoryBlock &block) const override {\n return entity_.get_vector_by_key(key, block);\n }\n\n //! Fetch vector by id\n virtual const void *get_vector_by_id(uint32_t id) const override {\n return entity_.get_vector(id);\n }\n\n virtual int get_vector_by_id(\n const uint32_t id, IndexStorage::MemoryBlock &block) const override {\n return entity_.get_vector(id, block);\n }\n\n //! Open index from file path\n virtual int open(IndexStorage::Pointer stg) override;\n\n //! Close file\n virtual int close(void) override;\n\n //! flush file\n virtual int flush(uint64_t checkpoint) override;\n\n //! Dump index into storage\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n virtual const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve meta of index\n virtual const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n virtual void print_debug_info() override;\n\n private:\n inline int check_params(const void *query,\n const IndexQueryMeta &qmeta) const {\n if (ailego_unlikely(!query)) {\n LOG_ERROR(\"null query\");\n return IndexError_InvalidArgument;\n }\n if (ailego_unlikely(qmeta.dimension() != meta_.dimension() ||\n qmeta.data_type() != meta_.data_type() ||\n qmeta.element_size() != meta_.element_size())) {\n LOG_ERROR(\"Unsupported query meta\");\n return IndexError_Mismatch;\n }\n return 0;\n }\n\n inline int check_sparse_count_is_zero(const uint32_t *sparse_count,\n uint32_t count) const {\n for (uint32_t i = 0; i < count; ++i) {\n if (sparse_count[i] != 0)\n LOG_ERROR(\"Sparse cout is not empty. Index: %u, Sparse Count: %u\", i,\n sparse_count[i]);\n return IndexError_InvalidArgument;\n }\n\n return 0;\n }\n\n private:\n //! To share ctx across streamer/searcher, we need to update the context for\n //! current streamer/searcher\n int update_context(HnswContext *ctx) const;\n\n private:\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 };\n class Stats : public IndexStreamer::Stats {\n public:\n void clear(void) {\n set_revision_id(0u);\n set_loaded_count(0u);\n set_added_count(0u);\n set_discarded_count(0u);\n set_index_size(0u);\n set_dumped_size(0u);\n set_check_point(0u);\n set_create_time(0u);\n set_update_time(0u);\n clear_attributes();\n }\n };\n\n HnswStreamerEntity entity_;\n HnswAlgorithm::UPointer alg_;\n IndexMeta meta_{};\n IndexMetric::Pointer metric_{};\n\n IndexMetric::MatrixDistance add_distance_{};\n IndexMetric::MatrixDistance search_distance_{};\n\n IndexMetric::MatrixBatchDistance add_batch_distance_{};\n IndexMetric::MatrixBatchDistance search_batch_distance_{};\n\n Stats stats_{};\n std::mutex mutex_{};\n\n size_t max_index_size_{0UL};\n size_t chunk_size_{HnswEntity::kDefaultChunkSize};\n size_t docs_hard_limit_{HnswEntity::kDefaultDocsHardLimit};\n size_t docs_soft_limit_{0UL};\n uint32_t min_neighbor_cnt_{0u};\n uint32_t upper_max_neighbor_cnt_{HnswEntity::kDefaultUpperMaxNeighborCnt};\n uint32_t l0_max_neighbor_cnt_{HnswEntity::kDefaultL0MaxNeighborCnt};\n uint32_t ef_{HnswEntity::kDefaultEf};\n uint32_t ef_construction_{HnswEntity::kDefaultEfConstruction};\n uint32_t scaling_factor_{HnswEntity::kDefaultScalingFactor};\n size_t bruteforce_threshold_{HnswEntity::kDefaultBruteForceThreshold};\n size_t max_scan_limit_{HnswEntity::kDefaultMaxScanLimit};\n size_t min_scan_limit_{HnswEntity::kDefaultMinScanLimit};\n float bf_negative_prob_{HnswEntity::kDefaultBFNegativeProbability};\n float max_scan_ratio_{HnswEntity::kDefaultScanRatio};\n\n uint32_t magic_{0U};\n State state_{STATE_INIT};\n bool bf_enabled_{false};\n bool check_crc_enabled_{false};\n bool filter_same_key_{false};\n bool get_vector_enabled_{false};\n bool force_padding_topk_enabled_{false};\n bool use_id_map_{true};\n\n //! avoid add vector while dumping index\n ailego::SharedMutex shared_mutex_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_streamer_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"hnsw_streamer_entity.h\"\n#include \n\n// #define DEBUG_PRINT\n\nnamespace zvec {\nnamespace core {\n\nHnswStreamerEntity::HnswStreamerEntity(IndexStreamer::Stats &stats)\n : stats_(stats) {}\n\nHnswStreamerEntity::~HnswStreamerEntity() {}\n\nint HnswStreamerEntity::init(size_t max_doc_cnt) {\n if (std::pow(scaling_factor(), kMaxGraphLayers) < max_doc_cnt) {\n LOG_ERROR(\"scalingFactor=%zu is too small\", scaling_factor());\n return IndexError_InvalidArgument;\n }\n\n std::lock_guard lock(mutex_);\n broker_ = std::make_shared(stats_);\n upper_neighbor_index_ = std::make_shared();\n keys_map_lock_ = std::make_shared();\n keys_map_ = std::make_shared>();\n if (!keys_map_ || !upper_neighbor_index_ || !broker_ || !keys_map_lock_) {\n LOG_ERROR(\"HnswStreamerEntity new object failed\");\n return IndexError_NoMemory;\n }\n keys_map_->set_empty_key(kInvalidKey);\n\n neighbor_size_ = neighbors_size();\n upper_neighbor_size_ = upper_neighbors_size();\n\n //! vector + key + level 0 neighbors\n size_t size = vector_size() + sizeof(key_t) + neighbor_size_;\n\n size = AlignSize(size);\n set_node_size(size);\n return 0;\n}\n\nint HnswStreamerEntity::cleanup() {\n std::lock_guard lock(mutex_);\n mutable_header()->clear();\n chunk_size_ = kDefaultChunkSize;\n node_index_mask_bits_ = 0U;\n node_index_mask_ = 0U;\n node_cnt_per_chunk_ = 0U;\n neighbor_size_ = 0U;\n upper_neighbor_size_ = 0U;\n if (upper_neighbor_index_) {\n upper_neighbor_index_->cleanup();\n }\n if (keys_map_) {\n keys_map_->clear();\n }\n node_chunks_.clear();\n upper_neighbor_chunks_.clear();\n filter_same_key_ = false;\n get_vector_enabled_ = false;\n broker_.reset();\n\n return 0;\n}\n\nint HnswStreamerEntity::update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) {\n std::vector buffer(neighbor_size_);\n NeighborsHeader *hd = reinterpret_cast(buffer.data());\n hd->neighbor_cnt = neighbors.size();\n size_t i = 0;\n for (; i < neighbors.size(); ++i) {\n hd->neighbors[i] = neighbors[i].first;\n }\n\n auto loc = get_neighbor_chunk_loc(level, id);\n size_t size = reinterpret_cast(&hd->neighbors[i]) - &buffer[0];\n size_t ret = loc.first->write(loc.second, hd, size);\n if (ailego_unlikely(ret != size)) {\n LOG_ERROR(\"Write neighbor header failed, ret=%zu\", ret);\n\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nconst Neighbors HnswStreamerEntity::get_neighbors(level_t level,\n node_id_t id) const {\n Chunk *chunk = nullptr;\n size_t offset = 0UL;\n size_t neighbor_size = neighbor_size_;\n if (level == 0UL) {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n offset =\n (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t);\n\n sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_);\n ailego_assert_with(chunk_idx < node_chunks_.size(), \"invalid chunk idx\");\n chunk = node_chunks_[chunk_idx].get();\n } else {\n auto p = get_upper_neighbor_chunk_loc(level, id);\n chunk = upper_neighbor_chunks_[p.first].get();\n offset = p.second;\n neighbor_size = upper_neighbor_size_;\n }\n\n ailego_assert_with(offset < chunk->data_size(), \"invalid chunk offset\");\n IndexStorage::MemoryBlock neighbor_block;\n size_t size = chunk->read(offset, neighbor_block, neighbor_size);\n if (ailego_unlikely(size != neighbor_size)) {\n LOG_ERROR(\"Read neighbor header failed, ret=%zu\", size);\n return Neighbors();\n }\n return Neighbors(neighbor_block);\n}\n\n//! Get vector data by key\nconst void *HnswStreamerEntity::get_vector(node_id_t id) const {\n auto loc = get_vector_chunk_loc(id);\n const void *vec = nullptr;\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = vector_size();\n\n size_t ret = node_chunks_[loc.first]->read(loc.second, &vec, read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%u, read size=%zu, ret=%zu\",\n loc.second, read_size, ret);\n }\n\n return vec;\n}\n\nint HnswStreamerEntity::get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const {\n for (auto i = 0U; i < count; ++i) {\n auto loc = get_vector_chunk_loc(ids[i]);\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = vector_size();\n\n size_t ret = node_chunks_[loc.first]->read(loc.second, &vecs[i], read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%u, read size=%zu, ret=%zu\",\n loc.second, read_size, ret);\n return IndexError_ReadData;\n }\n }\n return 0;\n}\n\nint HnswStreamerEntity::get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const {\n auto loc = get_vector_chunk_loc(id);\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = vector_size();\n\n size_t ret = node_chunks_[loc.first]->read(loc.second, block, read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%u, read size=%zu, ret=%zu\",\n loc.second, read_size, ret);\n return IndexError_ReadData;\n }\n return 0;\n}\n\nint HnswStreamerEntity::get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const {\n vec_blocks.resize(count);\n for (auto i = 0U; i < count; ++i) {\n auto loc = get_vector_chunk_loc(ids[i]);\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = vector_size();\n\n size_t ret =\n node_chunks_[loc.first]->read(loc.second, vec_blocks[i], read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%u, read size=%zu, ret=%zu\",\n loc.second, read_size, ret);\n return IndexError_ReadData;\n }\n }\n return 0;\n}\n\nkey_t HnswStreamerEntity::get_key(node_id_t id) const {\n if (use_key_info_map_) {\n auto loc = get_key_chunk_loc(id);\n IndexStorage::MemoryBlock key_block;\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n size_t ret =\n node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t));\n if (ailego_unlikely(ret != sizeof(key_t))) {\n LOG_ERROR(\"Read vector failed, ret=%zu\", ret);\n return kInvalidKey;\n }\n\n return *reinterpret_cast(key_block.data());\n } else {\n return id;\n }\n}\n\nvoid HnswStreamerEntity::add_neighbor(level_t level, node_id_t id,\n uint32_t size, node_id_t neighbor_id) {\n auto loc = get_neighbor_chunk_loc(level, id);\n size_t offset =\n loc.second + sizeof(NeighborsHeader) + size * sizeof(node_id_t);\n ailego_assert_with(size < neighbor_cnt(level), \"invalid neighbor size\");\n ailego_assert_with(offset < loc.first->data_size(), \"invalid chunk offset\");\n size_t ret = loc.first->write(offset, &neighbor_id, sizeof(node_id_t));\n if (ailego_unlikely(ret != sizeof(node_id_t))) {\n LOG_ERROR(\"Write neighbor id failed, ret=%zu\", ret);\n return;\n }\n\n uint32_t neighbors = size + 1;\n ret = loc.first->write(loc.second, &neighbors, sizeof(uint32_t));\n if (ailego_unlikely(ret != sizeof(uint32_t))) {\n LOG_ERROR(\"Write neighbor cnt failed, ret=%zu\", ret);\n }\n\n return;\n}\n\nint HnswStreamerEntity::init_chunks(const Chunk::Pointer &header_chunk) {\n if (header_chunk->data_size() < header_size()) {\n LOG_ERROR(\"Invalid header chunk size\");\n return IndexError_InvalidFormat;\n }\n IndexStorage::MemoryBlock header_block;\n size_t size = header_chunk->read(0UL, header_block, header_size());\n if (ailego_unlikely(size != header_size())) {\n LOG_ERROR(\"Read header chunk failed\");\n return IndexError_ReadData;\n }\n *mutable_header() =\n *reinterpret_cast(header_block.data());\n\n int ret = check_hnsw_index(&header());\n if (ret != 0) {\n broker_->close();\n return ret;\n }\n\n node_chunks_.resize(broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NODE));\n for (auto seq = 0UL; seq < node_chunks_.size(); ++seq) {\n node_chunks_[seq] = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_NODE, seq);\n if (!node_chunks_[seq]) {\n LOG_ERROR(\"Missing hnsw streamer data chunk %zu th of %zu\", seq,\n node_chunks_.size());\n return IndexError_InvalidFormat;\n }\n }\n\n upper_neighbor_chunks_.resize(\n broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR));\n for (auto seq = 0UL; seq < upper_neighbor_chunks_.size(); ++seq) {\n upper_neighbor_chunks_[seq] =\n broker_->get_chunk(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, seq);\n if (!upper_neighbor_chunks_[seq]) {\n LOG_ERROR(\"Missing hnsw streamer index chunk %zu th of %zu\", seq,\n upper_neighbor_chunks_.size());\n return IndexError_InvalidFormat;\n }\n }\n\n return 0;\n}\n\nint HnswStreamerEntity::open(IndexStorage::Pointer stg, uint64_t max_index_size,\n bool check_crc) {\n std::lock_guard lock(mutex_);\n bool huge_page = stg->isHugePage();\n LOG_DEBUG(\"huge_page: %d\", (int)huge_page);\n int ret = init_chunk_params(max_index_size, huge_page);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"init_chunk_params failed for %s\", IndexError::What(ret));\n return ret;\n }\n ret = broker_->open(std::move(stg), max_index_size_, chunk_size_, check_crc);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Open index failed for %s\", IndexError::What(ret));\n return ret;\n }\n ret = upper_neighbor_index_->init(broker_, upper_neighbor_chunk_size_,\n scaling_factor(), estimate_doc_capacity(),\n kUpperHashMemoryInflateRatio);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Init neighbor hash map failed\");\n return ret;\n }\n\n //! init header\n auto header_chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_HEADER,\n ChunkBroker::kDefaultChunkSeqId);\n if (!header_chunk) { // open empty index, create one\n auto p =\n broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_HEADER,\n ChunkBroker::kDefaultChunkSeqId, header_size());\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc header chunk failed\");\n return p.first;\n }\n size_t size = p.second->write(0UL, &header(), header_size());\n if (ailego_unlikely(size != header_size())) {\n LOG_ERROR(\"Write header chunk failed\");\n return IndexError_WriteData;\n }\n return 0;\n }\n\n //! Open an exist hnsw index\n ret = init_chunks(header_chunk);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n //! total docs including features wrote in index but neighbors may not ready\n node_id_t total_vecs = 0;\n if (node_chunks_.size() > 0) {\n size_t last_idx = node_chunks_.size() - 1;\n auto last_chunk = node_chunks_[last_idx];\n if (last_chunk->data_size() % node_size()) {\n LOG_WARN(\"The index may broken\");\n return IndexError_InvalidFormat;\n }\n total_vecs = last_idx * node_cnt_per_chunk_ +\n node_chunks_[last_idx]->data_size() / node_size();\n }\n\n LOG_INFO(\n \"Open index, l0NeighborCnt=%zu upperNeighborCnt=%zu \"\n \"efConstruction=%zu curDocCnt=%u totalVecs=%u maxLevel=%u\",\n l0_neighbor_cnt(), upper_neighbor_cnt(), ef_construction(), doc_cnt(),\n total_vecs, cur_max_level());\n //! try to correct the docCnt if index not fully flushed\n if (doc_cnt() != total_vecs) {\n LOG_WARN(\"Index closed abnormally, using totalVecs as curDocCnt\");\n *mutable_doc_cnt() = total_vecs;\n }\n if (filter_same_key_ || get_vector_enabled_) {\n if (use_key_info_map_) {\n for (node_id_t id = 0U; id < doc_cnt(); ++id) {\n if (get_key(id) == kInvalidKey) {\n continue;\n }\n (*keys_map_)[get_key(id)] = id;\n }\n }\n }\n\n stats_.set_loaded_count(doc_cnt());\n\n return 0;\n}\n\nint HnswStreamerEntity::close() {\n LOG_DEBUG(\"close index\");\n\n std::lock_guard lock(mutex_);\n flush_header();\n mutable_header()->reset();\n upper_neighbor_index_->cleanup();\n keys_map_->clear();\n header_.clear();\n node_chunks_.clear();\n upper_neighbor_chunks_.clear();\n\n return broker_->close();\n}\n\nint HnswStreamerEntity::flush(uint64_t checkpoint) {\n LOG_INFO(\"Flush index, curDocs=%u\", doc_cnt());\n\n std::lock_guard lock(mutex_);\n flush_header();\n int ret = broker_->flush(checkpoint);\n if (ret != 0) {\n return ret;\n }\n\n return 0;\n}\n\nint HnswStreamerEntity::dump(const IndexDumper::Pointer &dumper) {\n LOG_INFO(\"Dump index, curDocs=%u\", doc_cnt());\n\n //! sort by keys, to support get_vector by key in searcher\n std::vector keys(doc_cnt());\n for (node_id_t i = 0; i < doc_cnt(); ++i) {\n keys[i] = get_key(i);\n }\n\n //! dump neighbors\n auto get_level = [&](node_id_t id) {\n auto it = upper_neighbor_index_->find(id);\n if (it == upper_neighbor_index_->end()) {\n return 0U;\n };\n auto meta = reinterpret_cast(&it->second);\n return meta->level;\n };\n auto ret = dump_segments(dumper, keys.data(), get_level);\n if (ailego_unlikely(ret < 0)) {\n return ret;\n }\n *stats_.mutable_dumped_size() += ret;\n\n return 0;\n}\n\nint HnswStreamerEntity::check_hnsw_index(const HNSWHeader *hd) const {\n if (l0_neighbor_cnt() != hd->l0_neighbor_cnt() ||\n upper_neighbor_cnt() != hd->upper_neighbor_cnt()) {\n LOG_ERROR(\"Param neighbor cnt: %zu:%zu mismatch index previous %zu:%zu\",\n l0_neighbor_cnt(), upper_neighbor_cnt(), hd->l0_neighbor_cnt(),\n hd->upper_neighbor_cnt());\n return IndexError_Mismatch;\n }\n if (vector_size() != hd->vector_size()) {\n LOG_ERROR(\"vector size %zu mismatch index previous %zu\", vector_size(),\n hd->vector_size());\n return IndexError_Mismatch;\n }\n if (ef_construction() != hd->ef_construction()) {\n LOG_WARN(\"Param efConstruction %zu mismatch index previous %zu\",\n ef_construction(), hd->ef_construction());\n }\n if (scaling_factor() != hd->scaling_factor()) {\n LOG_WARN(\"Param scalingFactor %zu mismatch index previous %zu\",\n scaling_factor(), hd->scaling_factor());\n return IndexError_Mismatch;\n }\n if (prune_cnt() != hd->neighbor_prune_cnt()) {\n LOG_WARN(\"Param pruneCnt %zu mismatch index previous %zu\", prune_cnt(),\n hd->neighbor_prune_cnt());\n return IndexError_Mismatch;\n }\n if ((hd->entry_point() != kInvalidNodeId &&\n hd->entry_point() >= hd->doc_cnt()) ||\n (hd->entry_point() == kInvalidNodeId && hd->doc_cnt() > 0U)) {\n LOG_WARN(\"Invalid entryPoint %u, docCnt %u\", hd->entry_point(),\n hd->doc_cnt());\n return IndexError_InvalidFormat;\n }\n if (hd->entry_point() == kInvalidNodeId &&\n broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NODE) > 0) {\n LOG_WARN(\"The index is broken, maybe it haven't flush\");\n return IndexError_InvalidFormat;\n }\n\n return 0;\n}\n\nint HnswStreamerEntity::add_vector(level_t level, key_t key, const void *vec,\n node_id_t *id) {\n Chunk::Pointer node_chunk;\n size_t chunk_offset = -1UL;\n\n std::lock_guard lock(mutex_);\n // duplicate check\n if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) {\n LOG_WARN(\"Try to add duplicate key, ignore it\");\n return IndexError_Duplicate;\n }\n\n node_id_t local_id = static_cast(doc_cnt());\n uint32_t chunk_index = node_chunks_.size() - 1U;\n if (chunk_index == -1U ||\n (node_chunks_[chunk_index]->data_size() >=\n node_cnt_per_chunk_ * node_size())) { // no space left and need to alloc\n if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n chunk_index++;\n auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NODE, chunk_index,\n chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n node_chunk = p.second;\n chunk_offset = 0UL;\n node_chunks_.emplace_back(node_chunk);\n } else {\n node_chunk = node_chunks_[chunk_index];\n chunk_offset = node_chunk->data_size();\n }\n\n size_t size = node_chunk->write(chunk_offset, vec, vector_size());\n if (ailego_unlikely(size != vector_size())) {\n LOG_ERROR(\"Chunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t));\n if (ailego_unlikely(size != sizeof(key_t))) {\n LOG_ERROR(\"Chunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n //! level 0 neighbors is inited to zero by default\n\n int ret = add_upper_neighbor(level, local_id);\n if (ret != 0) {\n return ret;\n }\n\n chunk_offset += node_size();\n if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"Chunk resize to %zu failed\", chunk_offset);\n return IndexError_Runtime;\n }\n if (filter_same_key_ || get_vector_enabled_) {\n if (use_key_info_map_) {\n keys_map_lock_->lock();\n (*keys_map_)[key] = local_id;\n keys_map_lock_->unlock();\n }\n }\n\n *mutable_doc_cnt() += 1;\n broker_->mark_dirty();\n *id = local_id;\n\n return 0;\n}\n\nint HnswStreamerEntity::add_vector_with_id(level_t level, node_id_t id,\n const void *vec) {\n Chunk::Pointer node_chunk;\n size_t chunk_offset = -1UL;\n key_t key = id;\n\n std::lock_guard lock(mutex_);\n\n // duplicate check\n if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) {\n LOG_WARN(\"Try to add duplicate key, ignore it\");\n return IndexError_Duplicate;\n }\n\n // set node_chunk & chunk_offset if succeed\n auto func_get_node_chunk_and_offset = [&](node_id_t node_id) -> int {\n uint32_t chunk_index = node_id >> node_index_mask_bits_;\n ailego_assert_with(chunk_index <= node_chunks_.size(), \"invalid chunk idx\");\n // belongs to next chunk\n if (chunk_index == node_chunks_.size()) {\n if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NODE, chunk_index,\n chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n node_chunk = p.second;\n node_chunks_.emplace_back(node_chunk);\n }\n\n node_chunk = node_chunks_[chunk_index];\n chunk_offset = (node_id & node_index_mask_) * node_size();\n return 0;\n };\n\n for (size_t start_id = doc_cnt(); start_id < id; ++start_id) {\n if (auto ret = func_get_node_chunk_and_offset(start_id); ret != 0) {\n LOG_ERROR(\"func_get_node_chunk_and_offset failed\");\n return ret;\n }\n size_t size = node_chunk->write(chunk_offset + vector_size(), &kInvalidKey,\n sizeof(key_t));\n if (ailego_unlikely(size != sizeof(key_t))) {\n LOG_ERROR(\"Chunk write key failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n chunk_offset += node_size();\n if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"Chunk resize to %zu failed\", chunk_offset);\n return IndexError_Runtime;\n }\n }\n\n if (auto ret = func_get_node_chunk_and_offset(id); ret != 0) {\n LOG_ERROR(\"func_get_node_chunk_and_offset failed\");\n return ret;\n }\n\n size_t size = node_chunk->write(chunk_offset, vec, vector_size());\n if (ailego_unlikely(size != vector_size())) {\n LOG_ERROR(\"Chunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t));\n if (ailego_unlikely(size != sizeof(key_t))) {\n LOG_ERROR(\"Chunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n //! level 0 neighbors is inited to zero by default\n\n int ret = add_upper_neighbor(level, id);\n if (ret != 0) {\n return ret;\n }\n\n if (*mutable_doc_cnt() <= id) {\n *mutable_doc_cnt() = id + 1;\n chunk_offset += node_size();\n if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"Chunk resize to %zu failed\", chunk_offset);\n return IndexError_Runtime;\n }\n }\n\n if (filter_same_key_ || get_vector_enabled_) {\n if (use_key_info_map_) {\n keys_map_lock_->lock();\n (*keys_map_)[key] = id;\n keys_map_lock_->unlock();\n }\n }\n\n broker_->mark_dirty();\n\n return 0;\n}\n\nvoid HnswStreamerEntity::update_ep_and_level(node_id_t ep, level_t level) {\n HnswEntity::update_ep_and_level(ep, level);\n flush_header();\n\n return;\n}\n\nconst HnswEntity::Pointer HnswStreamerEntity::clone() const {\n std::vector node_chunks;\n node_chunks.reserve(node_chunks_.size());\n for (size_t i = 0UL; i < node_chunks_.size(); ++i) {\n node_chunks.emplace_back(node_chunks_[i]->clone());\n if (ailego_unlikely(!node_chunks[i])) {\n LOG_ERROR(\"HnswStreamerEntity get chunk failed in clone\");\n return HnswEntity::Pointer();\n }\n }\n\n std::vector upper_neighbor_chunks;\n upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size());\n for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) {\n upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone());\n if (ailego_unlikely(!upper_neighbor_chunks[i])) {\n LOG_ERROR(\"HnswStreamerEntity get chunk failed in clone\");\n return HnswEntity::Pointer();\n }\n }\n\n HnswStreamerEntity *entity = new (std::nothrow) HnswStreamerEntity(\n stats_, header(), chunk_size_, node_index_mask_bits_,\n upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_,\n upper_neighbor_index_, keys_map_lock_, keys_map_, use_key_info_map_,\n std::move(node_chunks), std::move(upper_neighbor_chunks), broker_);\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"HnswStreamerEntity new failed\");\n }\n return HnswEntity::Pointer(entity);\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw/hnsw_streamer_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"hnsw_chunk.h\"\n#include \"hnsw_entity.h\"\n#include \"hnsw_index_hash.h\"\n#include \"hnsw_params.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! HnswStreamerEntity manage vector data, pkey, and node's neighbors\nclass HnswStreamerEntity : public HnswEntity {\n public:\n //! Cleanup\n //! return 0 on success, or errCode in failure\n virtual int cleanup() override;\n\n //! Make a copy of streamer entity, to support thread-safe operation.\n //! The segment in container cannot be read concurrenly\n virtual const HnswEntity::Pointer clone() const override;\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const override;\n\n //! Get vector feature data by key\n virtual const void *get_vector(node_id_t id) const override;\n\n //! Get vectors feature data by local ids\n virtual int get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const override;\n\n virtual int get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n\n virtual int get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const override;\n\n //! Get the node id's neighbors on graph level\n //! Note: the neighbors cannot be modified, using the following\n //! method to get WritableNeighbors if want to\n virtual const Neighbors get_neighbors(level_t level,\n node_id_t id) const override;\n\n //! Add vector and key to hnsw entity, and local id will be saved in id\n virtual int add_vector(level_t level, key_t key, const void *vec,\n node_id_t *id) override;\n\n //! Add vector and id to hnsw entity\n virtual int add_vector_with_id(level_t level, node_id_t id,\n const void *vec) override;\n\n virtual int update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) override;\n\n //! Append neighbor_id to node id neighbors on level\n //! Notice: the caller must be ensure the neighbors not full\n virtual void add_neighbor(level_t level, node_id_t id, uint32_t size,\n node_id_t neighbor_id) override;\n\n //! Dump index by dumper\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n virtual void update_ep_and_level(node_id_t ep, level_t level) override;\n\n void set_use_key_info_map(bool use_id_map) {\n use_key_info_map_ = use_id_map;\n LOG_DEBUG(\"use_key_info_map_: %d\", (int)use_key_info_map_);\n }\n\n public:\n //! Constructor\n HnswStreamerEntity(IndexStreamer::Stats &stats);\n\n //! Destructor\n ~HnswStreamerEntity();\n\n //! Get vector feature data by key\n virtual const void *get_vector_by_key(key_t key) const override {\n auto id = get_id(key);\n return id == kInvalidNodeId ? nullptr : get_vector(id);\n }\n\n virtual int get_vector_by_key(\n const key_t key, IndexStorage::MemoryBlock &block) const override {\n auto id = get_id(key);\n if (id != kInvalidNodeId) {\n return get_vector(id, block);\n } else {\n return IndexError_InvalidArgument;\n }\n }\n\n //! Init entity\n int init(size_t max_doc_cnt);\n\n //! Flush graph entity to disk\n //! return 0 on success, or errCode in failure\n int flush(uint64_t checkpoint);\n\n //! Open entity from storage\n //! return 0 on success, or errCode in failure\n int open(IndexStorage::Pointer stg, uint64_t max_index_size, bool check_crc);\n\n //! Close entity\n //! return 0 on success, or errCode in failure\n int close();\n\n //! Set meta information from entity\n int set_index_meta(const IndexMeta &meta) const {\n return IndexHelper::SerializeToStorage(meta, broker_->storage().get());\n }\n\n //! Get meta information from entity\n int get_index_meta(IndexMeta *meta) const {\n return IndexHelper::DeserializeFromStorage(broker_->storage().get(), meta);\n }\n\n //! Set params: chunk size\n inline void set_chunk_size(size_t val) {\n chunk_size_ = val;\n }\n\n //! Set params\n inline void set_filter_same_key(bool val) {\n filter_same_key_ = val;\n }\n\n //! Set params\n inline void set_get_vector(bool val) {\n get_vector_enabled_ = val;\n }\n\n //! Get vector local id by key\n inline node_id_t get_id(key_t key) const {\n if (use_key_info_map_) {\n keys_map_lock_->lock_shared();\n auto it = keys_map_->find(key);\n keys_map_lock_->unlock_shared();\n return it == keys_map_->end() ? kInvalidNodeId : it->second;\n } else {\n return key;\n }\n }\n\n void print_key_map() const {\n std::cout << \"key map begins\" << std::endl;\n\n auto iter = keys_map_->begin();\n while (iter != keys_map_->end()) {\n std::cout << \"key: \" << iter->first << \", id: \" << iter->second\n << std::endl;\n ;\n iter++;\n }\n\n std::cout << \"key map ends\" << std::endl;\n }\n\n //! Get l0 neighbors size\n inline size_t neighbors_size() const {\n return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t);\n }\n\n //! Get neighbors size for level > 0\n inline size_t upper_neighbors_size() const {\n return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t);\n }\n\n\n private:\n union UpperNeighborIndexMeta {\n struct {\n uint32_t level : 4;\n uint32_t index : 28; // index is composite type: chunk idx, and the\n // N th neighbors in chunk, they two composite\n // the 28 bits location\n };\n uint32_t data;\n };\n\n template \n using HashMap = google::dense_hash_map>;\n template \n using HashMapPointer = std::shared_ptr>;\n\n template \n using HashSet = google::dense_hash_set>;\n template \n using HashSetPointer = std::shared_ptr>;\n\n //! upper neighbor index hashmap\n using NIHashMap = HnswIndexHashMap;\n using NIHashMapPointer = std::shared_ptr;\n\n //! Private construct, only be called by clone method\n HnswStreamerEntity(IndexStreamer::Stats &stats, const HNSWHeader &hd,\n size_t chunk_size, uint32_t node_index_mask_bits,\n uint32_t upper_neighbor_mask_bits, bool filter_same_key,\n bool get_vector_enabled,\n const NIHashMapPointer &upper_neighbor_index,\n std::shared_ptr &keys_map_lock,\n const HashMapPointer &keys_map,\n bool use_key_info_map,\n std::vector &&node_chunks,\n std::vector &&upper_neighbor_chunks,\n const ChunkBroker::Pointer &broker)\n : stats_(stats),\n chunk_size_(chunk_size),\n node_index_mask_bits_(node_index_mask_bits),\n node_cnt_per_chunk_(1UL << node_index_mask_bits_),\n node_index_mask_(node_cnt_per_chunk_ - 1),\n upper_neighbor_mask_bits_(upper_neighbor_mask_bits),\n upper_neighbor_mask_((1U << upper_neighbor_mask_bits_) - 1),\n filter_same_key_(filter_same_key),\n get_vector_enabled_(get_vector_enabled),\n use_key_info_map_(use_key_info_map),\n upper_neighbor_index_(upper_neighbor_index),\n keys_map_lock_(keys_map_lock),\n keys_map_(keys_map),\n node_chunks_(std::move(node_chunks)),\n upper_neighbor_chunks_(std::move(upper_neighbor_chunks)),\n broker_(broker) {\n *mutable_header() = hd;\n\n neighbor_size_ = neighbors_size();\n upper_neighbor_size_ = upper_neighbors_size();\n }\n\n //! Called only in searching procedure per context, so no need to lock\n void sync_chunks(ChunkBroker::CHUNK_TYPE type, size_t idx,\n std::vector *chunks) const {\n if (ailego_likely(idx < chunks->size())) {\n return;\n }\n for (size_t i = chunks->size(); i <= idx; ++i) {\n auto chunk = broker_->get_chunk(type, i);\n // the storage can ensure get chunk will success after the first get\n ailego_assert_with(!!chunk, \"get chunk failed\");\n chunks->emplace_back(std::move(chunk));\n }\n }\n\n //! return pair: chunk index + chunk offset\n inline std::pair get_vector_chunk_loc(\n node_id_t id) const {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n uint32_t offset = (id & node_index_mask_) * node_size();\n\n sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_);\n return std::make_pair(chunk_idx, offset);\n }\n\n //! return pair: chunk index + chunk offset\n inline std::pair get_key_chunk_loc(node_id_t id) const {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n uint32_t offset = (id & node_index_mask_) * node_size() + vector_size();\n\n sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_);\n return std::make_pair(chunk_idx, offset);\n }\n\n inline std::pair get_upper_neighbor_chunk_loc(\n level_t level, node_id_t id) const {\n auto it = upper_neighbor_index_->find(id);\n ailego_assert_abort(it != upper_neighbor_index_->end(),\n \"Get upper neighbor header failed\");\n auto meta = reinterpret_cast(&it->second);\n uint32_t chunk_idx = (meta->index) >> upper_neighbor_mask_bits_;\n uint32_t offset = (((meta->index) & upper_neighbor_mask_) + level - 1) *\n upper_neighbor_size_;\n sync_chunks(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, chunk_idx,\n &upper_neighbor_chunks_);\n ailego_assert_abort(chunk_idx < upper_neighbor_chunks_.size(),\n \"invalid chunk idx\");\n ailego_assert_abort(offset < upper_neighbor_chunks_[chunk_idx]->data_size(),\n \"invalid chunk offset\");\n return std::make_pair(chunk_idx, offset);\n }\n\n //! return pair: chunk + chunk offset\n inline std::pair get_neighbor_chunk_loc(level_t level,\n node_id_t id) const {\n if (level == 0UL) {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n uint32_t offset =\n (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t);\n\n sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_);\n ailego_assert_abort(chunk_idx < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_abort(offset < node_chunks_[chunk_idx]->data_size(),\n \"invalid chunk offset\");\n return std::make_pair(node_chunks_[chunk_idx].get(), offset);\n } else {\n auto p = get_upper_neighbor_chunk_loc(level, id);\n return std::make_pair(upper_neighbor_chunks_[p.first].get(), p.second);\n }\n }\n\n //! Chunk hnsw index valid\n int check_hnsw_index(const HNSWHeader *hd) const;\n\n size_t get_total_upper_neighbors_size(level_t level) const {\n return level * upper_neighbor_size_;\n }\n\n //! Add upper neighbor header and reserve space for upper neighbor\n int add_upper_neighbor(level_t level, node_id_t id) {\n if (level == 0) {\n return 0;\n }\n Chunk::Pointer chunk;\n uint64_t chunk_offset = -1UL;\n size_t neighbors_size = get_total_upper_neighbors_size(level);\n uint64_t chunk_index = upper_neighbor_chunks_.size() - 1UL;\n if (chunk_index == -1UL ||\n (upper_neighbor_chunks_[chunk_index]->padding_size() <\n neighbors_size)) { // no space left and need to alloc\n chunk_index++;\n if (ailego_unlikely(upper_neighbor_chunks_.capacity() ==\n upper_neighbor_chunks_.size())) {\n LOG_ERROR(\"add upper neighbor failed for no memory quota\");\n return IndexError_IndexFull;\n }\n auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR,\n chunk_index, upper_neighbor_chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n chunk = p.second;\n chunk_offset = 0UL;\n upper_neighbor_chunks_.emplace_back(chunk);\n } else {\n chunk = upper_neighbor_chunks_[chunk_index];\n chunk_offset = chunk->data_size();\n }\n ailego_assert_with((size_t)level < kMaxGraphLayers, \"invalid level\");\n ailego_assert_with(chunk_offset % upper_neighbor_size_ == 0,\n \"invalid offset\");\n ailego_assert_with((chunk_offset / upper_neighbor_size_) <\n (1U << upper_neighbor_mask_bits_),\n \"invalid offset\");\n ailego_assert_with(chunk_index < (1U << (28 - upper_neighbor_mask_bits_)),\n \"invalid chunk index\");\n UpperNeighborIndexMeta meta;\n meta.level = level;\n meta.index = (chunk_index << upper_neighbor_mask_bits_) |\n (chunk_offset / upper_neighbor_size_);\n chunk_offset += upper_neighbor_size_ * level;\n if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) {\n LOG_ERROR(\"HashMap insert value failed\");\n return IndexError_Runtime;\n }\n\n if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"Chunk resize to %zu failed\", (size_t)chunk_offset);\n return IndexError_Runtime;\n }\n\n return 0;\n }\n\n size_t estimate_doc_capacity() const {\n return node_chunks_.capacity() * node_cnt_per_chunk_;\n }\n\n int init_chunk_params(size_t max_index_size, bool huge_page) {\n node_cnt_per_chunk_ = std::max(1, chunk_size_ / node_size());\n //! align node cnt per chunk to pow of 2\n node_index_mask_bits_ = std::ceil(std::log2(node_cnt_per_chunk_));\n node_cnt_per_chunk_ = 1UL << node_index_mask_bits_;\n if (huge_page) {\n chunk_size_ = AlignHugePageSize(node_cnt_per_chunk_ * node_size());\n } else {\n chunk_size_ = AlignPageSize(node_cnt_per_chunk_ * node_size());\n }\n node_index_mask_ = node_cnt_per_chunk_ - 1;\n\n if (max_index_size == 0UL) {\n max_index_size_ = chunk_size_ * kDefaultMaxChunkCnt;\n } else {\n max_index_size_ = max_index_size;\n }\n\n //! To get a balanced upper neighbor chunk size.\n //! If the upper chunk size is equal to node chunk size, it may waste\n //! upper neighbor chunk space; if the upper neighbor chunk size is too\n //! small, the will need large upper neighbor chunks index space. So to\n //! get a balanced ratio be sqrt of the node/neighbor size ratio\n float ratio =\n std::sqrt(node_size() * scaling_factor() * 1.0f / upper_neighbor_size_);\n if (huge_page) {\n upper_neighbor_chunk_size_ = AlignHugePageSize(\n std::max(get_total_upper_neighbors_size(kMaxGraphLayers),\n static_cast(chunk_size_ / ratio)));\n } else {\n upper_neighbor_chunk_size_ = AlignPageSize(\n std::max(get_total_upper_neighbors_size(kMaxGraphLayers),\n static_cast(chunk_size_ / ratio)));\n }\n upper_neighbor_mask_bits_ =\n std::ceil(std::log2(upper_neighbor_chunk_size_ / upper_neighbor_size_));\n upper_neighbor_mask_ = (1 << upper_neighbor_mask_bits_) - 1;\n\n size_t max_node_chunk_cnt = std::ceil(max_index_size_ / chunk_size_);\n size_t max_upper_chunk_cnt = std::ceil(\n (max_node_chunk_cnt * node_cnt_per_chunk_ * 1.0f / scaling_factor()) /\n (upper_neighbor_chunk_size_ / upper_neighbor_size_));\n max_upper_chunk_cnt =\n max_upper_chunk_cnt + std::ceil(max_upper_chunk_cnt / scaling_factor());\n\n //! reserve space to avoid memmove in chunks vector emplace chunk, so\n //! as to lock-free in reading chunk\n node_chunks_.reserve(max_node_chunk_cnt);\n upper_neighbor_chunks_.reserve(max_upper_chunk_cnt);\n\n LOG_DEBUG(\n \"Settings: nodeSize=%zu chunkSize=%u upperNeighborSize=%u \"\n \"upperNeighborChunkSize=%u \"\n \"nodeCntPerChunk=%u maxChunkCnt=%zu maxNeighborChunkCnt=%zu \"\n \"maxIndexSize=%zu ratio=%.3f\",\n node_size(), chunk_size_, upper_neighbor_size_,\n upper_neighbor_chunk_size_, node_cnt_per_chunk_, max_node_chunk_cnt,\n max_upper_chunk_cnt, max_index_size_, ratio);\n\n return 0;\n }\n\n //! Init node chunk and neighbor chunks\n int init_chunks(const Chunk::Pointer &header_chunk);\n\n int flush_header(void) {\n if (!broker_->dirty()) {\n // do not need to flush\n return 0;\n }\n auto header_chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_HEADER,\n ChunkBroker::kDefaultChunkSeqId);\n if (ailego_unlikely(!header_chunk)) {\n LOG_ERROR(\"get header chunk failed\");\n return IndexError_Runtime;\n }\n size_t size = header_chunk->write(0UL, &header(), header_size());\n if (ailego_unlikely(size != header_size())) {\n LOG_ERROR(\"Write header chunk failed\");\n return IndexError_WriteData;\n }\n\n return 0;\n }\n\n private:\n HnswStreamerEntity(const HnswStreamerEntity &) = delete;\n HnswStreamerEntity &operator=(const HnswStreamerEntity &) = delete;\n static constexpr uint64_t kUpperHashMemoryInflateRatio = 2.0f;\n\n private:\n IndexStreamer::Stats &stats_;\n HNSWHeader header_{};\n std::mutex mutex_{};\n size_t max_index_size_{0UL};\n uint32_t chunk_size_{kDefaultChunkSize};\n uint32_t upper_neighbor_chunk_size_{kDefaultChunkSize};\n uint32_t node_index_mask_bits_{0U};\n uint32_t node_cnt_per_chunk_{0U};\n uint32_t node_index_mask_{0U};\n uint32_t neighbor_size_{0U};\n uint32_t upper_neighbor_size_{0U};\n //! UpperNeighborIndex.index composite chunkIdx and offset in chunk by the\n //! following mask\n uint32_t upper_neighbor_mask_bits_{0U};\n uint32_t upper_neighbor_mask_{0U};\n bool filter_same_key_{false};\n bool get_vector_enabled_{false};\n bool use_key_info_map_{true};\n\n NIHashMapPointer upper_neighbor_index_{};\n\n mutable std::shared_ptr keys_map_lock_{};\n HashMapPointer keys_map_{};\n\n //! the chunks will be changed in searcher, so need mutable\n //! data chunk include: vector, key, level 0 neighbors\n mutable std::vector node_chunks_{};\n\n //! upper neighbor chunk inlude: UpperNeighborHeader + (1~level) neighbors\n mutable std::vector upper_neighbor_chunks_{};\n\n ChunkBroker::Pointer broker_{}; // chunk broker\n};\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_rabitq/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\nif(AUTO_DETECT_ARCH)\n foreach(FILE ${HNSW_RABITQ_FILES})\n set_source_files_properties(\n ${FILE}\n PROPERTIES\n COMPILE_FLAGS \"${RABITQ_ARCH_FLAG}\"\n )\n endforeach()\nendif()\n\ncc_library(\n NAME core_knn_hnsw_rabitq\n STATIC SHARED STRICT ALWAYS_LINK\n SRCS *.cc\n LIBS core_framework rabitqlib sparsehash\n INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_algorithm.h\"\n#include \n#include \"hnsw_rabitq_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswRabitqAlgorithm::HnswRabitqAlgorithm(HnswRabitqEntity &entity)\n : entity_(entity),\n mt_(std::chrono::system_clock::now().time_since_epoch().count()),\n lock_pool_(kLockCnt) {}\n\nint HnswRabitqAlgorithm::cleanup() {\n return 0;\n}\n\nint HnswRabitqAlgorithm::add_node(node_id_t id, level_t level,\n HnswRabitqContext *ctx) {\n spin_lock_.lock();\n\n // std::cout << \"id: \" << id << \", level: \" << level << std::endl;\n\n auto cur_max_level = entity_.cur_max_level();\n auto entry_point = entity_.entry_point();\n if (ailego_unlikely(entry_point == kInvalidNodeId)) {\n entity_.update_ep_and_level(id, level);\n spin_lock_.unlock();\n return 0;\n }\n spin_lock_.unlock();\n\n if (ailego_unlikely(level > cur_max_level)) {\n mutex_.lock();\n // re-check max level\n cur_max_level = entity_.cur_max_level();\n entry_point = entity_.entry_point();\n if (level <= cur_max_level) {\n mutex_.unlock();\n }\n }\n\n level_t cur_level = cur_max_level;\n ResultRecord dist = ctx->dist_calculator()(entry_point);\n for (; cur_level > level; --cur_level) {\n select_entry_point(cur_level, &entry_point, &dist, ctx);\n }\n\n for (; cur_level >= 0; --cur_level) {\n search_neighbors(cur_level, &entry_point, &dist, ctx->level_topk(cur_level),\n ctx);\n }\n\n // add neighbors from down level to top level, to avoid upper level visible\n // to knn_search but the under layer level not ready\n for (cur_level = 0; cur_level <= level; ++cur_level) {\n add_neighbors(id, cur_level, ctx->level_topk(cur_level), ctx);\n ctx->level_topk(cur_level).clear();\n }\n\n if (ailego_unlikely(level > cur_max_level)) {\n spin_lock_.lock();\n entity_.update_ep_and_level(id, level);\n spin_lock_.unlock();\n mutex_.unlock();\n }\n\n return 0;\n}\n\n//! select_entry_point on hnsw level, ef = 1\nvoid HnswRabitqAlgorithm::select_entry_point(level_t level,\n node_id_t *entry_point,\n ResultRecord *dist,\n HnswRabitqContext *ctx) const {\n auto &entity = ctx->get_entity();\n HnswRabitqAddDistCalculator &dc = ctx->dist_calculator();\n while (true) {\n const Neighbors neighbors = entity.get_neighbors(level, *entry_point);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n uint32_t size = neighbors.size();\n if (size == 0) {\n break;\n }\n\n std::vector neighbor_vec_blocks;\n int ret = dc.get_vector(&neighbors[0], size, neighbor_vec_blocks);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_vector())++;\n }\n if (ailego_unlikely(ret != 0)) {\n break;\n }\n\n bool find_closer = false;\n\n std::vector dists(size);\n std::vector neighbor_vecs(size);\n for (uint32_t i = 0; i < size; ++i) {\n neighbor_vecs[i] = neighbor_vec_blocks[i].data();\n }\n\n dc.batch_dist(neighbor_vecs.data(), size, dists.data());\n\n for (uint32_t i = 0; i < size; ++i) {\n ResultRecord cur_dist = dists[i];\n\n if (cur_dist < *dist) {\n *entry_point = neighbors[i];\n *dist = cur_dist;\n find_closer = true;\n }\n }\n\n if (!find_closer) {\n break;\n }\n }\n\n return;\n}\n\nvoid HnswRabitqAlgorithm::add_neighbors(node_id_t id, level_t level,\n TopkHeap &topk_heap,\n HnswRabitqContext *ctx) {\n if (ailego_unlikely(topk_heap.size() == 0)) {\n return;\n }\n\n HnswRabitqAddDistCalculator &dc = ctx->dist_calculator();\n\n update_neighbors(dc, id, level, topk_heap);\n\n // reverse update neighbors\n for (size_t i = 0; i < topk_heap.size(); ++i) {\n reverse_update_neighbors(dc, topk_heap[i].first, level, id,\n topk_heap[i].second, ctx->update_heap());\n }\n\n return;\n}\n\nvoid HnswRabitqAlgorithm::search_neighbors(level_t level,\n node_id_t *entry_point,\n ResultRecord *dist, TopkHeap &topk,\n HnswRabitqContext *ctx) const {\n const auto &entity = ctx->get_entity();\n HnswRabitqAddDistCalculator &dc = ctx->dist_calculator();\n VisitFilter &visit = ctx->visit_filter();\n CandidateHeap &candidates = ctx->candidates();\n std::function filter = [](node_id_t) { return false; };\n if (ctx->filter().is_valid()) {\n filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); };\n }\n\n candidates.clear();\n visit.clear();\n visit.set_visited(*entry_point);\n if (!filter(*entry_point)) {\n topk.emplace(*entry_point, *dist);\n }\n\n candidates.emplace(*entry_point, *dist);\n while (!candidates.empty() && !ctx->reach_scan_limit()) {\n auto top = candidates.begin();\n node_id_t main_node = top->first;\n ResultRecord main_dist = top->second;\n\n if (topk.full() && main_dist > topk[0].second) {\n break;\n }\n\n candidates.pop();\n const Neighbors neighbors = entity.get_neighbors(level, main_node);\n ailego_prefetch(neighbors.data);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n\n std::vector neighbor_ids(neighbors.size());\n uint32_t size = 0;\n for (uint32_t i = 0; i < neighbors.size(); ++i) {\n node_id_t node = neighbors[i];\n if (visit.visited(node)) {\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_visit_dup_cnt())++;\n }\n continue;\n }\n visit.set_visited(node);\n neighbor_ids[size++] = node;\n }\n if (size == 0) {\n continue;\n }\n\n std::vector neighbor_vec_blocks;\n int ret = dc.get_vector(neighbor_ids.data(), size, neighbor_vec_blocks);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_vector())++;\n }\n if (ailego_unlikely(ret != 0)) {\n break;\n }\n\n // do prefetch\n static constexpr node_id_t BATCH_SIZE = 12;\n static constexpr node_id_t PREFETCH_STEP = 2;\n for (uint32_t i = 0; i < std::min(BATCH_SIZE * PREFETCH_STEP, size); ++i) {\n ailego_prefetch(neighbor_vec_blocks[i].data());\n }\n // done\n\n std::vector dists(size);\n std::vector neighbor_vecs(size);\n\n for (uint32_t i = 0; i < size; ++i) {\n neighbor_vecs[i] = neighbor_vec_blocks[i].data();\n }\n\n dc.batch_dist(neighbor_vecs.data(), size, dists.data());\n\n for (uint32_t i = 0; i < size; ++i) {\n node_id_t node = neighbor_ids[i];\n ResultRecord cur_dist = dists[i];\n\n if ((!topk.full()) || cur_dist < topk[0].second) {\n candidates.emplace(node, cur_dist);\n // update entry_point for next level scan\n if (cur_dist < *dist) {\n *entry_point = node;\n *dist = cur_dist;\n }\n if (!filter(node)) {\n topk.emplace(node, cur_dist);\n }\n } // end if\n } // end for\n } // while\n\n return;\n}\n\nvoid HnswRabitqAlgorithm::update_neighbors(HnswRabitqAddDistCalculator &dc,\n node_id_t id, level_t level,\n TopkHeap &topk_heap) {\n topk_heap.sort();\n\n uint32_t max_neighbor_cnt = entity_.neighbor_cnt(level);\n if (topk_heap.size() <= static_cast(entity_.prune_cnt())) {\n if (topk_heap.size() <= static_cast(max_neighbor_cnt)) {\n entity_.update_neighbors(level, id, topk_heap);\n return;\n }\n }\n\n uint32_t cur_size = 0;\n for (size_t i = 0; i < topk_heap.size(); ++i) {\n node_id_t cur_node = topk_heap[i].first;\n ResultRecord cur_node_dist = topk_heap[i].second;\n bool good = true;\n for (uint32_t j = 0; j < cur_size; ++j) {\n ResultRecord tmp_dist = dc.dist(cur_node, topk_heap[j].first);\n if (tmp_dist <= cur_node_dist) {\n good = false;\n break;\n }\n }\n\n if (good) {\n topk_heap[cur_size].first = cur_node;\n topk_heap[cur_size].second = cur_node_dist;\n cur_size++;\n if (cur_size >= max_neighbor_cnt) {\n break;\n }\n }\n }\n\n // when after-prune neighbor count is too seldom,\n // we use this strategy to make-up enough edges\n // not only just make-up out-degrees\n // we also make-up enough in-degrees\n uint32_t min_neighbors = entity_.min_neighbor_cnt();\n for (size_t k = cur_size; cur_size < min_neighbors && k < topk_heap.size();\n ++k) {\n bool exist = false;\n for (size_t j = 0; j < cur_size; ++j) {\n if (topk_heap[j].first == topk_heap[k].first) {\n exist = true;\n break;\n }\n }\n if (!exist) {\n topk_heap[cur_size].first = topk_heap[k].first;\n topk_heap[cur_size].second = topk_heap[k].second;\n cur_size++;\n }\n }\n\n topk_heap.resize(cur_size);\n entity_.update_neighbors(level, id, topk_heap);\n\n return;\n}\n\nvoid HnswRabitqAlgorithm::reverse_update_neighbors(\n HnswRabitqAddDistCalculator &dc, node_id_t id, level_t level,\n node_id_t link_id, ResultRecord dist, TopkHeap &update_heap) {\n const size_t max_neighbor_cnt = entity_.neighbor_cnt(level);\n\n uint32_t lock_idx = id & kLockMask;\n lock_pool_[lock_idx].lock();\n const Neighbors neighbors = entity_.get_neighbors(level, id);\n size_t size = neighbors.size();\n ailego_assert_with(size <= max_neighbor_cnt, \"invalid neighbor size\");\n if (size < max_neighbor_cnt) {\n entity_.add_neighbor(level, id, size, link_id);\n lock_pool_[lock_idx].unlock();\n return;\n }\n\n update_heap.emplace(link_id, dist);\n\n for (size_t i = 0; i < size; ++i) {\n node_id_t node = neighbors[i];\n ResultRecord cur_dist = dc.dist(id, node);\n update_heap.emplace(node, cur_dist);\n }\n\n //! TODO: optimize prune\n //! prune edges\n update_heap.sort();\n size_t cur_size = 0;\n for (size_t i = 0; i < update_heap.size(); ++i) {\n node_id_t cur_node = update_heap[i].first;\n ResultRecord cur_node_dist = update_heap[i].second;\n bool good = true;\n for (size_t j = 0; j < cur_size; ++j) {\n ResultRecord tmp_dist = dc.dist(cur_node, update_heap[j].first);\n if (tmp_dist <= cur_node_dist) {\n good = false;\n break;\n }\n }\n\n if (good) {\n update_heap[cur_size].first = cur_node;\n update_heap[cur_size].second = cur_node_dist;\n cur_size++;\n if (cur_size >= max_neighbor_cnt) {\n break;\n }\n }\n }\n\n update_heap.resize(cur_size);\n entity_.update_neighbors(level, id, update_heap);\n\n lock_pool_[lock_idx].unlock();\n\n update_heap.clear();\n\n return;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_algorithm.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"hnsw_rabitq_context.h\"\n#include \"hnsw_rabitq_dist_calculator.h\"\n#include \"hnsw_rabitq_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! hnsw graph algorithm implement\nclass HnswRabitqAlgorithm {\n public:\n typedef std::unique_ptr UPointer;\n\n public:\n //! Constructor\n explicit HnswRabitqAlgorithm(HnswRabitqEntity &entity);\n\n //! Destructor\n ~HnswRabitqAlgorithm() = default;\n\n //! Cleanup HnswRabitqAlgorithm\n int cleanup();\n\n //! Add a node to hnsw graph\n //! @id: the node unique id\n //! @level: a node will be add to graph in each level [0, level]\n //! return 0 on success, or errCode in failure\n int add_node(node_id_t id, level_t level, HnswRabitqContext *ctx);\n\n //! Initiate HnswRabitqAlgorithm\n int init() {\n level_probas_.clear();\n double level_mult =\n 1 / std::log(static_cast(entity_.scaling_factor()));\n for (int level = 0;; level++) {\n // refers faiss get_random_level alg\n double proba =\n std::exp(-level / level_mult) * (1 - std::exp(-1 / level_mult));\n if (proba < 1e-9) {\n break;\n }\n level_probas_.push_back(proba);\n }\n\n return 0;\n }\n\n //! Generate a random level\n //! return graph level\n uint32_t get_random_level() const {\n // gen rand float (0, 1)\n double f = mt_() / static_cast(mt_.max());\n for (size_t level = 0; level < level_probas_.size(); level++) {\n if (f < level_probas_[level]) {\n return level;\n }\n f -= level_probas_[level];\n }\n return level_probas_.size() - 1;\n }\n\n private:\n //! Select in upper layer to get entry point for next layer search\n void select_entry_point(level_t level, node_id_t *entry_point,\n ResultRecord *dist, HnswRabitqContext *ctx) const;\n\n //! update node id neighbors from topkHeap, and reverse link is also updated\n void add_neighbors(node_id_t id, level_t level, TopkHeap &topk_heap,\n HnswRabitqContext *ctx);\n\n //! Given a node id and level, search the nearest neighbors in graph\n //! Note: the nearest neighbors result keeps in topk, and entry_point and\n //! dist will be updated to current level nearest node id and distance\n void search_neighbors(level_t level, node_id_t *entry_point,\n ResultRecord *dist, TopkHeap &topk,\n HnswRabitqContext *ctx) const;\n\n //! Update the node's neighbors\n void update_neighbors(HnswRabitqAddDistCalculator &dc, node_id_t id,\n level_t level, TopkHeap &topk_heap);\n\n //! Checking linkId could be id's new neighbor, and add as neighbor if true\n //! @dc distance calculator\n //! @updateHeap temporary heap in updating neighbors\n void reverse_update_neighbors(HnswRabitqAddDistCalculator &dc, node_id_t id,\n level_t level, node_id_t link_id,\n ResultRecord dist, TopkHeap &update_heap);\n\n private:\n HnswRabitqAlgorithm(const HnswRabitqAlgorithm &) = delete;\n HnswRabitqAlgorithm &operator=(const HnswRabitqAlgorithm &) = delete;\n\n private:\n static constexpr uint32_t kLockCnt{1U << 8};\n static constexpr uint32_t kLockMask{kLockCnt - 1U};\n\n HnswRabitqEntity &entity_;\n mutable std::mt19937 mt_{};\n std::vector level_probas_{};\n\n mutable ailego::SpinMutex spin_lock_{}; // global spin lock\n std::mutex mutex_{}; // global mutex\n // TODO: spin lock?\n std::vector lock_pool_{};\n};\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_builder.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"zvec/core/framework/index_error.h\"\n#include \"zvec/core/framework/index_factory.h\"\n#include \"zvec/core/framework/index_logger.h\"\n#include \"zvec/core/framework/index_memory.h\"\n#include \"zvec/core/framework/index_meta.h\"\n#include \"zvec/core/framework/index_provider.h\"\n#include \"hnsw_rabitq_algorithm.h\"\n#include \"hnsw_rabitq_entity.h\"\n#include \"hnsw_rabitq_params.h\"\n#include \"rabitq_converter.h\"\n#include \"rabitq_params.h\"\n#include \"rabitq_reformer.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswRabitqBuilder::HnswRabitqBuilder() {}\n\nint HnswRabitqBuilder::init(const IndexMeta &meta,\n const ailego::Params ¶ms) {\n LOG_INFO(\"Begin HnswRabitqBuilder::init\");\n\n meta_ = meta;\n auto params_copy = params;\n meta_.set_builder(\"HnswRabitqBuilder\", HnswRabitqEntity::kRevision,\n std::move(params_copy));\n\n size_t memory_quota = 0UL;\n params.get(PARAM_HNSW_RABITQ_BUILDER_MEMORY_QUOTA, &memory_quota);\n params.get(PARAM_HNSW_RABITQ_BUILDER_THREAD_COUNT, &thread_cnt_);\n params.get(PARAM_HNSW_RABITQ_BUILDER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_);\n params.get(PARAM_HNSW_RABITQ_BUILDER_EFCONSTRUCTION, &ef_construction_);\n params.get(PARAM_HNSW_RABITQ_BUILDER_CHECK_INTERVAL_SECS,\n &check_interval_secs_);\n\n params.get(PARAM_HNSW_RABITQ_BUILDER_MAX_NEIGHBOR_COUNT,\n &upper_max_neighbor_cnt_);\n float multiplier = HnswRabitqEntity::kDefaultL0MaxNeighborCntMultiplier;\n params.get(PARAM_HNSW_RABITQ_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER,\n &multiplier);\n l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_;\n scaling_factor_ = upper_max_neighbor_cnt_;\n params.get(PARAM_HNSW_RABITQ_BUILDER_SCALING_FACTOR, &scaling_factor_);\n\n multiplier = HnswRabitqEntity::kDefaultNeighborPruneMultiplier;\n params.get(PARAM_HNSW_RABITQ_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier);\n size_t prune_cnt = multiplier * upper_max_neighbor_cnt_;\n\n if (ef_construction_ == 0) {\n ef_construction_ = HnswRabitqEntity::kDefaultEfConstruction;\n }\n if (upper_max_neighbor_cnt_ == 0) {\n upper_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt;\n }\n if (upper_max_neighbor_cnt_ > kMaxNeighborCnt) {\n LOG_ERROR(\"[%s] must be in range (0,%d]\",\n PARAM_HNSW_RABITQ_BUILDER_MAX_NEIGHBOR_COUNT.c_str(),\n kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) {\n LOG_ERROR(\"[%s]-[%d] must be <= [%s]-[%d]\",\n PARAM_HNSW_RABITQ_BUILDER_MIN_NEIGHBOR_COUNT.c_str(),\n min_neighbor_cnt_,\n PARAM_HNSW_RABITQ_BUILDER_MAX_NEIGHBOR_COUNT.c_str(),\n upper_max_neighbor_cnt_);\n return IndexError_InvalidArgument;\n }\n if (l0_max_neighbor_cnt_ == 0) {\n l0_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt;\n }\n if (l0_max_neighbor_cnt_ > HnswRabitqEntity::kMaxNeighborCnt) {\n LOG_ERROR(\"L0MaxNeighborCnt must be in range (0,%d)\",\n HnswRabitqEntity::kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (scaling_factor_ == 0U) {\n scaling_factor_ = HnswRabitqEntity::kDefaultScalingFactor;\n }\n if (scaling_factor_ < 5 || scaling_factor_ > 1000) {\n LOG_ERROR(\"[%s] must be in range [5,1000]\",\n PARAM_HNSW_RABITQ_BUILDER_SCALING_FACTOR.c_str());\n return IndexError_InvalidArgument;\n }\n if (thread_cnt_ == 0) {\n thread_cnt_ = std::thread::hardware_concurrency();\n }\n if (thread_cnt_ > std::thread::hardware_concurrency()) {\n LOG_WARN(\"[%s] greater than cpu cores %zu\",\n PARAM_HNSW_RABITQ_BUILDER_THREAD_COUNT.c_str(),\n static_cast(std::thread::hardware_concurrency()));\n }\n if (prune_cnt == 0UL) {\n prune_cnt = upper_max_neighbor_cnt_;\n }\n\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"CreateMetric failed, name: %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n int ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"IndexMetric init failed, ret=%d\", ret);\n return ret;\n }\n\n uint32_t total_bits = 0;\n params.get(PARAM_RABITQ_TOTAL_BITS, &total_bits);\n if (total_bits == 0) {\n total_bits = kDefaultRabitqTotalBits;\n }\n if (total_bits < 1 || total_bits > 9) {\n LOG_ERROR(\"Invalid total_bits: %zu, must be in [1, 9]\", (size_t)total_bits);\n return IndexError_InvalidArgument;\n }\n uint8_t ex_bits = total_bits - 1;\n entity_.set_ex_bits(ex_bits);\n\n uint32_t dimension = 0;\n params.get(PARAM_HNSW_RABITQ_GENERAL_DIMENSION, &dimension);\n if (dimension == 0) {\n LOG_ERROR(\"%s not set\", PARAM_HNSW_RABITQ_GENERAL_DIMENSION.c_str());\n return IndexError_InvalidArgument;\n }\n if (dimension < kMinRabitqDimSize || dimension > kMaxRabitqDimSize) {\n LOG_ERROR(\"Invalid dimension: %u, must be in [%d, %d]\", dimension,\n kMinRabitqDimSize, kMaxRabitqDimSize);\n return IndexError_InvalidArgument;\n }\n entity_.update_rabitq_params_and_vector_size(dimension);\n\n entity_.set_ef_construction(ef_construction_);\n entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_);\n entity_.set_min_neighbor_cnt(min_neighbor_cnt_);\n entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_);\n entity_.set_scaling_factor(scaling_factor_);\n entity_.set_memory_quota(memory_quota);\n entity_.set_prune_cnt(prune_cnt);\n\n ret = entity_.init();\n if (ret != 0) {\n return ret;\n }\n\n alg_ = HnswRabitqAlgorithm::UPointer(new HnswRabitqAlgorithm(entity_));\n\n ret = alg_->init();\n if (ret != 0) {\n return ret;\n }\n\n // Create and initialize RaBitQ converter\n converter_ = std::make_shared();\n\n IndexMeta converter_meta = meta_;\n converter_meta.set_dimension(dimension);\n ret = converter_->init(converter_meta, params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize RabitqConverter: %d\", ret);\n return ret;\n }\n\n state_ = BUILD_STATE_INITED;\n LOG_INFO(\n \"End HnswRabitqBuilder::init, params: rawVectorSize=%u vectorSize=%zu \"\n \"efConstruction=%u \"\n \"l0NeighborCnt=%u upperNeighborCnt=%u scalingFactor=%u \"\n \"memoryQuota=%zu neighborPruneCnt=%zu metricName=%s \",\n meta_.element_size(), entity_.vector_size(), ef_construction_,\n l0_max_neighbor_cnt_, upper_max_neighbor_cnt_, scaling_factor_,\n memory_quota, prune_cnt, meta_.metric_name().c_str());\n\n return 0;\n}\n\nint HnswRabitqBuilder::cleanup(void) {\n LOG_INFO(\"Begin HnswRabitqBuilder::cleanup\");\n\n l0_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultL0MaxNeighborCnt;\n min_neighbor_cnt_ = 0;\n upper_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt;\n ef_construction_ = HnswRabitqEntity::kDefaultEfConstruction;\n scaling_factor_ = HnswRabitqEntity::kDefaultScalingFactor;\n check_interval_secs_ = kDefaultLogIntervalSecs;\n errcode_ = 0;\n error_ = false;\n entity_.cleanup();\n if (alg_) {\n alg_->cleanup();\n }\n meta_.clear();\n metric_.reset();\n stats_.clear_attributes();\n stats_.set_trained_count(0UL);\n stats_.set_built_count(0UL);\n stats_.set_dumped_count(0UL);\n stats_.set_discarded_count(0UL);\n stats_.set_trained_costtime(0UL);\n stats_.set_built_costtime(0UL);\n stats_.set_dumped_costtime(0UL);\n state_ = BUILD_STATE_INIT;\n\n LOG_INFO(\"End HnswRabitqBuilder::cleanup\");\n\n return 0;\n}\n\nint HnswRabitqBuilder::train(IndexThreads::Pointer,\n IndexHolder::Pointer holder) {\n if (state_ != BUILD_STATE_INITED) {\n LOG_ERROR(\"Init the builder before HnswRabitqBuilder::train\");\n return IndexError_NoReady;\n }\n\n if (!holder) {\n LOG_ERROR(\"Input holder is nullptr while training index\");\n return IndexError_InvalidArgument;\n }\n if (!holder->is_matched(meta_)) {\n LOG_ERROR(\"Input holder doesn't match index meta while training index\");\n return IndexError_Mismatch;\n }\n LOG_INFO(\"Begin HnswRabitqBuilder::train\");\n size_t trained_cost_time = 0;\n size_t trained_count = 0;\n\n int ret = train_converter_and_load_reformer(holder);\n if (ret != 0) {\n return ret;\n }\n\n if (metric_->support_train()) {\n auto start_time = ailego::Monotime::MilliSeconds();\n auto iter = holder->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Create iterator for holder failed\");\n return IndexError_Runtime;\n }\n while (iter->is_valid()) {\n ret = metric_->train(iter->data(), meta_.dimension());\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw build measure train failed, ret=%d\", ret);\n return ret;\n }\n iter->next();\n ++trained_count;\n }\n trained_cost_time = ailego::Monotime::MilliSeconds() - start_time;\n }\n stats_.set_trained_count(trained_count);\n stats_.set_trained_costtime(trained_cost_time);\n state_ = BUILD_STATE_TRAINED;\n\n LOG_INFO(\"End HnswRabitqBuilder::train\");\n\n return 0;\n}\n\nint HnswRabitqBuilder::train_converter_and_load_reformer(\n IndexHolder::Pointer holder) {\n // Train converter (KMeans clustering)\n int ret = converter_->train(holder);\n if (ret != 0) {\n LOG_ERROR(\"Failed to train RabitqConverter: %d\", ret);\n return ret;\n }\n auto memory_dumper = IndexFactory::CreateDumper(\"MemoryDumper\");\n memory_dumper->init(ailego::Params());\n std::string file_id = ailego::StringHelper::Concat(\n \"rabitq_converter_\", ailego::Monotime::MilliSeconds(), rand());\n ret = memory_dumper->create(file_id);\n if (ret != 0) {\n LOG_ERROR(\"Failed to create memory dumper: %d\", ret);\n return ret;\n }\n // Release memory\n AILEGO_DEFER([&file_id]() { IndexMemory::Instance()->remove(file_id); });\n ret = converter_->dump(memory_dumper);\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump RabitqConverter: %d\", ret);\n return ret;\n }\n ret = memory_dumper->close();\n if (ret != 0) {\n LOG_ERROR(\"Failed to close memory dumper: %d\", ret);\n return ret;\n }\n\n reformer_ = std::make_shared();\n ailego::Params reformer_params;\n reformer_params.set(PARAM_RABITQ_METRIC_NAME, meta_.metric_name());\n ret = reformer_->init(reformer_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize RabitqReformer: %d\", ret);\n return ret;\n }\n auto memory_storage = IndexFactory::CreateStorage(\"MemoryReadStorage\");\n ret = memory_storage->open(file_id, false);\n if (ret != 0) {\n LOG_ERROR(\"Failed to open memory storage: %d\", ret);\n return ret;\n }\n ret = reformer_->load(memory_storage);\n if (ret != 0) {\n LOG_ERROR(\"Failed to load RabitqReformer: %d\", ret);\n return ret;\n }\n return 0;\n}\n\nint HnswRabitqBuilder::train(const IndexTrainer::Pointer & /*trainer*/) {\n if (state_ != BUILD_STATE_INITED) {\n LOG_ERROR(\"Init the builder before HnswRabitqBuilder::train\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"Begin HnswRabitqBuilder::train by trainer\");\n\n stats_.set_trained_count(0UL);\n stats_.set_trained_costtime(0UL);\n state_ = BUILD_STATE_TRAINED;\n\n LOG_INFO(\"End HnswRabitqBuilder::train by trainer\");\n\n return 0;\n}\n\nint HnswRabitqBuilder::build(IndexThreads::Pointer threads,\n IndexHolder::Pointer holder) {\n if (state_ != BUILD_STATE_TRAINED) {\n LOG_ERROR(\"Train the index before HnswRabitqBuilder::build\");\n return IndexError_NoReady;\n }\n\n if (!holder) {\n LOG_ERROR(\"Input holder is nullptr while building index\");\n return IndexError_InvalidArgument;\n }\n if (!holder->is_matched(meta_)) {\n LOG_ERROR(\"Input holder doesn't match index meta while building index\");\n return IndexError_Mismatch;\n }\n IndexProvider::Pointer provider =\n std::dynamic_pointer_cast(holder);\n if (!provider) {\n LOG_ERROR(\"Rabitq builder expect IndexProvider\");\n return IndexError_InvalidArgument;\n }\n\n if (!threads) {\n threads = std::make_shared(thread_cnt_, false);\n }\n\n auto start_time = ailego::Monotime::MilliSeconds();\n LOG_INFO(\"Begin HnswRabitqBuilder::build\");\n\n if (holder->count() != static_cast(-1)) {\n LOG_DEBUG(\"HnswRabitqBuilder holder documents count %lu\", holder->count());\n int ret = entity_.reserve_space(holder->count());\n if (ret != 0) {\n LOG_ERROR(\"HnswBuilde reserver space failed\");\n return ret;\n }\n }\n auto iter = holder->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Create iterator for holder failed\");\n return IndexError_Runtime;\n }\n int ret;\n error_ = false;\n IndexQueryMeta ometa;\n ometa.set_meta(holder->data_type(), holder->dimension());\n while (iter->is_valid()) {\n const void *vec = iter->data();\n // quantize vector\n std::string converted_vector;\n IndexQueryMeta converted_meta;\n ret = reformer_->convert(vec, ometa, &converted_vector, &converted_meta);\n if (ret != 0) {\n LOG_ERROR(\"Rabitq hnsw convert failed, ret=%d\", ret);\n return ret;\n }\n\n\n level_t level = alg_->get_random_level();\n node_id_t id;\n\n if (converted_vector.size() != entity_.vector_size()) {\n LOG_ERROR(\n \"Converted vector size %zu is not equal to entity vector size %zu\",\n converted_vector.size(), entity_.vector_size());\n return IndexError_InvalidArgument;\n }\n ret = entity_.add_vector(level, iter->key(), converted_vector.data(), &id);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n iter->next();\n }\n\n LOG_INFO(\"Finished save vector, start build graph...\");\n\n auto task_group = threads->make_group();\n if (!task_group) {\n LOG_ERROR(\"Failed to create task group\");\n return IndexError_Runtime;\n }\n\n std::atomic finished{0};\n for (size_t i = 0; i < threads->count(); ++i) {\n task_group->submit(ailego::Closure ::New(this, &HnswRabitqBuilder::do_build,\n i, threads->count(), provider,\n &finished));\n }\n\n while (!task_group->is_finished()) {\n std::unique_lock lk(mutex_);\n cond_.wait_until(lk, std::chrono::system_clock::now() +\n std::chrono::seconds(check_interval_secs_));\n if (error_.load(std::memory_order_acquire)) {\n LOG_ERROR(\"Failed to build index while waiting finish\");\n return errcode_;\n }\n LOG_INFO(\"Built cnt %zu, finished percent %.3f%%\",\n static_cast(finished.load()),\n finished.load() * 100.0f / entity_.doc_cnt());\n }\n if (error_.load(std::memory_order_acquire)) {\n LOG_ERROR(\"Failed to build index while waiting finish\");\n return errcode_;\n }\n task_group->wait_finish();\n\n stats_.set_built_count(finished.load());\n stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time);\n\n state_ = BUILD_STATE_BUILT;\n LOG_INFO(\"End HnswRabitqBuilder::build with RaBitQ quantization\");\n return 0;\n}\n\nvoid HnswRabitqBuilder::do_build(node_id_t idx, size_t step_size,\n IndexProvider::Pointer provider,\n std::atomic *finished) {\n AILEGO_DEFER([&]() {\n std::lock_guard latch(mutex_);\n cond_.notify_one();\n });\n HnswRabitqContext *ctx = new (std::nothrow) HnswRabitqContext(\n meta_.dimension(), metric_,\n std::shared_ptr(&entity_, [](HnswRabitqEntity *) {}));\n if (ailego_unlikely(ctx == nullptr)) {\n if (!error_.exchange(true)) {\n LOG_ERROR(\"Failed to create context\");\n errcode_ = IndexError_NoMemory;\n }\n return;\n }\n HnswRabitqContext::Pointer auto_ptr(ctx);\n ctx->set_provider(std::move(provider));\n ctx->set_max_scan_num(entity_.doc_cnt());\n int ret = ctx->init(HnswRabitqContext::kBuilderContext);\n if (ret != 0) {\n if (!error_.exchange(true)) {\n LOG_ERROR(\"Failed to init context\");\n errcode_ = IndexError_Runtime;\n }\n return;\n }\n\n for (node_id_t id = idx; id < entity_.doc_cnt(); id += step_size) {\n ctx->reset_query(ctx->dist_calculator().get_vector(id));\n ret = alg_->add_node(id, entity_.get_level(id), ctx);\n if (ailego_unlikely(ret != 0)) {\n if (!error_.exchange(true)) {\n LOG_ERROR(\"Hnsw graph add node failed\");\n errcode_ = ret;\n }\n return;\n }\n ctx->clear();\n (*finished)++;\n }\n}\n\nint HnswRabitqBuilder::dump(const IndexDumper::Pointer &dumper) {\n if (state_ != BUILD_STATE_BUILT) {\n LOG_INFO(\"Build the index before HnswRabitqBuilder::dump\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"Begin HnswRabitqBuilder::dump\");\n\n meta_.set_searcher(\"HnswRabitqSearcher\", HnswRabitqEntity::kRevision,\n ailego::Params());\n auto start_time = ailego::Monotime::MilliSeconds();\n\n int ret = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return ret;\n }\n\n // Dump RaBitQ centroids first\n if (converter_) {\n ret = converter_->dump(dumper);\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump RabitqConverter: %d\", ret);\n return ret;\n }\n LOG_INFO(\"RaBitQ centroids dumped: %zu bytes, cost %zu ms\",\n converter_->stats().dumped_size(),\n static_cast(converter_->stats().dumped_costtime()));\n }\n\n ret = entity_.dump(dumper);\n if (ret != 0) {\n LOG_ERROR(\"HnswRabitqBuilder dump index failed\");\n return ret;\n }\n\n stats_.set_dumped_count(entity_.doc_cnt());\n stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time);\n\n LOG_INFO(\"End HnswRabitqBuilder::dump\");\n return 0;\n}\n\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"zvec/core/framework/index_builder.h\"\n#include \"zvec/core/framework/index_converter.h\"\n#include \"zvec/core/framework/index_reformer.h\"\n#include \"hnsw_rabitq_algorithm.h\"\n#include \"hnsw_rabitq_builder_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswRabitqBuilder : public IndexBuilder {\n public:\n //! Constructor\n HnswRabitqBuilder();\n\n //! Initialize the builder\n virtual int init(const IndexMeta &meta,\n const ailego::Params ¶ms) override;\n\n //! Cleanup the builder\n virtual int cleanup(void) override;\n\n //! Train the data\n virtual int train(IndexThreads::Pointer,\n IndexHolder::Pointer holder) override;\n\n //! Train the data\n virtual int train(const IndexTrainer::Pointer &trainer) override;\n\n\n //! Build the index\n virtual int build(IndexThreads::Pointer threads,\n IndexHolder::Pointer holder) override;\n\n //! Dump index into storage\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n virtual const Stats &stats(void) const override {\n return stats_;\n }\n\n private:\n void do_build(node_id_t idx, size_t step_size,\n IndexProvider::Pointer provider,\n std::atomic *finished);\n\n int train_converter_and_load_reformer(IndexHolder::Pointer holder);\n\n constexpr static uint32_t kDefaultLogIntervalSecs = 15U;\n constexpr static uint32_t kMaxNeighborCnt = 65535;\n\n private:\n enum BUILD_STATE {\n BUILD_STATE_INIT = 0,\n BUILD_STATE_INITED = 1,\n BUILD_STATE_TRAINED = 2,\n BUILD_STATE_BUILT = 3\n };\n\n HnswRabitqBuilderEntity entity_{};\n HnswRabitqAlgorithm::UPointer alg_; // impl graph algorithm\n uint32_t thread_cnt_{0};\n uint32_t min_neighbor_cnt_{0};\n uint32_t upper_max_neighbor_cnt_{\n HnswRabitqEntity::kDefaultUpperMaxNeighborCnt};\n uint32_t l0_max_neighbor_cnt_{HnswRabitqEntity::kDefaultL0MaxNeighborCnt};\n uint32_t ef_construction_{HnswRabitqEntity::kDefaultEfConstruction};\n uint32_t scaling_factor_{HnswRabitqEntity::kDefaultScalingFactor};\n uint32_t check_interval_secs_{kDefaultLogIntervalSecs};\n\n int errcode_{0};\n std::atomic_bool error_{false};\n IndexMeta meta_{};\n IndexMetric::Pointer metric_{};\n IndexConverter::Pointer converter_{}; // RaBitQ converter\n IndexReformer::Pointer reformer_{}; // RaBitQ reformer\n std::mutex mutex_{};\n std::condition_variable cond_{};\n Stats stats_{};\n\n BUILD_STATE state_{BUILD_STATE_INIT};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_builder_entity.h\"\n#include \n#include \n#include \"utility/sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswRabitqBuilderEntity::HnswRabitqBuilderEntity() {\n update_ep_and_level(kInvalidNodeId, 0U);\n}\n\nint HnswRabitqBuilderEntity::cleanup() {\n memory_quota_ = 0UL;\n neighbors_size_ = 0U;\n upper_neighbors_size_ = 0U;\n padding_size_ = 0U;\n vectors_buffer_.clear();\n keys_buffer_.clear();\n neighbors_buffer_.clear();\n upper_neighbors_buffer_.clear();\n neighbors_index_.clear();\n\n vectors_buffer_.shrink_to_fit();\n keys_buffer_.shrink_to_fit();\n neighbors_buffer_.shrink_to_fit();\n upper_neighbors_buffer_.shrink_to_fit();\n neighbors_index_.shrink_to_fit();\n\n this->HnswRabitqEntity::cleanup();\n\n return 0;\n}\n\nint HnswRabitqBuilderEntity::init() {\n size_t size = vector_size();\n\n //! aligned size to 32\n set_node_size(AlignSize(size));\n //! if node size is aligned to 1k, the build performance will downgrade\n if (node_size() % 1024 == 0) {\n set_node_size(AlignSize(node_size() + 1));\n }\n\n padding_size_ = node_size() - size;\n\n neighbors_size_ = neighbors_size();\n upper_neighbors_size_ = upper_neighbors_size();\n\n return 0;\n}\n\nint HnswRabitqBuilderEntity::reserve_space(size_t docs) {\n if (memory_quota_ > 0 && (node_size() * docs + neighbors_size_ * docs +\n sizeof(NeighborIndex) * docs >\n memory_quota_)) {\n return IndexError_NoMemory;\n }\n\n vectors_buffer_.reserve(node_size() * docs);\n keys_buffer_.reserve(sizeof(key_t) * docs);\n neighbors_buffer_.reserve(neighbors_size_ * docs);\n neighbors_index_.reserve(docs);\n\n return 0;\n}\n\nint HnswRabitqBuilderEntity::add_vector(level_t level, key_t key,\n const void *vec, node_id_t *id) {\n if (memory_quota_ > 0 &&\n (vectors_buffer_.capacity() + keys_buffer_.capacity() +\n neighbors_buffer_.capacity() + upper_neighbors_buffer_.capacity() +\n neighbors_index_.capacity() * sizeof(NeighborIndex)) > memory_quota_) {\n LOG_ERROR(\"Add vector failed, used memory exceed quota, cur_doc=%zu\",\n static_cast(doc_cnt()));\n return IndexError_NoMemory;\n }\n\n vectors_buffer_.append(reinterpret_cast(vec), vector_size());\n vectors_buffer_.append(padding_size_, '\\0');\n keys_buffer_.append(reinterpret_cast(&key), sizeof(key));\n\n // init level 0 neighbors\n neighbors_buffer_.append(neighbors_size_, '\\0');\n\n neighbors_index_.emplace_back(upper_neighbors_buffer_.size(), level);\n\n // init upper layer neighbors\n for (level_t cur_level = 1; cur_level <= level; ++cur_level) {\n upper_neighbors_buffer_.append(upper_neighbors_size_, '\\0');\n }\n\n *id = (*mutable_doc_cnt())++;\n\n return 0;\n}\n\nkey_t HnswRabitqBuilderEntity::get_key(node_id_t id) const {\n return *(reinterpret_cast(keys_buffer_.data() +\n id * sizeof(key_t)));\n}\n\nconst void *HnswRabitqBuilderEntity::get_vector(node_id_t id) const {\n return vectors_buffer_.data() + id * node_size();\n}\n\nint HnswRabitqBuilderEntity::get_vector(\n const node_id_t id, IndexStorage::MemoryBlock &block) const {\n const void *vec = get_vector(id);\n block.reset((void *)vec);\n return 0;\n}\n\nint HnswRabitqBuilderEntity::get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const {\n for (uint32_t i = 0; i < count; ++i) {\n vecs[i] = vectors_buffer_.data() + ids[i] * node_size();\n }\n\n return 0;\n}\n\nint HnswRabitqBuilderEntity::get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const {\n std::vector vecs(count);\n get_vector(ids, count, vecs.data());\n for (uint32_t i = 0; i < count; ++i) {\n vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i]));\n }\n return 0;\n}\n\nconst Neighbors HnswRabitqBuilderEntity::get_neighbors(level_t level,\n node_id_t id) const {\n const NeighborsHeader *hd = get_neighbor_header(level, id);\n return {hd->neighbor_cnt, hd->neighbors};\n}\n\nint HnswRabitqBuilderEntity::update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) {\n NeighborsHeader *hd =\n const_cast(get_neighbor_header(level, id));\n for (size_t i = 0; i < neighbors.size(); ++i) {\n hd->neighbors[i] = neighbors[i].first;\n }\n hd->neighbor_cnt = neighbors.size();\n\n // std::cout << \"id: \" << id << \", neighbour, id: \";\n // for (size_t i = 0; i < neighbors.size(); ++i) {\n // if (i == neighbors.size()-1)\n // std::cout << neighbors[i].first << \", score:\" << neighbors[i].second <<\n // std::endl;\n // else\n // std::cout << neighbors[i].first << \", score:\" << neighbors[i].second <<\n // \", id: \";\n // }\n\n return 0;\n}\n\nvoid HnswRabitqBuilderEntity::add_neighbor(level_t level, node_id_t id,\n uint32_t /*size*/,\n node_id_t neighbor_id) {\n NeighborsHeader *hd =\n const_cast(get_neighbor_header(level, id));\n hd->neighbors[hd->neighbor_cnt++] = neighbor_id;\n\n return;\n}\n\nint HnswRabitqBuilderEntity::dump(const IndexDumper::Pointer &dumper) {\n key_t *keys =\n reinterpret_cast(const_cast(keys_buffer_.data()));\n auto ret =\n dump_segments(dumper, keys, [&](node_id_t id) { return get_level(id); });\n if (ailego_unlikely(ret < 0)) {\n return ret;\n }\n\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"hnsw_rabitq_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswRabitqBuilderEntity : public HnswRabitqEntity {\n public:\n //! Add vector and key to hnsw entity, and local id will be saved to id\n virtual int add_vector(level_t level, key_t key, const void *vec,\n node_id_t *id) override;\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const override;\n\n //! Get vector feature data by key\n virtual const void *get_vector(node_id_t id) const override;\n\n //! Batch get vectors feature data by keys\n virtual int get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const override;\n\n virtual int get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n virtual int get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const override;\n\n //! Get the node id's neighbors on graph level\n const NeighborsHeader *get_neighbor_header(level_t level,\n node_id_t id) const {\n if (level == 0) {\n return reinterpret_cast(\n neighbors_buffer_.data() + neighbors_size_ * id);\n } else {\n size_t offset = neighbors_index_[id].offset;\n return reinterpret_cast(\n upper_neighbors_buffer_.data() + offset +\n (level - 1) * upper_neighbors_size_);\n }\n }\n\n //! Get the node id's neighbors on graph level\n virtual const Neighbors get_neighbors(level_t level,\n node_id_t id) const override;\n\n //! Replace node id in level's neighbors\n virtual int update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors)\n override;\n\n //! add a neighbor to id in graph level\n virtual void add_neighbor(level_t level, node_id_t id, uint32_t size,\n node_id_t neighbor_id) override;\n\n //! Dump the hnsw graph to dumper\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Cleanup the entity\n virtual int cleanup(void) override;\n\n public:\n //! Constructor\n HnswRabitqBuilderEntity();\n\n //! Get the node graph level by id\n level_t get_level(node_id_t id) const {\n return neighbors_index_[id].level;\n }\n\n //! Init builerEntity\n int init();\n\n //! reserve buffer space for documents\n //! @param docs number of documents\n int reserve_space(size_t docs);\n\n //! Set memory quota params\n inline void set_memory_quota(size_t memory_quota) {\n memory_quota_ = memory_quota;\n }\n\n //! Get neighbors size\n inline size_t neighbors_size() const {\n return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t);\n }\n\n //! Get upper neighbors size\n inline size_t upper_neighbors_size() const {\n return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t);\n }\n\n public:\n HnswRabitqBuilderEntity(const HnswRabitqBuilderEntity &) = delete;\n HnswRabitqBuilderEntity &operator=(const HnswRabitqBuilderEntity &) = delete;\n\n private:\n friend class HnswRabitqSearcherEntity;\n //! class internal used only\n struct NeighborIndex {\n NeighborIndex(size_t off, level_t l) : offset(off), level(l) {}\n uint64_t offset : 48;\n uint64_t level : 16;\n };\n\n std::string vectors_buffer_{}; // aligned vectors\n std::string keys_buffer_{}; // aligned vectors\n std::string neighbors_buffer_{}; // level 0 neighbors buffer\n std::string upper_neighbors_buffer_{}; // upper layer neighbors buffer\n\n std::string sparse_data_buffer_{}; // aligned spase data buffer\n size_t sparse_data_offset_{0}; //\n\n // upper layer offset + level in upper_neighbors_buffer_\n std::vector neighbors_index_{};\n size_t memory_quota_{0UL};\n size_t neighbors_size_{0U}; // level 0 neighbors size\n size_t upper_neighbors_size_{0U}; // level 0 neighbors size\n size_t padding_size_{}; // padding size for each vector element\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_chunk.h\"\n#include \n#include \n#include \n#include \n#include \"zvec/core/framework/index_error.h\"\n#include \"zvec/core/framework/index_helper.h\"\n#include \"zvec/core/framework/index_logger.h\"\n#include \"zvec/core/framework/index_streamer.h\"\n\nnamespace zvec {\nnamespace core {\n\nint HnswRabitqChunkBroker::init_storage(size_t chunk_size) {\n chunk_meta_.clear();\n chunk_meta_.chunk_size = chunk_size;\n chunk_meta_.create_time = ailego::Realtime::Seconds();\n stats_.set_create_time(chunk_meta_.create_time);\n chunk_meta_.update_time = ailego::Realtime::Seconds();\n stats_.set_update_time(chunk_meta_.update_time);\n\n //! alloc meta chunk\n size_t size = sizeof(HnswChunkMeta);\n size = (size + page_mask_) & (~page_mask_);\n const std::string segment_id =\n make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId);\n int ret = stg_->append(segment_id, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Storage append segment failed for %s\", IndexError::What(ret));\n return ret;\n }\n chunk_meta_segment_ = get_chunk(CHUNK_TYPE_META, kDefaultChunkSeqId);\n if (ailego_unlikely(!chunk_meta_segment_)) {\n LOG_ERROR(\"Get meta segment failed\");\n return IndexError_Runtime;\n }\n\n //! update meta info and write to storage\n chunk_meta_.chunk_cnts[CHUNK_TYPE_META] += 1;\n chunk_meta_.total_size += size;\n (*stats_.mutable_index_size()) += size;\n size = chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta));\n if (ailego_unlikely(size != sizeof(HnswChunkMeta))) {\n LOG_ERROR(\"Storage write data failed, wsize=%zu\", size);\n return IndexError_WriteData;\n }\n\n return 0;\n}\n\nint HnswRabitqChunkBroker::load_storage(size_t chunk_size) {\n IndexStorage::MemoryBlock data_block;\n size_t size = chunk_meta_segment_->read(0UL, data_block,\n chunk_meta_segment_->data_size());\n if (size != sizeof(HnswChunkMeta)) {\n LOG_ERROR(\"Invalid hnsw meta chunk, read size=%zu chunk size=%zu\", size,\n chunk_meta_segment_->data_size());\n return IndexError_InvalidFormat;\n }\n std::memcpy(&chunk_meta_, data_block.data(), size);\n if (chunk_meta_.chunk_size != chunk_size) {\n LOG_ERROR(\n \"Params hnsw chunk size=%zu mismatch from previous %zu \"\n \"in index\",\n chunk_size, (size_t)chunk_meta_.chunk_size);\n return IndexError_Mismatch;\n }\n\n *stats_.mutable_check_point() = stg_->check_point();\n stats_.set_revision_id(chunk_meta_.revision_id);\n stats_.set_update_time(chunk_meta_.update_time);\n stats_.set_create_time(chunk_meta_.create_time);\n\n char create_time[32];\n char update_time[32];\n ailego::Realtime::Gmtime(chunk_meta_.create_time, \"%Y-%m-%d %H:%M:%S\",\n create_time, sizeof(create_time));\n ailego::Realtime::Gmtime(chunk_meta_.update_time, \"%Y-%m-%d %H:%M:%S\",\n update_time, sizeof(update_time));\n LOG_DEBUG(\n \"Load index, indexSize=%zu chunkSize=%zu nodeChunks=%zu \"\n \"upperNeighborChunks=%zu revisionId=%zu \"\n \"createTime=%s updateTime=%s\",\n (size_t)chunk_meta_.total_size, (size_t)chunk_meta_.chunk_size,\n (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_NODE],\n (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_UPPER_NEIGHBOR],\n (size_t)chunk_meta_.revision_id, create_time, update_time);\n\n return 0;\n}\n\nint HnswRabitqChunkBroker::open(IndexStorage::Pointer stg,\n size_t max_index_size, size_t chunk_size,\n bool check_crc) {\n if (ailego_unlikely(stg_)) {\n LOG_ERROR(\"An storage instance is already opened\");\n return IndexError_Duplicate;\n }\n stg_ = std::move(stg);\n if (stg_->isHugePage()) {\n page_mask_ = ailego::MemoryHelper::HugePageSize() - 1;\n } else {\n page_mask_ = ailego::MemoryHelper::PageSize() - 1;\n }\n check_crc_ = check_crc;\n max_chunks_size_ = max_index_size;\n dirty_ = false;\n\n const std::string segment_id =\n make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId);\n chunk_meta_segment_ = stg_->get(segment_id);\n if (!chunk_meta_segment_) {\n LOG_DEBUG(\"Create new index\");\n return init_storage(chunk_size);\n }\n\n return load_storage(chunk_size);\n}\n\nint HnswRabitqChunkBroker::close(void) {\n flush(0UL);\n\n stg_.reset();\n check_crc_ = false;\n dirty_ = false;\n\n return 0;\n}\n\nint HnswRabitqChunkBroker::flush(uint64_t checkpoint) {\n ailego_assert_with(chunk_meta_segment_, \"invalid meta segment\");\n\n chunk_meta_.update_time = ailego::Realtime::Seconds();\n stats_.set_update_time(chunk_meta_.update_time);\n\n size_t size =\n chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta));\n if (ailego_unlikely(size != sizeof(HnswChunkMeta))) {\n LOG_ERROR(\"Storage write data failed, wsize=%zu\", size);\n }\n\n stg_->refresh(checkpoint);\n int ret = stg_->flush();\n if (ret == 0) {\n (*stats_.mutable_check_point()) = checkpoint;\n } else {\n LOG_ERROR(\"Storage flush failed for %s\", IndexError::What(ret));\n }\n return ret;\n}\n\nstd::pair HnswRabitqChunkBroker::alloc_chunk(\n int type, uint64_t seq_id, size_t size) {\n ailego_assert_with(type < CHUNK_TYPE_MAX, \"chunk type overflow\");\n\n Chunk::Pointer chunk;\n if (ailego_unlikely(!stg_)) {\n LOG_ERROR(\"Init storage first\");\n return std::make_pair(IndexError_Uninitialized, chunk);\n }\n\n //! check exist a empty chunk with the same name\n chunk = get_chunk(type, seq_id);\n if (chunk) {\n if (ailego_unlikely(chunk->capacity() == size &&\n chunk->data_size() == 0UL)) {\n LOG_ERROR(\"Exist invalid chunk size %zu, expect size %zu\",\n chunk->capacity(), size);\n chunk.reset();\n return std::make_pair(IndexError_Runtime, chunk);\n }\n return std::make_pair(0, chunk);\n }\n //! align to page size\n size = (size + page_mask_) & (~page_mask_);\n if (ailego_unlikely(chunk_meta_.total_size + size >= max_chunks_size_)) {\n LOG_ERROR(\"No space to new a chunk, curIndexSize=%zu allocSize=%zu\",\n (size_t)chunk_meta_.total_size, size);\n return std::make_pair(IndexError_IndexFull, chunk);\n }\n\n std::string segment_id = make_segment_id(type, seq_id);\n int ret = stg_->append(segment_id, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Storage append segment failed for %s\", IndexError::What(ret));\n return std::make_pair(ret, chunk);\n }\n chunk_meta_.chunk_cnts[type] += 1;\n chunk_meta_.total_size += size;\n (*stats_.mutable_index_size()) += size;\n\n size = chunk_meta_segment_->write(0UL, &chunk_meta_, sizeof(HnswChunkMeta));\n if (ailego_unlikely(size != sizeof(HnswChunkMeta))) {\n LOG_ERROR(\"Storage append segment failed, wsize=%zu\", size);\n }\n\n chunk = get_chunk(type, seq_id);\n return std::make_pair(chunk ? 0 : IndexError_NoMemory, chunk);\n}\n\nChunk::Pointer HnswRabitqChunkBroker::get_chunk(int type,\n uint64_t seq_id) const {\n std::string segment_id = make_segment_id(type, seq_id);\n return stg_->get(segment_id);\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_chunk.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"zvec/core/framework/index_error.h\"\n#include \"zvec/core/framework/index_logger.h\"\n#include \"zvec/core/framework/index_storage.h\"\n#include \"zvec/core/framework/index_streamer.h\"\n\nnamespace zvec {\nnamespace core {\n\nusing Chunk = IndexStorage::Segment;\n\nclass HnswRabitqChunkBroker {\n public:\n typedef std::shared_ptr Pointer;\n\n enum CHUNK_TYPE {\n CHUNK_TYPE_HEADER = 1,\n CHUNK_TYPE_META = 2,\n CHUNK_TYPE_NODE = 3,\n CHUNK_TYPE_UPPER_NEIGHBOR = 4,\n CHUNK_TYPE_NEIGHBOR_INDEX = 5,\n CHUNK_TYPE_SPARSE_NODE = 6,\n CHUNK_TYPE_MAX = 8\n };\n static constexpr size_t kDefaultChunkSeqId = 0UL;\n\n HnswRabitqChunkBroker(IndexStreamer::Stats &stats) : stats_(stats) {}\n\n //! Open storage\n int open(IndexStorage::Pointer stg, size_t max_index_size, size_t chunk_size,\n bool check_crc);\n\n int close(void);\n\n int flush(uint64_t checkpoint);\n\n //! alloc a new chunk with size, not thread-safe\n std::pair alloc_chunk(int type, uint64_t seq_id,\n size_t size);\n\n //! alloc a new chunk with chunk size\n inline std::pair alloc_chunk(int type, uint64_t seq_id) {\n return alloc_chunk(type, seq_id, chunk_meta_.chunk_size);\n }\n\n Chunk::Pointer get_chunk(int type, uint64_t seq_id) const;\n\n inline size_t get_chunk_cnt(int type) const {\n ailego_assert_with(type < CHUNK_TYPE_MAX, \"chunk type overflow\");\n return chunk_meta_.chunk_cnts[type];\n }\n\n inline bool dirty(void) const {\n return dirty_;\n }\n\n inline void mark_dirty(void) {\n if (!dirty_) {\n dirty_ = true;\n chunk_meta_.revision_id += 1;\n stats_.set_revision_id(chunk_meta_.revision_id);\n }\n }\n\n const IndexStorage::Pointer storage(void) const {\n return stg_;\n }\n\n private:\n HnswRabitqChunkBroker(const HnswRabitqChunkBroker &) = delete;\n HnswRabitqChunkBroker &operator=(const HnswRabitqChunkBroker &) = delete;\n\n struct HnswChunkMeta {\n HnswChunkMeta(void) {\n memset(this, 0, sizeof(HnswChunkMeta));\n }\n void clear() {\n memset(this, 0, sizeof(HnswChunkMeta));\n }\n\n uint64_t chunk_cnts[CHUNK_TYPE_MAX];\n uint64_t chunk_size; // size of per chunk\n uint64_t total_size; // total size of allocated chunk\n uint64_t revision_id; // index revision\n uint64_t create_time;\n uint64_t update_time;\n uint64_t reserved[3];\n };\n\n static_assert(sizeof(HnswChunkMeta) % 32 == 0,\n \"HnswChunkMeta must be aligned with 32 bytes\");\n\n //! Init the storage after open an empty index\n int init_storage(size_t chunk_size);\n\n //! Load index from storage\n int load_storage(size_t chunk_size);\n\n static inline const std::string make_segment_id(int type, uint64_t seq_id) {\n return \"HnswT\" + ailego::StringHelper::ToString(type) + \"S\" +\n ailego::StringHelper::ToString(seq_id);\n }\n\n private:\n IndexStreamer::Stats &stats_;\n HnswChunkMeta chunk_meta_{};\n size_t page_mask_{0UL};\n size_t max_chunks_size_{0UL};\n IndexStorage::Pointer stg_{};\n IndexStorage::Segment::Pointer chunk_meta_segment_{};\n bool check_crc_{false};\n bool dirty_{false}; // set as true if index is modified , the flag\n // will not be cleared even if flushed\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_context.h\"\n#include \n#include \"hnsw_rabitq_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswRabitqContext::HnswRabitqContext(size_t dimension,\n const IndexMetric::Pointer &metric,\n const HnswRabitqEntity::Pointer &entity)\n : IndexContext(metric),\n entity_(entity),\n add_dc_(entity_.get(), metric, dimension) {}\n\nHnswRabitqContext::HnswRabitqContext(const IndexMetric::Pointer &metric,\n const HnswRabitqEntity::Pointer &entity)\n : IndexContext(metric), entity_(entity), add_dc_(entity_.get(), metric) {}\n\nHnswRabitqContext::~HnswRabitqContext() {\n visit_filter_.destroy();\n}\n\nint HnswRabitqContext::init(ContextType type) {\n int ret;\n uint32_t doc_cnt;\n\n type_ = type;\n\n switch (type) {\n case kBuilderContext:\n ret = visit_filter_.init(VisitFilter::ByteMap, entity_->doc_cnt(),\n max_scan_num_, negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n candidates_.limit(max_scan_num_);\n update_heap_.limit(entity_->l0_neighbor_cnt() + 1);\n break;\n\n case kSearcherContext:\n ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_,\n negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n candidates_.limit(max_scan_num_);\n break;\n\n case kStreamerContext:\n // maxScanNum is unknown if inited from streamer, so the docCnt may\n // change. we need to compute maxScanNum by scan ratio, and preserve\n // max_doc_cnt space from visit filter\n doc_cnt = entity_->doc_cnt();\n max_scan_num_ = compute_max_scan_num(doc_cnt);\n reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt);\n ret = visit_filter_.init(filter_mode_, reserve_max_doc_cnt_,\n max_scan_num_, negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n\n update_heap_.limit(entity_->l0_neighbor_cnt() + 1);\n candidates_.limit(max_scan_num_);\n\n check_need_adjuct_ctx();\n break;\n\n default:\n LOG_ERROR(\"Init context failed\");\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nint HnswRabitqContext::update(const ailego::Params ¶ms) {\n auto update_visit_filter_param = [&]() {\n bool need_update = false;\n std::string p;\n switch (type_) {\n case kSearcherContext:\n p = PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_ENABLE;\n break;\n case kStreamerContext:\n p = PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_ENABLE;\n break;\n }\n\n if (params.has(p)) {\n bool bf_enabled;\n params.get(p, &bf_enabled);\n if (bf_enabled ^ (filter_mode_ == VisitFilter::BloomFilter)) {\n need_update = true;\n filter_mode_ =\n bf_enabled ? VisitFilter::BloomFilter : VisitFilter::ByteMap;\n }\n }\n\n float prob = negative_probability_;\n p.clear();\n switch (type_) {\n case kSearcherContext:\n p = PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB;\n break;\n case kStreamerContext:\n p = PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB;\n break;\n }\n params.get(p, &prob);\n if (filter_mode_ == VisitFilter::BloomFilter &&\n std::abs(prob - negative_probability_) > 1e-6) {\n need_update = true;\n }\n if (need_update) {\n visit_filter_.destroy();\n int max_doc_cnt = 0;\n if (type_ == kSearcherContext) {\n max_doc_cnt = entity_->doc_cnt();\n } else {\n max_doc_cnt = reserve_max_doc_cnt_;\n }\n int ret = visit_filter_.init(filter_mode_, max_doc_cnt, max_scan_num_,\n negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n }\n return 0;\n };\n\n switch (type_) {\n case kSearcherContext:\n if (params.has(PARAM_HNSW_RABITQ_SEARCHER_EF)) {\n params.get(PARAM_HNSW_RABITQ_SEARCHER_EF, &ef_);\n topk_heap_.limit(std::max(topk_, ef_));\n }\n\n if (params.has(PARAM_HNSW_RABITQ_SEARCHER_MAX_SCAN_RATIO)) {\n params.get(PARAM_HNSW_RABITQ_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_);\n max_scan_num_ =\n static_cast(max_scan_ratio_ * entity_->doc_cnt());\n max_scan_num_ = std::max(10000U, max_scan_num_);\n }\n\n if (params.has(PARAM_HNSW_RABITQ_SEARCHER_BRUTE_FORCE_THRESHOLD)) {\n params.get(PARAM_HNSW_RABITQ_SEARCHER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n }\n\n return update_visit_filter_param();\n\n case kStreamerContext:\n if (params.has(PARAM_HNSW_RABITQ_STREAMER_EF)) {\n params.get(PARAM_HNSW_RABITQ_STREAMER_EF, &ef_);\n topk_heap_.limit(std::max(topk_, ef_));\n }\n params.get(PARAM_HNSW_RABITQ_STREAMER_EF, &ef_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_);\n if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) {\n LOG_ERROR(\"[%s] must be in range (0.0f,1.0f]\",\n PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO.c_str());\n return IndexError_InvalidArgument;\n }\n if (max_scan_limit_ < min_scan_limit_) {\n LOG_ERROR(\"[%s] must be >= [%s]\",\n PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT.c_str(),\n PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (params.has(PARAM_HNSW_RABITQ_STREAMER_BRUTE_FORCE_THRESHOLD)) {\n params.get(PARAM_HNSW_RABITQ_STREAMER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n }\n\n return update_visit_filter_param();\n\n default:\n LOG_ERROR(\"update context failed, type=%zu\", static_cast(type_));\n return IndexError_Runtime;\n }\n}\n\nint HnswRabitqContext::update_context(ContextType type, const IndexMeta &meta,\n const IndexMetric::Pointer &metric,\n const HnswRabitqEntity::Pointer &entity,\n uint32_t magic_num) {\n uint32_t doc_cnt;\n\n if (ailego_unlikely(type != type_)) {\n LOG_ERROR(\n \"HnswRabitqContext doesn't support shared by different type, \"\n \"src=%u dst=%u\",\n type_, type);\n return IndexError_Unsupported;\n }\n\n magic_ = kInvalidMgic;\n\n // TODO: support change filter mode?\n switch (type) {\n case kBuilderContext:\n LOG_ERROR(\"BuildContext doesn't support update\");\n return IndexError_NotImplemented;\n\n case kSearcherContext:\n if (!visit_filter_.reset(entity->doc_cnt(), max_scan_num_)) {\n LOG_ERROR(\"Reset filter failed, mode %d\", visit_filter_.get_mode());\n return IndexError_Runtime;\n }\n\n candidates_.limit(max_scan_num_);\n topk_heap_.limit(std::max(topk_, ef_));\n break;\n\n case kStreamerContext:\n doc_cnt = entity->doc_cnt();\n max_scan_num_ = compute_max_scan_num(doc_cnt);\n reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt);\n if (!visit_filter_.reset(reserve_max_doc_cnt_, max_scan_num_)) {\n LOG_ERROR(\"Reset filter failed, mode %d\", visit_filter_.get_mode());\n return IndexError_Runtime;\n }\n\n update_heap_.limit(entity->l0_neighbor_cnt() + 1);\n candidates_.limit(max_scan_num_);\n topk_heap_.limit(std::max(topk_, ef_));\n break;\n\n default:\n LOG_ERROR(\"update context failed\");\n return IndexError_Runtime;\n }\n\n entity_ = entity;\n dc().update(entity_.get(), metric, meta.dimension());\n magic_ = magic_num;\n level_topks_.clear();\n\n return 0;\n}\n\nvoid HnswRabitqContext::fill_random_to_topk_full(void) {\n static std::mt19937 mt(\n std::chrono::system_clock::now().time_since_epoch().count());\n std::uniform_int_distribution dt(0, entity_->doc_cnt() - 1);\n std::function gen;\n node_id_t seqid;\n std::function myfilter = [](node_id_t) { return false; };\n if (this->filter().is_valid()) {\n myfilter = [&](node_id_t id) {\n return this->filter()(entity_->get_key(id));\n };\n }\n\n if (topk_heap_.limit() < entity_->doc_cnt() / 2) {\n gen = [&](void) { return dt(mt); };\n } else {\n // If topk limit is big value, gen sequential id from an random initial\n seqid = dt(mt);\n gen = [&](void) {\n seqid = seqid == (entity_->doc_cnt() - 1) ? 0 : (seqid + 1);\n return seqid;\n };\n }\n\n for (size_t i = 0; !topk_heap_.full() && i < entity_->doc_cnt(); ++i) {\n const auto id = gen();\n if (!visit_filter_.visited(id) && !myfilter(id)) {\n visit_filter_.set_visited(id);\n topk_heap_.emplace(id, dc().dist(id));\n }\n }\n return;\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_context.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"utility/visit_filter.h\"\n#include \"zvec/core/framework/index_context.h\"\n#include \"zvec/core/framework/index_provider.h\"\n#include \"hnsw_rabitq_dist_calculator.h\"\n#include \"hnsw_rabitq_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswRabitqContext : public IndexContext {\n public:\n //! Index Context Pointer\n typedef std::unique_ptr Pointer;\n\n enum ContextType {\n kUnknownContext = 0,\n kSearcherContext = 1,\n kBuilderContext = 2,\n kStreamerContext = 3\n };\n\n //! Construct\n HnswRabitqContext(size_t dimension, const IndexMetric::Pointer &metric,\n const HnswRabitqEntity::Pointer &entity);\n\n //! Construct\n HnswRabitqContext(const IndexMetric::Pointer &metric,\n const HnswRabitqEntity::Pointer &entity);\n\n //! Destructor\n virtual ~HnswRabitqContext();\n\n public:\n //! Set topk of search result\n virtual void set_topk(uint32_t val) override {\n topk_ = val;\n topk_heap_.limit(std::max(val, ef_));\n }\n\n //! Retrieve search result\n virtual const IndexDocumentList &result(void) const override {\n return results_[0];\n }\n\n //! Retrieve search result\n virtual const IndexDocumentList &result(size_t idx) const override {\n return results_[idx];\n }\n\n //! Retrieve result object for output\n virtual IndexDocumentList *mutable_result(size_t idx) override {\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n return &results_[idx];\n }\n\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(void) const override {\n return group_results_[0];\n }\n\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(\n size_t idx) const override {\n return group_results_[idx];\n }\n\n virtual uint32_t magic(void) const override {\n return magic_;\n }\n\n //! Set mode of debug\n virtual void set_debug_mode(bool enable) override {\n debug_mode_ = enable;\n }\n\n //! Retrieve mode of debug\n virtual bool debug_mode(void) const override {\n return this->debugging();\n }\n\n //! Retrieve string of debug\n virtual std::string debug_string(void) const override {\n char buf[4096];\n size_t size = snprintf(\n buf, sizeof(buf),\n \"scan_cnt=%zu,get_vector_cnt=%u,get_neighbors_cnt=%u,dup_node=%u\",\n get_scan_num(), stats_get_vector_cnt_, stats_get_neighbors_cnt_,\n stats_visit_dup_cnt_);\n return std::string(buf, size);\n }\n\n //! Update the parameters of context\n virtual int update(const ailego::Params ¶ms) override;\n\n public:\n //! Init context\n int init(ContextType type);\n\n //! Update context, the context may be shared by different searcher/streamer\n int update_context(ContextType type, const IndexMeta &meta,\n const IndexMetric::Pointer &metric,\n const HnswRabitqEntity::Pointer &entity,\n uint32_t magic_num);\n\n inline const HnswRabitqEntity &get_entity() const {\n return *entity_;\n }\n\n inline void resize_results(size_t size) {\n if (group_by_search()) {\n group_results_.resize(size);\n } else {\n results_.resize(size);\n }\n }\n\n inline void topk_to_result() {\n return topk_to_result(0);\n }\n\n //! Construct result from topk heap, result will be normalized\n inline void topk_to_result(uint32_t idx) {\n if (group_by_search()) {\n topk_to_group_result(idx);\n } else {\n topk_to_single_result(idx);\n }\n }\n\n inline void topk_to_single_result(uint32_t idx) {\n if (force_padding_topk_ && !topk_heap_.full() &&\n topk_heap_.size() < entity_->doc_cnt()) {\n this->fill_random_to_topk_full();\n }\n if (ailego_unlikely(topk_heap_.size() == 0)) {\n return;\n }\n\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n int size = std::min(topk_, static_cast(topk_heap_.size()));\n topk_heap_.sort();\n results_[idx].clear();\n\n for (int i = 0; i < size; ++i) {\n auto score = topk_heap_[i].second;\n if (score.est_dist > this->threshold()) {\n break;\n }\n\n node_id_t id = topk_heap_[i].first;\n if (fetch_vector_) {\n results_[idx].emplace_back(entity_->get_key(id), score.est_dist, id,\n entity_->get_vector(id));\n } else {\n results_[idx].emplace_back(entity_->get_key(id), score.est_dist, id);\n }\n }\n\n return;\n }\n\n //! Construct result from topk heap, result will be normalized\n inline void topk_to_group_result(uint32_t idx) {\n ailego_assert_with(idx < group_results_.size(), \"invalid idx\");\n\n group_results_[idx].clear();\n\n std::vector> group_topk_list;\n std::vector> best_score_in_groups;\n for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end();\n itr++) {\n const std::string &group_id = (*itr).first;\n auto &heap = (*itr).second;\n heap.sort();\n\n if (heap.size() > 0) {\n ResultRecord best_score = heap[0].second;\n best_score_in_groups.push_back(std::make_pair(group_id, best_score));\n }\n }\n\n std::sort(best_score_in_groups.begin(), best_score_in_groups.end(),\n [](const std::pair &a,\n const std::pair &b) -> int {\n return a.second < b.second;\n });\n\n // truncate to group num\n for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size();\n ++i) {\n const std::string &group_id = best_score_in_groups[i].first;\n\n group_topk_list.emplace_back(\n std::make_pair(group_id, group_topk_heaps_[group_id]));\n }\n\n group_results_[idx].resize(group_topk_list.size());\n\n for (uint32_t i = 0; i < group_topk_list.size(); ++i) {\n const std::string &group_id = group_topk_list[i].first;\n group_results_[idx][i].set_group_id(group_id);\n\n uint32_t size = std::min(\n group_topk_, static_cast(group_topk_list[i].second.size()));\n\n for (uint32_t j = 0; j < size; ++j) {\n auto score = group_topk_list[i].second[j].second;\n if (score > this->threshold()) {\n break;\n }\n\n node_id_t id = group_topk_list[i].second[j].first;\n\n if (fetch_vector_) {\n group_results_[idx][i].mutable_docs()->emplace_back(\n entity_->get_key(id), score.est_dist, id,\n entity_->get_vector(id));\n } else {\n group_results_[idx][i].mutable_docs()->emplace_back(\n entity_->get_key(id), score.est_dist, id);\n }\n }\n }\n }\n\n inline void reset_query(const void *query) {\n if (auto query_preprocess_func = index_metric_->get_query_preprocess_func();\n query_preprocess_func != nullptr) {\n size_t dim = dc().dimension();\n preprocess_buffer_.resize(dim);\n memcpy(preprocess_buffer_.data(), query, dim);\n query_preprocess_func(preprocess_buffer_.data(), dim);\n query = preprocess_buffer_.data();\n }\n\n dc().reset_query(query);\n dc().clear_compare_cnt();\n query_ = query;\n }\n\n inline HnswRabitqAddDistCalculator &dist_calculator() {\n return dc();\n }\n\n inline TopkHeap &topk_heap() {\n return topk_heap_;\n }\n\n inline TopkHeap &update_heap() {\n return update_heap_;\n }\n\n inline VisitFilter &visit_filter() {\n return visit_filter_;\n }\n\n inline CandidateHeap &candidates() {\n return candidates_;\n }\n\n inline void set_max_scan_num(uint32_t max_scan_num) {\n max_scan_num_ = max_scan_num;\n }\n\n inline void set_max_scan_limit(uint32_t max_scan_limit) {\n max_scan_limit_ = max_scan_limit;\n }\n\n inline void set_min_scan_limit(uint32_t min_scan_limit) {\n min_scan_limit_ = min_scan_limit;\n }\n\n inline void set_ef(uint32_t v) {\n ef_ = v;\n }\n\n inline void set_filter_mode(uint32_t v) {\n filter_mode_ = v;\n }\n\n inline void set_filter_negative_probability(float v) {\n negative_probability_ = v;\n }\n\n inline void set_max_scan_ratio(float v) {\n max_scan_ratio_ = v;\n }\n\n virtual void set_magic(uint32_t v) {\n magic_ = v;\n }\n\n virtual void set_force_padding_topk(bool v) {\n force_padding_topk_ = v;\n }\n\n void set_bruteforce_threshold(uint32_t v) override {\n bruteforce_threshold_ = v;\n }\n\n inline uint32_t get_bruteforce_threshold() const {\n return bruteforce_threshold_;\n }\n\n void set_fetch_vector(bool v) override {\n fetch_vector_ = v;\n }\n\n bool fetch_vector() const override {\n return fetch_vector_;\n }\n\n //! Reset context\n void reset(void) override {\n set_filter(nullptr);\n reset_threshold();\n set_fetch_vector(false);\n set_group_params(0, 0);\n reset_group_by();\n }\n\n inline std::map &group_topk_heaps() {\n return group_topk_heaps_;\n }\n\n inline TopkHeap &level_topk(int level) {\n if (ailego_unlikely(level_topks_.size() <= static_cast(level))) {\n int cur_level = level_topks_.size();\n level_topks_.resize(level + 1);\n for (; cur_level <= level; ++cur_level) {\n size_t heap_size = std::max(entity_->neighbor_cnt(cur_level),\n entity_->ef_construction());\n level_topks_[cur_level].clear();\n level_topks_[cur_level].limit(heap_size);\n }\n }\n\n return level_topks_[level];\n }\n\n inline void check_need_adjuct_ctx(void) {\n check_need_adjuct_ctx(entity_->doc_cnt());\n }\n\n inline size_t compute_reserve_cnt(uint32_t cur_doc) const {\n if (cur_doc > kMaxReserveDocCnt) {\n return kMaxReserveDocCnt;\n } else if (cur_doc < kMinReserveDocCnt) {\n return kMinReserveDocCnt;\n }\n return cur_doc;\n }\n\n //! candidates heap and visitfilter need to resize as doc cnt growing up\n inline void check_need_adjuct_ctx(uint32_t doc_cnt) {\n if (ailego_unlikely(doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_)) {\n while (doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_) {\n reserve_max_doc_cnt_ =\n reserve_max_doc_cnt_ + compute_reserve_cnt(reserve_max_doc_cnt_);\n }\n uint32_t max_scan_cnt = compute_max_scan_num(reserve_max_doc_cnt_);\n max_scan_num_ = max_scan_cnt;\n visit_filter_.reset(reserve_max_doc_cnt_, max_scan_cnt);\n candidates_.clear();\n candidates_.limit(max_scan_num_);\n }\n }\n\n inline uint32_t compute_max_scan_num(uint32_t max_doc_cnt) const {\n uint32_t max_scan = max_doc_cnt * max_scan_ratio_;\n if (max_scan < min_scan_limit_) {\n max_scan = min_scan_limit_;\n } else if (max_scan > max_scan_limit_) {\n max_scan = max_scan_limit_;\n }\n return max_scan;\n }\n\n inline size_t get_scan_num() const {\n return dc().compare_cnt();\n }\n\n inline uint64_t reach_scan_limit() const {\n return dc().compare_cnt() >= max_scan_num_;\n }\n\n inline bool error() const {\n return dc().error();\n }\n\n inline void clear() {\n add_dc_.clear();\n if (ailego_unlikely(this->debugging())) {\n stats_get_neighbors_cnt_ = 0u;\n stats_get_vector_cnt_ = 0u;\n stats_visit_dup_cnt_ = 0u;\n }\n // do not clear results_ for the next query will need it\n for (auto &it : results_) {\n it.clear();\n }\n }\n\n uint32_t *mutable_stats_get_neighbors() {\n return &stats_get_neighbors_cnt_;\n }\n\n uint32_t *mutable_stats_get_vector() {\n return &stats_get_vector_cnt_;\n }\n\n uint32_t *mutable_stats_visit_dup_cnt() {\n return &stats_visit_dup_cnt_;\n }\n\n inline bool debugging(void) const {\n return debug_mode_;\n }\n\n inline void update_dist_caculator_distance(\n const IndexMetric::MatrixDistance &distance,\n const IndexMetric::MatrixBatchDistance &batch_distance) {\n dc().update_distance(distance, batch_distance);\n }\n\n //! Get topk\n inline uint32_t topk() const override {\n return topk_;\n }\n\n //! Get group topk\n inline uint32_t group_topk() const {\n return group_topk_;\n }\n\n //! Get group num\n inline uint32_t group_num() const {\n return group_num_;\n }\n\n //! Get if group by search\n inline bool group_by_search() {\n return group_num_ > 0;\n }\n\n //! Set group params\n void set_group_params(uint32_t group_num, uint32_t group_topk) override {\n group_num_ = group_num;\n group_topk_ = group_topk;\n\n topk_ = group_topk_ * group_num_;\n\n topk_heap_.limit(std::max(topk_, ef_));\n\n group_topk_heaps_.clear();\n }\n\n void set_provider(IndexProvider::Pointer provider) {\n add_dc_.set_provider(std::move(provider));\n }\n\n const void *query() const {\n return query_;\n }\n\n private:\n inline HnswRabitqAddDistCalculator &dc() {\n return add_dc_;\n }\n\n inline const HnswRabitqAddDistCalculator &dc() const {\n return add_dc_;\n }\n\n private:\n // Filling random nodes if topk not full\n void fill_random_to_topk_full(void);\n\n constexpr static uint32_t kTriggerReserveCnt = 4096UL;\n constexpr static uint32_t kMinReserveDocCnt = 4096UL;\n constexpr static uint32_t kMaxReserveDocCnt = 128 * 1024UL;\n constexpr static uint32_t kInvalidMgic = -1U;\n\n private:\n HnswRabitqEntity::Pointer entity_;\n HnswRabitqAddDistCalculator add_dc_;\n IndexMetric::Pointer metric_;\n\n bool debug_mode_{false};\n bool force_padding_topk_{false};\n uint32_t max_scan_num_{0};\n uint32_t max_scan_limit_{0};\n uint32_t min_scan_limit_{0};\n uint32_t reserve_max_doc_cnt_{kMinReserveDocCnt};\n uint32_t topk_{0};\n uint32_t group_topk_{0};\n uint32_t filter_mode_{VisitFilter::ByteMap};\n float negative_probability_{HnswRabitqEntity::kDefaultBFNegativeProbability};\n uint32_t ef_{HnswRabitqEntity::kDefaultEf};\n float max_scan_ratio_{HnswRabitqEntity::kDefaultScanRatio};\n uint32_t magic_{0U};\n std::vector results_{};\n std::vector group_results_{};\n TopkHeap topk_heap_{};\n TopkHeap update_heap_{};\n std::vector level_topks_{};\n CandidateHeap candidates_{};\n VisitFilter visit_filter_{};\n uint32_t bruteforce_threshold_{};\n bool fetch_vector_{false};\n\n uint32_t group_num_{0};\n std::map group_topk_heaps_{};\n\n uint32_t type_{kUnknownContext};\n //! debug stats info\n uint32_t stats_get_neighbors_cnt_{0u};\n uint32_t stats_get_vector_cnt_{0u};\n uint32_t stats_visit_dup_cnt_{0u};\n std::string preprocess_buffer_;\n const void *query_{nullptr};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.cc", "content": "// Copyright 2025-present the centaurdb project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License\n\n#include \"core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.h\"\n#include \"zvec/core/framework/index_error.h\"\n\nnamespace zvec::core {\n\nint HnswRabitqAddDistCalculator::get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const {\n for (uint32_t i = 0; i < count; ++i) {\n const node_id_t id = ids[i];\n key_t key = entity_->get_key(id);\n if (key == kInvalidKey) {\n return IndexError_NoExist;\n }\n IndexStorage::MemoryBlock block;\n int ret = provider_->get_vector(key, block);\n if (ret != 0) {\n return ret;\n }\n vec_blocks.push_back(std::move(block));\n }\n return 0;\n}\n\n} // namespace zvec::core\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_dist_calculator.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"zvec/core/framework/index_meta.h\"\n#include \"zvec/core/framework/index_metric.h\"\n#include \"zvec/core/framework/index_provider.h\"\n#include \"hnsw_rabitq_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! HnswRabitqAddDistCalculator is only used for index construction\nclass HnswRabitqAddDistCalculator {\n public:\n typedef std::shared_ptr Pointer;\n\n public:\n enum DistType {\n DIST_NONE = 0,\n DIST_DENSE = 1,\n DIST_HYBRID = 2,\n DIST_SPARSE = 3\n };\n\n public:\n //! Constructor\n HnswRabitqAddDistCalculator(const HnswRabitqEntity *entity,\n const IndexMetric::Pointer &metric, uint32_t dim)\n : entity_(entity),\n distance_(metric->distance()),\n batch_distance_(metric->batch_distance()),\n query_(nullptr),\n dim_(dim),\n compare_cnt_(0) {}\n\n //! Constructor\n HnswRabitqAddDistCalculator(const HnswRabitqEntity *entity,\n const IndexMetric::Pointer &metric, uint32_t dim,\n const void *query)\n : entity_(entity),\n distance_(metric->distance()),\n batch_distance_(metric->batch_distance()),\n query_(query),\n dim_(dim),\n compare_cnt_(0) {}\n\n //! Constructor\n HnswRabitqAddDistCalculator(const HnswRabitqEntity *entity,\n const IndexMetric::Pointer &metric)\n : entity_(entity),\n distance_(metric->distance()),\n batch_distance_(metric->batch_distance()),\n query_(nullptr),\n dim_(0),\n compare_cnt_(0) {}\n\n void update(const HnswRabitqEntity *entity,\n const IndexMetric::Pointer &metric) {\n entity_ = entity;\n distance_ = metric->distance();\n batch_distance_ = metric->batch_distance();\n }\n\n void update(const HnswRabitqEntity *entity,\n const IndexMetric::Pointer &metric, uint32_t dim) {\n entity_ = entity;\n distance_ = metric->distance();\n batch_distance_ = metric->batch_distance();\n dim_ = dim;\n }\n\n inline void update_distance(\n const IndexMetric::MatrixDistance &distance,\n const IndexMetric::MatrixBatchDistance &batch_distance) {\n distance_ = distance;\n batch_distance_ = batch_distance;\n }\n\n //! Reset query vector data\n inline void reset_query(const void *query) {\n error_ = false;\n query_ = query;\n }\n\n //! Returns distance\n inline dist_t dist(const void *vec_lhs, const void *vec_rhs) {\n if (ailego_unlikely(vec_lhs == nullptr || vec_rhs == nullptr)) {\n LOG_ERROR(\"Nullptr of dense vector\");\n error_ = true;\n return 0.0f;\n }\n\n float score{0.0f};\n\n distance_(vec_lhs, vec_rhs, dim_, &score);\n\n return score;\n }\n\n //! Returns distance between query and vec.\n inline dist_t dist(const void *vec) {\n compare_cnt_++;\n\n return dist(vec, query_);\n }\n\n //! Return distance between query and node id.\n inline dist_t dist(node_id_t id) {\n compare_cnt_++;\n\n const void *feat = get_vector(id);\n if (ailego_unlikely(feat == nullptr)) {\n LOG_ERROR(\"Get nullptr vector, id=%u\", id);\n error_ = true;\n return 0.0f;\n }\n\n return dist(feat, query_);\n }\n\n //! Return dist node lhs between node rhs\n inline dist_t dist(node_id_t lhs, node_id_t rhs) {\n compare_cnt_++;\n\n const void *feat = get_vector(lhs);\n const void *query = get_vector(rhs);\n if (ailego_unlikely(feat == nullptr || query == nullptr)) {\n LOG_ERROR(\"Get nullptr vector\");\n error_ = true;\n return 0.0f;\n }\n\n return dist(feat, query);\n }\n\n dist_t operator()(const void *vec) {\n return dist(vec);\n }\n\n dist_t operator()(id_t i) {\n return dist(i);\n }\n\n dist_t operator()(id_t lhs, id_t rhs) {\n return dist(lhs, rhs);\n }\n\n void batch_dist(const void **vecs, size_t num, dist_t *distances) {\n compare_cnt_++;\n\n batch_distance_(vecs, query_, num, dim_, distances);\n }\n\n inline dist_t batch_dist(node_id_t id) {\n compare_cnt_++;\n\n const void *feat = get_vector(id);\n if (ailego_unlikely(feat == nullptr)) {\n LOG_ERROR(\"Get nullptr vector, id=%u\", id);\n error_ = true;\n return 0.0f;\n }\n dist_t score = 0;\n batch_distance_(&feat, query_, 1, dim_, &score);\n\n return score;\n }\n\n inline void clear() {\n compare_cnt_ = 0;\n error_ = false;\n }\n\n inline void clear_compare_cnt() {\n compare_cnt_ = 0;\n }\n\n inline bool error() const {\n return error_;\n }\n\n //! Get distances compute times\n inline uint32_t compare_cnt() const {\n return compare_cnt_;\n }\n\n inline uint32_t dimension() const {\n return dim_;\n }\n\n void set_provider(IndexProvider::Pointer provider) {\n provider_ = std::move(provider);\n }\n\n int get_vector(const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const;\n\n const void *get_vector(node_id_t id) const {\n key_t key = entity_->get_key(id);\n if (key == kInvalidKey) {\n return nullptr;\n }\n return provider_->get_vector(key);\n }\n\n private:\n HnswRabitqAddDistCalculator(const HnswRabitqAddDistCalculator &) = delete;\n HnswRabitqAddDistCalculator &operator=(const HnswRabitqAddDistCalculator &) =\n delete;\n\n private:\n const HnswRabitqEntity *entity_;\n IndexMetric::MatrixDistance distance_;\n IndexMetric::MatrixBatchDistance batch_distance_;\n\n const void *query_;\n uint32_t dim_;\n\n uint32_t compare_cnt_; // record distance compute times\n uint32_t compare_cnt_batch_; // record batch distance compute time\n bool error_{false};\n\n // get raw vector\n IndexProvider::Pointer provider_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_entity.h\"\n#include \n#include \"utility/sparse_utility.h\"\n#include \"zvec/core/framework/index_stats.h\"\n\nnamespace zvec {\nnamespace core {\n\nconst std::string HnswRabitqEntity::kGraphHeaderSegmentId = \"graph.header\";\nconst std::string HnswRabitqEntity::kGraphFeaturesSegmentId = \"graph.features\";\nconst std::string HnswRabitqEntity::kGraphKeysSegmentId = \"graph.keys\";\nconst std::string HnswRabitqEntity::kGraphNeighborsSegmentId =\n \"graph.neighbors\";\nconst std::string HnswRabitqEntity::kGraphOffsetsSegmentId = \"graph.offsets\";\nconst std::string HnswRabitqEntity::kGraphMappingSegmentId = \"graph.mapping\";\nconst std::string HnswRabitqEntity::kHnswHeaderSegmentId = \"hnsw.header\";\nconst std::string HnswRabitqEntity::kHnswNeighborsSegmentId = \"hnsw.neighbors\";\nconst std::string HnswRabitqEntity::kHnswOffsetsSegmentId = \"hnsw.offsets\";\n\nvoid HnswRabitqEntity::update_rabitq_params_and_vector_size(\n uint32_t dimension) {\n uint32_t padded_dim = ((dimension + 63) / 64) * 64;\n header_.graph.padded_dim = padded_dim;\n // BinDataMap layout: bin_code (padded_dim/8) + f_add + f_rescale + f_error\n header_.graph.size_bin_data =\n rabitqlib::BinDataMap::data_bytes(padded_dim);\n // ExDataMap layout: ex_code (padded_dim*ex_bits/8) + f_add_ex + f_rescale_ex\n header_.graph.size_ex_data = rabitqlib::ExDataMap::data_bytes(\n padded_dim, header_.graph.ex_bits);\n // quantized vector format: cluster_id + bin_data + ex_data\n header_.graph.vector_size =\n sizeof(uint32_t) + size_bin_data() + size_ex_data();\n}\n\nint HnswRabitqEntity::CalcAndAddPadding(const IndexDumper::Pointer &dumper,\n size_t data_size,\n size_t *padding_size) {\n *padding_size = AlignSize(data_size) - data_size;\n if (*padding_size == 0) {\n return 0;\n }\n\n std::string padding(*padding_size, '\\0');\n if (dumper->write(padding.data(), *padding_size) != *padding_size) {\n LOG_ERROR(\"Append padding failed, size %zu\", *padding_size);\n return IndexError_WriteData;\n }\n return 0;\n}\n\nint64_t HnswRabitqEntity::dump_segment(const IndexDumper::Pointer &dumper,\n const std::string &segment_id,\n const void *data, size_t size) const {\n size_t len = dumper->write(data, size);\n if (len != size) {\n LOG_ERROR(\"Dump segment %s data failed, expect: %zu, actual: %zu\",\n segment_id.c_str(), size, len);\n return IndexError_WriteData;\n }\n\n size_t padding_size = AlignSize(size) - size;\n if (padding_size > 0) {\n std::string padding(padding_size, '\\0');\n if (dumper->write(padding.data(), padding_size) != padding_size) {\n LOG_ERROR(\"Append padding failed, size %zu\", padding_size);\n return IndexError_WriteData;\n }\n }\n\n uint32_t crc = ailego::Crc32c::Hash(data, size);\n int ret = dumper->append(segment_id, size, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s meta failed, ret=%d\", segment_id.c_str(), ret);\n return ret;\n }\n\n return len + padding_size;\n}\n\nint64_t HnswRabitqEntity::dump_header(const IndexDumper::Pointer &dumper,\n const HNSWHeader &hd) const {\n //! dump basic graph header. header is aligned and does not need padding\n int64_t graph_hd_size =\n dump_segment(dumper, kGraphHeaderSegmentId, &hd.graph, hd.graph.size);\n if (graph_hd_size < 0) {\n return graph_hd_size;\n }\n\n //! dump basic graph header. header is aligned and does not need padding\n int64_t hnsw_hd_size =\n dump_segment(dumper, kHnswHeaderSegmentId, &hd.hnsw, hd.hnsw.size);\n if (hnsw_hd_size < 0) {\n return hnsw_hd_size;\n }\n\n return graph_hd_size + hnsw_hd_size;\n}\n\nvoid HnswRabitqEntity::reshuffle_vectors(\n const std::function & /*get_level*/,\n std::vector * /*n2o_mapping*/,\n std::vector * /*o2n_mapping*/, key_t * /*keys*/) const {\n // TODO\n return;\n}\n\nint64_t HnswRabitqEntity::dump_mapping_segment(\n const IndexDumper::Pointer &dumper, const key_t *keys) const {\n std::vector mapping(doc_cnt());\n\n std::iota(mapping.begin(), mapping.end(), 0U);\n std::sort(mapping.begin(), mapping.end(),\n [&](node_id_t i, node_id_t j) { return keys[i] < keys[j]; });\n\n size_t size = mapping.size() * sizeof(node_id_t);\n\n return dump_segment(dumper, kGraphMappingSegmentId, mapping.data(), size);\n}\n\nint64_t HnswRabitqEntity::dump_segments(\n const IndexDumper::Pointer &dumper, key_t *keys,\n const std::function &get_level) const {\n HNSWHeader dump_hd(header());\n\n dump_hd.graph.node_size = AlignSize(vector_size());\n\n std::vector n2o_mapping; // map new id to origin id\n std::vector o2n_mapping; // map origin id to new id\n reshuffle_vectors(get_level, &n2o_mapping, &o2n_mapping, keys);\n if (!o2n_mapping.empty()) {\n dump_hd.hnsw.entry_point = o2n_mapping[entry_point()];\n }\n\n //! Dump header\n int64_t hd_size = dump_header(dumper, dump_hd);\n if (hd_size < 0) {\n return hd_size;\n }\n\n //! Dump vectors\n int64_t vecs_size = dump_vectors(dumper, n2o_mapping);\n if (vecs_size < 0) {\n return vecs_size;\n }\n\n //! Dump neighbors\n auto neighbors_size =\n dump_neighbors(dumper, get_level, n2o_mapping, o2n_mapping);\n if (neighbors_size < 0) {\n return neighbors_size;\n }\n //! free memory\n n2o_mapping = std::vector();\n o2n_mapping = std::vector();\n\n //! Dump keys\n size_t key_segment_size = doc_cnt() * sizeof(key_t);\n int64_t keys_size =\n dump_segment(dumper, kGraphKeysSegmentId, keys, key_segment_size);\n if (keys_size < 0) {\n return keys_size;\n }\n\n //! Dump mapping\n int64_t mapping_size = dump_mapping_segment(dumper, keys);\n if (mapping_size < 0) {\n return mapping_size;\n }\n\n return hd_size + keys_size + vecs_size + neighbors_size + mapping_size;\n}\n\nint64_t HnswRabitqEntity::dump_vectors(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping) const {\n size_t vector_dump_size = vector_size();\n\n size_t padding_size = AlignSize(vector_dump_size) - vector_dump_size;\n\n char padding[padding_size];\n memset(padding, 0, sizeof(padding));\n const void *data = nullptr;\n uint32_t crc = 0U;\n size_t vecs_size = 0UL;\n\n //! dump vectors\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n data = get_vector(reorder_mapping.empty() ? id : reorder_mapping[id]);\n if (ailego_unlikely(!data)) {\n return IndexError_ReadData;\n }\n size_t len = dumper->write(data, vector_size());\n if (len != vector_size()) {\n LOG_ERROR(\"Dump vectors failed, write=%zu expect=%zu\", len,\n vector_size());\n return IndexError_WriteData;\n }\n\n crc = ailego::Crc32c::Hash(data, vector_size(), crc);\n vecs_size += vector_size();\n\n if (padding_size == 0) {\n continue;\n }\n\n len = dumper->write(padding, padding_size);\n if (len != padding_size) {\n LOG_ERROR(\"Dump vectors failed, write=%zu expect=%zu\", len, padding_size);\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(padding, padding_size, crc);\n vecs_size += padding_size;\n }\n\n int ret = dumper->append(kGraphFeaturesSegmentId, vecs_size, 0UL, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump vectors segment meta failed, ret %d\", ret);\n return ret;\n }\n\n return vecs_size;\n}\n\nint64_t HnswRabitqEntity::dump_graph_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const {\n std::vector graph_meta;\n graph_meta.reserve(doc_cnt());\n size_t offset = 0;\n uint32_t crc = 0;\n node_id_t mapping[l0_neighbor_cnt()];\n\n uint32_t min_neighbor_count = 10000;\n uint32_t max_neighbor_count = 0;\n size_t sum_neighbor_count = 0;\n\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n const Neighbors neighbors =\n get_neighbors(0, reorder_mapping.empty() ? id : reorder_mapping[id]);\n ailego_assert_with(!!neighbors.data, \"invalid neighbors\");\n ailego_assert_with(neighbors.size() <= l0_neighbor_cnt(),\n \"invalid neighbors\");\n\n uint32_t neighbor_count = neighbors.size();\n if (neighbor_count < min_neighbor_count) {\n min_neighbor_count = neighbor_count;\n }\n if (neighbor_count > max_neighbor_count) {\n max_neighbor_count = neighbor_count;\n }\n sum_neighbor_count += neighbor_count;\n\n graph_meta.emplace_back(offset, neighbor_count);\n size_t size = neighbors.size() * sizeof(node_id_t);\n const node_id_t *data = &neighbors[0];\n if (!neighbor_mapping.empty()) {\n for (node_id_t i = 0; i < neighbors.size(); ++i) {\n mapping[i] = neighbor_mapping[neighbors[i]];\n }\n data = mapping;\n }\n if (dumper->write(data, size) != size) {\n LOG_ERROR(\"Dump graph neighbor id=%zu failed, size %zu\",\n static_cast(id), size);\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(data, size, crc);\n offset += size;\n }\n\n uint32_t average_neighbor_count = 0;\n if (doc_cnt() > 0) {\n average_neighbor_count = sum_neighbor_count / doc_cnt();\n }\n LOG_INFO(\n \"Dump hnsw graph: min_neighbor_count[%u] max_neighbor_count[%u] \"\n \"average_neighbor_count[%u]\",\n min_neighbor_count, max_neighbor_count, average_neighbor_count);\n\n size_t padding_size = 0;\n int ret = CalcAndAddPadding(dumper, offset, &padding_size);\n if (ret != 0) {\n return ret;\n }\n ret = dumper->append(kGraphNeighborsSegmentId, offset, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s failed, ret %d\",\n kGraphNeighborsSegmentId.c_str(), ret);\n return ret;\n }\n\n //! dump level 0 neighbors meta\n auto len = dump_segment(dumper, kGraphOffsetsSegmentId, graph_meta.data(),\n graph_meta.size() * sizeof(GraphNeighborMeta));\n if (len < 0) {\n return len;\n }\n\n return len + offset + padding_size;\n}\n\nint64_t HnswRabitqEntity::dump_upper_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::function &get_level,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const {\n std::vector hnsw_meta;\n hnsw_meta.reserve(doc_cnt());\n size_t offset = 0;\n uint32_t crc = 0;\n node_id_t buffer[upper_neighbor_cnt() + 1];\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n node_id_t new_id = reorder_mapping.empty() ? id : reorder_mapping[id];\n auto level = get_level(new_id);\n if (level == 0) {\n hnsw_meta.emplace_back(0U, 0U);\n continue;\n }\n hnsw_meta.emplace_back(offset, level);\n ailego_assert_with((size_t)level < kMaxGraphLayers, \"invalid level\");\n for (level_t cur_level = 1; cur_level <= level; ++cur_level) {\n const Neighbors neighbors = get_neighbors(cur_level, new_id);\n ailego_assert_with(!!neighbors.data, \"invalid neighbors\");\n ailego_assert_with(neighbors.size() <= neighbor_cnt(cur_level),\n \"invalid neighbors\");\n memset(buffer, 0, sizeof(buffer));\n buffer[0] = neighbors.size();\n if (neighbor_mapping.empty()) {\n memcpy(&buffer[1], &neighbors[0], neighbors.size() * sizeof(node_id_t));\n } else {\n for (node_id_t i = 0; i < neighbors.size(); ++i) {\n buffer[i + 1] = neighbor_mapping[neighbors[i]];\n }\n }\n if (dumper->write(buffer, sizeof(buffer)) != sizeof(buffer)) {\n LOG_ERROR(\"Dump graph neighbor id=%zu failed, size %zu\",\n static_cast(id), sizeof(buffer));\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(buffer, sizeof(buffer), crc);\n offset += sizeof(buffer);\n }\n }\n size_t padding_size = 0;\n int ret = CalcAndAddPadding(dumper, offset, &padding_size);\n if (ret != 0) {\n return ret;\n }\n\n ret = dumper->append(kHnswNeighborsSegmentId, offset, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s failed, ret %d\", kHnswNeighborsSegmentId.c_str(),\n ret);\n return ret;\n }\n\n //! dump level 0 neighbors meta\n auto len = dump_segment(dumper, kHnswOffsetsSegmentId, hnsw_meta.data(),\n hnsw_meta.size() * sizeof(HnswNeighborMeta));\n if (len < 0) {\n return len;\n }\n\n return len + offset + padding_size;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"zvec/core/framework/index_dumper.h\"\n#include \"zvec/core/framework/index_error.h\"\n#include \"zvec/core/framework/index_storage.h\"\n\nnamespace zvec {\nnamespace core {\n\nusing node_id_t = uint32_t;\nusing key_t = uint64_t;\nusing level_t = int32_t;\nusing dist_t = float;\nstruct EstimateRecord {\n float ip_x0_qr;\n float est_dist;\n float low_dist;\n\n bool operator<(const EstimateRecord &other) const {\n return this->est_dist < other.est_dist;\n }\n};\nstruct ResultRecord {\n float est_dist;\n float low_dist;\n ResultRecord() : est_dist(0.0f), low_dist(0.0f) {}\n ResultRecord(float dist) : est_dist(dist), low_dist(dist) {}\n explicit ResultRecord(const EstimateRecord &other)\n : est_dist(other.est_dist), low_dist(other.low_dist) {}\n ResultRecord(float est_dist, float low_dist)\n : est_dist(est_dist), low_dist(low_dist) {}\n bool operator<(const ResultRecord &other) const {\n return this->est_dist < other.est_dist;\n }\n bool operator<=(const ResultRecord &other) const {\n return this->est_dist <= other.est_dist;\n }\n bool operator>(const ResultRecord &other) const {\n return this->est_dist > other.est_dist;\n }\n};\nusing TopkHeap = ailego::KeyValueHeap;\nusing CandidateHeap =\n ailego::KeyValueHeap>;\nconstexpr node_id_t kInvalidNodeId = static_cast(-1);\nconstexpr key_t kInvalidKey = static_cast(-1);\nclass DistCalculator;\n\nstruct GraphHeader {\n uint32_t size;\n uint32_t version;\n uint32_t graph_type;\n uint32_t doc_count;\n uint32_t vector_size;\n uint32_t node_size;\n uint32_t l0_neighbor_count;\n uint32_t prune_type;\n uint32_t prune_neighbor_count;\n uint32_t ef_construction;\n uint32_t options;\n uint32_t min_neighbor_count;\n uint32_t padded_dim;\n uint32_t size_bin_data;\n uint32_t size_ex_data;\n uint8_t ex_bits;\n uint8_t reserved_[4067];\n};\n\nstatic_assert(sizeof(GraphHeader) % 32 == 0,\n \"GraphHeader must be aligned with 32 bytes\");\n\n//! Hnsw upper neighbor header\nstruct HnswHeader {\n uint32_t size; // header size\n uint32_t revision; // current total docs of the graph\n uint32_t upper_neighbor_count;\n uint32_t ef_construction;\n uint32_t scaling_factor;\n uint32_t max_level;\n uint32_t entry_point;\n uint32_t options;\n uint8_t reserved_[30];\n};\n\nstatic_assert(sizeof(HnswHeader) % 32 == 0,\n \"GraphHeader must be aligned with 32 bytes\");\n\n//! Hnsw common header and upper neighbor header\nstruct HNSWHeader {\n HNSWHeader() {\n clear();\n }\n\n HNSWHeader(const HNSWHeader &header) {\n memcpy(this, &header, sizeof(header));\n }\n\n HNSWHeader &operator=(const HNSWHeader &header) {\n memcpy(this, &header, sizeof(header));\n return *this;\n }\n\n //! Reset state to zero, and the params is untouched\n void inline reset() {\n graph.doc_count = 0U;\n hnsw.entry_point = kInvalidNodeId;\n hnsw.max_level = 0;\n }\n\n //! Clear all fields to init value\n void inline clear() {\n memset(this, 0, sizeof(HNSWHeader));\n hnsw.entry_point = kInvalidNodeId;\n graph.size = sizeof(GraphHeader);\n hnsw.size = sizeof(HnswHeader);\n }\n\n size_t l0_neighbor_cnt() const {\n return graph.l0_neighbor_count;\n }\n\n size_t upper_neighbor_cnt() const {\n return hnsw.upper_neighbor_count;\n }\n\n size_t vector_size() const {\n return graph.vector_size;\n }\n\n uint8_t ex_bits() const {\n return graph.ex_bits;\n }\n\n uint32_t padded_dim() const {\n return graph.padded_dim;\n }\n\n size_t ef_construction() const {\n return graph.ef_construction;\n }\n\n size_t scaling_factor() const {\n return hnsw.scaling_factor;\n }\n\n size_t neighbor_prune_cnt() const {\n return graph.prune_neighbor_count;\n }\n\n node_id_t entry_point() const {\n return hnsw.entry_point;\n }\n\n node_id_t doc_cnt() const {\n return graph.doc_count;\n }\n\n GraphHeader graph;\n HnswHeader hnsw;\n};\n\nstruct NeighborsHeader {\n uint32_t neighbor_cnt;\n node_id_t neighbors[0];\n};\n\nstruct Neighbors {\n Neighbors() : cnt{0}, data{nullptr} {}\n\n Neighbors(uint32_t cnt_in, const node_id_t *data_in)\n : cnt{cnt_in}, data{data_in} {}\n\n Neighbors(IndexStorage::MemoryBlock &&mem_block)\n : neighbor_block{std::move(mem_block)} {\n auto hd = reinterpret_cast(neighbor_block.data());\n cnt = hd->neighbor_cnt;\n data = hd->neighbors;\n }\n\n size_t size(void) const {\n return cnt;\n }\n\n const node_id_t &operator[](size_t idx) const {\n return data[idx];\n }\n\n uint32_t cnt;\n const node_id_t *data;\n IndexStorage::MemoryBlock neighbor_block;\n};\n\n//! level 0 neighbors offset\nstruct GraphNeighborMeta {\n GraphNeighborMeta(size_t o, size_t cnt) : offset(o), neighbor_cnt(cnt) {}\n\n uint64_t offset : 48;\n uint64_t neighbor_cnt : 16;\n};\n\n//! hnsw upper neighbors meta\nstruct HnswNeighborMeta {\n HnswNeighborMeta(size_t o, size_t l) : offset(o), level(l) {}\n\n uint64_t offset : 48; // offset = idx * upper neighors size\n uint64_t level : 16;\n};\n\nclass HnswRabitqEntity {\n public:\n //! Constructor\n HnswRabitqEntity() {}\n\n //! Constructor\n HnswRabitqEntity(const HNSWHeader &hd) {\n header_ = hd;\n }\n\n //! Destructor\n virtual ~HnswRabitqEntity() {}\n\n //! HnswRabitqEntity Pointerd;\n typedef std::shared_ptr Pointer;\n\n //! Get max neighbor size of graph level\n inline size_t neighbor_cnt(level_t level) const {\n return level == 0 ? header_.graph.l0_neighbor_count\n : header_.hnsw.upper_neighbor_count;\n }\n\n //! get max neighbor size of graph level 0\n inline size_t l0_neighbor_cnt() const {\n return header_.graph.l0_neighbor_count;\n }\n\n //! get min neighbor size of graph\n inline size_t min_neighbor_cnt() const {\n return header_.graph.min_neighbor_count;\n }\n\n //! get upper neighbor size of graph level other than 0\n inline size_t upper_neighbor_cnt() const {\n return header_.hnsw.upper_neighbor_count;\n }\n\n //! Get current total doc of the hnsw graph\n inline node_id_t *mutable_doc_cnt() {\n return &header_.graph.doc_count;\n }\n\n inline node_id_t doc_cnt() const {\n return header_.graph.doc_count;\n }\n\n //! Get hnsw graph scaling params\n inline size_t scaling_factor() const {\n return header_.hnsw.scaling_factor;\n }\n\n //! Get prune_size\n inline size_t prune_cnt() const {\n return header_.graph.prune_neighbor_count;\n }\n\n //! Current entity of top level graph\n inline node_id_t entry_point() const {\n return header_.hnsw.entry_point;\n }\n\n //! Current max graph level\n inline level_t cur_max_level() const {\n return header_.hnsw.max_level;\n }\n\n //! Retrieve index vector size\n size_t vector_size() const {\n return header_.graph.vector_size;\n }\n\n //! Retrieve node size\n size_t node_size() const {\n return header_.graph.node_size;\n }\n\n //! Retrieve ef constuction\n size_t ef_construction() const {\n return header_.graph.ef_construction;\n }\n\n uint8_t ex_bits() const {\n return header_.graph.ex_bits;\n }\n\n uint32_t padded_dim() const {\n return header_.graph.padded_dim;\n }\n\n uint32_t size_bin_data() const {\n return header_.graph.size_bin_data;\n }\n\n uint32_t size_ex_data() const {\n return header_.graph.size_ex_data;\n }\n\n void update_rabitq_params_and_vector_size(uint32_t dimension);\n\n void set_ex_bits(uint8_t ex_bits) {\n header_.graph.ex_bits = ex_bits;\n }\n\n void set_prune_cnt(size_t v) {\n header_.graph.prune_neighbor_count = v;\n }\n\n void set_scaling_factor(size_t val) {\n header_.hnsw.scaling_factor = val;\n }\n\n void set_l0_neighbor_cnt(size_t cnt) {\n header_.graph.l0_neighbor_count = cnt;\n }\n\n void set_min_neighbor_cnt(size_t cnt) {\n header_.graph.min_neighbor_count = cnt;\n }\n\n void set_upper_neighbor_cnt(size_t cnt) {\n header_.hnsw.upper_neighbor_count = cnt;\n }\n\n void set_ef_construction(size_t ef) {\n header_.graph.ef_construction = ef;\n }\n\n protected:\n inline const HNSWHeader &header() const {\n return header_;\n }\n\n inline HNSWHeader *mutable_header() {\n return &header_;\n }\n\n inline size_t header_size() const {\n return sizeof(header_);\n }\n\n void set_node_size(size_t size) {\n header_.graph.node_size = size;\n }\n\n //! Dump all segment by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_segments(\n const IndexDumper::Pointer &dumper, key_t *keys,\n const std::function &get_level) const;\n\n private:\n //! dump mapping segment, for get_vector_by_key in provider\n int64_t dump_mapping_segment(const IndexDumper::Pointer &dumper,\n const key_t *keys) const;\n\n //! dump hnsw head by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_header(const IndexDumper::Pointer &dumper,\n const HNSWHeader &hd) const;\n\n //! dump vectors by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_vectors(const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping) const;\n\n //! dump hnsw neighbors by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_neighbors(const IndexDumper::Pointer &dumper,\n const std::function &get_level,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const {\n auto len1 = dump_graph_neighbors(dumper, reorder_mapping, neighbor_mapping);\n if (len1 < 0) {\n return len1;\n }\n auto len2 = dump_upper_neighbors(dumper, get_level, reorder_mapping,\n neighbor_mapping);\n if (len2 < 0) {\n return len2;\n }\n\n return len1 + len2;\n }\n\n //! dump segment by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_segment(const IndexDumper::Pointer &dumper,\n const std::string &segment_id, const void *data,\n size_t size) const;\n\n //! Dump level 0 neighbors\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_graph_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const;\n\n //! Dump upper level neighbors\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_upper_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::function &get_level,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const;\n\n public:\n //! Cleanup the entity\n virtual int cleanup(void) {\n header_.clear();\n return 0;\n }\n\n //! Make a copy of searcher entity, to support thread-safe operation.\n //! The segment in container cannot be read concurrenly\n virtual const HnswRabitqEntity::Pointer clone() const {\n LOG_ERROR(\"Update neighbors not implemented\");\n return HnswRabitqEntity::Pointer();\n }\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const = 0;\n\n //! Get vector feature data by key\n virtual const void *get_vector(node_id_t id) const = 0;\n\n //! Get vectors feature data by keys\n virtual int get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const = 0;\n\n virtual int get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const = 0;\n virtual int get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const = 0;\n\n //! Retrieve a vector using a primary key\n virtual const void *get_vector_by_key(uint64_t /*key*/) const {\n LOG_ERROR(\"get vector not implemented\");\n return nullptr;\n }\n\n virtual int get_vector_by_key(const key_t /*key*/,\n IndexStorage::MemoryBlock & /*block*/) const {\n return IndexError_NotImplemented;\n }\n\n //! Get the node id's neighbors on graph level\n //! Note: the neighbors cannot be modified, using the following\n //! method to get WritableNeighbors if want to\n virtual const Neighbors get_neighbors(level_t level, node_id_t id) const = 0;\n\n //! Add vector and key to hnsw entity, and local id will be saved in id\n virtual int add_vector(level_t /*level*/, key_t /*key*/, const void * /*vec*/,\n node_id_t * /*id*/) {\n return IndexError_NotImplemented;\n }\n\n //! Add vector and id to hnsw entity\n virtual int add_vector_with_id(level_t /*level*/, node_id_t /*id*/,\n const void * /*vec*/) {\n return IndexError_NotImplemented;\n }\n\n virtual int update_neighbors(\n level_t /*level*/, node_id_t /*id*/,\n const std::vector> & /*neighbors*/) {\n LOG_ERROR(\"Update neighbors dense not implemented\");\n\n return 0;\n }\n\n //! Append neighbor_id to node id neighbors on level, size is the current\n //! neighbors size. Notice: the caller must be ensure the neighbors not full\n virtual void add_neighbor(level_t /*level*/, node_id_t /*id*/,\n uint32_t /*size*/, node_id_t /*neighbor_id*/) {\n LOG_ERROR(\"Add neighbor not implemented\");\n }\n\n //! Update entry point and max level\n virtual void update_ep_and_level(node_id_t ep, level_t level) {\n header_.hnsw.entry_point = ep;\n header_.hnsw.max_level = level;\n }\n\n virtual int load(const IndexStorage::Pointer & /*container*/,\n bool /*check_crc*/) {\n LOG_ERROR(\"Load not implemented\");\n return IndexError_NotImplemented;\n }\n\n virtual int dump(const IndexDumper::Pointer & /*dumper*/) {\n LOG_ERROR(\"Dump not implemented\");\n return IndexError_NotImplemented;\n }\n\n static int CalcAndAddPadding(const IndexDumper::Pointer &dumper,\n size_t data_size, size_t *padding_size);\n\n uint32_t get_cluster_id(const void *vec) const {\n return *reinterpret_cast(\n reinterpret_cast(vec) + cluster_id_offset());\n }\n\n const char *get_bin_data(const void *vec) const {\n return reinterpret_cast(vec) + bin_data_offset();\n }\n\n const char *get_ex_data(const void *vec) const {\n return reinterpret_cast(vec) + ex_data_offset();\n }\n\n uint32_t cluster_id_offset() const {\n return 0;\n }\n\n uint32_t bin_data_offset() const {\n return cluster_id_offset() + sizeof(uint32_t);\n }\n\n uint32_t ex_data_offset() const {\n return bin_data_offset() + size_bin_data();\n }\n\n protected:\n static inline size_t AlignSize(size_t size) {\n return (size + 0x1F) & (~0x1F);\n }\n\n static inline size_t AlignPageSize(size_t size) {\n size_t page_mask = ailego::MemoryHelper::PageSize() - 1;\n return (size + page_mask) & (~page_mask);\n }\n\n static inline size_t AlignHugePageSize(size_t size) {\n size_t page_mask = ailego::MemoryHelper::HugePageSize() - 1;\n return (size + page_mask) & (~page_mask);\n }\n\n //! rearrange vectors to improve cache locality\n void reshuffle_vectors(const std::function &get_level,\n std::vector *n2o_mapping,\n std::vector *o2n_mapping,\n key_t *keys) const;\n\n public:\n const static std::string kGraphHeaderSegmentId;\n const static std::string kGraphFeaturesSegmentId;\n const static std::string kGraphKeysSegmentId;\n const static std::string kGraphNeighborsSegmentId;\n const static std::string kGraphOffsetsSegmentId;\n const static std::string kGraphMappingSegmentId;\n const static std::string kHnswHeaderSegmentId;\n const static std::string kHnswNeighborsSegmentId;\n const static std::string kHnswOffsetsSegmentId;\n\n constexpr static uint32_t kRevision = 0U;\n constexpr static size_t kMaxGraphLayers = 15;\n constexpr static uint32_t kDefaultEfConstruction = 500;\n constexpr static uint32_t kDefaultEf = 500;\n constexpr static uint32_t kDefaultUpperMaxNeighborCnt = 50; // M of HNSW\n constexpr static uint32_t kDefaultL0MaxNeighborCnt = 100;\n constexpr static uint32_t kMaxNeighborCnt = 65535;\n constexpr static float kDefaultScanRatio = 0.1f;\n constexpr static uint32_t kDefaultMinScanLimit = 10000;\n constexpr static uint32_t kDefaultMaxScanLimit =\n std::numeric_limits::max();\n constexpr static float kDefaultBFNegativeProbability = 0.001f;\n constexpr static uint32_t kDefaultScalingFactor = 50U;\n constexpr static uint32_t kDefaultBruteForceThreshold = 1000U;\n constexpr static uint32_t kDefaultDocsHardLimit = 1 << 30U; // 1 billion\n constexpr static float kDefaultDocsSoftLimitRatio = 0.9f;\n constexpr static size_t kMaxChunkSize = 0xFFFFFFFF;\n constexpr static size_t kDefaultChunkSize = 2UL * 1024UL * 1024UL;\n constexpr static size_t kDefaultMaxChunkCnt = 50000UL;\n constexpr static float kDefaultNeighborPruneMultiplier =\n 1.0f; // prune_cnt = upper_max_neighbor_cnt * multiplier\n constexpr static float kDefaultL0MaxNeighborCntMultiplier =\n 2.0f; // l0_max_neighbor_cnt = upper_max_neighbor_cnt * multiplier\n\n protected:\n HNSWHeader header_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_hash.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"hnsw_rabitq_chunk.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! Persistent hashmap implement through open addressing algorithm\ntemplate ::value>::type>\nclass HnswIndexHashMap {\n using key_type = Key;\n using val_type = Val;\n\n struct Iterator {\n key_type first;\n val_type second;\n };\n typedef Iterator *iterator;\n typedef Iterator Item;\n typedef const Iterator *const_iterator;\n\n class Slot {\n public:\n Slot(Chunk::Pointer &&chunk, const void *data)\n : chunk_(std::move(chunk)),\n items_(reinterpret_cast(data)) {}\n //! Return a empty loc or the key item loc\n\n Slot(Chunk::Pointer &&chunk, IndexStorage::MemoryBlock &&mem_block)\n : chunk_(std::move(chunk)), items_block_(std::move(mem_block)) {\n items_ = reinterpret_cast(items_block_.data());\n }\n const_iterator find(key_type key, uint32_t max_items, uint32_t mask) const {\n auto it = &items_[key & mask];\n for (auto i = 0U; i < max_items; ++i) {\n if (it->first == key || it->second == EmptyVal) {\n // LOG_DEBUG(\"i=%u\", i);\n return it;\n }\n ++it;\n if (it == &items_[max_items]) {\n it = &items_[0];\n }\n }\n return nullptr;\n }\n\n bool update(const_iterator it) {\n uint32_t offset = reinterpret_cast(it) -\n reinterpret_cast(&items_[0]);\n if (ailego_unlikely(chunk_->write(offset, it, sizeof(Item)) !=\n sizeof(Item))) {\n LOG_ERROR(\"Chunk write failed\");\n return false;\n }\n return true;\n }\n\n private:\n Chunk::Pointer chunk_{};\n const Item *items_{nullptr}; // point to chunk data\n IndexStorage::MemoryBlock items_block_{};\n };\n\n public:\n //! Init the hash\n //! broker the index allocator\n //! chunk_size the size of per chunk allocated, actual size may greater\n //! factor factor = 1/ratio, ratio is the probability of a squence\n //! number inserted to this container\n //! max the max number key can be inserted\n //! expansion_ratio memory expansion ratio\n int init(HnswRabitqChunkBroker::Pointer &broker, uint32_t chunk_size,\n uint32_t factor, size_t max, float expansion_ratio) {\n ailego_assert_with(expansion_ratio > 1.0f, \"ratio must > 1.0f\");\n broker_ = broker;\n\n size_t items = std::ceil(chunk_size * 1.0f / sizeof(Item));\n slot_items_ = 1UL << static_cast((std::ceil(std::log2(items))));\n size_t range = slot_items_ * factor / expansion_ratio;\n mask_bits_ = std::floor(std::log2(range));\n range = 1UL << mask_bits_;\n size_t max_slots = std::ceil(max * 1.0f / range);\n slots_.reserve(max_slots);\n slot_loc_mask_ = slot_items_ - 1U;\n int ret = load();\n if (ret != 0) {\n return ret;\n }\n\n LOG_DEBUG(\n \"HnswRabitqIndexHash init, chunkSize=%u factor=%u max=%zu \"\n \"ratio=%f slotItems=%u maxSlots=%zu maskBits=%u \"\n \"range=%zu\",\n chunk_size, factor, max, expansion_ratio, slot_items_, max_slots,\n mask_bits_, range);\n\n return 0;\n }\n\n int cleanup(void) {\n broker_.reset();\n slots_.clear();\n slots_.shrink_to_fit();\n mask_bits_ = 0U;\n slot_items_ = 0U;\n slot_loc_mask_ = 0U;\n\n return 0;\n }\n\n const_iterator end(void) const {\n return nullptr;\n }\n\n const_iterator find(const key_type key) const {\n auto idx = key >> mask_bits_;\n if (idx >= slots_.size()) {\n return end();\n }\n auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_);\n return it && it->second != EmptyVal ? it : nullptr;\n }\n\n bool insert(key_type key, val_type val) {\n auto idx = key >> mask_bits_;\n if (idx >= slots_.size()) {\n if (ailego_unlikely(idx >= slots_.capacity())) {\n LOG_ERROR(\"no space to insert\");\n return false;\n }\n for (auto i = slots_.size(); i <= idx; ++i) {\n if (ailego_unlikely(!alloc_slot(i))) {\n return false;\n }\n }\n }\n auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_);\n if (ailego_unlikely(it == nullptr)) {\n LOG_ERROR(\"no space to insert\");\n return false;\n }\n\n //! TODO: write memory is ok?\n const_cast(it)->first = key;\n const_cast(it)->second = val;\n\n return slots_[idx].update(it);\n }\n\n private:\n bool alloc_slot(size_t idx) {\n ailego_assert_with(idx == slots_.size(), \"invalid idx\");\n\n size_t size = slot_items_ * sizeof(Item);\n auto p = broker_->alloc_chunk(\n HnswRabitqChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, idx, size);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return false;\n }\n Chunk::Pointer chunk = p.second;\n if (ailego_unlikely(chunk->resize(size) != size)) {\n LOG_ERROR(\"Chunk resize failed, size=%zu\", size);\n return false;\n }\n //! Read the whole data to memory\n IndexStorage::MemoryBlock data_block;\n if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {\n LOG_ERROR(\"Chunk read failed, size=%zu\", size);\n return false;\n }\n\n slots_.emplace_back(std::move(chunk), std::move(data_block));\n return true;\n }\n\n int load(void) {\n size_t slots_cnt = broker_->get_chunk_cnt(\n HnswRabitqChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX);\n for (size_t i = 0UL; i < slots_cnt; ++i) {\n auto chunk = broker_->get_chunk(\n HnswRabitqChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, i);\n if (!chunk) {\n LOG_ERROR(\"Get chunk failed, seq=%zu\", i);\n return IndexError_InvalidFormat;\n }\n size_t size = sizeof(Item) * slot_items_;\n if (chunk->data_size() < size) {\n LOG_ERROR(\n \"Hash params may be mismatch, seq=%zu, data_size=%zu \"\n \"expect=%zu\",\n i, chunk->data_size(), size);\n return IndexError_InvalidFormat;\n }\n //! Read the whole data to memory\n IndexStorage::MemoryBlock data_block;\n if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {\n LOG_ERROR(\"Chunk read failed, size=%zu\", size);\n return false;\n }\n slots_.emplace_back(std::move(chunk), std::move(data_block));\n }\n return 0;\n }\n\n private:\n HnswRabitqChunkBroker::Pointer broker_{}; // chunk broker\n std::vector slots_{};\n uint32_t mask_bits_{0U};\n uint32_t slot_items_{}; // must be a power of 2\n uint32_t slot_loc_mask_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_index_provider.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"zvec/core/framework/index_provider.h\"\n#include \"zvec/core/framework/index_searcher.h\"\n#include \"zvec/core/framework/index_streamer.h\"\n#include \"hnsw_rabitq_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswRabitqIndexProvider : public IndexProvider {\n public:\n HnswRabitqIndexProvider(const IndexMeta &meta,\n const HnswRabitqEntity::Pointer &entity,\n const std::string &owner)\n : meta_(meta), entity_(entity), owner_class_(owner) {}\n\n HnswRabitqIndexProvider(const HnswRabitqIndexProvider &) = delete;\n HnswRabitqIndexProvider &operator=(const HnswRabitqIndexProvider &) = delete;\n\n public: // holder interface\n //! Create a new iterator\n IndexProvider::Iterator::Pointer create_iterator() override {\n return HnswRabitqIndexProvider::Iterator::Pointer(new (std::nothrow)\n Iterator(entity_));\n }\n\n //! Retrieve count of vectors\n size_t count(void) const override {\n return entity_->doc_cnt();\n }\n\n //! Retrieve dimension of vector\n size_t dimension(void) const override {\n return meta_.dimension();\n }\n\n //! Retrieve type of vector\n IndexMeta::DataType data_type(void) const override {\n return meta_.data_type();\n }\n\n //! Retrieve vector size in bytes\n size_t element_size(void) const override {\n return meta_.element_size();\n }\n\n public: // provider's unique interface\n //! Retrieve a vector using a primary key\n const void *get_vector(uint64_t key) const override {\n return entity_->get_vector_by_key(key);\n }\n\n int get_vector(const uint64_t key,\n IndexStorage::MemoryBlock &block) const override {\n return entity_->get_vector_by_key(key, block);\n }\n\n //! Retrieve the owner class\n const std::string &owner_class(void) const override {\n return owner_class_;\n }\n\n private:\n class Iterator : public IndexProvider::Iterator {\n public:\n Iterator(const HnswRabitqEntity::Pointer &entity)\n : entity_(entity), cur_id_(0U) {}\n\n //! Retrieve pointer of data\n //! NOTICE: the vec feature will be changed after iterating to next, so\n //! the caller need to keep a copy of it before iterator to next vector\n virtual const void *data(void) const override {\n return entity_->get_vector(cur_id_);\n }\n\n //! Test if the iterator is valid\n virtual bool is_valid(void) const override {\n return cur_id_ < entity_->doc_cnt();\n }\n\n //! Retrieve primary key\n virtual uint64_t key(void) const override {\n return entity_->get_key(cur_id_);\n }\n\n //! Next iterator\n virtual void next(void) override {\n // cur_id_ += 1;\n cur_id_ = get_next_valid_id(cur_id_ + 1);\n }\n\n //! Reset the iterator\n void reset(void) {\n cur_id_ = get_next_valid_id(0);\n }\n\n private:\n node_id_t get_next_valid_id(node_id_t start_id) {\n for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) {\n if (entity_->get_key(i) != kInvalidNodeId) {\n cur_id_ = i;\n return i;\n }\n }\n return kInvalidNodeId;\n }\n\n private:\n const HnswRabitqEntity::Pointer entity_;\n node_id_t cur_id_;\n };\n\n private:\n const IndexMeta &meta_;\n const HnswRabitqEntity::Pointer entity_;\n const std::string owner_class_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_params.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\ninline const std::string PARAM_HNSW_RABITQ_GENERAL_DIMENSION(\n \"proxima.hnsw_rabitq.general.dimension\");\n\ninline const std::string PARAM_HNSW_RABITQ_BUILDER_THREAD_COUNT(\n \"proxima.hnsw_rabitq.builder.thread_count\");\ninline const std::string PARAM_HNSW_RABITQ_BUILDER_MEMORY_QUOTA(\n \"proxima.hnsw_rabitq.builder.memory_quota\");\ninline const std::string PARAM_HNSW_RABITQ_BUILDER_EFCONSTRUCTION(\n \"proxima.hnsw_rabitq.builder.efconstruction\");\ninline const std::string PARAM_HNSW_RABITQ_BUILDER_SCALING_FACTOR(\n \"proxima.hnsw_rabitq.builder.scaling_factor\");\ninline const std::string PARAM_HNSW_RABITQ_BUILDER_CHECK_INTERVAL_SECS(\n \"proxima.hnsw_rabitq.builder.check_interval_secs\");\ninline const std::string PARAM_HNSW_RABITQ_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER(\n \"proxima.hnsw_rabitq.builder.neighbor_prune_multiplier\");\ninline const std::string PARAM_HNSW_RABITQ_BUILDER_MIN_NEIGHBOR_COUNT(\n \"proxima.hnsw_rabitq.builder.min_neighbor_count\");\ninline const std::string PARAM_HNSW_RABITQ_BUILDER_MAX_NEIGHBOR_COUNT(\n \"proxima.hnsw_rabitq.builder.max_neighbor_count\");\ninline const std::string\n PARAM_HNSW_RABITQ_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER(\n \"proxima.hnsw_rabitq.builder.l0_max_neighbor_count_multiplier\");\n\ninline const std::string PARAM_HNSW_RABITQ_SEARCHER_EF(\n \"proxima.hnsw_rabitq.searcher.ef\");\ninline const std::string PARAM_HNSW_RABITQ_SEARCHER_BRUTE_FORCE_THRESHOLD(\n \"proxima.hnsw_rabitq.searcher.brute_force_threshold\");\ninline const std::string PARAM_HNSW_RABITQ_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE(\n \"proxima.hnsw_rabitq.searcher.neighbors_in_memory_enable\");\ninline const std::string PARAM_HNSW_RABITQ_SEARCHER_MAX_SCAN_RATIO(\n \"proxima.hnsw_rabitq.searcher.max_scan_ratio\");\ninline const std::string PARAM_HNSW_RABITQ_SEARCHER_CHECK_CRC_ENABLE(\n \"proxima.hnsw_rabitq.searcher.check_crc_enable\");\ninline const std::string PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_ENABLE(\n \"proxima.hnsw_rabitq.searcher.visit_bloomfilter_enable\");\ninline const std::string\n PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB(\n \"proxima.hnsw_rabitq.searcher.visit_bloomfilter_negative_prob\");\ninline const std::string PARAM_HNSW_RABITQ_SEARCHER_FORCE_PADDING_RESULT_ENABLE(\n \"proxima.hnsw_rabitq.searcher.force_padding_result_enable\");\n\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO(\n \"proxima.hnsw_rabitq.streamer.max_scan_ratio\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT(\n \"proxima.hnsw_rabitq.streamer.min_scan_limit\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT(\n \"proxima.hnsw_rabitq.streamer.max_scan_limit\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_EF(\n \"proxima.hnsw_rabitq.streamer.ef\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_EFCONSTRUCTION(\n \"proxima.hnsw_rabitq.streamer.efconstruction\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT(\n \"proxima.hnsw_rabitq.streamer.max_neighbor_count\");\ninline const std::string\n PARAM_HNSW_RABITQ_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER(\n \"proxima.hnsw_rabitq.streamer.l0_max_neighbor_count_multiplier\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_SCALING_FACTOR(\n \"proxima.hnsw_rabitq.streamer.scaling_factor\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_BRUTE_FORCE_THRESHOLD(\n \"proxima.hnsw_rabitq.streamer.brute_force_threshold\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT(\n \"proxima.hnsw_rabitq.streamer.docs_hard_limit\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT(\n \"proxima.hnsw_rabitq.streamer.docs_soft_limit\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_MAX_INDEX_SIZE(\n \"proxima.hnsw_rabitq.streamer.max_index_size\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_ENABLE(\n \"proxima.hnsw_rabitq.streamer.visit_bloomfilter_enable\");\ninline const std::string\n PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB(\n \"proxima.hnsw_rabitq.streamer.visit_bloomfilter_negative_prob\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_CHECK_CRC_ENABLE(\n \"proxima.hnsw_rabitq.streamer.check_crc_enable\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER(\n \"proxima.hnsw_rabitq.streamer.neighbor_prune_multiplier\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_CHUNK_SIZE(\n \"proxima.hnsw_rabitq.streamer.chunk_size\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_FILTER_SAME_KEY(\n \"proxima.hnsw_rabitq.streamer.filter_same_key\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_GET_VECTOR_ENABLE(\n \"proxima.hnsw_rabitq.streamer.get_vector_enable\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_MIN_NEIGHBOR_COUNT(\n \"proxima.hnsw_rabitq.streamer.min_neighbor_count\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_FORCE_PADDING_RESULT_ENABLE(\n \"proxima.hnsw_rabitq.streamer.force_padding_result_enable\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_ESTIMATE_DOC_COUNT(\n \"proxima.hnsw_rabitq.streamer.estimate_doc_count\");\ninline const std::string PARAM_HNSW_RABITQ_STREAMER_USE_ID_MAP(\n \"proxima.hnsw_rabitq.streamer.use_id_map\");\n\ninline const std::string PARAM_HNSW_RABITQ_REDUCER_WORKING_PATH(\n \"proxima.hnsw_rabitq.reducer.working_path\");\ninline const std::string PARAM_HNSW_RABITQ_REDUCER_NUM_OF_ADD_THREADS(\n \"proxima.hnsw_rabitq.reducer.num_of_add_threads\");\ninline const std::string PARAM_HNSW_RABITQ_REDUCER_INDEX_NAME(\n \"proxima.hnsw_rabitq.reducer.index_name\");\ninline const std::string PARAM_HNSW_RABITQ_REDUCER_EFCONSTRUCTION(\n \"proxima.hnsw_rabitq.reducer.efconstruction\");\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_query_algorithm.h\"\n#include \n#include \n#include \n#include \n#include \"zvec/ailego/internal/platform.h\"\n#include \"hnsw_rabitq_entity.h\"\n#include \"hnsw_rabitq_query_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswRabitqQueryAlgorithm::HnswRabitqQueryAlgorithm(HnswRabitqEntity &entity,\n size_t num_clusters,\n RabitqMetricType metric_type)\n : entity_(entity),\n mt_(std::chrono::system_clock::now().time_since_epoch().count()),\n lock_pool_(kLockCnt),\n num_clusters_(num_clusters),\n metric_type_(metric_type) {\n ex_bits_ = entity_.ex_bits();\n padded_dim_ = entity_.padded_dim();\n ip_func_ = rabitqlib::select_excode_ipfunc(ex_bits_);\n LOG_INFO(\n \"Create query algorithm. num_clusters=%zu ex_bits=%zu padded_dim=%zu\",\n num_clusters_, ex_bits_, padded_dim_);\n}\n\nint HnswRabitqQueryAlgorithm::cleanup() {\n return 0;\n}\n\nint HnswRabitqQueryAlgorithm::search(HnswRabitqQueryEntity *entity,\n HnswRabitqContext *ctx) const {\n spin_lock_.lock();\n auto maxLevel = entity_.cur_max_level();\n auto entry_point = entity_.entry_point();\n spin_lock_.unlock();\n\n if (ailego_unlikely(entry_point == kInvalidNodeId)) {\n return 0;\n }\n\n EstimateRecord curest;\n get_bin_est(entity_.get_vector(entry_point), curest, *entity);\n\n for (level_t cur_level = maxLevel; cur_level >= 1; --cur_level) {\n select_entry_point(cur_level, &entry_point, &curest, ctx, entity);\n }\n\n auto &topk_heap = ctx->topk_heap();\n topk_heap.clear();\n search_neighbors(0, &entry_point, &curest, topk_heap, ctx, entity);\n\n if (ctx->group_by_search()) {\n expand_neighbors_by_group(topk_heap, ctx, entity);\n }\n\n return 0;\n}\n\n\n//! select_entry_point on hnsw level, ef = 1\nvoid HnswRabitqQueryAlgorithm::select_entry_point(\n level_t level, node_id_t *entry_point, EstimateRecord *curest,\n HnswRabitqContext *ctx, HnswRabitqQueryEntity *query_entity) const {\n auto &entity = ctx->get_entity();\n while (true) {\n const Neighbors neighbors = entity.get_neighbors(level, *entry_point);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n ailego_prefetch(neighbors.data);\n uint32_t size = neighbors.size();\n if (size == 0) {\n break;\n }\n\n bool find_closer = false;\n for (uint32_t i = 0; i < size; ++i) {\n EstimateRecord candest;\n get_bin_est(entity_.get_vector(neighbors[i]), candest, *query_entity);\n\n if (candest.est_dist < curest->est_dist) {\n *curest = candest;\n *entry_point = neighbors[i];\n find_closer = true;\n }\n }\n\n if (!find_closer) {\n break;\n }\n }\n\n return;\n}\n\nvoid HnswRabitqQueryAlgorithm::search_neighbors(\n level_t level, node_id_t *entry_point, EstimateRecord *dist, TopkHeap &topk,\n HnswRabitqContext *ctx, HnswRabitqQueryEntity *query_entity) const {\n const auto &entity = ctx->get_entity();\n VisitFilter &visit = ctx->visit_filter();\n CandidateHeap &candidates = ctx->candidates();\n std::function filter = [](node_id_t) { return false; };\n if (ctx->filter().is_valid()) {\n filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); };\n }\n\n candidates.clear();\n visit.clear();\n visit.set_visited(*entry_point);\n if (!filter(*entry_point)) {\n topk.emplace(*entry_point, ResultRecord(*dist));\n }\n\n candidates.emplace(*entry_point, ResultRecord(*dist));\n while (!candidates.empty() && !ctx->reach_scan_limit()) {\n auto top = candidates.begin();\n node_id_t main_node = top->first;\n auto main_dist = top->second;\n\n if (topk.full() && main_dist.est_dist > topk[0].second.est_dist) {\n break;\n }\n\n candidates.pop();\n const Neighbors neighbors = entity.get_neighbors(level, main_node);\n ailego_prefetch(neighbors.data);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n\n std::vector neighbor_ids(neighbors.size());\n uint32_t size = 0;\n for (uint32_t i = 0; i < neighbors.size(); ++i) {\n node_id_t node = neighbors[i];\n if (visit.visited(node)) {\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_visit_dup_cnt())++;\n }\n continue;\n }\n visit.set_visited(node);\n neighbor_ids[size++] = node;\n }\n if (size == 0) {\n continue;\n }\n\n for (uint32_t i = 0; i < size; ++i) {\n node_id_t node = neighbor_ids[i];\n EstimateRecord candest;\n auto *cand_vector = entity_.get_vector(node);\n ailego_prefetch(cand_vector);\n get_bin_est(cand_vector, candest, *query_entity);\n\n if (ex_bits_ > 0) {\n // Check preliminary score against current worst full estimate.\n bool flag_update_KNNs =\n (!topk.full()) || candest.low_dist < topk[0].second.est_dist;\n\n if (flag_update_KNNs) {\n // Compute the full estimate if promising.\n get_full_est(cand_vector, candest, *query_entity);\n } else {\n continue;\n }\n } else {\n // ex_bits_ == 0: est_dist is already the best estimate\n if (topk.full() && candest.est_dist >= topk[0].second.est_dist) {\n continue;\n }\n }\n candidates.emplace(node, ResultRecord(candest));\n // update entry_point for next level scan\n if (candest < *dist) {\n *entry_point = node;\n *dist = candest;\n }\n if (!filter(node)) {\n topk.emplace(node, ResultRecord(candest));\n }\n } // end for\n } // while\n\n return;\n}\n\nvoid HnswRabitqQueryAlgorithm::expand_neighbors_by_group(\n TopkHeap &topk, HnswRabitqContext *ctx,\n HnswRabitqQueryEntity *query_entity) const {\n if (!ctx->group_by().is_valid()) {\n return;\n }\n\n const auto &entity = ctx->get_entity();\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity.get_key(id));\n };\n\n // devide into groups\n std::map &group_topk_heaps = ctx->group_topk_heaps();\n for (uint32_t i = 0; i < topk.size(); ++i) {\n node_id_t id = topk[i].first;\n auto score = topk[i].second;\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = group_topk_heaps[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, score);\n }\n\n // stage 2, expand to reach group num as possible\n if (group_topk_heaps.size() < ctx->group_num()) {\n VisitFilter &visit = ctx->visit_filter();\n CandidateHeap &candidates = ctx->candidates();\n\n std::function filter = [](node_id_t) { return false; };\n if (ctx->filter().is_valid()) {\n filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); };\n }\n\n // refill to get enough groups\n candidates.clear();\n visit.clear();\n for (uint32_t i = 0; i < topk.size(); ++i) {\n node_id_t id = topk[i].first;\n auto score = topk[i].second;\n\n visit.set_visited(id);\n candidates.emplace_back(id, score);\n }\n\n // do expand\n while (!candidates.empty() && !ctx->reach_scan_limit()) {\n auto top = candidates.begin();\n node_id_t main_node = top->first;\n\n candidates.pop();\n const Neighbors neighbors = entity.get_neighbors(0, main_node);\n ailego_prefetch(neighbors.data);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n\n std::vector neighbor_ids(neighbors.size());\n uint32_t size = 0;\n for (uint32_t i = 0; i < neighbors.size(); ++i) {\n node_id_t node = neighbors[i];\n if (visit.visited(node)) {\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_visit_dup_cnt())++;\n }\n continue;\n }\n visit.set_visited(node);\n neighbor_ids[size++] = node;\n }\n if (size == 0) {\n continue;\n }\n\n for (uint32_t i = 0; i < size; ++i) {\n node_id_t node = neighbor_ids[i];\n EstimateRecord candest;\n auto *cand_vector = entity_.get_vector(node);\n ailego_prefetch(cand_vector);\n get_full_est(cand_vector, candest, *query_entity);\n\n if (!filter(node)) {\n std::string group_id = group_by(node);\n\n auto &topk_heap = group_topk_heaps[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(node, ResultRecord(candest));\n\n if (group_topk_heaps.size() >= ctx->group_num()) {\n break;\n }\n }\n candidates.emplace(node, ResultRecord(candest));\n } // end for\n } // end while\n } // end if\n}\n\nvoid HnswRabitqQueryAlgorithm::get_bin_est(\n const void *vector, EstimateRecord &res,\n HnswRabitqQueryEntity &entity) const {\n const auto &q_to_centroids = entity.q_to_centroids;\n auto &query_wrapper = *entity.query_wrapper;\n uint32_t cluster_id = entity_.get_cluster_id(vector);\n const char *bin_data = entity_.get_bin_data(vector);\n if (metric_type_ == RabitqMetricType::kIP) {\n float norm = q_to_centroids[cluster_id];\n float error = q_to_centroids[cluster_id + num_clusters_];\n rabitqlib::split_single_estdist(bin_data, query_wrapper, padded_dim_,\n res.ip_x0_qr, res.est_dist, res.low_dist,\n -norm, error);\n } else {\n // L2 distance\n float norm = q_to_centroids[cluster_id];\n rabitqlib::split_single_estdist(bin_data, query_wrapper, padded_dim_,\n res.ip_x0_qr, res.est_dist, res.low_dist,\n norm * norm, norm);\n }\n}\n\nvoid HnswRabitqQueryAlgorithm::get_full_est(\n const void *vector, EstimateRecord &res,\n HnswRabitqQueryEntity &entity) const {\n const auto &q_to_centroids = entity.q_to_centroids;\n auto &query_wrapper = *entity.query_wrapper;\n uint32_t cluster_id = entity_.get_cluster_id(vector);\n const char *bin_data = entity_.get_bin_data(vector);\n const char *ex_data = entity_.get_ex_data(vector);\n\n if (metric_type_ == RabitqMetricType::kIP) {\n float norm = q_to_centroids[cluster_id];\n float error = q_to_centroids[cluster_id + num_clusters_];\n rabitqlib::split_single_fulldist(bin_data, ex_data, ip_func_, query_wrapper,\n padded_dim_, ex_bits_, res.est_dist,\n res.low_dist, res.ip_x0_qr, -norm, error);\n } else {\n // L2 distance\n float norm = q_to_centroids[cluster_id];\n rabitqlib::split_single_fulldist(\n bin_data, ex_data, ip_func_, query_wrapper, padded_dim_, ex_bits_,\n res.est_dist, res.low_dist, res.ip_x0_qr, norm * norm, norm);\n }\n}\n\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_algorithm.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"hnsw_rabitq_context.h\"\n#include \"hnsw_rabitq_dist_calculator.h\"\n#include \"hnsw_rabitq_entity.h\"\n#include \"rabitq_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswRabitqQueryEntity;\n\n//! hnsw graph algorithm implement\nclass HnswRabitqQueryAlgorithm {\n public:\n typedef std::unique_ptr UPointer;\n\n public:\n //! Constructor\n explicit HnswRabitqQueryAlgorithm(HnswRabitqEntity &entity,\n size_t num_clusters,\n RabitqMetricType metric_type);\n\n //! Destructor\n ~HnswRabitqQueryAlgorithm() = default;\n\n //! Cleanup HnswRabitqQueryAlgorithm\n int cleanup();\n\n //! do knn search in graph\n //! return 0 on success, or errCode in failure. results saved in ctx\n int search(HnswRabitqQueryEntity *entity, HnswRabitqContext *ctx) const;\n\n //! Initiate HnswRabitqQueryAlgorithm\n int init() {\n level_probas_.clear();\n double level_mult =\n 1 / std::log(static_cast(entity_.scaling_factor()));\n for (int level = 0;; level++) {\n // refers faiss get_random_level alg\n double proba =\n std::exp(-level / level_mult) * (1 - std::exp(-1 / level_mult));\n if (proba < 1e-9) {\n break;\n }\n level_probas_.push_back(proba);\n }\n\n return 0;\n }\n\n //! Generate a random level\n //! return graph level\n uint32_t get_random_level() const {\n // gen rand float (0, 1)\n double f = mt_() / static_cast(mt_.max());\n for (size_t level = 0; level < level_probas_.size(); level++) {\n if (f < level_probas_[level]) {\n return level;\n }\n f -= level_probas_[level];\n }\n return level_probas_.size() - 1;\n }\n void get_full_est(node_id_t id, EstimateRecord &res,\n HnswRabitqQueryEntity &entity) const {\n return get_full_est(entity_.get_vector(id), res, entity);\n }\n\n private:\n //! Select in upper layer to get entry point for next layer search\n void select_entry_point(level_t level, node_id_t *entry_point,\n EstimateRecord *dist, HnswRabitqContext *ctx,\n HnswRabitqQueryEntity *entity) const;\n\n\n //! Given a node id and level, search the nearest neighbors in graph\n //! Note: the nearest neighbors result keeps in topk, and entry_point and\n //! dist will be updated to current level nearest node id and distance\n void search_neighbors(level_t level, node_id_t *entry_point,\n EstimateRecord *dist, TopkHeap &topk,\n HnswRabitqContext *ctx,\n HnswRabitqQueryEntity *entity) const;\n\n\n //! expand neighbors until group nums are reached\n void expand_neighbors_by_group(TopkHeap &topk, HnswRabitqContext *ctx,\n HnswRabitqQueryEntity *query_entity) const;\n\n void get_full_est(const void *vector, EstimateRecord &res,\n HnswRabitqQueryEntity &entity) const;\n void get_bin_est(const void *vector, EstimateRecord &res,\n HnswRabitqQueryEntity &entity) const;\n\n private:\n HnswRabitqQueryAlgorithm(const HnswRabitqQueryAlgorithm &) = delete;\n HnswRabitqQueryAlgorithm &operator=(const HnswRabitqQueryAlgorithm &) =\n delete;\n\n\n private:\n static constexpr uint32_t kLockCnt{1U << 8};\n static constexpr uint32_t kLockMask{kLockCnt - 1U};\n\n HnswRabitqEntity &entity_;\n mutable std::mt19937 mt_{};\n std::vector level_probas_{};\n\n mutable ailego::SpinMutex spin_lock_{}; // global spin lock\n std::mutex mutex_{}; // global mutex\n // TODO: spin lock?\n std::vector lock_pool_{};\n size_t num_clusters_{0};\n RabitqMetricType metric_type_{RabitqMetricType::kL2};\n size_t padded_dim_{0};\n size_t ex_bits_{0};\n float (*ip_func_)(const float *, const uint8_t *, size_t);\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_query_entity.h", "content": "// Copyright 2025-present the centaurdb project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License\n\n#pragma once\n\n#include \n#include \n#include \nnamespace zvec::core {\n\nstruct HnswRabitqQueryEntity {\n std::vector rotated_query;\n std::vector q_to_centroids;\n std::unique_ptr> query_wrapper;\n};\n\n} // namespace zvec::core" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License\n#include \"hnsw_rabitq_builder.h\"\n#include \"hnsw_rabitq_searcher.h\"\n#include \"hnsw_rabitq_streamer.h\"\n#include \"rabitq_converter.h\"\n#include \"rabitq_reformer.h\"\n\nnamespace zvec::core {\n\nINDEX_FACTORY_REGISTER_STREAMER(HnswRabitqStreamer);\nINDEX_FACTORY_REGISTER_REFORMER_ALIAS(RabitqReformer, RabitqReformer);\nINDEX_FACTORY_REGISTER_SEARCHER(HnswRabitqSearcher);\nINDEX_FACTORY_REGISTER_CONVERTER_ALIAS(RabitqConverter, RabitqConverter);\nINDEX_FACTORY_REGISTER_BUILDER(HnswRabitqBuilder);\n\n} // namespace zvec::core" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_searcher.h\"\n#include \n#include \"hnsw_rabitq_algorithm.h\"\n#include \"hnsw_rabitq_entity.h\"\n#include \"hnsw_rabitq_index_provider.h\"\n#include \"hnsw_rabitq_params.h\"\n#include \"hnsw_rabitq_query_entity.h\"\n#include \"hnsw_rabitq_searcher_entity.h\"\n#include \"rabitq_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswRabitqSearcher::HnswRabitqSearcher() {}\n\nHnswRabitqSearcher::~HnswRabitqSearcher() {}\n\nint HnswRabitqSearcher::init(const ailego::Params &search_params) {\n params_ = search_params;\n params_.get(PARAM_HNSW_RABITQ_SEARCHER_EF, &ef_);\n params_.get(PARAM_HNSW_RABITQ_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_);\n params_.get(PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_ENABLE,\n &bf_enabled_);\n params_.get(PARAM_HNSW_RABITQ_SEARCHER_CHECK_CRC_ENABLE, &check_crc_enabled_);\n params_.get(PARAM_HNSW_RABITQ_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE,\n &neighbors_in_memory_enabled_);\n params_.get(PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB,\n &bf_negative_probability_);\n params_.get(PARAM_HNSW_RABITQ_SEARCHER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n params_.get(PARAM_HNSW_RABITQ_SEARCHER_FORCE_PADDING_RESULT_ENABLE,\n &force_padding_topk_enabled_);\n\n if (ef_ == 0) {\n ef_ = HnswRabitqEntity::kDefaultEf;\n }\n if (bf_negative_probability_ <= 0.0f || bf_negative_probability_ >= 1.0f) {\n LOG_ERROR(\n \"[%s] must be in range (0,1)\",\n PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str());\n return IndexError_InvalidArgument;\n }\n\n entity_.set_neighbors_in_memory(neighbors_in_memory_enabled_);\n\n ailego::Params reformer_params;\n reformer_params.set(PARAM_RABITQ_METRIC_NAME, meta_.metric_name());\n int ret = reformer_.init(reformer_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize RabitqReformer: %d\", ret);\n return ret;\n }\n\n state_ = STATE_INITED;\n\n LOG_DEBUG(\n \"Init params: ef=%u maxScanRatio=%f bfEnabled=%u checkCrcEnabled=%u \"\n \"neighborsInMemoryEnabled=%u bfNagtiveProb=%f bruteForceThreshold=%u \"\n \"forcePadding=%u\",\n ef_, max_scan_ratio_, bf_enabled_, check_crc_enabled_,\n neighbors_in_memory_enabled_, bf_negative_probability_,\n bruteforce_threshold_, force_padding_topk_enabled_);\n\n return 0;\n}\n\nvoid HnswRabitqSearcher::print_debug_info() {\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n Neighbors neighbours = entity_.get_neighbors(0, id);\n std::cout << \"node: \" << id << \"; \";\n for (uint32_t i = 0; i < neighbours.size(); ++i) {\n std::cout << neighbours[i];\n\n if (i == neighbours.size() - 1) {\n std::cout << std::endl;\n } else {\n std::cout << \", \";\n }\n }\n }\n}\n\nint HnswRabitqSearcher::cleanup() {\n LOG_INFO(\"Begin HnswRabitqSearcher:cleanup\");\n\n metric_.reset();\n meta_.clear();\n stats_.clear_attributes();\n stats_.set_loaded_count(0UL);\n stats_.set_loaded_costtime(0UL);\n max_scan_ratio_ = HnswRabitqEntity::kDefaultScanRatio;\n max_scan_num_ = 0U;\n ef_ = HnswRabitqEntity::kDefaultEf;\n bf_enabled_ = false;\n bf_negative_probability_ = HnswRabitqEntity::kDefaultBFNegativeProbability;\n bruteforce_threshold_ = HnswRabitqEntity::kDefaultBruteForceThreshold;\n check_crc_enabled_ = false;\n neighbors_in_memory_enabled_ = false;\n entity_.cleanup();\n state_ = STATE_INIT;\n\n LOG_INFO(\"End HnswRabitqSearcher:cleanup\");\n\n return 0;\n}\n\nint HnswRabitqSearcher::load(IndexStorage::Pointer container,\n IndexMetric::Pointer metric) {\n if (state_ != STATE_INITED) {\n LOG_ERROR(\"Init the searcher first before load index\");\n return IndexError_Runtime;\n }\n\n LOG_INFO(\"Begin HnswRabitqSearcher:load\");\n\n auto start_time = ailego::Monotime::MilliSeconds();\n\n int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to deserialize meta from container\");\n return ret;\n }\n\n ret = reformer_.load(container);\n if (ret != 0) {\n LOG_ERROR(\"Failed to load reformer from container: %d\", ret);\n return ret;\n }\n\n ret = entity_.load(container, check_crc_enabled_);\n if (ret != 0) {\n LOG_ERROR(\"HnswRabitqSearcher load index failed\");\n return ret;\n }\n\n alg_ = HnswRabitqQueryAlgorithm::UPointer(new HnswRabitqQueryAlgorithm(\n entity_, reformer_.num_clusters(), reformer_.rabitq_metric_type()));\n\n if (metric) {\n metric_ = metric;\n } else {\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"CreateMetric failed, name: %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"IndexMetric init failed, ret=%d\", ret);\n return ret;\n }\n if (metric_->query_metric()) {\n metric_ = metric_->query_metric();\n }\n }\n\n if (!metric_->is_matched(meta_)) {\n LOG_ERROR(\"IndexMetric not match index meta\");\n return IndexError_Mismatch;\n }\n\n max_scan_num_ = static_cast(max_scan_ratio_ * entity_.doc_cnt());\n max_scan_num_ = std::max(4096U, max_scan_num_);\n\n stats_.set_loaded_count(entity_.doc_cnt());\n stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time);\n state_ = STATE_LOADED;\n magic_ = IndexContext::GenerateMagic();\n\n LOG_INFO(\"End HnswRabitqSearcher::load\");\n\n return 0;\n}\n\nint HnswRabitqSearcher::unload() {\n LOG_INFO(\"HnswRabitqSearcher unload index\");\n\n meta_.clear();\n entity_.cleanup();\n metric_.reset();\n max_scan_num_ = 0;\n stats_.set_loaded_count(0UL);\n stats_.set_loaded_costtime(0UL);\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint HnswRabitqSearcher::update_context(HnswRabitqContext *ctx) const {\n const HnswRabitqEntity::Pointer entity = entity_.clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone search context entity\");\n return IndexError_Runtime;\n }\n ctx->set_max_scan_num(max_scan_num_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n\n return ctx->update_context(HnswRabitqContext::kSearcherContext, meta_,\n metric_, entity, magic_);\n}\n\nint HnswRabitqSearcher::search_impl(const void *query,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n if (ailego_unlikely(!query || !context)) {\n LOG_ERROR(\"The context is not created by this searcher\");\n return IndexError_Mismatch;\n }\n HnswRabitqContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswRabitqContext failed\");\n return IndexError_Cast;\n }\n\n if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) {\n return search_bf_impl(query, qmeta, count, context);\n }\n\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n int ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->resize_results(count);\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n int ret = reformer_.transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ret = alg_->search(&entity, ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher fast search failed\");\n return ret;\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nint HnswRabitqSearcher::search_bf_impl(const void *query,\n const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const {\n if (ailego_unlikely(!query || !context)) {\n LOG_ERROR(\"The context is not created by this searcher\");\n return IndexError_Mismatch;\n }\n HnswRabitqContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswRabitqContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n int ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->resize_results(count);\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n int ret = reformer_.transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ctx->group_topk_heaps().clear();\n\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n EstimateRecord dist;\n alg_->get_full_est(id, dist, entity);\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n } else {\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n int ret = reformer_.transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ctx->topk_heap().clear();\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n EstimateRecord dist;\n alg_->get_full_est(id, dist, entity);\n ctx->topk_heap().emplace(id, dist);\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nint HnswRabitqSearcher::search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n if (ailego_unlikely(!query || !context)) {\n LOG_ERROR(\"The context is not created by this searcher\");\n return IndexError_Mismatch;\n }\n\n if (ailego_unlikely(p_keys.size() != count)) {\n LOG_ERROR(\"The size of p_keys is not equal to count\");\n return IndexError_InvalidArgument;\n }\n\n HnswRabitqContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswRabitqContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n int ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->resize_results(count);\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n int ret = reformer_.transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ctx->group_topk_heaps().clear();\n\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n EstimateRecord dist;\n alg_->get_full_est(id, dist, entity);\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n } else {\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n int ret = reformer_.transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ctx->topk_heap().clear();\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n EstimateRecord dist;\n alg_->get_full_est(id, dist, entity);\n ctx->topk_heap().emplace(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nIndexSearcher::Context::Pointer HnswRabitqSearcher::create_context() const {\n if (ailego_unlikely(state_ != STATE_LOADED)) {\n LOG_ERROR(\"Load the index first before create context\");\n return Context::Pointer();\n }\n const HnswRabitqEntity::Pointer search_ctx_entity = entity_.clone();\n if (!search_ctx_entity) {\n LOG_ERROR(\"Failed to create search context entity\");\n return Context::Pointer();\n }\n HnswRabitqContext *ctx = new (std::nothrow)\n HnswRabitqContext(meta_.dimension(), metric_, search_ctx_entity);\n if (ailego_unlikely(ctx == nullptr)) {\n LOG_ERROR(\"Failed to new HnswRabitqContext\");\n return Context::Pointer();\n }\n ctx->set_ef(ef_);\n ctx->set_max_scan_num(max_scan_num_);\n uint32_t filter_mode =\n bf_enabled_ ? VisitFilter::BloomFilter : VisitFilter::ByteMap;\n ctx->set_filter_mode(filter_mode);\n ctx->set_filter_negative_probability(bf_negative_probability_);\n ctx->set_magic(magic_);\n ctx->set_force_padding_topk(force_padding_topk_enabled_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n if (ailego_unlikely(ctx->init(HnswRabitqContext::kSearcherContext)) != 0) {\n LOG_ERROR(\"Init HnswRabitqContext failed\");\n delete ctx;\n return Context::Pointer();\n }\n\n return Context::Pointer(ctx);\n}\n\nIndexProvider::Pointer HnswRabitqSearcher::create_provider(void) const {\n LOG_DEBUG(\"HnswRabitqSearcher create provider\");\n\n auto entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"Clone HnswRabitqEntity failed\");\n return Provider::Pointer();\n }\n return Provider::Pointer(new (std::nothrow) HnswRabitqIndexProvider(\n meta_, entity, \"HnswRabitqSearcher\"));\n}\n\nconst void *HnswRabitqSearcher::get_vector(uint64_t key) const {\n return entity_.get_vector_by_key(key);\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"zvec/core/framework/index_framework.h\"\n#include \"hnsw_rabitq_query_algorithm.h\"\n#include \"hnsw_rabitq_searcher_entity.h\"\n#include \"rabitq_reformer.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswRabitqSearcher : public IndexSearcher {\n public:\n using ContextPointer = IndexSearcher::Context::Pointer;\n\n public:\n HnswRabitqSearcher(void);\n ~HnswRabitqSearcher(void);\n\n HnswRabitqSearcher(const HnswRabitqSearcher &) = delete;\n HnswRabitqSearcher &operator=(const HnswRabitqSearcher &) = delete;\n\n protected:\n //! Initialize Searcher\n virtual int init(const ailego::Params ¶ms) override;\n\n //! Cleanup Searcher\n virtual int cleanup(void) override;\n\n //! Load Index from storage\n virtual int load(IndexStorage::Pointer container,\n IndexMetric::Pointer metric) override;\n\n //! Unload index from storage\n virtual int unload(void) override;\n\n //! KNN Search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n ContextPointer &context) const override {\n return search_impl(query, qmeta, 1, context);\n }\n\n //! KNN Search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n ContextPointer &context) const override;\n\n //! Linear Search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n ContextPointer &context) const override {\n return search_bf_impl(query, qmeta, 1, context);\n }\n\n //! Linear Search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n ContextPointer &context) const override;\n\n //! Linear search by primary keys\n virtual int search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, ContextPointer &context) const override {\n return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context);\n }\n\n //! Linear search by primary keys\n virtual int search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const override;\n\n //! Fetch vector by key\n virtual const void *get_vector(uint64_t key) const override;\n\n //! Create a searcher context\n virtual ContextPointer create_context() const override;\n\n //! Create a new iterator\n virtual IndexProvider::Pointer create_provider(void) const override;\n\n //! Retrieve statistics\n virtual const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve meta of index\n virtual const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n //! Retrieve params of index\n virtual const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n virtual void print_debug_info() override;\n\n private:\n //! To share ctx across streamer/searcher, we need to update the context for\n //! current streamer/searcher\n int update_context(HnswRabitqContext *ctx) const;\n\n private:\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 };\n\n HnswRabitqSearcherEntity entity_{};\n HnswRabitqQueryAlgorithm::UPointer alg_; // impl graph algorithm\n\n IndexMetric::Pointer metric_{};\n IndexMeta meta_{};\n ailego::Params params_{};\n Stats stats_;\n uint32_t ef_{HnswRabitqEntity::kDefaultEf};\n uint32_t max_scan_num_{0U};\n uint32_t bruteforce_threshold_{HnswRabitqEntity::kDefaultBruteForceThreshold};\n float max_scan_ratio_{HnswRabitqEntity::kDefaultScanRatio};\n bool bf_enabled_{false};\n bool check_crc_enabled_{false};\n bool neighbors_in_memory_enabled_{false};\n bool force_padding_topk_enabled_{false};\n float bf_negative_probability_{\n HnswRabitqEntity::kDefaultBFNegativeProbability};\n uint32_t magic_{0U};\n RabitqReformer reformer_;\n\n State state_{STATE_INIT};\n};\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_searcher_entity.h\"\n#include \n#include \"utility/sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswRabitqSearcherEntity::HnswRabitqSearcherEntity() {}\n\nint HnswRabitqSearcherEntity::cleanup(void) {\n storage_.reset();\n vectors_.reset();\n keys_.reset();\n neighbors_.reset();\n neighbors_meta_.reset();\n neighbors_in_memory_enabled_ = false;\n loaded_ = false;\n\n this->HnswRabitqEntity::cleanup();\n\n return 0;\n}\n\nkey_t HnswRabitqSearcherEntity::get_key(node_id_t id) const {\n const void *key;\n if (ailego_unlikely(keys_->read(id * sizeof(key_t), &key, sizeof(key_t)) !=\n sizeof(key_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidKey;\n }\n return *(reinterpret_cast(key));\n}\n\n//! Get vector local id by key\nnode_id_t HnswRabitqSearcherEntity::get_id(key_t key) const {\n if (ailego_unlikely(!mapping_)) {\n LOG_ERROR(\"Index missing mapping segment\");\n return kInvalidNodeId;\n }\n\n //! Do binary search\n node_id_t start = 0UL;\n node_id_t end = doc_cnt();\n const void *data;\n node_id_t idx = 0u;\n while (start < end) {\n idx = start + (end - start) / 2;\n if (ailego_unlikely(\n mapping_->read(idx * sizeof(node_id_t), &data, sizeof(node_id_t)) !=\n sizeof(node_id_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidNodeId;\n }\n const key_t *mkey;\n node_id_t local_id = *reinterpret_cast(data);\n if (ailego_unlikely(keys_->read(local_id * sizeof(key_t),\n (const void **)(&mkey),\n sizeof(key_t)) != sizeof(key_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidNodeId;\n }\n if (*mkey < key) {\n start = idx + 1;\n } else if (*mkey > key) {\n end = idx;\n } else {\n return local_id;\n }\n }\n return kInvalidNodeId;\n}\n\nconst void *HnswRabitqSearcherEntity::get_vector_by_key(key_t key) const {\n node_id_t local_id = get_id(key);\n if (ailego_unlikely(local_id == kInvalidNodeId)) {\n return nullptr;\n }\n\n return get_vector(local_id);\n}\n\nconst void *HnswRabitqSearcherEntity::get_vector(node_id_t id) const {\n size_t read_size = vector_size();\n size_t offset = node_size() * id;\n\n const void *vec;\n if (ailego_unlikely(vectors_->read(offset, &vec, read_size) != read_size)) {\n LOG_ERROR(\"Read vector from segment failed\");\n return nullptr;\n }\n return vec;\n}\n\nint HnswRabitqSearcherEntity::get_vector(\n const node_id_t id, IndexStorage::MemoryBlock &block) const {\n const void *vec = get_vector(id);\n block.reset((void *)vec);\n return 0;\n}\n\nconst void *HnswRabitqSearcherEntity::get_vectors() const {\n const void *vec;\n size_t len = node_size() * doc_cnt();\n if (vectors_->read(0, &vec, len) != len) {\n LOG_ERROR(\"Read vectors from segment failed\");\n return nullptr;\n }\n return vec;\n}\n\nint HnswRabitqSearcherEntity::get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const {\n ailego_assert_with(count <= segment_datas_.size(), \"invalid count\");\n\n size_t read_size = vector_size();\n\n for (uint32_t i = 0; i < count; ++i) {\n segment_datas_[i].offset = node_size() * ids[i];\n segment_datas_[i].length = read_size;\n\n ailego_assert_with(segment_datas_[i].offset < vectors_->data_size(),\n \"invalid offset\");\n }\n if (ailego_unlikely(!vectors_->read(&segment_datas_[0], count))) {\n LOG_ERROR(\"Read vectors from segment failed\");\n return IndexError_ReadData;\n }\n for (uint32_t i = 0; i < count; ++i) {\n vecs[i] = segment_datas_[i].data;\n }\n\n return 0;\n}\n\nint HnswRabitqSearcherEntity::get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const {\n const void *vecs[count];\n get_vector(ids, count, vecs);\n for (uint32_t i = 0; i < count; ++i) {\n vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i]));\n }\n return 0;\n}\n\nconst Neighbors HnswRabitqSearcherEntity::get_neighbors(level_t level,\n node_id_t id) const {\n if (level == 0) {\n if (neighbors_in_memory_enabled_) {\n auto hd = reinterpret_cast(\n fixed_neighbors_.get() + neighbors_size() * id);\n return {hd->neighbor_cnt, hd->neighbors};\n }\n\n const GraphNeighborMeta *m;\n if (ailego_unlikely(neighbors_meta_->read(id * sizeof(GraphNeighborMeta),\n (const void **)(&m),\n sizeof(GraphNeighborMeta)) !=\n sizeof(GraphNeighborMeta))) {\n LOG_ERROR(\"Read neighbors meta from segment failed\");\n return {0, nullptr};\n }\n\n const void *data;\n if (ailego_unlikely(neighbors_->read(m->offset, &data,\n m->neighbor_cnt * sizeof(node_id_t)) !=\n m->neighbor_cnt * sizeof(node_id_t))) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return {0, nullptr};\n }\n return {static_cast(m->neighbor_cnt),\n reinterpret_cast(data)};\n }\n\n //! Read level > 0 neighbors\n const HnswNeighborMeta *m;\n if (ailego_unlikely(upper_neighbors_meta_->read(id * sizeof(HnswNeighborMeta),\n (const void **)(&m),\n sizeof(HnswNeighborMeta)) !=\n sizeof(HnswNeighborMeta))) {\n LOG_ERROR(\"Read neighbors meta from segment failed\");\n return {0, nullptr};\n }\n\n ailego_assert_with(level <= m->level, \"invalid level\");\n size_t offset = m->offset + (level - 1) * upper_neighbors_size();\n ailego_assert_with(offset <= upper_neighbors_->data_size(), \"invalid offset\");\n const void *data;\n if (ailego_unlikely(\n upper_neighbors_->read(offset, &data, upper_neighbors_size()) !=\n upper_neighbors_size())) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return {0, nullptr};\n }\n\n auto hd = reinterpret_cast(data);\n return {hd->neighbor_cnt, hd->neighbors};\n}\n\nint HnswRabitqSearcherEntity::load(const IndexStorage::Pointer &container,\n bool check_crc) {\n storage_ = container;\n\n int ret = load_segments(check_crc);\n if (ret != 0) {\n return ret;\n }\n\n loaded_ = true;\n\n LOG_INFO(\n \"Index info: docCnt=%u entryPoint=%u maxLevel=%d efConstruct=%zu \"\n \"l0NeighborCnt=%zu upperNeighborCnt=%zu scalingFactor=%zu \"\n \"vectorSize=%zu nodeSize=%zu vectorSegmentSize=%zu keySegmentSize=%zu \"\n \"neighborsSegmentSize=%zu neighborsMetaSegmentSize=%zu \",\n doc_cnt(), entry_point(), cur_max_level(), ef_construction(),\n l0_neighbor_cnt(), upper_neighbor_cnt(), scaling_factor(), vector_size(),\n node_size(), vectors_->data_size(), keys_->data_size(),\n neighbors_ == nullptr ? 0 : neighbors_->data_size(),\n neighbors_meta_ == nullptr ? 0 : neighbors_meta_->data_size());\n\n return 0;\n}\n\nint HnswRabitqSearcherEntity::load_segments(bool check_crc) {\n //! load header\n const void *data = nullptr;\n HNSWHeader hd;\n auto graph_hd_segment = storage_->get(kGraphHeaderSegmentId);\n if (!graph_hd_segment || graph_hd_segment->data_size() < sizeof(hd.graph)) {\n LOG_ERROR(\"Miss or invalid segment %s\", kGraphHeaderSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n if (graph_hd_segment->read(0, reinterpret_cast(&data),\n sizeof(hd.graph)) != sizeof(hd.graph)) {\n LOG_ERROR(\"Read segment %s failed\", kGraphHeaderSegmentId.c_str());\n return IndexError_ReadData;\n }\n memcpy(&hd.graph, data, sizeof(hd.graph));\n\n auto hnsw_hd_segment = storage_->get(kHnswHeaderSegmentId);\n if (!hnsw_hd_segment || hnsw_hd_segment->data_size() < sizeof(hd.hnsw)) {\n LOG_ERROR(\"Miss or invalid segment %s\", kHnswHeaderSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n if (hnsw_hd_segment->read(0, reinterpret_cast(&data),\n sizeof(hd.hnsw)) != sizeof(hd.hnsw)) {\n LOG_ERROR(\"Read segment %s failed\", kHnswHeaderSegmentId.c_str());\n return IndexError_ReadData;\n }\n memcpy(&hd.hnsw, data, sizeof(hd.hnsw));\n *mutable_header() = hd;\n segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt()));\n\n vectors_ = storage_->get(kGraphFeaturesSegmentId);\n if (!vectors_) {\n LOG_ERROR(\"IndexStorage get segment %s failed\",\n kGraphFeaturesSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n keys_ = storage_->get(kGraphKeysSegmentId);\n if (!keys_) {\n LOG_ERROR(\"IndexStorage get segment %s failed\",\n kGraphKeysSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n\n neighbors_ = storage_->get(kGraphNeighborsSegmentId);\n if (!neighbors_ || (neighbors_->data_size() == 0 && doc_cnt() > 1)) {\n LOG_ERROR(\"IndexStorage get segment %s failed or empty\",\n kGraphNeighborsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n neighbors_meta_ = storage_->get(kGraphOffsetsSegmentId);\n if (!neighbors_meta_ ||\n neighbors_meta_->data_size() < sizeof(GraphNeighborMeta) * doc_cnt()) {\n LOG_ERROR(\"IndexStorage get segment %s failed or invalid size\",\n kGraphOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n upper_neighbors_ = storage_->get(kHnswNeighborsSegmentId);\n if (!upper_neighbors_ ||\n (upper_neighbors_->data_size() == 0 && cur_max_level() > 0)) {\n LOG_ERROR(\"IndexStorage get segment %s failed or empty\",\n kHnswNeighborsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n upper_neighbors_meta_ = storage_->get(kHnswOffsetsSegmentId);\n if (!upper_neighbors_meta_ || upper_neighbors_meta_->data_size() <\n sizeof(HnswNeighborMeta) * doc_cnt()) {\n LOG_ERROR(\"IndexStorage get segment %s failed or invalid size\",\n kHnswOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n mapping_ = storage_->get(kGraphMappingSegmentId);\n if (!mapping_ || mapping_->data_size() < sizeof(node_id_t) * doc_cnt()) {\n LOG_ERROR(\"IndexStorage get segment %s failed or invalid size\",\n kGraphMappingSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (check_crc) {\n std::vector segments;\n segments.emplace_back(graph_hd_segment);\n segments.emplace_back(hnsw_hd_segment);\n segments.emplace_back(vectors_);\n segments.emplace_back(keys_);\n\n segments.emplace_back(neighbors_);\n segments.emplace_back(neighbors_meta_);\n segments.emplace_back(upper_neighbors_);\n segments.emplace_back(upper_neighbors_meta_);\n\n if (!do_crc_check(segments)) {\n LOG_ERROR(\"Check index crc failed, the index may broken\");\n return IndexError_Runtime;\n }\n }\n\n if (neighbors_in_memory_enabled_) {\n int ret = load_and_flat_neighbors();\n if (ret != 0) {\n return ret;\n }\n }\n\n return 0;\n}\n\nint HnswRabitqSearcherEntity::load_and_flat_neighbors() {\n fixed_neighbors_.reset(\n new (std::nothrow) char[neighbors_size() * doc_cnt()]{},\n std::default_delete());\n if (!fixed_neighbors_) {\n LOG_ERROR(\"Malloc memory failed\");\n return IndexError_NoMemory;\n }\n\n //! Get a new segemnt to release the buffer after loading neighbors\n auto neighbors_meta = storage_->get(kGraphOffsetsSegmentId);\n if (!neighbors_meta) {\n LOG_ERROR(\"IndexStorage get segment graph.offsets failed\");\n return IndexError_InvalidArgument;\n }\n\n const GraphNeighborMeta *neighbors_index = nullptr;\n if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index),\n neighbors_meta->data_size()) !=\n neighbors_meta->data_size()) {\n LOG_ERROR(\"Read segment %s data failed\", kGraphOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n const char *neighbor_data;\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n size_t rd_size = neighbors_index[id].neighbor_cnt * sizeof(node_id_t);\n if (ailego_unlikely(\n neighbors_->read(neighbors_index[id].offset,\n reinterpret_cast(&neighbor_data),\n rd_size) != rd_size)) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return IndexError_ReadData;\n }\n // copy level 0 neighbors to fixed size neighbors memory\n char *dst = fixed_neighbors_.get() + neighbors_size() * id;\n *reinterpret_cast(dst) = neighbors_index[id].neighbor_cnt;\n memcpy(dst + sizeof(uint32_t), neighbor_data, rd_size);\n }\n\n return 0;\n}\n\nint HnswRabitqSearcherEntity::get_fixed_neighbors(\n std::vector *fixed_neighbors) const {\n //! Get a new segemnt to release the buffer after loading neighbors\n auto neighbors_meta = storage_->get(kGraphOffsetsSegmentId);\n if (!neighbors_meta) {\n LOG_ERROR(\"IndexStorage get segment graph.offsets failed\");\n return IndexError_InvalidArgument;\n }\n\n const GraphNeighborMeta *neighbors_index = nullptr;\n size_t meta_size = neighbors_meta->data_size();\n if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index),\n meta_size) != meta_size) {\n LOG_ERROR(\"Read segment %s data failed\", kGraphOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n size_t fixed_neighbor_cnt = l0_neighbor_cnt();\n fixed_neighbors->resize((fixed_neighbor_cnt + 1) * doc_cnt(), kInvalidNodeId);\n\n size_t neighbors_cnt_offset = fixed_neighbor_cnt * doc_cnt();\n size_t total_neighbor_cnt = 0;\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n size_t cur_neighbor_cnt = neighbors_index[id].neighbor_cnt;\n if (cur_neighbor_cnt == 0) {\n (*fixed_neighbors)[neighbors_cnt_offset + id] = 0;\n continue;\n }\n size_t rd_size = cur_neighbor_cnt * sizeof(node_id_t);\n const uint32_t *neighbors;\n if (neighbors_->read(neighbors_index[id].offset,\n reinterpret_cast(&neighbors),\n rd_size) != rd_size) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return IndexError_ReadData;\n }\n\n // copy level 0 neighbors to fixed size neighbors memory\n auto it = fixed_neighbors->begin() + id * fixed_neighbor_cnt;\n std::copy(neighbors, neighbors + cur_neighbor_cnt, it);\n\n (*fixed_neighbors)[neighbors_cnt_offset + id] = cur_neighbor_cnt;\n total_neighbor_cnt += cur_neighbor_cnt;\n }\n LOG_INFO(\"total neighbor cnt: %zu, average neighbor cnt: %zu\",\n total_neighbor_cnt, total_neighbor_cnt / doc_cnt());\n\n return 0;\n}\n\nbool HnswRabitqSearcherEntity::do_crc_check(\n std::vector &segments) const {\n constexpr size_t blk_size = 4096;\n const void *data;\n for (auto &segment : segments) {\n size_t offset = 0;\n size_t rd_size;\n uint32_t crc = 0;\n while (offset < segment->data_size()) {\n size_t size = std::min(blk_size, segment->data_size() - offset);\n if ((rd_size = segment->read(offset, &data, size)) <= 0) {\n break;\n }\n offset += rd_size;\n crc = ailego::Crc32c::Hash(data, rd_size, crc);\n }\n if (crc != segment->data_crc()) {\n return false;\n }\n }\n return true;\n}\n\nconst HnswRabitqEntity::Pointer HnswRabitqSearcherEntity::clone() const {\n auto vectors = vectors_->clone();\n if (ailego_unlikely(!vectors)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphFeaturesSegmentId.c_str());\n return HnswRabitqEntity::Pointer();\n }\n auto keys = keys_->clone();\n if (ailego_unlikely(!keys)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphKeysSegmentId.c_str());\n return HnswRabitqEntity::Pointer();\n }\n\n auto mapping = mapping_->clone();\n if (ailego_unlikely(!mapping)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphMappingSegmentId.c_str());\n return HnswRabitqEntity::Pointer();\n }\n\n auto neighbors = neighbors_->clone();\n if (ailego_unlikely(!neighbors)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphNeighborsSegmentId.c_str());\n return HnswRabitqEntity::Pointer();\n }\n auto upper_neighbors = upper_neighbors_->clone();\n if (ailego_unlikely(!neighbors)) {\n LOG_ERROR(\"clone segment %s failed\", kHnswNeighborsSegmentId.c_str());\n return HnswRabitqEntity::Pointer();\n }\n auto neighbors_meta = neighbors_meta_->clone();\n if (ailego_unlikely(!neighbors_meta)) {\n LOG_ERROR(\"clone segment %s failed\", kGraphOffsetsSegmentId.c_str());\n return HnswRabitqEntity::Pointer();\n }\n auto upper_neighbors_meta = upper_neighbors_meta_->clone();\n if (ailego_unlikely(!upper_neighbors_meta)) {\n LOG_ERROR(\"clone segment %s failed\", kHnswOffsetsSegmentId.c_str());\n return HnswRabitqEntity::Pointer();\n }\n\n SegmentGroupParam neighbor_group{neighbors, neighbors_meta, upper_neighbors,\n upper_neighbors_meta};\n\n HnswRabitqSearcherEntity *entity = new (std::nothrow)\n HnswRabitqSearcherEntity(header(), vectors, keys, mapping, neighbor_group,\n fixed_neighbors_, neighbors_in_memory_enabled_);\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"HnswRabitqSearcherEntity new failed\");\n }\n\n return HnswRabitqEntity::Pointer(entity);\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"hnsw_rabitq_builder_entity.h\"\n#include \"hnsw_rabitq_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswRabitqSearcherEntity : public HnswRabitqEntity {\n public:\n using Pointer = std::shared_ptr;\n using SegmentPointer = IndexStorage::Segment::Pointer;\n\n public:\n struct SegmentGroupParam {\n SegmentGroupParam(SegmentPointer neighbors_in,\n SegmentPointer neighbors_meta_in,\n SegmentPointer upper_neighbors_in,\n SegmentPointer upper_neighbors_meta_in)\n : neighbors{neighbors_in},\n neighbors_meta{neighbors_meta_in},\n upper_neighbors{upper_neighbors_in},\n upper_neighbors_meta{upper_neighbors_meta_in} {}\n\n SegmentPointer neighbors{nullptr};\n SegmentPointer neighbors_meta{nullptr};\n SegmentPointer upper_neighbors{nullptr};\n SegmentPointer upper_neighbors_meta{nullptr};\n };\n\n //! Constructor\n HnswRabitqSearcherEntity();\n\n //! Make a copy of searcher entity, to support thread-safe operation.\n //! The segment in container cannot be read concurrenly\n virtual const HnswRabitqEntity::Pointer clone() const override;\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const override;\n\n //! Get vector local id by key\n node_id_t get_id(key_t key) const;\n\n //! Get vector feature data by key\n virtual const void *get_vector_by_key(key_t key) const override;\n\n //! Get vector feature data by id\n virtual const void *get_vector(node_id_t id) const override;\n\n //! Get vector feature data by id\n virtual int get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const override;\n\n virtual int get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n virtual int get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const override;\n\n //! Get all vectors\n const void *get_vectors() const;\n\n //! Get the node id's neighbors on graph level\n virtual const Neighbors get_neighbors(level_t level,\n node_id_t id) const override;\n\n virtual int load(const IndexStorage::Pointer &container,\n bool check_crc) override;\n\n int load_segments(bool check_crc);\n\n virtual int cleanup(void) override;\n\n public:\n bool is_loaded() const {\n return loaded_;\n }\n\n void set_neighbors_in_memory(bool enabled) {\n neighbors_in_memory_enabled_ = enabled;\n }\n\n //! get fixed length neighbors data\n int get_fixed_neighbors(std::vector *fixed_neighbors) const;\n\n private:\n //! Constructor\n HnswRabitqSearcherEntity(const HNSWHeader &hd, const SegmentPointer &vectors,\n const SegmentPointer &keys,\n const SegmentPointer &mapping,\n const SegmentGroupParam &neighbor_group,\n const std::shared_ptr &fixed_neighbors,\n bool neighbors_in_memory_enabled)\n : HnswRabitqEntity(hd),\n vectors_(vectors),\n keys_(keys),\n mapping_(mapping),\n neighbors_(neighbor_group.neighbors),\n neighbors_meta_(neighbor_group.neighbors_meta),\n upper_neighbors_(neighbor_group.upper_neighbors),\n upper_neighbors_meta_(neighbor_group.upper_neighbors_meta),\n neighbors_in_memory_enabled_(neighbors_in_memory_enabled) {\n segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt()),\n IndexStorage::SegmentData(0U, 0U));\n fixed_neighbors_ = fixed_neighbors;\n }\n\n bool do_crc_check(std::vector &segments) const;\n\n inline size_t neighbors_size() const {\n return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t);\n }\n\n inline size_t upper_neighbors_size() const {\n return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t);\n }\n\n //! If neighbors_in_memory_enabled, load the level0 neighbors to memory\n int load_and_flat_neighbors(void);\n\n public:\n HnswRabitqSearcherEntity(const HnswRabitqSearcherEntity &) = delete;\n HnswRabitqSearcherEntity &operator=(const HnswRabitqSearcherEntity &) =\n delete;\n\n private:\n IndexStorage::Pointer storage_{};\n\n SegmentPointer vectors_{};\n SegmentPointer keys_{};\n SegmentPointer mapping_{};\n\n SegmentPointer neighbors_{};\n SegmentPointer neighbors_meta_{};\n SegmentPointer upper_neighbors_{};\n SegmentPointer upper_neighbors_meta_{};\n\n mutable std::vector segment_datas_{};\n std::shared_ptr fixed_neighbors_{}; // level 0 fixed size neighbors\n bool neighbors_in_memory_enabled_{false};\n bool loaded_{false};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_rabitq_streamer.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"algorithm/hnsw_rabitq/rabitq_reformer.h\"\n#include \"zvec/ailego/container/params.h\"\n#include \"zvec/ailego/logger/logger.h\"\n#include \"hnsw_rabitq_algorithm.h\"\n#include \"hnsw_rabitq_context.h\"\n#include \"hnsw_rabitq_dist_calculator.h\"\n#include \"hnsw_rabitq_index_provider.h\"\n#include \"hnsw_rabitq_query_entity.h\"\n#include \"rabitq_params.h\"\n#include \"rabitq_utils.h\"\n\nnamespace zvec {\nnamespace core {\nHnswRabitqStreamer::HnswRabitqStreamer() : entity_(stats_) {}\n\nHnswRabitqStreamer::HnswRabitqStreamer(IndexProvider::Pointer provider,\n RabitqReformer::Pointer reformer)\n : entity_(stats_),\n reformer_(std::move(reformer)),\n provider_(std::move(provider)) {}\n\nHnswRabitqStreamer::~HnswRabitqStreamer() {\n if (state_ == STATE_INITED) {\n this->cleanup();\n }\n}\n\nint HnswRabitqStreamer::init(const IndexMeta &imeta,\n const ailego::Params ¶ms) {\n meta_ = imeta;\n meta_.set_streamer(\"HnswRabitqStreamer\", HnswRabitqEntity::kRevision, params);\n\n params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_INDEX_SIZE, &max_index_size_);\n\n params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT,\n &upper_max_neighbor_cnt_);\n float multiplier = HnswRabitqEntity::kDefaultL0MaxNeighborCntMultiplier;\n params.get(PARAM_HNSW_RABITQ_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER,\n &multiplier);\n l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_;\n\n multiplier = HnswRabitqEntity::kDefaultNeighborPruneMultiplier;\n params.get(PARAM_HNSW_RABITQ_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier);\n size_t prune_cnt = multiplier * upper_max_neighbor_cnt_;\n scaling_factor_ = upper_max_neighbor_cnt_;\n params.get(PARAM_HNSW_RABITQ_STREAMER_SCALING_FACTOR, &scaling_factor_);\n\n params.get(PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT, &docs_hard_limit_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_EF, &ef_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_EFCONSTRUCTION, &ef_construction_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_ENABLE, &bf_enabled_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB,\n &bf_negative_prob_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_CHECK_CRC_ENABLE, &check_crc_enabled_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_CHUNK_SIZE, &chunk_size_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_FILTER_SAME_KEY, &filter_same_key_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_GET_VECTOR_ENABLE,\n &get_vector_enabled_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_FORCE_PADDING_RESULT_ENABLE,\n &force_padding_topk_enabled_);\n params.get(PARAM_HNSW_RABITQ_STREAMER_USE_ID_MAP, &use_id_map_);\n entity_.set_use_key_info_map(use_id_map_);\n\n params.get(PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit_);\n if (docs_soft_limit_ > 0 && docs_soft_limit_ > docs_hard_limit_) {\n LOG_ERROR(\"[%s] must be >= [%s]\",\n PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT.c_str(),\n PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT.c_str());\n return IndexError_InvalidArgument;\n } else if (docs_soft_limit_ == 0UL) {\n docs_soft_limit_ =\n docs_hard_limit_ * HnswRabitqEntity::kDefaultDocsSoftLimitRatio;\n }\n\n if (ef_ == 0U) {\n ef_ = HnswRabitqEntity::kDefaultEf;\n }\n if (ef_construction_ == 0U) {\n ef_construction_ = HnswRabitqEntity::kDefaultEfConstruction;\n }\n if (upper_max_neighbor_cnt_ == 0U) {\n upper_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt;\n }\n if (upper_max_neighbor_cnt_ > HnswRabitqEntity::kMaxNeighborCnt) {\n LOG_ERROR(\"[%s] must be in range (0,%d)\",\n PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT.c_str(),\n HnswRabitqEntity::kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (l0_max_neighbor_cnt_ == 0U) {\n l0_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultL0MaxNeighborCnt;\n }\n if (l0_max_neighbor_cnt_ > HnswRabitqEntity::kMaxNeighborCnt) {\n LOG_ERROR(\"MaxL0NeighborCnt must be in range (0,%d)\",\n HnswRabitqEntity::kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) {\n LOG_ERROR(\"[%s]-[%zu] must be <= [%s]-[%zu]\",\n PARAM_HNSW_RABITQ_STREAMER_MIN_NEIGHBOR_COUNT.c_str(),\n static_cast(min_neighbor_cnt_),\n PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT.c_str(),\n static_cast(upper_max_neighbor_cnt_));\n return IndexError_InvalidArgument;\n }\n\n if (bf_negative_prob_ <= 0.0f || bf_negative_prob_ >= 1.0f) {\n LOG_ERROR(\n \"[%s] must be in range (0,1)\",\n PARAM_HNSW_RABITQ_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (scaling_factor_ == 0U) {\n scaling_factor_ = HnswRabitqEntity::kDefaultScalingFactor;\n }\n if (scaling_factor_ < 5 || scaling_factor_ > 1000) {\n LOG_ERROR(\"[%s] must be in range [5,1000]\",\n PARAM_HNSW_RABITQ_STREAMER_SCALING_FACTOR.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) {\n LOG_ERROR(\"[%s] must be in range (0.0f,1.0f]\",\n PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_RATIO.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (max_scan_limit_ < min_scan_limit_) {\n LOG_ERROR(\"[%s] must be >= [%s]\",\n PARAM_HNSW_RABITQ_STREAMER_MAX_SCAN_LIMIT.c_str(),\n PARAM_HNSW_RABITQ_STREAMER_MIN_SCAN_LIMIT.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (prune_cnt == 0UL) {\n prune_cnt = upper_max_neighbor_cnt_;\n }\n if (chunk_size_ == 0UL) {\n chunk_size_ = HnswRabitqEntity::kDefaultChunkSize;\n }\n if (chunk_size_ > HnswRabitqEntity::kMaxChunkSize) {\n LOG_ERROR(\"[%s] must be < %zu\",\n PARAM_HNSW_RABITQ_STREAMER_CHUNK_SIZE.c_str(),\n HnswRabitqEntity::kMaxChunkSize);\n return IndexError_InvalidArgument;\n }\n uint32_t total_bits = 0;\n params.get(PARAM_RABITQ_TOTAL_BITS, &total_bits);\n if (total_bits == 0) {\n total_bits = kDefaultRabitqTotalBits;\n }\n if (total_bits < 1 || total_bits > 9) {\n LOG_ERROR(\"Invalid total_bits: %zu, must be in [1, 9]\", (size_t)total_bits);\n return IndexError_InvalidArgument;\n }\n uint8_t ex_bits = total_bits - 1;\n entity_.set_ex_bits(ex_bits);\n\n uint32_t dimension = 0;\n params.get(PARAM_HNSW_RABITQ_GENERAL_DIMENSION, &dimension);\n if (dimension == 0) {\n LOG_ERROR(\"%s not set\", PARAM_HNSW_RABITQ_GENERAL_DIMENSION.c_str());\n return IndexError_InvalidArgument;\n }\n if (dimension < kMinRabitqDimSize || dimension > kMaxRabitqDimSize) {\n LOG_ERROR(\"Invalid dimension: %u, must be in [%d, %d]\", dimension,\n kMinRabitqDimSize, kMaxRabitqDimSize);\n return IndexError_InvalidArgument;\n }\n entity_.update_rabitq_params_and_vector_size(dimension);\n\n entity_.set_ef_construction(ef_construction_);\n entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_);\n entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_);\n entity_.set_scaling_factor(scaling_factor_);\n entity_.set_prune_cnt(prune_cnt);\n\n entity_.set_chunk_size(chunk_size_);\n entity_.set_filter_same_key(filter_same_key_);\n entity_.set_get_vector(get_vector_enabled_);\n entity_.set_min_neighbor_cnt(min_neighbor_cnt_);\n\n int ret = entity_.init(docs_hard_limit_);\n if (ret != 0) {\n LOG_ERROR(\"Hnsw entity init failed for %s\", IndexError::What(ret));\n return ret;\n }\n\n LOG_DEBUG(\n \"Init params: maxIndexSize=%zu docsHardLimit=%zu docsSoftLimit=%zu \"\n \"efConstruction=%u ef=%u upperMaxNeighborCnt=%u l0MaxNeighborCnt=%u \"\n \"scalingFactor=%u maxScanRatio=%.3f minScanLimit=%zu maxScanLimit=%zu \"\n \"bfEnabled=%d bruteFoceThreshold=%zu bfNegativeProbability=%.5f \"\n \"checkCrcEnabled=%d pruneSize=%zu vectorSize=%u chunkSize=%zu \"\n \"filterSameKey=%u getVectorEnabled=%u minNeighborCount=%u \"\n \"forcePadding=%u \",\n max_index_size_, docs_hard_limit_, docs_soft_limit_, ef_construction_,\n ef_, upper_max_neighbor_cnt_, l0_max_neighbor_cnt_, scaling_factor_,\n max_scan_ratio_, min_scan_limit_, max_scan_limit_, bf_enabled_,\n bruteforce_threshold_, bf_negative_prob_, check_crc_enabled_, prune_cnt,\n meta_.element_size(), chunk_size_, filter_same_key_, get_vector_enabled_,\n min_neighbor_cnt_, force_padding_topk_enabled_);\n\n alg_ = HnswRabitqAlgorithm::UPointer(new HnswRabitqAlgorithm(entity_));\n\n ret = alg_->init();\n if (ret != 0) {\n return ret;\n }\n\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint HnswRabitqStreamer::cleanup(void) {\n if (state_ == STATE_OPENED) {\n this->close();\n }\n\n LOG_INFO(\"HnswRabitqStreamer cleanup\");\n\n meta_.clear();\n metric_.reset();\n stats_.clear();\n entity_.cleanup();\n\n if (alg_) {\n alg_->cleanup();\n }\n\n max_index_size_ = 0UL;\n docs_hard_limit_ = HnswRabitqEntity::kDefaultDocsHardLimit;\n docs_soft_limit_ = 0UL;\n upper_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt;\n l0_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultL0MaxNeighborCnt;\n ef_ = HnswRabitqEntity::kDefaultEf;\n ef_construction_ = HnswRabitqEntity::kDefaultEfConstruction;\n bf_enabled_ = false;\n scaling_factor_ = HnswRabitqEntity::kDefaultScalingFactor;\n bruteforce_threshold_ = HnswRabitqEntity::kDefaultBruteForceThreshold;\n max_scan_limit_ = HnswRabitqEntity::kDefaultMaxScanLimit;\n min_scan_limit_ = HnswRabitqEntity::kDefaultMinScanLimit;\n chunk_size_ = HnswRabitqEntity::kDefaultChunkSize;\n bf_negative_prob_ = HnswRabitqEntity::kDefaultBFNegativeProbability;\n max_scan_ratio_ = HnswRabitqEntity::kDefaultScanRatio;\n state_ = STATE_INIT;\n check_crc_enabled_ = false;\n filter_same_key_ = false;\n get_vector_enabled_ = false;\n\n return 0;\n}\n\nint HnswRabitqStreamer::open(IndexStorage::Pointer stg) {\n LOG_INFO(\"HnswRabitqStreamer open\");\n\n if (ailego_unlikely(state_ != STATE_INITED)) {\n LOG_ERROR(\"Open storage failed, init streamer first!\");\n return IndexError_NoReady;\n }\n\n // try to load reformer\n if (reformer_ == nullptr) {\n reformer_ = std::make_shared();\n ailego::Params reformer_params;\n reformer_params.set(PARAM_RABITQ_METRIC_NAME, meta_.metric_name());\n int ret = reformer_->init(reformer_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize RabitqReformer: %d\", ret);\n return ret;\n }\n\n ret = reformer_->load(stg);\n if (ret != 0) {\n LOG_ERROR(\"Failed to load reformer, ret=%d\", ret);\n return ret;\n }\n } else {\n if (!stg->has(RABITQ_CONVERTER_SEG_ID)) {\n int ret = reformer_->dump(stg);\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump reformer, ret=%d\", ret);\n return ret;\n }\n LOG_INFO(\"Dump reformer success.\");\n }\n }\n\n int ret = entity_.open(std::move(stg), max_index_size_, check_crc_enabled_);\n if (ret != 0) {\n return ret;\n }\n IndexMeta index_meta;\n ret = entity_.get_index_meta(&index_meta);\n if (ret == IndexError_NoExist) {\n // Set IndexMeta for the new index\n ret = entity_.set_index_meta(meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to set index meta for %s\", IndexError::What(ret));\n return ret;\n }\n } else if (ret != 0) {\n LOG_ERROR(\"Failed to get index meta for %s\", IndexError::What(ret));\n return ret;\n } else {\n if (index_meta.dimension() != meta_.dimension() ||\n index_meta.element_size() != meta_.element_size() ||\n index_meta.metric_name() != meta_.metric_name() ||\n index_meta.data_type() != meta_.data_type()) {\n LOG_ERROR(\"IndexMeta mismatch from the previous in index\");\n return IndexError_Mismatch;\n }\n // The IndexMetric Params may be updated like MipsSquaredEuclidean\n auto metric_params = index_meta.metric_params();\n metric_params.merge(meta_.metric_params());\n meta_.set_metric(index_meta.metric_name(), 0, metric_params);\n }\n\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"Failed to create metric %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"Failed to init metric, ret=%d\", ret);\n return ret;\n }\n\n if (!metric_->distance()) {\n LOG_ERROR(\"Invalid metric distance\");\n return IndexError_InvalidArgument;\n }\n\n if (!metric_->batch_distance()) {\n LOG_ERROR(\"Invalid metric batch distance\");\n return IndexError_InvalidArgument;\n }\n\n add_distance_ = metric_->distance();\n add_batch_distance_ = metric_->batch_distance();\n\n search_distance_ = add_distance_;\n search_batch_distance_ = add_batch_distance_;\n\n if (metric_->query_metric() && metric_->query_metric()->distance() &&\n metric_->query_metric()->batch_distance()) {\n search_distance_ = metric_->query_metric()->distance();\n search_batch_distance_ = metric_->query_metric()->batch_distance();\n }\n\n state_ = STATE_OPENED;\n magic_ = IndexContext::GenerateMagic();\n\n query_alg_ = HnswRabitqQueryAlgorithm::UPointer(new HnswRabitqQueryAlgorithm(\n entity_, reformer_->num_clusters(), reformer_->rabitq_metric_type()));\n\n return 0;\n}\n\nint HnswRabitqStreamer::close(void) {\n LOG_INFO(\"HnswRabitqStreamer close\");\n\n stats_.clear();\n meta_.set_metric(metric_->name(), 0, metric_->params());\n entity_.set_index_meta(meta_);\n int ret = entity_.close();\n if (ret != 0) {\n return ret;\n }\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint HnswRabitqStreamer::flush(uint64_t checkpoint) {\n LOG_INFO(\"HnswRabitqStreamer flush checkpoint=%zu\", (size_t)checkpoint);\n\n meta_.set_metric(metric_->name(), 0, metric_->params());\n entity_.set_index_meta(meta_);\n return entity_.flush(checkpoint);\n}\n\nint HnswRabitqStreamer::dump(const IndexDumper::Pointer &dumper) {\n LOG_INFO(\"HnswRabitqStreamer dump\");\n\n shared_mutex_.lock();\n AILEGO_DEFER([&]() { shared_mutex_.unlock(); });\n\n meta_.set_searcher(\"HnswRabitqSearcher\", HnswRabitqEntity::kRevision,\n ailego::Params());\n\n int ret = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return ret;\n }\n ret = reformer_->dump(dumper);\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump reformer into dumper.\");\n return ret;\n }\n return entity_.dump(dumper);\n}\n\nIndexStreamer::Context::Pointer HnswRabitqStreamer::create_context(void) const {\n if (ailego_unlikely(state_ != STATE_OPENED)) {\n LOG_ERROR(\"Create context failed, open storage first!\");\n return Context::Pointer();\n }\n\n HnswRabitqEntity::Pointer entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"CreateContext clone init failed\");\n return Context::Pointer();\n }\n HnswRabitqContext *ctx =\n new (std::nothrow) HnswRabitqContext(meta_.dimension(), metric_, entity);\n if (ailego_unlikely(ctx == nullptr)) {\n LOG_ERROR(\"Failed to new HnswRabitqContext\");\n return Context::Pointer();\n }\n ctx->set_ef(ef_);\n ctx->set_max_scan_limit(max_scan_limit_);\n ctx->set_min_scan_limit(min_scan_limit_);\n ctx->set_max_scan_ratio(max_scan_ratio_);\n ctx->set_filter_mode(bf_enabled_ ? VisitFilter::BloomFilter\n : VisitFilter::ByteMap);\n ctx->set_filter_negative_probability(bf_negative_prob_);\n ctx->set_magic(magic_);\n ctx->set_force_padding_topk(force_padding_topk_enabled_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n\n if (ailego_unlikely(ctx->init(HnswRabitqContext::kStreamerContext)) != 0) {\n LOG_ERROR(\"Init HnswRabitqContext failed\");\n delete ctx;\n return Context::Pointer();\n }\n uint32_t estimate_doc_count = 0;\n if (meta_.streamer_params().get(PARAM_HNSW_RABITQ_STREAMER_ESTIMATE_DOC_COUNT,\n &estimate_doc_count)) {\n LOG_DEBUG(\"HnswRabitqStreamer doc_count[%zu] estimate[%zu]\",\n (size_t)entity_.doc_cnt(), (size_t)estimate_doc_count);\n }\n ctx->check_need_adjuct_ctx(std::max(entity_.doc_cnt(), estimate_doc_count));\n\n return Context::Pointer(ctx);\n}\n\nIndexProvider::Pointer HnswRabitqStreamer::create_provider(void) const {\n LOG_DEBUG(\"HnswRabitqStreamer create provider\");\n\n auto entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"Clone HnswRabitqEntity failed\");\n return nullptr;\n }\n return Provider::Pointer(\n new HnswRabitqIndexProvider(meta_, entity, \"HnswRabitqStreamer\"));\n}\n\nint HnswRabitqStreamer::update_context(HnswRabitqContext *ctx) const {\n const HnswRabitqEntity::Pointer entity = entity_.clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone search context entity\");\n return IndexError_Runtime;\n }\n ctx->set_max_scan_limit(max_scan_limit_);\n ctx->set_min_scan_limit(min_scan_limit_);\n ctx->set_max_scan_ratio(max_scan_ratio_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n return ctx->update_context(HnswRabitqContext::kStreamerContext, meta_,\n metric_, entity, magic_);\n}\n\n//! Add a vector with id into index\nint HnswRabitqStreamer::add_with_id_impl(\n uint32_t id, const void *query, const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) {\n if (!provider_) {\n LOG_ERROR(\"Provider is nullptr, cannot add vector\");\n return IndexError_InvalidArgument;\n }\n\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n HnswRabitqContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswRabitqContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) {\n if (entity_.doc_cnt() >= docs_hard_limit_) {\n LOG_ERROR(\"Current docs %zu exceed [%s]\",\n static_cast(entity_.doc_cnt()),\n PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT.c_str());\n const std::lock_guard lk(mutex_);\n (*stats_.mutable_discarded_count())++;\n return IndexError_IndexFull;\n } else {\n LOG_WARN(\"Current docs %zu exceed [%s]\",\n static_cast(entity_.doc_cnt()),\n PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT.c_str());\n }\n }\n if (ailego_unlikely(!shared_mutex_.try_lock_shared())) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); });\n\n ctx->clear();\n ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_);\n ctx->reset_query(query);\n ctx->check_need_adjuct_ctx(entity_.doc_cnt());\n ctx->set_provider(provider_);\n\n if (metric_->support_train()) {\n const std::lock_guard lk(mutex_);\n ret = metric_->train(query, meta_.dimension());\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer metric train failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n }\n\n std::string converted_vector;\n IndexQueryMeta converted_meta;\n ret = reformer_->convert(query, qmeta, &converted_vector, &converted_meta);\n if (ret != 0) {\n LOG_ERROR(\"Rabitq hnsw convert failed, ret=%d\", ret);\n return ret;\n }\n\n level_t level = alg_->get_random_level();\n ret = entity_.add_vector_with_id(level, id, converted_vector.data());\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer add vector failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n ret = alg_->add_node(id, level, ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw steamer add node failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n if (ailego_unlikely(ctx->error())) {\n (*stats_.mutable_discarded_count())++;\n return IndexError_Runtime;\n }\n (*stats_.mutable_added_count())++;\n\n return 0;\n}\n\n//! Add a vector into index\nint HnswRabitqStreamer::add_impl(uint64_t pkey, const void *query,\n const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) {\n if (!provider_) {\n LOG_ERROR(\"Provider is nullptr, cannot add vector\");\n return IndexError_InvalidArgument;\n }\n\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n HnswRabitqContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswRabitqContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) {\n if (entity_.doc_cnt() >= docs_hard_limit_) {\n LOG_ERROR(\"Current docs %zu exceed [%s]\",\n static_cast(entity_.doc_cnt()),\n PARAM_HNSW_RABITQ_STREAMER_DOCS_HARD_LIMIT.c_str());\n const std::lock_guard lk(mutex_);\n (*stats_.mutable_discarded_count())++;\n return IndexError_IndexFull;\n } else {\n LOG_WARN(\"Current docs %zu exceed [%s]\",\n static_cast(entity_.doc_cnt()),\n PARAM_HNSW_RABITQ_STREAMER_DOCS_SOFT_LIMIT.c_str());\n }\n }\n if (ailego_unlikely(!shared_mutex_.try_lock_shared())) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); });\n\n ctx->clear();\n ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_);\n ctx->reset_query(query);\n ctx->check_need_adjuct_ctx(entity_.doc_cnt());\n ctx->set_provider(provider_);\n\n if (metric_->support_train()) {\n const std::lock_guard lk(mutex_);\n ret = metric_->train(query, meta_.dimension());\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer metric train failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n }\n\n std::string converted_vector;\n IndexQueryMeta converted_meta;\n ret = reformer_->convert(query, qmeta, &converted_vector, &converted_meta);\n if (ret != 0) {\n LOG_ERROR(\"Rabitq hnsw convert failed, ret=%d\", ret);\n return ret;\n }\n\n level_t level = alg_->get_random_level();\n node_id_t id;\n ret = entity_.add_vector(level, pkey, converted_vector.data(), &id);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer add vector failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n ret = alg_->add_node(id, level, ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw steamer add node failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n if (ailego_unlikely(ctx->error())) {\n (*stats_.mutable_discarded_count())++;\n return IndexError_Runtime;\n }\n (*stats_.mutable_added_count())++;\n\n return 0;\n}\n\n\nint HnswRabitqStreamer::search_impl(\n const void *query, const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) const {\n return search_impl(query, qmeta, 1, context);\n}\n\n//! Similarity search\nint HnswRabitqStreamer::search_impl(\n const void *query, const IndexQueryMeta &qmeta, uint32_t count,\n IndexStreamer::Context::Pointer &context) const {\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n HnswRabitqContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswRabitqContext failed\");\n return IndexError_Cast;\n }\n\n if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) {\n return search_bf_impl(query, qmeta, count, context);\n }\n\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_);\n ctx->resize_results(count);\n ctx->check_need_adjuct_ctx(entity_.doc_cnt());\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n ret = reformer_->transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ret = query_alg_->search(&entity, ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher fast search failed\");\n return ret;\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nvoid HnswRabitqStreamer::print_debug_info() {\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n Neighbors neighbours = entity_.get_neighbors(0, id);\n std::cout << \"node: \" << id << \"; \";\n if (neighbours.size() == 0) std::cout << std::endl;\n for (uint32_t i = 0; i < neighbours.size(); ++i) {\n std::cout << neighbours[i];\n\n if (i == neighbours.size() - 1) {\n std::cout << std::endl;\n } else {\n std::cout << \", \";\n }\n }\n }\n\n // entity_.print_key_map();\n}\n\nint HnswRabitqStreamer::search_bf_impl(\n const void *query, const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) const {\n return search_bf_impl(query, qmeta, 1, context);\n}\n\nint HnswRabitqStreamer::search_bf_impl(\n const void *query, const IndexQueryMeta &qmeta, uint32_t count,\n IndexStreamer::Context::Pointer &context) const {\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n HnswRabitqContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswRabitqContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_);\n ctx->resize_results(count);\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n ret = reformer_->transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw rabitq streamer transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ctx->group_topk_heaps().clear();\n\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n EstimateRecord dist;\n query_alg_->get_full_est(id, dist, entity);\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n } else {\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n ret = reformer_->transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw rabitq streamer transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ctx->topk_heap().clear();\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n EstimateRecord dist;\n query_alg_->get_full_est(id, dist, entity);\n ctx->topk_heap().emplace(id, dist);\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nint HnswRabitqStreamer::search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n int ret = check_params(query, qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n if (ailego_unlikely(p_keys.size() != count)) {\n LOG_ERROR(\"The size of p_keys is not equal to count\");\n return IndexError_InvalidArgument;\n }\n\n HnswRabitqContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswRabitqContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_);\n ctx->resize_results(count);\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_InvalidArgument;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n ret = reformer_->transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw rabitq streamer transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ctx->group_topk_heaps().clear();\n\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n EstimateRecord dist;\n query_alg_->get_full_est(id, dist, entity);\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n } else {\n for (size_t q = 0; q < count; ++q) {\n HnswRabitqQueryEntity entity;\n ret = reformer_->transform_to_entity(query, &entity);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw rabitq streamer transform failed\");\n return ret;\n }\n ctx->reset_query(query);\n ctx->topk_heap().clear();\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n key_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n EstimateRecord dist;\n query_alg_->get_full_est(id, dist, entity);\n ctx->topk_heap().emplace(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n query = static_cast(query) + qmeta.element_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"algorithm/hnsw_rabitq/rabitq_reformer.h\"\n#include \"zvec/core/framework/index_framework.h\"\n#include \"zvec/core/framework/index_provider.h\"\n#include \"zvec/core/framework/index_reformer.h\"\n#include \"hnsw_rabitq_algorithm.h\"\n#include \"hnsw_rabitq_query_algorithm.h\"\n#include \"hnsw_rabitq_streamer_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswRabitqStreamer : public IndexStreamer {\n public:\n using ContextPointer = IndexStreamer::Context::Pointer;\n\n HnswRabitqStreamer();\n explicit HnswRabitqStreamer(IndexProvider::Pointer provider,\n RabitqReformer::Pointer reformer = nullptr);\n virtual ~HnswRabitqStreamer(void);\n\n HnswRabitqStreamer(const HnswRabitqStreamer &streamer) = delete;\n HnswRabitqStreamer &operator=(const HnswRabitqStreamer &streamer) = delete;\n\n void set_provider(IndexProvider::Pointer provider) {\n provider_ = std::move(provider);\n }\n\n void set_reformer(IndexReformer::Pointer reformer) {\n reformer_ = std::dynamic_pointer_cast(reformer);\n }\n\n protected:\n //! Initialize Streamer\n virtual int init(const IndexMeta &imeta,\n const ailego::Params ¶ms) override;\n\n //! Cleanup Streamer\n virtual int cleanup(void) override;\n\n //! Create a context\n virtual Context::Pointer create_context(void) const override;\n\n //! Create a new iterator\n virtual IndexProvider::Pointer create_provider(void) const override;\n\n //! Add a vector into index\n virtual int add_impl(uint64_t pkey, const void *query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) override;\n\n //! Add a vector with id into index\n virtual int add_with_id_impl(uint32_t id, const void *query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) override;\n\n //! Similarity search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const override;\n\n //! Similarity brute force search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity brute force search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const override;\n\n //! Linear search by primary keys\n virtual int search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, ContextPointer &context) const override {\n return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context);\n }\n\n //! Linear search by primary keys\n virtual int search_bf_by_p_keys_impl(\n const void *query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const override;\n\n //! Fetch vector by key\n virtual const void *get_vector(uint64_t key) const override {\n return entity_.get_vector_by_key(key);\n }\n\n virtual int get_vector(const uint64_t key,\n IndexStorage::MemoryBlock &block) const override {\n return entity_.get_vector_by_key(key, block);\n }\n\n //! Fetch vector by id\n virtual const void *get_vector_by_id(uint32_t id) const override {\n return entity_.get_vector(id);\n }\n\n virtual int get_vector_by_id(\n const uint32_t id, IndexStorage::MemoryBlock &block) const override {\n return entity_.get_vector(id, block);\n }\n\n //! Open index from file path\n virtual int open(IndexStorage::Pointer stg) override;\n\n //! Close file\n virtual int close(void) override;\n\n //! flush file\n virtual int flush(uint64_t checkpoint) override;\n\n //! Dump index into storage\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n virtual const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve meta of index\n virtual const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n virtual void print_debug_info() override;\n\n private:\n inline int check_params(const void *query,\n const IndexQueryMeta &qmeta) const {\n if (ailego_unlikely(!query)) {\n LOG_ERROR(\"null query\");\n return IndexError_InvalidArgument;\n }\n if (ailego_unlikely(qmeta.dimension() != meta_.dimension() ||\n qmeta.data_type() != meta_.data_type() ||\n qmeta.element_size() != meta_.element_size())) {\n LOG_ERROR(\"Unsupported query meta\");\n return IndexError_Mismatch;\n }\n return 0;\n }\n\n inline int check_sparse_count_is_zero(const uint32_t *sparse_count,\n uint32_t count) const {\n for (uint32_t i = 0; i < count; ++i) {\n if (sparse_count[i] != 0)\n LOG_ERROR(\"Sparse cout is not empty. Index: %u, Sparse Count: %u\", i,\n sparse_count[i]);\n return IndexError_InvalidArgument;\n }\n\n return 0;\n }\n\n private:\n //! To share ctx across streamer/searcher, we need to update the context for\n //! current streamer/searcher\n int update_context(HnswRabitqContext *ctx) const;\n\n private:\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 };\n class Stats : public IndexStreamer::Stats {\n public:\n void clear(void) {\n set_revision_id(0u);\n set_loaded_count(0u);\n set_added_count(0u);\n set_discarded_count(0u);\n set_index_size(0u);\n set_dumped_size(0u);\n set_check_point(0u);\n set_create_time(0u);\n set_update_time(0u);\n clear_attributes();\n }\n };\n\n HnswRabitqStreamerEntity entity_;\n HnswRabitqAlgorithm::UPointer alg_;\n IndexMeta meta_{};\n IndexMetric::Pointer metric_{};\n\n IndexMetric::MatrixDistance add_distance_{};\n IndexMetric::MatrixDistance search_distance_{};\n\n IndexMetric::MatrixBatchDistance add_batch_distance_{};\n IndexMetric::MatrixBatchDistance search_batch_distance_{};\n\n RabitqReformer::Pointer reformer_{}; // RaBitQ reformer\n HnswRabitqQueryAlgorithm::UPointer query_alg_; // query algorithm\n // provider_ provides raw vector, which is used to build graph\n IndexProvider::Pointer provider_{};\n\n Stats stats_{};\n std::mutex mutex_{};\n\n size_t max_index_size_{0UL};\n size_t chunk_size_{HnswRabitqEntity::kDefaultChunkSize};\n size_t docs_hard_limit_{HnswRabitqEntity::kDefaultDocsHardLimit};\n size_t docs_soft_limit_{0UL};\n uint32_t min_neighbor_cnt_{0u};\n uint32_t upper_max_neighbor_cnt_{\n HnswRabitqEntity::kDefaultUpperMaxNeighborCnt};\n uint32_t l0_max_neighbor_cnt_{HnswRabitqEntity::kDefaultL0MaxNeighborCnt};\n uint32_t ef_{HnswRabitqEntity::kDefaultEf};\n uint32_t ef_construction_{HnswRabitqEntity::kDefaultEfConstruction};\n uint32_t scaling_factor_{HnswRabitqEntity::kDefaultScalingFactor};\n size_t bruteforce_threshold_{HnswRabitqEntity::kDefaultBruteForceThreshold};\n size_t max_scan_limit_{HnswRabitqEntity::kDefaultMaxScanLimit};\n size_t min_scan_limit_{HnswRabitqEntity::kDefaultMinScanLimit};\n float bf_negative_prob_{HnswRabitqEntity::kDefaultBFNegativeProbability};\n float max_scan_ratio_{HnswRabitqEntity::kDefaultScanRatio};\n\n uint32_t magic_{0U};\n State state_{STATE_INIT};\n bool bf_enabled_{false};\n bool check_crc_enabled_{false};\n bool filter_same_key_{false};\n bool get_vector_enabled_{false};\n bool force_padding_topk_enabled_{false};\n bool use_id_map_{true};\n\n //! avoid add vector while dumping index\n ailego::SharedMutex shared_mutex_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"hnsw_rabitq_streamer_entity.h\"\n#include \n\n// #define DEBUG_PRINT\n\nnamespace zvec {\nnamespace core {\n\nHnswRabitqStreamerEntity::HnswRabitqStreamerEntity(IndexStreamer::Stats &stats)\n : stats_(stats) {}\n\nHnswRabitqStreamerEntity::~HnswRabitqStreamerEntity() {}\n\nint HnswRabitqStreamerEntity::init(size_t max_doc_cnt) {\n if (std::pow(scaling_factor(), kMaxGraphLayers) < max_doc_cnt) {\n LOG_ERROR(\"scalingFactor=%zu is too small\", scaling_factor());\n return IndexError_InvalidArgument;\n }\n\n std::lock_guard lock(mutex_);\n broker_ = std::make_shared(stats_);\n upper_neighbor_index_ = std::make_shared();\n keys_map_lock_ = std::make_shared();\n keys_map_ = std::make_shared>();\n if (!keys_map_ || !upper_neighbor_index_ || !broker_ || !keys_map_lock_) {\n LOG_ERROR(\"HnswRabitqStreamerEntity new object failed\");\n return IndexError_NoMemory;\n }\n keys_map_->set_empty_key(kInvalidKey);\n\n neighbor_size_ = neighbors_size();\n upper_neighbor_size_ = upper_neighbors_size();\n\n //! vector + key + level 0 neighbors\n size_t size = vector_size() + sizeof(key_t) + neighbor_size_;\n\n size = AlignSize(size);\n set_node_size(size);\n return 0;\n}\n\nint HnswRabitqStreamerEntity::cleanup() {\n std::lock_guard lock(mutex_);\n mutable_header()->clear();\n chunk_size_ = kDefaultChunkSize;\n node_index_mask_bits_ = 0U;\n node_index_mask_ = 0U;\n node_cnt_per_chunk_ = 0U;\n neighbor_size_ = 0U;\n upper_neighbor_size_ = 0U;\n if (upper_neighbor_index_) {\n upper_neighbor_index_->cleanup();\n }\n if (keys_map_) {\n keys_map_->clear();\n }\n node_chunks_.clear();\n upper_neighbor_chunks_.clear();\n filter_same_key_ = false;\n get_vector_enabled_ = false;\n broker_.reset();\n\n return 0;\n}\n\nint HnswRabitqStreamerEntity::update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) {\n char buffer[neighbor_size_];\n NeighborsHeader *hd = reinterpret_cast(buffer);\n hd->neighbor_cnt = neighbors.size();\n size_t i = 0;\n for (; i < neighbors.size(); ++i) {\n hd->neighbors[i] = neighbors[i].first;\n }\n\n auto loc = get_neighbor_chunk_loc(level, id);\n size_t size = reinterpret_cast(&hd->neighbors[i]) - &buffer[0];\n size_t ret = loc.first->write(loc.second, hd, size);\n if (ailego_unlikely(ret != size)) {\n LOG_ERROR(\"Write neighbor header failed, ret=%zu\", ret);\n\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nconst Neighbors HnswRabitqStreamerEntity::get_neighbors(level_t level,\n node_id_t id) const {\n Chunk *chunk = nullptr;\n size_t offset = 0UL;\n size_t neighbor_size = neighbor_size_;\n if (level == 0UL) {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n offset =\n (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t);\n\n sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, chunk_idx,\n &node_chunks_);\n ailego_assert_with(chunk_idx < node_chunks_.size(), \"invalid chunk idx\");\n chunk = node_chunks_[chunk_idx].get();\n } else {\n auto p = get_upper_neighbor_chunk_loc(level, id);\n chunk = upper_neighbor_chunks_[p.first].get();\n offset = p.second;\n neighbor_size = upper_neighbor_size_;\n }\n\n ailego_assert_with(offset < chunk->data_size(), \"invalid chunk offset\");\n IndexStorage::MemoryBlock neighbor_block;\n size_t size = chunk->read(offset, neighbor_block, neighbor_size);\n if (ailego_unlikely(size != neighbor_size)) {\n LOG_ERROR(\"Read neighbor header failed, ret=%zu\", size);\n return Neighbors();\n }\n return Neighbors(std::move(neighbor_block));\n}\n\n//! Get vector data by key\nconst void *HnswRabitqStreamerEntity::get_vector(node_id_t id) const {\n auto loc = get_vector_chunk_loc(id);\n const void *vec = nullptr;\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = vector_size();\n\n size_t ret = node_chunks_[loc.first]->read(loc.second, &vec, read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%zu, read size=%zu, ret=%zu\",\n static_cast(loc.second), read_size, ret);\n }\n\n return vec;\n}\n\nint HnswRabitqStreamerEntity::get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const {\n for (auto i = 0U; i < count; ++i) {\n auto loc = get_vector_chunk_loc(ids[i]);\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = vector_size();\n\n size_t ret = node_chunks_[loc.first]->read(loc.second, &vecs[i], read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%zu, read size=%zu, ret=%zu\",\n static_cast(loc.second), read_size, ret);\n return IndexError_ReadData;\n }\n }\n return 0;\n}\n\nint HnswRabitqStreamerEntity::get_vector(\n const node_id_t id, IndexStorage::MemoryBlock &block) const {\n auto loc = get_vector_chunk_loc(id);\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = vector_size();\n\n size_t ret = node_chunks_[loc.first]->read(loc.second, block, read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%zu, read size=%zu, ret=%zu\",\n static_cast(loc.second), read_size, ret);\n return IndexError_ReadData;\n }\n return 0;\n}\n\nint HnswRabitqStreamerEntity::get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const {\n vec_blocks.resize(count);\n for (auto i = 0U; i < count; ++i) {\n auto loc = get_vector_chunk_loc(ids[i]);\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = vector_size();\n\n size_t ret =\n node_chunks_[loc.first]->read(loc.second, vec_blocks[i], read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%zu, read size=%zu, ret=%zu\",\n static_cast(loc.second), read_size, ret);\n return IndexError_ReadData;\n }\n }\n return 0;\n}\n\nkey_t HnswRabitqStreamerEntity::get_key(node_id_t id) const {\n if (use_key_info_map_) {\n auto loc = get_key_chunk_loc(id);\n IndexStorage::MemoryBlock key_block;\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n size_t ret =\n node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t));\n if (ailego_unlikely(ret != sizeof(key_t))) {\n LOG_ERROR(\"Read vector failed, ret=%zu\", ret);\n return kInvalidKey;\n }\n\n return *reinterpret_cast(key_block.data());\n } else {\n return id;\n }\n}\n\nvoid HnswRabitqStreamerEntity::add_neighbor(level_t level, node_id_t id,\n uint32_t size,\n node_id_t neighbor_id) {\n auto loc = get_neighbor_chunk_loc(level, id);\n size_t offset =\n loc.second + sizeof(NeighborsHeader) + size * sizeof(node_id_t);\n ailego_assert_with(size < neighbor_cnt(level), \"invalid neighbor size\");\n ailego_assert_with(offset < loc.first->data_size(), \"invalid chunk offset\");\n size_t ret = loc.first->write(offset, &neighbor_id, sizeof(node_id_t));\n if (ailego_unlikely(ret != sizeof(node_id_t))) {\n LOG_ERROR(\"Write neighbor id failed, ret=%zu\", ret);\n return;\n }\n\n uint32_t neighbors = size + 1;\n ret = loc.first->write(loc.second, &neighbors, sizeof(uint32_t));\n if (ailego_unlikely(ret != sizeof(uint32_t))) {\n LOG_ERROR(\"Write neighbor cnt failed, ret=%zu\", ret);\n }\n\n return;\n}\n\nint HnswRabitqStreamerEntity::init_chunks(const Chunk::Pointer &header_chunk) {\n if (header_chunk->data_size() < header_size()) {\n LOG_ERROR(\"Invalid header chunk size\");\n return IndexError_InvalidFormat;\n }\n IndexStorage::MemoryBlock header_block;\n size_t size = header_chunk->read(0UL, header_block, header_size());\n if (ailego_unlikely(size != header_size())) {\n LOG_ERROR(\"Read header chunk failed\");\n return IndexError_ReadData;\n }\n *mutable_header() =\n *reinterpret_cast(header_block.data());\n\n int ret = check_hnsw_index(&header());\n if (ret != 0) {\n broker_->close();\n return ret;\n }\n\n node_chunks_.resize(\n broker_->get_chunk_cnt(HnswRabitqChunkBroker::CHUNK_TYPE_NODE));\n for (auto seq = 0UL; seq < node_chunks_.size(); ++seq) {\n node_chunks_[seq] =\n broker_->get_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, seq);\n if (!node_chunks_[seq]) {\n LOG_ERROR(\"Missing hnsw streamer data chunk %zu th of %zu\", seq,\n node_chunks_.size());\n return IndexError_InvalidFormat;\n }\n }\n\n upper_neighbor_chunks_.resize(\n broker_->get_chunk_cnt(HnswRabitqChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR));\n for (auto seq = 0UL; seq < upper_neighbor_chunks_.size(); ++seq) {\n upper_neighbor_chunks_[seq] = broker_->get_chunk(\n HnswRabitqChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, seq);\n if (!upper_neighbor_chunks_[seq]) {\n LOG_ERROR(\"Missing hnsw streamer index chunk %zu th of %zu\", seq,\n upper_neighbor_chunks_.size());\n return IndexError_InvalidFormat;\n }\n }\n\n return 0;\n}\n\nint HnswRabitqStreamerEntity::open(IndexStorage::Pointer stg,\n uint64_t max_index_size, bool check_crc) {\n std::lock_guard lock(mutex_);\n bool huge_page = stg->isHugePage();\n LOG_DEBUG(\"huge_page: %d\", (int)huge_page);\n int ret = init_chunk_params(max_index_size, huge_page);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"init_chunk_params failed for %s\", IndexError::What(ret));\n return ret;\n }\n ret = broker_->open(std::move(stg), max_index_size_, chunk_size_, check_crc);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Open index failed for %s\", IndexError::What(ret));\n return ret;\n }\n ret = upper_neighbor_index_->init(broker_, upper_neighbor_chunk_size_,\n scaling_factor(), estimate_doc_capacity(),\n kUpperHashMemoryInflateRatio);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Init neighbor hash map failed\");\n return ret;\n }\n\n //! init header\n auto header_chunk =\n broker_->get_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_HEADER,\n HnswRabitqChunkBroker::kDefaultChunkSeqId);\n if (!header_chunk) { // open empty index, create one\n auto p = broker_->alloc_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_HEADER,\n HnswRabitqChunkBroker::kDefaultChunkSeqId,\n header_size());\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc header chunk failed\");\n return p.first;\n }\n size_t size = p.second->write(0UL, &header(), header_size());\n if (ailego_unlikely(size != header_size())) {\n LOG_ERROR(\"Write header chunk failed\");\n return IndexError_WriteData;\n }\n return 0;\n }\n\n //! Open an exist hnsw index\n ret = init_chunks(header_chunk);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n //! total docs including features wrote in index but neighbors may not ready\n node_id_t total_vecs = 0;\n if (node_chunks_.size() > 0) {\n size_t last_idx = node_chunks_.size() - 1;\n auto last_chunk = node_chunks_[last_idx];\n if (last_chunk->data_size() % node_size()) {\n LOG_WARN(\"The index may broken\");\n return IndexError_InvalidFormat;\n }\n total_vecs = last_idx * node_cnt_per_chunk_ +\n node_chunks_[last_idx]->data_size() / node_size();\n }\n\n LOG_INFO(\n \"Open index, l0NeighborCnt=%zu upperNeighborCnt=%zu \"\n \"efConstruction=%zu curDocCnt=%u totalVecs=%u maxLevel=%u\",\n l0_neighbor_cnt(), upper_neighbor_cnt(), ef_construction(), doc_cnt(),\n total_vecs, cur_max_level());\n //! try to correct the docCnt if index not fully flushed\n if (doc_cnt() != total_vecs) {\n LOG_WARN(\"Index closed abnormally, using totalVecs as curDocCnt\");\n *mutable_doc_cnt() = total_vecs;\n }\n if (filter_same_key_ || get_vector_enabled_) {\n if (use_key_info_map_) {\n for (node_id_t id = 0U; id < doc_cnt(); ++id) {\n if (get_key(id) == kInvalidKey) {\n continue;\n }\n (*keys_map_)[get_key(id)] = id;\n }\n }\n }\n\n stats_.set_loaded_count(doc_cnt());\n\n return 0;\n}\n\nint HnswRabitqStreamerEntity::close() {\n LOG_DEBUG(\"close index\");\n\n std::lock_guard lock(mutex_);\n flush_header();\n mutable_header()->reset();\n upper_neighbor_index_->cleanup();\n keys_map_->clear();\n header_.clear();\n node_chunks_.clear();\n upper_neighbor_chunks_.clear();\n\n return broker_->close();\n}\n\nint HnswRabitqStreamerEntity::flush(uint64_t checkpoint) {\n LOG_INFO(\"Flush index, curDocs=%zu\", static_cast(doc_cnt()));\n\n std::lock_guard lock(mutex_);\n flush_header();\n int ret = broker_->flush(checkpoint);\n if (ret != 0) {\n return ret;\n }\n\n return 0;\n}\n\nint HnswRabitqStreamerEntity::dump(const IndexDumper::Pointer &dumper) {\n LOG_INFO(\"Dump index, curDocs=%zu\", static_cast(doc_cnt()));\n\n //! sort by keys, to support get_vector by key in searcher\n std::vector keys(doc_cnt());\n for (node_id_t i = 0; i < doc_cnt(); ++i) {\n keys[i] = get_key(i);\n }\n\n //! dump neighbors\n auto get_level = [&](node_id_t id) {\n auto it = upper_neighbor_index_->find(id);\n if (it == upper_neighbor_index_->end()) {\n return 0U;\n };\n auto meta = reinterpret_cast(&it->second);\n return meta->level;\n };\n auto ret = dump_segments(dumper, keys.data(), get_level);\n if (ailego_unlikely(ret < 0)) {\n return ret;\n }\n *stats_.mutable_dumped_size() += ret;\n\n return 0;\n}\n\nint HnswRabitqStreamerEntity::check_hnsw_index(const HNSWHeader *hd) const {\n if (l0_neighbor_cnt() != hd->l0_neighbor_cnt() ||\n upper_neighbor_cnt() != hd->upper_neighbor_cnt()) {\n LOG_ERROR(\"Param neighbor cnt: %zu:%zu mismatch index previous %zu:%zu\",\n l0_neighbor_cnt(), upper_neighbor_cnt(), hd->l0_neighbor_cnt(),\n hd->upper_neighbor_cnt());\n return IndexError_Mismatch;\n }\n if (vector_size() != hd->vector_size()) {\n LOG_ERROR(\"vector size %zu mismatch index previous %zu\", vector_size(),\n hd->vector_size());\n return IndexError_Mismatch;\n }\n if (ef_construction() != hd->ef_construction()) {\n LOG_WARN(\"Param efConstruction %zu mismatch index previous %zu\",\n ef_construction(), hd->ef_construction());\n }\n if (scaling_factor() != hd->scaling_factor()) {\n LOG_WARN(\"Param scalingFactor %zu mismatch index previous %zu\",\n scaling_factor(), hd->scaling_factor());\n return IndexError_Mismatch;\n }\n if (prune_cnt() != hd->neighbor_prune_cnt()) {\n LOG_WARN(\"Param pruneCnt %zu mismatch index previous %zu\", prune_cnt(),\n hd->neighbor_prune_cnt());\n return IndexError_Mismatch;\n }\n if ((hd->entry_point() != kInvalidNodeId &&\n hd->entry_point() >= hd->doc_cnt()) ||\n (hd->entry_point() == kInvalidNodeId && hd->doc_cnt() > 0U)) {\n LOG_WARN(\"Invalid entryPoint %zu, docCnt %zu\",\n static_cast(hd->entry_point()),\n static_cast(hd->doc_cnt()));\n return IndexError_InvalidFormat;\n }\n if (hd->entry_point() == kInvalidNodeId &&\n broker_->get_chunk_cnt(HnswRabitqChunkBroker::CHUNK_TYPE_NODE) > 0) {\n LOG_WARN(\"The index is broken, maybe it haven't flush\");\n return IndexError_InvalidFormat;\n }\n\n return 0;\n}\n\nint HnswRabitqStreamerEntity::add_vector(level_t level, key_t key,\n const void *vec, node_id_t *id) {\n Chunk::Pointer node_chunk;\n size_t chunk_offset = -1UL;\n\n std::lock_guard lock(mutex_);\n // duplicate check\n if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) {\n LOG_WARN(\"Try to add duplicate key, ignore it\");\n return IndexError_Duplicate;\n }\n\n node_id_t local_id = static_cast(doc_cnt());\n uint32_t chunk_index = node_chunks_.size() - 1U;\n if (chunk_index == -1U ||\n (node_chunks_[chunk_index]->data_size() >=\n node_cnt_per_chunk_ * node_size())) { // no space left and need to alloc\n if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n chunk_index++;\n auto p = broker_->alloc_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_NODE,\n chunk_index, chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n node_chunk = p.second;\n chunk_offset = 0UL;\n node_chunks_.emplace_back(node_chunk);\n } else {\n node_chunk = node_chunks_[chunk_index];\n chunk_offset = node_chunk->data_size();\n }\n\n size_t size = node_chunk->write(chunk_offset, vec, vector_size());\n if (ailego_unlikely(size != vector_size())) {\n LOG_ERROR(\"Chunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t));\n if (ailego_unlikely(size != sizeof(key_t))) {\n LOG_ERROR(\"Chunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n //! level 0 neighbors is inited to zero by default\n\n int ret = add_upper_neighbor(level, local_id);\n if (ret != 0) {\n return ret;\n }\n\n chunk_offset += node_size();\n if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"Chunk resize to %zu failed\", chunk_offset);\n return IndexError_Runtime;\n }\n if (filter_same_key_ || get_vector_enabled_) {\n if (use_key_info_map_) {\n keys_map_lock_->lock();\n (*keys_map_)[key] = local_id;\n keys_map_lock_->unlock();\n }\n }\n\n *mutable_doc_cnt() += 1;\n broker_->mark_dirty();\n *id = local_id;\n\n return 0;\n}\n\nint HnswRabitqStreamerEntity::add_vector_with_id(level_t level, node_id_t id,\n const void *vec) {\n Chunk::Pointer node_chunk;\n size_t chunk_offset = -1UL;\n key_t key = id;\n\n std::lock_guard lock(mutex_);\n\n // duplicate check\n if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) {\n LOG_WARN(\"Try to add duplicate key, ignore it\");\n return IndexError_Duplicate;\n }\n\n // set node_chunk & chunk_offset if succeed\n auto func_get_node_chunk_and_offset = [&](node_id_t node_id) -> int {\n uint32_t chunk_index = node_id >> node_index_mask_bits_;\n ailego_assert_with(chunk_index <= node_chunks_.size(), \"invalid chunk idx\");\n // belongs to next chunk\n if (chunk_index == node_chunks_.size()) {\n if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n auto p = broker_->alloc_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_NODE,\n chunk_index, chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n node_chunk = p.second;\n node_chunks_.emplace_back(node_chunk);\n }\n\n node_chunk = node_chunks_[chunk_index];\n chunk_offset = (node_id & node_index_mask_) * node_size();\n return 0;\n };\n\n for (size_t start_id = doc_cnt(); start_id < id; ++start_id) {\n if (auto ret = func_get_node_chunk_and_offset(start_id); ret != 0) {\n LOG_ERROR(\"func_get_node_chunk_and_offset failed\");\n return ret;\n }\n size_t size = node_chunk->write(chunk_offset + vector_size(), &kInvalidKey,\n sizeof(key_t));\n if (ailego_unlikely(size != sizeof(key_t))) {\n LOG_ERROR(\"Chunk write key failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n chunk_offset += node_size();\n if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"Chunk resize to %zu failed\", chunk_offset);\n return IndexError_Runtime;\n }\n }\n\n if (auto ret = func_get_node_chunk_and_offset(id); ret != 0) {\n LOG_ERROR(\"func_get_node_chunk_and_offset failed\");\n return ret;\n }\n\n size_t size = node_chunk->write(chunk_offset, vec, vector_size());\n if (ailego_unlikely(size != vector_size())) {\n LOG_ERROR(\"Chunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t));\n if (ailego_unlikely(size != sizeof(key_t))) {\n LOG_ERROR(\"Chunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n //! level 0 neighbors is inited to zero by default\n\n int ret = add_upper_neighbor(level, id);\n if (ret != 0) {\n return ret;\n }\n\n if (*mutable_doc_cnt() <= id) {\n *mutable_doc_cnt() = id + 1;\n chunk_offset += node_size();\n if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"Chunk resize to %zu failed\", chunk_offset);\n return IndexError_Runtime;\n }\n }\n\n if (filter_same_key_ || get_vector_enabled_) {\n if (use_key_info_map_) {\n keys_map_lock_->lock();\n (*keys_map_)[key] = id;\n keys_map_lock_->unlock();\n }\n }\n\n broker_->mark_dirty();\n\n return 0;\n}\n\nvoid HnswRabitqStreamerEntity::update_ep_and_level(node_id_t ep,\n level_t level) {\n HnswRabitqEntity::update_ep_and_level(ep, level);\n flush_header();\n\n return;\n}\n\nconst HnswRabitqEntity::Pointer HnswRabitqStreamerEntity::clone() const {\n std::vector node_chunks;\n node_chunks.reserve(node_chunks_.size());\n for (size_t i = 0UL; i < node_chunks_.size(); ++i) {\n node_chunks.emplace_back(node_chunks_[i]->clone());\n if (ailego_unlikely(!node_chunks[i])) {\n LOG_ERROR(\"HnswRabitqStreamerEntity get chunk failed in clone\");\n return HnswRabitqEntity::Pointer();\n }\n }\n\n std::vector upper_neighbor_chunks;\n upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size());\n for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) {\n upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone());\n if (ailego_unlikely(!upper_neighbor_chunks[i])) {\n LOG_ERROR(\"HnswRabitqStreamerEntity get chunk failed in clone\");\n return HnswRabitqEntity::Pointer();\n }\n }\n\n HnswRabitqStreamerEntity *entity =\n new (std::nothrow) HnswRabitqStreamerEntity(\n stats_, header(), chunk_size_, node_index_mask_bits_,\n upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_,\n upper_neighbor_index_, keys_map_lock_, keys_map_, use_key_info_map_,\n std::move(node_chunks), std::move(upper_neighbor_chunks), broker_);\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"HnswRabitqStreamerEntity new failed\");\n }\n return HnswRabitqEntity::Pointer(entity);\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \"zvec/core/framework/index_framework.h\"\n#include \"hnsw_rabitq_chunk.h\"\n#include \"hnsw_rabitq_entity.h\"\n#include \"hnsw_rabitq_index_hash.h\"\n#include \"hnsw_rabitq_params.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! HnswRabitqStreamerEntity manage vector data, pkey, and node's neighbors\nclass HnswRabitqStreamerEntity : public HnswRabitqEntity {\n public:\n //! Cleanup\n //! return 0 on success, or errCode in failure\n virtual int cleanup() override;\n\n //! Make a copy of streamer entity, to support thread-safe operation.\n //! The segment in container cannot be read concurrenly\n virtual const HnswRabitqEntity::Pointer clone() const override;\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const override;\n\n //! Get vector feature data by key\n virtual const void *get_vector(node_id_t id) const override;\n\n //! Get vectors feature data by local ids\n virtual int get_vector(const node_id_t *ids, uint32_t count,\n const void **vecs) const override;\n\n virtual int get_vector(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n\n virtual int get_vector(\n const node_id_t *ids, uint32_t count,\n std::vector &vec_blocks) const override;\n\n //! Get the node id's neighbors on graph level\n //! Note: the neighbors cannot be modified, using the following\n //! method to get WritableNeighbors if want to\n virtual const Neighbors get_neighbors(level_t level,\n node_id_t id) const override;\n\n //! Add vector and key to hnsw entity, and local id will be saved in id\n virtual int add_vector(level_t level, key_t key, const void *vec,\n node_id_t *id) override;\n\n //! Add vector and id to hnsw entity\n virtual int add_vector_with_id(level_t level, node_id_t id,\n const void *vec) override;\n\n virtual int update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors)\n override;\n\n //! Append neighbor_id to node id neighbors on level\n //! Notice: the caller must be ensure the neighbors not full\n virtual void add_neighbor(level_t level, node_id_t id, uint32_t size,\n node_id_t neighbor_id) override;\n\n //! Dump index by dumper\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n virtual void update_ep_and_level(node_id_t ep, level_t level) override;\n\n void set_use_key_info_map(bool use_id_map) {\n use_key_info_map_ = use_id_map;\n LOG_DEBUG(\"use_key_info_map_: %d\", (int)use_key_info_map_);\n }\n\n public:\n //! Constructor\n HnswRabitqStreamerEntity(IndexStreamer::Stats &stats);\n\n //! Destructor\n ~HnswRabitqStreamerEntity();\n\n //! Get vector feature data by key\n virtual const void *get_vector_by_key(key_t key) const override {\n auto id = get_id(key);\n return id == kInvalidNodeId ? nullptr : get_vector(id);\n }\n\n virtual int get_vector_by_key(\n const key_t key, IndexStorage::MemoryBlock &block) const override {\n auto id = get_id(key);\n if (id != kInvalidNodeId) {\n return get_vector(id, block);\n } else {\n return IndexError_InvalidArgument;\n }\n }\n\n //! Init entity\n int init(size_t max_doc_cnt);\n\n //! Flush graph entity to disk\n //! return 0 on success, or errCode in failure\n int flush(uint64_t checkpoint);\n\n //! Open entity from storage\n //! return 0 on success, or errCode in failure\n int open(IndexStorage::Pointer stg, uint64_t max_index_size, bool check_crc);\n\n //! Close entity\n //! return 0 on success, or errCode in failure\n int close();\n\n //! Set meta information from entity\n int set_index_meta(const IndexMeta &meta) const {\n return IndexHelper::SerializeToStorage(meta, broker_->storage().get());\n }\n\n //! Get meta information from entity\n int get_index_meta(IndexMeta *meta) const {\n return IndexHelper::DeserializeFromStorage(broker_->storage().get(), meta);\n }\n\n //! Set params: chunk size\n inline void set_chunk_size(size_t val) {\n chunk_size_ = val;\n }\n\n //! Set params\n inline void set_filter_same_key(bool val) {\n filter_same_key_ = val;\n }\n\n //! Set params\n inline void set_get_vector(bool val) {\n get_vector_enabled_ = val;\n }\n\n //! Get vector local id by key\n inline node_id_t get_id(key_t key) const {\n if (use_key_info_map_) {\n keys_map_lock_->lock_shared();\n auto it = keys_map_->find(key);\n keys_map_lock_->unlock_shared();\n return it == keys_map_->end() ? kInvalidNodeId : it->second;\n } else {\n return key;\n }\n }\n\n void print_key_map() const {\n std::cout << \"key map begins\" << std::endl;\n\n auto iter = keys_map_->begin();\n while (iter != keys_map_->end()) {\n std::cout << \"key: \" << iter->first << \", id: \" << iter->second\n << std::endl;\n ;\n iter++;\n }\n\n std::cout << \"key map ends\" << std::endl;\n }\n\n //! Get l0 neighbors size\n inline size_t neighbors_size() const {\n return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t);\n }\n\n //! Get neighbors size for level > 0\n inline size_t upper_neighbors_size() const {\n return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t);\n }\n\n\n private:\n union UpperNeighborIndexMeta {\n struct {\n uint32_t level : 4;\n uint32_t index : 28; // index is composite type: chunk idx, and the\n // N th neighbors in chunk, they two composite\n // the 28 bits location\n };\n uint32_t data;\n };\n\n template \n using HashMap = google::dense_hash_map>;\n template \n using HashMapPointer = std::shared_ptr>;\n\n template \n using HashSet = google::dense_hash_set>;\n template \n using HashSetPointer = std::shared_ptr>;\n\n //! upper neighbor index hashmap\n using NIHashMap = HnswIndexHashMap;\n using NIHashMapPointer = std::shared_ptr;\n\n //! Private construct, only be called by clone method\n HnswRabitqStreamerEntity(IndexStreamer::Stats &stats, const HNSWHeader &hd,\n size_t chunk_size, uint32_t node_index_mask_bits,\n uint32_t upper_neighbor_mask_bits,\n bool filter_same_key, bool get_vector_enabled,\n const NIHashMapPointer &upper_neighbor_index,\n std::shared_ptr &keys_map_lock,\n const HashMapPointer &keys_map,\n bool use_key_info_map,\n std::vector &&node_chunks,\n std::vector &&upper_neighbor_chunks,\n const HnswRabitqChunkBroker::Pointer &broker)\n : stats_(stats),\n chunk_size_(chunk_size),\n node_index_mask_bits_(node_index_mask_bits),\n node_cnt_per_chunk_(1UL << node_index_mask_bits_),\n node_index_mask_(node_cnt_per_chunk_ - 1),\n upper_neighbor_mask_bits_(upper_neighbor_mask_bits),\n upper_neighbor_mask_((1U << upper_neighbor_mask_bits_) - 1),\n filter_same_key_(filter_same_key),\n get_vector_enabled_(get_vector_enabled),\n use_key_info_map_(use_key_info_map),\n upper_neighbor_index_(upper_neighbor_index),\n keys_map_lock_(keys_map_lock),\n keys_map_(keys_map),\n node_chunks_(std::move(node_chunks)),\n upper_neighbor_chunks_(std::move(upper_neighbor_chunks)),\n broker_(broker) {\n *mutable_header() = hd;\n\n neighbor_size_ = neighbors_size();\n upper_neighbor_size_ = upper_neighbors_size();\n }\n\n //! Called only in searching procedure per context, so no need to lock\n void sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE type, size_t idx,\n std::vector *chunks) const {\n if (ailego_likely(idx < chunks->size())) {\n return;\n }\n for (size_t i = chunks->size(); i <= idx; ++i) {\n auto chunk = broker_->get_chunk(type, i);\n // the storage can ensure get chunk will success after the first get\n ailego_assert_with(!!chunk, \"get chunk failed\");\n chunks->emplace_back(std::move(chunk));\n }\n }\n\n //! return pair: chunk index + chunk offset\n inline std::pair get_vector_chunk_loc(\n node_id_t id) const {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n uint32_t offset = (id & node_index_mask_) * node_size();\n\n sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, chunk_idx,\n &node_chunks_);\n return std::make_pair(chunk_idx, offset);\n }\n\n //! return pair: chunk index + chunk offset\n inline std::pair get_key_chunk_loc(node_id_t id) const {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n uint32_t offset = (id & node_index_mask_) * node_size() + vector_size();\n\n sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, chunk_idx,\n &node_chunks_);\n return std::make_pair(chunk_idx, offset);\n }\n\n inline std::pair get_upper_neighbor_chunk_loc(\n level_t level, node_id_t id) const {\n auto it = upper_neighbor_index_->find(id);\n ailego_assert_abort(it != upper_neighbor_index_->end(),\n \"Get upper neighbor header failed\");\n auto meta = reinterpret_cast(&it->second);\n uint32_t chunk_idx = (meta->index) >> upper_neighbor_mask_bits_;\n uint32_t offset = (((meta->index) & upper_neighbor_mask_) + level - 1) *\n upper_neighbor_size_;\n sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, chunk_idx,\n &upper_neighbor_chunks_);\n ailego_assert_abort(chunk_idx < upper_neighbor_chunks_.size(),\n \"invalid chunk idx\");\n ailego_assert_abort(offset < upper_neighbor_chunks_[chunk_idx]->data_size(),\n \"invalid chunk offset\");\n return std::make_pair(chunk_idx, offset);\n }\n\n //! return pair: chunk + chunk offset\n inline std::pair get_neighbor_chunk_loc(level_t level,\n node_id_t id) const {\n if (level == 0UL) {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n uint32_t offset =\n (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t);\n\n sync_chunks(HnswRabitqChunkBroker::CHUNK_TYPE_NODE, chunk_idx,\n &node_chunks_);\n ailego_assert_abort(chunk_idx < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_abort(offset < node_chunks_[chunk_idx]->data_size(),\n \"invalid chunk offset\");\n return std::make_pair(node_chunks_[chunk_idx].get(), offset);\n } else {\n auto p = get_upper_neighbor_chunk_loc(level, id);\n return std::make_pair(upper_neighbor_chunks_[p.first].get(), p.second);\n }\n }\n\n //! Chunk hnsw index valid\n int check_hnsw_index(const HNSWHeader *hd) const;\n\n size_t get_total_upper_neighbors_size(level_t level) const {\n return level * upper_neighbor_size_;\n }\n\n //! Add upper neighbor header and reserve space for upper neighbor\n int add_upper_neighbor(level_t level, node_id_t id) {\n if (level == 0) {\n return 0;\n }\n Chunk::Pointer chunk;\n uint64_t chunk_offset = -1UL;\n size_t neighbors_size = get_total_upper_neighbors_size(level);\n uint64_t chunk_index = upper_neighbor_chunks_.size() - 1UL;\n if (chunk_index == -1UL ||\n (upper_neighbor_chunks_[chunk_index]->padding_size() <\n neighbors_size)) { // no space left and need to alloc\n chunk_index++;\n if (ailego_unlikely(upper_neighbor_chunks_.capacity() ==\n upper_neighbor_chunks_.size())) {\n LOG_ERROR(\"add upper neighbor failed for no memory quota\");\n return IndexError_IndexFull;\n }\n auto p =\n broker_->alloc_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR,\n chunk_index, upper_neighbor_chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n chunk = p.second;\n chunk_offset = 0UL;\n upper_neighbor_chunks_.emplace_back(chunk);\n } else {\n chunk = upper_neighbor_chunks_[chunk_index];\n chunk_offset = chunk->data_size();\n }\n ailego_assert_with((size_t)level < kMaxGraphLayers, \"invalid level\");\n ailego_assert_with(chunk_offset % upper_neighbor_size_ == 0,\n \"invalid offset\");\n ailego_assert_with((chunk_offset / upper_neighbor_size_) <\n (1U << upper_neighbor_mask_bits_),\n \"invalid offset\");\n ailego_assert_with(chunk_index < (1U << (28 - upper_neighbor_mask_bits_)),\n \"invalid chunk index\");\n UpperNeighborIndexMeta meta;\n meta.level = level;\n meta.index = (chunk_index << upper_neighbor_mask_bits_) |\n (chunk_offset / upper_neighbor_size_);\n chunk_offset += upper_neighbor_size_ * level;\n if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) {\n LOG_ERROR(\"HashMap insert value failed\");\n return IndexError_Runtime;\n }\n\n if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"Chunk resize to %zu failed\", (size_t)chunk_offset);\n return IndexError_Runtime;\n }\n\n return 0;\n }\n\n size_t estimate_doc_capacity() const {\n return node_chunks_.capacity() * node_cnt_per_chunk_;\n }\n\n int init_chunk_params(size_t max_index_size, bool huge_page) {\n node_cnt_per_chunk_ = std::max(1, chunk_size_ / node_size());\n //! align node cnt per chunk to pow of 2\n node_index_mask_bits_ = std::ceil(std::log2(node_cnt_per_chunk_));\n node_cnt_per_chunk_ = 1UL << node_index_mask_bits_;\n if (huge_page) {\n chunk_size_ = AlignHugePageSize(node_cnt_per_chunk_ * node_size());\n } else {\n chunk_size_ = AlignPageSize(node_cnt_per_chunk_ * node_size());\n }\n node_index_mask_ = node_cnt_per_chunk_ - 1;\n\n if (max_index_size == 0UL) {\n max_index_size_ = chunk_size_ * kDefaultMaxChunkCnt;\n } else {\n max_index_size_ = max_index_size;\n }\n\n //! To get a balanced upper neighbor chunk size.\n //! If the upper chunk size is equal to node chunk size, it may waste\n //! upper neighbor chunk space; if the upper neighbor chunk size is too\n //! small, the will need large upper neighbor chunks index space. So to\n //! get a balanced ratio be sqrt of the node/neighbor size ratio\n float ratio =\n std::sqrt(node_size() * scaling_factor() * 1.0f / upper_neighbor_size_);\n if (huge_page) {\n upper_neighbor_chunk_size_ = AlignHugePageSize(\n std::max(get_total_upper_neighbors_size(kMaxGraphLayers),\n static_cast(chunk_size_ / ratio)));\n } else {\n upper_neighbor_chunk_size_ = AlignPageSize(\n std::max(get_total_upper_neighbors_size(kMaxGraphLayers),\n static_cast(chunk_size_ / ratio)));\n }\n upper_neighbor_mask_bits_ =\n std::ceil(std::log2(upper_neighbor_chunk_size_ / upper_neighbor_size_));\n upper_neighbor_mask_ = (1 << upper_neighbor_mask_bits_) - 1;\n\n size_t max_node_chunk_cnt = std::ceil(max_index_size_ / chunk_size_);\n size_t max_upper_chunk_cnt = std::ceil(\n (max_node_chunk_cnt * node_cnt_per_chunk_ * 1.0f / scaling_factor()) /\n (upper_neighbor_chunk_size_ / upper_neighbor_size_));\n max_upper_chunk_cnt =\n max_upper_chunk_cnt + std::ceil(max_upper_chunk_cnt / scaling_factor());\n\n //! reserve space to avoid memmove in chunks vector emplace chunk, so\n //! as to lock-free in reading chunk\n node_chunks_.reserve(max_node_chunk_cnt);\n upper_neighbor_chunks_.reserve(max_upper_chunk_cnt);\n\n LOG_DEBUG(\n \"Settings: nodeSize=%zu chunkSize=%u upperNeighborSize=%u \"\n \"upperNeighborChunkSize=%u \"\n \"nodeCntPerChunk=%u maxChunkCnt=%zu maxNeighborChunkCnt=%zu \"\n \"maxIndexSize=%zu ratio=%.3f\",\n node_size(), chunk_size_, upper_neighbor_size_,\n upper_neighbor_chunk_size_, node_cnt_per_chunk_, max_node_chunk_cnt,\n max_upper_chunk_cnt, max_index_size_, ratio);\n\n return 0;\n }\n\n //! Init node chunk and neighbor chunks\n int init_chunks(const Chunk::Pointer &header_chunk);\n\n int flush_header(void) {\n if (!broker_->dirty()) {\n // do not need to flush\n return 0;\n }\n auto header_chunk =\n broker_->get_chunk(HnswRabitqChunkBroker::CHUNK_TYPE_HEADER,\n HnswRabitqChunkBroker::kDefaultChunkSeqId);\n if (ailego_unlikely(!header_chunk)) {\n LOG_ERROR(\"get header chunk failed\");\n return IndexError_Runtime;\n }\n size_t size = header_chunk->write(0UL, &header(), header_size());\n if (ailego_unlikely(size != header_size())) {\n LOG_ERROR(\"Write header chunk failed\");\n return IndexError_WriteData;\n }\n\n return 0;\n }\n\n private:\n HnswRabitqStreamerEntity(const HnswRabitqStreamerEntity &) = delete;\n HnswRabitqStreamerEntity &operator=(const HnswRabitqStreamerEntity &) =\n delete;\n static constexpr uint64_t kUpperHashMemoryInflateRatio = 2.0f;\n\n private:\n IndexStreamer::Stats &stats_;\n HNSWHeader header_{};\n std::mutex mutex_{};\n size_t max_index_size_{0UL};\n uint32_t chunk_size_{kDefaultChunkSize};\n uint32_t upper_neighbor_chunk_size_{kDefaultChunkSize};\n uint32_t node_index_mask_bits_{0U};\n uint32_t node_cnt_per_chunk_{0U};\n uint32_t node_index_mask_{0U};\n uint32_t neighbor_size_{0U};\n uint32_t upper_neighbor_size_{0U};\n //! UpperNeighborIndex.index composite chunkIdx and offset in chunk by the\n //! following mask\n uint32_t upper_neighbor_mask_bits_{0U};\n uint32_t upper_neighbor_mask_{0U};\n bool filter_same_key_{false};\n bool get_vector_enabled_{false};\n bool use_key_info_map_{true};\n\n NIHashMapPointer upper_neighbor_index_{};\n\n mutable std::shared_ptr keys_map_lock_{};\n HashMapPointer keys_map_{};\n\n //! the chunks will be changed in searcher, so need mutable\n //! data chunk include: vector, key, level 0 neighbors\n mutable std::vector node_chunks_{};\n\n //! upper neighbor chunk inlude: UpperNeighborHeader + (1~level) neighbors\n mutable std::vector upper_neighbor_chunks_{};\n\n HnswRabitqChunkBroker::Pointer broker_{}; // chunk broker\n};\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_rabitq/rabitq_converter.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"rabitq_converter.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"ailego/pattern/defer.h\"\n#include \"algorithm/hnsw_rabitq/rabitq_reformer.h\"\n#include \"zvec/core/framework/index_cluster.h\"\n#include \"zvec/core/framework/index_error.h\"\n#include \"zvec/core/framework/index_factory.h\"\n#include \"zvec/core/framework/index_features.h\"\n#include \"zvec/core/framework/index_holder.h\"\n#include \"zvec/core/framework/index_memory.h\"\n#include \"zvec/core/framework/index_meta.h\"\n#include \"rabitq_params.h\"\n#include \"rabitq_utils.h\"\n\n#ifdef _MSC_VER\n#define strncasecmp _strnicmp\n#endif\n\nnamespace zvec {\nnamespace core {\n\nRabitqConverter::~RabitqConverter() {\n this->cleanup();\n}\n\nint RabitqConverter::init(const IndexMeta &meta, const ailego::Params ¶ms) {\n // Copy meta and ensure it has metric information\n meta_ = meta;\n dimension_ = meta.dimension();\n\n if (meta_.metric_name().empty()) {\n LOG_ERROR(\"Meta metric is empty\");\n return IndexError_InvalidArgument;\n }\n\n // Round up dimension to multiple of 64\n padded_dim_ = ((dimension_ + 63) / 64) * 64;\n\n // Get RaBitQ parameters with defaults\n uint32_t total_bits = 0;\n params.get(PARAM_RABITQ_TOTAL_BITS, &total_bits);\n if (total_bits == 0) {\n total_bits = kDefaultRabitqTotalBits;\n }\n if (total_bits < 1 || total_bits > 9) {\n LOG_ERROR(\"Invalid total_bits: %zu, must be in [1, 9]\", (size_t)total_bits);\n return IndexError_InvalidArgument;\n }\n ex_bits_ = total_bits - 1;\n\n params.get(PARAM_RABITQ_NUM_CLUSTERS, &num_clusters_);\n if (num_clusters_ == 0) {\n num_clusters_ = kDefaultNumClusters;\n }\n\n if (ex_bits_ > 8) {\n LOG_ERROR(\"Invalid ex_bits: %zu, must be <= 8\", ex_bits_);\n return IndexError_InvalidArgument;\n }\n\n if (meta.data_type() != IndexMeta::DataType::DT_FP32) {\n LOG_ERROR(\"RaBitQ only supports FP32 data type\");\n return IndexError_Unsupported;\n }\n params.get(PARAM_RABITQ_SAMPLE_COUNT, &sample_count_);\n\n std::string rotator_type_str;\n params.get(PARAM_RABITQ_ROTATOR_TYPE, &rotator_type_str);\n if (rotator_type_str.empty()) {\n rotator_type_ = rabitqlib::RotatorType::FhtKacRotator;\n } else if (strncasecmp(rotator_type_str.c_str(), \"fht\", 3) == 0) {\n rotator_type_ = rabitqlib::RotatorType::FhtKacRotator;\n } else if (strncasecmp(rotator_type_str.c_str(), \"matrix\", 6) == 0) {\n rotator_type_ = rabitqlib::RotatorType::MatrixRotator;\n } else {\n LOG_ERROR(\"Invalid rotator_type: %s\", rotator_type_str.c_str());\n return IndexError_InvalidArgument;\n }\n\n // Create rotator\n rotator_.reset(\n rabitqlib::choose_rotator(dimension_, rotator_type_, padded_dim_));\n\n LOG_INFO(\n \"RabitqConverter initialized: dim=%zu, padded_dim=%zu, \"\n \"num_clusters=%zu, ex_bits=%zu, rotator_type=%d[%s] sample_count[%zu]\",\n dimension_, padded_dim_, num_clusters_, ex_bits_, (int)rotator_type_,\n rotator_type_str.c_str(), sample_count_);\n\n return 0;\n}\n\nint RabitqConverter::cleanup() {\n centroids_.clear();\n rotated_centroids_.clear();\n result_holder_.reset();\n rotator_.reset();\n return 0;\n}\n\nint RabitqConverter::train(IndexHolder::Pointer holder) {\n if (!holder) {\n LOG_ERROR(\"Null holder for training\");\n return IndexError_InvalidArgument;\n }\n\n ailego::ElapsedTime timer;\n\n size_t vector_count = holder->count();\n if (vector_count == 0) {\n LOG_ERROR(\"No vectors for training\");\n return IndexError_InvalidArgument;\n }\n\n // do sampling from all data\n size_t sample_count = vector_count;\n if (sample_count_ > 0) {\n sample_count = std::min(sample_count_, vector_count);\n }\n LOG_INFO(\"Training with %zu vectors from %zu of holder\", sample_count,\n vector_count);\n auto sampler = std::make_shared>(\n meta_, sample_count);\n auto iter = holder->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Create iterator error\");\n return IndexError_Runtime;\n }\n for (; iter->is_valid(); iter->next()) {\n sampler->emplace(iter->data());\n }\n\n // Holder is not needed, cleanup it.\n holder.reset();\n\n if (sampler->count() == 0) {\n LOG_ERROR(\"Load training data error\");\n return IndexError_InvalidLength;\n }\n\n\n // Create KmeansCluster for training centroids\n auto cluster = IndexFactory::CreateCluster(\"OptKmeansCluster\");\n if (!cluster) {\n LOG_ERROR(\"Failed to create OptKmeansCluster\");\n return IndexError_NoExist;\n }\n\n // Initialize cluster\n LOG_INFO(\n \"Initializing KmeansCluster with meta: dim=%u, data_type=%d, metric=%s\",\n meta_.dimension(), (int)meta_.data_type(), meta_.metric_name().c_str());\n ailego::Params cluster_params;\n int ret = cluster->init(meta_, cluster_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize KmeansCluster: %d\", ret);\n return ret;\n }\n\n ret = cluster->mount(sampler);\n if (ret != 0) {\n LOG_ERROR(\"Failed to mount training data: %d\", ret);\n return ret;\n }\n cluster->suggest(num_clusters_);\n\n // Perform clustering\n IndexCluster::CentroidList cents;\n // TODO: support specify threads with argument\n auto threads = std::make_shared(0, false);\n ret = cluster->cluster(threads, cents);\n if (ret != 0) {\n LOG_ERROR(\"Failed to perform clustering: %d\", ret);\n return ret;\n }\n\n if (cents.size() != num_clusters_) {\n LOG_WARN(\"Expected %zu clusters, got %zu\", num_clusters_, cents.size());\n num_clusters_ = cents.size();\n }\n // Extract original centroids (for LinearSeeker query)\n centroids_.resize(num_clusters_ * dimension_);\n // Extract rotated centroids (for quantization)\n rotated_centroids_.resize(num_clusters_ * padded_dim_);\n for (uint32_t i = 0; i < num_clusters_; ++i) {\n const float *cent_data = static_cast(cents[i].feature());\n // Save original centroids\n std::memcpy(¢roids_[i * dimension_], cent_data,\n dimension_ * sizeof(float));\n // Save rotated centroids\n this->rotator_->rotate(cent_data, &rotated_centroids_[i * padded_dim_]);\n }\n\n stats_.set_trained_count(sampler->count());\n stats_.set_trained_costtime(timer.milli_seconds());\n\n LOG_INFO(\"Training completed: %zu centroids, cost %zu ms\", num_clusters_,\n static_cast(timer.milli_seconds()));\n\n return 0;\n}\n\n\nint RabitqConverter::transform(IndexHolder::Pointer holder) {\n if (!holder) {\n LOG_ERROR(\"Null holder for transformation\");\n return IndexError_InvalidArgument;\n }\n\n if (rotated_centroids_.empty()) {\n LOG_ERROR(\"Centroids not trained yet\");\n return IndexError_NoReady;\n }\n\n LOG_ERROR(\"Not implemented\");\n return IndexError_NotImplemented;\n}\n\nint RabitqConverter::dump(const IndexDumper::Pointer &dumper) {\n if (!dumper) {\n LOG_ERROR(\"Null dumper\");\n return IndexError_InvalidArgument;\n }\n\n if (rotated_centroids_.empty() || centroids_.empty()) {\n LOG_ERROR(\"No centroids to dump\");\n return IndexError_NoReady;\n }\n\n ailego::ElapsedTime timer;\n size_t dumped_size = 0;\n\n int ret = dump_rabitq_centroids(\n dumper, dimension_, padded_dim_, ex_bits_, num_clusters_, rotator_type_,\n rotated_centroids_, centroids_, rotator_, &dumped_size);\n if (ret != 0) {\n return ret;\n }\n\n stats_.set_dumped_size(dumped_size);\n stats_.set_dumped_costtime(timer.milli_seconds());\n\n LOG_INFO(\"Dump completed: %zu bytes, cost %zu ms\", stats_.dumped_size(),\n static_cast(timer.milli_seconds()));\n return 0;\n}\n\nint RabitqConverter::to_reformer(IndexReformer::Pointer *reformer) {\n auto memory_dumper = IndexFactory::CreateDumper(\"MemoryDumper\");\n memory_dumper->init(ailego::Params());\n std::string file_id = ailego::StringHelper::Concat(\n \"rabitq_converter_\", ailego::Monotime::MilliSeconds(), rand());\n int ret = memory_dumper->create(file_id);\n if (ret != 0) {\n LOG_ERROR(\"Failed to create memory dumper: %d\", ret);\n return ret;\n }\n // Release memory\n AILEGO_DEFER([&file_id]() { IndexMemory::Instance()->remove(file_id); });\n ret = this->dump(memory_dumper);\n if (ret != 0) {\n LOG_ERROR(\"Failed to dump RabitqConverter: %d\", ret);\n return ret;\n }\n ret = memory_dumper->close();\n if (ret != 0) {\n LOG_ERROR(\"Failed to close memory dumper: %d\", ret);\n return ret;\n }\n\n auto res = std::make_shared();\n ailego::Params reformer_params;\n reformer_params.set(PARAM_RABITQ_METRIC_NAME, meta_.metric_name());\n ret = res->init(reformer_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize RabitqReformer: %d\", ret);\n return ret;\n }\n auto memory_storage = IndexFactory::CreateStorage(\"MemoryReadStorage\");\n ret = memory_storage->open(file_id, false);\n if (ret != 0) {\n LOG_ERROR(\"Failed to open memory storage: %d\", ret);\n return ret;\n }\n ret = res->load(memory_storage);\n if (ret != 0) {\n LOG_ERROR(\"Failed to load RabitqReformer: %d\", ret);\n return ret;\n }\n *reformer = std::move(res);\n return 0;\n}\n\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/rabitq_converter.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \"zvec/core/framework/index_cluster.h\"\n#include \"zvec/core/framework/index_converter.h\"\n#include \"zvec/core/framework/index_reformer.h\"\n#include \"zvec/core/framework/index_threads.h\"\n#include \"rabitq_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass RabitqReformer;\n\n/*! RaBitQ Converter\n * Trains KMeans centroids and quantizes vectors using RaBitQ\n */\nclass RabitqConverter : public IndexConverter {\n public:\n //! Constructor\n RabitqConverter() = default;\n\n //! Destructor\n ~RabitqConverter() override;\n\n //! Initialize Converter\n int init(const IndexMeta &meta, const ailego::Params ¶ms) override;\n\n //! Cleanup Converter\n int cleanup(void) override;\n\n //! Train the data - perform KMeans clustering\n int train(IndexHolder::Pointer holder) override;\n\n //! Transform the data - quantize vectors using RaBitQ\n int transform(IndexHolder::Pointer holder) override;\n\n //! Dump centroids and config into storage\n int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve a holder as result\n IndexHolder::Pointer result(void) const override {\n return result_holder_;\n }\n\n //! Retrieve Index Meta\n const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n int to_reformer(IndexReformer::Pointer *reformer) override;\n\n private:\n static inline size_t AlignSize(size_t size) {\n return (size + 0x1F) & (~0x1F);\n }\n\n private:\n IndexMeta meta_;\n IndexHolder::Pointer result_holder_;\n Stats stats_;\n size_t sample_count_{0};\n\n // RaBitQ parameters\n size_t num_clusters_{0};\n size_t ex_bits_{0};\n size_t dimension_{0};\n size_t padded_dim_{0};\n\n // Original centroids: num_clusters * dimension (for LinearSeeker query)\n std::vector centroids_;\n // Rotated centroids: num_clusters * padded_dim (for quantization)\n std::vector rotated_centroids_;\n\n // Rotator for vector transformation\n rabitqlib::RotatorType rotator_type_{rabitqlib::RotatorType::FhtKacRotator};\n std::unique_ptr> rotator_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/rabitq_params.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\n// Local metric type enum that mirrors rabitqlib::MetricType,\n// without exposing rabitqlib headers to consumers of this file.\nenum class RabitqMetricType {\n kL2 = 0,\n kIP = 1,\n};\n\n// RaBitQ Converter parameters\nstatic const std::string PARAM_RABITQ_NUM_CLUSTERS(\n \"proxima.rabitq.num_clusters\");\nstatic const std::string PARAM_RABITQ_TOTAL_BITS(\"proxima.rabitq.total_bits\");\nstatic const std::string PARAM_RABITQ_METRIC_NAME(\"proxima.rabitq.metric_name\");\nstatic const std::string PARAM_RABITQ_ROTATOR_TYPE(\n \"proxima.rabitq.rotator.type\");\nstatic const std::string PARAM_RABITQ_SAMPLE_COUNT(\n \"proxima.rabitq.sample_count\");\n\n// Default values\nconstexpr size_t kDefaultNumClusters = 16;\n// 4-bit, 5-bit, and 7-bit quantization typically achieve 90%, 95%, and 99%\n// recall, respectively—without accessing raw vectors for reranking\nconstexpr size_t kDefaultRabitqTotalBits = 7;\n\nconstexpr int kMinRabitqDimSize = 64;\nconstexpr int kMaxRabitqDimSize = 4095;\n\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/rabitq_reformer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"rabitq_reformer.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"core/algorithm/cluster/linear_seeker.h\"\n#include \"zvec/core/framework/index_error.h\"\n#include \"zvec/core/framework/index_factory.h\"\n#include \"zvec/core/framework/index_features.h\"\n#include \"zvec/core/framework/index_meta.h\"\n#include \"zvec/core/framework/index_storage.h\"\n#include \"hnsw_rabitq_query_entity.h\"\n#include \"rabitq_converter.h\"\n#include \"rabitq_utils.h\"\n\nnamespace zvec {\nnamespace core {\n\n// All rabitqlib types are confined to this translation unit via pimpl.\nstruct RabitqReformer::Impl {\n // RaBitQ parameters\n size_t num_clusters{0};\n size_t ex_bits{0};\n size_t dimension{0};\n size_t padded_dim{0};\n size_t size_bin_data{0};\n size_t size_ex_data{0};\n bool loaded{false};\n\n // Original centroids: num_clusters * dimension (for LinearSeeker query)\n std::vector centroids;\n // Rotated centroids: num_clusters * padded_dim (for quantization)\n std::vector rotated_centroids;\n\n rabitqlib::RotatorType rotator_type{rabitqlib::RotatorType::FhtKacRotator};\n std::unique_ptr> rotator;\n rabitqlib::quant::RabitqConfig query_config;\n rabitqlib::quant::RabitqConfig config;\n rabitqlib::MetricType metric_type{rabitqlib::METRIC_L2};\n\n LinearSeeker::Pointer centroid_seeker;\n CoherentIndexFeatures::Pointer centroid_features;\n\n // Translate local enum to rabitqlib enum (used only inside this .cc).\n static rabitqlib::MetricType to_rabitq(RabitqMetricType m) {\n return m == RabitqMetricType::kIP ? rabitqlib::METRIC_IP\n : rabitqlib::METRIC_L2;\n }\n\n // Translate rabitqlib enum to local enum.\n static RabitqMetricType from_rabitq(rabitqlib::MetricType m) {\n return m == rabitqlib::METRIC_IP ? RabitqMetricType::kIP\n : RabitqMetricType::kL2;\n }\n\n int quantize_vector(const float *raw_vector, uint32_t cluster_id,\n std::string *quantized_data) const;\n};\n\nRabitqReformer::RabitqReformer() : impl_(std::make_unique()) {}\n\nRabitqReformer::~RabitqReformer() {\n this->cleanup();\n}\n\nsize_t RabitqReformer::num_clusters() const {\n return impl_->num_clusters;\n}\n\nRabitqMetricType RabitqReformer::rabitq_metric_type() const {\n return Impl::from_rabitq(impl_->metric_type);\n}\n\nint RabitqReformer::init(const ailego::Params ¶ms) {\n std::string metric_name = params.get_as_string(PARAM_RABITQ_METRIC_NAME);\n if (metric_name == \"SquaredEuclidean\") {\n impl_->metric_type = rabitqlib::METRIC_L2;\n } else if (metric_name == \"InnerProduct\") {\n impl_->metric_type = rabitqlib::METRIC_IP;\n } else if (metric_name == \"Cosine\") {\n impl_->metric_type = rabitqlib::METRIC_IP;\n } else {\n LOG_ERROR(\"Unsupported metric name: %s\", metric_name.c_str());\n return IndexError_InvalidArgument;\n }\n LOG_DEBUG(\"Rabitq reformer init done. metric_name=%s metric_type=%d\",\n metric_name.c_str(), static_cast(impl_->metric_type));\n return 0;\n}\n\nint RabitqReformer::cleanup() {\n impl_->centroids.clear();\n impl_->rotated_centroids.clear();\n impl_->centroid_seeker.reset();\n impl_->centroid_features.reset();\n impl_->loaded = false;\n impl_->rotator.reset();\n return 0;\n}\n\nint RabitqReformer::unload() {\n return this->cleanup();\n}\n\nint RabitqReformer::load(IndexStorage::Pointer storage) {\n if (!storage) {\n LOG_ERROR(\"Invalid storage for load\");\n return IndexError_InvalidArgument;\n }\n\n auto segment = storage->get(RABITQ_CONVERTER_SEG_ID);\n if (!segment) {\n LOG_ERROR(\"Failed to get segment %s\", RABITQ_CONVERTER_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n\n size_t offset = 0;\n RabitqConverterHeader header;\n IndexStorage::MemoryBlock block;\n size_t size = segment->read(offset, block, sizeof(header));\n if (size != sizeof(header)) {\n LOG_ERROR(\"Failed to read header\");\n return IndexError_InvalidFormat;\n }\n memcpy(&header, block.data(), sizeof(header));\n impl_->dimension = header.dim;\n impl_->padded_dim = header.padded_dim;\n impl_->ex_bits = header.ex_bits;\n impl_->num_clusters = header.num_clusters;\n impl_->rotator_type =\n static_cast(header.rotator_type);\n offset += sizeof(header);\n\n // Read rotated centroids\n size_t rotated_centroids_size =\n sizeof(float) * header.num_clusters * header.padded_dim;\n size = segment->read(offset, block, rotated_centroids_size);\n if (size != rotated_centroids_size) {\n LOG_ERROR(\"Failed to read rotated centroids\");\n return IndexError_InvalidFormat;\n }\n impl_->rotated_centroids.resize(header.num_clusters * header.padded_dim);\n memcpy(impl_->rotated_centroids.data(), block.data(), rotated_centroids_size);\n offset += size;\n\n // Read original centroids (for LinearSeeker query)\n size_t centroids_size = sizeof(float) * header.num_clusters * header.dim;\n size = segment->read(offset, block, centroids_size);\n if (size != centroids_size) {\n LOG_ERROR(\"Failed to read centroids\");\n return IndexError_InvalidFormat;\n }\n impl_->centroids.resize(header.num_clusters * header.dim);\n memcpy(impl_->centroids.data(), block.data(), centroids_size);\n offset += size;\n\n // Read rotator\n size_t rotator_size = header.rotator_size;\n size = segment->read(offset, block, rotator_size);\n if (size != rotator_size) {\n LOG_ERROR(\"Failed to read rotator\");\n return IndexError_InvalidFormat;\n }\n impl_->rotator.reset(rabitqlib::choose_rotator(\n impl_->dimension, impl_->rotator_type, impl_->padded_dim));\n impl_->rotator->load(reinterpret_cast(block.data()));\n offset += size;\n\n impl_->query_config = rabitqlib::quant::faster_config(\n impl_->padded_dim, rabitqlib::SplitSingleQuery::kNumBits);\n impl_->config =\n rabitqlib::quant::faster_config(impl_->padded_dim, impl_->ex_bits + 1);\n\n impl_->size_bin_data =\n rabitqlib::BinDataMap::data_bytes(impl_->padded_dim);\n impl_->size_ex_data = rabitqlib::ExDataMap::data_bytes(\n impl_->padded_dim, impl_->ex_bits);\n\n // Initialize LinearSeeker for centroid search\n IndexMeta centroid_meta;\n centroid_meta.set_data_type(IndexMeta::DataType::DT_FP32);\n centroid_meta.set_dimension(static_cast(impl_->dimension));\n // Note:\n // 1. spherical kmeans is used for InnerProduct and Cosine, so centroids are\n // normalized.\n // 2. for Cosine metric, `transform_to_entity` input is normalized, need to\n // use InnerProduct metric as Cosine metric requires extra dimension which is\n // unsuitable for centroids.\n centroid_meta.set_metric(impl_->metric_type == rabitqlib::METRIC_L2\n ? \"SquaredEuclidean\"\n : \"InnerProduct\",\n 0, ailego::Params());\n\n impl_->centroid_features = std::make_shared();\n impl_->centroid_features->mount(centroid_meta, impl_->centroids.data(),\n impl_->centroids.size() * sizeof(float));\n\n impl_->centroid_seeker = std::make_shared();\n int ret = impl_->centroid_seeker->init(centroid_meta);\n if (ret != 0) {\n LOG_ERROR(\"Failed to init centroid seeker. ret[%d]\", ret);\n return ret;\n }\n ret = impl_->centroid_seeker->mount(impl_->centroid_features);\n if (ret != 0) {\n LOG_ERROR(\"Failed to mount centroid features. ret[%d]\", ret);\n return ret;\n }\n\n LOG_INFO(\n \"Rabitq reformer load done. dimension=%zu, padded_dim=%zu, \"\n \"ex_bits=%zu, num_clusters=%zu, size_bin_data=%zu, size_ex_data=%zu \"\n \"rotator_type=%d\",\n impl_->dimension, impl_->padded_dim, impl_->ex_bits, impl_->num_clusters,\n impl_->size_bin_data, impl_->size_ex_data, (int)impl_->rotator_type);\n impl_->loaded = true;\n return 0;\n}\n\nint RabitqReformer::convert(const void *record, const IndexQueryMeta &rmeta,\n std::string *out, IndexQueryMeta *ometa) const {\n if (!impl_->loaded) {\n LOG_ERROR(\"Centroids not loaded yet\");\n return IndexError_NoReady;\n }\n\n if (!record || !out) {\n LOG_ERROR(\"Invalid arguments for convert\");\n return IndexError_InvalidArgument;\n }\n\n // input may be transformed, require rmeta.dimension >= dimension\n if (rmeta.dimension() < impl_->dimension ||\n rmeta.data_type() != IndexMeta::DataType::DT_FP32) {\n LOG_ERROR(\"Invalid record meta: dimension=%zu, data_type=%d\",\n static_cast(rmeta.dimension()), (int)rmeta.data_type());\n return IndexError_InvalidArgument;\n }\n\n // Find nearest centroid using LinearSeeker\n Seeker::Document doc;\n int ret = impl_->centroid_seeker->seek(\n record, impl_->dimension * sizeof(float), &doc);\n if (ret != 0) {\n LOG_ERROR(\"Failed to seek centroid. ret[%d]\", ret);\n return ret;\n }\n uint32_t cluster_id = doc.index;\n\n const float *vector = static_cast(record);\n ret = impl_->quantize_vector(vector, cluster_id, out);\n if (ret != 0) {\n LOG_ERROR(\"Failed to quantize vector\");\n return ret;\n }\n\n ometa->set_meta(IndexMeta::DataType::DT_INT8, (uint32_t)out->size());\n return 0;\n}\n\nint RabitqReformer::transform(const void *, const IndexQueryMeta &,\n std::string *, IndexQueryMeta *) const {\n return IndexError_NotImplemented;\n}\n\nint RabitqReformer::transform_to_entity(const void *query,\n HnswRabitqQueryEntity *entity) const {\n if (!impl_->loaded) {\n LOG_ERROR(\"Centroids not loaded yet\");\n return IndexError_NoReady;\n }\n\n if (!query) {\n LOG_ERROR(\"Invalid arguments for transform\");\n return IndexError_InvalidArgument;\n }\n\n const float *query_vector = static_cast(query);\n\n // Apply rotator\n entity->rotated_query.resize(impl_->padded_dim);\n impl_->rotator->rotate(query_vector, entity->rotated_query.data());\n\n // Quantize query to 4-bit representation\n entity->query_wrapper = std::make_unique>(\n entity->rotated_query.data(), impl_->padded_dim, impl_->ex_bits,\n impl_->query_config, impl_->metric_type);\n\n // Preprocess - get the distance from query to all centroids\n entity->q_to_centroids.resize(impl_->num_clusters);\n\n if (impl_->metric_type == rabitqlib::METRIC_L2) {\n for (size_t i = 0; i < impl_->num_clusters; i++) {\n entity->q_to_centroids[i] = std::sqrt(rabitqlib::euclidean_sqr(\n entity->rotated_query.data(),\n impl_->rotated_centroids.data() + (i * impl_->padded_dim),\n impl_->padded_dim));\n }\n } else if (impl_->metric_type == rabitqlib::METRIC_IP) {\n entity->q_to_centroids.resize(impl_->num_clusters * 2);\n // first half as g_add, second half as g_error\n for (size_t i = 0; i < impl_->num_clusters; i++) {\n entity->q_to_centroids[i] = rabitqlib::dot_product(\n entity->rotated_query.data(),\n impl_->rotated_centroids.data() + (i * impl_->padded_dim),\n impl_->padded_dim);\n entity->q_to_centroids[i + impl_->num_clusters] =\n std::sqrt(rabitqlib::euclidean_sqr(\n entity->rotated_query.data(),\n impl_->rotated_centroids.data() + (i * impl_->padded_dim),\n impl_->padded_dim));\n }\n }\n\n return 0;\n}\n\nint RabitqReformer::Impl::quantize_vector(const float *raw_vector,\n uint32_t cluster_id,\n std::string *quantized_data) const {\n std::vector rotated_data(padded_dim);\n rotator->rotate(raw_vector, rotated_data.data());\n\n // quantized format: cluster_id + bin_data + ex_data\n quantized_data->resize(sizeof(cluster_id) + size_bin_data + size_ex_data);\n memcpy(&(*quantized_data)[0], &cluster_id, sizeof(cluster_id));\n int bin_data_offset = sizeof(cluster_id);\n int ex_data_offset = bin_data_offset + size_bin_data;\n rabitqlib::quant::quantize_split_single(\n rotated_data.data(), rotated_centroids.data() + (cluster_id * padded_dim),\n padded_dim, ex_bits, &(*quantized_data)[bin_data_offset],\n &(*quantized_data)[ex_data_offset], metric_type, config);\n\n return 0;\n}\n\nint RabitqReformer::dump(const IndexDumper::Pointer &dumper) {\n if (!dumper) {\n LOG_ERROR(\"Null dumper\");\n return IndexError_InvalidArgument;\n }\n\n if (!impl_->loaded || impl_->rotated_centroids.empty() ||\n impl_->centroids.empty()) {\n LOG_ERROR(\"No centroids to dump\");\n return IndexError_NoReady;\n }\n\n size_t dumped_size = 0;\n int ret = dump_rabitq_centroids(\n dumper, impl_->dimension, impl_->padded_dim, impl_->ex_bits,\n impl_->num_clusters, impl_->rotator_type, impl_->rotated_centroids,\n impl_->centroids, impl_->rotator, &dumped_size);\n if (ret != 0) {\n return ret;\n }\n\n LOG_INFO(\"RabitqReformer dump completed: %zu bytes\", dumped_size);\n return 0;\n}\n\nint RabitqReformer::dump(const IndexStorage::Pointer &storage) {\n if (!storage) {\n LOG_ERROR(\"Null storage\");\n return IndexError_InvalidArgument;\n }\n\n if (!impl_->loaded || impl_->rotated_centroids.empty() ||\n impl_->centroids.empty()) {\n LOG_ERROR(\"No centroids to dump\");\n return IndexError_NoReady;\n }\n\n auto align_size = [](size_t size) -> size_t {\n return (size + 0x1F) & (~0x1F);\n };\n\n size_t header_size = sizeof(RabitqConverterHeader);\n size_t rotated_centroids_size =\n impl_->rotated_centroids.size() * sizeof(float);\n size_t centroids_size = impl_->centroids.size() * sizeof(float);\n size_t rotator_size = impl_->rotator->dump_bytes();\n size_t data_size =\n header_size + rotated_centroids_size + centroids_size + rotator_size;\n size_t total_size = align_size(data_size);\n\n int ret = storage->append(RABITQ_CONVERTER_SEG_ID, total_size);\n if (ret != 0) {\n LOG_ERROR(\"Failed to append segment %s, ret=%d\",\n RABITQ_CONVERTER_SEG_ID.c_str(), ret);\n return ret;\n }\n\n auto segment = storage->get(RABITQ_CONVERTER_SEG_ID);\n if (!segment) {\n LOG_ERROR(\"Failed to get segment %s\", RABITQ_CONVERTER_SEG_ID.c_str());\n return IndexError_ReadData;\n }\n\n size_t offset = 0;\n\n RabitqConverterHeader header;\n header.dim = static_cast(impl_->dimension);\n header.padded_dim = static_cast(impl_->padded_dim);\n header.num_clusters = static_cast(impl_->num_clusters);\n header.ex_bits = static_cast(impl_->ex_bits);\n header.rotator_type = static_cast(impl_->rotator_type);\n header.rotator_size = static_cast(rotator_size);\n size_t written = segment->write(offset, &header, header_size);\n if (written != header_size) {\n LOG_ERROR(\"Failed to write header: written=%zu, expected=%zu\", written,\n header_size);\n return IndexError_WriteData;\n }\n offset += header_size;\n\n written = segment->write(offset, impl_->rotated_centroids.data(),\n rotated_centroids_size);\n if (written != rotated_centroids_size) {\n LOG_ERROR(\"Failed to write rotated centroids: written=%zu, expected=%zu\",\n written, rotated_centroids_size);\n return IndexError_WriteData;\n }\n offset += rotated_centroids_size;\n\n written = segment->write(offset, impl_->centroids.data(), centroids_size);\n if (written != centroids_size) {\n LOG_ERROR(\"Failed to write centroids: written=%zu, expected=%zu\", written,\n centroids_size);\n return IndexError_WriteData;\n }\n offset += centroids_size;\n\n std::vector buffer(rotator_size);\n impl_->rotator->save(buffer.data());\n written = segment->write(offset, buffer.data(), rotator_size);\n if (written != rotator_size) {\n LOG_ERROR(\"Failed to write rotator data: written=%zu, expected=%zu\",\n written, rotator_size);\n return IndexError_WriteData;\n }\n\n LOG_INFO(\"RabitqReformer dump to storage completed: %zu bytes\", data_size);\n return 0;\n}\n\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/rabitq_reformer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n//\n#pragma once\n\n#include \n#include \"zvec/core/framework/index_dumper.h\"\n#include \"zvec/core/framework/index_reformer.h\"\n#include \"zvec/core/framework/index_storage.h\"\n#include \"rabitq_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswRabitqQueryEntity;\n\n/*! RaBitQ Reformer\n * Loads centroids and performs query transformation and vector quantization.\n *\n * All rabitqlib types are hidden behind a pimpl to avoid leaking rabitqlib\n * headers to consumers of this class.\n */\nclass RabitqReformer : public IndexReformer {\n public:\n typedef std::shared_ptr Pointer;\n\n RabitqReformer();\n ~RabitqReformer() override;\n\n // Non-copyable\n RabitqReformer(const RabitqReformer &) = delete;\n RabitqReformer &operator=(const RabitqReformer &) = delete;\n\n int init(const ailego::Params ¶ms) override;\n int cleanup(void) override;\n int load(IndexStorage::Pointer storage) override;\n int unload(void) override;\n\n // transform() is not implemented for RabitqReformer; use transform_to_entity.\n int transform(const void *query, const IndexQueryMeta &qmeta,\n std::string *out, IndexQueryMeta *ometa) const override;\n\n int convert(const void *record, const IndexQueryMeta &rmeta, std::string *out,\n IndexQueryMeta *ometa) const override;\n\n int dump(const IndexDumper::Pointer &dumper);\n int dump(const IndexStorage::Pointer &storage);\n\n int transform_to_entity(const void *query,\n HnswRabitqQueryEntity *entity) const;\n\n size_t num_clusters() const;\n RabitqMetricType rabitq_metric_type() const;\n\n private:\n struct Impl;\n std::unique_ptr impl_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/rabitq_utils.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"rabitq_utils.h\"\n#include \n#include \n#include \"zvec/core/framework/index_error.h\"\n#include \"zvec/core/framework/index_logger.h\"\n\nnamespace zvec {\nnamespace core {\n\nint dump_rabitq_centroids(\n const IndexDumper::Pointer &dumper, size_t dimension, size_t padded_dim,\n size_t ex_bits, size_t num_clusters, rabitqlib::RotatorType rotator_type,\n const std::vector &rotated_centroids,\n const std::vector ¢roids,\n const std::unique_ptr> &rotator,\n size_t *out_dumped_size) {\n auto align_size = [](size_t size) -> size_t {\n return (size + 0x1F) & (~0x1F);\n };\n\n uint32_t crc = 0;\n size_t dumped_size = 0;\n\n // Write header\n RabitqConverterHeader header;\n header.dim = static_cast(dimension);\n header.padded_dim = static_cast(padded_dim);\n header.num_clusters = static_cast(num_clusters);\n header.ex_bits = static_cast(ex_bits);\n header.rotator_type = static_cast(rotator_type);\n header.rotator_size = static_cast(rotator->dump_bytes());\n size_t size = dumper->write(&header, sizeof(header));\n if (size != sizeof(header)) {\n LOG_ERROR(\"Failed to write header: written=%zu, expected=%zu\", size,\n sizeof(header));\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(&header, sizeof(header), crc);\n dumped_size += size;\n\n // Write rotated centroids\n size = dumper->write(rotated_centroids.data(),\n rotated_centroids.size() * sizeof(float));\n if (size != rotated_centroids.size() * sizeof(float)) {\n LOG_ERROR(\"Failed to write rotated centroids: written=%zu, expected=%zu\",\n size, rotated_centroids.size() * sizeof(float));\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(rotated_centroids.data(),\n rotated_centroids.size() * sizeof(float), crc);\n dumped_size += size;\n\n // Write original centroids\n size = dumper->write(centroids.data(), centroids.size() * sizeof(float));\n if (size != centroids.size() * sizeof(float)) {\n LOG_ERROR(\"Failed to write centroids: written=%zu, expected=%zu\", size,\n centroids.size() * sizeof(float));\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(centroids.data(), centroids.size() * sizeof(float),\n crc);\n dumped_size += size;\n\n // Write rotator data\n std::vector buffer(rotator->dump_bytes());\n rotator->save(buffer.data());\n size = dumper->write(buffer.data(), buffer.size());\n if (size != buffer.size()) {\n LOG_ERROR(\"Failed to write rotator data: written=%zu, expected=%zu\", size,\n buffer.size());\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(buffer.data(), buffer.size(), crc);\n dumped_size += size;\n\n // Write padding\n size_t padding_size = align_size(dumped_size) - dumped_size;\n if (padding_size > 0) {\n std::string padding(padding_size, '\\0');\n if (dumper->write(padding.data(), padding_size) != padding_size) {\n LOG_ERROR(\"Append padding failed, size %lu\", padding_size);\n return IndexError_WriteData;\n }\n }\n\n int ret =\n dumper->append(RABITQ_CONVERTER_SEG_ID, dumped_size, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s meta failed, ret=%d\",\n RABITQ_CONVERTER_SEG_ID.c_str(), ret);\n return ret;\n }\n\n if (out_dumped_size) {\n *out_dumped_size = dumped_size;\n }\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_rabitq/rabitq_utils.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \"zvec/core/framework/index_dumper.h\"\n\nnamespace zvec {\nnamespace core {\n\ninline const std::string RABITQ_CONVERTER_SEG_ID{\"rabitq.converter\"};\n\nstruct RabitqConverterHeader {\n uint32_t num_clusters;\n uint32_t dim;\n uint32_t padded_dim;\n uint32_t rotator_size;\n uint8_t ex_bits;\n uint8_t rotator_type;\n uint8_t padding[2];\n uint32_t reserve[3];\n\n RabitqConverterHeader() {\n memset(this, 0, sizeof(RabitqConverterHeader));\n }\n};\nstatic_assert(sizeof(RabitqConverterHeader) % 32 == 0,\n \"RabitqConverterHeader must be aligned with 32 bytes\");\n\n// Common dump implementation for RabitqConverter and RabitqReformer\nint dump_rabitq_centroids(\n const IndexDumper::Pointer &dumper, size_t dimension, size_t padded_dim,\n size_t ex_bits, size_t num_clusters, rabitqlib::RotatorType rotator_type,\n const std::vector &rotated_centroids,\n const std::vector ¢roids,\n const std::unique_ptr> &rotator,\n size_t *out_dumped_size = nullptr);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\ncc_library(\n NAME core_knn_hnsw_sparse \n STATIC SHARED STRICT ALWAYS_LINK\n SRCS *.cc\n LIBS core_framework sparsehash\n INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.cc", "content": "\n// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_sparse_algorithm.h\"\n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nHnswSparseAlgorithm::HnswSparseAlgorithm(HnswSparseEntity &entity)\n : entity_(entity),\n mt_(std::chrono::system_clock::now().time_since_epoch().count()),\n lock_pool_(kLockCnt) {}\n\nint HnswSparseAlgorithm::cleanup() {\n return 0;\n}\n\nint HnswSparseAlgorithm::add_node(node_id_t id, level_t level,\n HnswSparseContext *ctx) {\n spin_lock_.lock();\n\n // std::cout << \"id: \" << id << \", level: \" << level << std::endl;\n\n auto cur_max_level = entity_.cur_max_level();\n auto entry_point = entity_.entry_point();\n if (ailego_unlikely(entry_point == kInvalidNodeId)) {\n entity_.update_ep_and_level(id, level);\n spin_lock_.unlock();\n return 0;\n }\n spin_lock_.unlock();\n\n if (ailego_unlikely(level > cur_max_level)) {\n mutex_.lock();\n // re-check max level\n cur_max_level = entity_.cur_max_level();\n entry_point = entity_.entry_point();\n if (level <= cur_max_level) {\n mutex_.unlock();\n }\n }\n\n level_t cur_level = cur_max_level;\n dist_t dist = ctx->dist_calculator()(entry_point);\n for (; cur_level > level; --cur_level) {\n select_entry_point(cur_level, &entry_point, &dist, ctx);\n }\n\n for (; cur_level >= 0; --cur_level) {\n search_neighbors(cur_level, &entry_point, &dist, ctx->level_topk(cur_level),\n ctx);\n }\n\n // add neighbors from down level to top level, to avoid upper level visible\n // to knn_search but the under layer level not ready\n for (cur_level = 0; cur_level <= level; ++cur_level) {\n add_neighbors(id, cur_level, ctx->level_topk(cur_level), ctx);\n ctx->level_topk(cur_level).clear();\n }\n\n if (ailego_unlikely(level > cur_max_level)) {\n spin_lock_.lock();\n entity_.update_ep_and_level(id, level);\n spin_lock_.unlock();\n mutex_.unlock();\n }\n\n return 0;\n}\n\nint HnswSparseAlgorithm::search(HnswSparseContext *ctx) const {\n spin_lock_.lock();\n auto maxLevel = entity_.cur_max_level();\n auto entry_point = entity_.entry_point();\n spin_lock_.unlock();\n\n if (ailego_unlikely(entry_point == kInvalidNodeId)) {\n return 0;\n }\n\n dist_t dist = ctx->dist_calculator().dist(entry_point);\n for (level_t cur_level = maxLevel; cur_level >= 1; --cur_level) {\n select_entry_point(cur_level, &entry_point, &dist, ctx);\n }\n\n auto &topk_heap = ctx->topk_heap();\n topk_heap.clear();\n search_neighbors(0, &entry_point, &dist, topk_heap, ctx);\n\n if (ctx->group_by_search()) {\n expand_neighbors_by_group(topk_heap, ctx);\n }\n\n return 0;\n}\n\n//! select_entry_point on hnsw level, ef = 1\nvoid HnswSparseAlgorithm::select_entry_point(level_t level,\n node_id_t *entry_point,\n dist_t *dist,\n HnswSparseContext *ctx) const {\n auto &entity = ctx->get_entity();\n HnswSparseDistCalculator &dc = ctx->dist_calculator();\n while (true) {\n const Neighbors neighbors = entity.get_neighbors(level, *entry_point);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n uint32_t size = neighbors.size();\n if (size == 0) {\n break;\n }\n\n std::vector neighbor_block_vecs;\n int ret = entity.get_vector_metas(&neighbors[0], size, neighbor_block_vecs);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_vector())++;\n }\n if (ailego_unlikely(ret != 0)) {\n break;\n }\n bool find_closer = false;\n for (uint32_t i = 0; i < size; ++i) {\n dist_t cur_dist = dc.dist(neighbor_block_vecs[i].data());\n if (cur_dist < *dist) {\n *entry_point = neighbors[i];\n *dist = cur_dist;\n find_closer = true;\n }\n }\n\n if (!find_closer) {\n break;\n }\n }\n\n return;\n}\n\nvoid HnswSparseAlgorithm::add_neighbors(node_id_t id, level_t level,\n TopkHeap &topk_heap,\n HnswSparseContext *ctx) {\n if (ailego_unlikely(topk_heap.size() == 0)) {\n return;\n }\n\n HnswSparseDistCalculator &dc = ctx->dist_calculator();\n\n update_neighbors(dc, id, level, topk_heap);\n\n // reverse update neighbors\n for (size_t i = 0; i < topk_heap.size(); ++i) {\n reverse_update_neighbors(dc, topk_heap[i].first, level, id,\n topk_heap[i].second, ctx->update_heap());\n }\n\n return;\n}\n\nvoid HnswSparseAlgorithm::search_neighbors(level_t level,\n node_id_t *entry_point, dist_t *dist,\n TopkHeap &topk,\n HnswSparseContext *ctx) const {\n const auto &entity = ctx->get_entity();\n HnswSparseDistCalculator &dc = ctx->dist_calculator();\n VisitFilter &visit = ctx->visit_filter();\n CandidateHeap &candidates = ctx->candidates();\n std::function filter = [](node_id_t) { return false; };\n if (ctx->filter().is_valid()) {\n filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); };\n }\n\n candidates.clear();\n visit.clear();\n visit.set_visited(*entry_point);\n if (!filter(*entry_point)) {\n topk.emplace(*entry_point, *dist);\n }\n\n candidates.emplace(*entry_point, *dist);\n while (!candidates.empty() && !ctx->reach_scan_limit()) {\n auto top = candidates.begin();\n node_id_t main_node = top->first;\n dist_t main_dist = top->second;\n\n if (topk.full() && main_dist > topk[0].second) {\n break;\n }\n\n candidates.pop();\n const Neighbors neighbors = entity.get_neighbors(level, main_node);\n ailego_prefetch(neighbors.data);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n\n std::vector neighbor_ids(neighbors.size());\n uint32_t size = 0;\n for (uint32_t i = 0; i < neighbors.size(); ++i) {\n node_id_t node = neighbors[i];\n if (visit.visited(node)) {\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_visit_dup_cnt())++;\n }\n continue;\n }\n visit.set_visited(node);\n neighbor_ids[size++] = node;\n }\n if (size == 0) {\n continue;\n }\n\n std::vector neighbor_block_vecs;\n int ret =\n entity.get_vector_metas(neighbor_ids.data(), size, neighbor_block_vecs);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_vector())++;\n }\n if (ailego_unlikely(ret != 0)) {\n break;\n }\n\n static constexpr node_id_t PREFETCH_STEP = 2;\n static constexpr node_id_t SPARSE_PREFETCH_STEP = 1;\n\n for (uint32_t i = 0; i < std::min(PREFETCH_STEP, size); ++i) {\n ailego_prefetch(neighbor_block_vecs[i].data());\n }\n for (uint32_t i = 0; i < size; ++i) {\n node_id_t node = neighbor_ids[i];\n node_id_t prefetch_id = i + PREFETCH_STEP;\n if (prefetch_id < size) {\n ailego_prefetch(neighbor_block_vecs[prefetch_id].data());\n }\n\n node_id_t sparse_prefetch_id = i + SPARSE_PREFETCH_STEP;\n if (sparse_prefetch_id < size) {\n IndexStorage::MemoryBlock sparse_block;\n int sparse_length = 0;\n entity.get_sparse_data_from_vector(\n neighbor_block_vecs[sparse_prefetch_id].data(), sparse_block,\n sparse_length);\n auto sparse_data = std::make_pair(sparse_block.data(), sparse_length);\n if (sparse_data.first != nullptr) {\n ailego_prefetch(sparse_data.first);\n }\n }\n\n dist_t cur_dist = dc.dist(neighbor_block_vecs[i].data());\n if ((!topk.full()) || cur_dist < topk[0].second) {\n candidates.emplace(node, cur_dist);\n // update entry_point for next level scan\n if (cur_dist < *dist) {\n *entry_point = node;\n *dist = cur_dist;\n }\n if (!filter(node)) {\n topk.emplace(node, cur_dist);\n }\n } // end if\n } // end for\n } // while\n\n return;\n}\n\nvoid HnswSparseAlgorithm::expand_neighbors_by_group(\n TopkHeap &topk, HnswSparseContext *ctx) const {\n if (!ctx->group_by().is_valid()) {\n return;\n }\n\n const auto &entity = ctx->get_entity();\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity.get_key(id));\n };\n\n // devide into groups\n std::map &group_topk_heaps = ctx->group_topk_heaps();\n for (uint32_t i = 0; i < topk.size(); ++i) {\n node_id_t id = topk[i].first;\n auto score = topk[i].second;\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = group_topk_heaps[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, score);\n }\n\n // stage 2, expand to reach group num as possible\n if (group_topk_heaps.size() < ctx->group_num()) {\n VisitFilter &visit = ctx->visit_filter();\n CandidateHeap &candidates = ctx->candidates();\n HnswSparseDistCalculator &dc = ctx->dist_calculator();\n\n std::function filter = [](node_id_t) { return false; };\n if (ctx->filter().is_valid()) {\n filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); };\n }\n\n // refill to get enough groups\n candidates.clear();\n visit.clear();\n for (uint32_t i = 0; i < topk.size(); ++i) {\n node_id_t id = topk[i].first;\n float score = topk[i].second;\n\n visit.set_visited(id);\n candidates.emplace_back(id, score);\n }\n\n // do expand\n while (!candidates.empty() && !ctx->reach_scan_limit()) {\n auto top = candidates.begin();\n node_id_t main_node = top->first;\n\n candidates.pop();\n const Neighbors neighbors = entity.get_neighbors(0, main_node);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_neighbors())++;\n }\n\n std::vector neighbor_ids(neighbors.size());\n uint32_t size = 0;\n for (uint32_t i = 0; i < neighbors.size(); ++i) {\n node_id_t node = neighbors[i];\n if (visit.visited(node)) {\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_visit_dup_cnt())++;\n }\n continue;\n }\n visit.set_visited(node);\n neighbor_ids[size++] = node;\n }\n if (size == 0) {\n continue;\n }\n\n std::vector neighbor_block_vecs;\n int ret = entity.get_vector_metas(neighbor_ids.data(), size,\n neighbor_block_vecs);\n if (ailego_unlikely(ctx->debugging())) {\n (*ctx->mutable_stats_get_vector())++;\n }\n if (ailego_unlikely(ret != 0)) {\n break;\n }\n\n static constexpr node_id_t PREFETCH_STEP = 2;\n for (uint32_t i = 0; i < size; ++i) {\n node_id_t node = neighbor_ids[i];\n node_id_t prefetch_id = i + PREFETCH_STEP;\n if (prefetch_id < size) {\n ailego_prefetch(neighbor_block_vecs[prefetch_id].data());\n }\n dist_t cur_dist = dc.dist(neighbor_block_vecs[i].data());\n\n if (!filter(node)) {\n std::string group_id = group_by(node);\n\n auto &topk_heap = group_topk_heaps[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(node, cur_dist);\n\n if (group_topk_heaps.size() >= ctx->group_num()) {\n break;\n }\n }\n\n candidates.emplace(node, cur_dist);\n } // end for\n } // end while\n } // end if\n}\n\nvoid HnswSparseAlgorithm::update_neighbors(HnswSparseDistCalculator &dc,\n node_id_t id, level_t level,\n TopkHeap &topk_heap) {\n topk_heap.sort();\n\n uint32_t max_neighbor_cnt = entity_.neighbor_cnt(level);\n if (topk_heap.size() <= static_cast(entity_.prune_cnt())) {\n if (topk_heap.size() <= static_cast(max_neighbor_cnt)) {\n entity_.update_neighbors(level, id, topk_heap);\n return;\n }\n }\n\n uint32_t cur_size = 0;\n for (size_t i = 0; i < topk_heap.size(); ++i) {\n node_id_t cur_node = topk_heap[i].first;\n dist_t cur_node_dist = topk_heap[i].second;\n bool good = true;\n for (uint32_t j = 0; j < cur_size; ++j) {\n dist_t tmp_dist = dc.dist(cur_node, topk_heap[j].first);\n if (tmp_dist <= cur_node_dist) {\n good = false;\n break;\n }\n }\n\n if (good) {\n topk_heap[cur_size].first = cur_node;\n topk_heap[cur_size].second = cur_node_dist;\n cur_size++;\n if (cur_size >= max_neighbor_cnt) {\n break;\n }\n }\n }\n\n // when after-prune neighbor count is too seldom,\n // we use this strategy to make-up enough edges\n // not only just make-up out-degrees\n // we also make-up enough in-degrees\n uint32_t min_neighbors = entity_.min_neighbor_cnt();\n for (size_t k = cur_size; cur_size < min_neighbors && k < topk_heap.size();\n ++k) {\n bool exist = false;\n for (size_t j = 0; j < cur_size; ++j) {\n if (topk_heap[j].first == topk_heap[k].first) {\n exist = true;\n break;\n }\n }\n if (!exist) {\n topk_heap[cur_size].first = topk_heap[k].first;\n topk_heap[cur_size].second = topk_heap[k].second;\n cur_size++;\n }\n }\n\n topk_heap.resize(cur_size);\n entity_.update_neighbors(level, id, topk_heap);\n\n return;\n}\n\nvoid HnswSparseAlgorithm::reverse_update_neighbors(HnswSparseDistCalculator &dc,\n node_id_t id, level_t level,\n node_id_t link_id,\n dist_t dist,\n TopkHeap &update_heap) {\n const size_t max_neighbor_cnt = entity_.neighbor_cnt(level);\n\n uint32_t lock_idx = id & kLockMask;\n lock_pool_[lock_idx].lock();\n const Neighbors neighbors = entity_.get_neighbors(level, id);\n size_t size = neighbors.size();\n ailego_assert_with(size <= max_neighbor_cnt, \"invalid neighbor size\");\n if (size < max_neighbor_cnt) {\n entity_.add_neighbor(level, id, size, link_id);\n lock_pool_[lock_idx].unlock();\n return;\n }\n\n update_heap.emplace(link_id, dist);\n\n for (size_t i = 0; i < size; ++i) {\n node_id_t node = neighbors[i];\n dist_t cur_dist = dc.dist(id, node);\n update_heap.emplace(node, cur_dist);\n }\n\n //! TODO: optimize prune\n //! prune edges\n update_heap.sort();\n size_t cur_size = 0;\n for (size_t i = 0; i < update_heap.size(); ++i) {\n node_id_t cur_node = update_heap[i].first;\n dist_t cur_node_dist = update_heap[i].second;\n bool good = true;\n for (size_t j = 0; j < cur_size; ++j) {\n dist_t tmp_dist = dc.dist(cur_node, update_heap[j].first);\n if (tmp_dist <= cur_node_dist) {\n good = false;\n break;\n }\n }\n\n if (good) {\n update_heap[cur_size].first = cur_node;\n update_heap[cur_size].second = cur_node_dist;\n cur_size++;\n if (cur_size >= max_neighbor_cnt) {\n break;\n }\n }\n }\n\n update_heap.resize(cur_size);\n entity_.update_neighbors(level, id, update_heap);\n\n lock_pool_[lock_idx].unlock();\n\n update_heap.clear();\n\n return;\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_algorithm.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"hnsw_sparse_context.h\"\n#include \"hnsw_sparse_dist_calculator.h\"\n#include \"hnsw_sparse_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! hnsw graph algorithm implement\nclass HnswSparseAlgorithm {\n public:\n typedef std::unique_ptr UPointer;\n\n public:\n //! Constructor\n explicit HnswSparseAlgorithm(HnswSparseEntity &entity);\n\n //! Cleanup HnswSparseAlgorithm\n int cleanup();\n\n //! Add a node to hnsw graph\n //! @id: the node unique id\n //! @level: a node will be add to graph in each level [0, level]\n //! return 0 on success, or errCode in failure\n int add_node(node_id_t id, level_t level, HnswSparseContext *ctx);\n\n //! do knn search in graph\n //! return 0 on success, or errCode in failure. results saved in ctx\n int search(HnswSparseContext *ctx) const;\n\n //! Initiate HnswAlgorithm\n int init() {\n level_probas_.clear();\n double level_mult =\n 1 / std::log(static_cast(entity_.scaling_factor()));\n for (int level = 0;; level++) {\n // refers faiss get_random_level alg\n double proba =\n std::exp(-level / level_mult) * (1 - std::exp(-1 / level_mult));\n if (proba < 1e-9) {\n break;\n }\n level_probas_.push_back(proba);\n }\n\n return 0;\n }\n\n //! Generate a random level\n //! return graph level\n uint32_t get_random_level() const {\n // gen rand float (0, 1)\n double f = mt_() / static_cast(mt_.max());\n for (size_t level = 0; level < level_probas_.size(); level++) {\n if (f < level_probas_[level]) {\n return level;\n }\n f -= level_probas_[level];\n }\n return level_probas_.size() - 1;\n }\n\n private:\n //! Select in upper layer to get entry point for next layer search\n void select_entry_point(level_t level, node_id_t *entry_point, dist_t *dist,\n HnswSparseContext *ctx) const;\n\n //! update node id neighbors from topkHeap, and reverse link is also updated\n void add_neighbors(node_id_t id, level_t level, TopkHeap &topk_heap,\n HnswSparseContext *ctx);\n\n //! Given a node id and level, search the nearest neighbors in graph\n //! Note: the nearest neighbors result keeps in topk, and entry_point and\n //! dist will be updated to current level nearest node id and distance\n void search_neighbors(level_t level, node_id_t *entry_point, dist_t *dist,\n TopkHeap &topk, HnswSparseContext *ctx) const;\n\n //! Update the node's neighbors\n void update_neighbors(HnswSparseDistCalculator &dc, node_id_t id,\n level_t level, TopkHeap &topk_heap);\n\n //! Checking linkId could be id's new neighbor, and add as neighbor if true\n //! @dc distance calculator\n //! @updateHeap temporary heap in updating neighbors\n void reverse_update_neighbors(HnswSparseDistCalculator &dc, node_id_t id,\n level_t level, node_id_t link_id, dist_t dist,\n TopkHeap &update_heap);\n\n //! expand neighbors until group nums are reached\n void expand_neighbors_by_group(TopkHeap &topk, HnswSparseContext *ctx) const;\n\n private:\n HnswSparseAlgorithm(const HnswSparseAlgorithm &) = delete;\n HnswSparseAlgorithm &operator=(const HnswSparseAlgorithm &) = delete;\n\n private:\n static constexpr uint32_t kLockCnt{1U << 8};\n static constexpr uint32_t kLockMask{kLockCnt - 1U};\n\n HnswSparseEntity &entity_;\n mutable std::mt19937 mt_{};\n std::vector level_probas_{};\n\n mutable ailego::SpinMutex spin_lock_{}; // global spin lock\n std::mutex mutex_{}; // global mutex\n // TODO: spin lock?\n std::vector lock_pool_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_sparse_builder.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"hnsw_sparse_algorithm.h\"\n#include \"hnsw_sparse_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswSparseBuilder::HnswSparseBuilder() {}\n\nint HnswSparseBuilder::init(const IndexMeta &meta,\n const ailego::Params ¶ms) {\n LOG_INFO(\"Begin HnswSparseBuilder::init\");\n\n meta_ = meta;\n auto params_copy = params;\n meta_.set_builder(\"HnswSparseBuilder\", HnswSparseEntity::kRevision,\n std::move(params_copy));\n\n size_t memory_quota = 0UL;\n params.get(PARAM_HNSW_SPARSE_BUILDER_MEMORY_QUOTA, &memory_quota);\n params.get(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, &thread_cnt_);\n params.get(PARAM_HNSW_SPARSE_BUILDER_EFCONSTRUCTION, &ef_construction_);\n params.get(PARAM_HNSW_SPARSE_BUILDER_CHECK_INTERVAL_SECS,\n &check_interval_secs_);\n\n params.get(PARAM_HNSW_SPARSE_BUILDER_MAX_NEIGHBOR_COUNT,\n &upper_max_neighbor_cnt_);\n float multiplier = HnswSparseEntity::kDefaultL0MaxNeighborCntMultiplier;\n params.get(PARAM_HNSW_SPARSE_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER,\n &multiplier);\n l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_;\n scaling_factor_ = upper_max_neighbor_cnt_;\n params.get(PARAM_HNSW_SPARSE_BUILDER_SCALING_FACTOR, &scaling_factor_);\n\n multiplier = HnswSparseEntity::kDefaultNeighborPruneMultiplier;\n params.get(PARAM_HNSW_SPARSE_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier);\n size_t prune_cnt = multiplier * upper_max_neighbor_cnt_;\n\n if (ef_construction_ == 0) {\n ef_construction_ = HnswSparseEntity::kDefaultEfConstruction;\n }\n if (upper_max_neighbor_cnt_ == 0) {\n upper_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt;\n }\n if (upper_max_neighbor_cnt_ > kMaxNeighborCnt) {\n LOG_ERROR(\"[%s] must be in range (0,%d]\",\n PARAM_HNSW_SPARSE_BUILDER_MAX_NEIGHBOR_COUNT.c_str(),\n kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) {\n LOG_ERROR(\"[%s]-[%d] must be <= [%s]-[%d]\",\n PARAM_HNSW_SPARSE_BUILDER_MIN_NEIGHBOR_COUNT.c_str(),\n min_neighbor_cnt_,\n PARAM_HNSW_SPARSE_BUILDER_MAX_NEIGHBOR_COUNT.c_str(),\n upper_max_neighbor_cnt_);\n return IndexError_InvalidArgument;\n }\n if (l0_max_neighbor_cnt_ == 0) {\n l0_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt;\n }\n if (l0_max_neighbor_cnt_ > HnswSparseEntity::kMaxNeighborCnt) {\n LOG_ERROR(\"L0MaxNeighborCnt must be in range (0,%d)\",\n HnswSparseEntity::kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (scaling_factor_ == 0U) {\n scaling_factor_ = HnswSparseEntity::kDefaultScalingFactor;\n }\n if (scaling_factor_ < 5 || scaling_factor_ > 1000) {\n LOG_ERROR(\"[%s] must be in range [5,1000]\",\n PARAM_HNSW_SPARSE_BUILDER_SCALING_FACTOR.c_str());\n return IndexError_InvalidArgument;\n }\n if (thread_cnt_ == 0) {\n thread_cnt_ = std::thread::hardware_concurrency();\n }\n if (thread_cnt_ > std::thread::hardware_concurrency()) {\n LOG_WARN(\"[%s] greater than cpu cores %u\",\n PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT.c_str(),\n std::thread::hardware_concurrency());\n }\n if (prune_cnt == 0UL) {\n prune_cnt = upper_max_neighbor_cnt_;\n }\n\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"CreateMeasure failed, name: %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n int ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"IndexMeasure init failed, ret=%d\", ret);\n return ret;\n }\n\n entity_.set_ef_construction(ef_construction_);\n entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_);\n entity_.set_min_neighbor_cnt(min_neighbor_cnt_);\n entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_);\n entity_.set_scaling_factor(scaling_factor_);\n entity_.set_memory_quota(memory_quota);\n entity_.set_prune_cnt(prune_cnt);\n\n entity_.set_sparse_meta_size(HnswSparseEntity::kSparseMetaSize);\n entity_.set_sparse_unit_size(meta.unit_size());\n\n ret = entity_.init();\n if (ret != 0) {\n return ret;\n }\n\n alg_ = HnswSparseAlgorithm::UPointer(new HnswSparseAlgorithm(entity_));\n\n ret = alg_->init();\n if (ret != 0) {\n return ret;\n }\n\n state_ = BUILD_STATE_INITED;\n LOG_INFO(\n \"End HnswSparseBuilder::init, params: efConstruction=%u \"\n \"l0NeighborCnt=%u upperNeighborCnt=%u scalingFactor=%u \"\n \"memoryQuota=%zu neighborPruneCnt=%zu measureName=%s \",\n ef_construction_, l0_max_neighbor_cnt_, upper_max_neighbor_cnt_,\n scaling_factor_, memory_quota, prune_cnt, meta_.metric_name().c_str());\n\n return 0;\n}\n\nint HnswSparseBuilder::cleanup(void) {\n LOG_INFO(\"Begin HnswSparseBuilder::cleanup\");\n\n l0_max_neighbor_cnt_ = HnswSparseEntity::kDefaultL0MaxNeighborCnt;\n min_neighbor_cnt_ = 0;\n upper_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt;\n ef_construction_ = HnswSparseEntity::kDefaultEfConstruction;\n scaling_factor_ = HnswSparseEntity::kDefaultScalingFactor;\n check_interval_secs_ = kDefaultLogIntervalSecs;\n errcode_ = 0;\n error_ = false;\n entity_.cleanup();\n alg_->cleanup();\n meta_.clear();\n metric_.reset();\n stats_.clear_attributes();\n stats_.set_trained_count(0UL);\n stats_.set_built_count(0UL);\n stats_.set_dumped_count(0UL);\n stats_.set_discarded_count(0UL);\n stats_.set_trained_costtime(0UL);\n stats_.set_built_costtime(0UL);\n stats_.set_dumped_costtime(0UL);\n state_ = BUILD_STATE_INIT;\n\n LOG_INFO(\"End HnswSparseBuilder::cleanup\");\n\n return 0;\n}\n\nint HnswSparseBuilder::train(IndexThreads::Pointer,\n IndexSparseHolder::Pointer /*holder*/) {\n if (state_ != BUILD_STATE_INITED) {\n LOG_ERROR(\"Init the builder before HnswSparseBuilder::train\");\n return IndexError_NoReady;\n }\n\n stats_.set_trained_count(0UL);\n stats_.set_trained_costtime(0UL);\n state_ = BUILD_STATE_TRAINED;\n\n LOG_INFO(\"End HnswSparseBuilder::train\");\n\n return 0;\n}\n\nint HnswSparseBuilder::train(const IndexTrainer::Pointer & /*trainer*/) {\n if (state_ != BUILD_STATE_INITED) {\n LOG_ERROR(\"Init the builder before HnswSparseBuilder::train\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"Begin HnswSparseBuilder::train by trainer\");\n\n stats_.set_trained_count(0UL);\n stats_.set_trained_costtime(0UL);\n state_ = BUILD_STATE_TRAINED;\n\n LOG_INFO(\"End HnswSparseBuilder::train by trainer\");\n\n return 0;\n}\n\nint HnswSparseBuilder::build(IndexThreads::Pointer threads,\n IndexSparseHolder::Pointer holder) {\n if (!holder) {\n LOG_ERROR(\"Input holder is nullptr while building index\");\n return IndexError_InvalidArgument;\n }\n\n if (!holder->is_matched(meta_)) {\n LOG_ERROR(\"Input holder doesn't match index meta while building index\");\n return IndexError_Mismatch;\n }\n if (!threads) {\n threads = std::make_shared(thread_cnt_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n\n auto start_time = ailego::Monotime::MilliSeconds();\n\n LOG_INFO(\"Begin HnswSparseBuilder::build sparse\");\n\n // holder should be hybrid holder\n auto sparse_holder = std::dynamic_pointer_cast(holder);\n\n if (sparse_holder == nullptr) {\n LOG_ERROR(\"HnswSparseBuilder failed to cast holder\");\n return IndexError_Runtime;\n }\n\n if (sparse_holder->count() != static_cast(-1)) {\n LOG_DEBUG(\"HnswSparseBuilder holder documents count %lu\",\n sparse_holder->count());\n\n int ret = entity_.reserve_space(sparse_holder->count(),\n sparse_holder->total_sparse_count());\n if (ret != 0) {\n LOG_ERROR(\"HnswBuilde reserver space failed\");\n return ret;\n }\n }\n auto iter = sparse_holder->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Create iterator for holder failed\");\n return IndexError_Runtime;\n }\n\n int ret;\n error_ = false;\n while (iter->is_valid()) {\n level_t level = alg_->get_random_level();\n node_id_t id;\n\n ret = entity_.add_vector(level, iter->key(), iter->sparse_count(),\n iter->sparse_indices(), iter->sparse_data(), &id);\n\n if (ailego_unlikely(ret != 0) && ret != IndexError_InvalidValue) {\n return ret;\n }\n\n iter->next();\n }\n // Holder is not needed, cleanup it.\n sparse_holder.reset();\n\n LOG_INFO(\"Finished save vector, start build graph...\");\n\n std::atomic finished{0};\n\n ret = build_graph(threads, finished);\n if (ret != 0) {\n LOG_ERROR(\"Failed to build graph\");\n return ret;\n }\n\n stats_.set_built_count(finished.load());\n stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time);\n state_ = BUILD_STATE_BUILT;\n\n LOG_INFO(\"End HnswSparseBuilder::build\");\n return 0;\n}\n\nint HnswSparseBuilder::build_graph(IndexThreads::Pointer threads,\n std::atomic &finished) {\n auto task_group = threads->make_group();\n if (!task_group) {\n LOG_ERROR(\"Failed to create task group\");\n return IndexError_Runtime;\n }\n\n for (size_t i = 0; i < threads->count(); ++i) {\n task_group->submit(ailego::Closure ::New(this, &HnswSparseBuilder::do_build,\n i, threads->count(), &finished));\n }\n\n while (!task_group->is_finished()) {\n std::unique_lock lk(mutex_);\n cond_.wait_until(lk, std::chrono::system_clock::now() +\n std::chrono::seconds(check_interval_secs_));\n if (error_.load(std::memory_order_acquire)) {\n LOG_ERROR(\"Failed to build index while waiting finish\");\n return errcode_;\n }\n LOG_INFO(\"Built cnt %u, finished percent %.3f%%\", finished.load(),\n finished.load() * 100.0f / entity_.doc_cnt());\n }\n if (error_.load(std::memory_order_acquire)) {\n LOG_ERROR(\"Failed to build index while waiting finish\");\n return errcode_;\n }\n task_group->wait_finish();\n\n return 0;\n}\n\nvoid HnswSparseBuilder::do_build(node_id_t idx, size_t step_size,\n std::atomic *finished) {\n AILEGO_DEFER([&]() {\n std::lock_guard latch(mutex_);\n cond_.notify_one();\n });\n\n HnswSparseContext *ctx = new (std::nothrow) HnswSparseContext(\n metric_,\n std::shared_ptr(&entity_, [](HnswSparseEntity *) {}));\n if (ailego_unlikely(ctx == nullptr)) {\n if (!error_.exchange(true)) {\n LOG_ERROR(\"Failed to create context\");\n errcode_ = IndexError_NoMemory;\n }\n return;\n }\n HnswSparseContext::Pointer auto_ptr(ctx);\n ctx->set_max_scan_num(entity_.doc_cnt());\n int ret = ctx->init(HnswSparseContext::kSparseBuilderContext);\n if (ret != 0) {\n if (!error_.exchange(true)) {\n LOG_ERROR(\"Failed to init context\");\n errcode_ = IndexError_Runtime;\n }\n return;\n }\n\n IndexQueryMeta qmeta(meta_.data_type());\n for (node_id_t id = idx; id < entity_.doc_cnt(); id += step_size) {\n const void *vec = entity_.get_vector_meta(id);\n\n auto sparse_data = entity_.get_sparse_data_from_vector(vec);\n\n ctx->reset_query(sparse_data.first);\n\n ret = alg_->add_node(id, entity_.get_level(id), ctx);\n if (ailego_unlikely(ret != 0)) {\n if (!error_.exchange(true)) {\n LOG_ERROR(\"Hnsw graph add node failed\");\n errcode_ = ret;\n }\n return;\n }\n ctx->clear();\n (*finished)++;\n }\n}\n\nint HnswSparseBuilder::dump(const IndexDumper::Pointer &dumper) {\n if (state_ != BUILD_STATE_BUILT) {\n LOG_INFO(\"Build the index before HnswSparseBuilder::dump\");\n return IndexError_NoReady;\n }\n\n LOG_INFO(\"Begin HnswSparseBuilder::dump\");\n\n meta_.set_searcher(\"HnswSparseSearcher\", HnswSparseEntity::kRevision,\n ailego::Params());\n auto start_time = ailego::Monotime::MilliSeconds();\n\n int ret = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return ret;\n }\n\n ret = entity_.dump(dumper);\n if (ret != 0) {\n LOG_ERROR(\"HnswSparseBuilder dump index failed\");\n return ret;\n }\n\n stats_.set_dumped_count(entity_.doc_cnt());\n stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time);\n\n LOG_INFO(\"EndHnswSparseBuilder::dump\");\n return 0;\n}\n\nint HnswSparseBuilder::build(IndexThreads::Pointer threads, size_t count,\n const uint64_t *keys,\n const uint64_t *sparse_indptr,\n const uint32_t *sparse_indices,\n const void *sparse_data) {\n IndexQueryMeta qmeta(meta_.data_type());\n\n return build(threads, qmeta, count, keys, sparse_indptr, sparse_indices,\n sparse_data);\n}\n\nint HnswSparseBuilder::build(IndexThreads::Pointer threads,\n const IndexQueryMeta &qmeta, size_t count,\n const uint64_t *keys,\n const uint64_t *sparse_indptr,\n const uint32_t *sparse_indices,\n const void *sparse_data) {\n if (!threads) {\n threads = std::make_shared(thread_cnt_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n\n auto start_time = ailego::Monotime::MilliSeconds();\n\n LOG_INFO(\"Begin HnswSparseBuilder::build sparse, documents count %lu\", count);\n\n size_t total_sparse_count = sparse_indptr[count];\n\n int ret = entity_.reserve_space(count, total_sparse_count);\n if (ret != 0) {\n LOG_ERROR(\"HnswBuilde reserver space failed\");\n return ret;\n }\n\n if (qmeta.data_type() == meta_.data_type()) {\n for (size_t i = 0; i < count; i++) {\n level_t level = alg_->get_random_level();\n node_id_t id;\n\n uint32_t sparse_count = sparse_indptr[i + 1] - sparse_indptr[i];\n const uint32_t *sparse_indices_temp = sparse_indices + sparse_indptr[i];\n\n const void *sparse_data_temp = static_cast(sparse_data) +\n sparse_indptr[i] * qmeta.unit_size();\n\n ret = entity_.add_vector(level, keys[i], sparse_count,\n sparse_indices_temp, sparse_data_temp, &id);\n if (ailego_unlikely(ret != 0) && ret != IndexError_InvalidValue) {\n return ret;\n }\n }\n } else if (meta_.data_type() == IndexMeta::DataType::DT_FP16 &&\n qmeta.data_type() == IndexMeta::DataType::DT_FP32) {\n // transform from float 32 to float 16\n auto reformer = IndexFactory::CreateReformer(\"HalfFloatSparseReformer\");\n if (!reformer) {\n LOG_ERROR(\"Sparse reformer not existed.\");\n\n return IndexError_NoExist;\n }\n\n meta_.set_converter(\"HalfFloatSparseConverter\", 0, ailego::Params());\n meta_.set_reformer(\"HalfFloatSparseReformer\", 0, ailego::Params());\n\n for (size_t i = 0; i < count; i++) {\n level_t level = alg_->get_random_level();\n node_id_t id;\n\n uint32_t sparse_count = sparse_indptr[i + 1] - sparse_indptr[i];\n const uint32_t *sparse_indices_temp = sparse_indices + sparse_indptr[i];\n\n const void *sparse_data_temp = static_cast(sparse_data) +\n sparse_indptr[i] * qmeta.unit_size();\n\n std::string query_fp16;\n IndexQueryMeta ometa;\n\n reformer->transform(sparse_count, sparse_indices_temp, sparse_data_temp,\n qmeta, &query_fp16, &ometa);\n\n ret = entity_.add_vector(level, keys[i], sparse_count,\n sparse_indices_temp, query_fp16.data(), &id);\n if (ailego_unlikely(ret != 0) && ret != IndexError_InvalidValue) {\n return ret;\n }\n }\n } else {\n LOG_ERROR(\"Format not supported.\");\n\n return IndexError_Unsupported;\n }\n\n LOG_INFO(\"Finished save vector, start build graph...\");\n\n std::atomic finished{0};\n\n ret = build_graph(threads, finished);\n if (ret != 0) {\n LOG_ERROR(\"Failed to build graph\");\n return ret;\n }\n\n stats_.set_built_count(finished.load());\n stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time);\n state_ = BUILD_STATE_BUILT;\n\n LOG_INFO(\"End HnswSparseBuilder::build\");\n return 0;\n}\n\nINDEX_FACTORY_REGISTER_BUILDER(HnswSparseBuilder);\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"hnsw_sparse_algorithm.h\"\n#include \"hnsw_sparse_builder_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSparseBuilder : public IndexBuilder {\n public:\n //! Constructor\n HnswSparseBuilder();\n\n //! Initialize the builder\n int init(const IndexMeta &meta, const ailego::Params ¶ms) override;\n\n //! Cleanup the builder\n int cleanup(void) override;\n\n //! Train the data\n int train(IndexThreads::Pointer, IndexSparseHolder::Pointer holder) override;\n\n //! Train the data\n int train(const IndexTrainer::Pointer &trainer) override;\n\n //! Build the index\n int build(IndexThreads::Pointer threads,\n IndexSparseHolder::Pointer holder) override;\n\n //! Build the index with indptr format\n int build(IndexThreads::Pointer threads, const IndexQueryMeta &qmeta,\n size_t count, const uint64_t *keys, const uint64_t *sparse_indptr,\n const uint32_t *sparse_indices, const void *sparse_data) override;\n\n //! Build the index with indptr format\n int build(IndexThreads::Pointer threads, size_t count, const uint64_t *keys,\n const uint64_t *sparse_indptr, const uint32_t *sparse_indices,\n const void *sparse_data) override;\n\n //! Dump index into storage\n int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n const Stats &stats(void) const override {\n return stats_;\n }\n\n private:\n int build_graph(IndexThreads::Pointer threads,\n std::atomic &finished);\n void do_build(node_id_t idx, size_t step_size,\n std::atomic *finished);\n\n constexpr static uint32_t kDefaultLogIntervalSecs = 15U;\n constexpr static uint32_t kMaxNeighborCnt = 65535;\n\n private:\n enum BUILD_STATE {\n BUILD_STATE_INIT = 0,\n BUILD_STATE_INITED = 1,\n BUILD_STATE_TRAINED = 2,\n BUILD_STATE_BUILT = 3\n };\n\n HnswSparseBuilderEntity entity_{};\n HnswSparseAlgorithm::UPointer alg_; // impl graph algorithm\n uint32_t thread_cnt_{0};\n uint32_t l0_max_neighbor_cnt_{HnswSparseEntity::kDefaultL0MaxNeighborCnt};\n uint32_t min_neighbor_cnt_{0};\n uint32_t upper_max_neighbor_cnt_{\n HnswSparseEntity::kDefaultUpperMaxNeighborCnt};\n uint32_t ef_construction_{HnswSparseEntity::kDefaultEfConstruction};\n uint32_t scaling_factor_{HnswSparseEntity::kDefaultScalingFactor};\n uint32_t check_interval_secs_{kDefaultLogIntervalSecs};\n\n int errcode_{0};\n std::atomic_bool error_{false};\n IndexMeta meta_{};\n IndexMetric::Pointer metric_{};\n std::mutex mutex_{};\n std::condition_variable cond_{};\n Stats stats_{};\n\n BUILD_STATE state_{BUILD_STATE_INIT};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_sparse_builder_entity.h\"\n#include \n#include \"utility/sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswSparseBuilderEntity::HnswSparseBuilderEntity() {\n update_ep_and_level(kInvalidNodeId, 0U);\n}\n\nint HnswSparseBuilderEntity::cleanup() {\n memory_quota_ = 0UL;\n neighbors_size_ = 0U;\n upper_neighbors_size_ = 0U;\n padding_size_ = 0U;\n vectors_buffer_.clear();\n keys_buffer_.clear();\n neighbors_buffer_.clear();\n upper_neighbors_buffer_.clear();\n neighbors_index_.clear();\n\n vectors_buffer_.shrink_to_fit();\n keys_buffer_.shrink_to_fit();\n neighbors_buffer_.shrink_to_fit();\n upper_neighbors_buffer_.shrink_to_fit();\n neighbors_index_.shrink_to_fit();\n\n this->HnswSparseEntity::cleanup();\n\n return 0;\n}\n\nint HnswSparseBuilderEntity::init() {\n size_t size = vector_size();\n\n size += sparse_meta_size();\n\n //! aligned size to 32\n set_node_size(AlignSize(size));\n //! if node size is aligned to 1k, the build performance will downgrade\n if (node_size() % 1024 == 0) {\n set_node_size(AlignSize(node_size() + 1));\n }\n\n padding_size_ = node_size() - size;\n\n neighbors_size_ = neighbors_size();\n upper_neighbors_size_ = upper_neighbors_size();\n\n return 0;\n}\n\nint HnswSparseBuilderEntity::reserve_space(size_t docs,\n size_t total_sparse_count) {\n if (memory_quota_ > 0 && (node_size() * docs + neighbors_size_ * docs +\n sizeof(SparseNeighborIndex) * docs >\n memory_quota_)) {\n return IndexError_NoMemory;\n }\n\n vectors_buffer_.reserve(node_size() * docs);\n keys_buffer_.reserve(sizeof(key_t) * docs);\n neighbors_buffer_.reserve(neighbors_size_ * docs);\n neighbors_index_.reserve(docs);\n\n sparse_data_buffer_.reserve(sizeof(uint32_t) * docs +\n (sizeof(uint32_t)) * total_sparse_count +\n sparse_unit_size() * total_sparse_count);\n\n return 0;\n}\n\nint HnswSparseBuilderEntity::add_vector(level_t level, key_t key,\n const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_vec, node_id_t *id) {\n if (ailego_unlikely(sparse_count > HnswSparseEntity::kSparseMaxDimSize)) {\n LOG_WARN(\n \"Failed to add sparse vector: number of non-zero elements (%u) exceeds \"\n \"maximum allowed (%u), key=%zu\",\n sparse_count, HnswSparseEntity::kSparseMaxDimSize, (size_t)key);\n return IndexError_InvalidValue;\n }\n\n std::string sparse_buffer;\n SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_vec,\n sparse_unit_size(), sparse_buffer);\n\n uint32_t sparse_len = sparse_buffer.size();\n\n if (memory_quota_ > 0 &&\n (vectors_buffer_.capacity() + keys_buffer_.capacity() +\n neighbors_buffer_.capacity() + upper_neighbors_buffer_.capacity() +\n neighbors_index_.capacity() * sizeof(SparseNeighborIndex) +\n sparse_len >\n memory_quota_)) {\n LOG_ERROR(\"Add vector failed, used memory exceed quota, cur_doc=%u\",\n doc_cnt());\n return IndexError_NoMemory;\n }\n\n vectors_buffer_.append(reinterpret_cast(&sparse_data_offset_),\n sizeof(uint64_t));\n vectors_buffer_.append(reinterpret_cast(&sparse_len),\n sizeof(uint32_t));\n vectors_buffer_.append(sizeof(uint32_t),\n '\\0'); // reserve to make it up to meta size\n vectors_buffer_.append(padding_size_, '\\0');\n\n keys_buffer_.append(reinterpret_cast(&key), sizeof(key));\n\n sparse_data_buffer_.append(sparse_buffer.data(), sparse_len);\n sparse_data_offset_ += sparse_len;\n\n // init level 0 neighbors\n neighbors_buffer_.append(neighbors_size_, '\\0');\n\n neighbors_index_.emplace_back(upper_neighbors_buffer_.size(), level);\n\n // init upper layer neighbors\n for (level_t cur_level = 1; cur_level <= level; ++cur_level) {\n upper_neighbors_buffer_.append(upper_neighbors_size_, '\\0');\n }\n\n *id = (*mutable_doc_cnt())++;\n\n return 0;\n}\n\nkey_t HnswSparseBuilderEntity::get_key(node_id_t id) const {\n return *(reinterpret_cast(keys_buffer_.data() +\n id * sizeof(key_t)));\n}\n\nconst void *HnswSparseBuilderEntity::get_vector_meta(node_id_t id) const {\n return vectors_buffer_.data() + id * node_size();\n}\n\nint HnswSparseBuilderEntity::get_vector_meta(\n const node_id_t id, IndexStorage::MemoryBlock &block) const {\n const void *vec = get_vector_meta(id);\n block.reset((void *)vec);\n return 0;\n}\n\nint HnswSparseBuilderEntity::get_vector_metas(const node_id_t *ids,\n uint32_t count,\n const void **vecs) const {\n for (uint32_t i = 0; i < count; ++i) {\n vecs[i] = vectors_buffer_.data() + ids[i] * node_size();\n }\n\n return 0;\n}\n\nint HnswSparseBuilderEntity::get_vector_metas(\n const node_id_t *ids, uint32_t count,\n std::vector &block_vecs) const {\n const void *vecs[count];\n get_vector_metas(ids, count, vecs);\n for (uint32_t i = 0; i < count; ++i) {\n block_vecs.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i]));\n }\n return 0;\n}\n\n//! Get vector feature data by key\nconst void *HnswSparseBuilderEntity::get_sparse_data(uint64_t offset,\n uint32_t /*len*/) const {\n return reinterpret_cast(sparse_data_buffer_.data()) + offset;\n}\n\nint HnswSparseBuilderEntity::get_sparse_data(\n uint64_t offset, uint32_t len, IndexStorage::MemoryBlock &block) const {\n const void *vec = get_sparse_data(offset, len);\n block.reset((void *)vec);\n return 0;\n}\n\n//! Get sparse data from id\nconst void *HnswSparseBuilderEntity::get_sparse_data(node_id_t id) const {\n auto sparse_data = get_sparse_data_from_vector(get_vector_meta(id));\n\n return sparse_data.first;\n}\n\nint HnswSparseBuilderEntity::get_sparse_data(\n const node_id_t id, IndexStorage::MemoryBlock &block) const {\n const void *vec = get_sparse_data(id);\n block.reset((void *)vec);\n return 0;\n}\n\n//! Get sparse data from vector\nstd::pair\nHnswSparseBuilderEntity::get_sparse_data_from_vector(const void *vec) const {\n uint32_t vec_size = vector_size();\n const char *vec_ptr = reinterpret_cast(vec);\n\n uint64_t offset = *((uint64_t *)(vec_ptr + vec_size));\n uint32_t sparse_vector_len =\n *((uint32_t *)(vec_ptr + vec_size + sizeof(uint64_t)));\n\n const void *sparse_data = get_sparse_data(offset, sparse_vector_len);\n if (ailego_unlikely(sparse_data == nullptr)) {\n LOG_ERROR(\"Get nullptr sparse, offset=%zu, len=%u\", (size_t)offset,\n sparse_vector_len);\n\n return std::make_pair(nullptr, 0);\n }\n\n return std::make_pair(sparse_data, sparse_vector_len);\n}\n\nint HnswSparseBuilderEntity::get_sparse_data_from_vector(\n const void *vec, IndexStorage::MemoryBlock &block,\n int &sparse_length) const {\n std::pair sparse_data =\n get_sparse_data_from_vector(vec);\n block.reset((void *)sparse_data.first);\n sparse_length = sparse_data.second;\n return 0;\n}\n\nconst Neighbors HnswSparseBuilderEntity::get_neighbors(level_t level,\n node_id_t id) const {\n const NeighborsHeader *hd = get_neighbor_header(level, id);\n return {hd->neighbor_cnt, hd->neighbors};\n}\n\nint HnswSparseBuilderEntity::update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) {\n NeighborsHeader *hd =\n const_cast(get_neighbor_header(level, id));\n for (size_t i = 0; i < neighbors.size(); ++i) {\n hd->neighbors[i] = neighbors[i].first;\n }\n hd->neighbor_cnt = neighbors.size();\n\n // std::cout << \"id: \" << id << \", neighbour, id: \";\n // for (size_t i = 0; i < neighbors.size(); ++i) {\n // if (i == neighbors.size()-1)\n // std::cout << neighbors[i].first << \", score:\" << neighbors[i].second <<\n // std::endl;\n // else\n // std::cout << neighbors[i].first << \", score:\" << neighbors[i].second <<\n // \", id: \";\n // }\n\n return 0;\n}\n\nvoid HnswSparseBuilderEntity::add_neighbor(level_t level, node_id_t id,\n uint32_t /*size*/,\n node_id_t neighbor_id) {\n NeighborsHeader *hd =\n const_cast(get_neighbor_header(level, id));\n hd->neighbors[hd->neighbor_cnt++] = neighbor_id;\n\n return;\n}\n\nint HnswSparseBuilderEntity::dump(const IndexDumper::Pointer &dumper) {\n key_t *keys =\n reinterpret_cast(const_cast(keys_buffer_.data()));\n auto ret =\n dump_segments(dumper, keys, [&](node_id_t id) { return get_level(id); });\n if (ailego_unlikely(ret < 0)) {\n return ret;\n }\n\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_builder_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"hnsw_sparse_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSparseBuilderEntity : public HnswSparseEntity {\n public:\n //! Add vector and key to hnsw entity, and local id will be saved in id\n virtual int add_vector(level_t level, key_t key, const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_vec,\n node_id_t *id) override;\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const override;\n\n //! Get vector feature data by key\n virtual const void *get_vector_meta(node_id_t id) const override;\n\n virtual int get_vector_meta(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n\n //! Batch get vectors feature data by keys\n virtual int get_vector_metas(const node_id_t *ids, uint32_t count,\n const void **vecs) const override;\n virtual int get_vector_metas(\n const node_id_t *ids, uint32_t count,\n std::vector &block_vecs) const override;\n\n //! Get the node id's neighbors on graph level\n const NeighborsHeader *get_neighbor_header(level_t level,\n node_id_t id) const {\n if (level == 0) {\n return reinterpret_cast(\n neighbors_buffer_.data() + neighbors_size_ * id);\n } else {\n size_t offset = neighbors_index_[id].offset;\n return reinterpret_cast(\n upper_neighbors_buffer_.data() + offset +\n (level - 1) * upper_neighbors_size_);\n }\n }\n\n //! Get the node id's neighbors on graph level\n virtual const Neighbors get_neighbors(level_t level,\n node_id_t id) const override;\n\n //! Replace node id in level's neighbors\n virtual int update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) override;\n\n //! add a neighbor to id in graph level\n virtual void add_neighbor(level_t level, node_id_t id, uint32_t size,\n node_id_t neighbor_id) override;\n\n //! Get vector sparse feature data by chunk index and offset\n virtual const void *get_sparse_data(uint64_t offset,\n uint32_t len) const override;\n //! Get sparse data from id\n virtual const void *get_sparse_data(node_id_t id) const override;\n\n virtual int get_sparse_data(uint64_t offset, uint32_t len,\n IndexStorage::MemoryBlock &block) const override;\n\n virtual int get_sparse_data(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n\n //! Get sparse data from vector\n virtual std::pair get_sparse_data_from_vector(\n const void *vec) const override;\n\n virtual int get_sparse_data_from_vector(const void *vec,\n IndexStorage::MemoryBlock &block,\n int &sparse_length) const override;\n\n //! Dump the hnsw graph to dumper\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Cleanup the entity\n virtual int cleanup(void) override;\n\n public:\n //! Constructor\n HnswSparseBuilderEntity();\n\n //! Get the node graph level by id\n level_t get_level(node_id_t id) const {\n return neighbors_index_[id].level;\n }\n\n //! Init builerEntity\n int init();\n\n //! reserve buffer space for documents\n //! @param docs number of documents\n //! @param total_sparse_count total dim of sparse count\n int reserve_space(size_t docs, size_t total_sparse_count);\n\n //! Set memory quota params\n inline void set_memory_quota(size_t memory_quota) {\n memory_quota_ = memory_quota;\n }\n\n //! Get neighbors size\n inline size_t neighbors_size() const {\n return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t);\n }\n\n //! Get upper neighbors size\n inline size_t upper_neighbors_size() const {\n return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t);\n }\n\n public:\n HnswSparseBuilderEntity(const HnswSparseBuilderEntity &) = delete;\n HnswSparseBuilderEntity &operator=(const HnswSparseBuilderEntity &) = delete;\n\n private:\n friend class HnswSparseSearcherEntity;\n\n //! class internal used only\n struct SparseNeighborIndex {\n SparseNeighborIndex(size_t off, level_t l) : offset(off), level(l) {}\n uint64_t offset : 48;\n uint64_t level : 16;\n };\n\n std::string vectors_buffer_{}; // aligned vectors\n std::string keys_buffer_{}; // aligned vectors\n std::string neighbors_buffer_{}; // level 0 neighbors buffer\n std::string upper_neighbors_buffer_{}; // upper layer neighbors buffer\n\n std::string sparse_data_buffer_{}; // aligned spase data buffer\n size_t sparse_data_offset_{0}; //\n\n // upper layer offset + level in upper_neighbors_buffer_\n std::vector neighbors_index_{};\n size_t memory_quota_{0UL};\n size_t neighbors_size_{0U}; // level 0 neighbors size\n size_t upper_neighbors_size_{0U}; // level 0 neighbors size\n size_t padding_size_{}; // padding size for each vector element\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_sparse_chunk.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nint SparseChunkBroker::init_storage(size_t chunk_size) {\n chunk_meta_.clear();\n chunk_meta_.chunk_size = chunk_size;\n chunk_meta_.create_time = ailego::Realtime::Seconds();\n stats_.set_create_time(chunk_meta_.create_time);\n chunk_meta_.update_time = ailego::Realtime::Seconds();\n stats_.set_update_time(chunk_meta_.update_time);\n\n //! alloc meta chunk\n size_t size = sizeof(HnswSparseChunkMeta);\n size = (size + page_mask_) & (~page_mask_);\n const std::string segment_id =\n make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId);\n int ret = stg_->append(segment_id, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Storage append segment failed for %s\", IndexError::What(ret));\n return ret;\n }\n chunk_meta_segment_ = get_chunk(CHUNK_TYPE_META, kDefaultChunkSeqId);\n if (ailego_unlikely(!chunk_meta_segment_)) {\n LOG_ERROR(\"Get meta segment failed\");\n return IndexError_Runtime;\n }\n\n //! update meta info and write to storage\n chunk_meta_.chunk_cnts[CHUNK_TYPE_META] += 1;\n chunk_meta_.total_size += size;\n (*stats_.mutable_index_size()) += size;\n size = chunk_meta_segment_->write(0UL, &chunk_meta_,\n sizeof(HnswSparseChunkMeta));\n if (ailego_unlikely(size != sizeof(HnswSparseChunkMeta))) {\n LOG_ERROR(\"Storage write data failed, wsize=%zu\", size);\n return IndexError_WriteData;\n }\n\n return 0;\n}\n\nint SparseChunkBroker::load_storage(size_t chunk_size) {\n IndexStorage::MemoryBlock data_block;\n size_t size = chunk_meta_segment_->read(0UL, data_block,\n chunk_meta_segment_->data_size());\n if (size != sizeof(HnswSparseChunkMeta)) {\n LOG_ERROR(\"Invalid hnsw meta chunk, read size=%zu chunk size=%zu\", size,\n chunk_meta_segment_->data_size());\n return IndexError_InvalidFormat;\n }\n std::memcpy(&chunk_meta_, data_block.data(), size);\n if (chunk_meta_.chunk_size != chunk_size) {\n LOG_ERROR(\n \"Params hnsw chunk size=%zu mismatch from previous %zu \"\n \"in index\",\n chunk_size, (size_t)chunk_meta_.chunk_size);\n return IndexError_Mismatch;\n }\n\n *stats_.mutable_check_point() = stg_->check_point();\n stats_.set_revision_id(chunk_meta_.revision_id);\n stats_.set_update_time(chunk_meta_.update_time);\n stats_.set_create_time(chunk_meta_.create_time);\n\n char create_time[32];\n char update_time[32];\n ailego::Realtime::Gmtime(chunk_meta_.create_time, \"%Y-%m-%d %H:%M:%S\",\n create_time, sizeof(create_time));\n ailego::Realtime::Gmtime(chunk_meta_.update_time, \"%Y-%m-%d %H:%M:%S\",\n update_time, sizeof(update_time));\n LOG_DEBUG(\n \"Load index, indexSize=%zu chunkSize=%zu nodeChunks=%zu \"\n \"upperNeighborChunks=%zu revisionId=%zu \"\n \"createTime=%s updateTime=%s\",\n (size_t)chunk_meta_.total_size, (size_t)chunk_meta_.chunk_size,\n (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_NODE],\n (size_t)chunk_meta_.chunk_cnts[CHUNK_TYPE_UPPER_NEIGHBOR],\n (size_t)chunk_meta_.revision_id, create_time, update_time);\n\n return 0;\n}\n\nint SparseChunkBroker::open(IndexStorage::Pointer stg, size_t max_index_size,\n size_t chunk_size, bool check_crc) {\n if (ailego_unlikely(stg_)) {\n LOG_ERROR(\"An storage instance is already opened\");\n return IndexError_Duplicate;\n }\n stg_ = std::move(stg);\n check_crc_ = check_crc;\n max_chunks_size_ = max_index_size;\n dirty_ = false;\n\n const std::string segment_id =\n make_segment_id(CHUNK_TYPE_META, kDefaultChunkSeqId);\n chunk_meta_segment_ = stg_->get(segment_id);\n if (!chunk_meta_segment_) {\n LOG_DEBUG(\"Create new index\");\n return init_storage(chunk_size);\n }\n\n return load_storage(chunk_size);\n}\n\nint SparseChunkBroker::close(void) {\n flush(0UL);\n\n stg_.reset();\n check_crc_ = false;\n dirty_ = false;\n\n return 0;\n}\n\nint SparseChunkBroker::flush(uint64_t checkpoint) {\n ailego_assert_with(chunk_meta_segment_, \"invalid meta segment\");\n\n chunk_meta_.update_time = ailego::Realtime::Seconds();\n stats_.set_update_time(chunk_meta_.update_time);\n\n size_t size = chunk_meta_segment_->write(0UL, &chunk_meta_,\n sizeof(HnswSparseChunkMeta));\n if (ailego_unlikely(size != sizeof(HnswSparseChunkMeta))) {\n LOG_ERROR(\"Storage write data failed, wsize=%zu\", size);\n }\n\n stg_->refresh(checkpoint);\n int ret = stg_->flush();\n if (ret == 0) {\n (*stats_.mutable_check_point()) = checkpoint;\n } else {\n LOG_ERROR(\"Storage flush failed for %s\", IndexError::What(ret));\n }\n return ret;\n}\n\nstd::pair SparseChunkBroker::alloc_chunk(\n int type, uint64_t seq_id, size_t size) {\n ailego_assert_with(type < CHUNK_TYPE_MAX, \"chunk type overflow\");\n\n SparseChunk::Pointer chunk;\n if (ailego_unlikely(!stg_)) {\n LOG_ERROR(\"Init storage first\");\n return std::make_pair(IndexError_Uninitialized, chunk);\n }\n\n //! check exist a empty chunk with the same name\n chunk = get_chunk(type, seq_id);\n if (chunk) {\n if (ailego_unlikely(chunk->capacity() == size &&\n chunk->data_size() == 0UL)) {\n LOG_ERROR(\"Exist invalid chunk size %zu, expect size %zu\",\n chunk->capacity(), size);\n chunk.reset();\n return std::make_pair(IndexError_Runtime, chunk);\n }\n return std::make_pair(0, chunk);\n }\n //! align to page size\n size = (size + page_mask_) & (~page_mask_);\n if (ailego_unlikely(chunk_meta_.total_size + size >= max_chunks_size_)) {\n LOG_ERROR(\"No space to new a chunk, curIndexSize=%zu allocSize=%zu\",\n (size_t)chunk_meta_.total_size, size);\n return std::make_pair(IndexError_IndexFull, chunk);\n }\n\n std::string segment_id = make_segment_id(type, seq_id);\n int ret = stg_->append(segment_id, size);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Storage append segment failed for %s\", IndexError::What(ret));\n return std::make_pair(ret, chunk);\n }\n chunk_meta_.chunk_cnts[type] += 1;\n chunk_meta_.total_size += size;\n (*stats_.mutable_index_size()) += size;\n\n size = chunk_meta_segment_->write(0UL, &chunk_meta_,\n sizeof(HnswSparseChunkMeta));\n if (ailego_unlikely(size != sizeof(HnswSparseChunkMeta))) {\n LOG_ERROR(\"Storage append segment failed, wsize=%zu\", size);\n }\n\n chunk = get_chunk(type, seq_id);\n return std::make_pair(chunk ? 0 : IndexError_NoMemory, chunk);\n}\n\nSparseChunk::Pointer SparseChunkBroker::get_chunk(int type,\n uint64_t seq_id) const {\n std::string segment_id = make_segment_id(type, seq_id);\n return stg_->get(segment_id);\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_chunk.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nusing SparseChunk = IndexStorage::Segment;\n\nclass SparseChunkBroker {\n public:\n typedef std::shared_ptr Pointer;\n\n enum CHUNK_TYPE {\n CHUNK_TYPE_HEADER = 1,\n CHUNK_TYPE_META = 2,\n CHUNK_TYPE_NODE = 3,\n CHUNK_TYPE_UPPER_NEIGHBOR = 4,\n CHUNK_TYPE_NEIGHBOR_INDEX = 5,\n CHUNK_TYPE_SPARSE_NODE = 6,\n CHUNK_TYPE_MAX = 8\n };\n static constexpr size_t kDefaultChunkSeqId = 0UL;\n\n SparseChunkBroker(IndexStreamer::Stats &stats) : stats_(stats) {\n page_mask_ = ailego::MemoryHelper::PageSize() - 1;\n }\n\n //! Open storage\n int open(IndexStorage::Pointer stg, size_t max_index_size, size_t chunk_size,\n bool check_crc);\n\n int close(void);\n\n int flush(uint64_t checkpoint);\n\n //! alloc a new chunk with size, not thread-safe\n std::pair alloc_chunk(int type, uint64_t seq_id,\n size_t size);\n\n //! alloc a new chunk with chunk size\n inline std::pair alloc_chunk(int type,\n uint64_t seq_id) {\n return alloc_chunk(type, seq_id, chunk_meta_.chunk_size);\n }\n\n SparseChunk::Pointer get_chunk(int type, uint64_t seq_id) const;\n\n inline size_t get_chunk_cnt(int type) const {\n ailego_assert_with(type < CHUNK_TYPE_MAX, \"chunk type overflow\");\n return chunk_meta_.chunk_cnts[type];\n }\n\n inline bool dirty(void) const {\n return dirty_;\n }\n\n inline void mark_dirty(void) {\n if (!dirty_) {\n dirty_ = true;\n chunk_meta_.revision_id += 1;\n stats_.set_revision_id(chunk_meta_.revision_id);\n }\n }\n\n const IndexStorage::Pointer storage(void) const {\n return stg_;\n }\n\n private:\n SparseChunkBroker(const SparseChunkBroker &) = delete;\n SparseChunkBroker &operator=(const SparseChunkBroker &) = delete;\n\n struct HnswSparseChunkMeta {\n HnswSparseChunkMeta(void) {\n memset(this, 0, sizeof(HnswSparseChunkMeta));\n }\n void clear() {\n memset(this, 0, sizeof(HnswSparseChunkMeta));\n }\n\n uint64_t chunk_cnts[CHUNK_TYPE_MAX];\n uint64_t chunk_size; // size of per chunk\n uint64_t total_size; // total size of allocated chunk\n uint64_t revision_id; // index revision\n uint64_t create_time;\n uint64_t update_time;\n uint64_t reserved[3];\n };\n\n static_assert(sizeof(HnswSparseChunkMeta) % 32 == 0,\n \"HnswSparseChunkMeta must be aligned with 32 bytes\");\n\n //! Init the storage after open an empty index\n int init_storage(size_t chunk_size);\n\n //! Load index from storage\n int load_storage(size_t chunk_size);\n\n static inline const std::string make_segment_id(int type, uint64_t seq_id) {\n return \"HnswT\" + ailego::StringHelper::ToString(type) + \"S\" +\n ailego::StringHelper::ToString(seq_id);\n }\n\n private:\n IndexStreamer::Stats &stats_;\n HnswSparseChunkMeta chunk_meta_{};\n size_t page_mask_{0UL};\n size_t max_chunks_size_{0UL};\n IndexStorage::Pointer stg_{};\n IndexStorage::Segment::Pointer chunk_meta_segment_{};\n bool check_crc_{false};\n bool dirty_{false}; // set as true if index is modified , the flag\n // will not be cleared even if flushed\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_context.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_sparse_context.h\"\n#include \n#include \"hnsw_sparse_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswSparseContext::HnswSparseContext(const IndexMetric::Pointer &metric,\n const HnswSparseEntity::Pointer &entity)\n : IndexContext(metric), entity_(entity), dc_(entity_.get(), metric) {}\n\nHnswSparseContext::~HnswSparseContext() {\n visit_filter_.destroy();\n}\n\nint HnswSparseContext::init(ContextType type) {\n int ret;\n uint32_t doc_cnt;\n\n type_ = type;\n\n switch (type) {\n case kSparseBuilderContext:\n ret = visit_filter_.init(VisitFilter::ByteMap, entity_->doc_cnt(),\n max_scan_num_, negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n candidates_.limit(max_scan_num_);\n update_heap_.limit(entity_->l0_neighbor_cnt() + 1);\n break;\n\n case kSparseSearcherContext:\n ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_,\n negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n candidates_.limit(max_scan_num_);\n break;\n\n case kSparseStreamerContext:\n // maxScanNum is unknown if inited from streamer, so the docCnt may\n // change. we need to compute maxScanNum by scan ratio, and preserve\n // max_doc_cnt space from visit filter\n doc_cnt = entity_->doc_cnt();\n max_scan_num_ = compute_max_scan_num(doc_cnt);\n reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt);\n ret = visit_filter_.init(filter_mode_, reserve_max_doc_cnt_,\n max_scan_num_, negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n\n update_heap_.limit(entity_->l0_neighbor_cnt() + 1);\n candidates_.limit(max_scan_num_);\n\n check_need_adjuct_ctx();\n break;\n\n default:\n LOG_ERROR(\"Init context failed\");\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nint HnswSparseContext::update(const ailego::Params ¶ms) {\n LOG_DEBUG(\"Update hnsw context params\");\n\n auto update_visit_filter_param = [&]() {\n bool need_update = false;\n std::string p;\n switch (type_) {\n case kSparseSearcherContext:\n p = PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_ENABLE;\n break;\n case kSparseStreamerContext:\n p = PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_ENABLE;\n break;\n }\n\n if (params.has(p)) {\n bool bf_enabled;\n params.get(p, &bf_enabled);\n if (bf_enabled ^ (filter_mode_ == VisitFilter::BloomFilter)) {\n need_update = true;\n filter_mode_ =\n bf_enabled ? VisitFilter::BloomFilter : VisitFilter::ByteMap;\n }\n }\n\n float prob = negative_probability_;\n p.clear();\n switch (type_) {\n case kSparseSearcherContext:\n p = PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB;\n break;\n case kSparseStreamerContext:\n p = PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB;\n break;\n }\n params.get(p, &prob);\n if (filter_mode_ == VisitFilter::BloomFilter &&\n std::abs(prob - negative_probability_) > 1e-6) {\n need_update = true;\n }\n if (need_update) {\n visit_filter_.destroy();\n int max_doc_cnt = 0;\n if (type_ == kSparseSearcherContext) {\n max_doc_cnt = entity_->doc_cnt();\n } else {\n max_doc_cnt = reserve_max_doc_cnt_;\n }\n int ret = visit_filter_.init(filter_mode_, max_doc_cnt, max_scan_num_,\n negative_probability_);\n if (ret != 0) {\n LOG_ERROR(\"Create filter failed, mode %d\", filter_mode_);\n return ret;\n }\n }\n return 0;\n };\n\n switch (type_) {\n case kSparseSearcherContext:\n if (params.has(PARAM_HNSW_SPARSE_SEARCHER_EF)) {\n params.get(PARAM_HNSW_SPARSE_SEARCHER_EF, &ef_);\n topk_heap_.limit(std::max(topk_, ef_));\n }\n\n if (params.has(PARAM_HNSW_SPARSE_SEARCHER_MAX_SCAN_RATIO)) {\n params.get(PARAM_HNSW_SPARSE_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_);\n max_scan_num_ =\n static_cast(max_scan_ratio_ * entity_->doc_cnt());\n max_scan_num_ = std::max(10000U, max_scan_num_);\n }\n\n if (params.has(PARAM_HNSW_SPARSE_SEARCHER_BRUTE_FORCE_THRESHOLD)) {\n params.get(PARAM_HNSW_SPARSE_SEARCHER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n }\n\n return update_visit_filter_param();\n\n case kSparseStreamerContext:\n if (params.has(PARAM_HNSW_SPARSE_STREAMER_EF)) {\n params.get(PARAM_HNSW_SPARSE_STREAMER_EF, &ef_);\n topk_heap_.limit(std::max(topk_, ef_));\n }\n params.get(PARAM_HNSW_SPARSE_STREAMER_EF, &ef_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_);\n if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) {\n LOG_ERROR(\"[%s] must be in range (0.0f,1.0f]\",\n PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO.c_str());\n return IndexError_InvalidArgument;\n }\n if (max_scan_limit_ < min_scan_limit_) {\n LOG_ERROR(\"[%s] must be >= [%s]\",\n PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT.c_str(),\n PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (params.has(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD)) {\n params.get(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n }\n\n return update_visit_filter_param();\n\n default:\n LOG_ERROR(\"update context failed, type=%u\", type_);\n return IndexError_Runtime;\n }\n}\n\n\nint HnswSparseContext::update_context(ContextType type,\n const IndexMeta & /*meta*/,\n const IndexMetric::Pointer &metric,\n const HnswSparseEntity::Pointer &entity,\n uint32_t magic_num) {\n uint32_t doc_cnt;\n\n if (ailego_unlikely(type != type_)) {\n LOG_ERROR(\n \"HnswSparseContext doesn't support shared by different type, \"\n \"src=%u dst=%u\",\n type_, type);\n return IndexError_Unsupported;\n }\n\n magic_ = kInvalidMgic;\n\n // TODO: support change filter mode?\n switch (type) {\n case kSparseBuilderContext:\n LOG_ERROR(\"BuildContext doesn't support update\");\n return IndexError_NotImplemented;\n\n case kSparseSearcherContext:\n if (!visit_filter_.reset(entity->doc_cnt(), max_scan_num_)) {\n LOG_ERROR(\"Reset filter failed, mode %d\", visit_filter_.get_mode());\n return IndexError_Runtime;\n }\n\n candidates_.limit(max_scan_num_);\n topk_heap_.limit(std::max(topk_, ef_));\n break;\n\n case kSparseStreamerContext:\n doc_cnt = entity->doc_cnt();\n max_scan_num_ = compute_max_scan_num(doc_cnt);\n reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt);\n if (!visit_filter_.reset(reserve_max_doc_cnt_, max_scan_num_)) {\n LOG_ERROR(\"Reset filter failed, mode %d\", visit_filter_.get_mode());\n return IndexError_Runtime;\n }\n\n update_heap_.limit(entity->l0_neighbor_cnt() + 1);\n candidates_.limit(max_scan_num_);\n topk_heap_.limit(std::max(topk_, ef_));\n break;\n\n default:\n LOG_ERROR(\"update context failed\");\n return IndexError_Runtime;\n }\n\n entity_ = entity;\n dc_.update(entity_.get(), metric);\n magic_ = magic_num;\n level_topks_.clear();\n\n return 0;\n}\n\nvoid HnswSparseContext::fill_random_to_topk_full(void) {\n static std::mt19937 mt(\n std::chrono::system_clock::now().time_since_epoch().count());\n std::uniform_int_distribution dt(0, entity_->doc_cnt() - 1);\n std::function gen;\n node_id_t seqid;\n std::function myfilter = [](node_id_t) { return false; };\n if (this->filter().is_valid()) {\n myfilter = [&](node_id_t id) {\n return this->filter()(entity_->get_key(id));\n };\n }\n\n if (topk_heap_.limit() < entity_->doc_cnt() / 2) {\n gen = [&](void) { return dt(mt); };\n } else {\n // If topk limit is big value, gen sequential id from an random initial\n seqid = dt(mt);\n gen = [&](void) {\n seqid = seqid == (entity_->doc_cnt() - 1) ? 0 : (seqid + 1);\n return seqid;\n };\n }\n\n for (size_t i = 0; !topk_heap_.full() && i < entity_->doc_cnt(); ++i) {\n const auto id = gen();\n if (!visit_filter_.visited(id) && !myfilter(id)) {\n visit_filter_.set_visited(id);\n topk_heap_.emplace(id, dc_.dist(id));\n }\n }\n return;\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_context.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"utility/sparse_utility.h\"\n#include \"utility/visit_filter.h\"\n#include \"hnsw_sparse_dist_calculator.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSparseContext : public IndexContext {\n public:\n //! Index Context Pointer\n typedef std::unique_ptr Pointer;\n\n enum ContextType {\n kUnknownContext = 0,\n kSparseSearcherContext = 1,\n kSparseBuilderContext = 2,\n kSparseStreamerContext = 3,\n };\n\n //! Construct\n HnswSparseContext(const IndexMetric::Pointer &metric,\n const HnswSparseEntity::Pointer &entity);\n\n //! Destructor\n virtual ~HnswSparseContext();\n\n public:\n //! Set topk of search result\n virtual void set_topk(uint32_t val) override {\n topk_ = val;\n topk_heap_.limit(std::max(val, ef_));\n }\n\n //! Retrieve search result\n virtual const IndexDocumentList &result(void) const override {\n return results_[0];\n }\n\n //! Retrieve search result\n virtual const IndexDocumentList &result(size_t idx) const override {\n return results_[idx];\n }\n\n //! Retrieve result object for output\n virtual IndexDocumentList *mutable_result(size_t idx) override {\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n return &results_[idx];\n }\n\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(void) const override {\n return group_results_[0];\n }\n\n //! Retrieve search group result with index\n virtual const IndexGroupDocumentList &group_result(\n size_t idx) const override {\n return group_results_[idx];\n }\n\n virtual uint32_t magic(void) const override {\n return magic_;\n }\n\n //! Set mode of debug\n virtual void set_debug_mode(bool enable) override {\n debug_mode_ = enable;\n }\n\n //! Retrieve mode of debug\n virtual bool debug_mode(void) const override {\n return this->debugging();\n }\n\n //! Retrieve string of debug\n virtual std::string debug_string(void) const override {\n char buf[4096];\n size_t size = snprintf(\n buf, sizeof(buf),\n \"scan_cnt=%zu,get_vector_cnt=%u,get_neighbors_cnt=%u,dup_node=%u\",\n get_scan_num(), stats_get_vector_cnt_, stats_get_neighbors_cnt_,\n stats_visit_dup_cnt_);\n return std::string(buf, size);\n }\n\n //! Update the parameters of context\n virtual int update(const ailego::Params ¶ms) override;\n\n public:\n //! Init context\n int init(ContextType type);\n\n //! Update context, the context may be shared by different searcher/streamer\n int update_context(ContextType type, const IndexMeta &meta,\n const IndexMetric::Pointer &metric,\n const HnswSparseEntity::Pointer &entity,\n uint32_t magic_num);\n\n inline const HnswSparseEntity &get_entity() const {\n return *entity_;\n }\n\n inline void resize_results(size_t size) {\n if (group_by_search()) {\n group_results_.resize(size);\n } else {\n results_.resize(size);\n }\n }\n\n inline void topk_to_result() {\n return topk_to_result(0);\n }\n\n //! Construct result from topk heap, result will be normalized\n inline void topk_to_result(uint32_t idx) {\n if (group_by_search()) {\n topk_to_group_result(idx);\n } else {\n topk_to_single_result(idx);\n }\n }\n\n inline void recal_topk_dist() {\n TopkHeap heap(topk_heap_);\n topk_heap_.clear();\n\n for (size_t i = 0; i < heap.size(); ++i) {\n node_id_t id = heap[i].first;\n dist_t dist = dc_.dist(id);\n topk_heap_.emplace_back(id, dist);\n }\n }\n\n inline void topk_to_single_result(uint32_t idx) {\n if (force_padding_topk_ && !topk_heap_.full() &&\n topk_heap_.size() < entity_->doc_cnt()) {\n this->fill_random_to_topk_full();\n }\n if (ailego_unlikely(topk_heap_.size() == 0)) {\n return;\n }\n\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n int size = std::min(topk_, static_cast(topk_heap_.size()));\n topk_heap_.sort();\n results_[idx].clear();\n\n for (int i = 0; i < size; ++i) {\n auto score = topk_heap_[i].second;\n if (score > this->threshold()) {\n break;\n }\n\n node_id_t id = topk_heap_[i].first;\n if (fetch_vector_) {\n IndexSparseDocument sparse_doc;\n IndexStorage::MemoryBlock vec_block;\n entity_->get_sparse_data(id, vec_block);\n const void *sparse_data = vec_block.data();\n if (sparse_data != nullptr) {\n SparseUtility::ReverseSparseFormat(sparse_data, sparse_doc,\n entity_->sparse_unit_size());\n }\n\n results_[idx].emplace_back(entity_->get_key(id), score, id,\n entity_->get_vector_meta(id), sparse_doc);\n } else {\n results_[idx].emplace_back(entity_->get_key(id), score, id);\n }\n }\n\n return;\n }\n\n //! Construct result from topk heap, result will be normalized\n inline void topk_to_group_result(uint32_t idx) {\n ailego_assert_with(idx < group_results_.size(), \"invalid idx\");\n\n group_results_[idx].clear();\n\n std::vector> group_topk_list;\n std::vector> best_score_in_groups;\n for (auto itr = group_topk_heaps_.begin(); itr != group_topk_heaps_.end();\n itr++) {\n const std::string &group_id = (*itr).first;\n auto &heap = (*itr).second;\n heap.sort();\n\n if (heap.size() > 0) {\n float best_score = heap[0].second;\n best_score_in_groups.push_back(std::make_pair(group_id, best_score));\n }\n }\n\n std::sort(best_score_in_groups.begin(), best_score_in_groups.end(),\n [](const std::pair &a,\n const std::pair &b) -> int {\n return a.second < b.second;\n });\n\n // truncate to group num\n for (uint32_t i = 0; i < group_num() && i < best_score_in_groups.size();\n ++i) {\n const std::string &group_id = best_score_in_groups[i].first;\n\n group_topk_list.emplace_back(\n std::make_pair(group_id, group_topk_heaps_[group_id]));\n }\n\n group_results_[idx].resize(group_topk_list.size());\n\n for (uint32_t i = 0; i < group_topk_list.size(); ++i) {\n const std::string &group_id = group_topk_list[i].first;\n group_results_[idx][i].set_group_id(group_id);\n\n uint32_t size = std::min(\n group_topk_, static_cast(group_topk_list[i].second.size()));\n\n for (uint32_t j = 0; j < size; ++j) {\n auto score = group_topk_list[i].second[j].second;\n if (score > this->threshold()) {\n break;\n }\n\n node_id_t id = group_topk_list[i].second[j].first;\n\n if (fetch_vector_) {\n group_results_[idx][i].mutable_docs()->emplace_back(\n entity_->get_key(id), score, id, entity_->get_vector_meta(id));\n } else {\n group_results_[idx][i].mutable_docs()->emplace_back(\n entity_->get_key(id), score, id);\n }\n }\n }\n }\n\n inline void reset_query(const void *query) {\n dc_.reset_query(query);\n dc_.clear_compare_cnt();\n }\n\n inline HnswSparseDistCalculator &dist_calculator() {\n return dc_;\n }\n\n inline TopkHeap &topk_heap() {\n return topk_heap_;\n }\n\n inline TopkHeap &update_heap() {\n return update_heap_;\n }\n\n inline VisitFilter &visit_filter() {\n return visit_filter_;\n }\n\n inline CandidateHeap &candidates() {\n return candidates_;\n }\n\n inline void set_max_scan_num(uint32_t max_scan_num) {\n max_scan_num_ = max_scan_num;\n }\n\n inline void set_max_scan_limit(uint32_t max_scan_limit) {\n max_scan_limit_ = max_scan_limit;\n }\n\n inline void set_min_scan_limit(uint32_t min_scan_limit) {\n min_scan_limit_ = min_scan_limit;\n }\n\n inline void set_ef(uint32_t v) {\n ef_ = v;\n }\n\n inline void set_filter_mode(uint32_t v) {\n filter_mode_ = v;\n }\n\n inline void set_filter_negative_probability(float v) {\n negative_probability_ = v;\n }\n\n inline void set_max_scan_ratio(float v) {\n max_scan_ratio_ = v;\n }\n\n virtual void set_magic(uint32_t v) {\n magic_ = v;\n }\n\n virtual void set_force_padding_topk(bool v) {\n force_padding_topk_ = v;\n }\n\n virtual void set_bruteforce_threshold(uint32_t v) override {\n bruteforce_threshold_ = v;\n }\n\n inline uint32_t get_bruteforce_threshold() const {\n return bruteforce_threshold_;\n }\n\n virtual void set_fetch_vector(bool v) override {\n fetch_vector_ = v;\n }\n\n virtual bool fetch_vector() const override {\n return fetch_vector_;\n }\n\n //! Reset context\n void reset(void) override {\n set_filter(nullptr);\n reset_threshold();\n set_fetch_vector(false);\n set_group_params(0, 0);\n reset_group_by();\n }\n\n inline std::map &group_topk_heaps() {\n return group_topk_heaps_;\n }\n\n inline TopkHeap &level_topk(int level) {\n if (ailego_unlikely(level_topks_.size() <= static_cast(level))) {\n int cur_level = level_topks_.size();\n level_topks_.resize(level + 1);\n for (; cur_level <= level; ++cur_level) {\n size_t heap_size = std::max(entity_->neighbor_cnt(cur_level),\n entity_->ef_construction());\n level_topks_[cur_level].clear();\n level_topks_[cur_level].limit(heap_size);\n }\n }\n\n return level_topks_[level];\n }\n\n inline void check_need_adjuct_ctx(void) {\n check_need_adjuct_ctx(entity_->doc_cnt());\n }\n\n inline size_t compute_reserve_cnt(uint32_t cur_doc) const {\n if (cur_doc > kMaxReserveDocCnt) {\n return kMaxReserveDocCnt;\n } else if (cur_doc < kMinReserveDocCnt) {\n return kMinReserveDocCnt;\n }\n return cur_doc;\n }\n\n //! candidates heap and visitfilter need to resize as doc cnt growing up\n inline void check_need_adjuct_ctx(uint32_t doc_cnt) {\n if (ailego_unlikely(doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_)) {\n while (doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_) {\n reserve_max_doc_cnt_ =\n reserve_max_doc_cnt_ + compute_reserve_cnt(reserve_max_doc_cnt_);\n }\n uint32_t max_scan_cnt = compute_max_scan_num(reserve_max_doc_cnt_);\n max_scan_num_ = max_scan_cnt;\n visit_filter_.reset(reserve_max_doc_cnt_, max_scan_cnt);\n candidates_.clear();\n candidates_.limit(max_scan_num_);\n }\n }\n\n inline uint32_t compute_max_scan_num(uint32_t max_doc_cnt) const {\n uint32_t max_scan = max_doc_cnt * max_scan_ratio_;\n if (max_scan < min_scan_limit_) {\n max_scan = min_scan_limit_;\n } else if (max_scan > max_scan_limit_) {\n max_scan = max_scan_limit_;\n }\n return max_scan;\n }\n\n inline size_t get_scan_num() const {\n return dc_.compare_cnt();\n }\n\n inline uint64_t reach_scan_limit() const {\n return dc_.compare_cnt() >= max_scan_num_;\n }\n\n inline bool error() const {\n return dc_.error();\n }\n\n inline void clear() {\n dc_.clear();\n if (ailego_unlikely(this->debugging())) {\n stats_get_neighbors_cnt_ = 0u;\n stats_get_vector_cnt_ = 0u;\n stats_visit_dup_cnt_ = 0u;\n }\n // do not clear results_ for the next query will need it\n for (auto &it : results_) {\n it.clear();\n }\n }\n\n uint32_t *mutable_stats_get_neighbors() {\n return &stats_get_neighbors_cnt_;\n }\n\n uint32_t *mutable_stats_get_vector() {\n return &stats_get_vector_cnt_;\n }\n\n uint32_t *mutable_stats_visit_dup_cnt() {\n return &stats_visit_dup_cnt_;\n }\n\n inline bool debugging(void) const {\n return debug_mode_;\n }\n\n inline void update_dist_caculator_distance(\n const IndexMetric::MatrixSparseDistance &distance) {\n dc_.update_distance(distance);\n }\n\n //! Get topk\n inline uint32_t topk() const override {\n return topk_;\n }\n\n //! Get group topk\n inline uint32_t group_topk() const {\n return group_topk_;\n }\n\n //! Get group num\n inline uint32_t group_num() const {\n return group_num_;\n }\n\n //! Get if group by search\n inline bool group_by_search() {\n return group_num_ > 0;\n }\n\n //! Set group params\n void set_group_params(uint32_t group_num, uint32_t group_topk) override {\n group_num_ = group_num;\n group_topk_ = group_topk;\n\n topk_ = group_topk_ * group_num_;\n\n topk_heap_.limit(std::max(topk_, ef_));\n\n group_topk_heaps_.clear();\n }\n\n private:\n // Filling random nodes if topk not full\n void fill_random_to_topk_full(void);\n\n constexpr static uint32_t kTriggerReserveCnt = 4096UL;\n constexpr static uint32_t kMinReserveDocCnt = 4096UL;\n constexpr static uint32_t kMaxReserveDocCnt = 128 * 1024UL;\n constexpr static uint32_t kInvalidMgic = -1U;\n\n private:\n HnswSparseEntity::Pointer entity_;\n HnswSparseDistCalculator dc_;\n bool debug_mode_{false};\n bool force_padding_topk_{false};\n uint32_t max_scan_num_{0};\n uint32_t max_scan_limit_{0};\n uint32_t min_scan_limit_{0};\n uint32_t reserve_max_doc_cnt_{kMinReserveDocCnt};\n uint32_t topk_{0};\n uint32_t group_topk_{0};\n uint32_t filter_mode_{VisitFilter::ByteMap};\n float negative_probability_{HnswSparseEntity::kDefaultBFNegativeProbability};\n uint32_t ef_{HnswSparseEntity::kDefaultEf};\n float max_scan_ratio_{HnswSparseEntity::kDefaultScanRatio};\n uint32_t magic_{0U};\n std::vector results_{};\n std::vector group_results_{};\n TopkHeap topk_heap_{};\n TopkHeap update_heap_{};\n std::vector level_topks_{};\n CandidateHeap candidates_{};\n VisitFilter visit_filter_{};\n uint32_t bruteforce_threshold_{};\n bool fetch_vector_{false};\n\n uint32_t group_num_{0};\n std::map group_topk_heaps_{};\n\n uint32_t type_{kUnknownContext};\n //! debug stats info\n uint32_t stats_get_neighbors_cnt_{0u};\n uint32_t stats_get_vector_cnt_{0u};\n uint32_t stats_visit_dup_cnt_{0u};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_dist_calculator.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"hnsw_sparse_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSparseDistCalculator {\n public:\n typedef std::shared_ptr Pointer;\n\n public:\n //! Constructor\n HnswSparseDistCalculator(const HnswSparseEntity *entity,\n const IndexMetric::Pointer &metric)\n : entity_(entity),\n distance_(metric->sparse_distance()),\n query_{nullptr},\n compare_cnt_(0) {}\n\n //! Constructor\n HnswSparseDistCalculator(const HnswSparseEntity *entity,\n const IndexMetric::Pointer &metric,\n const void *query)\n : entity_(entity),\n distance_(metric->sparse_distance()),\n query_(query),\n compare_cnt_(0) {}\n\n void update(const HnswSparseEntity *entity,\n const IndexMetric::Pointer &metric) {\n entity_ = entity;\n distance_ = metric->sparse_distance();\n }\n\n inline void update_distance(\n const IndexMetric::MatrixSparseDistance &distance) {\n distance_ = distance;\n }\n\n //! Reset query vector data\n inline void reset_query(const void *query) {\n error_ = false;\n query_ = query;\n }\n\n //! Returns distance\n inline dist_t dist(const void *sparse_data_lhs, const void *sparse_data_rhs) {\n float score{0.0f};\n\n if (ailego_unlikely(sparse_data_lhs == nullptr ||\n sparse_data_rhs == nullptr)) {\n // LOG_WARN(\"Nullptr of sparse vector. Return dense score only\");\n // error_ = true;\n return score;\n }\n\n distance_(sparse_data_lhs, sparse_data_rhs, &score);\n\n return score;\n }\n\n //! Returns distance between query and vec.\n inline dist_t dist(const void *vec) {\n compare_cnt_++;\n\n auto sparse_data = entity_->get_sparse_data_from_vector(vec);\n if (sparse_data.first == nullptr) {\n error_ = true;\n return 0.0f;\n }\n\n return dist(sparse_data.first, query_);\n }\n\n //! Return distance between query and node id.\n inline dist_t dist(node_id_t id) {\n compare_cnt_++;\n\n const void *feat = entity_->get_vector_meta(id);\n if (ailego_unlikely(feat == nullptr)) {\n LOG_ERROR(\"Get nullptr vector, id=%u\", id);\n error_ = true;\n return 0.0f;\n }\n\n auto sparse_data = entity_->get_sparse_data_from_vector(feat);\n if (sparse_data.first == nullptr) {\n error_ = true;\n return 0.0f;\n }\n\n return dist(sparse_data.first, query_);\n }\n\n //! Return dist node lhs between node rhs\n inline dist_t dist(node_id_t lhs, node_id_t rhs) {\n compare_cnt_++;\n\n const void *feat = entity_->get_vector_meta(lhs);\n const void *query = entity_->get_vector_meta(rhs);\n if (ailego_unlikely(feat == nullptr || query == nullptr)) {\n LOG_ERROR(\"Get nullptr vector\");\n error_ = true;\n return 0.0f;\n }\n\n auto feat_sparse_data = entity_->get_sparse_data_from_vector(feat);\n if (feat_sparse_data.first == nullptr) {\n error_ = true;\n return 0.0f;\n }\n\n auto query_sparse_data = entity_->get_sparse_data_from_vector(query);\n if (query_sparse_data.first == nullptr) {\n error_ = true;\n return 0.0f;\n }\n\n return dist(feat_sparse_data.first, query_sparse_data.first);\n }\n\n dist_t operator()(const void *vec) {\n return dist(vec);\n }\n\n dist_t operator()(id_t i) {\n return dist(i);\n }\n\n dist_t operator()(id_t lhs, id_t rhs) {\n return dist(lhs, rhs);\n }\n\n inline void clear() {\n compare_cnt_ = 0;\n error_ = false;\n }\n\n inline void clear_compare_cnt() {\n compare_cnt_ = 0;\n }\n\n inline bool error() const {\n return error_;\n }\n\n //! Get distances compute times\n inline uint32_t compare_cnt() const {\n return compare_cnt_;\n }\n\n private:\n HnswSparseDistCalculator(const HnswSparseDistCalculator &) = delete;\n HnswSparseDistCalculator &operator=(const HnswSparseDistCalculator &) =\n delete;\n\n private:\n const HnswSparseEntity *entity_;\n\n IndexMetric::MatrixSparseDistance distance_;\n\n const void *query_;\n\n uint32_t compare_cnt_; // record distance compute times\n bool error_{false};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \"hnsw_sparse_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nconst std::string HnswSparseEntity::kSparseGraphHeaderSegmentId =\n \"sparse_graph.header\";\nconst std::string HnswSparseEntity::kSparseGraphFeaturesSegmentId =\n \"sparse_graph.features\";\nconst std::string HnswSparseEntity::kSparseGraphKeysSegmentId =\n \"sparse_graph.keys\";\nconst std::string HnswSparseEntity::kSparseGraphNeighborsSegmentId =\n \"sparse_graph.neighbors\";\nconst std::string HnswSparseEntity::kSparseGraphOffsetsSegmentId =\n \"sparse_graph.offsets\";\nconst std::string HnswSparseEntity::kSparseGraphMappingSegmentId =\n \"sparse_graph.mapping\";\nconst std::string HnswSparseEntity::kSparseHnswHeaderSegmentId =\n \"sparse_hnsw.header\";\nconst std::string HnswSparseEntity::kSparseHnswNeighborsSegmentId =\n \"sparse_hnsw.neighbors\";\nconst std::string HnswSparseEntity::kSparseHnswOffsetsSegmentId =\n \"sparse_hnsw.offsets\";\nconst std::string HnswSparseEntity::kSparseGraphVectorsSegmentId =\n \"sparse_graph.vectors\";\nconst std::string HnswSparseEntity::kSparseGraphVectorMetaSegmentId =\n \"sparse_graph.vector_meta\";\n\nint HnswSparseEntity::CalcAndAddPadding(const IndexDumper::Pointer &dumper,\n size_t data_size,\n size_t *padding_size) {\n *padding_size = AlignSize(data_size) - data_size;\n if (*padding_size == 0) {\n return 0;\n }\n\n std::string padding(*padding_size, '\\0');\n if (dumper->write(padding.data(), *padding_size) != *padding_size) {\n LOG_ERROR(\"Append padding failed, size %lu\", *padding_size);\n return IndexError_WriteData;\n }\n return 0;\n}\n\nint64_t HnswSparseEntity::dump_segment(const IndexDumper::Pointer &dumper,\n const std::string &segment_id,\n const void *data, size_t size) const {\n size_t len = dumper->write(data, size);\n if (len != size) {\n LOG_ERROR(\"Dump segment %s data failed, expect: %lu, actual: %lu\",\n segment_id.c_str(), size, len);\n return IndexError_WriteData;\n }\n\n size_t padding_size = AlignSize(size) - size;\n if (padding_size > 0) {\n std::string padding(padding_size, '\\0');\n if (dumper->write(padding.data(), padding_size) != padding_size) {\n LOG_ERROR(\"Append padding failed, size %lu\", padding_size);\n return IndexError_WriteData;\n }\n }\n\n uint32_t crc = ailego::Crc32c::Hash(data, size);\n int ret = dumper->append(segment_id, size, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s meta failed, ret=%d\", segment_id.c_str(), ret);\n return ret;\n }\n\n return len + padding_size;\n}\n\nint64_t HnswSparseEntity::dump_header(const IndexDumper::Pointer &dumper,\n const HNSWSparseHeader &hd) const {\n //! dump basic graph header. header is aligned and does not need padding\n int64_t graph_hd_size = dump_segment(dumper, kSparseGraphHeaderSegmentId,\n &hd.graph, hd.graph.size);\n if (graph_hd_size < 0) {\n return graph_hd_size;\n }\n\n //! dump basic graph header. header is aligned and does not need padding\n int64_t hnsw_hd_size =\n dump_segment(dumper, kSparseHnswHeaderSegmentId, &hd.hnsw, hd.hnsw.size);\n if (hnsw_hd_size < 0) {\n return hnsw_hd_size;\n }\n\n return graph_hd_size + hnsw_hd_size;\n}\n\nvoid HnswSparseEntity::reshuffle_vectors(\n const std::function & /*get_level*/,\n std::vector * /*n2o_mapping*/,\n std::vector * /*o2n_mapping*/, key_t * /*keys*/) const {\n // TODO\n return;\n}\n\nint64_t HnswSparseEntity::dump_mapping_segment(\n const IndexDumper::Pointer &dumper, const key_t *keys) const {\n std::vector mapping(doc_cnt());\n\n std::iota(mapping.begin(), mapping.end(), 0U);\n std::sort(mapping.begin(), mapping.end(),\n [&](node_id_t i, node_id_t j) { return keys[i] < keys[j]; });\n\n size_t size = mapping.size() * sizeof(node_id_t);\n return dump_segment(dumper, kSparseGraphMappingSegmentId, mapping.data(),\n size);\n}\n\nint64_t HnswSparseEntity::dump_segments(\n const IndexDumper::Pointer &dumper, key_t *keys,\n const std::function &get_level) const {\n HNSWSparseHeader dump_hd(header());\n\n dump_hd.graph.node_size = sparse_meta_size();\n\n std::vector n2o_mapping; // map new id to origin id\n std::vector o2n_mapping; // map origin id to new id\n reshuffle_vectors(get_level, &n2o_mapping, &o2n_mapping, keys);\n if (!o2n_mapping.empty()) {\n dump_hd.hnsw.entry_point = o2n_mapping[entry_point()];\n }\n\n //! Dump header\n int64_t hd_size = dump_header(dumper, dump_hd);\n if (hd_size < 0) {\n return hd_size;\n }\n\n //! Dump vectors\n int64_t sparse_vector_meta_size =\n dump_sparse_vector_meta(dumper, n2o_mapping);\n if (sparse_vector_meta_size < 0) {\n return sparse_vector_meta_size;\n }\n\n int64_t sparse_vecs_size = dump_sparse_vector(dumper, n2o_mapping);\n if (sparse_vecs_size < 0) {\n return sparse_vecs_size;\n }\n\n //! Dump neighbors\n auto neighbors_size =\n dump_neighbors(dumper, get_level, n2o_mapping, o2n_mapping);\n if (neighbors_size < 0) {\n return neighbors_size;\n }\n //! free memory\n n2o_mapping = std::vector();\n o2n_mapping = std::vector();\n\n //! Dump keys\n size_t key_segment_size = doc_cnt() * sizeof(key_t);\n int64_t keys_size =\n dump_segment(dumper, kSparseGraphKeysSegmentId, keys, key_segment_size);\n if (keys_size < 0) {\n return keys_size;\n }\n\n //! Dump mapping\n int64_t mapping_size = dump_mapping_segment(dumper, keys);\n if (mapping_size < 0) {\n return mapping_size;\n }\n\n return hd_size + keys_size + sparse_vector_meta_size + sparse_vecs_size +\n neighbors_size + mapping_size;\n}\n\n\nint64_t HnswSparseEntity::dump_sparse_vector_meta(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping) const {\n const void *data = nullptr;\n uint32_t crc = 0U;\n size_t dump_size = 0UL;\n\n uint64_t sparse_data_offset = 0UL;\n uint64_t sparse_data_len = 0UL;\n\n //! dump vectors\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n data = get_vector_meta(reorder_mapping.empty() ? id : reorder_mapping[id]);\n if (ailego_unlikely(!data)) {\n return IndexError_ReadData;\n }\n\n const char *data_ptr = reinterpret_cast(data);\n sparse_data_len = *((uint32_t *)(data_ptr + sizeof(uint64_t)));\n\n size_t len = dumper->write(&sparse_data_offset, sizeof(uint64_t));\n if (len != sizeof(uint64_t)) {\n LOG_ERROR(\"Dump sparse data offset failed, write=%zu expect=%zu\", len,\n sizeof(uint64_t));\n return IndexError_WriteData;\n }\n\n crc = ailego::Crc32c::Hash(&sparse_data_offset, sizeof(uint64_t), crc);\n dump_size += sizeof(uint64_t);\n\n len = dumper->write(&sparse_data_len, sizeof(uint64_t));\n if (len != sizeof(uint64_t)) {\n LOG_ERROR(\"Dump sparse data len failed, write=%zu expect=%zu\", len,\n sizeof(uint64_t));\n return IndexError_WriteData;\n }\n\n crc = ailego::Crc32c::Hash(&sparse_data_len, sizeof(uint64_t), crc);\n dump_size += sizeof(uint64_t);\n\n sparse_data_offset += sparse_data_len;\n }\n\n int ret =\n dumper->append(kSparseGraphVectorMetaSegmentId, dump_size, 0UL, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump vectors segment meta failed, ret %d\", ret);\n return ret;\n }\n\n return dump_size;\n}\n\nint64_t HnswSparseEntity::dump_sparse_vector(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping) const {\n uint32_t crc = 0U;\n size_t data_size = 0UL;\n const void *data = nullptr;\n\n uint64_t sparse_data_len = 0UL;\n uint32_t sparse_chunk_index = 0U;\n uint32_t sparse_chunk_offset = 0U;\n\n //! dump vectors\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n data = get_vector_meta(reorder_mapping.empty() ? id : reorder_mapping[id]);\n if (ailego_unlikely(!data)) {\n return IndexError_ReadData;\n }\n\n const char *data_ptr = reinterpret_cast(data);\n\n sparse_data_len = *((uint32_t *)(data_ptr + sizeof(uint64_t)));\n\n uint64_t sparse_offset = *((uint64_t *)(data_ptr));\n\n const void *sparse = get_sparse_data(sparse_offset, sparse_data_len);\n if (ailego_unlikely(sparse == nullptr)) {\n LOG_ERROR(\"Get nullptr sparse, chunk index=%u, chunk offset=%u, len=%zu\",\n sparse_chunk_index, sparse_chunk_offset,\n (size_t)sparse_data_len);\n return IndexError_ReadData;\n }\n\n size_t len = dumper->write(sparse, sparse_data_len);\n if (len != sparse_data_len) {\n LOG_ERROR(\"Dump sparse data failed, write=%zu expect=%zu\", len,\n (size_t)sparse_data_len);\n return IndexError_WriteData;\n }\n\n crc = ailego::Crc32c::Hash(sparse, sparse_data_len, crc);\n data_size += sparse_data_len;\n }\n\n int ret = dumper->append(kSparseGraphVectorsSegmentId, data_size, 0UL, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump vectors segment meta failed, ret %d\", ret);\n return ret;\n }\n\n return data_size;\n}\n\nint64_t HnswSparseEntity::dump_graph_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const {\n std::vector graph_meta;\n graph_meta.reserve(doc_cnt());\n size_t offset = 0;\n uint32_t crc = 0;\n std::vector mapping(l0_neighbor_cnt());\n\n uint32_t min_neighbor_count = 10000;\n uint32_t max_neighbor_count = 0;\n size_t sum_neighbor_count = 0;\n\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n const Neighbors neighbors =\n get_neighbors(0, reorder_mapping.empty() ? id : reorder_mapping[id]);\n ailego_assert_with(!!neighbors.data, \"invalid neighbors\");\n ailego_assert_with(neighbors.size() <= l0_neighbor_cnt(),\n \"invalid neighbors\");\n\n uint32_t neighbor_count = neighbors.size();\n if (neighbor_count < min_neighbor_count) {\n min_neighbor_count = neighbor_count;\n }\n if (neighbor_count > max_neighbor_count) {\n max_neighbor_count = neighbor_count;\n }\n sum_neighbor_count += neighbor_count;\n\n graph_meta.emplace_back(offset, neighbor_count);\n size_t size = neighbors.size() * sizeof(node_id_t);\n const node_id_t *data = &neighbors[0];\n if (!neighbor_mapping.empty()) {\n for (node_id_t i = 0; i < neighbors.size(); ++i) {\n mapping[i] = neighbor_mapping[neighbors[i]];\n }\n data = mapping.data();\n }\n if (dumper->write(data, size) != size) {\n LOG_ERROR(\"Dump graph neighbor id=%u failed, size %lu\", id, size);\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(data, size, crc);\n offset += size;\n }\n\n uint32_t average_neighbor_count = 0;\n if (doc_cnt() > 0) {\n average_neighbor_count = sum_neighbor_count / doc_cnt();\n }\n LOG_INFO(\n \"Dump hnsw graph: min_neighbor_count[%u] max_neighbor_count[%u] \"\n \"average_neighbor_count[%u]\",\n min_neighbor_count, max_neighbor_count, average_neighbor_count);\n\n size_t padding_size = 0;\n int ret = CalcAndAddPadding(dumper, offset, &padding_size);\n if (ret != 0) {\n return ret;\n }\n ret =\n dumper->append(kSparseGraphNeighborsSegmentId, offset, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s failed, ret %d\",\n kSparseGraphNeighborsSegmentId.c_str(), ret);\n return ret;\n }\n\n //! dump level 0 neighbors meta\n auto len =\n dump_segment(dumper, kSparseGraphOffsetsSegmentId, graph_meta.data(),\n graph_meta.size() * sizeof(SparseGraphNeighborMeta));\n if (len < 0) {\n return len;\n }\n\n return len + offset + padding_size;\n}\n\nint64_t HnswSparseEntity::dump_upper_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::function &get_level,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const {\n std::vector hnsw_meta;\n hnsw_meta.reserve(doc_cnt());\n size_t offset = 0;\n uint32_t crc = 0;\n std::vector buffer(upper_neighbor_cnt() + 1);\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n node_id_t new_id = reorder_mapping.empty() ? id : reorder_mapping[id];\n auto level = get_level(new_id);\n if (level == 0) {\n hnsw_meta.emplace_back(0U, 0U);\n continue;\n }\n hnsw_meta.emplace_back(offset, level);\n ailego_assert_with((size_t)level < kMaxGraphLayers, \"invalid level\");\n for (level_t cur_level = 1; cur_level <= level; ++cur_level) {\n const Neighbors neighbors = get_neighbors(cur_level, new_id);\n ailego_assert_with(!!neighbors.data, \"invalid neighbors\");\n ailego_assert_with(neighbors.size() <= neighbor_cnt(cur_level),\n \"invalid neighbors\");\n memset(buffer.data(), 0, sizeof(node_id_t) * buffer.size());\n buffer[0] = neighbors.size();\n if (neighbor_mapping.empty()) {\n memcpy(&buffer[1], &neighbors[0], neighbors.size() * sizeof(node_id_t));\n } else {\n for (node_id_t i = 0; i < neighbors.size(); ++i) {\n buffer[i + 1] = neighbor_mapping[neighbors[i]];\n }\n }\n if (dumper->write(buffer.data(), sizeof(node_id_t) * buffer.size()) !=\n sizeof(node_id_t) * buffer.size()) {\n LOG_ERROR(\"Dump graph neighbor id=%u failed, size %lu\", id,\n sizeof(node_id_t) * buffer.size());\n return IndexError_WriteData;\n }\n crc = ailego::Crc32c::Hash(buffer.data(),\n sizeof(node_id_t) * buffer.size(), crc);\n offset += sizeof(node_id_t) * buffer.size();\n }\n }\n size_t padding_size = 0;\n int ret = CalcAndAddPadding(dumper, offset, &padding_size);\n if (ret != 0) {\n return ret;\n }\n\n ret =\n dumper->append(kSparseHnswNeighborsSegmentId, offset, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s failed, ret %d\",\n kSparseHnswNeighborsSegmentId.c_str(), ret);\n return ret;\n }\n\n //! dump level 0 neighbors meta\n auto len = dump_segment(dumper, kSparseHnswOffsetsSegmentId, hnsw_meta.data(),\n hnsw_meta.size() * sizeof(HnswSparseNeighborMeta));\n if (len < 0) {\n return len;\n }\n\n return len + offset + padding_size;\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nusing node_id_t = uint32_t;\nusing key_t = uint64_t;\nusing level_t = int32_t;\nusing dist_t = float;\nusing TopkHeap = ailego::KeyValueHeap;\nusing CandidateHeap =\n ailego::KeyValueHeap>;\nconstexpr node_id_t kInvalidNodeId = static_cast(-1);\nconstexpr key_t kInvalidKey = static_cast(-1);\nclass HnswSparseDistCalculator;\n\nstruct SparseGraphHeader {\n uint32_t size;\n uint32_t version;\n uint32_t graph_type;\n uint32_t doc_count;\n uint32_t vector_size;\n uint32_t node_size;\n uint32_t l0_neighbor_count;\n uint32_t prune_type;\n uint32_t prune_neighbor_count;\n uint32_t ef_construction;\n uint32_t options;\n uint32_t min_neighbor_count;\n uint32_t sparse_meta_size;\n uint32_t sparse_unit_size;\n uint32_t total_sparse_count;\n uint8_t reserved[868];\n};\n\nstatic_assert(sizeof(SparseGraphHeader) % 32 == 0,\n \"SparseGraphHeader must be aligned with 32 bytes\");\n\n//! Hnsw upper neighbor header\nstruct HnswSparseHeader {\n uint32_t size; // header size\n uint32_t revision; // current total docs of the graph\n uint32_t upper_neighbor_count;\n uint32_t ef_construction;\n uint32_t scaling_factor;\n uint32_t max_level;\n uint32_t entry_point;\n uint32_t options;\n uint8_t reserved[30];\n};\n\nstruct SparseData {\n public:\n SparseData() {};\n\n SparseData(uint32_t sparse_count, const uint32_t *sparse_indices,\n const void *sparse_vec)\n : count(sparse_count), indices(sparse_indices), vec(sparse_vec) {}\n\n uint32_t count{0};\n const uint32_t *indices{nullptr};\n const void *vec{nullptr};\n};\n\nstatic_assert(sizeof(HnswSparseHeader) % 32 == 0,\n \"SparseGraphHeader must be aligned with 32 bytes\");\n\n//! Hnsw common header and upper neighbor header\nstruct HNSWSparseHeader {\n HNSWSparseHeader() {\n clear();\n }\n\n HNSWSparseHeader(const HNSWSparseHeader &header) {\n memcpy(this, &header, sizeof(header));\n }\n\n HNSWSparseHeader &operator=(const HNSWSparseHeader &header) {\n memcpy(this, &header, sizeof(header));\n return *this;\n }\n\n //! Reset state to zero, and the params is untouched\n void inline reset() {\n graph.doc_count = 0U;\n hnsw.entry_point = kInvalidNodeId;\n hnsw.max_level = 0;\n graph.total_sparse_count = 0U;\n }\n\n //! Clear all fields to init value\n void inline clear() {\n memset(this, 0, sizeof(HNSWSparseHeader));\n hnsw.entry_point = kInvalidNodeId;\n graph.size = sizeof(SparseGraphHeader);\n hnsw.size = sizeof(HnswSparseHeader);\n graph.total_sparse_count = 0U;\n }\n\n size_t neighbor_cnt() const {\n return graph.l0_neighbor_count;\n }\n\n size_t upper_neighbor_cnt() const {\n return hnsw.upper_neighbor_count;\n }\n\n size_t vector_size() const {\n return graph.vector_size;\n }\n\n size_t ef_construction() const {\n return graph.ef_construction;\n }\n\n size_t scaling_factor() const {\n return hnsw.scaling_factor;\n }\n\n size_t neighbor_prune_cnt() const {\n return graph.prune_neighbor_count;\n }\n\n node_id_t entry_point() const {\n return hnsw.entry_point;\n }\n\n node_id_t doc_cnt() const {\n return graph.doc_count;\n }\n\n uint32_t total_sparse_count() const {\n return graph.total_sparse_count;\n }\n\n SparseGraphHeader graph;\n HnswSparseHeader hnsw;\n};\n\nstruct NeighborsHeader {\n uint32_t neighbor_cnt;\n node_id_t neighbors[0];\n};\n\nstruct Neighbors {\n Neighbors() : cnt{0}, data{nullptr} {}\n\n Neighbors(uint32_t cnt_in, const node_id_t *data_in)\n : cnt{cnt_in}, data{data_in} {}\n\n Neighbors(IndexStorage::MemoryBlock &&mem_block)\n : neighbor_block{std::move(mem_block)} {\n auto hd = reinterpret_cast(neighbor_block.data());\n cnt = hd->neighbor_cnt;\n data = hd->neighbors;\n }\n\n size_t size(void) const {\n return cnt;\n }\n\n const node_id_t &operator[](size_t idx) const {\n return data[idx];\n }\n\n uint32_t cnt;\n const node_id_t *data;\n IndexStorage::MemoryBlock neighbor_block;\n};\n\n//! level 0 neighbors offset\nstruct SparseGraphNeighborMeta {\n SparseGraphNeighborMeta(size_t o, size_t cnt)\n : offset(o), neighbor_cnt(cnt) {}\n\n uint64_t offset : 48;\n uint64_t neighbor_cnt : 16;\n};\n\n//! hnsw upper neighbors meta\nstruct HnswSparseNeighborMeta {\n HnswSparseNeighborMeta(size_t o, size_t l) : offset(o), level(l) {}\n\n uint64_t offset : 48; // offset = idx * upper neighors size\n uint64_t level : 16;\n};\n\nclass HnswSparseEntity {\n public:\n //! Constructor\n HnswSparseEntity() {}\n\n //! Constructor\n HnswSparseEntity(const HNSWSparseHeader &hd) {\n header_ = hd;\n }\n\n //! Destructor\n virtual ~HnswSparseEntity() {}\n\n //! HnswSparseEntity Pointerd;\n typedef std::shared_ptr Pointer;\n\n //! Get max neighbor size of graph level\n inline size_t neighbor_cnt(level_t level) const {\n return level == 0 ? header_.graph.l0_neighbor_count\n : header_.hnsw.upper_neighbor_count;\n }\n\n //! get max neighbor size of graph level 0\n inline size_t l0_neighbor_cnt() const {\n return header_.graph.l0_neighbor_count;\n }\n\n //! get min neighbor size of graph\n inline size_t min_neighbor_cnt() const {\n return header_.graph.min_neighbor_count;\n }\n\n //! get upper neighbor size of graph level other than 0\n inline size_t upper_neighbor_cnt() const {\n return header_.hnsw.upper_neighbor_count;\n }\n\n //! Get current total doc of the hnsw graph\n inline node_id_t *mutable_doc_cnt() {\n return &header_.graph.doc_count;\n }\n\n inline node_id_t doc_cnt() const {\n return header_.graph.doc_count;\n }\n\n inline uint32_t *mutable_total_sparse_count() {\n return &header_.graph.total_sparse_count;\n }\n\n uint32_t total_sparse_count() const {\n return header_.graph.total_sparse_count;\n }\n\n //! Get hnsw graph scaling params\n inline size_t scaling_factor() const {\n return header_.hnsw.scaling_factor;\n }\n\n //! Get prune_size\n inline size_t prune_cnt() const {\n return header_.graph.prune_neighbor_count;\n }\n\n //! Current entity of top level graph\n inline node_id_t entry_point() const {\n return header_.hnsw.entry_point;\n }\n\n //! Current max graph level\n inline level_t cur_max_level() const {\n return header_.hnsw.max_level;\n }\n\n //! Retrieve index vector size\n size_t vector_size() const {\n return header_.graph.vector_size;\n }\n\n //! Retrieve node size\n size_t node_size() const {\n return header_.graph.node_size;\n }\n\n //! Retrieve ef constuction\n size_t ef_construction() const {\n return header_.graph.ef_construction;\n }\n\n //! Retrieve sparse meta size\n size_t sparse_meta_size() const {\n return header_.graph.sparse_meta_size;\n }\n\n //! Retrieve sparse unit size\n size_t sparse_unit_size() const {\n return header_.graph.sparse_unit_size;\n }\n\n void set_vector_size(size_t size) {\n header_.graph.vector_size = size;\n }\n\n void set_prune_cnt(size_t v) {\n header_.graph.prune_neighbor_count = v;\n }\n\n void set_scaling_factor(size_t val) {\n header_.hnsw.scaling_factor = val;\n }\n\n void set_l0_neighbor_cnt(size_t cnt) {\n header_.graph.l0_neighbor_count = cnt;\n }\n\n void set_min_neighbor_cnt(size_t cnt) {\n header_.graph.min_neighbor_count = cnt;\n }\n\n void set_upper_neighbor_cnt(size_t cnt) {\n header_.hnsw.upper_neighbor_count = cnt;\n }\n\n void set_ef_construction(size_t ef) {\n header_.graph.ef_construction = ef;\n }\n\n void set_sparse_meta_size(size_t size) {\n header_.graph.sparse_meta_size = size;\n }\n\n void set_sparse_unit_size(size_t size) {\n header_.graph.sparse_unit_size = size;\n }\n\n protected:\n inline const HNSWSparseHeader &header() const {\n return header_;\n }\n\n inline HNSWSparseHeader *mutable_header() {\n return &header_;\n }\n\n inline size_t header_size() const {\n return sizeof(header_);\n }\n\n void set_node_size(size_t size) {\n header_.graph.node_size = size;\n }\n\n //! Dump all segment by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_segments(\n const IndexDumper::Pointer &dumper, key_t *keys,\n const std::function &get_level) const;\n\n private:\n //! dump mapping segment, for get_vector_by_key in provider\n int64_t dump_mapping_segment(const IndexDumper::Pointer &dumper,\n const key_t *keys) const;\n\n //! dump hnsw head by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_header(const IndexDumper::Pointer &dumper,\n const HNSWSparseHeader &hd) const;\n\n //! dump vectors by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_sparse_vector_meta(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping) const;\n\n //! dump sparse vectors by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_sparse_vector(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping) const;\n\n //! dump hnsw neighbors by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_neighbors(const IndexDumper::Pointer &dumper,\n const std::function &get_level,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const {\n auto len1 = dump_graph_neighbors(dumper, reorder_mapping, neighbor_mapping);\n if (len1 < 0) {\n return len1;\n }\n auto len2 = dump_upper_neighbors(dumper, get_level, reorder_mapping,\n neighbor_mapping);\n if (len2 < 0) {\n return len2;\n }\n\n return len1 + len2;\n }\n\n //! dump segment by dumper\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_segment(const IndexDumper::Pointer &dumper,\n const std::string &segment_id, const void *data,\n size_t size) const;\n\n //! Dump level 0 neighbors\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_graph_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const;\n\n //! Dump upper level neighbors\n //! Return dump size if success, errno(<0) in failure\n int64_t dump_upper_neighbors(\n const IndexDumper::Pointer &dumper,\n const std::function &get_level,\n const std::vector &reorder_mapping,\n const std::vector &neighbor_mapping) const;\n\n public:\n //! Cleanup the entity\n virtual int cleanup(void) {\n header_.clear();\n return 0;\n }\n\n //! Make a copy of searcher entity, to support thread-safe operation.\n //! The segment in container cannot be read concurrenly\n virtual const HnswSparseEntity::Pointer clone() const {\n LOG_ERROR(\"Update neighbors not implemented\");\n return HnswSparseEntity::Pointer();\n }\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const = 0;\n\n //! Get vector feature data by key\n virtual const void *get_vector_meta(node_id_t id) const = 0;\n\n virtual int get_vector_meta(const node_id_t id,\n IndexStorage::MemoryBlock &block) const = 0;\n\n //! Get vectors feature data by keys\n virtual int get_vector_metas(const node_id_t *ids, uint32_t count,\n const void **vecs) const = 0;\n virtual int get_vector_metas(\n const node_id_t *ids, uint32_t count,\n std::vector &block_vecs) const = 0;\n\n //! Retrieve a sparse vector using a primary key\n virtual int get_sparse_vector_by_key(\n uint64_t /*key*/, uint32_t * /*sparse_count*/,\n std::string * /*sparse_indices_buffer*/,\n std::string * /*sparse_values_buffer*/) const {\n LOG_ERROR(\"get sparse vector not implemented\");\n return IndexError_NotImplemented;\n }\n\n //! Retrieve a sparse vector using a primary key\n virtual int get_sparse_vector_by_id(\n node_id_t /*id*/, uint32_t * /*sparse_count*/,\n std::string * /*sparse_indices_buffer*/,\n std::string * /*sparse_values_buffer*/) const {\n LOG_ERROR(\"get sparse vector not implemented\");\n return IndexError_NotImplemented;\n }\n\n //! Get vector sparse feature data by chunk index and offset\n virtual const void *get_sparse_data(uint64_t offset, uint32_t len) const = 0;\n\n //! Get sparse data from id\n virtual const void *get_sparse_data(node_id_t id) const = 0;\n\n virtual int get_sparse_data(uint64_t offset, uint32_t len,\n IndexStorage::MemoryBlock &block) const = 0;\n\n virtual int get_sparse_data(const node_id_t id,\n IndexStorage::MemoryBlock &block) const = 0;\n\n //! Get sparse data from vector\n virtual std::pair get_sparse_data_from_vector(\n const void *vec) const = 0;\n virtual int get_sparse_data_from_vector(const void *vec,\n IndexStorage::MemoryBlock &block,\n int &sparse_length) const = 0;\n\n //! Get the node id's neighbors on graph level\n //! Note: the neighbors cannot be modified, using the following\n //! method to get WritableNeighbors if want to\n virtual const Neighbors get_neighbors(level_t level, node_id_t id) const = 0;\n\n //! Add vector and key to hnsw entity, and local id will be saved in id\n virtual int add_vector(level_t /*level*/, key_t /*key*/,\n const std::string & /*vec*/, uint32_t /*sparse_count*/,\n node_id_t * /*id*/) {\n return IndexError_NotImplemented;\n }\n\n virtual int add_vector(level_t /*level*/, key_t /*key*/,\n const uint32_t /*sparse_count*/,\n const uint32_t * /*sparse_indices*/,\n const void * /*sparse_vec*/, node_id_t * /*id*/) {\n return IndexError_NotImplemented;\n }\n\n //! Add vector and id\n virtual int add_vector_with_id(level_t /*level*/, node_id_t /*id*/,\n const std::string & /*vec*/,\n uint32_t /*sparse_count*/) {\n return IndexError_NotImplemented;\n }\n\n virtual int update_neighbors(\n level_t /*level*/, node_id_t /*id*/,\n const std::vector> & /*neighbors*/) {\n LOG_ERROR(\"Update neighbors dense not implemented\");\n\n return 0;\n }\n\n //! Append neighbor_id to node id neighbors on level, size is the current\n //! neighbors size. Notice: the caller must be ensure the neighbors not full\n virtual void add_neighbor(level_t /*level*/, node_id_t /*id*/,\n uint32_t /*size*/, node_id_t /*neighbor_id*/) {\n LOG_ERROR(\"Add neighbor not implemented\");\n }\n\n //! Update entry point and max level\n virtual void update_ep_and_level(node_id_t ep, level_t level) {\n header_.hnsw.entry_point = ep;\n header_.hnsw.max_level = level;\n }\n\n virtual int load(const IndexStorage::Pointer & /*container*/,\n bool /*check_crc*/) {\n LOG_ERROR(\"Load not implemented\");\n return IndexError_NotImplemented;\n }\n\n virtual int dump(const IndexDumper::Pointer & /*dumper*/) {\n LOG_ERROR(\"Dump not implemented\");\n return IndexError_NotImplemented;\n }\n\n static int CalcAndAddPadding(const IndexDumper::Pointer &dumper,\n size_t data_size, size_t *padding_size);\n\n protected:\n static inline size_t AlignSize(size_t size) {\n return (size + 0x1F) & (~0x1F);\n }\n\n static inline size_t AlignPageSize(size_t size) {\n size_t page_mask = ailego::MemoryHelper::PageSize() - 1;\n return (size + page_mask) & (~page_mask);\n }\n\n //! rearrange vectors to improve cache locality\n void reshuffle_vectors(const std::function &get_level,\n std::vector *n2o_mapping,\n std::vector *o2n_mapping,\n key_t *keys) const;\n\n public:\n const static std::string kSparseGraphHeaderSegmentId;\n const static std::string kSparseGraphFeaturesSegmentId;\n const static std::string kSparseGraphKeysSegmentId;\n const static std::string kSparseGraphNeighborsSegmentId;\n const static std::string kSparseGraphOffsetsSegmentId;\n const static std::string kSparseGraphMappingSegmentId;\n const static std::string kSparseHnswHeaderSegmentId;\n const static std::string kSparseHnswNeighborsSegmentId;\n const static std::string kSparseHnswOffsetsSegmentId;\n const static std::string kSparseGraphVectorsSegmentId;\n const static std::string kSparseGraphVectorMetaSegmentId;\n\n constexpr static uint32_t kRevision = 0U;\n constexpr static size_t kMaxGraphLayers = 15;\n constexpr static uint32_t kDefaultEfConstruction = 500;\n constexpr static uint32_t kDefaultEf = 500;\n constexpr static uint32_t kDefaultUpperMaxNeighborCnt = 50; // M of HNSW\n constexpr static uint32_t kDefaultL0MaxNeighborCnt = 100;\n constexpr static uint32_t kMaxNeighborCnt = 65535;\n constexpr static float kDefaultScanRatio = 0.1f;\n constexpr static uint32_t kDefaultMinScanLimit = 10000;\n constexpr static uint32_t kDefaultMaxScanLimit =\n std::numeric_limits::max();\n constexpr static float kDefaultBFNegativeProbability = 0.001f;\n constexpr static uint32_t kDefaultScalingFactor = 50U;\n constexpr static uint32_t kDefaultBruteForceThreshold = 1000U;\n constexpr static uint32_t kDefaultDocsHardLimit = 1 << 30U; // 1 billion\n constexpr static float kDefaultDocsSoftLimitRatio = 0.9f;\n constexpr static size_t kMaxChunkSize = 0xFFFFFFFF;\n constexpr static size_t kDefaultChunkSize = 2UL * 1024UL * 1024UL;\n constexpr static size_t kDefaultMaxChunkCnt = 50000UL;\n constexpr static float kDefaultNeighborPruneMultiplier =\n 1.0f; // prune_cnt = upper_max_neighbor_cnt * multiplier\n constexpr static float kDefaultL0MaxNeighborCntMultiplier =\n 2.0f; // l0_max_neighbor_cnt = upper_max_neighbor_cnt * multiplier\n\n constexpr static uint32_t kSparseMetaSize = 2u * sizeof(uint64_t);\n constexpr static float kDefaultSparseNeighborRatio = 0.5f;\n constexpr static uint32_t kSparseMaxDimSize = 16384;\n constexpr static float kDefaultQueryFilteringRatio = 0.0f; // turn off\n\n protected:\n HNSWSparseHeader header_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_index_hash.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"hnsw_sparse_chunk.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! Persistent hashmap implement through open addressing algorithm\ntemplate ::value>::type>\nclass HnswSparseIndexHashMap {\n using key_type = Key;\n using val_type = Val;\n\n struct Iterator {\n key_type first;\n val_type second;\n };\n typedef Iterator *iterator;\n typedef Iterator Item;\n typedef const Iterator *const_iterator;\n\n class Slot {\n public:\n Slot(SparseChunk::Pointer &&chunk, const void *data)\n : chunk_(std::move(chunk)),\n items_(reinterpret_cast(data)) {}\n //! Return a empty loc or the key item loc\n\n Slot(SparseChunk::Pointer &&chunk, IndexStorage::MemoryBlock &&mem_block)\n : chunk_(std::move(chunk)), items_block_(std::move(mem_block)) {\n items_ = reinterpret_cast(items_block_.data());\n }\n const_iterator find(key_type key, uint32_t max_items, uint32_t mask) const {\n auto it = &items_[key & mask];\n for (auto i = 0U; i < max_items; ++i) {\n if (it->first == key || it->second == EmptyVal) {\n // LOG_DEBUG(\"i=%u\", i);\n return it;\n }\n ++it;\n if (it == &items_[max_items]) {\n it = &items_[0];\n }\n }\n return nullptr;\n }\n\n bool update(const_iterator it) {\n uint32_t offset = reinterpret_cast(it) -\n reinterpret_cast(&items_[0]);\n if (ailego_unlikely(chunk_->write(offset, it, sizeof(Item)) !=\n sizeof(Item))) {\n LOG_ERROR(\"Chunk write failed\");\n return false;\n }\n return true;\n }\n\n private:\n SparseChunk::Pointer chunk_{};\n const Item *items_{nullptr}; // point to chunk data\n IndexStorage::MemoryBlock items_block_{};\n };\n\n public:\n //! Init the hash\n //! broker the index allocator\n //! chunk_size the size of per chunk allocated, actual size may greater\n //! factor factor = 1/ratio, ratio is the probability of a squence\n //! number inserted to this container\n //! max the max number key can be inserted\n //! expansion_ratio memory expansion ratio\n int init(SparseChunkBroker::Pointer &broker, uint32_t chunk_size,\n uint32_t factor, size_t max, float expansion_ratio) {\n ailego_assert_with(expansion_ratio > 1.0f, \"ratio must > 1.0f\");\n broker_ = broker;\n\n size_t items = std::ceil(chunk_size * 1.0f / sizeof(Item));\n slot_items_ = 1UL << static_cast((std::ceil(std::log2(items))));\n size_t range = slot_items_ * factor / expansion_ratio;\n mask_bits_ = std::floor(std::log2(range));\n range = 1UL << mask_bits_;\n size_t max_slots = std::ceil(max * 1.0f / range);\n slots_.reserve(max_slots);\n slot_loc_mask_ = slot_items_ - 1U;\n\n int ret = load();\n if (ret != 0) {\n return ret;\n }\n\n LOG_DEBUG(\n \"HnswIndexHash init, chunkSize=%u factor=%u max=%zu \"\n \"ratio=%f slotItems=%u maxSlots=%zu maskBits=%u \"\n \"range=%zu\",\n chunk_size, factor, max, expansion_ratio, slot_items_, max_slots,\n mask_bits_, range);\n\n return 0;\n }\n\n int cleanup(void) {\n broker_.reset();\n slots_.clear();\n slots_.shrink_to_fit();\n mask_bits_ = 0U;\n slot_items_ = 0U;\n slot_loc_mask_ = 0U;\n\n return 0;\n }\n\n const_iterator end(void) const {\n return nullptr;\n }\n\n const_iterator find(const key_type key) const {\n auto idx = key >> mask_bits_;\n if (idx >= slots_.size()) {\n return end();\n }\n auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_);\n return it && it->second != EmptyVal ? it : nullptr;\n }\n\n bool insert(key_type key, val_type val) {\n auto idx = key >> mask_bits_;\n if (idx >= slots_.size()) {\n if (ailego_unlikely(idx >= slots_.capacity())) {\n LOG_ERROR(\"no space to insert\");\n return false;\n }\n for (auto i = slots_.size(); i <= idx; ++i) {\n if (ailego_unlikely(!alloc_slot(i))) {\n return false;\n }\n }\n }\n auto it = slots_[idx].find(key, slot_items_, slot_loc_mask_);\n if (ailego_unlikely(it == nullptr)) {\n LOG_ERROR(\"no space to insert\");\n return false;\n }\n\n //! TODO: write memory is ok?\n const_cast(it)->first = key;\n const_cast(it)->second = val;\n\n return slots_[idx].update(it);\n }\n\n private:\n bool alloc_slot(size_t idx) {\n ailego_assert_with(idx == slots_.size(), \"invalid idx\");\n\n size_t size = slot_items_ * sizeof(Item);\n auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX,\n idx, size);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return false;\n }\n SparseChunk::Pointer chunk = p.second;\n if (ailego_unlikely(chunk->resize(size) != size)) {\n LOG_ERROR(\"Chunk resize failed, size=%zu\", size);\n return false;\n }\n //! Read the whole data to memory\n IndexStorage::MemoryBlock data_block;\n if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {\n LOG_ERROR(\"Chunk read failed, size=%zu\", size);\n return false;\n }\n\n slots_.emplace_back(std::move(chunk), std::move(data_block));\n return true;\n }\n\n int load(void) {\n size_t slots_cnt =\n broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX);\n for (size_t i = 0UL; i < slots_cnt; ++i) {\n auto chunk =\n broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_NEIGHBOR_INDEX, i);\n if (!chunk) {\n LOG_ERROR(\"Get chunk failed, seq=%zu\", i);\n return IndexError_InvalidFormat;\n }\n size_t size = sizeof(Item) * slot_items_;\n if (chunk->data_size() < size) {\n LOG_ERROR(\n \"Hash params may be mismatch, seq=%zu, data_size=%zu \"\n \"expect=%zu\",\n i, chunk->data_size(), size);\n return IndexError_InvalidFormat;\n }\n //! Read the whole data to memory\n IndexStorage::MemoryBlock data_block;\n if (ailego_unlikely(chunk->read(0U, data_block, size) != size)) {\n LOG_ERROR(\"Chunk read failed, size=%zu\", size);\n return false;\n }\n slots_.emplace_back(std::move(chunk), std::move(data_block));\n }\n return 0;\n }\n\n private:\n SparseChunkBroker::Pointer broker_{}; // chunk broker\n std::vector slots_{};\n uint32_t mask_bits_{0U};\n uint32_t slot_items_{}; // must be a power of 2\n uint32_t slot_loc_mask_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_index_provider.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"hnsw_sparse_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSparseIndexProvider : public IndexSparseProvider {\n public:\n HnswSparseIndexProvider(const IndexMeta &meta,\n const HnswSparseEntity::Pointer &entity,\n const std::string &owner)\n : meta_(meta), entity_(entity), owner_class_(owner) {}\n\n HnswSparseIndexProvider(const HnswSparseIndexProvider &) = delete;\n HnswSparseIndexProvider &operator=(const HnswSparseIndexProvider &) = delete;\n\n public:\n //! Create a new iterator\n IndexSparseProvider::Iterator::Pointer create_iterator(void) override {\n return IndexSparseProvider::Iterator::Pointer(new (std::nothrow)\n Iterator(entity_));\n }\n\n //! Retrieve count of vectors\n size_t count(void) const override {\n return entity_->doc_cnt();\n }\n\n size_t total_sparse_count(void) const override {\n return entity_->total_sparse_count();\n }\n\n //! Retrieve type of vector\n IndexMeta::DataType data_type(void) const override {\n return meta_.data_type();\n }\n\n //! Retrieve a vector using a primary key\n int get_sparse_vector(uint64_t key, uint32_t *sparse_count,\n std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override {\n return entity_->get_sparse_vector_by_key(\n key, sparse_count, sparse_indices_buffer, sparse_values_buffer);\n }\n\n //! Retrieve the owner class\n const std::string &owner_class(void) const override {\n return owner_class_;\n }\n\n private:\n class Iterator : public IndexSparseProvider::Iterator {\n public:\n Iterator(const HnswSparseEntity::Pointer &entity)\n : entity_(entity), cur_id_(0U), valid_(false) {\n const void *sparse_data = entity_->get_sparse_data(cur_id_);\n if (sparse_data != nullptr) {\n valid_ = true;\n\n sparse_indices_buffer_.clear();\n sparse_data_buffer_.clear();\n\n SparseUtility::ReverseSparseFormat(\n sparse_data, &sparse_count_, &sparse_indices_buffer_,\n &sparse_data_buffer_, entity_->sparse_unit_size());\n }\n }\n\n //! Retrieve sparse count\n virtual uint32_t sparse_count() const override {\n return sparse_count_;\n }\n\n //! Retrieve sparse indices\n virtual const uint32_t *sparse_indices() const override {\n return reinterpret_cast(sparse_indices_buffer_.data());\n }\n\n //! Retrieve sparse data\n virtual const void *sparse_data() const override {\n return reinterpret_cast(sparse_data_buffer_.data());\n }\n\n //! Test if the iterator is valid\n virtual bool is_valid(void) const override {\n return cur_id_ < entity_->doc_cnt() && valid_;\n }\n\n //! Retrieve primary key\n virtual uint64_t key(void) const override {\n return entity_->get_key(cur_id_);\n }\n\n //! Next iterator\n virtual void next(void) override {\n cur_id_ = get_next_valid_id(cur_id_ + 1);\n\n if (cur_id_ < entity_->doc_cnt()) {\n const void *sparse_data = entity_->get_sparse_data(cur_id_);\n if (sparse_data != nullptr) {\n valid_ = true;\n\n sparse_indices_buffer_.clear();\n sparse_data_buffer_.clear();\n\n SparseUtility::ReverseSparseFormat(\n sparse_data, &sparse_count_, &sparse_indices_buffer_,\n &sparse_data_buffer_, entity_->sparse_unit_size());\n } else {\n valid_ = false;\n }\n }\n }\n\n //! Reset the iterator\n void reset(void) {\n cur_id_ = get_next_valid_id(0);\n const void *sparse_data = entity_->get_sparse_data(cur_id_);\n if (sparse_data != nullptr) {\n valid_ = true;\n\n SparseUtility::ReverseSparseFormat(\n sparse_data, &sparse_count_, &sparse_indices_buffer_,\n &sparse_data_buffer_, entity_->sparse_unit_size());\n }\n }\n\n private:\n node_id_t get_next_valid_id(node_id_t start_id) {\n for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) {\n if (entity_->get_key(i) != kInvalidNodeId) {\n return i;\n }\n }\n return kInvalidNodeId;\n }\n\n private:\n const HnswSparseEntity::Pointer entity_;\n node_id_t cur_id_;\n uint32_t sparse_count_{0};\n std::string sparse_indices_buffer_;\n std::string sparse_data_buffer_;\n bool valid_{false};\n };\n\n private:\n const IndexMeta &meta_;\n const HnswSparseEntity::Pointer entity_;\n const std::string owner_class_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_params.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\nstatic const std::string PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT(\n \"proxima.hnsw.sparse_builder.thread_count\");\nstatic const std::string PARAM_HNSW_SPARSE_BUILDER_MEMORY_QUOTA(\n \"proxima.hnsw.sparse_builder.memory_quota\");\nstatic const std::string PARAM_HNSW_SPARSE_BUILDER_EFCONSTRUCTION(\n \"proxima.hnsw.sparse_builder.efconstruction\");\nstatic const std::string PARAM_HNSW_SPARSE_BUILDER_SCALING_FACTOR(\n \"proxima.hnsw.sparse_builder.scaling_factor\");\nstatic const std::string PARAM_HNSW_SPARSE_BUILDER_CHECK_INTERVAL_SECS(\n \"proxima.hnsw.sparse_builder.check_interval_secs\");\nstatic const std::string PARAM_HNSW_SPARSE_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER(\n \"proxima.hnsw.sparse_builder.neighbor_prune_multiplier\");\nstatic const std::string PARAM_HNSW_SPARSE_BUILDER_MIN_NEIGHBOR_COUNT(\n \"proxima.hnsw.sparse_builder.min_neighbor_count\");\nstatic const std::string PARAM_HNSW_SPARSE_BUILDER_MAX_NEIGHBOR_COUNT(\n \"proxima.hnsw.sparse_builder.max_neighbor_count\");\nstatic const std::string\n PARAM_HNSW_SPARSE_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER(\n \"proxima.hnsw.sparse_builder.l0_max_neighbor_count_multiplier\");\n\nstatic const std::string PARAM_HNSW_SPARSE_SEARCHER_EF(\n \"proxima.hnsw.sparse_searcher.ef\");\nstatic const std::string PARAM_HNSW_SPARSE_SEARCHER_BRUTE_FORCE_THRESHOLD(\n \"proxima.hnsw.sparse_searcher.brute_force_threshold\");\nstatic const std::string PARAM_HNSW_SPARSE_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE(\n \"proxima.hnsw.sparse_searcher.neighbors_in_memory_enable\");\nstatic const std::string PARAM_HNSW_SPARSE_SEARCHER_MAX_SCAN_RATIO(\n \"proxima.hnsw.sparse_searcher.max_scan_ratio\");\nstatic const std::string PARAM_HNSW_SPARSE_SEARCHER_CHECK_CRC_ENABLE(\n \"proxima.hnsw.sparse_searcher.check_crc_enable\");\nstatic const std::string PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_ENABLE(\n \"proxima.hnsw.sparse_searcher.visit_bloomfilter_enable\");\nstatic const std::string\n PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB(\n \"proxima.hnsw.sparse_searcher.visit_bloomfilter_negative_prob\");\nstatic const std::string PARAM_HNSW_SPARSE_SEARCHER_FORCE_PADDING_RESULT_ENABLE(\n \"proxima.hnsw.sparse_searcher.force_padding_result_enable\");\nstatic const std::string PARAM_HNSW_SPARSE_SEARCHER_QUERY_FILTERING_RATIO(\n \"proxima.hnsw.sparse_searcher.query_filtering_ratio\");\n\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO(\n \"proxima.hnsw.sparse_streamer.max_scan_ratio\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT(\n \"proxima.hnsw.sparse_streamer.min_scan_limit\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT(\n \"proxima.hnsw.sparse_streamer.max_scan_limit\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_EF(\n \"proxima.hnsw.sparse_streamer.ef\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION(\n \"proxima.hnsw.sparse_streamer.efconstruction\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT(\n \"proxima.hnsw.sparse_streamer.max_neighbor_count\");\nstatic const std::string\n PARAM_HNSW_SPARSE_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER(\n \"proxima.hnsw.sparse_streamer.l0_max_neighbor_count_multiplier\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR(\n \"proxima.hnsw.sparse_streamer.scaling_factor\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD(\n \"proxima.hnsw.sparse_streamer.brute_force_threshold\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT(\n \"proxima.hnsw.sparse_streamer.docs_hard_limit\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT(\n \"proxima.hnsw.sparse_streamer.docs_soft_limit\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_MAX_INDEX_SIZE(\n \"proxima.hnsw.sparse_streamer.max_index_size\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_ENABLE(\n \"proxima.hnsw.sparse_streamer.visit_bloomfilter_enable\");\nstatic const std::string\n PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB(\n \"proxima.hnsw.sparse_streamer.visit_bloomfilter_negative_prob\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_CHECK_CRC_ENABLE(\n \"proxima.hnsw.sparse_streamer.check_crc_enable\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER(\n \"proxima.hnsw.sparse_streamer.neighbor_prune_multiplier\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_CHUNK_SIZE(\n \"proxima.hnsw.sparse_streamer.chunk_size\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_FILTER_SAME_KEY(\n \"proxima.hnsw.sparse_streamer.filter_same_key\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_GET_VECTOR_ENABLE(\n \"proxima.hnsw.sparse_streamer.get_vector_enable\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_MIN_NEIGHBOR_COUNT(\n \"proxima.hnsw.sparse_streamer.min_neighbor_count\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_FORCE_PADDING_RESULT_ENABLE(\n \"proxima.hnsw.sparse_streamer.force_padding_result_enable\");\nstatic const std::string PARAM_HNSW_SPARSE_STREAMER_QUERY_FILTERING_RATIO(\n \"proxima.hnsw.sparse_streamer.query_filtering_ratio\");\n\nstatic const std::string PARAM_HNSW_SPARSE_REDUCER_WORKING_PATH(\n \"proxima.hnsw.sparse_reducer.working_path\");\nstatic const std::string PARAM_HNSW_SPARSE_REDUCER_NUM_OF_ADD_THREADS(\n \"proxima.hnsw.sparse_reducer.num_of_add_threads\");\nstatic const std::string PARAM_HNSW_SPARSE_REDUCER_INDEX_NAME(\n \"proxima.hnsw.sparse_reducer.index_name\");\nstatic const std::string PARAM_HNSW_SPARSE_REDUCER_EFCONSTRUCTION(\n \"proxima.hnsw.sparse_reducer.efconstruction\");\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_sparse_searcher.h\"\n#include \"hnsw_sparse_algorithm.h\"\n#include \"hnsw_sparse_index_provider.h\"\n#include \"hnsw_sparse_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswSparseSearcher::HnswSparseSearcher() {}\n\nHnswSparseSearcher::~HnswSparseSearcher() {}\n\nint HnswSparseSearcher::init(const ailego::Params &search_params) {\n params_ = search_params;\n params_.get(PARAM_HNSW_SPARSE_SEARCHER_EF, &ef_);\n params_.get(PARAM_HNSW_SPARSE_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_);\n params_.get(PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_ENABLE,\n &bf_enabled_);\n params_.get(PARAM_HNSW_SPARSE_SEARCHER_CHECK_CRC_ENABLE, &check_crc_enabled_);\n params_.get(PARAM_HNSW_SPARSE_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE,\n &neighbors_in_memory_enabled_);\n params_.get(PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB,\n &bf_negative_probability_);\n params_.get(PARAM_HNSW_SPARSE_SEARCHER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n params_.get(PARAM_HNSW_SPARSE_SEARCHER_FORCE_PADDING_RESULT_ENABLE,\n &force_padding_topk_enabled_);\n\n query_filtering_enabled_ =\n params_.get(PARAM_HNSW_SPARSE_SEARCHER_QUERY_FILTERING_RATIO,\n &query_filtering_ratio_);\n\n if (ef_ == 0) {\n ef_ = HnswSparseEntity::kDefaultEf;\n }\n if (bf_negative_probability_ <= 0.0f || bf_negative_probability_ >= 1.0f) {\n LOG_ERROR(\n \"[%s] must be in range (0,1)\",\n PARAM_HNSW_SPARSE_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (query_filtering_enabled_ &&\n (query_filtering_ratio_ <= 0.0f || query_filtering_ratio_ >= 1.0f)) {\n LOG_ERROR(\"[%s] must be in range (0, 1)\",\n PARAM_HNSW_SPARSE_SEARCHER_QUERY_FILTERING_RATIO.c_str());\n return IndexError_InvalidArgument;\n }\n\n entity_.set_neighbors_in_memory(neighbors_in_memory_enabled_);\n\n state_ = STATE_INITED;\n\n LOG_DEBUG(\n \"Init params: ef=%u maxScanRatio=%f bfEnabled=%u checkCrcEnabled=%u \"\n \"neighborsInMemoryEnabled=%u bfNagtiveProb=%f bruteForceThreshold=%u \"\n \"forcePadding=%u filteringRatio=%f\",\n ef_, max_scan_ratio_, bf_enabled_, check_crc_enabled_,\n neighbors_in_memory_enabled_, bf_negative_probability_,\n bruteforce_threshold_, force_padding_topk_enabled_,\n query_filtering_ratio_);\n\n return 0;\n}\n\nvoid HnswSparseSearcher::print_debug_info() {\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n Neighbors neighbours = entity_.get_neighbors(0, id);\n std::cout << \"node: \" << id << \"; \";\n for (uint32_t i = 0; i < neighbours.size(); ++i) {\n std::cout << neighbours[i];\n\n if (i == neighbours.size() - 1) {\n std::cout << std::endl;\n } else {\n std::cout << \", \";\n }\n }\n }\n}\n\nint HnswSparseSearcher::cleanup() {\n LOG_INFO(\"Begin HnswSparseSearcher:cleanup\");\n\n metric_.reset();\n meta_.clear();\n stats_.clear_attributes();\n stats_.set_loaded_count(0UL);\n stats_.set_loaded_costtime(0UL);\n max_scan_ratio_ = HnswSparseEntity::kDefaultScanRatio;\n max_scan_num_ = 0U;\n ef_ = HnswSparseEntity::kDefaultEf;\n bf_enabled_ = false;\n bf_negative_probability_ = HnswSparseEntity::kDefaultBFNegativeProbability;\n bruteforce_threshold_ = HnswSparseEntity::kDefaultBruteForceThreshold;\n check_crc_enabled_ = false;\n neighbors_in_memory_enabled_ = false;\n entity_.cleanup();\n state_ = STATE_INIT;\n\n LOG_INFO(\"End HnswSparseSearcher:cleanup\");\n\n return 0;\n}\n\nint HnswSparseSearcher::load(IndexStorage::Pointer container,\n IndexMetric::Pointer metric) {\n if (state_ != STATE_INITED) {\n LOG_ERROR(\"Init the searcher first before load index\");\n return IndexError_Runtime;\n }\n\n LOG_INFO(\"Begin HnswSparseSearcher:load\");\n\n auto start_time = ailego::Monotime::MilliSeconds();\n\n int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to deserialize meta from container\");\n return ret;\n }\n\n ret = entity_.load(container, check_crc_enabled_);\n if (ret != 0) {\n LOG_ERROR(\"HnswSparseSearcher load index failed\");\n return ret;\n }\n\n alg_ = HnswSparseAlgorithm::UPointer(new HnswSparseAlgorithm(entity_));\n\n if (metric) {\n metric_ = metric;\n } else {\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"CreateMeasure failed, name: %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"IndexMetric init failed, ret=%d\", ret);\n return ret;\n }\n if (metric_->query_metric()) {\n metric_ = metric_->query_metric();\n }\n }\n\n // if (!metric_->is_matched(meta_)) {\n // LOG_ERROR(\"IndexMeasure not match index meta\");\n // return IndexError_Mismatch;\n // }\n\n max_scan_num_ = static_cast(max_scan_ratio_ * entity_.doc_cnt());\n max_scan_num_ = std::max(4096U, max_scan_num_);\n\n stats_.set_loaded_count(entity_.doc_cnt());\n stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time);\n state_ = STATE_LOADED;\n magic_ = IndexContext::GenerateMagic();\n\n LOG_INFO(\"End HnswSparseSearcher::load\");\n\n return 0;\n}\n\nint HnswSparseSearcher::unload() {\n LOG_INFO(\"HnswSparseSearcher unload index\");\n\n meta_.clear();\n entity_.cleanup();\n metric_.reset();\n max_scan_num_ = 0;\n stats_.set_loaded_count(0UL);\n stats_.set_loaded_costtime(0UL);\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint HnswSparseSearcher::update_context(HnswSparseContext *ctx) const {\n const HnswSparseEntity::Pointer entity = entity_.clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone search context entity\");\n return IndexError_Runtime;\n }\n ctx->set_max_scan_num(max_scan_num_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n\n return ctx->update_context(HnswSparseContext::kSparseSearcherContext, meta_,\n metric_, entity, magic_);\n}\n\n//! Similarity search with sparse inputs\nint HnswSparseSearcher::search_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n if (ailego_unlikely(!context)) {\n LOG_ERROR(\"The context is not created by this searcher\");\n return IndexError_Mismatch;\n }\n HnswSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswSparseContext failed\");\n return IndexError_Cast;\n }\n\n if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) {\n return search_bf_impl(sparse_count, sparse_indices, sparse_query, qmeta,\n count, context);\n }\n\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n int ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->resize_results(count);\n\n const uint32_t *sparse_indices_tmp = sparse_indices;\n const void *sparse_query_tmp = sparse_query;\n\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n std::string sparse_query_filtered_buffer;\n\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n if (query_filtering_enabled_) {\n if (!SparseUtility::FilterSparseQuery(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n qmeta.data_type(), entity_.sparse_unit_size(),\n query_filtering_ratio_, &sparse_query_filtered_buffer)) {\n LOG_ERROR(\"Hnsw filtering failed\");\n return IndexError_Runtime;\n }\n\n ctx->reset_query(sparse_query_filtered_buffer.data());\n } else {\n ctx->reset_query(sparse_query_buffer.data());\n }\n\n int ret = alg_->search(ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher fast search failed\");\n return ret;\n }\n\n if (query_filtering_enabled_) {\n ctx->reset_query(sparse_query_buffer.data());\n ctx->recal_topk_dist();\n }\n\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\n//! Similarity search with sparse inputs\nint HnswSparseSearcher::search_bf_impl(\n const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta, uint32_t count,\n IndexStreamer::Context::Pointer &context) const {\n if (ailego_unlikely(!context)) {\n LOG_ERROR(\"The context is not created by this searcher\");\n return IndexError_Mismatch;\n }\n HnswSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswSparseContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n int ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->resize_results(count);\n\n const uint32_t *sparse_indices_tmp = sparse_indices;\n const void *sparse_query_tmp = sparse_query;\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_Runtime;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n\n ctx->group_topk_heaps().clear();\n\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n dist_t dist = ctx->dist_calculator().dist(id);\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n } else {\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n\n ctx->topk_heap().clear();\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n dist_t dist = ctx->dist_calculator().dist(id);\n ctx->topk_heap().emplace(id, dist);\n }\n }\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\n//! Similarity search with sparse inputs\nint HnswSparseSearcher::search_bf_by_p_keys_impl(\n const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n if (ailego_unlikely(!context)) {\n LOG_ERROR(\"The context is not created by this searcher\");\n return IndexError_Mismatch;\n }\n\n if (ailego_unlikely(p_keys.size() != count)) {\n LOG_ERROR(\"The size of p_keys is not equal to count\");\n return IndexError_InvalidArgument;\n }\n\n HnswSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswSparseContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n int ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->resize_results(count);\n\n const uint32_t *sparse_indices_tmp = sparse_indices;\n const void *sparse_query_tmp = sparse_query;\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_Runtime;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n ctx->group_topk_heaps().clear();\n\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n dist_t dist = ctx->dist_calculator().dist(id);\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n } else {\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n ctx->topk_heap().clear();\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n dist_t dist = ctx->dist_calculator().dist(id);\n ctx->topk_heap().emplace(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nIndexSearcher::Context::Pointer HnswSparseSearcher::create_context() const {\n if (ailego_unlikely(state_ != STATE_LOADED)) {\n LOG_ERROR(\"Load the index first before create context\");\n return Context::Pointer();\n }\n const HnswSparseEntity::Pointer search_ctx_entity = entity_.clone();\n if (!search_ctx_entity) {\n LOG_ERROR(\"Failed to create search context entity\");\n return Context::Pointer();\n }\n HnswSparseContext *ctx =\n new (std::nothrow) HnswSparseContext(metric_, search_ctx_entity);\n if (ailego_unlikely(ctx == nullptr)) {\n LOG_ERROR(\"Failed to new HnswSparseContext\");\n return Context::Pointer();\n }\n ctx->set_ef(ef_);\n ctx->set_max_scan_num(max_scan_num_);\n uint32_t filter_mode =\n bf_enabled_ ? VisitFilter::BloomFilter : VisitFilter::ByteMap;\n ctx->set_filter_mode(filter_mode);\n ctx->set_filter_negative_probability(bf_negative_probability_);\n ctx->set_magic(magic_);\n ctx->set_force_padding_topk(force_padding_topk_enabled_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n if (ailego_unlikely(ctx->init(HnswSparseContext::kSparseSearcherContext)) !=\n 0) {\n LOG_ERROR(\"Init HnswSparseContext failed\");\n delete ctx;\n return Context::Pointer();\n }\n\n return Context::Pointer(ctx);\n}\n\nIndexSearcher::SparseProvider::Pointer\nHnswSparseSearcher::create_sparse_provider(void) const {\n LOG_DEBUG(\"HnswSparseSearcher create sparse provider\");\n\n auto entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"Clone HnswSparseEntity failed\");\n return SparseProvider::Pointer();\n }\n return SparseProvider::Pointer(new (std::nothrow) HnswSparseIndexProvider(\n meta_, entity, \"HnswSparseSearcher\"));\n}\n\nint HnswSparseSearcher::get_sparse_vector(\n uint64_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const {\n return entity_.get_sparse_vector_by_key(\n key, sparse_count, sparse_indices_buffer, sparse_values_buffer);\n}\n\nINDEX_FACTORY_REGISTER_SEARCHER(HnswSparseSearcher);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"hnsw_sparse_searcher_entity.h\"\n#include \"hnsw_sparse_streamer.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSparseSearcher : public IndexSearcher {\n public:\n using ContextPointer = IndexSearcher::Context::Pointer;\n\n public:\n HnswSparseSearcher(void);\n virtual ~HnswSparseSearcher(void);\n\n HnswSparseSearcher(const HnswSparseSearcher &) = delete;\n HnswSparseSearcher &operator=(const HnswSparseSearcher &) = delete;\n\n protected:\n //! Initialize Searcher\n int init(const ailego::Params ¶ms) override;\n\n //! Cleanup Searcher\n int cleanup(void) override;\n\n //! Load Index from storage\n int load(IndexStorage::Pointer container,\n IndexMetric::Pointer measure) override;\n\n //! Unload index from storage\n int unload(void) override;\n\n //! Similarity search with sparse inputs\n int search_impl(const uint32_t sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override {\n return search_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1,\n context);\n }\n\n //! Similarity search with sparse inputs\n int search_impl(const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const override;\n\n //! Similarity brute force search with sparse inputs\n int search_bf_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override {\n return search_bf_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1,\n context);\n }\n\n //! Similarity brute force search with sparse inputs\n int search_bf_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const override;\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta,\n ContextPointer &context) const override {\n return search_bf_by_p_keys_impl(&sparse_count, sparse_indices, sparse_query,\n p_keys, qmeta, 1, context);\n }\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const override;\n\n //! Fetch sparser vector by key\n int get_sparse_vector(uint64_t key, uint32_t *sparse_count,\n std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override;\n\n //! Create a searcher context\n ContextPointer create_context() const override;\n\n //! Create a new iterator\n IndexSearcher::SparseProvider::Pointer create_sparse_provider(\n void) const override;\n\n //! Retrieve statistics\n const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve meta of index\n const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n //! Retrieve params of index\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n void print_debug_info() override;\n\n private:\n //! To share ctx across streamer/searcher, we need to update the context for\n //! current streamer/searcher\n int update_context(HnswSparseContext *ctx) const;\n\n private:\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 };\n\n HnswSparseSearcherEntity entity_{};\n HnswSparseAlgorithm::UPointer alg_; // impl graph algorithm\n\n IndexMetric::Pointer metric_{};\n IndexMeta meta_{};\n ailego::Params params_{};\n Stats stats_;\n uint32_t ef_{HnswSparseEntity::kDefaultEf};\n uint32_t max_scan_num_{0U};\n uint32_t bruteforce_threshold_{HnswSparseEntity::kDefaultBruteForceThreshold};\n float max_scan_ratio_{HnswSparseEntity::kDefaultScanRatio};\n bool bf_enabled_{false};\n bool check_crc_enabled_{false};\n bool neighbors_in_memory_enabled_{false};\n bool force_padding_topk_enabled_{false};\n float bf_negative_probability_{\n HnswSparseEntity::kDefaultBFNegativeProbability};\n\n bool query_filtering_enabled_{false};\n float query_filtering_ratio_{HnswSparseEntity::kDefaultQueryFilteringRatio};\n\n uint32_t magic_{0U};\n\n State state_{STATE_INIT};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_sparse_searcher_entity.h\"\n#include \n#include \"utility/sparse_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswSparseSearcherEntity::HnswSparseSearcherEntity() {}\n\nint HnswSparseSearcherEntity::cleanup(void) {\n container_.reset();\n sparse_vector_meta_.reset();\n keys_.reset();\n neighbors_.reset();\n neighbors_meta_.reset();\n sparse_vectors_.reset();\n neighbors_in_memory_enabled_ = false;\n loaded_ = false;\n\n this->HnswSparseEntity::cleanup();\n\n return 0;\n}\n\nkey_t HnswSparseSearcherEntity::get_key(node_id_t id) const {\n const void *key;\n if (ailego_unlikely(keys_->read(id * sizeof(key_t), &key, sizeof(key_t)) !=\n sizeof(key_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidKey;\n }\n return *(reinterpret_cast(key));\n}\n\n//! Get vector local id by key\nnode_id_t HnswSparseSearcherEntity::get_id(key_t key) const {\n if (ailego_unlikely(!mapping_)) {\n LOG_ERROR(\"Index missing mapping segment\");\n return kInvalidNodeId;\n }\n\n //! Do binary search\n node_id_t start = 0UL;\n node_id_t end = doc_cnt();\n const void *data;\n node_id_t idx = 0u;\n while (start < end) {\n idx = start + (end - start) / 2;\n if (ailego_unlikely(\n mapping_->read(idx * sizeof(node_id_t), &data, sizeof(node_id_t)) !=\n sizeof(node_id_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidNodeId;\n }\n const key_t *mkey;\n node_id_t local_id = *reinterpret_cast(data);\n if (ailego_unlikely(keys_->read(local_id * sizeof(key_t),\n (const void **)(&mkey),\n sizeof(key_t)) != sizeof(key_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return kInvalidNodeId;\n }\n if (*mkey < key) {\n start = idx + 1;\n } else if (*mkey > key) {\n end = idx;\n } else {\n return local_id;\n }\n }\n return kInvalidNodeId;\n}\n\nint HnswSparseSearcherEntity::get_sparse_vector_by_key(\n key_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const {\n *sparse_count = 0;\n\n auto id = get_id(key);\n if (id == kInvalidNodeId) {\n return IndexError_NoExist;\n }\n\n const void *sparse_data = get_sparse_data(id);\n if (sparse_data == nullptr) {\n return IndexError_InvalidValue;\n }\n\n SparseUtility::ReverseSparseFormat(sparse_data, sparse_count,\n sparse_indices_buffer,\n sparse_values_buffer, sparse_unit_size());\n\n return 0;\n}\n\nconst void *HnswSparseSearcherEntity::get_vector_meta(node_id_t id) const {\n size_t read_size = sparse_meta_size();\n size_t offset = sparse_meta_size() * id;\n\n const void *vec;\n if (ailego_unlikely(sparse_vector_meta_->read(offset, &vec, read_size) !=\n read_size)) {\n LOG_ERROR(\"Read vector from segment failed\");\n return nullptr;\n }\n return vec;\n}\n\nint HnswSparseSearcherEntity::get_vector_meta(\n const node_id_t id, IndexStorage::MemoryBlock &block) const {\n const void *vec = get_vector_meta(id);\n block.reset((void *)vec);\n return 0;\n}\n\nint HnswSparseSearcherEntity::get_vector_metas(const node_id_t *ids,\n uint32_t count,\n const void **vecs) const {\n ailego_assert_with(count <= segment_datas_.size(), \"invalid count\");\n\n size_t read_size = sparse_meta_size();\n\n for (uint32_t i = 0; i < count; ++i) {\n segment_datas_[i].offset = sparse_meta_size() * ids[i];\n segment_datas_[i].length = read_size;\n\n ailego_assert_with(\n segment_datas_[i].offset < sparse_vector_meta_->data_size(),\n \"invalid offset\");\n }\n if (ailego_unlikely(!sparse_vector_meta_->read(&segment_datas_[0], count))) {\n LOG_ERROR(\"Read vectors from segment failed\");\n return IndexError_ReadData;\n }\n for (uint32_t i = 0; i < count; ++i) {\n vecs[i] = segment_datas_[i].data;\n }\n\n return 0;\n}\n\nint HnswSparseSearcherEntity::get_vector_metas(\n const node_id_t *ids, uint32_t count,\n std::vector &block_vecs) const {\n const void *vecs[count];\n get_vector_metas(ids, count, vecs);\n for (uint32_t i = 0; i < count; ++i) {\n block_vecs.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i]));\n }\n return 0;\n}\n\nconst Neighbors HnswSparseSearcherEntity::get_neighbors(level_t level,\n node_id_t id) const {\n if (level == 0) {\n if (neighbors_in_memory_enabled_) {\n auto hd = reinterpret_cast(\n fixed_neighbors_.get() + neighbors_size() * id);\n return {hd->neighbor_cnt, hd->neighbors};\n }\n\n const SparseGraphNeighborMeta *m;\n if (ailego_unlikely(\n neighbors_meta_->read(id * sizeof(SparseGraphNeighborMeta),\n (const void **)(&m),\n sizeof(SparseGraphNeighborMeta)) !=\n sizeof(SparseGraphNeighborMeta))) {\n LOG_ERROR(\"Read neighbors meta from segment failed\");\n return {0, nullptr};\n }\n\n const void *data;\n if (ailego_unlikely(neighbors_->read(m->offset, &data,\n m->neighbor_cnt * sizeof(node_id_t)) !=\n m->neighbor_cnt * sizeof(node_id_t))) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return {0, nullptr};\n }\n return {static_cast(m->neighbor_cnt),\n reinterpret_cast(data)};\n }\n\n //! Read level > 0 neighbors\n const HnswSparseNeighborMeta *m;\n if (ailego_unlikely(\n upper_neighbors_meta_->read(id * sizeof(HnswSparseNeighborMeta),\n (const void **)(&m),\n sizeof(HnswSparseNeighborMeta)) !=\n sizeof(HnswSparseNeighborMeta))) {\n LOG_ERROR(\"Read neighbors meta from segment failed\");\n return {0, nullptr};\n }\n\n ailego_assert_with(level <= m->level, \"invalid level\");\n size_t offset = m->offset + (level - 1) * upper_neighbors_size();\n ailego_assert_with(offset <= upper_neighbors_->data_size(), \"invalid offset\");\n const void *data;\n if (ailego_unlikely(\n upper_neighbors_->read(offset, &data, upper_neighbors_size()) !=\n upper_neighbors_size())) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return {0, nullptr};\n }\n\n auto hd = reinterpret_cast(data);\n return {hd->neighbor_cnt, hd->neighbors};\n}\n\nint HnswSparseSearcherEntity::load(const IndexStorage::Pointer &container,\n bool check_crc) {\n container_ = container;\n\n int ret = load_segments(check_crc);\n if (ret != 0) {\n return ret;\n }\n\n loaded_ = true;\n\n LOG_INFO(\n \"Index info: docCnt=%u entryPoint=%u maxLevel=%d efConstruct=%zu \"\n \"l0NeighborCnt=%zu upperNeighborCnt=%zu scalingFactor=%zu \"\n \"nodeSize=%zu sparesMetaSegmentSize=%zu keySegmentSize=%zu \"\n \"neighborsSegmentSize=%zu neighborsMetaSegmentSize=%zu \"\n \"sparseVectorSegmentSize=%zu\",\n doc_cnt(), entry_point(), cur_max_level(), ef_construction(),\n l0_neighbor_cnt(), upper_neighbor_cnt(), scaling_factor(), node_size(),\n sparse_vector_meta_->data_size(), keys_->data_size(),\n neighbors_->data_size(), neighbors_meta_->data_size(),\n sparse_vectors_->data_size());\n\n return 0;\n}\n\nint HnswSparseSearcherEntity::load_segments(bool check_crc) {\n //! load header\n const void *data = nullptr;\n HNSWSparseHeader hd;\n auto graph_hd_segment = container_->get(kSparseGraphHeaderSegmentId);\n if (!graph_hd_segment || graph_hd_segment->data_size() < sizeof(hd.graph)) {\n LOG_ERROR(\"Miss or invalid segment %s\",\n kSparseGraphHeaderSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n if (graph_hd_segment->read(0, reinterpret_cast(&data),\n sizeof(hd.graph)) != sizeof(hd.graph)) {\n LOG_ERROR(\"Read segment %s failed\", kSparseGraphHeaderSegmentId.c_str());\n return IndexError_ReadData;\n }\n memcpy(&hd.graph, data, sizeof(hd.graph));\n\n auto hnsw_hd_segment = container_->get(kSparseHnswHeaderSegmentId);\n if (!hnsw_hd_segment || hnsw_hd_segment->data_size() < sizeof(hd.hnsw)) {\n LOG_ERROR(\"Miss or invalid segment %s\", kSparseHnswHeaderSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n if (hnsw_hd_segment->read(0, reinterpret_cast(&data),\n sizeof(hd.hnsw)) != sizeof(hd.hnsw)) {\n LOG_ERROR(\"Read segment %s failed\", kSparseHnswHeaderSegmentId.c_str());\n return IndexError_ReadData;\n }\n memcpy(&hd.hnsw, data, sizeof(hd.hnsw));\n *mutable_header() = hd;\n segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt()));\n\n sparse_vector_meta_ = container_->get(kSparseGraphVectorMetaSegmentId);\n if (!sparse_vector_meta_) {\n LOG_ERROR(\"IndexStorage get segment %s failed\",\n kSparseGraphVectorMetaSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n\n keys_ = container_->get(kSparseGraphKeysSegmentId);\n if (!keys_) {\n LOG_ERROR(\"IndexStorage get segment %s failed\",\n kSparseGraphKeysSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n\n sparse_vectors_ = container_->get(kSparseGraphVectorsSegmentId);\n if (!sparse_vectors_) {\n LOG_ERROR(\"IndexStorage get segment %s failed\",\n kSparseGraphVectorsSegmentId.c_str());\n return IndexError_InvalidFormat;\n }\n\n neighbors_ = container_->get(kSparseGraphNeighborsSegmentId);\n if (!neighbors_ || (neighbors_->data_size() == 0 && doc_cnt() > 1)) {\n LOG_ERROR(\"IndexStorage get segment %s failed or empty\",\n kSparseGraphNeighborsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n neighbors_meta_ = container_->get(kSparseGraphOffsetsSegmentId);\n if (!neighbors_meta_ || neighbors_meta_->data_size() <\n sizeof(SparseGraphNeighborMeta) * doc_cnt()) {\n LOG_ERROR(\"IndexStorage get segment %s failed or invalid size\",\n kSparseGraphOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n upper_neighbors_ = container_->get(kSparseHnswNeighborsSegmentId);\n if (!upper_neighbors_ ||\n (upper_neighbors_->data_size() == 0 && cur_max_level() > 0)) {\n LOG_ERROR(\"IndexStorage get segment %s failed or empty\",\n kSparseHnswNeighborsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n upper_neighbors_meta_ = container_->get(kSparseHnswOffsetsSegmentId);\n if (!upper_neighbors_meta_ ||\n upper_neighbors_meta_->data_size() <\n sizeof(HnswSparseNeighborMeta) * doc_cnt()) {\n LOG_ERROR(\"IndexStorage get segment %s failed or invalid size\",\n kSparseHnswOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n mapping_ = container_->get(kSparseGraphMappingSegmentId);\n if (!mapping_ || mapping_->data_size() < sizeof(node_id_t) * doc_cnt()) {\n LOG_ERROR(\"IndexStorage get segment %s failed or invalid size\",\n kSparseGraphMappingSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (check_crc) {\n std::vector segments;\n segments.emplace_back(graph_hd_segment);\n segments.emplace_back(hnsw_hd_segment);\n segments.emplace_back(sparse_vector_meta_);\n segments.emplace_back(keys_);\n\n segments.emplace_back(sparse_vectors_);\n\n segments.emplace_back(neighbors_);\n segments.emplace_back(neighbors_meta_);\n segments.emplace_back(upper_neighbors_);\n segments.emplace_back(upper_neighbors_meta_);\n\n if (!do_crc_check(segments)) {\n LOG_ERROR(\"Check index crc failed, the index may broken\");\n return IndexError_Runtime;\n }\n }\n\n if (neighbors_in_memory_enabled_) {\n int ret = load_and_flat_neighbors();\n if (ret != 0) {\n return ret;\n }\n }\n\n return 0;\n}\n\nint HnswSparseSearcherEntity::load_and_flat_neighbors() {\n fixed_neighbors_.reset(\n new (std::nothrow) char[neighbors_size() * doc_cnt()]{},\n std::default_delete());\n if (!fixed_neighbors_) {\n LOG_ERROR(\"Malloc memory failed\");\n return IndexError_NoMemory;\n }\n\n //! Get a new segemnt to release the buffer after loading neighbors\n auto neighbors_meta = container_->get(kSparseGraphOffsetsSegmentId);\n if (!neighbors_meta) {\n LOG_ERROR(\"IndexStorage get segment graph.offsets failed\");\n return IndexError_InvalidArgument;\n }\n\n const SparseGraphNeighborMeta *neighbors_index = nullptr;\n if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index),\n neighbors_meta->data_size()) !=\n neighbors_meta->data_size()) {\n LOG_ERROR(\"Read segment %s data failed\",\n kSparseGraphOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n const char *neighbor_data;\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n size_t rd_size = neighbors_index[id].neighbor_cnt * sizeof(node_id_t);\n if (ailego_unlikely(\n neighbors_->read(neighbors_index[id].offset,\n reinterpret_cast(&neighbor_data),\n rd_size) != rd_size)) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return IndexError_ReadData;\n }\n // copy level 0 neighbors to fixed size neighbors memory\n char *dst = fixed_neighbors_.get() + neighbors_size() * id;\n *reinterpret_cast(dst) = neighbors_index[id].neighbor_cnt;\n memcpy(dst + sizeof(uint32_t), neighbor_data, rd_size);\n }\n\n return 0;\n}\n\nint HnswSparseSearcherEntity::get_fixed_neighbors(\n std::vector *fixed_neighbors) const {\n //! Get a new segemnt to release the buffer after loading neighbors\n auto neighbors_meta = container_->get(kSparseGraphOffsetsSegmentId);\n if (!neighbors_meta) {\n LOG_ERROR(\"IndexStorage get segment graph.offsets failed\");\n return IndexError_InvalidArgument;\n }\n\n const SparseGraphNeighborMeta *neighbors_index = nullptr;\n size_t meta_size = neighbors_meta->data_size();\n if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index),\n meta_size) != meta_size) {\n LOG_ERROR(\"Read segment %s data failed\",\n kSparseGraphOffsetsSegmentId.c_str());\n return IndexError_InvalidArgument;\n }\n\n size_t fixed_neighbor_cnt = l0_neighbor_cnt();\n fixed_neighbors->resize((fixed_neighbor_cnt + 1) * doc_cnt(), kInvalidNodeId);\n\n size_t neighbors_cnt_offset = fixed_neighbor_cnt * doc_cnt();\n size_t total_neighbor_cnt = 0;\n for (node_id_t id = 0; id < doc_cnt(); ++id) {\n size_t cur_neighbor_cnt = neighbors_index[id].neighbor_cnt;\n if (cur_neighbor_cnt == 0) {\n (*fixed_neighbors)[neighbors_cnt_offset + id] = 0;\n continue;\n }\n size_t rd_size = cur_neighbor_cnt * sizeof(node_id_t);\n const uint32_t *neighbors;\n if (neighbors_->read(neighbors_index[id].offset,\n reinterpret_cast(&neighbors),\n rd_size) != rd_size) {\n LOG_ERROR(\"Read neighbors from segment failed\");\n return IndexError_ReadData;\n }\n\n // copy level 0 neighbors to fixed size neighbors memory\n auto it = fixed_neighbors->begin() + id * fixed_neighbor_cnt;\n std::copy(neighbors, neighbors + cur_neighbor_cnt, it);\n\n (*fixed_neighbors)[neighbors_cnt_offset + id] = cur_neighbor_cnt;\n total_neighbor_cnt += cur_neighbor_cnt;\n }\n LOG_INFO(\"total neighbor cnt: %zu, average neighbor cnt: %zu\",\n total_neighbor_cnt, total_neighbor_cnt / doc_cnt());\n\n return 0;\n}\n\nbool HnswSparseSearcherEntity::do_crc_check(\n std::vector &segments) const {\n constexpr size_t blk_size = 4096;\n const void *data;\n for (auto &segment : segments) {\n size_t offset = 0;\n size_t rd_size;\n uint32_t crc = 0;\n while (offset < segment->data_size()) {\n size_t size = std::min(blk_size, segment->data_size() - offset);\n if ((rd_size = segment->read(offset, &data, size)) <= 0) {\n break;\n }\n offset += rd_size;\n crc = ailego::Crc32c::Hash(data, rd_size, crc);\n }\n if (crc != segment->data_crc()) {\n return false;\n }\n }\n return true;\n}\n\nconst HnswSparseEntity::Pointer HnswSparseSearcherEntity::clone() const {\n auto keys = keys_->clone();\n if (ailego_unlikely(!keys)) {\n LOG_ERROR(\"clone segment %s failed\", kSparseGraphKeysSegmentId.c_str());\n return HnswSparseEntity::Pointer();\n }\n\n auto mapping = mapping_->clone();\n if (ailego_unlikely(!mapping)) {\n LOG_ERROR(\"clone segment %s failed\", kSparseGraphMappingSegmentId.c_str());\n return HnswSparseEntity::Pointer();\n }\n\n auto sparse_vector_meta = sparse_vector_meta_->clone();\n if (ailego_unlikely(!sparse_vector_meta)) {\n LOG_ERROR(\"clone segment %s failed\",\n kSparseGraphVectorMetaSegmentId.c_str());\n return HnswSparseEntity::Pointer();\n }\n\n auto sparse_vectors = sparse_vectors_->clone();\n if (ailego_unlikely(!sparse_vectors)) {\n LOG_ERROR(\"clone segment %s failed\", kSparseGraphVectorsSegmentId.c_str());\n return HnswSparseEntity::Pointer();\n }\n\n auto neighbors = neighbors_->clone();\n if (ailego_unlikely(!neighbors)) {\n LOG_ERROR(\"clone segment %s failed\",\n kSparseGraphNeighborsSegmentId.c_str());\n return HnswSparseEntity::Pointer();\n }\n auto upper_neighbors = upper_neighbors_->clone();\n if (ailego_unlikely(!neighbors)) {\n LOG_ERROR(\"clone segment %s failed\", kSparseHnswNeighborsSegmentId.c_str());\n return HnswSparseEntity::Pointer();\n }\n auto neighbors_meta = neighbors_meta_->clone();\n if (ailego_unlikely(!neighbors_meta)) {\n LOG_ERROR(\"clone segment %s failed\", kSparseGraphOffsetsSegmentId.c_str());\n return HnswSparseEntity::Pointer();\n }\n auto upper_neighbors_meta = upper_neighbors_meta_->clone();\n if (ailego_unlikely(!upper_neighbors_meta)) {\n LOG_ERROR(\"clone segment %s failed\", kSparseHnswOffsetsSegmentId.c_str());\n return HnswSparseEntity::Pointer();\n }\n\n SegmentGroupParam neighbor_group{neighbors, neighbors_meta, upper_neighbors,\n upper_neighbors_meta};\n SegmentGroupParam dense_neighbor_group{nullptr, nullptr, nullptr, nullptr};\n SegmentGroupParam sparse_neighbor_group{nullptr, nullptr, nullptr, nullptr};\n\n HnswSparseSearcherEntity *entity = new (std::nothrow)\n HnswSparseSearcherEntity(header(), keys, mapping, neighbor_group,\n sparse_vector_meta, sparse_vectors,\n fixed_neighbors_, neighbors_in_memory_enabled_);\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"HnswSparseSearcherEntity new failed\");\n }\n\n return HnswSparseEntity::Pointer(entity);\n}\n\n//! Get vector sparse feature data by chunk index and offset\nconst void *HnswSparseSearcherEntity::get_sparse_data(uint64_t offset,\n uint32_t len) const {\n const void *sparse_data = nullptr;\n\n uint32_t real_length = sparse_vectors_->read(offset, &sparse_data, len);\n\n if (ailego_unlikely(real_length != len)) {\n LOG_ERROR(\"Read sparse data from segment failed, %u vs %u\", real_length,\n len);\n return nullptr;\n }\n\n return sparse_data;\n}\n\nint HnswSparseSearcherEntity::get_sparse_data(\n uint64_t offset, uint32_t len, IndexStorage::MemoryBlock &block) const {\n const void *vec = get_sparse_data(offset, len);\n block.reset((void *)vec);\n return 0;\n}\n\n//! Get sparse data from id\nconst void *HnswSparseSearcherEntity::get_sparse_data(node_id_t id) const {\n const void *vec = get_vector_meta(id);\n if (vec == nullptr) {\n LOG_ERROR(\"get vector failed, id: %u\", id);\n\n return nullptr;\n }\n\n auto sparse_data = get_sparse_data_from_vector(vec);\n\n return sparse_data.first;\n}\n\nint HnswSparseSearcherEntity::get_sparse_data(\n const node_id_t id, IndexStorage::MemoryBlock &block) const {\n const void *vec = get_sparse_data(id);\n block.reset((void *)vec);\n return 0;\n}\n\n//! Get sparse data from vector\nstd::pair\nHnswSparseSearcherEntity::get_sparse_data_from_vector(const void *vec) const {\n if (vec == nullptr) {\n LOG_ERROR(\"vec is nullptr\");\n\n return std::make_pair(nullptr, 0);\n }\n\n const char *vec_ptr = reinterpret_cast(vec);\n\n uint64_t offset = *((uint64_t *)(vec_ptr));\n uint32_t sparse_vector_len = *((uint32_t *)(vec_ptr + sizeof(uint64_t)));\n\n const void *sparse_data = get_sparse_data(offset, sparse_vector_len);\n if (ailego_unlikely(sparse_data == nullptr)) {\n LOG_ERROR(\"Get nullptr sparse, offset=%zu, len=%u\", (size_t)offset,\n sparse_vector_len);\n\n return std::make_pair(nullptr, 0);\n }\n\n return std::make_pair(sparse_data, sparse_vector_len);\n}\n\nint HnswSparseSearcherEntity::get_sparse_data_from_vector(\n const void *vec, IndexStorage::MemoryBlock &block,\n int &sparse_length) const {\n std::pair sparse_data =\n get_sparse_data_from_vector(vec);\n block.reset((void *)sparse_data.first);\n sparse_length = sparse_data.second;\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_searcher_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \"hnsw_sparse_builder_entity.h\"\n#include \"hnsw_sparse_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSparseSearcherEntity : public HnswSparseEntity {\n public:\n using Pointer = std::shared_ptr;\n using SegmentPointer = IndexStorage::Segment::Pointer;\n\n public:\n struct SegmentGroupParam {\n SegmentGroupParam(SegmentPointer neighbors_in,\n SegmentPointer neighbors_meta_in,\n SegmentPointer upper_neighbors_in,\n SegmentPointer upper_neighbors_meta_in)\n : neighbors{neighbors_in},\n neighbors_meta{neighbors_meta_in},\n upper_neighbors{upper_neighbors_in},\n upper_neighbors_meta{upper_neighbors_meta_in} {}\n\n SegmentPointer neighbors{nullptr};\n SegmentPointer neighbors_meta{nullptr};\n SegmentPointer upper_neighbors{nullptr};\n SegmentPointer upper_neighbors_meta{nullptr};\n };\n\n //! Constructor\n HnswSparseSearcherEntity();\n\n //! Make a copy of searcher entity, to support thread-safe operation.\n //! The segment in container cannot be read concurrenly\n virtual const HnswSparseEntity::Pointer clone() const override;\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const override;\n\n //! Get vector local id by key\n node_id_t get_id(key_t key) const;\n\n //! Get sparse vector feature data by key\n virtual int get_sparse_vector_by_key(\n key_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override;\n\n //! Get vector feature data by id\n virtual const void *get_vector_meta(node_id_t id) const override;\n\n virtual int get_vector_meta(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n\n //! Get vector feature data by id\n virtual int get_vector_metas(const node_id_t *ids, uint32_t count,\n const void **vecs) const override;\n\n virtual int get_vector_metas(\n const node_id_t *ids, uint32_t count,\n std::vector &block_vecs) const override;\n\n //! Get vector sparse feature data by chunk index and offset\n virtual const void *get_sparse_data(uint64_t offset,\n uint32_t len) const override;\n\n //! Get sparse data from id\n virtual const void *get_sparse_data(node_id_t id) const override;\n\n virtual int get_sparse_data(uint64_t offset, uint32_t len,\n IndexStorage::MemoryBlock &block) const override;\n\n virtual int get_sparse_data(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n\n //! Get sparse data from vector\n virtual std::pair get_sparse_data_from_vector(\n const void *vec) const override;\n\n virtual int get_sparse_data_from_vector(const void *vec,\n IndexStorage::MemoryBlock &block,\n int &sparse_length) const override;\n\n //! Get the node id's neighbors on graph level\n virtual const Neighbors get_neighbors(level_t level,\n node_id_t id) const override;\n\n virtual int load(const IndexStorage::Pointer &container,\n bool check_crc) override;\n\n int load_segments(bool check_crc);\n\n virtual int cleanup(void) override;\n\n public:\n bool is_loaded() const {\n return loaded_;\n }\n\n void set_neighbors_in_memory(bool enabled) {\n neighbors_in_memory_enabled_ = enabled;\n }\n\n //! get fixed length neighbors data\n int get_fixed_neighbors(std::vector *fixed_neighbors) const;\n\n private:\n //! Constructor\n HnswSparseSearcherEntity(const HNSWSparseHeader &hd,\n const SegmentPointer &keys,\n const SegmentPointer &mapping,\n const SegmentGroupParam &neighbor_group,\n const SegmentPointer &sparse_vector_meta,\n const SegmentPointer &sparse_vectors,\n const std::shared_ptr &fixed_neighbors,\n bool neighbors_in_memory_enabled)\n : HnswSparseEntity(hd),\n keys_(keys),\n mapping_(mapping),\n neighbors_(neighbor_group.neighbors),\n neighbors_meta_(neighbor_group.neighbors_meta),\n upper_neighbors_(neighbor_group.upper_neighbors),\n upper_neighbors_meta_(neighbor_group.upper_neighbors_meta),\n sparse_vector_meta_(sparse_vector_meta),\n sparse_vectors_(sparse_vectors),\n neighbors_in_memory_enabled_(neighbors_in_memory_enabled) {\n segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt()),\n IndexStorage::SegmentData(0U, 0U));\n fixed_neighbors_ = fixed_neighbors;\n }\n\n bool do_crc_check(std::vector &segments) const;\n\n inline size_t neighbors_size() const {\n return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t);\n }\n\n inline size_t upper_neighbors_size() const {\n return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t);\n }\n\n //! If neighbors_in_memory_enabled, load the level0 neighbors to memory\n int load_and_flat_neighbors(void);\n\n public:\n HnswSparseSearcherEntity(const HnswSparseSearcherEntity &) = delete;\n HnswSparseSearcherEntity &operator=(const HnswSparseSearcherEntity &) =\n delete;\n\n private:\n IndexStorage::Pointer container_{};\n\n SegmentPointer keys_{};\n SegmentPointer mapping_{};\n\n SegmentPointer neighbors_{};\n SegmentPointer neighbors_meta_{};\n SegmentPointer upper_neighbors_{};\n SegmentPointer upper_neighbors_meta_{};\n\n SegmentPointer sparse_vector_meta_{};\n SegmentPointer sparse_vectors_{};\n\n mutable std::vector segment_datas_{};\n std::shared_ptr fixed_neighbors_{}; // level 0 fixed size neighbors\n bool neighbors_in_memory_enabled_{false};\n bool loaded_{false};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_sparse_streamer.h\"\n#include \n#include \n#include \n#include \n#include \"hnsw_sparse_algorithm.h\"\n#include \"hnsw_sparse_context.h\"\n#include \"hnsw_sparse_dist_calculator.h\"\n#include \"hnsw_sparse_index_provider.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswSparseStreamer::HnswSparseStreamer() : entity_(stats_) {}\n\nHnswSparseStreamer::~HnswSparseStreamer() {\n if (state_ == STATE_INITED) {\n this->cleanup();\n }\n}\n\nint HnswSparseStreamer::init(const IndexMeta &imeta,\n const ailego::Params ¶ms) {\n meta_ = imeta;\n meta_.set_streamer(\"HnswSparseStreamer\", HnswSparseEntity::kRevision, params);\n\n params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_INDEX_SIZE, &max_index_size_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT,\n &upper_max_neighbor_cnt_);\n float multiplier = HnswSparseEntity::kDefaultL0MaxNeighborCntMultiplier;\n params.get(PARAM_HNSW_SPARSE_STREAMER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER,\n &multiplier);\n l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_;\n\n multiplier = HnswSparseEntity::kDefaultNeighborPruneMultiplier;\n params.get(PARAM_HNSW_SPARSE_STREAMER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier);\n size_t prune_cnt = multiplier * upper_max_neighbor_cnt_;\n scaling_factor_ = upper_max_neighbor_cnt_;\n params.get(PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR, &scaling_factor_);\n\n params.get(PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT, &docs_hard_limit_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_EF, &ef_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION, &ef_construction_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_ENABLE, &bf_enabled_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB,\n &bf_negative_prob_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_CHECK_CRC_ENABLE, &check_crc_enabled_);\n\n params.get(PARAM_HNSW_SPARSE_STREAMER_CHUNK_SIZE, &chunk_size_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_FILTER_SAME_KEY, &filter_same_key_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_GET_VECTOR_ENABLE,\n &get_vector_enabled_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_);\n params.get(PARAM_HNSW_SPARSE_STREAMER_FORCE_PADDING_RESULT_ENABLE,\n &force_padding_topk_enabled_);\n\n query_filtering_enabled_ =\n params.get(PARAM_HNSW_SPARSE_STREAMER_QUERY_FILTERING_RATIO,\n &query_filtering_ratio_);\n\n params.get(PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit_);\n if (docs_soft_limit_ > 0 && docs_soft_limit_ > docs_hard_limit_) {\n LOG_ERROR(\"[%s] must be >= [%s]\",\n PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT.c_str(),\n PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT.c_str());\n return IndexError_InvalidArgument;\n } else if (docs_soft_limit_ == 0UL) {\n docs_soft_limit_ =\n docs_hard_limit_ * HnswSparseEntity::kDefaultDocsSoftLimitRatio;\n }\n\n if (ef_ == 0U) {\n ef_ = HnswSparseEntity::kDefaultEf;\n }\n if (ef_construction_ == 0U) {\n ef_construction_ = HnswSparseEntity::kDefaultEfConstruction;\n }\n if (upper_max_neighbor_cnt_ == 0U) {\n upper_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt;\n }\n if (upper_max_neighbor_cnt_ > HnswSparseEntity::kMaxNeighborCnt) {\n LOG_ERROR(\"[%s] must be in range (0,%d)\",\n PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT.c_str(),\n HnswSparseEntity::kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (l0_max_neighbor_cnt_ == 0U) {\n l0_max_neighbor_cnt_ = HnswSparseEntity::kDefaultL0MaxNeighborCnt;\n }\n if (l0_max_neighbor_cnt_ > HnswSparseEntity::kMaxNeighborCnt) {\n LOG_ERROR(\"UpperNeighborCnt must be in range (0,%d)\",\n HnswSparseEntity::kMaxNeighborCnt);\n return IndexError_InvalidArgument;\n }\n if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) {\n LOG_ERROR(\"[%s]-[%u] must be <= [%s]-[%u]\",\n PARAM_HNSW_SPARSE_STREAMER_MIN_NEIGHBOR_COUNT.c_str(),\n min_neighbor_cnt_,\n PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT.c_str(),\n upper_max_neighbor_cnt_);\n return IndexError_InvalidArgument;\n }\n\n if (bf_negative_prob_ <= 0.0f || bf_negative_prob_ >= 1.0f) {\n LOG_ERROR(\n \"[%s] must be in range (0,1)\",\n PARAM_HNSW_SPARSE_STREAMER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (scaling_factor_ == 0U) {\n scaling_factor_ = HnswSparseEntity::kDefaultScalingFactor;\n }\n if (scaling_factor_ < 5 || scaling_factor_ > 1000) {\n LOG_ERROR(\"[%s] must be in range [5,1000]\",\n PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) {\n LOG_ERROR(\"[%s] must be in range (0.0f,1.0f]\",\n PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_RATIO.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (max_scan_limit_ < min_scan_limit_) {\n LOG_ERROR(\"[%s] must be >= [%s]\",\n PARAM_HNSW_SPARSE_STREAMER_MAX_SCAN_LIMIT.c_str(),\n PARAM_HNSW_SPARSE_STREAMER_MIN_SCAN_LIMIT.c_str());\n return IndexError_InvalidArgument;\n }\n\n if (prune_cnt == 0UL) {\n prune_cnt = upper_max_neighbor_cnt_;\n }\n if (chunk_size_ == 0UL) {\n chunk_size_ = HnswSparseEntity::kDefaultChunkSize;\n }\n if (chunk_size_ > HnswSparseEntity::kMaxChunkSize) {\n LOG_ERROR(\"[%s] must be < %zu\",\n PARAM_HNSW_SPARSE_STREAMER_CHUNK_SIZE.c_str(),\n HnswSparseEntity::kMaxChunkSize);\n return IndexError_InvalidArgument;\n }\n\n if (query_filtering_enabled_ &&\n (query_filtering_ratio_ <= 0.0f || query_filtering_ratio_ >= 1.0f)) {\n LOG_ERROR(\"[%s] must be in range (0, 1)\",\n PARAM_HNSW_SPARSE_SEARCHER_QUERY_FILTERING_RATIO.c_str());\n return IndexError_InvalidArgument;\n }\n\n entity_.set_ef_construction(ef_construction_);\n entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_);\n entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_);\n entity_.set_scaling_factor(scaling_factor_);\n entity_.set_prune_cnt(prune_cnt);\n\n entity_.set_chunk_size(chunk_size_);\n entity_.set_filter_same_key(filter_same_key_);\n entity_.set_get_vector(get_vector_enabled_);\n entity_.set_min_neighbor_cnt(min_neighbor_cnt_);\n\n entity_.set_sparse_meta_size(HnswSparseEntity::kSparseMetaSize);\n entity_.set_sparse_unit_size(meta_.unit_size());\n\n int ret = entity_.init(max_index_size_, docs_hard_limit_);\n if (ret != 0) {\n LOG_ERROR(\"Hnsw entity init failed for %s\", IndexError::What(ret));\n return ret;\n }\n LOG_DEBUG(\n \"Init params: maxIndexSize=%zu docsHardLimit=%zu docsSoftLimit=%zu \"\n \"efConstruction=%u ef=%u l0NeighborCnt=%u upperNeighborCnt=%u \"\n \"scalingFactor=%u maxScanRatio=%.3f minScanLimit=%zu maxScanLimit=%zu \"\n \"bfEnabled=%d bruteFoceThreshold=%zu bfNegativeProbability=%.5f \"\n \"checkCrcEnabled=%d pruneSize=%zu chunkSize=%zu \"\n \"filterSameKey=%u getVectorEnabled=%u \"\n \"minNeighborCount=%u forcePadding=%u filteringRatio=%f\",\n max_index_size_, docs_hard_limit_, docs_soft_limit_, ef_construction_,\n ef_, l0_max_neighbor_cnt_, upper_max_neighbor_cnt_, scaling_factor_,\n max_scan_ratio_, min_scan_limit_, max_scan_limit_, bf_enabled_,\n bruteforce_threshold_, bf_negative_prob_, check_crc_enabled_, prune_cnt,\n chunk_size_, filter_same_key_, get_vector_enabled_, min_neighbor_cnt_,\n force_padding_topk_enabled_, query_filtering_ratio_);\n\n alg_ = HnswSparseAlgorithm::UPointer(new HnswSparseAlgorithm(entity_));\n\n ret = alg_->init();\n if (ret != 0) {\n return ret;\n }\n\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint HnswSparseStreamer::cleanup(void) {\n if (state_ == STATE_OPENED) {\n this->close();\n }\n\n LOG_INFO(\"HnswSparseStreamer cleanup\");\n\n meta_.clear();\n metric_.reset();\n stats_.clear();\n entity_.cleanup();\n\n if (alg_) {\n alg_->cleanup();\n }\n\n max_index_size_ = 0UL;\n docs_hard_limit_ = HnswSparseEntity::kDefaultDocsHardLimit;\n docs_soft_limit_ = 0UL;\n upper_max_neighbor_cnt_ = HnswSparseEntity::kDefaultUpperMaxNeighborCnt;\n ef_ = HnswSparseEntity::kDefaultEf;\n ef_construction_ = HnswSparseEntity::kDefaultEfConstruction;\n bf_enabled_ = false;\n scaling_factor_ = HnswSparseEntity::kDefaultScalingFactor;\n bruteforce_threshold_ = HnswSparseEntity::kDefaultBruteForceThreshold;\n max_scan_limit_ = HnswSparseEntity::kDefaultMaxScanLimit;\n min_scan_limit_ = HnswSparseEntity::kDefaultMinScanLimit;\n chunk_size_ = HnswSparseEntity::kDefaultChunkSize;\n bf_negative_prob_ = HnswSparseEntity::kDefaultBFNegativeProbability;\n max_scan_ratio_ = HnswSparseEntity::kDefaultScanRatio;\n state_ = STATE_INIT;\n check_crc_enabled_ = false;\n filter_same_key_ = false;\n get_vector_enabled_ = false;\n\n sparse_neighbor_ratio_ = HnswSparseEntity::kDefaultSparseNeighborRatio;\n sparse_neighbor_cnt_ = 0UL;\n sparse_min_neighbor_cnt_ = 0UL;\n upper_sparse_neighbor_cnt_ = 0UL;\n\n return 0;\n}\n\nint HnswSparseStreamer::open(IndexStorage::Pointer stg) {\n LOG_INFO(\"HnswSparseStreamer open\");\n\n if (ailego_unlikely(state_ != STATE_INITED)) {\n LOG_ERROR(\"Open storage failed, init streamer first!\");\n return IndexError_NoReady;\n }\n int ret = entity_.open(std::move(stg), check_crc_enabled_);\n if (ret != 0) {\n return ret;\n }\n IndexMeta index_meta;\n ret = entity_.get_index_meta(&index_meta);\n if (ret == IndexError_NoExist) {\n // Set IndexMeta for the new index\n ret = entity_.set_index_meta(meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to set index meta for %s\", IndexError::What(ret));\n return ret;\n }\n } else if (ret != 0) {\n LOG_ERROR(\"Failed to get index meta for %s\", IndexError::What(ret));\n return ret;\n } else {\n if (index_meta.metric_name() != meta_.metric_name() ||\n index_meta.data_type() != meta_.data_type()) {\n LOG_ERROR(\"IndexMeta mismatch from the previous in index\");\n return IndexError_Mismatch;\n }\n // The IndexMetric Params may be updated like MipsSquaredEuclidean\n auto metric_params = index_meta.metric_params();\n metric_params.merge(meta_.metric_params());\n meta_.set_metric(index_meta.metric_name(), 0, metric_params);\n }\n\n metric_ = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric_) {\n LOG_ERROR(\"Failed to create metric %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n ret = metric_->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"Failled to init metric, ret=%d\", ret);\n return ret;\n }\n\n if (!metric_->sparse_distance()) {\n LOG_ERROR(\"Invalid metric distance\");\n return IndexError_InvalidArgument;\n }\n\n add_distance_ = metric_->sparse_distance();\n search_distance_ = add_distance_;\n\n if (metric_->query_metric() && metric_->query_metric()->distance()) {\n search_distance_ = metric_->query_metric()->sparse_distance();\n }\n\n state_ = STATE_OPENED;\n magic_ = IndexContext::GenerateMagic();\n\n return 0;\n}\n\nint HnswSparseStreamer::close(void) {\n LOG_INFO(\"HnswSparseStreamer close\");\n\n stats_.clear();\n meta_.set_metric(metric_->name(), 0, metric_->params());\n entity_.set_index_meta(meta_);\n int ret = entity_.close();\n if (ret != 0) {\n return ret;\n }\n state_ = STATE_INITED;\n\n return 0;\n}\n\nint HnswSparseStreamer::flush(uint64_t checkpoint) {\n LOG_INFO(\"HnswSparseStreamer flush checkpoint=%zu\", (size_t)checkpoint);\n\n meta_.set_metric(metric_->name(), 0, metric_->params());\n entity_.set_index_meta(meta_);\n return entity_.flush(checkpoint);\n}\n\nint HnswSparseStreamer::dump(const IndexDumper::Pointer &dumper) {\n LOG_INFO(\"HnswSparseStreamer dump\");\n\n shared_mutex_.lock();\n AILEGO_DEFER([&]() { shared_mutex_.unlock(); });\n\n meta_.set_searcher(\"HnswSparseSearcher\", HnswSparseEntity::kRevision,\n ailego::Params());\n\n int ret = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return ret;\n }\n return entity_.dump(dumper);\n}\n\nIndexStreamer::Context::Pointer HnswSparseStreamer::create_context(void) const {\n if (ailego_unlikely(state_ != STATE_OPENED)) {\n LOG_ERROR(\"Create context failed, open storage first!\");\n return Context::Pointer();\n }\n\n HnswSparseEntity::Pointer entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"CreateContext clone init failed\");\n return Context::Pointer();\n }\n HnswSparseContext *ctx =\n new (std::nothrow) HnswSparseContext(metric_, entity);\n if (ailego_unlikely(ctx == nullptr)) {\n LOG_ERROR(\"Failed to new HnswSparseContext\");\n return Context::Pointer();\n }\n ctx->set_ef(ef_);\n ctx->set_max_scan_limit(max_scan_limit_);\n ctx->set_min_scan_limit(min_scan_limit_);\n ctx->set_max_scan_ratio(max_scan_ratio_);\n ctx->set_filter_mode(bf_enabled_ ? VisitFilter::BloomFilter\n : VisitFilter::ByteMap);\n ctx->set_filter_negative_probability(bf_negative_prob_);\n ctx->set_magic(magic_);\n ctx->set_force_padding_topk(force_padding_topk_enabled_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n\n if (ailego_unlikely(ctx->init(HnswSparseContext::kSparseStreamerContext)) !=\n 0) {\n LOG_ERROR(\"Init HnswSparseContext failed\");\n delete ctx;\n return Context::Pointer();\n }\n\n return Context::Pointer(ctx);\n}\n\nIndexStreamer::SparseProvider::Pointer\nHnswSparseStreamer::create_sparse_provider(void) const {\n LOG_DEBUG(\"HnswSparseStreamer create sparse provider\");\n\n auto entity = entity_.clone();\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"Clone HnswSparseEntity failed\");\n return SparseProvider::Pointer();\n }\n return SparseProvider::Pointer(\n new HnswSparseIndexProvider(meta_, entity, \"HnswSparseStreamer\"));\n}\n\nint HnswSparseStreamer::update_context(HnswSparseContext *ctx) const {\n const HnswSparseEntity::Pointer entity = entity_.clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone search context entity\");\n return IndexError_Runtime;\n }\n ctx->set_max_scan_limit(max_scan_limit_);\n ctx->set_min_scan_limit(min_scan_limit_);\n ctx->set_max_scan_ratio(max_scan_ratio_);\n ctx->set_bruteforce_threshold(bruteforce_threshold_);\n return ctx->update_context(HnswSparseContext::kSparseStreamerContext, meta_,\n metric_, entity, magic_);\n}\n\n//! Add a vector with id into index with sparse inputs\nint HnswSparseStreamer::add_with_id_impl(uint32_t id,\n const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) {\n int ret = check_params(qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n if (ailego_unlikely(sparse_count > HnswSparseEntity::kSparseMaxDimSize)) {\n LOG_WARN(\n \"Failed to add sparse vector: number of non-zero elements (%u) exceeds \"\n \"maximum allowed (%u), id=%u\",\n sparse_count, HnswSparseEntity::kSparseMaxDimSize, id);\n return IndexError_InvalidValue;\n }\n\n HnswSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswSparseContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) {\n if (entity_.doc_cnt() >= docs_hard_limit_) {\n LOG_ERROR(\"Current docs %u exceed [%s]\", entity_.doc_cnt(),\n PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT.c_str());\n const std::lock_guard lk(mutex_);\n (*stats_.mutable_discarded_count())++;\n return IndexError_IndexFull;\n } else {\n LOG_WARN(\"Current docs %u exceed [%s]\", entity_.doc_cnt(),\n PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT.c_str());\n }\n }\n if (ailego_unlikely(!shared_mutex_.try_lock_shared())) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); });\n\n ctx->clear();\n ctx->update_dist_caculator_distance(add_distance_);\n\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_query,\n entity_.sparse_unit_size(),\n sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n ctx->check_need_adjuct_ctx(entity_.doc_cnt());\n\n level_t level = alg_->get_random_level();\n ret =\n entity_.add_vector_with_id(level, id, sparse_query_buffer, sparse_count);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer add vector failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n ret = alg_->add_node(id, level, ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw stramer add node failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n if (ailego_unlikely(ctx->error())) {\n (*stats_.mutable_discarded_count())++;\n return IndexError_Runtime;\n }\n (*stats_.mutable_added_count())++;\n\n return 0;\n}\n\n//! Add a vector into index with sparse inputs\nint HnswSparseStreamer::add_impl(uint64_t pkey, const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) {\n int ret = check_params(qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n if (ailego_unlikely(sparse_count > HnswSparseEntity::kSparseMaxDimSize)) {\n LOG_WARN(\n \"Failed to add sparse vector: number of non-zero elements (%u) exceeds \"\n \"maximum allowed (%u), key=%zu\",\n sparse_count, HnswSparseEntity::kSparseMaxDimSize, (size_t)pkey);\n return IndexError_InvalidValue;\n }\n\n HnswSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswSparseContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) {\n if (entity_.doc_cnt() >= docs_hard_limit_) {\n LOG_ERROR(\"Current docs %u exceed [%s]\", entity_.doc_cnt(),\n PARAM_HNSW_SPARSE_STREAMER_DOCS_HARD_LIMIT.c_str());\n const std::lock_guard lk(mutex_);\n (*stats_.mutable_discarded_count())++;\n return IndexError_IndexFull;\n } else {\n LOG_WARN(\"Current docs %u exceed [%s]\", entity_.doc_cnt(),\n PARAM_HNSW_SPARSE_STREAMER_DOCS_SOFT_LIMIT.c_str());\n }\n }\n if (ailego_unlikely(!shared_mutex_.try_lock_shared())) {\n LOG_ERROR(\"Cannot add vector while dumping index\");\n (*stats_.mutable_discarded_count())++;\n return IndexError_Unsupported;\n }\n AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); });\n\n ctx->clear();\n ctx->update_dist_caculator_distance(add_distance_);\n\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(sparse_count, sparse_indices, sparse_query,\n entity_.sparse_unit_size(),\n sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n ctx->check_need_adjuct_ctx(entity_.doc_cnt());\n\n level_t level = alg_->get_random_level();\n node_id_t id;\n ret = entity_.add_vector(level, pkey, sparse_query_buffer, sparse_count, &id);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw streamer add vector failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n ret = alg_->add_node(id, level, ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw stramer add node failed\");\n (*stats_.mutable_discarded_count())++;\n return ret;\n }\n\n if (ailego_unlikely(ctx->error())) {\n (*stats_.mutable_discarded_count())++;\n return IndexError_Runtime;\n }\n (*stats_.mutable_added_count())++;\n\n return 0;\n}\n\n//! Similarity search with sparse inputs\nint HnswSparseStreamer::search_impl(\n const uint32_t sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) const {\n return search_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1,\n context);\n}\n\n//! Similarity search with sparse inputs\nint HnswSparseStreamer::search_impl(\n const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta, uint32_t count,\n IndexStreamer::Context::Pointer &context) const {\n int ret = check_params(qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n HnswSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswSparseContext failed\");\n return IndexError_Cast;\n }\n\n if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) {\n return search_bf_impl(sparse_count, sparse_indices, sparse_query, qmeta,\n count, context);\n }\n\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->update_dist_caculator_distance(search_distance_);\n ctx->resize_results(count);\n ctx->check_need_adjuct_ctx(entity_.doc_cnt());\n\n const uint32_t *sparse_indices_tmp = sparse_indices;\n const void *sparse_query_tmp = sparse_query;\n\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n std::string sparse_query_filtered_buffer;\n\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n if (query_filtering_enabled_) {\n if (!SparseUtility::FilterSparseQuery(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n qmeta.data_type(), entity_.sparse_unit_size(),\n query_filtering_ratio_, &sparse_query_filtered_buffer)) {\n LOG_ERROR(\"Hnsw filtering failed\");\n\n return IndexError_Runtime;\n }\n\n ctx->reset_query(sparse_query_filtered_buffer.data());\n } else {\n ctx->reset_query(sparse_query_buffer.data());\n }\n\n ret = alg_->search(ctx);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Hnsw searcher fast search failed\");\n return ret;\n }\n\n if (query_filtering_enabled_) {\n ctx->reset_query(sparse_query_buffer.data());\n ctx->recal_topk_dist();\n }\n\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\n//! Similarity search with sparse inputs\nint HnswSparseStreamer::search_bf_impl(\n const uint32_t sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n IndexStreamer::Context::Pointer &context) const {\n return search_bf_impl(&sparse_count, sparse_indices, sparse_query, qmeta, 1,\n context);\n}\n\n//! Similarity search with sparse inputs\nint HnswSparseStreamer::search_bf_impl(\n const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta, uint32_t count,\n IndexStreamer::Context::Pointer &context) const {\n int ret = check_params(qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n HnswSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswSparseContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->update_dist_caculator_distance(search_distance_);\n ctx->resize_results(count);\n\n const uint32_t *sparse_indices_tmp = sparse_indices;\n const void *sparse_query_tmp = sparse_query;\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_Runtime;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n ctx->group_topk_heaps().clear();\n\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n\n if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) {\n dist_t dist = ctx->dist_calculator().dist(id);\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n } else {\n auto &filter = ctx->filter();\n auto &topk = ctx->topk_heap();\n\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n topk.clear();\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n if (entity_.get_key(id) == kInvalidKey) {\n continue;\n }\n\n if (!filter.is_valid() || !filter(entity_.get_key(id))) {\n dist_t dist = ctx->dist_calculator().dist(id);\n topk.emplace(id, dist);\n }\n }\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n }\n\n return 0;\n}\n\n//! Linear search by primary keys\nint HnswSparseStreamer::search_bf_by_p_keys_impl(\n const uint32_t sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, ContextPointer &context) const {\n return search_bf_by_p_keys_impl(&sparse_count, sparse_indices, sparse_query,\n p_keys, qmeta, 1, context);\n}\n\n//! Linear search by primary keys with sparse inputs\nint HnswSparseStreamer::search_bf_by_p_keys_impl(\n const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n int ret = check_params(qmeta);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n if (ailego_unlikely(p_keys.size() != count)) {\n LOG_ERROR(\"The size of p_keys is not equal to count\");\n return IndexError_InvalidArgument;\n }\n\n HnswSparseContext *ctx = dynamic_cast(context.get());\n ailego_do_if_false(ctx) {\n LOG_ERROR(\"Cast context to HnswSparseContext failed\");\n return IndexError_Cast;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher or streamer\n ret = update_context(ctx);\n if (ret != 0) {\n return ret;\n }\n }\n\n ctx->clear();\n ctx->update_dist_caculator_distance(search_distance_);\n ctx->resize_results(count);\n\n const uint32_t *sparse_indices_tmp = sparse_indices;\n const void *sparse_query_tmp = sparse_query;\n\n if (ctx->group_by_search()) {\n if (!ctx->group_by().is_valid()) {\n LOG_ERROR(\"Invalid group-by function\");\n return IndexError_Runtime;\n }\n\n std::function group_by = [&](node_id_t id) {\n return ctx->group_by()(entity_.get_key(id));\n };\n\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n ctx->group_topk_heaps().clear();\n\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n uint64_t pk = p_keys[q][idx];\n if (!ctx->filter().is_valid() || !ctx->filter()(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n dist_t dist = ctx->dist_calculator().dist(id);\n\n std::string group_id = group_by(id);\n\n auto &topk_heap = ctx->group_topk_heaps()[group_id];\n if (topk_heap.empty()) {\n topk_heap.limit(ctx->group_topk());\n }\n topk_heap.emplace_back(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n } else {\n auto &filter = ctx->filter();\n auto &topk = ctx->topk_heap();\n\n for (size_t q = 0; q < count; ++q) {\n std::string sparse_query_buffer;\n SparseUtility::TransSparseFormat(\n sparse_count[q], sparse_indices_tmp, sparse_query_tmp,\n entity_.sparse_unit_size(), sparse_query_buffer);\n\n ctx->reset_query(sparse_query_buffer.data());\n topk.clear();\n for (size_t idx = 0; idx < p_keys[q].size(); ++idx) {\n key_t pk = p_keys[q][idx];\n if (!filter.is_valid() || !filter(pk)) {\n node_id_t id = entity_.get_id(pk);\n if (id != kInvalidNodeId) {\n dist_t dist = ctx->dist_calculator().dist(id);\n topk.emplace(id, dist);\n }\n }\n }\n ctx->topk_to_result(q);\n\n sparse_indices_tmp += sparse_count[q];\n sparse_query_tmp = reinterpret_cast(sparse_query_tmp) +\n sparse_count[q] * qmeta.unit_size();\n }\n }\n\n if (ailego_unlikely(ctx->error())) {\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nvoid HnswSparseStreamer::print_debug_info() {\n for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) {\n Neighbors neighbours = entity_.get_neighbors(0, id);\n std::cout << \"node: \" << id << \"; \";\n for (uint32_t i = 0; i < neighbours.size(); ++i) {\n std::cout << neighbours[i];\n\n if (i == neighbours.size() - 1) {\n std::cout << std::endl;\n } else {\n std::cout << \", \";\n }\n }\n }\n\n // entity_.print_key_map();\n}\n\nINDEX_FACTORY_REGISTER_STREAMER(HnswSparseStreamer);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"hnsw_sparse_algorithm.h\"\n#include \"hnsw_sparse_streamer_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\nclass HnswSparseStreamer : public IndexStreamer {\n public:\n using ContextPointer = IndexStreamer::Context::Pointer;\n\n HnswSparseStreamer(void);\n virtual ~HnswSparseStreamer(void);\n\n HnswSparseStreamer(const HnswSparseStreamer &streamer) = delete;\n HnswSparseStreamer &operator=(const HnswSparseStreamer &streamer) = delete;\n\n protected:\n //! Initialize Streamer\n int init(const IndexMeta &imeta, const ailego::Params ¶ms) override;\n\n //! Cleanup Streamer\n int cleanup(void) override;\n\n //! Create a context\n Context::Pointer create_context(void) const override;\n\n //! Create a new sparse iterator\n IndexStreamer::SparseProvider::Pointer create_sparse_provider(\n void) const override;\n\n int add_impl(uint64_t pkey, const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta, Context::Pointer &context) override;\n\n int add_with_id_impl(uint32_t id, const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) override;\n\n //! Similarity search with sparse inputs\n int search_impl(const uint32_t sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity search with sparse inputs\n int search_impl(const uint32_t *sparse_count, const uint32_t *sparse_indices,\n const void *sparse_query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const override;\n\n //! Similarity brute force search with sparse inputs\n int search_bf_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity brute force search with sparse inputs\n int search_bf_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices, const void *sparse_query,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const override;\n\n //! Linear search by primary keys\n int search_bf_by_p_keys_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta,\n ContextPointer &context) const override;\n\n //! Linear search by primary keys with sparse inputs\n int search_bf_by_p_keys_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const std::vector> &p_keys,\n const IndexQueryMeta &qmeta, uint32_t count,\n ContextPointer &context) const override;\n\n //! Fetch sparse vector by key\n int get_sparse_vector(uint64_t key, uint32_t *sparse_count,\n std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override {\n return entity_.get_sparse_vector_by_key(\n key, sparse_count, sparse_indices_buffer, sparse_values_buffer);\n }\n\n //! Fetch vector by id\n int get_sparse_vector_by_id(\n uint32_t id, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override {\n return entity_.get_sparse_vector_by_id(\n id, sparse_count, sparse_indices_buffer, sparse_values_buffer);\n }\n\n //! Open index from file path\n int open(IndexStorage::Pointer stg) override;\n\n //! Close file\n int close(void) override;\n\n //! flush file\n int flush(uint64_t checkpoint) override;\n\n //! Dump index into storage\n int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n const Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve sparse meta of index\n const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n void print_debug_info() override;\n\n private:\n inline int check_params(const IndexQueryMeta &qmeta) const {\n if (ailego_unlikely(qmeta.data_type() != meta_.data_type())) {\n LOG_ERROR(\"Unsupported query meta\");\n return IndexError_Mismatch;\n }\n return 0;\n }\n\n inline int check_sparse_count_is_zero(const uint32_t *sparse_count,\n uint32_t count) const {\n for (uint32_t i = 0; i < count; ++i) {\n if (sparse_count[i] != 0)\n LOG_ERROR(\"Sparse cout is not empty. Index: %u, Sparse Count: %u\", i,\n sparse_count[i]);\n return IndexError_InvalidArgument;\n }\n\n return 0;\n }\n\n private:\n //! To share ctx across streamer/searcher, we need to update the context for\n //! current streamer/searcher\n int update_context(HnswSparseContext *ctx) const;\n\n private:\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 };\n class Stats : public IndexStreamer::Stats {\n public:\n void clear(void) {\n set_revision_id(0u);\n set_loaded_count(0u);\n set_added_count(0u);\n set_discarded_count(0u);\n set_index_size(0u);\n set_dumped_size(0u);\n set_check_point(0u);\n set_create_time(0u);\n set_update_time(0u);\n clear_attributes();\n }\n };\n\n HnswSparseStreamerEntity entity_;\n HnswSparseAlgorithm::UPointer alg_;\n IndexMeta meta_{};\n IndexMetric::Pointer metric_{};\n\n IndexMetric::MatrixSparseDistance add_distance_{};\n IndexMetric::MatrixSparseDistance search_distance_{};\n Stats stats_{};\n std::mutex mutex_{};\n\n size_t max_index_size_{0UL};\n size_t chunk_size_{HnswSparseEntity::kDefaultChunkSize};\n size_t docs_hard_limit_{HnswSparseEntity::kDefaultDocsHardLimit};\n size_t docs_soft_limit_{0UL};\n uint32_t min_neighbor_cnt_{0u};\n uint32_t upper_max_neighbor_cnt_{\n HnswSparseEntity::kDefaultUpperMaxNeighborCnt};\n uint32_t l0_max_neighbor_cnt_{HnswSparseEntity::kDefaultL0MaxNeighborCnt};\n uint32_t ef_{HnswSparseEntity::kDefaultEf};\n uint32_t ef_construction_{HnswSparseEntity::kDefaultEfConstruction};\n uint32_t scaling_factor_{HnswSparseEntity::kDefaultScalingFactor};\n size_t bruteforce_threshold_{HnswSparseEntity::kDefaultBruteForceThreshold};\n size_t max_scan_limit_{HnswSparseEntity::kDefaultMaxScanLimit};\n size_t min_scan_limit_{HnswSparseEntity::kDefaultMinScanLimit};\n float bf_negative_prob_{HnswSparseEntity::kDefaultBFNegativeProbability};\n float max_scan_ratio_{HnswSparseEntity::kDefaultScanRatio};\n float sparse_neighbor_ratio_{HnswSparseEntity::kDefaultSparseNeighborRatio};\n uint32_t sparse_neighbor_cnt_{0UL};\n uint32_t sparse_min_neighbor_cnt_{0UL};\n uint32_t upper_sparse_neighbor_cnt_{0UL};\n\n bool query_filtering_enabled_{false};\n float query_filtering_ratio_{HnswSparseEntity::kDefaultQueryFilteringRatio};\n\n uint32_t magic_{0U};\n State state_{STATE_INIT};\n bool bf_enabled_{false};\n bool check_crc_enabled_{false};\n bool filter_same_key_{false};\n bool get_vector_enabled_{false};\n bool force_padding_topk_enabled_{false};\n\n //! avoid add vector while dumping index\n ailego::SharedMutex shared_mutex_{};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"hnsw_sparse_streamer_entity.h\"\n#include \n#include \n#include \n#include \n#include \"utility/sparse_utility.h\"\n#include \"hnsw_sparse_dist_calculator.h\"\n\nnamespace zvec {\nnamespace core {\n\nHnswSparseStreamerEntity::HnswSparseStreamerEntity(IndexStreamer::Stats &stats)\n : stats_(stats) {}\n\nHnswSparseStreamerEntity::~HnswSparseStreamerEntity() {}\n\nint HnswSparseStreamerEntity::init(uint64_t max_index_size,\n size_t max_doc_cnt) {\n if (std::pow(scaling_factor(), kMaxGraphLayers) < max_doc_cnt) {\n LOG_ERROR(\"scalingFactor=%zu is too small\", scaling_factor());\n return IndexError_InvalidArgument;\n }\n\n std::lock_guard lock(mutex_);\n broker_ = std::make_shared(stats_);\n upper_neighbor_index_ = std::make_shared();\n keys_map_lock_ = std::make_shared();\n keys_map_ = std::make_shared>();\n if (!keys_map_ || !upper_neighbor_index_ || !broker_ || !keys_map_lock_) {\n LOG_ERROR(\"HnswSparseStreamerEntity new object failed\");\n return IndexError_NoMemory;\n }\n keys_map_->set_empty_key(kInvalidKey);\n\n neighbor_size_ = neighbors_size();\n upper_neighbor_size_ = upper_neighbors_size();\n\n //! vector + key + level 0 neighbors\n size_t size = sizeof(key_t) + neighbor_size_ + sparse_meta_size();\n\n size = AlignSize(size);\n set_node_size(size);\n\n return init_chunk_params(max_index_size);\n}\n\nint HnswSparseStreamerEntity::cleanup() {\n std::lock_guard lock(mutex_);\n mutable_header()->clear();\n chunk_size_ = kDefaultChunkSize;\n node_index_mask_bits_ = 0U;\n node_index_mask_ = 0U;\n node_cnt_per_chunk_ = 0U;\n neighbor_size_ = 0U;\n upper_neighbor_size_ = 0U;\n if (upper_neighbor_index_) {\n upper_neighbor_index_->cleanup();\n }\n if (keys_map_) {\n keys_map_->clear();\n }\n node_chunks_.clear();\n upper_neighbor_chunks_.clear();\n filter_same_key_ = false;\n get_vector_enabled_ = false;\n broker_.reset();\n\n return 0;\n}\n\nint HnswSparseStreamerEntity::update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) {\n std::vector buffer(neighbor_size_);\n NeighborsHeader *hd = reinterpret_cast(buffer.data());\n hd->neighbor_cnt = neighbors.size();\n size_t i = 0;\n for (; i < neighbors.size(); ++i) {\n hd->neighbors[i] = neighbors[i].first;\n }\n\n auto loc = get_neighbor_chunk_loc(level, id);\n size_t size = reinterpret_cast(&hd->neighbors[i]) - &buffer[0];\n size_t ret = loc.first->write(loc.second, hd, size);\n if (ailego_unlikely(ret != size)) {\n LOG_ERROR(\"Write neighbor header failed, ret=%zu\", ret);\n\n return IndexError_Runtime;\n }\n\n return 0;\n}\n\nconst Neighbors HnswSparseStreamerEntity::get_neighbors(level_t level,\n node_id_t id) const {\n SparseChunk *chunk = nullptr;\n size_t offset = 0UL;\n size_t neighbor_size = neighbor_size_;\n if (level == 0UL) {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n offset = (id & node_index_mask_) * node_size() + sizeof(key_t) +\n sparse_meta_size();\n\n sync_chunks(SparseChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_);\n ailego_assert_with(chunk_idx < node_chunks_.size(), \"invalid chunk idx\");\n chunk = node_chunks_[chunk_idx].get();\n } else {\n auto p = get_upper_neighbor_chunk_loc(level, id);\n chunk = upper_neighbor_chunks_[p.first].get();\n offset = p.second;\n neighbor_size = upper_neighbor_size_;\n }\n\n ailego_assert_with(offset < chunk->data_size(), \"invalid chunk offset\");\n IndexStorage::MemoryBlock neighbor_block;\n size_t size = chunk->read(offset, neighbor_block, neighbor_size);\n if (ailego_unlikely(size != neighbor_size)) {\n LOG_ERROR(\"Read neighbor header failed, ret=%zu\", size);\n return Neighbors();\n }\n return Neighbors(std::move(neighbor_block));\n}\n\n//! Get vector feature data by key\nconst void *HnswSparseStreamerEntity::get_vector_meta(node_id_t id) const {\n auto loc = get_vector_chunk_loc(id);\n const void *vec = nullptr;\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = sparse_meta_size();\n\n size_t ret = node_chunks_[loc.first]->read(loc.second, &vec, read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%u, read size=%zu, ret=%zu\",\n loc.second, read_size, ret);\n }\n\n return vec;\n}\n\nint HnswSparseStreamerEntity::get_vector_meta(\n const node_id_t id, IndexStorage::MemoryBlock &block) const {\n auto loc = get_vector_chunk_loc(id);\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = sparse_meta_size();\n\n size_t ret = node_chunks_[loc.first]->read(loc.second, block, read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%u, read size=%zu, ret=%zu\",\n loc.second, read_size, ret);\n return IndexError_ReadData;\n }\n\n return 0;\n}\n\nint HnswSparseStreamerEntity::get_vector_metas(const node_id_t *ids,\n uint32_t count,\n const void **vecs) const {\n for (auto i = 0U; i < count; ++i) {\n auto loc = get_vector_chunk_loc(ids[i]);\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = sparse_meta_size();\n\n size_t ret = node_chunks_[loc.first]->read(loc.second, &vecs[i], read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%u, read size=%zu, ret=%zu\",\n loc.second, read_size, ret);\n return IndexError_ReadData;\n }\n }\n\n return 0;\n}\n\nint HnswSparseStreamerEntity::get_vector_metas(\n const node_id_t *ids, uint32_t count,\n std::vector &block_vecs) const {\n block_vecs.resize(count);\n for (auto i = 0U; i < count; ++i) {\n auto loc = get_vector_chunk_loc(ids[i]);\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t read_size = sparse_meta_size();\n\n size_t ret =\n node_chunks_[loc.first]->read(loc.second, block_vecs[i], read_size);\n if (ailego_unlikely(ret != read_size)) {\n LOG_ERROR(\"Read vector failed, offset=%u, read size=%zu, ret=%zu\",\n loc.second, read_size, ret);\n return IndexError_ReadData;\n }\n }\n\n return 0;\n}\n\n//! Get vector feature data by key\nconst void *HnswSparseStreamerEntity::get_sparse_data(uint64_t offset,\n uint32_t len) const {\n uint32_t chunk_index = offset >> 32;\n uint32_t chunk_offset = offset & 0xFFFFFFFF;\n\n auto loc = get_sparse_chunk_loc(chunk_index, chunk_offset);\n const void *data = nullptr;\n\n ailego_assert_with(loc.first < sparse_node_chunks_.size(),\n \"invalid chunk idx\");\n ailego_assert_with(loc.second < sparse_node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t ret = sparse_node_chunks_[loc.first]->read(loc.second, &data, len);\n if (ailego_unlikely(ret != len)) {\n LOG_ERROR(\"Read sparse vector failed, offset=%zu, read size=%u, ret=%zu\",\n (size_t)offset, len, ret);\n }\n return data;\n}\n\nint HnswSparseStreamerEntity::get_sparse_data(\n uint64_t offset, uint32_t len, IndexStorage::MemoryBlock &block) const {\n uint32_t chunk_index = offset >> 32;\n uint32_t chunk_offset = offset & 0xFFFFFFFF;\n\n auto loc = get_sparse_chunk_loc(chunk_index, chunk_offset);\n ailego_assert_with(loc.first < sparse_node_chunks_.size(),\n \"invalid chunk idx\");\n ailego_assert_with(loc.second < sparse_node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n\n size_t ret = sparse_node_chunks_[loc.first]->read(loc.second, block, len);\n if (ailego_unlikely(ret != len)) {\n LOG_ERROR(\"Read sparse vector failed, offset=%zu, read size=%u, ret=%zu\",\n (size_t)offset, len, ret);\n return IndexError_ReadData;\n }\n return 0;\n}\n\n//! Get sparse data from id\nconst void *HnswSparseStreamerEntity::get_sparse_data(node_id_t id) const {\n auto sparse_data = get_sparse_data_from_vector(get_vector_meta(id));\n\n return sparse_data.first;\n}\n\nint HnswSparseStreamerEntity::get_sparse_data(\n node_id_t id, IndexStorage::MemoryBlock &block) const {\n IndexStorage::MemoryBlock meta_block;\n get_vector_meta(id, meta_block);\n int sparse_length = 0;\n return get_sparse_data_from_vector(meta_block.data(), block, sparse_length);\n}\n\n//! Get sparse data from vector\nstd::pair\nHnswSparseStreamerEntity::get_sparse_data_from_vector(const void *vec) const {\n const char *vec_ptr = reinterpret_cast(vec);\n\n uint64_t offset = *((uint64_t *)(vec_ptr));\n uint32_t sparse_vector_len = *((uint32_t *)(vec_ptr + sizeof(uint64_t)));\n\n if (sparse_vector_len > 0) {\n const void *sparse_data = get_sparse_data(offset, sparse_vector_len);\n if (ailego_unlikely(sparse_data == nullptr)) {\n LOG_ERROR(\"Get nullptr sparse, offset=%zu, len=%u\", (size_t)offset,\n sparse_vector_len);\n\n return std::make_pair(nullptr, 0);\n }\n\n return std::make_pair(sparse_data, sparse_vector_len);\n }\n\n return std::make_pair(nullptr, 0);\n}\n\nint HnswSparseStreamerEntity::get_sparse_data_from_vector(\n const void *vec, IndexStorage::MemoryBlock &block,\n int &sparse_length) const {\n const char *vec_ptr = reinterpret_cast(vec);\n\n uint64_t offset = *((uint64_t *)(vec_ptr));\n uint32_t sparse_vector_len = *((uint32_t *)(vec_ptr + sizeof(uint64_t)));\n\n if (sparse_vector_len > 0) {\n int ret = get_sparse_data(offset, sparse_vector_len, block);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Get nullptr sparse, offset=%zu, len=%u\", (size_t)offset,\n sparse_vector_len);\n return IndexError_ReadData;\n }\n sparse_length = sparse_vector_len;\n }\n return 0;\n}\n\nkey_t HnswSparseStreamerEntity::get_key(node_id_t id) const {\n auto loc = get_key_chunk_loc(id);\n IndexStorage::MemoryBlock key_block;\n ailego_assert_with(loc.first < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(),\n \"invalid chunk offset\");\n size_t ret =\n node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t));\n if (ailego_unlikely(ret != sizeof(key_t))) {\n LOG_ERROR(\"Read vector failed, ret=%zu\", ret);\n return kInvalidKey;\n }\n\n return *reinterpret_cast(key_block.data());\n}\n\nvoid HnswSparseStreamerEntity::add_neighbor(level_t level, node_id_t id,\n uint32_t size,\n node_id_t neighbor_id) {\n auto loc = get_neighbor_chunk_loc(level, id);\n size_t offset =\n loc.second + sizeof(NeighborsHeader) + size * sizeof(node_id_t);\n ailego_assert_with(size < neighbor_cnt(level), \"invalid neighbor size\");\n ailego_assert_with(offset < loc.first->data_size(), \"invalid chunk offset\");\n size_t ret = loc.first->write(offset, &neighbor_id, sizeof(node_id_t));\n if (ailego_unlikely(ret != sizeof(node_id_t))) {\n LOG_ERROR(\"Write neighbor id failed, ret=%zu\", ret);\n return;\n }\n\n uint32_t neighbors = size + 1;\n ret = loc.first->write(loc.second, &neighbors, sizeof(uint32_t));\n if (ailego_unlikely(ret != sizeof(uint32_t))) {\n LOG_ERROR(\"Write neighbor cnt failed, ret=%zu\", ret);\n }\n\n return;\n}\n\nint HnswSparseStreamerEntity::init_chunks(\n const SparseChunk::Pointer &header_chunk) {\n if (header_chunk->data_size() < header_size()) {\n LOG_ERROR(\"Invalid header chunk size\");\n return IndexError_InvalidFormat;\n }\n IndexStorage::MemoryBlock data_block;\n size_t size = header_chunk->read(0UL, data_block, header_size());\n if (ailego_unlikely(size != header_size())) {\n LOG_ERROR(\"Read header chunk failed\");\n return IndexError_ReadData;\n }\n *mutable_header() =\n *reinterpret_cast(data_block.data());\n\n int ret = check_hnsw_index(&header());\n if (ret != 0) {\n broker_->close();\n return ret;\n }\n\n node_chunks_.resize(\n broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_NODE));\n for (auto seq = 0UL; seq < node_chunks_.size(); ++seq) {\n node_chunks_[seq] =\n broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_NODE, seq);\n if (!node_chunks_[seq]) {\n LOG_ERROR(\"Missing hnsw streamer data chunk %zu th of %zu\", seq,\n node_chunks_.size());\n return IndexError_InvalidFormat;\n }\n }\n\n upper_neighbor_chunks_.resize(\n broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR));\n for (auto seq = 0UL; seq < upper_neighbor_chunks_.size(); ++seq) {\n upper_neighbor_chunks_[seq] =\n broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, seq);\n if (!upper_neighbor_chunks_[seq]) {\n LOG_ERROR(\"Missing hnsw streamer index chunk %zu th of %zu\", seq,\n upper_neighbor_chunks_.size());\n return IndexError_InvalidFormat;\n }\n }\n\n sparse_node_chunks_.resize(\n broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE));\n for (auto seq = 0UL; seq < sparse_node_chunks_.size(); ++seq) {\n sparse_node_chunks_[seq] =\n broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE, seq);\n if (!sparse_node_chunks_[seq]) {\n LOG_ERROR(\"Missing hnsw streamer sparse data chunk %zu th of %zu\", seq,\n sparse_node_chunks_.size());\n return IndexError_InvalidFormat;\n }\n }\n\n return 0;\n}\n\nint HnswSparseStreamerEntity::open(IndexStorage::Pointer stg, bool check_crc) {\n std::lock_guard lock(mutex_);\n int ret =\n broker_->open(std::move(stg), max_index_size_, chunk_size_, check_crc);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Open index failed for %s\", IndexError::What(ret));\n return ret;\n }\n ret = upper_neighbor_index_->init(broker_, upper_neighbor_chunk_size_,\n scaling_factor(), estimate_doc_capacity(),\n kUpperHashMemoryInflateRatio);\n if (ailego_unlikely(ret != 0)) {\n LOG_ERROR(\"Init neighbor hash map failed\");\n return ret;\n }\n\n //! init header\n auto header_chunk = broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_HEADER,\n SparseChunkBroker::kDefaultChunkSeqId);\n if (!header_chunk) { // open empty index, create one\n auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_HEADER,\n SparseChunkBroker::kDefaultChunkSeqId,\n header_size());\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc header chunk failed\");\n return p.first;\n }\n size_t size = p.second->write(0UL, &header(), header_size());\n if (ailego_unlikely(size != header_size())) {\n LOG_ERROR(\"Write header chunk failed\");\n return IndexError_WriteData;\n }\n return 0;\n }\n\n //! Open an exist hnsw index\n ret = init_chunks(header_chunk);\n if (ailego_unlikely(ret != 0)) {\n return ret;\n }\n\n //! total docs including features wrote in index but neighbors may not ready\n node_id_t total_vecs = 0;\n if (node_chunks_.size() > 0) {\n size_t last_idx = node_chunks_.size() - 1;\n auto last_chunk = node_chunks_[last_idx];\n if (last_chunk->data_size() % node_size()) {\n LOG_WARN(\"The index may broken\");\n return IndexError_InvalidFormat;\n }\n total_vecs = last_idx * node_cnt_per_chunk_ +\n node_chunks_[last_idx]->data_size() / node_size();\n }\n\n LOG_INFO(\n \"Open index, l0NeighborCnt=%zu upperneighborCnt=%zu \"\n \"efConstruction=%zu curDocCnt=%u totalVecs=%u maxLevel=%u\",\n l0_neighbor_cnt(), upper_neighbor_cnt(), ef_construction(), doc_cnt(),\n total_vecs, cur_max_level());\n //! try to correct the docCnt if index not fully flushed\n if (doc_cnt() != total_vecs) {\n LOG_WARN(\"Index closed abnormally, using totalVecs as curDocCnt\");\n *mutable_doc_cnt() = total_vecs;\n }\n if (filter_same_key_ || get_vector_enabled_) {\n for (node_id_t id = 0U; id < doc_cnt(); ++id) {\n (*keys_map_)[get_key(id)] = id;\n }\n }\n\n stats_.set_loaded_count(doc_cnt());\n\n return 0;\n}\n\nint HnswSparseStreamerEntity::close() {\n LOG_DEBUG(\"close index\");\n\n std::lock_guard lock(mutex_);\n flush_header();\n mutable_header()->reset();\n upper_neighbor_index_->cleanup();\n keys_map_->clear();\n header_.clear();\n node_chunks_.clear();\n upper_neighbor_chunks_.clear();\n\n sparse_node_chunks_.clear();\n\n return broker_->close();\n}\n\nint HnswSparseStreamerEntity::flush(uint64_t checkpoint) {\n LOG_INFO(\"Flush index, curDocs=%u\", doc_cnt());\n\n std::lock_guard lock(mutex_);\n flush_header();\n int ret = broker_->flush(checkpoint);\n if (ret != 0) {\n return ret;\n }\n\n return 0;\n}\n\nint HnswSparseStreamerEntity::dump(const IndexDumper::Pointer &dumper) {\n LOG_INFO(\"Dump index, curDocs=%u\", doc_cnt());\n\n //! sort by keys, to support get_vector by key in searcher\n std::vector keys(doc_cnt());\n for (node_id_t i = 0; i < doc_cnt(); ++i) {\n keys[i] = get_key(i);\n }\n\n //! dump neighbors\n auto get_level = [&](node_id_t id) {\n auto it = upper_neighbor_index_->find(id);\n if (it == upper_neighbor_index_->end()) {\n return 0U;\n };\n auto meta = reinterpret_cast(&it->second);\n return meta->level;\n };\n auto ret = dump_segments(dumper, keys.data(), get_level);\n if (ailego_unlikely(ret < 0)) {\n return ret;\n }\n *stats_.mutable_dumped_size() += ret;\n\n return 0;\n}\n\nint HnswSparseStreamerEntity::check_hnsw_index(\n const HNSWSparseHeader *hd) const {\n if (l0_neighbor_cnt() != hd->neighbor_cnt() ||\n upper_neighbor_cnt() != hd->upper_neighbor_cnt()) {\n LOG_ERROR(\"Param neighbors:%zu:%zu mismatch index previous %zu:%zu\",\n l0_neighbor_cnt(), upper_neighbor_cnt(), hd->neighbor_cnt(),\n hd->upper_neighbor_cnt());\n return IndexError_Mismatch;\n }\n if (ef_construction() != hd->ef_construction()) {\n LOG_WARN(\"Param efConstruction %zu mismatch index previous %zu\",\n ef_construction(), hd->ef_construction());\n }\n if (scaling_factor() != hd->scaling_factor()) {\n LOG_WARN(\"Param scalingFactor %zu mismatch index previous %zu\",\n scaling_factor(), hd->scaling_factor());\n return IndexError_Mismatch;\n }\n if (prune_cnt() != hd->neighbor_prune_cnt()) {\n LOG_WARN(\"Param pruneCnt %zu mismatch index previous %zu\", prune_cnt(),\n hd->neighbor_prune_cnt());\n return IndexError_Mismatch;\n }\n if ((hd->entry_point() != kInvalidNodeId &&\n hd->entry_point() >= hd->doc_cnt()) ||\n (hd->entry_point() == kInvalidNodeId && hd->doc_cnt() > 0U)) {\n LOG_WARN(\"Invalid entryPoint %u, docCnt %u\", hd->entry_point(),\n hd->doc_cnt());\n return IndexError_InvalidFormat;\n }\n if (hd->entry_point() == kInvalidNodeId &&\n broker_->get_chunk_cnt(SparseChunkBroker::CHUNK_TYPE_NODE) > 0) {\n LOG_WARN(\"The index is broken, maybe it haven't flush\");\n return IndexError_InvalidFormat;\n }\n\n return 0;\n}\n\nint HnswSparseStreamerEntity::add_vector(level_t level, key_t key,\n const std::string &sparse_vec,\n uint32_t sparse_count, node_id_t *id) {\n // allocat sparse chunk\n uint32_t sparse_vector_len = sparse_vec.size();\n\n sparse_vector_len = AlignSize(sparse_vector_len);\n\n if (sparse_vector_len > sparse_chunk_size_) {\n LOG_ERROR(\n \"Sparse Vector Length exceed the chunk size, sparse vec len: %u, chunk \"\n \"size: %u\",\n sparse_vector_len, sparse_chunk_size_);\n return IndexError_InvalidArgument;\n }\n\n SparseChunk::Pointer node_chunk;\n SparseChunk::Pointer sparse_node_chunk;\n\n size_t chunk_offset = -1UL;\n size_t sparse_chunk_offset = -1UL;\n\n std::lock_guard lock(mutex_);\n // duplicate check\n if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) {\n LOG_WARN(\"Try to add duplicate key, ignore it\");\n return IndexError_Duplicate;\n }\n\n node_id_t local_id = static_cast(doc_cnt());\n\n uint32_t chunk_index = node_chunks_.size() - 1U;\n if (chunk_index == -1U ||\n (node_chunks_[chunk_index]->data_size() >=\n node_cnt_per_chunk_ * node_size())) { // no space left and need to alloc\n if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n chunk_index++;\n auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_NODE,\n chunk_index, chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n node_chunk = p.second;\n chunk_offset = 0UL;\n node_chunks_.emplace_back(node_chunk);\n } else {\n node_chunk = node_chunks_[chunk_index];\n chunk_offset = node_chunk->data_size();\n }\n\n uint32_t sparse_chunk_index = sparse_node_chunks_.size() - 1U;\n if (sparse_chunk_index == -1U ||\n sparse_node_chunks_[sparse_chunk_index]->data_size() + sparse_vector_len >\n sparse_chunk_size_) {\n if (ailego_unlikely(sparse_node_chunks_.capacity() ==\n sparse_node_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n sparse_chunk_index++;\n auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE,\n sparse_chunk_index, sparse_chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n sparse_node_chunk = p.second;\n\n sparse_node_chunks_.emplace_back(sparse_node_chunk);\n\n sparse_chunk_offset = 0UL;\n } else {\n sparse_node_chunk = sparse_node_chunks_[sparse_chunk_index];\n sparse_chunk_offset = sparse_node_chunk->data_size();\n }\n\n // write sparse vector\n if (sparse_vec.size() > 0) {\n size_t size = sparse_node_chunk->write(\n sparse_chunk_offset, sparse_vec.data(), sparse_vec.size());\n if (ailego_unlikely(size != sparse_vec.size())) {\n LOG_ERROR(\"SparseChunk write sparse vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n }\n\n uint64_t sparse_offset = sparse_chunk_index;\n sparse_offset = (sparse_offset << 32) + sparse_chunk_offset;\n\n size_t size =\n node_chunk->write(chunk_offset, &sparse_offset, sizeof(uint64_t));\n if (ailego_unlikely(size != sizeof(uint64_t))) {\n LOG_ERROR(\"SparseChunk write sparse vec index failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n size = node_chunk->write(chunk_offset + sizeof(uint64_t), &sparse_vector_len,\n sizeof(uint32_t));\n if (ailego_unlikely(size != sizeof(uint32_t))) {\n LOG_ERROR(\"SparseChunk write sparse vec len failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n size =\n node_chunk->write(chunk_offset + sparse_meta_size(), &key, sizeof(key_t));\n if (ailego_unlikely(size != sizeof(key_t))) {\n LOG_ERROR(\"SparseChunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n //! level 0 neighbors is inited to zero by default\n int ret = add_upper_neighbor(level, local_id);\n if (ret != 0) {\n return ret;\n }\n\n if (sparse_vector_len > 0) {\n sparse_chunk_offset += sparse_vector_len;\n if (ailego_unlikely(sparse_node_chunk->resize(sparse_chunk_offset) !=\n sparse_chunk_offset)) {\n LOG_ERROR(\"SparseChunk resize to %zu failed\", sparse_chunk_offset);\n return IndexError_Runtime;\n }\n }\n\n chunk_offset += node_size();\n if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"SparseChunk resize to %zu failed\", chunk_offset);\n return IndexError_Runtime;\n }\n\n if (filter_same_key_ || get_vector_enabled_) {\n keys_map_lock_->lock();\n (*keys_map_)[key] = local_id;\n keys_map_lock_->unlock();\n }\n\n *mutable_doc_cnt() += 1;\n *mutable_total_sparse_count() += sparse_count;\n\n broker_->mark_dirty();\n *id = local_id;\n\n return 0;\n}\n\nint HnswSparseStreamerEntity::add_vector_with_id(level_t level, node_id_t id,\n const std::string &sparse_vec,\n uint32_t sparse_count) {\n key_t key = id;\n SparseChunk::Pointer node_chunk;\n SparseChunk::Pointer sparse_node_chunk;\n size_t chunk_offset = -1UL;\n size_t sparse_chunk_offset = -1UL;\n\n // allocat sparse chunk\n uint32_t sparse_vector_len = sparse_vec.size();\n\n sparse_vector_len = AlignSize(sparse_vector_len);\n\n if (sparse_vector_len > sparse_chunk_size_) {\n LOG_ERROR(\n \"Sparse Vector Length exceed the chunk size, sparse vec len: %u, chunk \"\n \"size: %u\",\n sparse_vector_len, sparse_chunk_size_);\n return IndexError_InvalidArgument;\n }\n\n\n std::lock_guard lock(mutex_);\n\n // duplicate check\n if (ailego_unlikely(filter_same_key_ && get_id(key) != kInvalidNodeId)) {\n LOG_WARN(\"Try to add duplicate key, ignore it\");\n return IndexError_Duplicate;\n }\n\n auto func_get_sparse_node_chunk_and_offset = [&](node_id_t node_id) -> int {\n uint32_t chunk_index = node_id >> node_index_mask_bits_;\n ailego_assert_with(chunk_index <= node_chunks_.size(), \"invalid chunk idx\");\n // belongs to next chunk\n if (chunk_index == node_chunks_.size()) {\n if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_NODE,\n chunk_index, chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n node_chunk = p.second;\n node_chunks_.emplace_back(node_chunk);\n }\n\n node_chunk = node_chunks_[chunk_index];\n chunk_offset = (node_id & node_index_mask_) * node_size();\n return 0;\n };\n\n for (size_t start_id = doc_cnt(); start_id < id; ++start_id) {\n if (auto ret = func_get_sparse_node_chunk_and_offset(start_id); ret != 0) {\n LOG_ERROR(\"func_get_sparse_node_chunk_and_offset failed\");\n return ret;\n }\n size_t size = node_chunk->write(chunk_offset + sparse_meta_size(),\n &kInvalidKey, sizeof(key_t));\n if (ailego_unlikely(size != sizeof(key_t))) {\n LOG_ERROR(\"SparseChunk write key failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n chunk_offset += node_size();\n if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"SparseChunk resize to %zu failed\", chunk_offset);\n return IndexError_Runtime;\n }\n }\n\n if (auto ret = func_get_sparse_node_chunk_and_offset(id); ret != 0) {\n LOG_ERROR(\"func_get_sparse_node_chunk_and_offset failed\");\n return ret;\n }\n\n uint32_t sparse_chunk_index = sparse_node_chunks_.size() - 1U;\n if (sparse_chunk_index == -1U ||\n sparse_node_chunks_[sparse_chunk_index]->data_size() + sparse_vector_len >\n sparse_chunk_size_) {\n if (ailego_unlikely(sparse_node_chunks_.capacity() ==\n sparse_node_chunks_.size())) {\n LOG_ERROR(\"add vector failed for no memory quota\");\n return IndexError_IndexFull;\n }\n sparse_chunk_index++;\n auto p = broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE,\n sparse_chunk_index, sparse_chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n sparse_node_chunk = p.second;\n\n sparse_node_chunks_.emplace_back(sparse_node_chunk);\n\n sparse_chunk_offset = 0UL;\n } else {\n sparse_node_chunk = sparse_node_chunks_[sparse_chunk_index];\n sparse_chunk_offset = sparse_node_chunk->data_size();\n }\n\n // write sparse vector\n if (sparse_vec.size() > 0) {\n size_t size = sparse_node_chunk->write(\n sparse_chunk_offset, sparse_vec.data(), sparse_vec.size());\n if (ailego_unlikely(size != sparse_vec.size())) {\n LOG_ERROR(\"SparseChunk write sparse vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n }\n\n uint64_t sparse_offset = sparse_chunk_index;\n sparse_offset = (sparse_offset << 32) + sparse_chunk_offset;\n\n size_t size =\n node_chunk->write(chunk_offset, &sparse_offset, sizeof(uint64_t));\n if (ailego_unlikely(size != sizeof(uint64_t))) {\n LOG_ERROR(\"SparseChunk write sparse vec index failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n size = node_chunk->write(chunk_offset + sizeof(uint64_t), &sparse_vector_len,\n sizeof(uint32_t));\n if (ailego_unlikely(size != sizeof(uint32_t))) {\n LOG_ERROR(\"SparseChunk write sparse vec len failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n size =\n node_chunk->write(chunk_offset + sparse_meta_size(), &key, sizeof(key_t));\n if (ailego_unlikely(size != sizeof(key_t))) {\n LOG_ERROR(\"SparseChunk write vec failed, ret=%zu\", size);\n return IndexError_WriteData;\n }\n\n //! level 0 neighbors is inited to zero by default\n int ret = add_upper_neighbor(level, id);\n if (ret != 0) {\n return ret;\n }\n\n if (sparse_vector_len > 0) {\n sparse_chunk_offset += sparse_vector_len;\n if (ailego_unlikely(sparse_node_chunk->resize(sparse_chunk_offset) !=\n sparse_chunk_offset)) {\n LOG_ERROR(\"SparseChunk resize to %zu failed\", sparse_chunk_offset);\n return IndexError_Runtime;\n }\n }\n\n if (*mutable_doc_cnt() <= id) {\n *mutable_doc_cnt() = id + 1;\n chunk_offset += node_size();\n if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"Chunk resize to %zu failed\", chunk_offset);\n return IndexError_Runtime;\n }\n }\n *mutable_total_sparse_count() += sparse_count;\n\n if (filter_same_key_ || get_vector_enabled_) {\n keys_map_lock_->lock();\n (*keys_map_)[key] = id;\n keys_map_lock_->unlock();\n }\n\n broker_->mark_dirty();\n\n return 0;\n}\n\nvoid HnswSparseStreamerEntity::update_ep_and_level(node_id_t ep,\n level_t level) {\n HnswSparseEntity::update_ep_and_level(ep, level);\n flush_header();\n\n return;\n}\n\nconst HnswSparseEntity::Pointer HnswSparseStreamerEntity::clone() const {\n std::vector node_chunks;\n node_chunks.reserve(node_chunks_.size());\n for (size_t i = 0UL; i < node_chunks_.size(); ++i) {\n node_chunks.emplace_back(node_chunks_[i]->clone());\n if (ailego_unlikely(!node_chunks[i])) {\n LOG_ERROR(\"HnswSparseStreamerEntity get chunk failed in clone\");\n return HnswSparseEntity::Pointer();\n }\n }\n\n std::vector sparse_node_chunks;\n sparse_node_chunks.reserve(sparse_node_chunks_.size());\n for (size_t i = 0UL; i < sparse_node_chunks_.size(); ++i) {\n sparse_node_chunks.emplace_back(sparse_node_chunks_[i]->clone());\n if (ailego_unlikely(!sparse_node_chunks[i])) {\n LOG_ERROR(\"HnswSparseStreamerEntity get sparse chunk failed in clone\");\n return HnswSparseEntity::Pointer();\n }\n }\n\n std::vector upper_neighbor_chunks;\n upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size());\n for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) {\n upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone());\n if (ailego_unlikely(!upper_neighbor_chunks[i])) {\n LOG_ERROR(\"HnswSparseStreamerEntity get chunk failed in clone\");\n return HnswSparseEntity::Pointer();\n }\n }\n\n HnswSparseStreamerEntity *entity =\n new (std::nothrow) HnswSparseStreamerEntity(\n stats_, header(), chunk_size_, node_index_mask_bits_,\n upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_,\n sparse_chunk_size_, upper_neighbor_index_, keys_map_lock_, keys_map_,\n std::move(node_chunks), std::move(upper_neighbor_chunks),\n std::move(sparse_node_chunks), broker_);\n if (ailego_unlikely(!entity)) {\n LOG_ERROR(\"HnswSparseStreamerEntity new failed\");\n }\n return HnswSparseEntity::Pointer(entity);\n}\n\n//! Get sparse vector feature data by key\nint HnswSparseStreamerEntity::get_sparse_vector_by_key(\n key_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const {\n *sparse_count = 0;\n\n auto id = get_id(key);\n if (id == kInvalidNodeId) {\n return IndexError_NoExist;\n }\n\n return get_sparse_vector_by_id(id, sparse_count, sparse_indices_buffer,\n sparse_values_buffer);\n}\n\nint HnswSparseStreamerEntity::get_sparse_vector_by_id(\n node_id_t id, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const {\n IndexStorage::MemoryBlock block;\n get_sparse_data(id, block);\n const void *sparse_data = block.data();\n if (sparse_data == nullptr) {\n return IndexError_InvalidValue;\n }\n\n SparseUtility::ReverseSparseFormat(sparse_data, sparse_count,\n sparse_indices_buffer,\n sparse_values_buffer, sparse_unit_size());\n\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/hnsw_sparse/hnsw_sparse_streamer_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"hnsw_sparse_chunk.h\"\n#include \"hnsw_sparse_entity.h\"\n#include \"hnsw_sparse_index_hash.h\"\n#include \"hnsw_sparse_params.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! HnswSparseStreamerEntity manage vector data, pkey, and node's neighbors\nclass HnswSparseStreamerEntity : public HnswSparseEntity {\n public:\n //! Cleanup\n //! return 0 on success, or errCode in failure\n virtual int cleanup() override;\n\n //! Make a copy of streamer entity, to support thread-safe operation.\n //! The segment in container cannot be read concurrenly\n virtual const HnswSparseEntity::Pointer clone() const override;\n\n //! Get primary key of the node id\n virtual key_t get_key(node_id_t id) const override;\n\n //! Get vector feature data by key\n virtual const void *get_vector_meta(node_id_t id) const override;\n\n virtual int get_vector_meta(const node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n\n //! Get vectors feature data by local ids\n virtual int get_vector_metas(const node_id_t *ids, uint32_t count,\n const void **vecs) const override;\n virtual int get_vector_metas(\n const node_id_t *ids, uint32_t count,\n std::vector &block_vecs) const override;\n\n //! Get vector sparse feature data by chunk index and offset\n virtual const void *get_sparse_data(uint64_t offset,\n uint32_t len) const override;\n\n //! Get sparse data from id\n virtual const void *get_sparse_data(node_id_t id) const override;\n\n virtual int get_sparse_data(uint64_t offset, uint32_t len,\n IndexStorage::MemoryBlock &block) const override;\n virtual int get_sparse_data(node_id_t id,\n IndexStorage::MemoryBlock &block) const override;\n\n //! Get sparse data from vector\n virtual std::pair get_sparse_data_from_vector(\n const void *vec) const override;\n virtual int get_sparse_data_from_vector(const void *vec,\n IndexStorage::MemoryBlock &block,\n int &sparse_length) const override;\n\n //! Get sparse vector feature data by key\n virtual int get_sparse_vector_by_key(\n key_t key, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override;\n\n //! Get sparse vector feature data by id\n virtual int get_sparse_vector_by_id(\n node_id_t id, uint32_t *sparse_count, std::string *sparse_indices_buffer,\n std::string *sparse_values_buffer) const override;\n\n //! Get the node id's neighbors on graph level\n //! Note: the neighbors cannot be modified, using the following\n //! method to get WritableNeighbors if want to\n virtual const Neighbors get_neighbors(level_t level,\n node_id_t id) const override;\n\n\n //! Add vector and key to hnsw entity, and local id will be saved in id\n virtual int add_vector(level_t level, key_t key,\n const std::string &sparse_vec_buffer,\n uint32_t sparse_count, node_id_t *id) override;\n\n //! Add vector and id to hnsw entity\n virtual int add_vector_with_id(level_t level, node_id_t id,\n const std::string &sparse_vec,\n uint32_t sparse_count) override;\n\n virtual int update_neighbors(\n level_t level, node_id_t id,\n const std::vector> &neighbors) override;\n\n //! Replace node id in level's neighbors\n int update_neighbors_dense(\n level_t level, node_id_t id,\n const std::vector> &neighbors);\n\n //! Replace node id in level's neighbors\n int update_neighbors_sparse(\n level_t level, node_id_t id,\n const std::vector> &neighbors);\n\n //! Append neighbor_id to node id neighbors on level\n //! Notice: the caller must be ensure the neighbors not full\n virtual void add_neighbor(level_t level, node_id_t id, uint32_t size,\n node_id_t neighbor_id) override;\n\n //! Dump index by dumper\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n virtual void update_ep_and_level(node_id_t ep, level_t level) override;\n\n public:\n //! Constructor\n HnswSparseStreamerEntity(IndexStreamer::Stats &stats);\n\n //! Destructor\n ~HnswSparseStreamerEntity();\n\n //! Init entity\n int init(uint64_t max_index_size, size_t max_doc_cnt);\n\n //! Flush graph entity to disk\n //! return 0 on success, or errCode in failure\n int flush(uint64_t checkpoint);\n\n //! Open entity from storage\n //! return 0 on success, or errCode in failure\n int open(IndexStorage::Pointer stg, bool check_crc);\n\n //! Close entity\n //! return 0 on success, or errCode in failure\n int close();\n\n //! Set meta information from entity\n int set_index_meta(const IndexMeta &meta) const {\n return IndexHelper::SerializeToStorage(meta, broker_->storage().get());\n }\n\n //! Get meta information from entity\n int get_index_meta(IndexMeta *meta) const {\n return IndexHelper::DeserializeFromStorage(broker_->storage().get(), meta);\n }\n\n //! Set params: chunk size\n inline void set_chunk_size(size_t val) {\n chunk_size_ = val;\n }\n\n //! Set params\n inline void set_filter_same_key(bool val) {\n filter_same_key_ = val;\n }\n\n //! Set params\n inline void set_get_vector(bool val) {\n get_vector_enabled_ = val;\n }\n\n //! Get vector local id by key\n inline node_id_t get_id(key_t key) const {\n keys_map_lock_->lock_shared();\n auto it = keys_map_->find(key);\n keys_map_lock_->unlock_shared();\n return it == keys_map_->end() ? kInvalidNodeId : it->second;\n }\n\n void print_key_map() {\n std::cout << \"key map begins\" << std::endl;\n\n auto iter = keys_map_->begin();\n while (iter != keys_map_->end()) {\n std::cout << \"key: \" << iter->first << \", id: \" << iter->second\n << std::endl;\n ;\n iter++;\n }\n\n std::cout << \"key map ends\" << std::endl;\n }\n\n //! Get neighbors size\n inline size_t neighbors_size() const {\n return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t);\n }\n\n //! Get upper neighbors size\n inline size_t upper_neighbors_size() const {\n return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t);\n }\n\n private:\n union UpperNeighborIndexMeta {\n struct {\n uint32_t level : 4;\n uint32_t index : 28; // index is composite type: chunk idx, and the\n // N th neighbors in chunk, they two composite\n // the 28 bits location\n };\n uint32_t data;\n };\n\n template \n using HashMap = google::dense_hash_map>;\n template \n using HashMapPointer = std::shared_ptr>;\n\n template \n using HashSet = google::dense_hash_set>;\n template \n using HashSetPointer = std::shared_ptr>;\n\n //! upper neighbor index hashmap\n using NIHashMap = HnswSparseIndexHashMap;\n using NIHashMapPointer = std::shared_ptr;\n\n //! Private construct, only be called by clone method\n HnswSparseStreamerEntity(\n IndexStreamer::Stats &stats, const HNSWSparseHeader &hd,\n size_t chunk_size, uint32_t node_index_mask_bits,\n uint32_t upper_neighbor_mask_bits, bool filter_same_key,\n bool get_vector_enabled, uint32_t sparse_chunk_size,\n const NIHashMapPointer &upper_neighbor_index,\n std::shared_ptr &keys_map_lock,\n const HashMapPointer &keys_map,\n std::vector &&node_chunks,\n std::vector &&upper_neighbor_chunks,\n std::vector &&sparse_node_chunks,\n const SparseChunkBroker::Pointer &broker)\n : stats_(stats),\n chunk_size_(chunk_size),\n node_index_mask_bits_(node_index_mask_bits),\n node_cnt_per_chunk_(1UL << node_index_mask_bits_),\n node_index_mask_(node_cnt_per_chunk_ - 1),\n upper_neighbor_mask_bits_(upper_neighbor_mask_bits),\n upper_neighbor_mask_((1U << upper_neighbor_mask_bits_) - 1),\n filter_same_key_(filter_same_key),\n get_vector_enabled_(get_vector_enabled),\n sparse_chunk_size_(sparse_chunk_size),\n upper_neighbor_index_(upper_neighbor_index),\n keys_map_lock_(keys_map_lock),\n keys_map_(keys_map),\n node_chunks_(std::move(node_chunks)),\n upper_neighbor_chunks_(std::move(upper_neighbor_chunks)),\n sparse_node_chunks_(std::move(sparse_node_chunks)),\n broker_(broker) {\n *mutable_header() = hd;\n\n neighbor_size_ = neighbors_size();\n upper_neighbor_size_ = upper_neighbors_size();\n }\n\n //! Called only in searching procedure per context, so no need to lock\n void sync_chunks(SparseChunkBroker::CHUNK_TYPE type, size_t idx,\n std::vector *chunks) const {\n if (ailego_likely(idx < chunks->size())) {\n return;\n }\n for (size_t i = chunks->size(); i <= idx; ++i) {\n auto chunk = broker_->get_chunk(type, i);\n // the storage can ensure get chunk will success after the first get\n ailego_assert_with(!!chunk, \"get chunk failed\");\n chunks->emplace_back(std::move(chunk));\n }\n }\n\n //! return pair: chunk index + chunk offset\n inline std::pair get_vector_chunk_loc(\n node_id_t id) const {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n uint32_t offset = (id & node_index_mask_) * node_size();\n\n sync_chunks(SparseChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_);\n return std::make_pair(chunk_idx, offset);\n }\n\n //! return pair: chunk index + chunk offset\n inline std::pair get_key_chunk_loc(node_id_t id) const {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n uint32_t offset = (id & node_index_mask_) * node_size() + vector_size();\n\n offset += sparse_meta_size();\n\n sync_chunks(SparseChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_);\n return std::make_pair(chunk_idx, offset);\n }\n\n //! return pair: chunk index + chunk offset\n inline std::pair get_sparse_chunk_loc(\n uint32_t chunk_index, uint32_t chunk_offset) const {\n sync_chunks(SparseChunkBroker::CHUNK_TYPE_SPARSE_NODE, chunk_index,\n &sparse_node_chunks_);\n\n return std::make_pair(chunk_index, chunk_offset);\n }\n\n inline std::pair get_upper_neighbor_chunk_loc(\n level_t level, node_id_t id) const {\n auto it = upper_neighbor_index_->find(id);\n ailego_assert_abort(it != upper_neighbor_index_->end(),\n \"Get upper neighbor header failed\");\n auto meta = reinterpret_cast(&it->second);\n uint32_t chunk_idx = (meta->index) >> upper_neighbor_mask_bits_;\n uint32_t offset = (((meta->index) & upper_neighbor_mask_) + level - 1) *\n upper_neighbor_size_;\n sync_chunks(SparseChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, chunk_idx,\n &upper_neighbor_chunks_);\n ailego_assert_abort(chunk_idx < upper_neighbor_chunks_.size(),\n \"invalid chunk idx\");\n ailego_assert_abort(offset < upper_neighbor_chunks_[chunk_idx]->data_size(),\n \"invalid chunk offset\");\n return std::make_pair(chunk_idx, offset);\n }\n\n //! return pair: chunk + chunk offset\n inline std::pair get_neighbor_chunk_loc(\n level_t level, node_id_t id) const {\n if (level == 0UL) {\n uint32_t chunk_idx = id >> node_index_mask_bits_;\n uint32_t offset =\n (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t);\n\n offset += sparse_meta_size();\n\n sync_chunks(SparseChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_);\n ailego_assert_abort(chunk_idx < node_chunks_.size(), \"invalid chunk idx\");\n ailego_assert_abort(offset < node_chunks_[chunk_idx]->data_size(),\n \"invalid chunk offset\");\n return std::make_pair(node_chunks_[chunk_idx].get(), offset);\n } else {\n auto p = get_upper_neighbor_chunk_loc(level, id);\n return std::make_pair(upper_neighbor_chunks_[p.first].get(), p.second);\n }\n }\n\n //! Chunk hnsw index valid\n int check_hnsw_index(const HNSWSparseHeader *hd) const;\n\n size_t get_total_upper_neighbors_size(level_t level) const {\n return level * upper_neighbor_size_;\n }\n\n //! Add upper neighbor header and reserve space for upper neighbor\n int add_upper_neighbor(level_t level, node_id_t id) {\n if (level == 0) {\n return 0;\n }\n SparseChunk::Pointer chunk;\n uint64_t chunk_offset = -1UL;\n size_t neighbors_size = get_total_upper_neighbors_size(level);\n uint64_t chunk_index = upper_neighbor_chunks_.size() - 1UL;\n if (chunk_index == -1UL ||\n (upper_neighbor_chunks_[chunk_index]->padding_size() <\n neighbors_size)) { // no space left and need to alloc\n chunk_index++;\n if (ailego_unlikely(upper_neighbor_chunks_.capacity() ==\n upper_neighbor_chunks_.size())) {\n LOG_ERROR(\"add upper neighbor failed for no memory quota\");\n return IndexError_IndexFull;\n }\n auto p =\n broker_->alloc_chunk(SparseChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR,\n chunk_index, upper_neighbor_chunk_size_);\n if (ailego_unlikely(p.first != 0)) {\n LOG_ERROR(\"Alloc data chunk failed\");\n return p.first;\n }\n chunk = p.second;\n chunk_offset = 0UL;\n upper_neighbor_chunks_.emplace_back(chunk);\n } else {\n chunk = upper_neighbor_chunks_[chunk_index];\n chunk_offset = chunk->data_size();\n }\n ailego_assert_with((size_t)level < kMaxGraphLayers, \"invalid level\");\n ailego_assert_with(chunk_offset % upper_neighbor_size_ == 0,\n \"invalid offset\");\n ailego_assert_with((chunk_offset / upper_neighbor_size_) <\n (1U << upper_neighbor_mask_bits_),\n \"invalid offset\");\n ailego_assert_with(chunk_index < (1U << (28 - upper_neighbor_mask_bits_)),\n \"invalid chunk index\");\n UpperNeighborIndexMeta meta;\n meta.level = level;\n meta.index = (chunk_index << upper_neighbor_mask_bits_) |\n (chunk_offset / upper_neighbor_size_);\n chunk_offset += upper_neighbor_size_ * level;\n if (ailego_unlikely(!upper_neighbor_index_->insert(id, meta.data))) {\n LOG_ERROR(\"HashMap insert value failed\");\n return IndexError_Runtime;\n }\n\n if (ailego_unlikely(chunk->resize(chunk_offset) != chunk_offset)) {\n LOG_ERROR(\"SparseChunk resize to %zu failed\", (size_t)chunk_offset);\n return IndexError_Runtime;\n }\n\n return 0;\n }\n\n size_t estimate_doc_capacity() const {\n return node_chunks_.capacity() * node_cnt_per_chunk_;\n }\n\n int init_chunk_params(size_t max_index_size) {\n sparse_chunk_size_ = AlignPageSize(chunk_size_);\n\n node_cnt_per_chunk_ = std::max(1, chunk_size_ / node_size());\n //! align node cnt per chunk to pow of 2\n node_index_mask_bits_ = std::ceil(std::log2(node_cnt_per_chunk_));\n node_cnt_per_chunk_ = 1UL << node_index_mask_bits_;\n chunk_size_ = AlignPageSize(node_cnt_per_chunk_ * node_size());\n node_index_mask_ = node_cnt_per_chunk_ - 1;\n\n if (max_index_size == 0UL) {\n max_index_size_ = chunk_size_ * kDefaultMaxChunkCnt;\n } else {\n max_index_size_ = max_index_size;\n }\n\n //! To get a balanced upper neighbor chunk size.\n //! If the upper chunk size is equal to node chunk size, it may waste\n //! upper neighbor chunk space; if the upper neighbor chunk size is too\n //! small, the will need large upper neighbor chunks index space. So to\n //! get a balanced ratio be sqrt of the node/neighbor size ratio\n float ratio =\n std::sqrt(node_size() * scaling_factor() * 1.0f / upper_neighbor_size_);\n upper_neighbor_chunk_size_ =\n AlignPageSize(std::max(get_total_upper_neighbors_size(kMaxGraphLayers),\n static_cast(chunk_size_ / ratio)));\n upper_neighbor_mask_bits_ =\n std::ceil(std::log2(upper_neighbor_chunk_size_ / upper_neighbor_size_));\n upper_neighbor_mask_ = (1 << upper_neighbor_mask_bits_) - 1;\n\n size_t max_node_chunk_cnt = std::ceil(max_index_size_ / chunk_size_);\n size_t max_upper_chunk_cnt = std::ceil(\n (max_node_chunk_cnt * node_cnt_per_chunk_ * 1.0f / scaling_factor()) /\n (upper_neighbor_chunk_size_ / upper_neighbor_size_));\n max_upper_chunk_cnt =\n max_upper_chunk_cnt + std::ceil(max_upper_chunk_cnt / scaling_factor());\n\n //! reserve space to avoid memmove in chunks vector emplace chunk, so\n //! as to lock-free in reading chunk\n node_chunks_.reserve(max_node_chunk_cnt);\n sparse_node_chunks_.reserve(max_node_chunk_cnt);\n upper_neighbor_chunks_.reserve(max_upper_chunk_cnt);\n\n LOG_DEBUG(\n \"Settings: nodeSize=%zu chunkSize=%u upperNeighborSize=%u \"\n \"upperNeighborChunkSize=%u \"\n \"nodeCntPerChunk=%u maxChunkCnt=%zu maxNeighborChunkCnt=%zu \"\n \"maxIndexSize=%zu ratio=%.3f\",\n node_size(), chunk_size_, upper_neighbor_size_,\n upper_neighbor_chunk_size_, node_cnt_per_chunk_, max_node_chunk_cnt,\n max_upper_chunk_cnt, max_index_size_, ratio);\n\n return 0;\n }\n\n //! Init node chunk and neighbor chunks\n int init_chunks(const SparseChunk::Pointer &header_chunk);\n\n int flush_header(void) {\n if (!broker_->dirty()) {\n // do not need to flush\n return 0;\n }\n auto header_chunk =\n broker_->get_chunk(SparseChunkBroker::CHUNK_TYPE_HEADER,\n SparseChunkBroker::kDefaultChunkSeqId);\n if (ailego_unlikely(!header_chunk)) {\n LOG_ERROR(\"get header chunk failed\");\n return IndexError_Runtime;\n }\n size_t size = header_chunk->write(0UL, &header(), header_size());\n if (ailego_unlikely(size != header_size())) {\n LOG_ERROR(\"Write header chunk failed\");\n return IndexError_WriteData;\n }\n\n return 0;\n }\n\n private:\n HnswSparseStreamerEntity(const HnswSparseStreamerEntity &) = delete;\n HnswSparseStreamerEntity &operator=(const HnswSparseStreamerEntity &) =\n delete;\n static constexpr uint64_t kUpperHashMemoryInflateRatio = 2.0f;\n\n private:\n IndexStreamer::Stats &stats_;\n HNSWSparseHeader header_{};\n std::mutex mutex_{};\n size_t max_index_size_{0UL};\n uint32_t chunk_size_{kDefaultChunkSize};\n uint32_t upper_neighbor_chunk_size_{kDefaultChunkSize};\n uint32_t node_index_mask_bits_{0U};\n uint32_t node_cnt_per_chunk_{0U};\n uint32_t node_index_mask_{0U};\n uint32_t neighbor_size_{0U};\n uint32_t upper_neighbor_size_{0U};\n //! UpperNeighborIndex.index composite chunkIdx and offset in chunk by the\n //! following mask\n uint32_t upper_neighbor_mask_bits_{0U};\n uint32_t upper_neighbor_mask_{0U};\n bool filter_same_key_{false};\n bool get_vector_enabled_{false};\n\n uint32_t sparse_chunk_size_{kDefaultChunkSize};\n\n NIHashMapPointer upper_neighbor_index_{};\n\n mutable std::shared_ptr keys_map_lock_{};\n HashMapPointer keys_map_{};\n\n //! the chunks will be changed in searcher, so need mutable\n //! data chunk include: vector, key, level 0 neighbors\n mutable std::vector node_chunks_{};\n //! upper neighbor chunk inlude: UpperNeighborHeader + (1~level) neighbors\n mutable std::vector upper_neighbor_chunks_{};\n //! chunk that holds up sparse part\n mutable std::vector sparse_node_chunks_{};\n\n SparseChunkBroker::Pointer broker_{}; // chunk broker\n};\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/ivf/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\ncc_library(\n NAME core_knn_ivf STATIC SHARED STRICT ALWAYS_LINK\n SRCS *.cc\n LIBS zvec_ailego core_framework core_knn_cluster\n INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )\n" }, { "path": "src/core/algorithm/ivf/ivf_builder.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"ivf_builder.h\"\n#include \n#include \n#include \"algorithm/cluster/cluster_params.h\"\n#include \"ivf_dumper.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! IndexHolder support filtered by vector labels\n */\nclass LabelFilteredIndexHolder : public IndexHolder {\n public:\n /*! Index Holder Iterator\n */\n class Iterator : public IndexHolder::Iterator {\n public:\n //! Index Holder Iterator Pointer\n typedef std::unique_ptr Pointer;\n\n //! Constructor\n Iterator(const IVFBuilder::RandomAccessIndexHolder::Pointer &holder,\n const std::vector *elems)\n : holder_(holder), elems_(elems) {}\n\n //! Destructor\n virtual ~Iterator(void) {}\n\n //! Retrieve pointer of data\n virtual const void *data(void) const override {\n return holder_->element((*elems_)[index_]);\n }\n\n //! Test if the iterator is valid\n virtual bool is_valid(void) const override {\n return index_ < elems_->size();\n }\n\n //! Retrieve primary key\n virtual uint64_t key(void) const override {\n return (*elems_)[index_];\n }\n\n //! Next iterator\n virtual void next(void) override {\n ++index_;\n }\n\n private:\n //! Members\n const IVFBuilder::RandomAccessIndexHolder::Pointer holder_{nullptr};\n const std::vector *elems_{nullptr};\n size_t index_{0};\n };\n\n //! Constructor\n LabelFilteredIndexHolder(\n const IVFBuilder::RandomAccessIndexHolder::Pointer &holder,\n const std::vector &items)\n : holder_(holder), elems_(&items) {}\n\n //! Retrieve count of elements in holder (-1 indicates unknown)\n virtual size_t count(void) const override {\n return elems_->size();\n }\n\n //! Retrieve dimension\n virtual size_t dimension(void) const override {\n return holder_->dimension();\n }\n\n //! Retrieve type information\n virtual IndexMeta::DataType data_type(void) const override {\n return holder_->data_type();\n }\n\n //! Retrieve element size in bytes\n virtual size_t element_size(void) const override {\n return holder_->element_size();\n }\n\n //! Retrieve if it can multi-pass\n virtual bool multipass(void) const override {\n return true;\n }\n\n //! Create a new iterator\n virtual IndexHolder::Iterator::Pointer create_iterator(void) override {\n return IndexHolder::Iterator::Pointer(\n new LabelFilteredIndexHolder::Iterator(holder_, elems_));\n }\n\n private:\n //! Members\n const IVFBuilder::RandomAccessIndexHolder::Pointer holder_{};\n const std::vector *elems_{};\n};\n\nIVFBuilder::IVFBuilder() {}\n\nIVFBuilder::~IVFBuilder() {\n this->cleanup();\n}\n\nint IVFBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) {\n LOG_INFO(\"Begin IVFBuilder::init!\");\n\n if (state_ != INIT) {\n LOG_ERROR(\"IVFBuilder state wrong. state=%d\", state_);\n return IndexError_Logic;\n }\n\n meta_ = meta;\n converted_meta_ = meta;\n quantized_meta_ = meta;\n // Clear the converter/reformer params for external transforms\n converted_meta_.set_reformer(std::string(), 0, ailego::Params());\n converted_meta_.set_converter(std::string(), 0, ailego::Params());\n quantized_meta_.set_reformer(std::string(), 0, ailego::Params());\n quantized_meta_.set_converter(std::string(), 0, ailego::Params());\n params_ = params;\n\n if (!IndexFactory::HasMetric(meta_.metric_name())) {\n LOG_ERROR(\"Metric %s not exist\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n\n int ret = parse_centroids_num(params);\n ivf_check_with_msg(ret, \"Failed to parse centroids, ret=%d\", ret);\n\n ret = parse_clustering_params(params);\n ivf_check_with_msg(ret, \"Failed to parse clustering params, ret=%d\", ret);\n\n ret = parse_general_params(params);\n ivf_check_with_msg(ret, \"Failed to parse general params, ret=%d\", ret);\n\n LOG_INFO(\"End IVFBuilder::init!\");\n\n LOG_DEBUG(\n \"Converter=%s Quantizer=%s Optimizer=%s \"\n \"OptimizerQuantizer=%s QuantizeByCentroid=%u StoreFeatures=%u \"\n \"ClusterClass=%s TrainSamplesCount=%u TrainSampleRatio=%f \"\n \"BlockVectorCount=%u\",\n params.get_as_string(PARAM_IVF_BUILDER_CONVERTER_CLASS).c_str(),\n params.get_as_string(PARAM_IVF_BUILDER_QUANTIZER_CLASS).c_str(),\n params.get_as_string(PARAM_IVF_BUILDER_OPTIMIZER_CLASS).c_str(),\n params.get_as_string(PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_CLASS).c_str(),\n params.get_as_bool(PARAM_IVF_BUILDER_QUANTIZE_BY_CENTROID),\n params.get_as_bool(PARAM_IVF_BUILDER_STORE_ORIGINAL_FEATURES),\n params.get_as_string(PARAM_IVF_BUILDER_CLUSTER_CLASS).c_str(),\n params.get_as_uint32(PARAM_IVF_BUILDER_TRAIN_SAMPLE_COUNT),\n params.get_as_float(PARAM_IVF_BUILDER_TRAIN_SAMPLE_RATIO),\n block_vector_count_);\n\n state_ = INITED;\n return 0;\n}\n\nint IVFBuilder::cleanup(void) {\n LOG_INFO(\"Begin IVFBuilder::cleanup\");\n\n state_ = INIT;\n stats_.clear_attributes();\n stats_.set_built_costtime(0u);\n stats_.set_built_count(0u);\n stats_.set_discarded_count(0u);\n stats_.set_dumped_costtime(0u);\n stats_.set_dumped_count(0u);\n stats_.set_trained_costtime(0u);\n stats_.set_trained_count(0u);\n\n centroid_num_vec_.clear();\n cluster_class_.clear();\n converter_class_.clear();\n cluster_params_.clear();\n\n labels_.clear();\n centroid_index_.reset();\n holder_.reset();\n converted_meta_ = meta_;\n converter_.reset();\n quantized_meta_ = meta_;\n quantizers_.clear();\n\n error_ = false;\n err_code_ = 0;\n\n thread_count_ = 0;\n sample_count_ = 0;\n cluster_auto_tuning_ = false;\n store_original_features_ = false;\n quantize_by_centroid_ = false;\n\n LOG_INFO(\"End IVFBuilder::cleanup\");\n\n return 0;\n}\n\nint IVFBuilder::train(IndexThreads::Pointer threads,\n IndexHolder::Pointer holder) {\n LOG_INFO(\"Begin IVFBuilder::train with holder\");\n if (state_ != INITED) {\n LOG_ERROR(\"IVFBuilder train failed, wrong state=%d\", state_);\n return IndexError_Runtime;\n }\n\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n ailego::ElapsedTime timer;\n if (!holder || holder->count() == 0) {\n LOG_ERROR(\"Input holder is nullptr or empty while train index\");\n return IndexError_InvalidArgument;\n }\n if (!holder->is_matched(meta_)) {\n LOG_ERROR(\"Input holder doesn't match index meta while train index\");\n return IndexError_Mismatch;\n }\n\n if (converter_) {\n int ret = IndexConverter::TrainAndTransform(converter_, std::move(holder));\n ivf_check_with_msg(ret, \"Failed to train or transform by converter %s\",\n converter_->name().c_str());\n converted_meta_ = converter_->meta();\n holder = converter_->result();\n }\n\n ailego::Params train_params;\n int ret = prepare_trainer_params(train_params);\n ivf_check_with_msg(ret, \"Failed to prepare trainer params, ret=%d\", ret);\n\n IndexTrainer::Pointer trainer =\n IndexFactory::CreateTrainer(\"StratifiedClusterTrainer\");\n ivf_assert_with_msg(trainer, IndexError_NoExist, \"Failed to create trainer\");\n\n ret = trainer->init(converted_meta_, train_params);\n ivf_check_with_msg(ret, \"Trainer init failed with ret %d\", ret);\n\n ret = trainer->train(std::move(threads), std::move(holder));\n ivf_check_with_msg(ret, \"Trainer train failed with ret %d\", ret);\n\n ret = this->train(trainer);\n ivf_check_error_code(ret);\n\n stats_.set_trained_costtime(timer.milli_seconds());\n\n LOG_INFO(\"End IVFBuilder::train with holder\");\n\n state_ = TRAINED;\n return 0;\n}\n\nint IVFBuilder::train(const IndexTrainer::Pointer &trainer) {\n LOG_DEBUG(\"Begin IVFBuilder::train by trainer\");\n ailego::ElapsedTime timer;\n\n if (state_ != INITED) {\n LOG_ERROR(\"IVFBuilder train failed, wrong state=%d\", state_);\n return IndexError_Runtime;\n }\n\n if (!trainer) {\n LOG_ERROR(\"Input trainer is nullptr while train index\");\n return IndexError_InvalidArgument;\n }\n\n IndexCluster::CentroidList centroid_list;\n IndexBundle::Pointer boundle = trainer->indexes();\n int ret = IndexCluster::Deserialize(trainer->meta(), boundle, ¢roid_list);\n ivf_check_with_msg(ret, \"Failed to deserialize index\");\n\n const IndexMeta &meta = trainer->meta();\n if (meta.data_type() != converted_meta_.data_type() ||\n meta.metric_name().compare(converted_meta_.metric_name()) != 0 ||\n meta.element_size() != converted_meta_.element_size()) {\n if (meta.converter_name() != converter_class_) {\n LOG_ERROR(\"Input trainer doesn't match index meta while train index\");\n return IndexError_Mismatch;\n }\n //! Create converter from trainer params\n LOG_INFO(\"Train IVFBuilder by trainer with converter\");\n converter_ = CreateAndInitConverter(meta_, meta.converter_name(),\n meta.converter_params());\n ivf_assert(converter_, IndexError_Runtime);\n converted_meta_ = meta;\n }\n\n centroid_index_ = std::make_shared();\n if (!centroid_index_) {\n return IndexError_NoMemory;\n }\n ret = centroid_index_->init(converted_meta_, params_);\n ivf_check_error_code(ret);\n\n ret = centroid_index_->build(centroid_list);\n ivf_check_with_msg(ret, \"Failed to build centroid index\");\n\n if (params_.has(PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_CLASS)) {\n //! Quantize the centroids for searcher\n searcher_centroid_index_ = std::make_shared();\n if (!searcher_centroid_index_) {\n return IndexError_NoMemory;\n }\n ailego::Params params;\n params_.get(PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_PARAMS, ¶ms);\n searcher_centroid_index_->set_quantizer(\n params_.get_as_string(PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_CLASS),\n params);\n ret = searcher_centroid_index_->init(converted_meta_, params_);\n ivf_check_error_code(ret);\n\n ret = searcher_centroid_index_->build(centroid_list);\n ivf_check_with_msg(ret, \"Failed to build centroid index\");\n }\n\n stats_.set_trained_costtime(timer.milli_seconds());\n\n LOG_DEBUG(\"End IVFBuilder::train by trainer\");\n\n state_ = TRAINED;\n return 0;\n}\n\nint IVFBuilder::build(IndexThreads::Pointer threads,\n IndexHolder::Pointer holder) {\n LOG_INFO(\"Begin IVFBuilder::build!\");\n\n if (state_ != TRAINED) {\n LOG_ERROR(\"Train the index first before build\");\n return IndexError_Runtime;\n }\n\n ailego::ElapsedTime timer;\n if (!holder || holder->count() == 0) {\n LOG_ERROR(\"Input holder is nullptr or empty while building index\");\n return IndexError_InvalidArgument;\n }\n\n if (!holder->is_matched(meta_)) {\n LOG_ERROR(\"Input holder doesn't match index meta while building index\");\n return IndexError_Mismatch;\n }\n if (!threads) {\n threads = std::make_shared(thread_count_, false);\n if (!threads) {\n return IndexError_NoMemory;\n }\n }\n\n holder_ = std::make_shared(meta_);\n if (!holder_) {\n return IndexError_NoMemory;\n }\n if (holder->count() > 0) {\n holder_->reserve(holder->count());\n }\n for (auto iter = holder->create_iterator(); iter && iter->is_valid();\n iter->next()) {\n holder_->emplace(iter->key(), iter->data());\n }\n\n // Holder is not needed, cleanup it.\n holder.reset();\n\n IndexHolder::Pointer converted_holder = holder_;\n if (converter_) {\n int ret = converter_->transform(holder_);\n ivf_check_with_msg(ret, \"Failed to transform by converter %s\",\n converter_->name().c_str());\n converted_holder = converter_->result();\n }\n\n labels_.resize(centroid_index_->centroids_count());\n int ret = this->build_label_index(threads.get(), converted_holder);\n ivf_check_with_msg(ret, \"Failed to build index for %s\",\n IndexError::What(ret));\n\n ret = this->prepare_quantizer(threads.get());\n ivf_check_error_code(ret);\n\n stats_.set_built_costtime(timer.milli_seconds());\n\n LOG_INFO(\"End IVFBuilder::build\");\n\n state_ = BUILT;\n return 0;\n}\n\nint IVFBuilder::dump(const IndexDumper::Pointer &dumper) {\n LOG_INFO(\"Begin IVFBuilder::dump\");\n\n if (state_ != BUILT) {\n LOG_ERROR(\"Build the index before dump QC Index\");\n return IndexError_Runtime;\n }\n\n ailego::ElapsedTime timer;\n int ret = this->dump_index(dumper);\n ivf_check_with_msg(ret, \"Failed to dump index with ret=%d\", ret);\n\n // the fitting function for the follow points: 1000000(0.02) 10000000(0.01)\n // 50000000(0.005) 100000000(0.001)\n float scan_ratio = -0.004 * std::log(holder_->count()) + 0.0751;\n scan_ratio = std::max(scan_ratio, 0.0001f);\n\n // Set Searcher Params\n ailego::Params params;\n params.set(PARAM_IVF_SEARCHER_SCAN_RATIO, scan_ratio);\n meta_.set_searcher(\"IVFSearcher\", 0, std::move(params));\n meta_.set_builder(\"IVFBuilder\", 0, std::move(params_));\n\n ret = IndexHelper::SerializeToDumper(meta_, dumper.get());\n if (ret != 0) {\n LOG_ERROR(\"Failed to serialize meta into dumper.\");\n return ret;\n }\n\n stats_.set_discarded_count(stats_.built_count() - stats_.dumped_count());\n stats_.set_dumped_costtime(timer.milli_seconds());\n\n LOG_INFO(\"End IVFBuilder::dump\");\n\n return 0;\n}\n\nint IVFBuilder::CheckAndUpdateMajorOrder(IndexMeta &meta) {\n const std::string &metric_name = meta.metric_name();\n auto metric = IndexFactory::CreateMetric(metric_name);\n if (!metric) {\n LOG_ERROR(\"CreateMetric %s failed\", metric_name.c_str());\n return IndexError_InvalidArgument;\n }\n int ret = metric->init(meta, meta.metric_params());\n ivf_check_with_msg(ret, \"IndexMetric %s init failed\", metric_name.c_str());\n\n bool support_column_major = true;\n for (size_t m = 32; m != 0; m /= 2) {\n for (size_t n = m; n != 0; n /= 2) {\n if (metric->distance_matrix(m, n) == nullptr) {\n support_column_major = false;\n break;\n }\n }\n if (!support_column_major) {\n break;\n }\n }\n support_column_major &=\n meta.element_size() % IndexMeta::AlignSizeof(meta.data_type()) == 0;\n\n if (meta.major_order() == IndexMeta::MO_UNDEFINED) {\n if (support_column_major && meta.dimension() <= 512) {\n meta.set_major_order(IndexMeta::MO_COLUMN);\n } else {\n meta.set_major_order(IndexMeta::MO_ROW);\n }\n } else {\n if (!support_column_major && meta.major_order() == IndexMeta::MO_COLUMN) {\n LOG_WARN(\n \"Index Metric %s Unsupported \"\n \"Column Major Order\",\n metric_name.c_str());\n return IndexError_Unsupported;\n }\n }\n\n if (block_vector_count_ * quantized_meta_.element_size() % 32 != 0) {\n LOG_ERROR(\n \"block_vector_count * quantized_element_size not align with 32 bytes.\");\n return IndexError_InvalidArgument;\n }\n\n return 0;\n}\n\nint IVFBuilder::parse_centroids_num(const ailego::Params ¶ms) {\n std::string centroids_num =\n params.get_as_string(PARAM_IVF_BUILDER_CENTROID_COUNT);\n if (centroids_num.empty()) {\n LOG_ERROR(\"Param %s is required\", PARAM_IVF_BUILDER_CENTROID_COUNT.c_str());\n return IndexError_InvalidArgument;\n }\n\n std::vector centroid_str_vec;\n ailego::StringHelper::Split(centroids_num, CENTROID_SEPERATOR,\n ¢roid_str_vec);\n size_t level_cnt = centroid_str_vec.size();\n if ((level_cnt <= 0) || (level_cnt > 2)) {\n LOG_ERROR(\"Centroids level count must be [1,2]\");\n return IndexError_InvalidArgument;\n }\n\n for (size_t idx = 0; idx < level_cnt; ++idx) {\n uint32_t centroid_cnt = 0;\n if (!ailego::StringHelper::ToUint32(centroid_str_vec[idx], ¢roid_cnt)) {\n LOG_ERROR(\"Invalid centroids count %s\", centroid_str_vec[idx].c_str());\n return IndexError_InvalidArgument;\n }\n centroid_num_vec_.push_back(centroid_cnt);\n }\n\n return 0;\n}\n\nint IVFBuilder::parse_clustering_params(const ailego::Params ¶ms) {\n params.get(PARAM_IVF_BUILDER_CLUSTER_AUTO_TUNING, &cluster_auto_tuning_);\n\n cluster_class_ = params.get_as_string(PARAM_IVF_BUILDER_CLUSTER_CLASS);\n if (cluster_class_.empty()) {\n // OptKmeansCluster does not support custom metric\n cluster_class_ = meta_.metric_name() == kMipsMetricName\n ? \"KmeansCluster\"\n : \"OptKmeansCluster\";\n LOG_INFO(\"Using [%s] as default cluster class\", cluster_class_.c_str());\n }\n for (size_t i = 1; i <= centroid_num_vec_.size(); ++i) {\n std::string level_params_key =\n PARAM_IVF_BUILDER_CLUSTER_PARAMS_IN_LEVEL_PREFIX + std::to_string(i);\n ailego::Params level_params;\n params.get(level_params_key, &level_params);\n cluster_params_.push_back(level_params);\n }\n\n return 0;\n}\n\nint IVFBuilder::parse_general_params(const ailego::Params ¶ms) {\n thread_count_ = params.get_as_uint32(PARAM_IVF_BUILDER_THREAD_COUNT);\n sample_count_ = params.get_as_uint32(PARAM_IVF_BUILDER_TRAIN_SAMPLE_COUNT);\n sample_ratio_ = params.get_as_float(PARAM_IVF_BUILDER_TRAIN_SAMPLE_RATIO);\n\n params.get(PARAM_IVF_BUILDER_QUANTIZE_BY_CENTROID, &quantize_by_centroid_);\n params.get(PARAM_IVF_BUILDER_STORE_ORIGINAL_FEATURES,\n &store_original_features_);\n\n //! Prepare Converter for training\n if (meta_.metric_name() == kIPMetricName) {\n converter_class_ = kMipsConverterName;\n }\n params.get(PARAM_IVF_BUILDER_CONVERTER_CLASS, &converter_class_);\n if (!converter_class_.empty()) {\n ailego::Params converter_params;\n params_.get(PARAM_IVF_BUILDER_CONVERTER_PARAMS, &converter_params);\n converter_ =\n CreateAndInitConverter(meta_, converter_class_, converter_params);\n ivf_assert(converter_, IndexError_NoExist);\n }\n\n params_.get(PARAM_IVF_BUILDER_BLOCK_VECTOR_COUNT, &block_vector_count_);\n if (block_vector_count_ == 0) {\n block_vector_count_ = kDefaultBlockCount;\n }\n if (block_vector_count_ > kDefaultBlockCount ||\n block_vector_count_ & (block_vector_count_ - 1)) {\n LOG_ERROR(\"block_vector_count only can be [1|2|4|8|16|32].\");\n return IndexError_InvalidArgument;\n }\n if (block_vector_count_ * meta_.element_size() % 32 != 0) {\n LOG_ERROR(\"block_vector_count * element_size not align with 32 bytes.\");\n return IndexError_InvalidArgument;\n }\n return 0;\n}\n\nint IVFBuilder::prepare_trainer_params(ailego::Params ¶ms) {\n params.set(STRATIFIED_TRAINER_SAMPLE_COUNT, sample_count_);\n params.set(STRATIFIED_TRAINER_SAMPLE_RATIO, sample_ratio_);\n params.set(STRATIFIED_TRAINER_THREAD_COUNT, thread_count_);\n params.set(STRATIFIED_TRAINER_AUTOAUNE, cluster_auto_tuning_);\n if (centroid_num_vec_.empty()) {\n LOG_ERROR(\"Centroids no specified.\");\n return IndexError_InvalidArgument;\n }\n std::string cluster_count = std::to_string(centroid_num_vec_[0]);\n if (centroid_num_vec_.size() > 1) {\n cluster_count +=\n (CENTROID_SEPERATOR + std::to_string(centroid_num_vec_[1]));\n }\n params.set(STRATIFIED_TRAINER_CLUSTER_COUNT, cluster_count);\n\n for (size_t i = 1; i <= cluster_params_.size(); ++i) {\n std::string level_params_key =\n STRATIFIED_TRAINER_PARAMS_IN_LEVEL_PREFIX + std::to_string(i);\n params.set(level_params_key, cluster_params_[i - 1]);\n }\n params.set(STRATIFIED_TRAINER_CLASS_NAME, cluster_class_);\n\n return 0;\n}\n\nint IVFBuilder::build_label_index(IndexThreads *threads,\n const IndexHolder::Pointer &holder) {\n auto iter = holder->create_iterator();\n if (!iter) {\n LOG_ERROR(\"Create iterator for holder failed\");\n return IndexError_Runtime;\n }\n\n auto task_group = threads->make_group();\n if (!task_group) {\n LOG_ERROR(\"Failed to create task group\");\n return IndexError_Runtime;\n }\n\n size_t id = 0UL;\n AILEGO_DEFER([&]() {\n task_group->wait_finish();\n stats_.set_built_count(id);\n LOG_INFO(\"Finished building, total=%zu\", id);\n });\n\n size_t elem_size = holder->element_size();\n std::shared_ptr vectors = std::make_shared();\n ivf_assert(vectors, IndexError_NoMemory);\n for (; iter && iter->is_valid(); iter->next()) {\n ivf_assert(!error_, err_code_);\n vectors->emplace_back(iter->data(), elem_size, id);\n id++;\n if (vectors->size() == kBatchSize || id == holder_->count()) {\n auto task = ailego::Closure ::New(const_cast(this),\n &IVFBuilder::label, vectors);\n task_group->submit(std::move(task));\n vectors = std::make_shared();\n ivf_assert(vectors, IndexError_NoMemory);\n vectors->reserve(kBatchSize);\n }\n if (!(id & 0xFFFFF)) {\n LOG_INFO(\"Current built count:%zu\", id);\n }\n }\n ailego_assert_with(vectors->size() == 0, \"invalid size\");\n\n return err_code_;\n}\n\nint IVFBuilder::dump_index(const IndexDumper::Pointer &dumper) {\n int ret = CheckAndUpdateMajorOrder(quantized_meta_);\n ivf_check_error_code(ret);\n\n IVFDumper::Pointer ivf_dumper = std::make_shared(\n quantized_meta_, dumper, centroid_index_->centroids_count(),\n block_vector_count_);\n if (!ivf_dumper) {\n LOG_ERROR(\"Alloc IVFDumper failed\");\n return IndexError_NoMemory;\n }\n\n //! Dump inverted vectors\n std::vector dumped_ids;\n std::function record_dumped_id = [&](uint32_t) {};\n if (store_original_features_) {\n dumped_ids.reserve(holder_->count());\n record_dumped_id = [&](uint32_t id) { dumped_ids.emplace_back(id); };\n }\n if (quantizers_.size() == 0) {\n //! No quantizer for inverted vectors\n for (size_t i = 0; i < centroid_index_->centroids_count(); ++i) {\n ailego_assert_with(i < labels_.size(), \"Index Overflow\");\n for (size_t j = 0; j < labels_[i].size(); ++j) {\n auto id = labels_[i][j];\n record_dumped_id(id);\n ret = ivf_dumper->dump_inverted_vector(i, holder_->key(id),\n holder_->element(id));\n ivf_check_error_code(ret);\n }\n }\n } else {\n for (size_t i = 0; i < centroid_index_->centroids_count(); ++i) {\n ailego_assert_with(i < labels_.size(), \"Index Overflow\");\n auto holder =\n std::make_shared(holder_, labels_[i]);\n if (!holder) {\n return IndexError_NoMemory;\n }\n auto quantizer = quantize_by_centroid_ ? quantizers_[i] : quantizers_[0];\n ret = quantizer->transform(holder);\n ivf_check_error_code(ret);\n\n auto iter = quantizer->result()->create_iterator();\n for (; iter->is_valid(); iter->next()) {\n uint32_t id = iter->key();\n record_dumped_id(id);\n ret =\n ivf_dumper->dump_inverted_vector(i, holder_->key(id), iter->data());\n ivf_check_error_code(ret);\n }\n }\n }\n\n ret = ivf_dumper->dump_inverted_vector_finished();\n ivf_check_error_code(ret);\n\n ret = ivf_dumper->dump_quantizer_params(quantizers_);\n ivf_check_error_code(ret);\n\n auto centroid_index =\n searcher_centroid_index_ ? searcher_centroid_index_ : centroid_index_;\n ret = ivf_dumper->dump_centroid_index(centroid_index->data(),\n centroid_index->size());\n ivf_check_with_msg(ret, \"Failed to dump CentroidIndex\");\n\n if (store_original_features_) {\n for (size_t i = 0; i < dumped_ids.size(); ++i) {\n ret = ivf_dumper->dump_original_vector(holder_->element(dumped_ids[i]),\n holder_->element_size());\n ivf_check_error_code(ret);\n }\n }\n\n stats_.set_dumped_count(stats_.dumped_count() + ivf_dumper->dumped_count());\n\n return 0;\n}\n\nint IVFBuilder::prepare_quantizer(IndexThreads *threads) {\n std::string quantizer_name;\n params_.get(PARAM_IVF_BUILDER_QUANTIZER_CLASS, &quantizer_name);\n if (quantizer_name.empty()) {\n return 0;\n }\n\n //! Prepare Quantizers for inverted index\n ailego::Params quantizer_params;\n params_.get(PARAM_IVF_BUILDER_QUANTIZER_PARAMS, &quantizer_params);\n if (((quantizer_name != kInt8QuantizerName &&\n quantizer_name != kInt4QuantizerName) ||\n meta_.metric_name() != kIPMetricName) &&\n quantize_by_centroid_) {\n LOG_WARN(\"%s is supported in InnerProduct only\",\n PARAM_IVF_BUILDER_QUANTIZE_BY_CENTROID.c_str());\n quantize_by_centroid_ = false;\n }\n if (quantizer_name == kInt4QuantizerName && meta_.dimension() & 0x1) {\n LOG_ERROR(\"Unsupport quantizer=%s for dim=%u\", kInt4QuantizerName,\n meta_.dimension());\n return IndexError_Unsupported;\n }\n\n int ret = 0;\n auto create_and_init_quantizer = [&]() {\n auto quantizer = IndexFactory::CreateConverter(quantizer_name);\n if (!quantizer) {\n LOG_ERROR(\"Failed to create converter %s\", quantizer_name.c_str());\n ret = IndexError_NoExist;\n return IndexConverter::Pointer();\n }\n ret = quantizer->init(meta_, quantizer_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize converter %s for %s\",\n quantizer_name.c_str(), IndexError::What(ret));\n return IndexConverter::Pointer();\n }\n return quantizer;\n };\n for (size_t i = 0; i < centroid_index_->centroids_count(); ++i) {\n quantizers_.emplace_back(create_and_init_quantizer());\n ivf_check_error_code(ret);\n if (!quantize_by_centroid_) {\n break;\n }\n }\n\n //! Train the quantizers\n auto train_data = [&](size_t i) {\n IndexHolder::Pointer holder = holder_;\n size_t idx = 0;\n if (quantize_by_centroid_) {\n holder = std::make_shared(holder_, labels_[i]);\n if (!holder && !error_.exchange(true)) {\n err_code_ = IndexError_NoMemory;\n return;\n }\n idx = i;\n }\n if (holder->count() == 0) {\n return;\n }\n ret = quantizers_[idx]->train(holder);\n if (ret != 0) {\n LOG_ERROR(\"Failed to train converter %s for %s\", quantizer_name.c_str(),\n IndexError::What(ret));\n if (!error_.exchange(true)) {\n err_code_ = IndexError_Runtime;\n }\n }\n };\n\n auto task_group = threads->make_group();\n if (!task_group) {\n LOG_ERROR(\"Failed to create task group\");\n return IndexError_Runtime;\n }\n\n for (size_t i = 0; i < quantizers_.size(); ++i) {\n if (error_) {\n task_group->wait_finish();\n return err_code_;\n }\n task_group->submit(ailego::Closure ::New(train_data, i));\n }\n\n task_group->wait_finish();\n if (quantizers_.size() > 0) {\n quantized_meta_ = quantizers_[0]->meta();\n }\n\n return 0;\n}\n\nINDEX_FACTORY_REGISTER_BUILDER(IVFBuilder);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_builder.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"ivf_centroid_index.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! IVF Builder\n */\nclass IVFBuilder : public IndexBuilder {\n public:\n //! Constructor\n IVFBuilder();\n\n //! Destructor\n ~IVFBuilder();\n\n //! Disable them\n IVFBuilder(const IVFBuilder &) = delete;\n IVFBuilder &operator=(const IVFBuilder &) = delete;\n\n public:\n //! Initialize the builder\n virtual int init(const IndexMeta &meta,\n const ailego::Params ¶ms) override;\n\n //! Cleanup the builder\n virtual int cleanup(void) override;\n\n //! Train the data\n virtual int train(IndexThreads::Pointer threads,\n IndexHolder::Pointer holder) override;\n\n //! Train the data\n virtual int train(const IndexTrainer::Pointer &trainer) override;\n\n //! Build the index\n virtual int build(IndexThreads::Pointer threads,\n IndexHolder::Pointer holder) override;\n\n //! Dump index into file system\n virtual int dump(const IndexDumper::Pointer &dumper) override;\n\n //! Retrieve statistics\n virtual const Stats &stats(void) const override {\n return stats_;\n }\n\n IVFCentroidIndex::Pointer centroid_index() const {\n return centroid_index_;\n }\n\n public:\n /*! Random Access Index Holder\n */\n class RandomAccessIndexHolder : public IndexHolder {\n public:\n //! Index Holder Iterator Pointer\n typedef std::shared_ptr Pointer;\n\n /*! Random Access Index Holder Iterator\n */\n class Iterator : public IndexHolder::Iterator {\n public:\n //! Index Holder Iterator Pointer\n typedef std::unique_ptr Pointer;\n\n //! Constructor\n Iterator(RandomAccessIndexHolder *owner) : holder_(owner) {}\n\n //! Destructor\n virtual ~Iterator(void) {}\n\n //! Retrieve pointer of data\n virtual const void *data(void) const override {\n return holder_->element(id_);\n }\n\n //! Test if the iterator is valid\n virtual bool is_valid(void) const override {\n return id_ < holder_->count();\n }\n\n //! Retrieve primary key\n virtual uint64_t key(void) const override {\n return holder_->key(id_);\n }\n\n //! Next iterator\n virtual void next(void) override {\n ++id_;\n }\n\n private:\n //! Members\n RandomAccessIndexHolder *holder_{nullptr};\n uint32_t id_{0};\n };\n\n //! Constructor\n RandomAccessIndexHolder(const IndexMeta &meta)\n : features_(std::make_shared(meta)) {}\n\n //! Retrieve count of elements in holder (-1 indicates unknown)\n virtual size_t count(void) const override {\n return features_->count();\n }\n\n //! Retrieve dimension\n virtual size_t dimension(void) const override {\n return features_->dimension();\n }\n\n //! Retrieve type information\n virtual IndexMeta::DataType data_type(void) const override {\n return features_->data_type();\n }\n\n //! Retrieve element size in bytes\n virtual size_t element_size(void) const override {\n return features_->element_size();\n }\n\n //! Retrieve if it can multi-pass\n virtual bool multipass(void) const override {\n return true;\n }\n\n //! Create a new iterator\n virtual IndexHolder::Iterator::Pointer create_iterator(void) override {\n return IndexHolder::Iterator::Pointer(\n new RandomAccessIndexHolder::Iterator(this));\n }\n\n void reserve(size_t elems) {\n features_->reserve(elems);\n keys_.reserve(elems);\n }\n\n //! Append an element into holder\n void emplace(uint64_t pkey, const void *vec) {\n features_->emplace(vec);\n keys_.emplace_back(pkey);\n }\n\n //! Retrieve feature via local id\n const void *element(size_t id) const {\n return features_->element(id);\n }\n\n //! Retrieve key via local id\n uint64_t key(size_t id) const {\n ailego_assert_with(id < keys_.size(), \"Index Overflow\");\n return keys_[id];\n }\n\n private:\n //! Disable them\n RandomAccessIndexHolder(void) = delete;\n\n //! Members\n CompactIndexFeatures::Pointer features_{};\n std::vector keys_{};\n };\n\n private:\n /*! Wrapper of feature\n */\n class Vector {\n public:\n typedef std::shared_ptr Pointer;\n\n Vector(const void *vec, size_t len, uint32_t idx)\n : vec_(reinterpret_cast(vec), len), id_{idx} {}\n\n const void *data() const {\n return vec_.data();\n }\n\n size_t size() const {\n return vec_.size();\n }\n\n uint32_t id(void) const {\n return id_;\n }\n\n private:\n std::string vec_{};\n uint32_t id_{0u};\n };\n\n using VectorList = std::vector;\n\n //! Check MajorOrder in meta, and update the major order if needed\n int CheckAndUpdateMajorOrder(IndexMeta &meta);\n\n //! Parse params\n int parse_centroids_num(const ailego::Params ¶ms);\n int parse_clustering_params(const ailego::Params ¶ms);\n int parse_general_params(const ailego::Params ¶ms);\n\n //! Prepare params for trainer\n int prepare_trainer_params(ailego::Params ¶ms);\n\n //! Build the index\n int build_label_index(IndexThreads *threads,\n const IndexHolder::Pointer &holder);\n\n //! Dump the index to dumper\n int dump_index(const IndexDumper::Pointer &dumper);\n\n //! Prepare the quantizer for inverted index\n int prepare_quantizer(IndexThreads *threads);\n\n //! Quantize the centrods list\n int quantize_centroids();\n\n //! Create converter and init with params\n static IndexConverter::Pointer CreateAndInitConverter(\n const IndexMeta &meta, const std::string &name,\n const ailego::Params ¶ms) {\n auto converter = IndexFactory::CreateConverter(name);\n if (!converter) {\n LOG_ERROR(\"Failed to create converter %s\", name.c_str());\n return IndexConverter::Pointer();\n }\n int ret = converter->init(meta, params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize converter %s for %s\", name.c_str(),\n IndexError::What(ret));\n return IndexConverter::Pointer();\n }\n return converter;\n }\n\n //! Select the nearest centroid id for the vector\n void label(const std::shared_ptr &vecs) {\n for (size_t i = 0; i < vecs->size(); ++i) {\n auto &vec = (*vecs)[i];\n\n uint32_t centroid_idx =\n centroid_index_->search_nearest_centroid(vec.data(), vec.size());\n if (centroid_idx == IVFCentroidIndex::kInvalidID) {\n LOG_ERROR(\"Failed to search nearest centroid in CentroidIndex\");\n if (!error_.exchange(true)) {\n err_code_ = IndexError_Runtime;\n }\n return;\n }\n ailego_assert_with(centroid_idx < labels_.size(), \"Index Overflow\");\n mutex_.lock();\n labels_[centroid_idx].emplace_back(vec.id());\n mutex_.unlock();\n }\n }\n\n\n private:\n //! Constants\n static constexpr size_t kThreadPoolQueueSize = 300u;\n static constexpr size_t kBatchSize = 10u;\n static constexpr size_t kDefaultBlockCount = 32u;\n\n enum BuilderState { INIT = 0, INITED = 1, TRAINED = 2, BUILT = 3 };\n\n //! Members\n BuilderState state_{INIT};\n Stats stats_{};\n ailego::Params params_{};\n IndexMeta meta_{};\n\n std::vector centroid_num_vec_{};\n std::string cluster_class_{};\n std::string converter_class_{};\n std::vector cluster_params_{};\n\n std::vector> labels_{};\n std::mutex mutex_{};\n IVFCentroidIndex::Pointer centroid_index_{};\n IVFCentroidIndex::Pointer searcher_centroid_index_{};\n RandomAccessIndexHolder::Pointer holder_{};\n IndexMeta converted_meta_{};\n IndexConverter::Pointer converter_{};\n IndexMeta quantized_meta_{};\n std::vector quantizers_{};\n\n std::atomic_bool error_{false};\n int err_code_{0};\n\n uint32_t thread_count_{0};\n uint32_t sample_count_{0};\n float sample_ratio_{0.0};\n uint32_t block_vector_count_{kDefaultBlockCount};\n bool cluster_auto_tuning_{false};\n bool store_original_features_{false};\n bool quantize_by_centroid_{false};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_centroid_index.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"ivf_centroid_index.h\"\n#include \n#include \n#include \"metric/metric_params.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Fake Trainer to supply centroids in bundle\n */\nclass FakeClusterTrainer : public IndexTrainer {\n public:\n //! Constructor\n FakeClusterTrainer(const IndexMeta &imeta, const IndexBundle::Pointer &bundle)\n : meta_(imeta), bundle_(bundle) {}\n\n //! Destructor\n ~FakeClusterTrainer(void) {}\n\n protected:\n //! Initialize Trainer\n virtual int init(const IndexMeta &, const ailego::Params &) override {\n return 0;\n }\n\n //! Cleanup Trainer\n virtual int cleanup(void) override {\n return 0;\n }\n\n //! Train the data\n virtual int train(IndexHolder::Pointer) override {\n return 0;\n }\n\n //! Train the data\n virtual int train(IndexThreads::Pointer, IndexHolder::Pointer) override {\n return 0;\n }\n\n //! Load index from file path or dir\n virtual int load(IndexStorage::Pointer) override {\n return 0;\n }\n\n //! Dump index into file path or dir\n virtual int dump(const IndexDumper::Pointer &) override {\n return 0;\n }\n\n //! Retrieve Index Meta\n virtual const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n //! Retrieve statistics\n virtual const IndexTrainer::Stats &stats(void) const override {\n return stats_;\n }\n\n //! Retrieve the output indexes\n virtual IndexBundle::Pointer indexes(void) const override {\n return bundle_;\n }\n\n private:\n //! Members\n IndexMeta meta_{};\n Stats stats_{};\n IndexBundle::Pointer bundle_{};\n};\n\n/*! Int8QuantizerReformer for InnerProduct Measure\n */\nclass Int8QuantizerReformer4IP : public IndexReformer {\n public:\n //! Initialize Reformer\n virtual int init(const ailego::Params &) override {\n return 0;\n }\n\n //! Cleanup Reformer\n virtual int cleanup(void) override {\n return 0;\n }\n\n //! Load index from container\n virtual int load(IndexStorage::Pointer) override {\n return 0;\n }\n\n //! Unload index\n virtual int unload(void) override {\n return 0;\n }\n\n //! Transform query\n virtual int transform(const void * /*query*/,\n const IndexQueryMeta & /*qmeta*/, std::string * /*out*/,\n IndexQueryMeta * /*ometa*/) const override {\n#if 0\n size_t dim = qmeta.dimension();\n out->resize(IndexMeta::ElementSizeof(\n IndexMeta::DataType::DT_INT8, dim));\n ometa->set_meta(IndexMeta::DataType::DT_INT8, dim);\n const float *ivec = reinterpret_cast(query);\n int8_t *ovec = reinterpret_cast(&(*out)[0]);\n float abs_max = 0.0f;\n for (size_t i = 0; i < dim; ++i) {\n auto abs = std::abs(ivec[i]);\n if (abs > abs_max) {\n abs_max = abs;\n }\n }\n if (abs_max > 0.0f) {\n float scale = 127 / abs_max;\n for (size_t i = 0; i < dim; ++i) {\n ovec[i] = static_cast(std::round(ivec[i] * scale));\n }\n } else {\n std::fill(ovec, ovec + dim, static_cast(1));\n }\n return 0;\n#else\n return IndexError_NotImplemented;\n#endif\n }\n\n //! Transform queries\n virtual int transform(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, std::string *oquery,\n IndexQueryMeta *ometa) const override {\n size_t dim = qmeta.dimension();\n oquery->resize(count *\n IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT8, dim));\n ometa->set_meta(IndexMeta::DataType::DT_INT8, dim);\n const float *ivec = reinterpret_cast(query);\n int8_t *ovec = reinterpret_cast(&(*oquery)[0]);\n for (size_t q = 0; q < count; ++q) {\n float abs_max = 0.0f;\n const float *in = &ivec[q * dim];\n int8_t *out = &ovec[q * dim];\n for (size_t i = 0; i < dim; ++i) {\n auto abs = std::abs(in[i]);\n if (abs > abs_max) {\n abs_max = abs;\n }\n }\n if (abs_max > 0.0f) {\n float scale = 127 / abs_max;\n for (size_t i = 0; i < dim; ++i) {\n out[i] = static_cast(std::round(in[i] * scale));\n }\n } else {\n std::fill(out, out + dim, static_cast(1));\n }\n }\n return 0;\n }\n\n //! Normalize results\n virtual int normalize(const void * /*query*/,\n const IndexQueryMeta & /*qmeta*/,\n IndexDocumentList & /*result*/) const override {\n return 0;\n }\n};\n\n/*! Int4QuantizerReformer for InnerProduct Metric\n */\nclass Int4QuantizerReformer4IP : public IndexReformer {\n public:\n //! Initialize Reformer\n virtual int init(const ailego::Params &) override {\n return 0;\n }\n\n //! Cleanup Reformer\n virtual int cleanup(void) override {\n return 0;\n }\n\n //! Load index from container\n virtual int load(IndexStorage::Pointer) override {\n return 0;\n }\n\n //! Unload index\n virtual int unload(void) override {\n return 0;\n }\n\n //! Transform query\n virtual int transform(const void * /*query*/,\n const IndexQueryMeta & /*qmeta*/, std::string * /*out*/,\n IndexQueryMeta * /*ometa*/) const override {\n return IndexError_NotImplemented;\n }\n\n //! Transform queries\n virtual int transform(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, std::string *oquery,\n IndexQueryMeta *ometa) const override {\n if (qmeta.dimension() & 0x1) {\n LOG_ERROR(\"Unsuuport dim=%u for transform\", qmeta.dimension());\n return IndexError_Unsupported;\n }\n\n size_t dim = qmeta.dimension();\n oquery->resize(count *\n IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT4, dim));\n ometa->set_meta(IndexMeta::DataType::DT_INT4, dim);\n const float *ivec = reinterpret_cast(query);\n uint8_t *ovec = reinterpret_cast(&(*oquery)[0]);\n for (size_t q = 0; q < count; ++q) {\n float abs_max = 0.0f;\n float max = -std::numeric_limits::max();\n const float *in = &ivec[q * dim];\n uint8_t *out = &ovec[q * dim / 2];\n for (size_t i = 0; i < dim; ++i) {\n float abs = std::abs(in[i]);\n abs_max = std::max(abs_max, abs);\n max = std::max(max, in[i]);\n }\n if (abs_max > 0.0f) {\n float scale = ((7 * abs_max > 8 * max) ? 8 : 7) / abs_max;\n for (size_t i = 0; i < dim; i += 2) {\n auto v1 = static_cast(std::round(in[i] * scale));\n auto v2 = static_cast(std::round(in[i + 1] * scale));\n out[i / 2] = (static_cast(v1) << 4) |\n (static_cast(v2) & 0xF);\n }\n } else {\n std::fill(out, out + dim / 2, static_cast(9));\n }\n }\n return 0;\n }\n\n //! Normalize results\n virtual int normalize(const void * /*query*/,\n const IndexQueryMeta & /*qmeta*/,\n IndexDocumentList & /*result*/) const override {\n return 0;\n }\n};\n\nint IVFCentroidIndex::init(const IndexMeta &meta,\n const ailego::Params ¶ms) {\n meta_ = meta;\n\n params.get(PARAM_IVF_BUILDER_OPTIMIZER_CLASS, &builder_class_);\n params.get(PARAM_IVF_BUILDER_OPTIMIZER_PARAMS, &builder_params_);\n params.get(PARAM_IVF_SEARCHER_OPTIMIZER, &searcher_class_);\n params.get(PARAM_IVF_SEARCHER_OPTIMIZER_PARAMS, &searcher_params_);\n\n return 0;\n}\n\nint IVFCentroidIndex::search(const void *query, const IndexQueryMeta &qmeta,\n size_t count,\n IndexSearcher::Context::Pointer &ctx) {\n int ret = 0;\n\n if (reformer_) {\n std::string buffer;\n IndexQueryMeta ometa;\n ret = reformer_->transform(query, qmeta, count, &buffer, &ometa);\n if (ret != 0) {\n LOG_ERROR(\"Failed to transform querys by reformer\");\n return ret;\n }\n ret = searcher_->search_impl(buffer.data(), ometa, count, ctx);\n } else {\n ret = searcher_->search_impl(query, qmeta, count, ctx);\n }\n\n ivf_check_with_msg(ret, \"Failed to search in centroid index for %s\",\n IndexError::What(ret));\n\n return 0;\n}\n\nuint32_t IVFCentroidIndex::search_nearest_centroid(const void *query,\n size_t len) {\n //! Called in building index precedure, so transform the query is needless\n if (len != meta_.element_size()) {\n LOG_ERROR(\"Invalid query size actual: %zu, expected: %u\", len,\n meta_.element_size());\n return kInvalidID;\n }\n\n thread_local IndexSearcher::Context::Pointer context(\n searcher_->create_context());\n context->set_topk(1);\n\n IndexQueryMeta qmeta(meta_.data_type(), meta_.dimension());\n int ret = searcher_->search_impl(query, qmeta, context);\n if (ret != 0 || context->result().empty()) {\n LOG_ERROR(\"Failed to search nearest centroid, with ret %d\", ret);\n return kInvalidID;\n }\n\n return static_cast(context->result()[0].key());\n}\n\nuint32_t IVFCentroidIndex::transform_and_search_nearest_centroid(\n const void *record, const IndexQueryMeta &rmeta,\n IndexSearcher::Context::Pointer &ctx) const {\n int ret = 0;\n if (reformer_) {\n std::string buffer;\n IndexQueryMeta ometa;\n ret = reformer_->convert(record, rmeta, &buffer, &ometa);\n if (ret != 0) {\n LOG_ERROR(\"Failed to transform querys by reformer\");\n return kInvalidID;\n }\n ret = searcher_->search_impl(buffer.data(), ometa, ctx);\n } else {\n ret = searcher_->search_impl(record, rmeta, ctx);\n }\n if (ret != 0 || ctx->result().empty()) {\n LOG_ERROR(\"Failed to search in centroid index for %s\",\n IndexError::What(ret));\n return kInvalidID;\n }\n\n return static_cast(ctx->result()[0].key());\n}\n\nIndexHolder::Pointer IVFCentroidIndex::quantize_holder(\n const IndexHolder::Pointer &holder) {\n auto input = holder;\n if (meta_.reformer_name() == kMipsReformerName &&\n meta_.metric_name() == kL2MetricName &&\n (quantizer_class_ == kInt8QuantizerName ||\n quantizer_class_ == kInt4QuantizerName)) {\n //! Reverse for Mips if do convert by integer quantizer\n auto reverse = IndexFactory::CreateConverter(kMipsRevConverterName);\n if (!reverse) {\n LOG_ERROR(\"Failed to create converter %s\", kMipsRevConverterName);\n return nullptr;\n }\n ailego::Params params;\n auto p = meta_.reformer_params();\n params.set(MIPS_REVERSE_CONVERTER_M_VALUE,\n p.get_as_uint32(MIPS_REFORMER_M_VALUE));\n params.set(MIPS_REVERSE_CONVERTER_U_VALUE,\n p.get_as_float(MIPS_REFORMER_U_VALUE));\n params.set(MIPS_REVERSE_CONVERTER_L2_NORM,\n p.get_as_uint32(MIPS_REFORMER_L2_NORM));\n params.set(MIPS_REVERSE_CONVERTER_FORCED_SINGLE_FLOAT,\n p.get_as_float(MIPS_REFORMER_FORCED_HALF_FLOAT));\n int ret = reverse->init(meta_, params);\n if (ret != 0) {\n LOG_ERROR(\"Fail to init converter %s\", kMipsRevConverterName);\n return nullptr;\n }\n ret = IndexConverter::TrainAndTransform(reverse, holder);\n if (ret != 0) {\n LOG_ERROR(\"Fail to transform converter %s\", kMipsRevConverterName);\n return nullptr;\n }\n input = reverse->result();\n meta_ = reverse->meta();\n meta_.set_metric(kIPMetricName, 0, ailego::Params());\n meta_.set_reformer(\"\", 0, ailego::Params());\n }\n\n auto converter = IndexFactory::CreateConverter(quantizer_class_);\n if (!converter) {\n LOG_ERROR(\"Failed to create converter %s\", quantizer_class_.c_str());\n return nullptr;\n }\n int ret = converter->init(meta_, quantizer_params_);\n if (ret != 0) {\n LOG_ERROR(\"Fail to init converter %s\", quantizer_class_.c_str());\n return nullptr;\n }\n\n ret = IndexConverter::TrainAndTransform(converter, input);\n if (ret != 0) {\n LOG_ERROR(\"Fail to tranform converter %s\", quantizer_class_.c_str());\n return nullptr;\n }\n\n meta_ = converter->meta();\n return converter->result();\n}\n\nint IVFCentroidIndex::build_index(\n const IndexCluster::CentroidList ¢roid_list,\n const IndexDumper::Pointer &dumper) {\n IndexBuilder::Pointer builder = IndexFactory::CreateBuilder(builder_class_);\n if (!builder) {\n LOG_ERROR(\"Failed to create builder %s\", builder_class_.c_str());\n return IndexError_NoExist;\n }\n\n IndexHolder::Pointer holder =\n std::make_shared(meta_, centroid_list);\n if (!holder) {\n return IndexError_NoMemory;\n }\n if (holder->count() == 0) {\n LOG_ERROR(\"No centroids to build\");\n return IndexError_InvalidArgument;\n }\n centroids_count_ = holder->count();\n\n //! Set default params if not given\n auto count = std::to_string(\n static_cast(std::ceil(std::sqrt(centroids_count_ / 10.0))));\n // if (IsHcBuilder(builder_class_) &&\n // !builder_params_.has(hc::PARAM_HC_BUILDER_CENTROID_COUNT)) {\n // builder_params_.set(hc::PARAM_HC_BUILDER_CENTROID_COUNT, count);\n // } else if (builder_class_ == \"GcBuilder\" &&\n // !builder_params_.has(hc::PARAM_GC_BUILDER_CENTROID_COUNT)) {\n // builder_params_.set(hc::PARAM_GC_BUILDER_CENTROID_COUNT, count);\n // }\n if (!quantizer_class_.empty()) {\n holder = this->quantize_holder(holder);\n if (!holder) {\n return IndexError_Runtime;\n }\n }\n\n const auto name = builder_class_.c_str();\n int ret = builder->init(meta_, builder_params_);\n ivf_check_with_msg(ret, \"%s init failed, ret=%d\", name, ret);\n\n // if (IsHcBuilder(builder_class_) && quantizer_class_.empty()) {\n // auto trainer = this->prepare_trainer(centroid_list);\n // ret = trainer ? builder->train(trainer) : builder->train(holder);\n // } else {\n // ret = builder->train(holder);\n // }\n\n ret = builder->train(holder);\n ivf_check_with_msg(ret, \"%s train failed, ret=%d\", name, ret);\n\n ret = builder->build(holder);\n ivf_check_with_msg(ret, \"%s build failed, ret=%d\", name, ret);\n\n ret = builder->dump(dumper);\n ivf_check_with_msg(ret, \"%s dump failed, ret=%d\", name, ret);\n\n ret = dumper->close();\n ivf_check_error_code(ret);\n\n return 0;\n}\n\nint IVFCentroidIndex::build(const IndexCluster::CentroidList ¢roid_list) {\n index_building_ = true;\n //! Build and dump the index\n IndexDumper::Pointer dumper = IndexFactory::CreateDumper(\"MemoryDumper\");\n if (!dumper) {\n LOG_ERROR(\"Failed to create MemoryDumper\");\n return IndexError_NoExist;\n }\n path_ = IVFUtility::GenerateRandomPath(kTempralPathPrefix);\n int ret = dumper->create(path_);\n if (ret != 0) {\n LOG_ERROR(\"IndexDumper create path %s failed\", path_.c_str());\n return ret;\n }\n ret = this->build_index(centroid_list, dumper);\n ivf_check_error_code(ret);\n\n auto rope = IndexMemory::Instance()->open(path_);\n if (!rope) {\n LOG_ERROR(\"Open memory path %s failed.\", path_.c_str());\n return ret;\n }\n if (rope->count() != 1) {\n LOG_ERROR(\"Graph Rope block count not equal with 1.\");\n return ret;\n }\n (*rope)[0].read(0, &data_, 0);\n size_ = (*rope)[0].size();\n\n //! Load the index\n IndexStorage::Pointer container =\n IndexFactory::CreateStorage(\"MemoryReadStorage\");\n if (!container) {\n LOG_ERROR(\"Failed to create MemoryReadStorage\");\n return IndexError_NoExist;\n }\n ret = container->init(ailego::Params());\n ivf_check_with_msg(ret, \"Failed to initialize MemoryReadStorage for %s\",\n IndexError::What(ret));\n ret = container->open(path_, false);\n ivf_check_with_msg(ret, \"Failed to load path in MemoryReadStorage for %s\",\n IndexError::What(ret));\n\n ailego::Params searcher_params;\n if (!searcher_class_.empty()) {\n searcher_params.set(PARAM_IVF_SEARCHER_OPTIMIZER, searcher_class_);\n }\n if (!searcher_params_.empty()) {\n searcher_params.set(PARAM_IVF_SEARCHER_OPTIMIZER_PARAMS, searcher_params_);\n }\n ret = this->load(container, searcher_params);\n ivf_check_with_msg(ret, \"IVFCentroidIndex load failed with %s\",\n IndexError::What(ret));\n\n return 0;\n}\n\nint IVFCentroidIndex::load(const IndexStorage::Pointer &container,\n const ailego::Params params) {\n if (!container) {\n LOG_ERROR(\"Invalid container\");\n return IndexError_InvalidArgument;\n }\n\n int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to deserialize meta from container\");\n return ret;\n }\n\n auto reformer_name = meta_.reformer_name();\n if (!reformer_name.empty()) {\n LOG_DEBUG(\"Load CentroidIndex with reformer %s, metric %s\",\n reformer_name.c_str(), meta_.metric_name().c_str());\n if ((reformer_name == kInt8ReformerName ||\n reformer_name == kInt4ReformerName) &&\n meta_.metric_name() == kIPMetricName) {\n if (reformer_name == kInt8ReformerName) {\n reformer_ = std::make_shared();\n } else {\n reformer_ = std::make_shared();\n }\n if (!reformer_) {\n return IndexError_NoMemory;\n }\n } else {\n reformer_ = IndexFactory::CreateReformer(reformer_name);\n if (!reformer_) {\n LOG_ERROR(\"Failed to create reformer %s\", reformer_name.c_str());\n return IndexError_NoExist;\n }\n }\n ret = reformer_->init(meta_.reformer_params());\n ivf_check_with_msg(ret, \"Failed to initialize reformer %s\",\n reformer_name.c_str());\n }\n\n searcher_class_ = meta_.searcher_name();\n params.get(PARAM_IVF_SEARCHER_OPTIMIZER, &searcher_class_);\n params.get(PARAM_IVF_SEARCHER_OPTIMIZER_PARAMS, &searcher_params_);\n searcher_ = IndexFactory::CreateSearcher(searcher_class_);\n if (!searcher_) {\n LOG_ERROR(\"Failed to create searcher %s\", searcher_class_.c_str());\n return IndexError_Runtime;\n }\n\n auto searcher_params = meta_.searcher_params();\n searcher_params.merge(searcher_params_);\n ret = searcher_->init(searcher_params);\n ivf_check_with_msg(ret, \"Failed to initialize searcher %s\",\n searcher_class_.c_str());\n\n IndexMetric::Pointer metric;\n if (index_building_) {\n // The searcher index metric should specified in building process,\n // otherwise the query_metric will be used in searching\n metric = IndexFactory::CreateMetric(meta_.metric_name());\n ivf_assert_with_msg(metric, IndexError_NoExist,\n \"Failed to create metric %s\",\n meta_.metric_name().c_str());\n ret = metric->init(meta_, meta_.metric_params());\n ivf_check_with_msg(ret, \"Failed to initialize metric\");\n }\n ret = searcher_->load(container, metric);\n ivf_check_with_msg(ret, \"Failed to load searcher %s\",\n searcher_class_.c_str());\n\n return 0;\n}\n\nIndexTrainer::Pointer IVFCentroidIndex::prepare_trainer(\n const IndexCluster::CentroidList ¢roid_list) {\n IndexCluster::CentroidList level1_centroids;\n bool two_level = false;\n for (auto &it : centroid_list) {\n auto centroid = it;\n if (!centroid.subitems().empty()) {\n two_level = true;\n }\n centroid.mutable_subitems()->clear();\n centroid.mutable_similars()->clear();\n level1_centroids.emplace_back(centroid);\n }\n if (!two_level) {\n return IndexTrainer::Pointer();\n }\n\n IndexBundle::Pointer bundle;\n IndexCluster::Serialize(meta_, level1_centroids, &bundle);\n return std::make_shared(meta_, bundle);\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_centroid_index.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"ivf_params.h\"\n#include \"ivf_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! IVF Centroid Index\n */\nclass IVFCentroidIndex {\n public:\n typedef std::shared_ptr Pointer;\n\n //! Constructor\n IVFCentroidIndex(void) {}\n\n //! Destructor\n ~IVFCentroidIndex(void) {\n IndexMemory *instance = IndexMemory::Instance();\n if (instance) {\n if (instance->has(path_)) {\n instance->remove(path_);\n }\n }\n }\n\n //! Initialize\n int init(const IndexMeta &meta, const ailego::Params ¶ms);\n\n //! Set Quantizer for the index\n void set_quantizer(const std::string &quantizer_name,\n ailego::Params &quantizer_params) {\n quantizer_class_ = quantizer_name;\n quantizer_params_ = quantizer_params;\n }\n\n //! Retrieve data address of the index\n const void *data(void) const {\n return data_;\n }\n\n //! Retrieve size of the index\n size_t size(void) const {\n return size_;\n }\n\n //! Create searcher context for centroid index\n IndexSearcher::Context::Pointer create_context() const {\n return searcher_ ? searcher_->create_context() : nullptr;\n }\n\n //! Similarity search\n int search(const void *query, const IndexQueryMeta &qmeta, size_t count,\n IndexSearcher::Context::Pointer &ctx);\n\n //! Search the nearest point, must be called in local thread pool\n uint32_t search_nearest_centroid(const void *query, size_t len);\n\n //! Transform Data and Search the nearest point, called while adding record\n uint32_t transform_and_search_nearest_centroid(\n const void *record, const IndexQueryMeta &rmeta,\n IndexSearcher::Context::Pointer &ctx) const;\n\n //! Build Centroid Index From Centroid List\n int build(const IndexCluster::CentroidList ¢roid_list);\n\n //! Load Centroid Index From container\n int load(const IndexStorage::Pointer &container, const ailego::Params params);\n\n //! Retrieve centroid count of the index\n size_t centroids_count(void) const {\n return centroids_count_;\n }\n\n //! Retrieve meta\n const IndexMeta &meta() const {\n return meta_;\n }\n\n //! Retrieve reformer of the index\n const IndexReformer::Pointer reformer(void) const {\n return reformer_;\n }\n\n static constexpr uint32_t kInvalidID = std::numeric_limits::max();\n\n private:\n /*! Centroids IndexHolder\n */\n class CentroidsIndexHolder : public IndexHolder {\n public:\n class Iterator : public IndexHolder::Iterator {\n public:\n //! Index Holder Iterator Pointer\n typedef std::unique_ptr Pointer;\n\n //! Constructor\n Iterator(std::vector *features) : features_(features) {}\n\n //! Destructor\n virtual ~Iterator(void) {}\n\n //! Retrieve pointer of data\n virtual const void *data(void) const override {\n return (*features_)[id_];\n }\n\n //! Test if the iterator is valid\n virtual bool is_valid(void) const override {\n return id_ < features_->size();\n }\n\n //! Retrieve primary key\n virtual uint64_t key(void) const override {\n return id_;\n }\n\n //! Next iterator\n virtual void next(void) override {\n ++id_;\n }\n\n private:\n //! Members\n std::vector *features_{nullptr};\n uint32_t id_{0};\n };\n\n //! Constructor\n CentroidsIndexHolder(const IndexMeta &meta,\n const IndexCluster::CentroidList ¢roid_list)\n : dimension_(meta.dimension()),\n element_size_(meta.element_size()),\n data_type_(meta.data_type()) {\n using CentroidList = IndexCluster::CentroidList;\n\n std::function get_leaf_features =\n [&](const CentroidList ¢s) {\n if (cents.empty()) {\n return;\n }\n for (const auto &it : cents) {\n if (it.subitems().empty()) {\n features_.emplace_back(it.feature());\n } else {\n get_leaf_features(it.subitems());\n }\n }\n };\n\n get_leaf_features(centroid_list);\n }\n\n //! Retrieve count of elements in holder (-1 indicates unknown)\n virtual size_t count(void) const override {\n return features_.size();\n }\n\n //! Retrieve dimension\n virtual size_t dimension(void) const override {\n return dimension_;\n }\n\n //! Retrieve type information\n virtual IndexMeta::DataType data_type(void) const override {\n return data_type_;\n }\n\n //! Retrieve element size in bytes\n virtual size_t element_size(void) const override {\n return element_size_;\n }\n\n //! Retrieve if it can multi-pass\n virtual bool multipass(void) const override {\n return true;\n }\n\n //! Create a new iterator\n virtual IndexHolder::Iterator::Pointer create_iterator(void) override {\n return IndexHolder::Iterator::Pointer(\n new CentroidsIndexHolder::Iterator(&features_));\n }\n\n private:\n //! Members\n std::vector features_{};\n size_t dimension_{0};\n size_t element_size_{0};\n IndexMeta::DataType data_type_{IndexMeta::DataType::DT_UNDEFINED};\n };\n\n int build_index(const IndexCluster::CentroidList ¢roid_list,\n const IndexDumper::Pointer &dumper);\n\n //! Prepare trainer for clustering index\n IndexTrainer::Pointer prepare_trainer(\n const IndexCluster::CentroidList ¢roid_list);\n\n //! Quantize the centroid vectors in holder\n IndexHolder::Pointer quantize_holder(const IndexHolder::Pointer &holder);\n\n\n private:\n //! Constants\n constexpr static const char *kDefaultBuilder = \"FlatBuilder\";\n constexpr static const char *kTempralPathPrefix = \"IVF\";\n\n //! Members\n IndexMeta meta_{};\n\n IndexSearcher::Pointer searcher_{};\n IndexReformer::Pointer reformer_{};\n std::string builder_class_{kDefaultBuilder};\n std::string searcher_class_{};\n std::string quantizer_class_{};\n\n std::string path_{};\n\n ailego::Params builder_params_{};\n ailego::Params searcher_params_{};\n ailego::Params quantizer_params_{};\n\n const void *data_{};\n size_t size_{};\n size_t centroids_count_{0};\n bool index_building_{false};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_distance_calculator.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"ivf_distance_calculator.h\"\n#include \n\nnamespace zvec {\nnamespace core {\n\nIVFDistanceCalculator::IVFDistanceCalculator(const IndexMeta &meta,\n const IndexMetric::Pointer &metric,\n uint32_t block_vec_cnt)\n : metric_ptr_(metric), block_vec_cnt_(block_vec_cnt) {\n row_distance_ = metric->distance();\n distanceXx1_ = metric->distance_matrix(block_vec_cnt, 1);\n distances_.resize(33);\n for (size_t b = 32; b != 0; b /= 2) {\n distances_[b] = metric->distance_matrix(block_vec_cnt, b);\n }\n element_size_ = meta.element_size();\n dimension_ = meta.dimension();\n if (meta.major_order() == IndexMeta::MajorOrder::MO_COLUMN) {\n column_major_order_ = true;\n } else {\n column_major_order_ = false;\n }\n}\n\nIVFDistanceCalculator::~IVFDistanceCalculator() {\n row_distance_ = nullptr;\n distanceXx1_ = nullptr;\n distances_.clear();\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_distance_calculator.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nclass IVFDistanceCalculator {\n public:\n typedef std::shared_ptr Pointer;\n\n //! Constructor\n IVFDistanceCalculator(const IndexMeta &meta,\n const IndexMetric::Pointer &metric,\n uint32_t block_vec_cnt);\n\n virtual ~IVFDistanceCalculator();\n\n public:\n inline void query_centroids_distance(const void *query, size_t qnum,\n const void *feature, size_t fnum,\n float *distances);\n\n inline void query_centroids_distance(const void *query, const void *feature,\n size_t fnum, float *distances);\n\n inline void query_features_distance(const void *query, const void *feature,\n size_t fnum, float *distances);\n\n inline void query_features_distance(const void *query, const void *feature,\n bool column_major, size_t fnum,\n float *distances);\n\n protected:\n //! Row Major Distances -> Online\n inline void row_major_distance(const void *query, size_t qnum,\n const void *feature, size_t fnum, float *out);\n\n inline void row_major_distance(const void *query, const void *feature,\n size_t fnum, float *out);\n\n template \n inline void batch_query_centroids_distance(const void *query,\n const void *feature, size_t fnum,\n float *distances);\n\n protected:\n IndexMetric::Pointer metric_ptr_{};\n IndexMetric::MatrixDistance row_distance_{nullptr};\n IndexMetric::MatrixDistance distanceXx1_{nullptr};\n std::vector distances_{};\n\n size_t element_size_{0};\n size_t dimension_{0};\n uint32_t block_vec_cnt_{0};\n bool column_major_order_{false};\n};\n\nvoid IVFDistanceCalculator::query_centroids_distance(const void *query,\n size_t qnum,\n const void *feature,\n size_t fnum,\n float *distances) {\n if (column_major_order_) {\n switch (qnum) {\n case 1:\n batch_query_centroids_distance<1>(query, feature, fnum, distances);\n break;\n case 16:\n batch_query_centroids_distance<16>(query, feature, fnum, distances);\n break;\n case 8:\n batch_query_centroids_distance<8>(query, feature, fnum, distances);\n break;\n case 4:\n batch_query_centroids_distance<4>(query, feature, fnum, distances);\n break;\n case 2:\n batch_query_centroids_distance<2>(query, feature, fnum, distances);\n break;\n case 32:\n batch_query_centroids_distance<32>(query, feature, fnum, distances);\n break;\n default:\n LOG_ERROR(\"Unsupported query num %zu.\", qnum);\n break;\n }\n } else {\n const uint8_t *cur_query = reinterpret_cast(query);\n for (size_t q = 0; q < qnum; ++q) {\n this->row_major_distance(cur_query, feature, fnum, distances);\n cur_query += element_size_;\n distances += block_vec_cnt_;\n }\n }\n}\n\nvoid IVFDistanceCalculator::query_centroids_distance(const void *query,\n const void *feature,\n size_t fnum,\n float *distances) {\n this->query_features_distance(query, feature, fnum, distances);\n}\n\nvoid IVFDistanceCalculator::query_features_distance(const void *query,\n const void *feature,\n size_t fnum,\n float *distances) {\n if (column_major_order_) {\n if (fnum == block_vec_cnt_) {\n distanceXx1_(feature, query, dimension_, distances);\n } else {\n this->row_major_distance(query, feature, fnum, distances);\n }\n } else {\n this->row_major_distance(query, feature, fnum, distances);\n }\n}\n\nvoid IVFDistanceCalculator::query_features_distance(const void *query,\n const void *feature,\n bool column_major,\n size_t fnum,\n float *distances) {\n if (column_major) {\n ailego_assert_with(fnum == block_vec_cnt_, \"Invalid Block\");\n distanceXx1_(feature, query, dimension_, distances);\n } else {\n this->row_major_distance(query, feature, fnum, distances);\n }\n}\n\ntemplate \nvoid IVFDistanceCalculator::batch_query_centroids_distance(const void *query,\n const void *feature,\n size_t fnum,\n float *distances) {\n if (fnum == block_vec_cnt_) {\n distances_[Q](feature, query, dimension_, distances);\n } else {\n row_major_distance(query, Q, feature, fnum, distances);\n }\n}\n\nvoid IVFDistanceCalculator::row_major_distance(const void *query, size_t qnum,\n const void *feature, size_t fnum,\n float *out) {\n const uint8_t *cur_query = reinterpret_cast(query);\n for (size_t q = 0; q < qnum; ++q) {\n const uint8_t *tmp_feature = reinterpret_cast(feature);\n float *cur_out = out + q * fnum;\n for (size_t f = 0; f < fnum; ++f) {\n row_distance_(cur_query, tmp_feature, dimension_, cur_out + f);\n tmp_feature += element_size_;\n }\n cur_query += element_size_;\n }\n}\n\nvoid IVFDistanceCalculator::row_major_distance(const void *query,\n const void *feature, size_t fnum,\n float *out) {\n const uint8_t *cur_feature = reinterpret_cast(feature);\n for (size_t f = 0; f < fnum; ++f) {\n row_distance_(query, cur_feature, dimension_, out + f);\n cur_feature += element_size_;\n }\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_dumper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"ivf_dumper.h\"\n\nnamespace zvec {\nnamespace core {\n\nint IVFDumper::dump_inverted_vector(uint32_t inverted_list_id, uint64_t key,\n const void *vec) {\n int ret = this->check_dump_inverted_list(inverted_list_id);\n ivf_check_error_code(ret);\n\n ++inverted_lists_meta_[cur_list_id_].vector_count;\n ++header_.total_vector_count;\n block_.emplace(key, vec, IndexMeta::MajorOrder::MO_ROW);\n if (block_.full()) {\n ret = this->dump_block();\n ivf_check_error_code(ret);\n }\n return 0;\n}\n\nint IVFDumper::dump_inverted_block(uint32_t inverted_list_id,\n const uint64_t *keys, const void *vecs,\n uint32_t vector_count, bool column_major) {\n int ret = this->check_dump_inverted_list(inverted_list_id);\n ivf_check_error_code(ret);\n\n if (block_.match_order(column_major ? IndexMeta::MajorOrder::MO_COLUMN\n : IndexMeta::MajorOrder::MO_ROW) &&\n vector_count == block_.capacity()) {\n // Dump the block directly\n size_t size = vector_count * meta_.element_size();\n size_t pd_size = ailego_align(size, 32) - size;\n if (dumper_->write(vecs, size) != size) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper_->name().c_str());\n return IndexError_WriteData;\n }\n if (pd_size > 0) {\n std::string padding(pd_size, '\\0');\n if (dumper_->write(padding.data(), pd_size) != pd_size) {\n return IndexError_WriteData;\n }\n }\n std::copy(keys, keys + vector_count, std::back_inserter(keys_));\n ++inverted_lists_meta_[cur_list_id_].block_count;\n ++header_.block_count;\n header_.inverted_body_size += size;\n } else {\n size_t step_size = meta_.element_size();\n if (column_major) {\n step_size = IndexMeta::AlignSizeof(meta_.data_type());\n }\n for (size_t i = 0; i < vector_count; ++i) {\n auto v = reinterpret_cast(vecs) + i * step_size;\n block_.emplace(keys[i], v,\n column_major ? IndexMeta::MajorOrder::MO_COLUMN\n : IndexMeta::MajorOrder::MO_ROW);\n if (block_.full()) {\n ret = this->dump_block();\n ivf_check_error_code(ret);\n }\n }\n }\n\n inverted_lists_meta_[cur_list_id_].vector_count += vector_count;\n header_.total_vector_count += vector_count;\n\n return 0;\n}\n\nint IVFDumper::dump_container_segment(const IndexStorage::Pointer &container,\n const std::string &segmemt_id) {\n auto seg = container->get(segmemt_id, 2);\n if (!seg) {\n LOG_ERROR(\"Failed to fetch segment %s from %s\", segmemt_id.c_str(),\n container->name().c_str());\n return IndexError_InvalidFormat;\n }\n\n const size_t batch_size = 32 * 1024;\n const size_t total_size = seg->data_size() + seg->padding_size();\n size_t off = 0;\n while (off < total_size) {\n const void *data = nullptr;\n size_t rd_size = std::min(batch_size, total_size - off);\n if (seg->read(off, &data, rd_size) != rd_size) {\n LOG_ERROR(\"Failed to read data, off=%zu size=%zu\", off, rd_size);\n return IndexError_ReadData;\n }\n if (dumper_->write(data, rd_size) != rd_size) {\n LOG_ERROR(\"Failed to write data, size=%zu\", rd_size);\n return IndexError_WriteData;\n }\n off += rd_size;\n }\n\n int ret = dumper_->append(segmemt_id, seg->data_size(), seg->padding_size(),\n seg->data_crc());\n ivf_check_with_msg(ret, \"Failed to append %s\", segmemt_id.c_str());\n\n dumped_size_ += total_size;\n\n return 0;\n}\n\nint IVFDumper::dump_inverted_vector_finished(void) {\n //! Dump Inverted Index Segment\n if (!block_.empty()) {\n int ret = this->dump_block();\n ivf_check_error_code(ret);\n }\n header_.block_size = block_.block_size();\n size_t segment_size = header_.inverted_body_size;\n int ret = dumper_->append(IVF_INVERTED_BODY_SEG_ID, segment_size, 0, 0);\n if (ret != 0) {\n LOG_ERROR(\"Failed to append to segment %s, ret=%d\",\n IVF_INVERTED_BODY_SEG_ID.c_str(), ret);\n return ret;\n }\n dumped_size_ += segment_size;\n\n //! Dump Inverted Index Header Segment\n std::string str;\n meta_.serialize(&str);\n header_.header_size = sizeof(header_) + str.size();\n header_.index_meta_size = str.size();\n header_.inverted_list_count = inverted_lists_meta_.size();\n if (dumper_->write(&header_, sizeof(header_)) != sizeof(header_)) {\n LOG_ERROR(\"Failed to write data, size %zu\", sizeof(header_));\n return IndexError_WriteData;\n }\n if (dumper_->write(str.data(), str.size()) != str.size()) {\n LOG_ERROR(\"Failed to write data, size %zu\", str.size());\n return IndexError_WriteData;\n }\n size_t padding_size = 0;\n ret = this->dump_padding(header_.header_size, &padding_size);\n ivf_check_error_code(ret);\n ret = dumper_->append(IVF_INVERTED_HEADER_SEG_ID, header_.header_size,\n padding_size, 0);\n if (ret != 0) {\n LOG_ERROR(\"Failed to append to segment %s, ret:%d\",\n IVF_INVERTED_HEADER_SEG_ID.c_str(), ret);\n return ret;\n }\n dumped_size_ += header_.header_size + padding_size;\n\n LOG_DEBUG(\n \"Dump header info: blocks=%u block_size=%u block_vec_count=%u \"\n \"inverted_list_count=%u total_vecs=%u inverted_size=%zu\",\n header_.block_count, header_.block_size, header_.block_vector_count,\n header_.inverted_list_count, header_.total_vector_count,\n static_cast(header_.inverted_body_size));\n\n //! Dump Inverted Lists Meta Segment\n segment_size = inverted_lists_meta_.size() * sizeof(InvertedListMeta);\n ret = this->dump_segment(IVF_INVERTED_META_SEG_ID,\n inverted_lists_meta_.data(), segment_size);\n ivf_check_error_code(ret);\n\n //! Dump Keys Segment\n ret = this->dump_segment(IVF_KEYS_SEG_ID, keys_.data(),\n keys_.size() * sizeof(keys_[0]));\n ivf_check_error_code(ret);\n\n //! Dump Mapping Segment\n auto mapping = std::make_shared>();\n IVFUtility::Sort(keys_.data(), mapping.get(), keys_.size());\n ret = this->dump_segment(IVF_MAPPING_SEG_ID, mapping->data(),\n mapping->size() * sizeof(uint32_t));\n ivf_check_error_code(ret);\n mapping.reset();\n\n //! Dump the Offsets Segment\n return this->dump_offsets_segment();\n}\n\nint IVFDumper::dump_centroid_index(const void *data, size_t size) {\n int ret = this->dump_segment(IVF_CENTROID_SEG_ID, data, size);\n ivf_check_error_code(ret);\n\n return 0;\n}\n\nint IVFDumper::dump_quantizer_params(\n const std::vector &quantizers) {\n if (meta_.reformer_name() != kInt8ReformerName &&\n meta_.reformer_name() != kInt4ReformerName) {\n // IntegerQuantizer params is support only\n return 0;\n }\n if (quantizers.size() == 1) {\n //! Donot dump, using reformer params in IndexMeta\n return 0;\n }\n\n if (quantizers.size() != header_.inverted_list_count) {\n LOG_ERROR(\"Mismatch size, quantizers=%zu, inverted_list_count=%u\",\n quantizers.size(), header_.inverted_list_count);\n return IndexError_Logic;\n }\n bool int8_quantizer = meta_.reformer_name() == kInt8ReformerName;\n std::vector params;\n params.resize(header_.inverted_list_count);\n for (size_t i = 0; i < quantizers.size(); ++i) {\n auto &p = quantizers[i]->meta().reformer_params();\n auto &scale_key = int8_quantizer ? INT8_QUANTIZER_REFORMER_SCALE\n : INT4_QUANTIZER_REFORMER_SCALE;\n auto &bias_key = int8_quantizer ? INT8_QUANTIZER_REFORMER_BIAS\n : INT4_QUANTIZER_REFORMER_BIAS;\n if (inverted_lists_meta_[i].vector_count > 0 &&\n (!p.has(scale_key) || !p.has(bias_key))) {\n LOG_ERROR(\"Miss reformer params %s or %s\", bias_key.c_str(),\n scale_key.c_str());\n return IndexError_Logic;\n }\n\n params[i].bias = p.get_as_float(bias_key);\n params[i].scale = p.get_as_float(scale_key);\n }\n\n return this->dump_segment(\n int8_quantizer ? IVF_INT8_QUANTIZED_PARAMS_SEG_ID\n : IVF_INT4_QUANTIZED_PARAMS_SEG_ID,\n params.data(), params.size() * sizeof(InvertedIntegerQuantizerParams));\n}\n\nint IVFDumper::dump_original_vector(const void *data, size_t size) {\n if (dumped_feature_count_ >= header_.total_vector_count) {\n LOG_ERROR(\"Dump too much orignal features, expect=%u\",\n header_.total_vector_count);\n return IndexError_Logic;\n }\n\n if (dumper_->write(data, size) != size) {\n LOG_ERROR(\"Dumper write features failed\");\n return IndexError_WriteData;\n }\n dumped_features_size_ += size;\n ++dumped_feature_count_;\n if (dumped_feature_count_ == header_.total_vector_count) {\n //! Dump features finished, dump the meta\n size_t padding_size = 0;\n int ret = this->dump_padding(size, &padding_size);\n ivf_check_error_code(ret);\n\n ret = dumper_->append(IVF_FEATURES_SEG_ID, dumped_features_size_,\n padding_size, 0);\n if (ret != 0) {\n LOG_ERROR(\"Dumper append segment %s failed, ret:%d\",\n IVF_FEATURES_SEG_ID.c_str(), ret);\n return ret;\n }\n dumped_size_ += dumped_features_size_;\n }\n\n return 0;\n}\n\nint IVFDumper::check_dump_inverted_list(uint32_t inverted_list_id) {\n if (inverted_list_id < cur_list_id_) {\n LOG_ERROR(\"Invalid backward vector dumping, want=%u cur=%u\",\n inverted_list_id, cur_list_id_);\n return IndexError_Logic;\n }\n if (inverted_list_id >= inverted_lists_meta_.size()) {\n LOG_ERROR(\"Invalid inverted_list_id=%u, lists_size=%zu\", inverted_list_id,\n inverted_lists_meta_.size());\n return IndexError_Logic;\n }\n if (inverted_list_id != cur_list_id_) {\n //! flush previous inverted_list block\n int ret = this->dump_block();\n ivf_check_error_code(ret);\n for (auto idx = cur_list_id_ + 1; idx <= inverted_list_id; ++idx) {\n inverted_lists_meta_[idx].offset = header_.inverted_body_size;\n inverted_lists_meta_[idx].id_offset = header_.total_vector_count;\n }\n cur_list_id_ = inverted_list_id;\n }\n\n return 0;\n}\n\nint IVFDumper::dump_offsets_segment(void) const {\n bool col_pri = meta_.major_order() == IndexMeta::MajorOrder::MO_COLUMN;\n size_t total_size = 0;\n for (size_t i = 0; i < inverted_lists_meta_.size(); ++i) {\n std::vector offsets;\n const auto &m = inverted_lists_meta_[i];\n size_t vec_cnt = m.vector_count;\n size_t idx = 0;\n uint64_t off = m.offset;\n size_t align_idx = vec_cnt - vec_cnt % block_vector_count_;\n for (size_t j = 0; j < vec_cnt; ++j) {\n if (col_pri && j < align_idx) {\n offsets.emplace_back(off + idx * block_.align_size(), true);\n } else {\n offsets.emplace_back(off + idx * block_.element_size(), false);\n }\n ++idx;\n if (idx == block_vector_count_) {\n off += header_.block_size;\n idx = 0;\n }\n }\n if (idx != 0) {\n off += (vec_cnt - align_idx) * meta_.element_size();\n }\n\n size_t len = offsets.size() * sizeof(offsets[0]);\n size_t actual_len = dumper_->write(offsets.data(), len);\n if (actual_len != len) {\n LOG_ERROR(\"Write offsets failed expect %zu, actual: %zu.\", len,\n actual_len);\n return IndexError_WriteData;\n }\n total_size += len;\n }\n\n size_t padding_size = 0;\n int ret = this->dump_padding(total_size, &padding_size);\n ivf_check_error_code(ret);\n\n ret = dumper_->append(IVF_OFFSETS_SEG_ID, total_size, padding_size, 0);\n if (ret != 0) {\n LOG_ERROR(\"Dumper append segment %s failed, ret:%d\",\n IVF_OFFSETS_SEG_ID.c_str(), ret);\n return ret;\n }\n\n dumped_size_ += total_size + padding_size;\n\n return 0;\n}\n\nint IVFDumper::dump_segment(const std::string &segment_id, const void *data,\n size_t size) const {\n size_t len = dumper_->write(data, size);\n if (len != size) {\n LOG_ERROR(\"Dump segment %s data failed, expect=%zu, actual=%zu\",\n segment_id.c_str(), size, len);\n return IndexError_WriteData;\n }\n\n size_t padding_size = 0;\n int ret = this->dump_padding(size, &padding_size);\n ivf_check_error_code(ret);\n\n uint32_t crc = ailego::Crc32c::Hash(data, size);\n ret = dumper_->append(segment_id, size, padding_size, crc);\n if (ret != 0) {\n LOG_ERROR(\"Dump segment %s meta failed, ret=%d\", segment_id.c_str(), ret);\n return ret;\n }\n dumped_size_ += size + padding_size;\n\n return 0;\n}\n\nint IVFDumper::dump_padding(size_t data_size, size_t *padding_size) const {\n *padding_size = IVFUtility::AlignedSize(data_size) - data_size;\n if (*padding_size == 0) {\n return 0;\n }\n\n std::string padding(*padding_size, '\\0');\n if (dumper_->write(padding.data(), *padding_size) != *padding_size) {\n LOG_ERROR(\"Append padding failed, size %lu\", *padding_size);\n return IndexError_WriteData;\n }\n\n return 0;\n}\n\nint IVFDumper::dump_block(void) {\n if (block_.empty()) {\n return 0;\n }\n\n size_t size = ailego_align(block_.bytes(), 32);\n if (dumper_->write(block_.data(), size) != size) {\n LOG_ERROR(\"Failed to write data into dumper %s\", dumper_->name().c_str());\n return IndexError_WriteData;\n }\n auto &keys = block_.keys();\n std::copy(keys.begin(), keys.end(), std::back_inserter(keys_));\n ++inverted_lists_meta_[cur_list_id_].block_count;\n ++header_.block_count;\n header_.inverted_body_size += size;\n block_.clear();\n\n return 0;\n}\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/algorithm/ivf/ivf_dumper.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"metric/metric_params.h\"\n#include \"ivf_index_format.h\"\n#include \"ivf_params.h\"\n#include \"ivf_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Quantized Clustering Dumper\n */\nclass IVFDumper {\n public:\n typedef std::shared_ptr Pointer;\n\n //! Vectors block\n class Block {\n public:\n //! Initialize block\n void init(const IndexMeta &meta, uint32_t max_vec_count) {\n element_size_ = meta.element_size();\n auto bsize = IVFUtility::AlignedSize(max_vec_count, element_size_);\n data_.resize(bsize);\n count_ = 0u;\n major_order_ = meta.major_order();\n align_size_ = IndexMeta::AlignSizeof(meta.data_type());\n units_ = element_size_ / align_size_;\n max_vec_count_ = max_vec_count;\n keys_.reserve(max_vec_count_);\n }\n\n //! Add a vector to the block in row major order\n //! If the block is full and the block order is column, make a\n //! transpose\n void emplace(uint64_t key, const void *vec, IndexMeta::MajorOrder order) {\n switch (align_size_) {\n case 2:\n do_emplace(vec, order);\n break;\n case 4:\n do_emplace(vec, order);\n break;\n case 8:\n do_emplace(vec, order);\n break;\n default:\n ailego_check_with(false, \"Unsupport Aligned Size\");\n }\n keys_.emplace_back(key);\n }\n\n bool full(void) const {\n return count_ == max_vec_count_;\n }\n\n const void *data(void) const {\n return data_.data();\n }\n\n void clear(void) {\n count_ = 0u;\n keys_.clear();\n }\n\n bool empty(void) const {\n return count_ == 0u;\n }\n\n size_t size(void) const {\n return count_;\n }\n\n size_t capacity(void) const {\n return max_vec_count_;\n }\n\n size_t align_size(void) const {\n return align_size_;\n }\n\n size_t element_size(void) const {\n return element_size_;\n }\n\n //! Retrieve block data size\n size_t bytes(void) const {\n return element_size_ * count_;\n }\n\n //! Retrieve max block size in bytes\n size_t block_size(void) const {\n return data_.size();\n }\n\n IndexMeta::MajorOrder major_order(void) const {\n return major_order_;\n }\n\n const std::vector &keys(void) const {\n return keys_;\n }\n\n bool match_order(IndexMeta::MajorOrder column_major) const {\n return major_order_ == column_major;\n }\n\n private:\n //! Transpose the block vectors\n void transpose() {\n std::vector buf(data_.size());\n IVFUtility::Transpose(align_size_, data_.data(), count_, units_,\n buf.data());\n data_.swap(buf);\n }\n\n template \n void do_emplace(const void *vec, IndexMeta::MajorOrder order) {\n ailego_assert_with(count_ < max_vec_count_, \"emplace a full block\");\n\n T *dst = reinterpret_cast(data_.data() + element_size_ * count_);\n const T *src = reinterpret_cast(vec);\n size_t step = order == IndexMeta::MO_ROW ? 1 : max_vec_count_;\n for (auto i = 0u; i < units_; ++i) {\n *dst = *src;\n dst++;\n src += step;\n }\n\n count_++;\n if (full() && major_order_ == IndexMeta::MO_COLUMN) {\n transpose();\n }\n }\n\n private:\n //! Members\n std::vector data_{};\n std::vector keys_{};\n uint32_t count_{0u};\n uint32_t units_{0u};\n uint32_t align_size_{0u};\n uint32_t element_size_{0u};\n uint32_t max_vec_count_{0u};\n IndexMeta::MajorOrder major_order_{};\n };\n\n //! Constructor\n IVFDumper(const IndexMeta &meta, const IndexDumper::Pointer &dumper,\n size_t inverted_list_count, size_t block_vector_count)\n : meta_(meta),\n dumper_(dumper),\n block_vector_count_(block_vector_count),\n inverted_lists_meta_(inverted_list_count) {\n block_.init(meta, block_vector_count_);\n header_.block_vector_count = block_vector_count_;\n }\n\n //! Constructor\n IVFDumper(const IndexMeta &meta, const IndexDumper::Pointer &dumper,\n size_t inverted_list_count)\n : IVFDumper(meta, dumper, inverted_list_count, kDefaultBlockCount) {}\n\n //! Destructor\n ~IVFDumper() {\n // Check the dumper status\n if (dumped_feature_count_ > 0 &&\n dumped_feature_count_ != header_.total_vector_count) {\n LOG_ERROR(\"Dumped features=%u mismatch from invertedVecs=%u\",\n dumped_feature_count_, header_.total_vector_count);\n ailego_assert_with(false, \"invalid status\");\n }\n }\n\n //! Dump a vector in row major order\n int dump_inverted_vector(uint32_t inverted_list_id, uint64_t key,\n const void *vec);\n\n int dump_inverted_block(uint32_t inverted_list_id, const uint64_t *keys,\n const void *vecs, uint32_t vector_count,\n bool column_major);\n\n //! Finish dump the inverted vectors\n int dump_inverted_vector_finished(void);\n\n //! Dump the centroids index\n int dump_centroid_index(const void *data, size_t size);\n\n //! Dump params for each inverted list quantizer\n int dump_quantizer_params(\n const std::vector &quantizers);\n\n //! Dump the original vector, which doesnot been quantized\n int dump_original_vector(const void *data, size_t size);\n\n //! Retrieve total dumped size\n size_t dumped_size(void) const {\n return dumped_size_;\n }\n\n //! Retrieve total dumped vector count\n size_t dumped_count(void) const {\n return header_.total_vector_count;\n }\n\n //! Dump the segment from container\n int dump_container_segment(const IndexStorage::Pointer &container,\n const std::string &segmemt_id);\n\n private:\n int check_dump_inverted_list(uint32_t inverted_list_id);\n\n //! Dump offsets segment\n int dump_offsets_segment(void) const;\n\n //! Dump a segment\n int dump_segment(const std::string &segment_id, const void *data,\n size_t size) const;\n\n //! Dump segment padding\n int dump_padding(size_t data_size, size_t *padding_size) const;\n\n //! Dump a vector block\n int dump_block(void);\n\n private:\n //! Constants\n static constexpr size_t kDefaultBlockCount = 32u;\n\n //! Members\n Block block_{}; // vectors grouped in block\n const IndexMeta meta_{}; // IndexMeta of the inverted index\n const IndexDumper::Pointer dumper_{};\n size_t block_vector_count_{kDefaultBlockCount};\n std::vector inverted_lists_meta_{};\n std::vector keys_{};\n InvertedIndexHeader header_{};\n uint32_t cur_list_id_{0};\n uint32_t dumped_feature_count_{0};\n size_t dumped_features_size_{0};\n mutable size_t dumped_size_{0};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_entity.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"ivf_entity.h\"\n#include \n#include \"ivf_utility.h\"\nnamespace zvec {\nnamespace core {\n\n//! Initialize\nint IVFEntity::IVFReformerWrapper::init(const IndexMeta &imeta) {\n auto &name = imeta.reformer_name();\n\n if (name.empty()) {\n type_ = kReformerTpNone;\n return 0;\n }\n\n auto reformer = IndexFactory::CreateReformer(name);\n if (!reformer) {\n LOG_ERROR(\"Failed to create reformer %s\", name.c_str());\n return IndexError_NoExist;\n }\n int ret = reformer->init(imeta.reformer_params());\n ivf_check_with_msg(ret, \"Failed to init reformer %s\", name.c_str());\n\n reformer_ = std::move(reformer);\n\n if (name == kInt8ReformerName) {\n if (imeta.metric_name() == kIPMetricName) {\n type_ = kReformerTpInnerProductInt8;\n return 0;\n }\n auto &key = INT8_QUANTIZER_REFORMER_SCALE;\n if (!imeta.reformer_params().has(key)) {\n LOG_ERROR(\"Missing param %s in reformer %s\", key.c_str(), name.c_str());\n return IndexError_InvalidArgument;\n };\n float scale = imeta.reformer_params().get_as_float(key);\n reciprocal_ = scale == 0.0 ? 1.0 : (1.0 / scale);\n type_ = kReformerTpInt8;\n } else if (name == kInt4ReformerName) {\n if (imeta.metric_name() == kIPMetricName) {\n type_ = kReformerTpInnerProductInt4;\n return 0;\n }\n auto &key = INT4_QUANTIZER_REFORMER_SCALE;\n if (!imeta.reformer_params().has(key)) {\n LOG_ERROR(\"Missing param %s in reformer %s\", key.c_str(), name.c_str());\n return IndexError_InvalidArgument;\n };\n float scale = imeta.reformer_params().get_as_float(key);\n reciprocal_ = scale == 0.0 ? 1.0 : (1.0 / scale);\n type_ = kReformerTpInt4;\n } else {\n type_ = kReformerTpDefault;\n }\n\n LOG_DEBUG(\"Init QcReformer with %s, type=%u\", name.c_str(), type_);\n\n return 0;\n}\n\n//! Update the params, Called by gpu searcher only\nint IVFEntity::IVFReformerWrapper::update(const IndexMeta &meta) {\n auto &name = meta.reformer_name();\n if (name == kInt4ReformerName && meta.metric_name() == kL2MetricName) {\n auto &key = INT4_QUANTIZER_REFORMER_SCALE;\n if (!meta.reformer_params().has(key)) {\n LOG_ERROR(\"Missing param %s in reformer %s\", key.c_str(), name.c_str());\n return IndexError_InvalidArgument;\n };\n float scale = meta.reformer_params().get_as_float(key);\n reciprocal_ = scale == 0.0 ? 1.0 : (1.0 / scale / kNormalizeScaleFactor);\n type_ = kReformerTpInt8;\n\n ailego::Params params;\n float int8_scale = scale * kNormalizeScaleFactor;\n params.set(INT8_QUANTIZER_REFORMER_SCALE, int8_scale);\n float bias =\n meta.reformer_params().get_as_float(INT4_QUANTIZER_REFORMER_BIAS);\n params.set(INT8_QUANTIZER_REFORMER_BIAS, bias);\n params.set(\n INT4_QUANTIZER_REFORMER_METRIC,\n meta.reformer_params().get_as_string(INT4_QUANTIZER_REFORMER_METRIC));\n\n auto reformer = IndexFactory::CreateReformer(kInt8ReformerName);\n if (!reformer) {\n LOG_ERROR(\"Failed to create reformer %s\", name.c_str());\n return IndexError_NoExist;\n }\n int ret = reformer->init(params);\n ivf_check_with_msg(ret, \"Failed to init reformer %s\", name.c_str());\n\n reformer_ = reformer;\n\n LOG_DEBUG(\"Init QcReformer with %s, type=%u\", name.c_str(), type_);\n }\n\n return 0;\n}\n\n//! Transform a query\nint IVFEntity::IVFReformerWrapper::transform(const void *query,\n const IndexQueryMeta &qmeta,\n const void **out,\n IndexQueryMeta *ometa) {\n int ret = 0;\n\n switch (type_) {\n case kReformerTpNone:\n *out = query;\n *ometa = qmeta;\n break;\n\n case kReformerTpInnerProductInt8:\n if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) {\n return IndexError_Unsupported;\n }\n scales_.resize(1);\n buffer_.resize(IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT8,\n qmeta.dimension()));\n this->transform(0, static_cast(query), qmeta.dimension(),\n reinterpret_cast(&buffer_[0]));\n *ometa = qmeta;\n ometa->set_meta(IndexMeta::DataType::DT_INT8, qmeta.dimension());\n *out = buffer_.data();\n break;\n\n case kReformerTpInnerProductInt4:\n if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) {\n return IndexError_Unsupported;\n }\n scales_.resize(1);\n buffer_.resize(IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT4,\n qmeta.dimension()));\n this->transform(0, static_cast(query), qmeta.dimension(),\n reinterpret_cast(&buffer_[0]));\n *ometa = qmeta;\n ometa->set_meta(IndexMeta::DataType::DT_INT4, qmeta.dimension());\n *out = buffer_.data();\n break;\n\n case kReformerTpInt8:\n case kReformerTpInt4:\n /* FALLTHRU */\n case kReformerTpDefault:\n ret = reformer_->transform(query, qmeta, &buffer_, ometa);\n *out = buffer_.data();\n break;\n\n default:\n ret = IndexError_Unsupported;\n break;\n }\n\n return ret;\n}\n\n//! Transform querys\nint IVFEntity::IVFReformerWrapper::transform(const void *query,\n const IndexQueryMeta &qmeta,\n uint32_t count, const void **out,\n IndexQueryMeta *ometa) {\n int ret = 0;\n\n switch (type_) {\n case kReformerTpNone:\n *out = query;\n *ometa = qmeta;\n break;\n\n case kReformerTpInnerProductInt8:\n if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) {\n return IndexError_Unsupported;\n }\n scales_.resize(count);\n buffer_.resize(count *\n IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT8,\n qmeta.dimension()));\n {\n const float *ivec = reinterpret_cast(query);\n int8_t *ovec = reinterpret_cast(&buffer_[0]);\n for (size_t i = 0; i < count; ++i) {\n this->transform(i, &ivec[i * qmeta.dimension()], qmeta.dimension(),\n &ovec[i * qmeta.dimension()]);\n }\n }\n *ometa = qmeta;\n ometa->set_meta(IndexMeta::DataType::DT_INT8, qmeta.dimension());\n *out = buffer_.data();\n break;\n\n case kReformerTpInnerProductInt4:\n if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) {\n return IndexError_Unsupported;\n }\n scales_.resize(count);\n buffer_.resize(count *\n IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT4,\n qmeta.dimension()));\n {\n const float *ivec = reinterpret_cast(query);\n uint8_t *ovec = reinterpret_cast(&buffer_[0]);\n for (size_t i = 0; i < count; ++i) {\n this->transform(i, &ivec[i * qmeta.dimension()], qmeta.dimension(),\n &ovec[i * qmeta.dimension() / 2]);\n }\n }\n *ometa = qmeta;\n ometa->set_meta(IndexMeta::DataType::DT_INT4, qmeta.dimension());\n *out = buffer_.data();\n break;\n\n case kReformerTpInt8:\n case kReformerTpInt4:\n /* FALLTHRU */\n case kReformerTpDefault:\n ret = reformer_->transform(query, qmeta, count, &buffer_, ometa);\n *out = buffer_.data();\n break;\n\n default:\n ret = IndexError_Unsupported;\n break;\n }\n\n return ret;\n}\n\n//! Transform querys\nint IVFEntity::IVFReformerWrapper::transform_gpu(const void *query,\n const IndexQueryMeta &qmeta,\n uint32_t count,\n const void **out,\n IndexQueryMeta *ometa) {\n int ret = 0;\n\n switch (type_) {\n case kReformerTpNone:\n case kReformerTpDefault:\n *out = query;\n *ometa = qmeta;\n break;\n\n case kReformerTpInnerProductInt4:\n case kReformerTpInnerProductInt8:\n if (qmeta.data_type() != IndexMeta::DataType::DT_FP32) {\n return IndexError_Unsupported;\n }\n scales_.resize(count);\n buffer_.resize(count *\n IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT8,\n qmeta.dimension()));\n {\n const float *ivec = reinterpret_cast(query);\n int8_t *ovec = reinterpret_cast(&buffer_[0]);\n for (size_t i = 0; i < count; ++i) {\n this->transform(i, &ivec[i * qmeta.dimension()], qmeta.dimension(),\n &ovec[i * qmeta.dimension()]);\n }\n }\n *ometa = qmeta;\n ometa->set_meta(IndexMeta::DataType::DT_INT8, qmeta.dimension());\n *out = buffer_.data();\n break;\n\n case kReformerTpInt8:\n case kReformerTpInt4:\n ret = reformer_->transform(query, qmeta, count, &buffer_, ometa);\n *out = buffer_.data();\n break;\n\n default:\n ret = IndexError_Unsupported;\n break;\n }\n\n return ret;\n}\n\n\n//! Convert a record\nint IVFEntity::IVFReformerWrapper::convert(const void *record,\n const IndexQueryMeta &rmeta,\n const void **out,\n IndexQueryMeta *ometa) {\n if (type_ == kReformerTpNone) {\n *out = record;\n *ometa = rmeta;\n return 0;\n }\n\n int ret = reformer_->convert(record, rmeta, &buffer_, ometa);\n *out = buffer_.data();\n return ret;\n}\n\n//! Convert records\nint IVFEntity::IVFReformerWrapper::convert(const void *records,\n const IndexQueryMeta &rmeta,\n uint32_t count, const void **out,\n IndexQueryMeta *ometa) {\n if (type_ == kReformerTpNone) {\n *out = records;\n *ometa = rmeta;\n return 0;\n }\n int ret = reformer_->convert(records, rmeta, count, &buffer_, ometa);\n *out = buffer_.data();\n return ret;\n}\n\n//! Normalize score\nvoid IVFEntity::IVFReformerWrapper::normalize(size_t qidx,\n IndexDocumentHeap *heap) const {\n switch (type_) {\n case kReformerTpNone:\n return;\n\n case kReformerTpInnerProductInt8:\n case kReformerTpInnerProductInt4:\n ailego_assert_with(qidx < scales_.size(), \"invalid index\");\n {\n auto reciprocal = 1.0f / scales_[qidx];\n for (auto &it : *heap) {\n *it.mutable_score() *= reciprocal;\n }\n }\n break;\n\n case kReformerTpInt8:\n case kReformerTpInt4:\n for (auto &it : *heap) {\n *it.mutable_score() *= reciprocal_;\n }\n break;\n\n default:\n // Not support\n break;\n }\n}\n\n//! Normalize score\nvoid IVFEntity::IVFReformerWrapper::normalize(size_t qidx, const void *query,\n const IndexQueryMeta &qmeta,\n IndexDocumentHeap *heap) const {\n switch (type_) {\n case kReformerTpNone:\n return;\n\n case kReformerTpInnerProductInt8:\n case kReformerTpInnerProductInt4:\n ailego_assert_with(qidx < scales_.size(), \"invalid index\");\n {\n auto reciprocal = 1.0f / scales_[qidx];\n for (auto &it : *heap) {\n *it.mutable_score() *= reciprocal;\n }\n }\n break;\n\n case kReformerTpInt8:\n case kReformerTpInt4:\n for (auto &it : *heap) {\n *it.mutable_score() *= reciprocal_;\n }\n break;\n\n case kReformerTpDefault:\n reformer_->normalize(query, qmeta, *heap);\n break;\n\n default:\n // Not support\n LOG_ERROR(\"Not a supported type in QC reformer, type: %u\", type_);\n break;\n }\n}\n\nvoid IVFEntity::IVFReformerWrapper::transform(size_t qidx, const float *in,\n size_t dim, int8_t *out) {\n ailego_assert_with(qidx < scales_.size(), \"invalid index\");\n\n float abs_max = 0.0f;\n for (size_t i = 0; i < dim; ++i) {\n auto abs = std::abs(in[i]);\n if (abs > abs_max) {\n abs_max = abs;\n }\n }\n\n if (abs_max > 0.0f) {\n float scale = 127 / abs_max;\n for (size_t i = 0; i < dim; ++i) {\n out[i] = static_cast(std::round(in[i] * scale));\n }\n scales_[qidx] = scale;\n } else {\n std::fill(out, out + dim, static_cast(1));\n scales_[qidx] = std::numeric_limits::max();\n }\n}\n\nvoid IVFEntity::IVFReformerWrapper::transform(size_t qidx, const float *in,\n size_t dim, uint8_t *out) {\n ailego_assert_with(qidx < scales_.size(), \"invalid index\");\n ailego_assert_with(dim % 2 == 0, \"invalid dim\");\n\n float abs_max = 0.0f;\n float max = -std::numeric_limits::max();\n for (size_t i = 0; i < dim; ++i) {\n float abs = std::abs(in[i]);\n abs_max = std::max(abs_max, abs);\n max = std::max(max, in[i]);\n }\n if (abs_max > 0.0f) {\n float scale = ((7 * abs_max > 8 * max) ? 8 : 7) / abs_max;\n for (size_t i = 0; i < dim; i += 2) {\n auto v1 = static_cast(std::round(in[i] * scale));\n auto v2 = static_cast(std::round(in[i + 1] * scale));\n out[i / 2] =\n (static_cast(v1) & 0xF) | (static_cast(v2) << 4);\n }\n scales_[qidx] = scale;\n } else {\n std::fill(out, out + dim / 2, static_cast(9));\n scales_[qidx] = std::numeric_limits::max();\n }\n}\n\nint IVFEntity::load_header(const IndexStorage::Pointer &container) {\n //! Load the Header Segment\n auto header = container->get(IVF_INVERTED_HEADER_SEG_ID);\n if (!header) {\n LOG_ERROR(\"Failed to get segment %s\", IVF_INVERTED_HEADER_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n if (header->data_size() < sizeof(header_)) {\n LOG_ERROR(\"Invalid format for segment %s\",\n IVF_INVERTED_HEADER_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n const void *data = nullptr;\n if (header->read(0, &data, header->data_size()) != header->data_size()) {\n LOG_ERROR(\"Failed to read data, segment %s\",\n IVF_INVERTED_HEADER_SEG_ID.c_str());\n return IndexError_ReadData;\n }\n std::memcpy(&header_, data, sizeof(header_));\n if (header_.header_size < sizeof(header_) + header_.index_meta_size ||\n header_.header_size > header->data_size()) {\n LOG_ERROR(\"Invalid header size %u\", header_.header_size);\n return IndexError_InvalidFormat;\n }\n\n //! Load the index meta\n if (!meta_.deserialize(\n reinterpret_cast(data) + sizeof(header_),\n header_.index_meta_size)) {\n LOG_ERROR(\"Failed to deserialize index meta\");\n return IndexError_InvalidFormat;\n }\n\n int ret = reformer_.init(meta_);\n ivf_check_error_code(ret);\n\n //! Create the distance calculator\n auto metric = IndexFactory::CreateMetric(meta_.metric_name());\n if (!metric) {\n LOG_ERROR(\"Failed to create metric %s\", meta_.metric_name().c_str());\n return IndexError_NoExist;\n }\n ret = metric->init(meta_, meta_.metric_params());\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize metric %s\", meta_.metric_name().c_str());\n return ret;\n }\n calculator_ = std::make_shared(\n meta_, metric->query_metric() ? metric->query_metric() : metric,\n header_.block_vector_count);\n if (!calculator_) {\n return IndexError_NoMemory;\n }\n\n return 0;\n}\n\nint IVFEntity::load(const IndexStorage::Pointer &container) {\n int ret = this->load_header(container);\n ivf_check_error_code(ret);\n\n //! Load the remaining segments\n container_ = container;\n size_t expect_size = header_.inverted_body_size;\n inverted_ = load_segment(IVF_INVERTED_BODY_SEG_ID, expect_size);\n if (!inverted_) {\n LOG_ERROR(\"Failed to load segment, inverted_size=%zu block_count=%u\",\n static_cast(header_.inverted_body_size),\n header_.block_count);\n return IndexError_InvalidFormat;\n }\n\n expect_size = header_.inverted_list_count * sizeof(InvertedListMeta);\n inverted_meta_ = load_segment(IVF_INVERTED_META_SEG_ID, expect_size);\n if (!inverted_meta_) {\n LOG_ERROR(\"Failed to load segment, inverted_lists=%u\",\n header_.inverted_list_count);\n return IndexError_InvalidFormat;\n }\n\n expect_size = header_.total_vector_count * sizeof(uint64_t);\n keys_ = load_segment(IVF_KEYS_SEG_ID, expect_size);\n if (!keys_) {\n return IndexError_InvalidFormat;\n }\n\n expect_size = header_.total_vector_count * sizeof(InvertedVecLocation);\n offsets_ = load_segment(IVF_OFFSETS_SEG_ID, expect_size);\n if (!offsets_) {\n return IndexError_InvalidFormat;\n }\n\n expect_size = header_.total_vector_count * sizeof(uint32_t);\n mapping_ = load_segment(IVF_MAPPING_SEG_ID, expect_size);\n if (!mapping_) {\n return IndexError_InvalidFormat;\n }\n\n norm_value_sqrt_ =\n meta_.metric_name() == \"Euclidean\" || meta_.metric_name() == \"Manhattan\";\n if (container_->get(IVF_INT8_QUANTIZED_PARAMS_SEG_ID) ||\n container->get(IVF_INT4_QUANTIZED_PARAMS_SEG_ID)) {\n expect_size =\n header_.inverted_list_count * sizeof(InvertedIntegerQuantizerParams);\n auto &seg_id = meta_.reformer_name() == kInt8ReformerName\n ? IVF_INT8_QUANTIZED_PARAMS_SEG_ID\n : IVF_INT4_QUANTIZED_PARAMS_SEG_ID;\n integer_quantizer_params_ = load_segment(seg_id, expect_size);\n if (!integer_quantizer_params_) {\n return IndexError_InvalidFormat;\n }\n norm_value_ = 0.0f;\n } else if (meta_.reformer_name() == kInt8ReformerName ||\n meta_.reformer_name() == kInt4ReformerName) {\n auto &scale_key = meta_.reformer_name() == kInt8ReformerName\n ? INT8_QUANTIZER_REFORMER_SCALE\n : INT4_QUANTIZER_REFORMER_SCALE;\n auto scale = meta_.reformer_params().get_as_float(scale_key);\n norm_value_ = this->convert_to_normalize_value(scale);\n } else {\n norm_value_ = 1.0f;\n }\n\n if (container_->get(IVF_FEATURES_SEG_ID)) {\n features_ = load_segment(IVF_FEATURES_SEG_ID, 0);\n if (!features_) {\n return IndexError_InvalidFormat;\n }\n if (features_->data_size() % vector_count() != 0) {\n LOG_ERROR(\"Invalid featureSegment size=%zu, totalVecs=%zu\",\n features_->data_size(), vector_count());\n return IndexError_InvalidFormat;\n }\n }\n\n LOG_DEBUG(\n \"Load inverted index done, docs=%u invertedListCnt=%u \"\n \"elementSize=%u metric=%s reformer=%s\",\n header_.total_vector_count, header_.inverted_list_count,\n meta_.element_size(), meta_.metric_name().c_str(),\n meta_.reformer_name().c_str());\n return 0;\n}\n\nint IVFEntity::search(size_t inverted_list_id, const void *query,\n const IndexFilter &filter, uint32_t *scan_count,\n IndexDocumentHeap *heap,\n IndexContext::Stats *context_stats) const {\n ailego_assert_with(inverted_list_id < header_.inverted_list_count,\n \"invalid id\");\n auto list_meta = this->inverted_list_meta(inverted_list_id);\n ivf_assert(list_meta, IndexError_ReadData);\n\n const void *data = nullptr;\n const size_t block_vecs = header_.block_vector_count;\n std::vector distances(block_vecs);\n const size_t batch_size = kBatchBlocks;\n const size_t block_size = header_.block_size;\n const auto norm_val = this->inverted_list_normalize_value(inverted_list_id);\n for (size_t i = 0; i < list_meta->block_count; i += batch_size) {\n //! Read vecs\n const size_t off = list_meta->offset + i * block_size;\n const size_t blocks = std::min(batch_size, list_meta->block_count - i);\n const size_t size =\n std::min(blocks * block_size,\n static_cast(header_.inverted_body_size - off));\n if (inverted_->read(off, &data, size) != size) {\n LOG_ERROR(\"Failed to read block, off=%zu, size=%zu\", off, size);\n return IndexError_ReadData;\n }\n\n //! Read keys\n size_t items = std::min(blocks * block_vecs,\n list_meta->vector_count - (i * block_vecs));\n auto keys = get_keys(list_meta->id_offset + i * block_vecs, items);\n if (!keys) {\n return IndexError_ReadData;\n }\n\n //! Compute distances for each block\n for (size_t b = 0; b < blocks; ++b) {\n const size_t vecs_count =\n std::min(block_vecs, list_meta->vector_count - (i + b) * block_vecs);\n auto block_keys = keys + b * block_vecs;\n size_t keeps = 0;\n ailego_assert_with(block_vecs < sizeof(keeps) * 8, \"bits overflow\");\n for (size_t k = 0; k < vecs_count; ++k) {\n if (!filter(block_keys[k])) {\n keeps |= (1 << k);\n } else {\n ++(*context_stats->mutable_filtered_count());\n }\n }\n if (keeps == 0) {\n continue;\n }\n\n const void *block_data = static_cast(data) + b * block_size;\n calculator_->query_features_distance(query, block_data, vecs_count,\n distances.data());\n\n *(context_stats->mutable_dist_calced_count()) += vecs_count;\n\n uint32_t id_off = list_meta->id_offset + (i + b) * block_vecs;\n for (size_t k = 0; k < vecs_count; ++k) {\n if (keeps & (1 << k)) {\n if (block_keys[k] != kInvalidKey) {\n heap->emplace(block_keys[k], distances[k] * norm_val, id_off + k);\n }\n }\n }\n }\n }\n\n *scan_count = list_meta->vector_count;\n return 0;\n}\n\n//! search in inverted list without filter\nint IVFEntity::search(size_t inverted_list_id, const void *query,\n uint32_t *scan_count, IndexDocumentHeap *heap,\n IndexContext::Stats *context_stats) const {\n ailego_assert_with(inverted_list_id < header_.inverted_list_count,\n \"invalid id\");\n auto list_meta = inverted_list_meta(inverted_list_id);\n ivf_assert(list_meta, IndexError_ReadData);\n\n const void *data = nullptr;\n const size_t block_vecs = header_.block_vector_count;\n std::vector distances(block_vecs);\n const size_t batch_size = kBatchBlocks;\n const size_t block_size = header_.block_size;\n const auto norm_val = this->inverted_list_normalize_value(inverted_list_id);\n for (size_t i = 0; i < list_meta->block_count; i += batch_size) {\n //! Read vecs\n const size_t off = list_meta->offset + i * block_size;\n const size_t blocks = std::min(batch_size, list_meta->block_count - i);\n const size_t size =\n std::min(blocks * block_size,\n static_cast(header_.inverted_body_size - off));\n if (inverted_->read(off, &data, size) != size) {\n LOG_ERROR(\"Failed to read block, off=%zu, size=%zu\", off, size);\n return IndexError_ReadData;\n }\n\n //! Read keys\n size_t items = std::min(blocks * block_vecs,\n list_meta->vector_count - (i * block_vecs));\n auto keys = get_keys(list_meta->id_offset + i * block_vecs, items);\n if (!keys) {\n return IndexError_ReadData;\n }\n\n //! Compute distances for each block\n for (size_t b = 0; b < blocks; ++b) {\n const size_t vecs_count =\n std::min(block_vecs, list_meta->vector_count - (i + b) * block_vecs);\n auto block_keys = keys + b * block_vecs;\n const void *block_data = static_cast(data) + b * block_size;\n calculator_->query_features_distance(query, block_data, vecs_count,\n distances.data());\n for (size_t k = 0; k < vecs_count; ++k) {\n if (block_keys[k] != kInvalidKey) {\n uint32_t id = list_meta->id_offset + (i + b) * block_vecs + k;\n heap->emplace(block_keys[k], distances[k] * norm_val, id);\n }\n }\n *(context_stats->mutable_dist_calced_count()) += vecs_count;\n }\n }\n\n *scan_count = list_meta->vector_count;\n return 0;\n}\n\n//! search all inverted list with filter\nint IVFEntity::search(const void *query, const IndexFilter &filter,\n IndexDocumentHeap *heap,\n IndexContext::Stats *context_stats) const {\n for (size_t i = 0; i < header_.inverted_list_count; ++i) {\n uint32_t scan_count;\n int ret = this->search(i, query, filter, &scan_count, heap, context_stats);\n if (ret != 0) {\n return ret;\n }\n }\n\n return 0;\n}\n\n//! search all inverted list without filter\nint IVFEntity::search(const void *query, IndexDocumentHeap *heap,\n IndexContext::Stats *context_stats) const {\n for (size_t i = 0; i < header_.inverted_list_count; ++i) {\n uint32_t scan_count;\n int ret = this->search(i, query, &scan_count, heap, context_stats);\n if (ret != 0) {\n return ret;\n }\n }\n\n return 0;\n}\n\nconst void *IVFEntity::get_vector(size_t id) const {\n if (features_) {\n const void *data = nullptr;\n size_t element_size = features_->data_size() / vector_count();\n size_t off = id * element_size;\n if (features_->read(off, &data, element_size) != element_size) {\n LOG_ERROR(\"Failed to read segment, off=%zu size=%zu\", off, element_size);\n return nullptr;\n }\n return data;\n }\n\n const void *data = nullptr;\n size_t size = sizeof(InvertedVecLocation);\n if (offsets_->read(id * size, &data, size) != size) {\n LOG_ERROR(\"Failed to read offsets segment, id=%zu\", id);\n return nullptr;\n }\n auto &loc = *reinterpret_cast(data);\n if (loc.column_major) {\n vector_.resize(meta_.element_size());\n auto unit_size = IndexMeta::AlignSizeof(meta_.data_type());\n size_t cols = meta_.element_size() / unit_size;\n size_t step = block_vector_count() * unit_size;\n size_t rd_size = step * (cols - 1) + unit_size;\n if (inverted_->read(loc.offset, &data, rd_size) != rd_size) {\n LOG_ERROR(\"Failed to read data, off=%zu size=%zu\",\n static_cast(loc.offset), rd_size);\n return nullptr;\n }\n for (size_t c = 0; c < cols; ++c) {\n vector_.replace(c * unit_size, unit_size,\n reinterpret_cast(data) + c * step,\n unit_size);\n }\n return vector_.data();\n } else {\n if (inverted_->read(loc.offset, &data, meta_.element_size()) !=\n meta_.element_size()) {\n LOG_ERROR(\"Failed to read data, off=%zu size=%u\",\n static_cast(loc.offset), meta_.element_size());\n return nullptr;\n }\n return data;\n }\n}\n\nint IVFEntity::get_vector(size_t id, IndexStorage::MemoryBlock &block) const {\n if (features_) {\n size_t element_size = features_->data_size() / vector_count();\n size_t off = id * element_size;\n if (features_->read(off, block, element_size) != element_size) {\n LOG_ERROR(\"Failed to read segment, off=%zu size=%zu\", off, element_size);\n return IndexError_Runtime;\n }\n return 0;\n }\n\n\n IndexStorage::MemoryBlock data_block;\n size_t size = sizeof(InvertedVecLocation);\n if (offsets_->read(id * size, data_block, size) != size) {\n LOG_ERROR(\"Failed to read offsets segment, id=%zu\", id);\n return IndexError_Runtime;\n }\n const void *data = data_block.data();\n auto &loc = *reinterpret_cast(data);\n if (loc.column_major) {\n vector_.resize(meta_.element_size());\n auto unit_size = IndexMeta::AlignSizeof(meta_.data_type());\n size_t cols = meta_.element_size() / unit_size;\n size_t step = block_vector_count() * unit_size;\n size_t rd_size = step * (cols - 1) + unit_size;\n if (inverted_->read(loc.offset, &data, rd_size) != rd_size) {\n LOG_ERROR(\"Failed to read data, off=%zu size=%zu\",\n static_cast(loc.offset), rd_size);\n return IndexError_Runtime;\n }\n for (size_t c = 0; c < cols; ++c) {\n vector_.replace(c * unit_size, unit_size,\n reinterpret_cast(data) + c * step,\n unit_size);\n }\n block.reset(vector_.data());\n return 0;\n } else {\n if (inverted_->read(loc.offset, block, meta_.element_size()) !=\n meta_.element_size()) {\n LOG_ERROR(\"Failed to read data, off=%zu size=%u\",\n static_cast(loc.offset), meta_.element_size());\n return IndexError_Runtime;\n }\n return 0;\n }\n}\n\nuint32_t IVFEntity::key_to_id(uint64_t key) const {\n //! Do binary search\n uint32_t start = 0UL;\n uint32_t end = vector_count();\n const void *data = nullptr;\n uint32_t idx = 0u;\n while (start < end) {\n idx = start + (end - start) / 2;\n if (ailego_unlikely(mapping_->read(idx * sizeof(uint32_t), &data,\n sizeof(uint32_t)) != sizeof(uint32_t))) {\n LOG_ERROR(\"Failed to read mapping segment, idx=%u\", idx);\n return std::numeric_limits::max();\n }\n const uint64_t *mkey;\n uint32_t local_id = *reinterpret_cast(data);\n if (ailego_unlikely(keys_->read(local_id * sizeof(uint64_t),\n (const void **)(&mkey),\n sizeof(uint64_t)) != sizeof(uint64_t))) {\n LOG_ERROR(\"Read key from segment failed\");\n return std::numeric_limits::max();\n }\n if (*mkey < key) {\n start = idx + 1;\n } else if (*mkey > key) {\n end = idx;\n } else {\n return local_id;\n }\n }\n return std::numeric_limits::max();\n}\n\nconst void *IVFEntity::get_vector_by_key(uint64_t key) const {\n uint32_t id = this->key_to_id(key);\n if (id != std::numeric_limits::max()) {\n return get_vector(id);\n } else {\n return nullptr;\n }\n}\n\nint IVFEntity::get_vector_by_key(uint64_t key,\n IndexStorage::MemoryBlock &block) const {\n uint32_t id = this->key_to_id(key);\n if (id != std::numeric_limits::max()) {\n return get_vector(id, block);\n } else {\n return IndexError_Runtime;\n }\n}\n\nIVFEntity::Pointer IVFEntity::clone(void) const {\n auto entity = std::make_shared();\n return clone(entity);\n}\n\nIVFEntity::Pointer IVFEntity::clone(const IVFEntity::Pointer &entity) const {\n if (!entity) {\n LOG_ERROR(\"Failed to alloc IVFEntity\");\n return nullptr;\n }\n\n auto inverted = inverted_->clone();\n ivf_assert_with_msg(inverted, nullptr, \"Failed to clone inverted segment\");\n\n auto inverted_meta = inverted_meta_->clone();\n ivf_assert_with_msg(inverted_meta, nullptr,\n \"Failed to clone inverted meta segment\");\n\n auto keys = keys_->clone();\n ivf_assert_with_msg(keys, nullptr, \"Failed to clone keys segment\");\n\n auto offsets = offsets_->clone();\n ivf_assert_with_msg(offsets, nullptr, \"Failed to clone offsets segment\");\n\n auto mapping = mapping_->clone();\n ivf_assert_with_msg(mapping, nullptr, \"Failed to clone mapping segment\");\n\n IndexStorage::Segment::Pointer integer_quantizer_params;\n if (integer_quantizer_params_) {\n integer_quantizer_params = integer_quantizer_params_->clone();\n if (!integer_quantizer_params) {\n LOG_ERROR(\"Failed to clone integer quantizer params segment\");\n return nullptr;\n }\n }\n IndexStorage::Segment::Pointer features;\n if (features_) {\n features = features_->clone();\n if (!features) {\n LOG_ERROR(\"Failed to clone features segment\");\n return nullptr;\n }\n }\n\n entity->meta_ = this->meta_;\n entity->reformer_ = this->reformer_;\n entity->calculator_ = this->calculator_;\n entity->header_ = this->header_;\n entity->container_ = this->container_;\n\n entity->inverted_ = inverted;\n entity->inverted_meta_ = inverted_meta;\n entity->keys_ = keys;\n entity->offsets_ = offsets;\n entity->mapping_ = mapping;\n entity->integer_quantizer_params_ = integer_quantizer_params;\n entity->features_ = features;\n entity->norm_value_ = this->norm_value_;\n entity->norm_value_sqrt_ = this->norm_value_sqrt_;\n\n return entity;\n}\n\nIndexStorage::Segment::Pointer IVFEntity::load_segment(\n const std::string &seg_id, size_t expect_size) const {\n auto segment = container_->get(seg_id);\n if (!segment) {\n LOG_ERROR(\"Failed to get segment %s\", seg_id.c_str());\n return nullptr;\n }\n if (expect_size && segment->data_size() != expect_size) {\n LOG_ERROR(\"Invalid segment %s size=%zu, total_vecs=%u\", seg_id.c_str(),\n segment->data_size(), header_.total_vector_count);\n return nullptr;\n }\n return segment;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_entity.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \"metric/metric_params.h\"\n#include \"ivf_distance_calculator.h\"\n#include \"ivf_index_format.h\"\n#include \"ivf_params.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! IVF Entity\n */\nclass IVFEntity {\n public:\n typedef std::shared_ptr Pointer;\n\n class IVFReformerWrapper;\n\n //! Constructor\n IVFEntity() {}\n\n //! Destructor\n virtual ~IVFEntity() {}\n\n //! Disable them\n IVFEntity(const IVFEntity &) = delete;\n IVFEntity &operator=(const IVFEntity &) = delete;\n\n //! load the index from container\n virtual int load(const IndexStorage::Pointer &container);\n\n //! search in inverted list with filter\n int search(size_t inverted_list_id, const void *query,\n const IndexFilter &filter, uint32_t *scan_count,\n IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const;\n\n //! search in inverted list without filter\n int search(size_t inverted_list_id, const void *query, uint32_t *scan_count,\n IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const;\n\n //! search all inverted list with filter\n int search(const void *query, const IndexFilter &filter,\n IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const;\n\n //! search all inverted list without filter\n int search(const void *query, IndexDocumentHeap *heap,\n IndexContext::Stats *context_stats) const;\n\n //! Clone the entity\n virtual IVFEntity::Pointer clone(void) const;\n\n //! Clone the entity\n IVFEntity::Pointer clone(const IVFEntity::Pointer &entity) const;\n\n //! Retrieve the primary keys by local id in heap\n int retrieve_keys(IndexDocumentHeap *heap) const {\n for (auto &it : (*heap)) {\n uint64_t key = this->get_key(it.index());\n if (key == kInvalidKey) {\n return IndexError_ReadData;\n }\n it.set_key(key);\n }\n\n return 0;\n }\n\n //! Retrieve the total vectors in the index\n size_t vector_count(void) const {\n return header_.total_vector_count;\n }\n\n //! Retrieve the inverted list count\n size_t inverted_list_count(void) const {\n return header_.inverted_list_count;\n }\n\n //! Retrieve block size of the inverted vector\n size_t inverted_block_size(void) const {\n return header_.block_size;\n }\n\n //! Retrieve the vectors count in one block\n size_t block_vector_count(void) const {\n return header_.block_vector_count;\n }\n\n //! Retrieve IndexMeta of the inverted index\n const IndexMeta &meta(void) const {\n return meta_;\n }\n\n //! Retrieve a block of vectors\n const void *read_block(size_t inverted_list_id, size_t local_block_id,\n size_t *vecs_count) const {\n auto iv_meta = this->inverted_list_meta(inverted_list_id);\n if (!iv_meta || local_block_id >= iv_meta->block_count) {\n LOG_ERROR(\"Failed to read inverted list, listId=%zu blockIdx=%zu\",\n inverted_list_id, local_block_id);\n return nullptr;\n }\n\n size_t block_vecs = header_.block_vector_count;\n *vecs_count = std::min(block_vecs,\n iv_meta->vector_count - local_block_id * block_vecs);\n ailego_assert_with(*vecs_count <= header_.block_vector_count,\n \"invalid vecs\");\n const size_t off = iv_meta->offset + local_block_id * header_.block_size;\n const size_t size = *vecs_count * meta_.element_size();\n const void *data = nullptr;\n if (inverted_->read(off, &data, size) != size) {\n LOG_ERROR(\"Failed to read block off=%zu size=%zu\", off, size);\n return nullptr;\n }\n\n return data;\n }\n\n //! Retrieve the inverted list meta\n const InvertedListMeta *inverted_list_meta(size_t inverted_list_id) const {\n const void *data = nullptr;\n const size_t size = sizeof(InvertedListMeta);\n const size_t offset = inverted_list_id * size;\n if (inverted_meta_->read(offset, &data, size) != size) {\n LOG_ERROR(\"Failed to read inverted meta, id=%zu, size=%zu\",\n inverted_list_id, size);\n return nullptr;\n }\n\n return static_cast(data);\n }\n\n //! Retrieve the keys by consecutive local ids\n const uint64_t *get_keys(size_t id, size_t count) const {\n const void *data = nullptr;\n const size_t offset = id * sizeof(uint64_t);\n const size_t size = count * sizeof(uint64_t);\n if (keys_->read(offset, &data, size) != size) {\n LOG_ERROR(\"Failed to read keys, id=%zu, size=%zu\", id, size);\n return nullptr;\n }\n\n return static_cast(data);\n }\n\n //! Retrieve the key by local id\n uint64_t get_key(size_t id) const {\n const void *data = nullptr;\n const size_t offset = id * sizeof(uint64_t);\n const size_t size = sizeof(uint64_t);\n if (keys_->read(offset, &data, size) != size) {\n LOG_ERROR(\"Failed to read key, id=%zu\", id);\n return kInvalidKey;\n }\n\n return *static_cast(data);\n }\n\n //! Retrieve vector by local id\n const void *get_vector(size_t id) const;\n\n //! Retrieve vector by local id\n const void *get_vector_by_key(uint64_t key) const;\n\n int get_vector(size_t id, IndexStorage::MemoryBlock &block) const;\n\n int get_vector_by_key(uint64_t key, IndexStorage::MemoryBlock &block) const;\n\n uint32_t key_to_id(uint64_t key) const;\n\n //! Transform a query\n int transform(const void *query, const IndexQueryMeta &qmeta,\n const void **out, IndexQueryMeta *ometa) const {\n return reformer_.transform(query, qmeta, out, ometa);\n }\n\n //! Transform queries\n int transform(const void *query, const IndexQueryMeta &qmeta, uint32_t count,\n const void **out, IndexQueryMeta *ometa) const {\n return reformer_.transform(query, qmeta, count, out, ometa);\n }\n\n //! Normalize the score in query part\n void normalize(size_t qidx, IndexDocumentHeap *heap) const {\n return reformer_.normalize(qidx, heap);\n }\n\n //! Retrieve the value for each inverted list to multiply for normalizing\n float inverted_list_normalize_value(size_t inverted_list_id) const {\n if (norm_value_ != 0.0f) {\n return norm_value_;\n }\n\n // ailego_assert_with(integer_quantizer_params_, \"nullptr\");\n if (integer_quantizer_params_ != nullptr) {\n const void *data = nullptr;\n size_t size = sizeof(InvertedIntegerQuantizerParams);\n size_t off = inverted_list_id * size;\n if (integer_quantizer_params_->read(off, &data, size) != size) {\n LOG_ERROR(\"Failed to read data from segment, off=%zu\", off);\n return 1.0f;\n }\n auto scale =\n static_cast(data)->scale;\n return this->convert_to_normalize_value(scale);\n }\n\n return norm_value_;\n }\n\n //! Check whether the feature segment exist\n bool has_orignal_feature() const {\n return !!features_;\n }\n\n //! Retrieve reformer\n const IVFReformerWrapper &reformer(void) const {\n return reformer_;\n }\n\n /*! Index Reformer Wrapper\n * To transform query in inverted index searching, and normalize the score\n */\n class IVFReformerWrapper {\n public:\n //! Constructor\n IVFReformerWrapper() {}\n\n //! Assignment\n IVFReformerWrapper &operator=(const IVFReformerWrapper &wrapper) {\n reformer_ = wrapper.reformer_;\n type_ = wrapper.type_;\n buffer_.clear();\n buffer_.shrink_to_fit();\n reciprocal_ = wrapper.reciprocal_;\n return *this;\n }\n\n //! Initialize\n int init(const IndexMeta &imeta);\n\n //! Update\n int update(const IndexMeta &meta);\n\n //! Transform a query\n int transform(const void *query, const IndexQueryMeta &qmeta,\n const void **out, IndexQueryMeta *ometa);\n\n //! Transform queries\n int transform(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, const void **out, IndexQueryMeta *ometa);\n\n //! Convert a record\n virtual int convert(const void *record, const IndexQueryMeta &rmeta,\n const void **out, IndexQueryMeta *ometa);\n\n //! Convert records\n virtual int convert(const void *records, const IndexQueryMeta &rmeta,\n uint32_t count, const void **out,\n IndexQueryMeta *ometa);\n\n //! Transform queries\n int transform_gpu(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, const void **out, IndexQueryMeta *ometa);\n\n //! Normalize the score in query part\n void normalize(size_t qidx, IndexDocumentHeap *heap) const;\n\n //! Normalize the score in query part\n void normalize(size_t qidx, const void *query, const IndexQueryMeta &qmeta,\n IndexDocumentHeap *heap) const;\n\n private:\n //! Transform query from fp32 to int8\n void transform(size_t qidx, const float *in, size_t dim, int8_t *out);\n\n //! Transform query from fp32 to int4\n void transform(size_t qidx, const float *in, size_t dim, uint8_t *out);\n\n private:\n //! Constants\n enum Type {\n kReformerTpNone = 0,\n kReformerTpInnerProductInt8 = 1,\n kReformerTpInnerProductInt4 = 2,\n kReformerTpInt8 = 3,\n kReformerTpInt4 = 4,\n kReformerTpDefault = 7,\n };\n\n //! Members\n Type type_{kReformerTpNone};\n IndexReformer::Pointer reformer_{};\n std::string buffer_{};\n float reciprocal_{0.0}; // for int8\n std::vector scales_{}; // for int8 IP\n };\n\n private:\n //! Load the segment by seg_id in expect_size segment size\n IndexStorage::Segment::Pointer load_segment(const std::string &seg_id,\n size_t expect_size) const;\n\n //! Load the header segment\n int load_header(const IndexStorage::Pointer &container);\n\n //! Convert the int8 quantizer scale to normalize value\n float convert_to_normalize_value(float scale) const {\n auto v = scale == 0.0 ? 1.0 : (1.0 / scale);\n return !norm_value_sqrt_ ? v : std::sqrt(v);\n }\n\n protected:\n //! Constants\n static constexpr size_t kBatchBlocks = 10u;\n\n //! Members\n IndexMeta meta_{};\n mutable IVFReformerWrapper reformer_{};\n IVFDistanceCalculator::Pointer calculator_{};\n InvertedIndexHeader header_{};\n IndexStorage::Pointer container_{};\n IndexStorage::Segment::Pointer inverted_{};\n IndexStorage::Segment::Pointer inverted_meta_{};\n IndexStorage::Segment::Pointer keys_{};\n IndexStorage::Segment::Pointer offsets_{};\n IndexStorage::Segment::Pointer mapping_{};\n IndexStorage::Segment::Pointer features_{};\n IndexStorage::Segment::Pointer integer_quantizer_params_{};\n mutable std::string vector_{}; // temporary buffer for colomn major order\n float norm_value_{0.0f}; // normalize the inverted vector to orignal score\n bool norm_value_sqrt_{false}; // does the norm value need to sqrt\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_index_format.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nstatic constexpr uint64_t kInvalidKey = std::numeric_limits::max();\n\n/*! Index Format of Inverted Index Header\n */\nstruct InvertedIndexHeader {\n uint32_t header_size{0};\n uint32_t total_vector_count{0};\n uint64_t inverted_body_size{0};\n uint32_t inverted_list_count{0};\n uint32_t block_vector_count{0};\n uint32_t block_size{0};\n uint32_t block_count{0};\n uint32_t index_meta_size{0};\n char reserved_[28];\n char index_meta[0];\n};\n\n/*! Index Format of Inverted Index Meta for each Inverted list\n */\nstruct InvertedListMeta {\n uint64_t offset{0};\n uint32_t block_count{0};\n uint32_t vector_count{0};\n uint32_t id_offset{0};\n char reserved_[16];\n};\n\n/*! Index Format of Location in Inverted Index for each vector\n */\nstruct InvertedVecLocation {\n InvertedVecLocation(size_t off, bool col)\n : offset(off), column_major(col), reserved(0u) {}\n\n uint64_t offset : 48; // feature offset in posting block segment\n uint64_t column_major : 1; // coloum major if true\n uint64_t reserved : 15;\n};\n\n/*! Index Format of Integer Quantizer params for each inverted list\n */\nstruct InvertedIntegerQuantizerParams {\n float scale{1.0};\n float bias{0.0};\n};\n\n/*! Location of Vectors Block in Storage Segment\n */\nstruct BlockLocation {\n uint16_t segment_id;\n uint16_t block_index;\n};\n\n/*! The Header of a Block in Storage Segment\n */\nstruct BlockHeader {\n BlockLocation next;\n uint16_t vector_count;\n uint16_t column_major : 1;\n uint16_t reserved_ : 15;\n};\n\nstruct DeletionMap {\n void set(uint32_t index) {\n bitset.set(index);\n }\n\n void reset(uint32_t index) {\n bitset.reset(index);\n }\n\n bool test(uint32_t index) const {\n return bitset.test(index);\n }\n\n bool is_dirty() const {\n return bitset.test_any();\n }\n\n ailego::FixedBitset<32> bitset{};\n};\n\nstatic_assert(sizeof(DeletionMap) == 4, \"DeletionMap must be 4 bytes\");\n\n/*! Meta Information of Streamer Entity\n */\nstruct StreamerInvertedMeta {\n uint64_t create_time{0};\n uint64_t update_time{0};\n uint64_t revision_id{0};\n uint32_t segment_count{0};\n uint32_t segment_size{0};\n uint8_t reserved_[32];\n InvertedIndexHeader header;\n};\n\n/*! Location of Vector in Storage Segment\n */\nstruct VectorLocation {\n //! Constructor\n VectorLocation(void) {}\n\n //! Constructor\n VectorLocation(uint16_t id, bool col, uint32_t off)\n : segment_id(id), column_major(col), offset(off) {}\n\n uint16_t segment_id;\n uint16_t column_major : 1;\n uint16_t reserved_ : 15;\n uint32_t offset;\n\n public:\n bool operator==(const VectorLocation &other) const {\n return segment_id == other.segment_id &&\n column_major == other.column_major && offset == other.offset;\n }\n};\n\nstatic_assert(sizeof(VectorLocation) == sizeof(uint64_t),\n \"VectorLocation must be size of 8 bytes\");\n\nstruct KeyInfo {\n KeyInfo(void) {}\n KeyInfo(uint32_t idx, const VectorLocation &loc)\n : centroid_idx(idx), location(loc) {}\n KeyInfo(VectorLocation loc) : location(loc) {}\n uint32_t centroid_idx;\n VectorLocation location;\n};\n\n// Segments ID\nconst std::string IVF_CENTROID_SEG_ID(\"ivf.centroid\");\nconst std::string IVF_INVERTED_BODY_SEG_ID(\"ivf.inverted_body\");\nconst std::string IVF_INVERTED_HEADER_SEG_ID(\"ivf.inverted_header\");\nconst std::string IVF_INVERTED_META_SEG_ID(\"ivf.inverted_meta\");\nconst std::string IVF_KEYS_SEG_ID(\"hc.keys\");\nconst std::string IVF_OFFSETS_SEG_ID(\"ivf.offsets\");\nconst std::string IVF_MAPPING_SEG_ID(\"ivf.mapping\");\nconst std::string IVF_FEATURES_SEG_ID(\"ivf.features\");\nconst std::string IVF_INT8_QUANTIZED_PARAMS_SEG_ID(\"ivf.int8_quantized_params\");\nconst std::string IVF_INT4_QUANTIZED_PARAMS_SEG_ID(\"ivf.int4_quantized_params\");\n\nconst std::string IVF_INVERTED_LIST_HEAD_SEG_ID(\"ivf.inverted_list_head\");\nconst std::string IVF_STORAGE_SEGMENT_ID(\"ivf.S\");\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_index_provider.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"ivf_entity.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! IVF IndexProvider\n */\nclass IVFIndexProvider : public IndexProvider {\n public:\n IVFIndexProvider(const IndexMeta &meta, const IVFEntity::Pointer &entity,\n const std::string &owner)\n : meta_(meta), entity_(entity), owner_class_(owner) {}\n\n IVFIndexProvider(const IVFIndexProvider &) = delete;\n IVFIndexProvider &operator=(const IVFIndexProvider &) = delete;\n\n public:\n //! Create a new iterator\n virtual Iterator::Pointer create_iterator(void) override {\n return Iterator::Pointer(new (std::nothrow) Iterator(entity_));\n }\n\n //! Retrieve count of vectors\n virtual size_t count(void) const override {\n return entity_->vector_count();\n }\n\n //! Retrieve dimension of vector\n virtual size_t dimension(void) const override {\n return meta_.dimension();\n }\n\n //! Retrieve type of vector\n virtual IndexMeta::DataType data_type(void) const override {\n return meta_.data_type();\n }\n\n //! Retrieve vector size in bytes\n virtual size_t element_size(void) const override {\n return meta_.element_size();\n }\n\n //! Retrieve a vector using a primary key\n virtual const void *get_vector(uint64_t key) const override {\n return entity_->get_vector_by_key(key);\n }\n\n //! Retrieve the owner class\n virtual const std::string &owner_class(void) const override {\n return owner_class_;\n }\n\n private:\n class Iterator : public IndexProvider::Iterator {\n public:\n Iterator(const IVFEntity::Pointer &entity) : entity_(entity) {}\n\n //! Retrieve pointer of data\n //! NOTICE: the vec feature will be changed after iterating to next, so\n //! the caller need to keep a copy of it before iterator to next vector\n virtual const void *data(void) const override {\n return entity_->get_vector(index_);\n }\n\n //! Test if the iterator is valid\n virtual bool is_valid(void) const override {\n return index_ < entity_->vector_count();\n }\n\n //! Retrieve primary key\n virtual uint64_t key(void) const override {\n return entity_->get_key(index_);\n }\n\n //! Next iterator\n virtual void next(void) override {\n ++index_;\n }\n\n private:\n //! Members\n IVFEntity::Pointer entity_;\n size_t index_{0};\n };\n\n private:\n //! Members\n const IndexMeta &meta_;\n IVFEntity::Pointer entity_;\n std::string owner_class_;\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_params.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\nstatic const std::string SEPARATOR(\"/\");\nstatic const std::string CENTROID_SEPERATOR = \"*\";\n\n// builder params\nstatic const std::string PARAM_IVF_BUILDER_CENTROID_COUNT(\n \"proxima.ivf.builder.centroid_count\");\nstatic const std::string PARAM_IVF_BUILDER_CLUSTER_CLASS(\n \"proxima.ivf.builder.cluster_class\");\nstatic const std::string PARAM_IVF_BUILDER_THREAD_COUNT(\n \"proxima.ivf.builder.thread_count\");\nstatic const std::string PARAM_IVF_BUILDER_CLUSTER_AUTO_TUNING(\n \"proxima.ivf.builder.cluster_auto_tuning\");\nstatic const std::string PARAM_IVF_BUILDER_TRAIN_SAMPLE_COUNT(\n \"proxima.ivf.builder.train_sample_count\");\nstatic const std::string PARAM_IVF_BUILDER_TRAIN_SAMPLE_RATIO(\n \"proxima.ivf.builder.train_sample_ratio\");\nstatic const std::string PARAM_IVF_BUILDER_CONVERTER_PARAMS(\n \"proxima.ivf.builder.converter_params\");\nstatic const std::string PARAM_IVF_BUILDER_CONVERTER_CLASS(\n \"proxima.ivf.builder.converter_class\");\nstatic const std::string PARAM_IVF_BUILDER_STORE_ORIGINAL_FEATURES(\n \"proxima.ivf.builder.store_original_features\");\nstatic const std::string PARAM_IVF_BUILDER_QUANTIZER_CLASS(\n \"proxima.ivf.builder.quantizer_class\");\nstatic const std::string PARAM_IVF_BUILDER_QUANTIZE_BY_CENTROID(\n \"proxima.ivf.builder.quantize_by_centroid\");\nstatic const std::string PARAM_IVF_BUILDER_QUANTIZER_PARAMS(\n \"proxima.ivf.builder.quantizer_params\");\nstatic const std::string PARAM_IVF_BUILDER_CLUSTER_PARAMS_IN_LEVEL_PREFIX(\n \"proxima.ivf.builder.cluster_params_in_level_\");\nstatic const std::string PARAM_IVF_BUILDER_OPTIMIZER_CLASS(\n \"proxima.ivf.builder.optimizer_class\");\nstatic const std::string PARAM_IVF_BUILDER_OPTIMIZER_PARAMS(\n \"proxima.ivf.builder.optimizer_params\");\nstatic const std::string PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_CLASS(\n \"proxima.ivf.builder.optimizer_quantizer_class\");\nstatic const std::string PARAM_IVF_BUILDER_OPTIMIZER_QUANTIZER_PARAMS(\n \"proxima.ivf.builder.optimizer_quantizer_params\");\nstatic const std::string PARAM_IVF_BUILDER_BLOCK_VECTOR_COUNT(\n \"proxima.ivf.builder.block_vector_count\");\n\n// searcher params\nstatic const std::string PARAM_IVF_SEARCHER_SCAN_RATIO(\n \"proxima.ivf.searcher.scan_ratio\");\nstatic const std::string PARAM_IVF_SEARCHER_BRUTE_FORCE_THRESHOLD(\n \"proxima.ivf.searcher.brute_force_threshold\");\nstatic const std::string PARAM_IVF_SEARCHER_OPTIMIZER(\n \"proxima.ivf.searcher.optimizer\");\nstatic const std::string PARAM_IVF_SEARCHER_OPTIMIZER_PARAMS(\n \"proxima.ivf.searcher.optimizer_params\");\nstatic const std::string PARAM_IVF_SEARCHER_CONVERTER_REFORMER(\n \"proxima.ivf.searcher.converter_reformer\");\n\n// Constants\nstatic constexpr char const *kIPMetricName = \"InnerProduct\";\nstatic constexpr char const *kMipsMetricName = \"MipsSquaredEuclidean\";\nstatic constexpr char const *kL2MetricName = \"SquaredEuclidean\";\nstatic constexpr char const *kMipsConverterName = \"MipsConverter\";\nstatic constexpr char const *kMipsRevConverterName = \"MipsReverseConverter\";\nstatic constexpr char const *kMipsReformerName = \"MipsReformer\";\nstatic constexpr char const *kInt8QuantizerName = \"Int8QuantizerConverter\";\nstatic constexpr char const *kInt4QuantizerName = \"Int4QuantizerConverter\";\nstatic constexpr char const *kInt8ReformerName = \"Int8QuantizerReformer\";\nstatic constexpr char const *kInt4ReformerName = \"Int4QuantizerReformer\";\nstatic constexpr float kNormalizeScaleFactor = 16.0f;\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_searcher.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"ivf_searcher.h\"\n#include \n#include \n#include \"ivf_centroid_index.h\"\n#include \"ivf_index_provider.h\"\n#include \"ivf_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nint IVFSearcher::init(const ailego::Params ¶meters) {\n params_ = parameters;\n\n params_.get(PARAM_IVF_SEARCHER_BRUTE_FORCE_THRESHOLD, &bruteforce_threshold_);\n\n searcher_state_ = STATE_INITED;\n\n return 0;\n}\n\nint IVFSearcher::cleanup(void) {\n this->unload();\n\n params_.clear();\n bruteforce_threshold_ = kDefaultBfThreshold;\n\n searcher_state_ = STATE_INIT;\n return 0;\n}\n\nint IVFSearcher::load(IndexStorage::Pointer container,\n IndexMetric::Pointer /*metric*/) {\n if (!container) {\n LOG_ERROR(\"Invalid container\");\n return IndexError_InvalidArgument;\n }\n if (searcher_state_ != STATE_INITED) {\n LOG_ERROR(\"Initalize the searcher first before load index\");\n return IndexError_Runtime;\n }\n\n ailego::ElapsedTime timer;\n int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to deserialize meta from container\");\n return ret;\n }\n\n //! Load centroid index\n centroid_index_ = std::make_shared();\n if (!centroid_index_) {\n return IndexError_NoMemory;\n }\n auto seg = container->get(IVF_CENTROID_SEG_ID, 0);\n if (!seg) {\n LOG_ERROR(\"Failed to get segment %s\", IVF_CENTROID_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n IndexStorage::Pointer seg_container =\n std::make_shared(seg);\n if (!seg_container) {\n return IndexError_NoMemory;\n }\n ret = seg_container->open(std::string(), false);\n if (ret != 0) {\n LOG_ERROR(\"IndexSegmentStorage load failed for %s\", IndexError::What(ret));\n return ret;\n }\n ret = centroid_index_->load(seg_container, params_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to load index for %s\", IndexError::What(ret));\n return ret;\n }\n\n auto reformer = centroid_index_->reformer();\n params_.set(PARAM_IVF_SEARCHER_CONVERTER_REFORMER, reformer);\n\n //! load iverted index\n entity_ = std::make_shared();\n if (!entity_) {\n return IndexError_NoMemory;\n }\n ret = entity_->load(container);\n ivf_check_error_code(ret);\n\n magic_ = IndexContext::GenerateMagic();\n\n stats_.set_loaded_count(entity_->vector_count());\n stats_.set_loaded_costtime(timer.milli_seconds());\n\n searcher_state_ = STATE_LOADED;\n return 0;\n}\n\nint IVFSearcher::unload(void) {\n magic_ = 0;\n centroid_index_.reset();\n entity_.reset();\n stats_.set_loaded_count(0UL);\n stats_.set_loaded_costtime(0UL);\n stats_.clear_attributes();\n searcher_state_ = STATE_INITED;\n\n return 0;\n}\n\nint IVFSearcher::search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n return search_bf_impl(query, qmeta, 1, context);\n}\n\nint IVFSearcher::search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const {\n if (!query || qmeta.element_size() != meta_.element_size()) {\n LOG_ERROR(\"Null query or invalid qmeta\");\n return IndexError_InvalidArgument;\n }\n IVFSearcherContext *ctx = dynamic_cast(context.get());\n if (!ctx || ctx->topk() == 0) {\n LOG_ERROR(\"Invalid context or topk not set yet\");\n return IndexError_InvalidArgument;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher\n int ret = this->update_context(ctx);\n ivf_check_error_code(ret);\n }\n\n ctx->reset_results(count);\n auto &entity = ctx->entity();\n auto &filter = ctx->filter();\n\n //! Transform the querys for querying in inverted vector index later\n IndexQueryMeta iv_qmeta;\n int ret = entity->transform(query, qmeta, count, &query, &iv_qmeta);\n ivf_check_with_msg(ret, \"Failed to transform querys\");\n\n // TODO: do batch search in matrix\n for (size_t q = 0; q < count; ++q) {\n auto &context_stats = ctx->mutable_stats(q);\n auto &heap = ctx->mutable_result_heap();\n heap.clear();\n if (!filter.is_valid()) {\n ret = entity->search(query, &heap, &context_stats);\n } else {\n ret = entity->search(query, filter, &heap, &context_stats);\n }\n ivf_check_with_msg(ret, \"Failed to search in entity for %s\",\n IndexError::What(ret));\n heap.sort(); // sort the results\n if (!filter.is_valid()) {\n // mapping the local id to key if query without filter\n ret = entity->retrieve_keys(&heap);\n ivf_check_error_code(ret);\n }\n entity->normalize(q, &heap);\n ctx->topk_to_result(q);\n query = static_cast(query) + iv_qmeta.element_size();\n }\n\n return 0;\n}\n\nint IVFSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n return this->search_impl(query, qmeta, 1, context);\n}\n\nint IVFSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const {\n if (entity_->vector_count() <= bruteforce_threshold_) {\n return this->search_bf_impl(query, qmeta, count, context);\n }\n if (!query || qmeta.element_size() != meta_.element_size()) {\n LOG_ERROR(\"Null query or invalid qmeta\");\n return IndexError_InvalidArgument;\n }\n\n IVFSearcherContext *ctx = dynamic_cast(context.get());\n if (!ctx || ctx->topk() == 0) {\n LOG_ERROR(\"Invalid context or topk not set yet\");\n return IndexError_InvalidArgument;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher\n int ret = update_context(ctx);\n ivf_check_error_code(ret);\n }\n\n ctx->reset_results(count);\n auto &entity = ctx->entity();\n auto &filter = ctx->filter();\n\n auto ¢roid_index_ctx = ctx->centroid_searcher_ctx();\n int ret = centroid_index_->search(query, qmeta, count, centroid_index_ctx);\n ivf_check_error_code(ret);\n\n //! Transform the querys for querying in inverted vector index later\n IndexQueryMeta iv_qmeta;\n ret = entity->transform(query, qmeta, count, &query, &iv_qmeta);\n ivf_check_with_msg(ret, \"Failed to transform querys\");\n\n for (size_t q = 0; q < count; ++q) {\n auto ¢roids = centroid_index_ctx->result(q);\n auto &context_stats = ctx->mutable_stats(q);\n auto &heap = ctx->mutable_result_heap();\n heap.clear();\n uint32_t total_scan_count = 0;\n for (size_t i = 0;\n i < centroids.size() && total_scan_count < ctx->max_scan_count();\n ++i) {\n auto cid = centroids[i].key();\n uint32_t scan_count = 0;\n if (!filter.is_valid()) {\n ret = entity->search(cid, query, &scan_count, &heap, &context_stats);\n } else {\n ret = entity->search(cid, query, filter, &scan_count, &heap,\n &context_stats);\n }\n ivf_check_with_msg(ret, \"Failed to search in entity for %s\",\n IndexError::What(ret));\n total_scan_count += scan_count;\n }\n heap.sort(); // sort the results\n if (!filter.is_valid()) {\n // mapping the local id to key if query without filter\n ret = entity->retrieve_keys(&heap);\n ivf_check_error_code(ret);\n }\n entity->normalize(q, &heap);\n ctx->topk_to_result(q);\n query = static_cast(query) + iv_qmeta.element_size();\n }\n\n return 0;\n}\n\nconst IndexSearcher::Stats &IVFSearcher::stats(void) const {\n return stats_;\n}\n\nIndexSearcher::Context::Pointer IVFSearcher::create_context() const {\n if (searcher_state_ != STATE_LOADED) {\n LOG_ERROR(\"Load the index first before create context\");\n return nullptr;\n }\n\n auto entity = entity_->clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone IVFEntity\");\n return nullptr;\n }\n\n auto centroid_index_ctx = centroid_index_->create_context();\n if (!centroid_index_ctx) {\n LOG_ERROR(\"Failed to create centroid index context\");\n return nullptr;\n }\n\n auto context =\n new (std::nothrow) IVFSearcherContext(entity, centroid_index_ctx);\n if (!context) {\n LOG_ERROR(\"Failed to alloc IVFSearcherContext\");\n return nullptr;\n }\n int ret = context->init(params_);\n if (ret != 0) {\n delete context;\n return nullptr;\n }\n\n context->set_magic(magic_);\n\n return Context::Pointer(context);\n}\n\nIndexProvider::Pointer IVFSearcher::create_provider(void) const {\n if (searcher_state_ != STATE_LOADED) {\n LOG_ERROR(\"Load the index first before create provider\");\n return nullptr;\n }\n\n auto entity = entity_->clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone IVFEntity\");\n return Provider::Pointer();\n }\n\n auto *provider = new (std::nothrow)\n IVFIndexProvider(entity->has_orignal_feature() ? meta_ : entity->meta(),\n entity, \"IVFSearcher\");\n if (!provider) {\n LOG_ERROR(\"Failed to alloc IVFIndexProvider\");\n return Provider::Pointer();\n }\n\n return Provider::Pointer(provider);\n}\n\nINDEX_FACTORY_REGISTER_SEARCHER(IVFSearcher);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_searcher.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"ivf_centroid_index.h\"\n#include \"ivf_entity.h\"\n#include \"ivf_searcher_context.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! IVF Searcher\n */\nclass IVFSearcher : public IndexSearcher {\n public:\n //! Initialize Searcher\n virtual int init(const ailego::Params ¶meters) override;\n\n //! Cleanup Searcher\n virtual int cleanup(void) override;\n\n //! Load index from container\n virtual int load(IndexStorage::Pointer container,\n IndexMetric::Pointer metric) override;\n\n //! Unload index\n virtual int unload(void) override;\n\n //! Similarity brute force search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity brute force search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const override;\n\n //! Similarity search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const override;\n\n //! Retrieve statistics\n virtual const Stats &stats(void) const override;\n\n //! Create a searcher context\n virtual Context::Pointer create_context(void) const override;\n\n //! Create a new iterator\n virtual IndexProvider::Pointer create_provider(void) const override;\n\n //! Retrieve meta of index\n virtual const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n //! Retrieve params of index\n virtual const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n protected:\n int update_context(IVFSearcherContext *ctx) const {\n auto entity = entity_->clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone QcEntity\");\n return IndexError_Runtime;\n }\n\n //! The centroid index searcher may be different, so need to create one\n auto centroid_ctx = centroid_index_->create_context();\n if (!centroid_ctx) {\n LOG_ERROR(\"Failed to create centroid index searcher context\");\n return IndexError_Runtime;\n }\n\n return ctx->update_context(entity, centroid_ctx, params_, magic_);\n }\n\n private:\n //! Constants\n static constexpr uint32_t kDefaultBfThreshold = 1000u;\n\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 };\n\n //! Members\n IndexMeta meta_{};\n ailego::Params params_{};\n IVFCentroidIndex::Pointer centroid_index_{};\n IVFEntity::Pointer entity_{};\n uint32_t bruteforce_threshold_{kDefaultBfThreshold};\n uint32_t magic_{0};\n Stats stats_{};\n State searcher_state_{STATE_INIT};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_searcher_context.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \"ivf_entity.h\"\n#include \"ivf_utility.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! IVF Searcher Context\n */\nclass IVFSearcherContext : public IndexSearcher::Context {\n public:\n IVFSearcherContext(const IVFEntity::Pointer &ivf_entity,\n IndexSearcher::Context::Pointer ¢roid_ctx)\n : entity_(ivf_entity), centroid_searcher_ctx_(std::move(centroid_ctx)) {}\n\n public:\n //! Set topk of search result\n virtual void set_topk(uint32_t k) override {\n topk_ = k;\n result_heap_.limit(topk_);\n result_heap_.set_threshold(this->threshold());\n }\n\n //! Retrieve search result\n virtual const IndexDocumentList &result(void) const override {\n return results_[0];\n }\n\n //! Retrieve search result with index\n virtual const IndexDocumentList &result(size_t idx) const override {\n ailego_assert_with(results_.size() > idx, \"invalid index\");\n return results_[idx];\n }\n\n //! Retrieve mutable result with index\n virtual IndexDocumentList *mutable_result(size_t idx) override {\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n return &results_[idx];\n }\n\n inline IndexDocumentHeap *result_heap() {\n return &result_heap_;\n }\n\n //! Update the parameters of context\n virtual int update(const ailego::Params ¶ms) override {\n params.get(PARAM_IVF_SEARCHER_BRUTE_FORCE_THRESHOLD,\n &bruteforce_threshold_);\n params.get(PARAM_IVF_SEARCHER_SCAN_RATIO, &scan_ratio_);\n if (scan_ratio_ <= 0.0) {\n LOG_ERROR(\"Invalid params %s=%f\", PARAM_IVF_SEARCHER_SCAN_RATIO.c_str(),\n scan_ratio_);\n return IndexError_InvalidArgument;\n }\n size_t topk_val =\n std::max(static_cast(\n std::round(entity_->inverted_list_count() * scan_ratio_)),\n 1u);\n centroid_searcher_ctx_->set_topk(topk_val);\n max_scan_count_ =\n static_cast(std::ceil(entity_->vector_count() * scan_ratio_));\n max_scan_count_ = std::max(bruteforce_threshold_, max_scan_count_);\n return 0;\n }\n\n //! Retrieve magic number\n virtual uint32_t magic(void) const override {\n return magic_;\n }\n\n public:\n //! Initialize the context\n int init(const ailego::Params ¶ms) {\n return this->update(params);\n }\n\n //! Update the magic number\n void set_magic(uint32_t mag) {\n magic_ = mag;\n }\n\n //! Get Topk Value\n uint32_t topk() const override {\n return topk_;\n }\n\n //! Retrieve scan ratio\n float scan_ratio(void) const {\n return scan_ratio_;\n }\n\n //! Retrieve max scan count\n uint32_t max_scan_count(void) const {\n return max_scan_count_;\n }\n\n uint32_t bruteforce_threshold() const {\n return bruteforce_threshold_;\n }\n\n //! Retrieve magic number\n const IVFEntity::Pointer &entity() const {\n return entity_;\n }\n\n //! Retrieve Mutable Query Result By Query Index\n IndexDocumentHeap &mutable_result_heap() {\n return result_heap_;\n }\n\n void set_fetch_vector(bool v) override {\n fetch_vector_ = v;\n }\n\n bool fetch_vector(void) const override {\n return fetch_vector_;\n }\n\n //! Reset all the query results\n void reset_results(size_t qnum) {\n results_.resize(qnum);\n stats_vec_.resize(qnum);\n for (size_t i = 0; i < qnum; ++i) {\n results_[i].clear();\n stats_vec_[i].clear();\n }\n result_heap_.clear();\n result_heap_.limit(topk_);\n result_heap_.set_threshold(this->threshold());\n }\n\n //! Update context, the context may be shared by different searcher\n int update_context(IVFEntity::Pointer &new_entity,\n IndexSearcher::Context::Pointer ¢roid_ctx,\n const ailego::Params ¶ms, uint32_t magic_num) {\n entity_ = new_entity;\n centroid_searcher_ctx_ = std::move(centroid_ctx);\n int ret = this->update(params);\n ivf_check_error_code(ret);\n\n magic_ = magic_num;\n\n return 0;\n }\n\n //! Retrieve the centroid index context\n IndexSearcher::Context::Pointer ¢roid_searcher_ctx(void) {\n return centroid_searcher_ctx_;\n }\n\n const Stats &stats(size_t idx = 0) const {\n ailego_assert_with(stats_vec_.size() > idx, \"invalid index\");\n return stats_vec_[idx];\n }\n\n Stats &mutable_stats(size_t idx = 0) {\n ailego_assert_with(stats_vec_.size() > idx, \"invalid index\");\n return stats_vec_[idx];\n }\n\n void topk_to_result(uint32_t idx) {\n if (ailego_unlikely(result_heap_.size() == 0)) {\n return;\n }\n\n ailego_assert_with(idx < results_.size(), \"invalid idx\");\n int size = std::min(topk_, static_cast(result_heap_.size()));\n result_heap_.sort();\n results_[idx].clear();\n for (int i = 0; i < size; ++i) {\n auto score = result_heap_[i].score();\n if (score > this->threshold()) {\n break;\n }\n\n key_t key = result_heap_[i].key();\n if (fetch_vector_) {\n IndexStorage::MemoryBlock block;\n entity_->get_vector_by_key(key, block);\n results_[idx].emplace_back(key, score, key, block);\n } else {\n results_[idx].emplace_back(key, score);\n }\n }\n }\n\n private:\n //! Constants\n static constexpr float kDefaultScanRatio = 0.1f;\n static constexpr uint32_t kDefaultBfThreshold = 1000u;\n\n //! Members\n IVFEntity::Pointer entity_{};\n IndexSearcher::Context::Pointer centroid_searcher_ctx_{};\n IndexDocumentHeap result_heap_;\n std::vector results_{};\n std::vector stats_vec_{};\n\n bool fetch_vector_{false};\n uint32_t topk_{0};\n uint32_t magic_{0};\n float scan_ratio_{kDefaultScanRatio};\n uint32_t max_scan_count_{0};\n uint32_t bruteforce_threshold_{kDefaultBfThreshold};\n};\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_streamer.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \"ivf_streamer.h\"\n#include \n#include \n#include \"ivf_centroid_index.h\"\n#include \"ivf_index_provider.h\"\n#include \"ivf_params.h\"\n\nnamespace zvec {\nnamespace core {\n\nint IVFStreamer::init(const IndexMeta &meta, const ailego::Params ¶meters) {\n meta_ = meta;\n params_ = parameters;\n\n params_.get(PARAM_IVF_SEARCHER_BRUTE_FORCE_THRESHOLD, &bruteforce_threshold_);\n\n searcher_state_ = STATE_INITED;\n\n return 0;\n}\n\nint IVFStreamer::cleanup(void) {\n this->unload();\n\n params_.clear();\n bruteforce_threshold_ = kDefaultBfThreshold;\n\n searcher_state_ = STATE_INIT;\n return 0;\n}\n\nint IVFStreamer::open(IndexStorage::Pointer storage) {\n if (!storage) {\n LOG_ERROR(\"Invalid storage\");\n return IndexError_InvalidArgument;\n }\n if (searcher_state_ != STATE_INITED) {\n LOG_ERROR(\"Initalize the searcher first before load index\");\n return IndexError_Runtime;\n }\n\n ailego::ElapsedTime timer;\n int ret = IndexHelper::DeserializeFromStorage(storage.get(), &meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to deserialize meta from storage\");\n return ret;\n }\n\n //! Load centroid index\n centroid_index_ = std::make_shared();\n if (!centroid_index_) {\n return IndexError_NoMemory;\n }\n auto seg = storage->get(IVF_CENTROID_SEG_ID, 0);\n if (!seg) {\n LOG_ERROR(\"Failed to get segment %s\", IVF_CENTROID_SEG_ID.c_str());\n return IndexError_InvalidFormat;\n }\n IndexStorage::Pointer seg_container =\n std::make_shared(seg);\n if (!seg_container) {\n return IndexError_NoMemory;\n }\n ret = seg_container->open(std::string(), false);\n if (ret != 0) {\n LOG_ERROR(\"IndexSegmentStorage load failed for %s\", IndexError::What(ret));\n return ret;\n }\n ret = centroid_index_->load(seg_container, params_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to load index for %s\", IndexError::What(ret));\n return ret;\n }\n\n auto reformer = centroid_index_->reformer();\n params_.set(PARAM_IVF_SEARCHER_CONVERTER_REFORMER, reformer);\n\n //! load iverted index\n entity_ = std::make_shared();\n if (!entity_) {\n return IndexError_NoMemory;\n }\n ret = entity_->load(storage);\n ivf_check_error_code(ret);\n\n magic_ = IndexContext::GenerateMagic();\n\n stats_.set_loaded_count(entity_->vector_count());\n stats_.set_loaded_costtime(timer.milli_seconds());\n\n searcher_state_ = STATE_LOADED;\n return 0;\n}\n\nint IVFStreamer::unload(void) {\n magic_ = 0;\n centroid_index_.reset();\n entity_.reset();\n stats_.set_loaded_count(0UL);\n stats_.set_loaded_costtime(0UL);\n stats_.clear_attributes();\n searcher_state_ = STATE_INITED;\n\n return 0;\n}\n\nint IVFStreamer::search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n return search_bf_impl(query, qmeta, 1, context);\n}\n\nint IVFStreamer::search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const {\n if (!query || qmeta.element_size() != meta_.element_size()) {\n LOG_ERROR(\"Null query or invalid qmeta\");\n return IndexError_InvalidArgument;\n }\n IVFSearcherContext *ctx = dynamic_cast(context.get());\n if (!ctx || ctx->topk() == 0) {\n LOG_ERROR(\"Invalid context or topk not set yet\");\n return IndexError_InvalidArgument;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher\n int ret = this->update_context(ctx);\n ivf_check_error_code(ret);\n }\n\n ctx->reset_results(count);\n auto &entity = ctx->entity();\n auto &filter = ctx->filter();\n\n //! Transform the querys for querying in inverted vector index later\n IndexQueryMeta iv_qmeta;\n int ret = entity->transform(query, qmeta, count, &query, &iv_qmeta);\n ivf_check_with_msg(ret, \"Failed to transform querys\");\n\n // TODO: do batch search in matrix\n for (size_t q = 0; q < count; ++q) {\n auto &context_stats = ctx->mutable_stats(q);\n auto &heap = ctx->mutable_result_heap();\n heap.clear();\n if (!filter.is_valid()) {\n ret = entity->search(query, &heap, &context_stats);\n } else {\n ret = entity->search(query, filter, &heap, &context_stats);\n }\n ivf_check_with_msg(ret, \"Failed to search in entity for %s\",\n IndexError::What(ret));\n heap.sort(); // sort the results\n if (!filter.is_valid()) {\n // mapping the local id to key if query without filter\n ret = entity->retrieve_keys(&heap);\n ivf_check_error_code(ret);\n }\n entity->normalize(q, &heap);\n ctx->topk_to_result(q);\n query = static_cast(query) + iv_qmeta.element_size();\n }\n\n return 0;\n}\n\nint IVFStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n return this->search_impl(query, qmeta, 1, context);\n}\n\nint IVFStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const {\n if (entity_->vector_count() <= bruteforce_threshold_) {\n return this->search_bf_impl(query, qmeta, count, context);\n }\n if (!query || qmeta.element_size() != meta_.element_size()) {\n LOG_ERROR(\"Null query or invalid qmeta\");\n return IndexError_InvalidArgument;\n }\n\n IVFSearcherContext *ctx = dynamic_cast(context.get());\n if (!ctx || ctx->topk() == 0) {\n LOG_ERROR(\"Invalid context or topk not set yet\");\n return IndexError_InvalidArgument;\n }\n if (ctx->magic() != magic_) {\n //! context is created by another searcher\n int ret = update_context(ctx);\n ivf_check_error_code(ret);\n }\n\n ctx->reset_results(count);\n auto &entity = ctx->entity();\n auto &filter = ctx->filter();\n\n auto ¢roid_index_ctx = ctx->centroid_searcher_ctx();\n int ret = centroid_index_->search(query, qmeta, count, centroid_index_ctx);\n ivf_check_error_code(ret);\n\n //! Transform the querys for querying in inverted vector index later\n IndexQueryMeta iv_qmeta;\n ret = entity->transform(query, qmeta, count, &query, &iv_qmeta);\n ivf_check_with_msg(ret, \"Failed to transform querys\");\n\n for (size_t q = 0; q < count; ++q) {\n auto ¢roids = centroid_index_ctx->result(q);\n auto &context_stats = ctx->mutable_stats(q);\n auto &heap = ctx->mutable_result_heap();\n heap.clear();\n uint32_t total_scan_count = 0;\n for (size_t i = 0;\n i < centroids.size() && total_scan_count < ctx->max_scan_count();\n ++i) {\n auto cid = centroids[i].key();\n uint32_t scan_count = 0;\n if (!filter.is_valid()) {\n ret = entity->search(cid, query, &scan_count, &heap, &context_stats);\n } else {\n ret = entity->search(cid, query, filter, &scan_count, &heap,\n &context_stats);\n }\n ivf_check_with_msg(ret, \"Failed to search in entity for %s\",\n IndexError::What(ret));\n total_scan_count += scan_count;\n }\n heap.sort(); // sort the results\n if (!filter.is_valid()) {\n // mapping the local id to key if query without filter\n ret = entity->retrieve_keys(&heap);\n ivf_check_error_code(ret);\n }\n entity->normalize(q, &heap);\n ctx->topk_to_result(q);\n query = static_cast(query) + iv_qmeta.element_size();\n }\n\n return 0;\n}\n\nconst IndexSearcher::Stats &IVFStreamer::stats(void) const {\n return stats_;\n}\n\nIndexSearcher::Context::Pointer IVFStreamer::create_context() const {\n if (searcher_state_ != STATE_LOADED) {\n LOG_ERROR(\"Load the index first before create context\");\n return nullptr;\n }\n\n auto entity = entity_->clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone IVFEntity\");\n return nullptr;\n }\n\n auto centroid_index_ctx = centroid_index_->create_context();\n if (!centroid_index_ctx) {\n LOG_ERROR(\"Failed to create centroid index context\");\n return nullptr;\n }\n\n auto context =\n new (std::nothrow) IVFSearcherContext(entity, centroid_index_ctx);\n if (!context) {\n LOG_ERROR(\"Failed to alloc IVFSearcherContext\");\n return nullptr;\n }\n int ret = context->init(params_);\n if (ret != 0) {\n delete context;\n return nullptr;\n }\n\n context->set_magic(magic_);\n\n return Context::Pointer(context);\n}\n\nIndexProvider::Pointer IVFStreamer::create_provider(void) const {\n if (searcher_state_ != STATE_LOADED) {\n LOG_ERROR(\"Load the index first before create provider\");\n return nullptr;\n }\n\n auto entity = entity_->clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone IVFEntity\");\n return Provider::Pointer();\n }\n\n auto *provider = new (std::nothrow)\n IVFIndexProvider(entity->has_orignal_feature() ? meta_ : entity->meta(),\n entity, \"IVFStreamer\");\n if (!provider) {\n LOG_ERROR(\"Failed to alloc IVFIndexProvider\");\n return Provider::Pointer();\n }\n\n return Provider::Pointer(provider);\n}\n\nINDEX_FACTORY_REGISTER_STREAMER(IVFStreamer);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/algorithm/ivf/ivf_streamer.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#ifndef __IVF_STREAMER_H__\n#define __IVF_STREAMER_H__\n\n#include \n#include \"ivf_centroid_index.h\"\n#include \"ivf_entity.h\"\n#include \"ivf_searcher_context.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! IVF Searcher\n */\nclass IVFStreamer : public IndexStreamer {\n public:\n //! Initialize Searcher\n virtual int init(const IndexMeta & /*meta*/,\n const ailego::Params & /*params*/) override;\n\n //! Cleanup Searcher\n virtual int cleanup(void) override;\n\n //! Load index from container\n virtual int open(IndexStorage::Pointer storage) override;\n\n virtual int flush(uint64_t /*check_point*/) override {\n return 0;\n }\n virtual int close(void) override {\n return this->unload();\n }\n\n //! Unload index\n virtual int unload(void) override;\n\n //! Similarity brute force search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity brute force search\n virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const override;\n\n //! Similarity search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const override;\n\n //! Similarity search\n virtual int search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count,\n Context::Pointer &context) const override;\n\n //! Retrieve statistics\n virtual const Stats &stats(void) const override;\n\n //! Create a searcher context\n virtual Context::Pointer create_context(void) const override;\n\n //! Create a new iterator\n virtual IndexProvider::Pointer create_provider(void) const override;\n\n //! Retrieve meta of index\n virtual const IndexMeta &meta(void) const override {\n return meta_;\n }\n\n virtual int get_vector_by_id(\n const uint32_t id, IndexStorage::MemoryBlock &block) const override {\n return entity_->get_vector_by_key(id, block);\n }\n\n protected:\n int update_context(IVFSearcherContext *ctx) const {\n auto entity = entity_->clone();\n if (!entity) {\n LOG_ERROR(\"Failed to clone QcEntity\");\n return IndexError_Runtime;\n }\n\n //! The centroid index searcher may be different, so need to create one\n auto centroid_ctx = centroid_index_->create_context();\n if (!centroid_ctx) {\n LOG_ERROR(\"Failed to create centroid index searcher context\");\n return IndexError_Runtime;\n }\n\n return ctx->update_context(entity, centroid_ctx, params_, magic_);\n }\n\n private:\n //! Constants\n static constexpr uint32_t kDefaultBfThreshold = 1000u;\n\n enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 };\n\n //! Members\n IndexMeta meta_{};\n ailego::Params params_{};\n IndexBuilder::Pointer builder_;\n IVFCentroidIndex::Pointer centroid_index_{};\n IVFEntity::Pointer entity_{};\n uint32_t bruteforce_threshold_{kDefaultBfThreshold};\n uint32_t magic_{0};\n Stats stats_{};\n State searcher_state_{STATE_INIT};\n};\n\n} // namespace core\n} // namespace zvec\n\n#endif //__IVF_STREAMER_H__\n" }, { "path": "src/core/algorithm/ivf/ivf_utility.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\n#ifndef ivf_check_error_code\n#define ivf_check_error_code(code) \\\n if (ailego_unlikely((code) != 0)) return code\n#endif\n\n#ifndef ivf_assert\n#define ivf_assert(cond, code) \\\n if (ailego_unlikely(!(cond))) return code\n#endif\n\n#ifndef ivf_check_with_msg\n#define ivf_check_with_msg(code, fmt, args...) \\\n do { \\\n if (ailego_unlikely((code) != 0)) { \\\n LOG_ERROR(fmt, ##args); \\\n return code; \\\n } \\\n } while (0)\n#endif\n\n#ifndef ivf_assert_with_msg\n#define ivf_assert_with_msg(cond, err, fmt, args...) \\\n do { \\\n if (ailego_unlikely(!(cond))) { \\\n LOG_ERROR(fmt, ##args); \\\n return err; \\\n } \\\n } while (0)\n#endif\n\n/*! Quantized Clustering Utility\n */\nclass IVFUtility {\n public:\n //! Generator a random path with specificed prefix\n static inline std::string GenerateRandomPath(const std::string &prefix) {\n uint64_t timestamp = ailego::Monotime::MicroSeconds();\n return prefix + std::to_string(timestamp);\n }\n\n //! Compute the default scan ratio for total vectors\n static inline float ComputeScanRatio(size_t vector_count) {\n // the fitting function for the follow points: 1000000(0.02)\n // 10000000(0.01) 50000000(0.005) 100000000(0.001)\n float scan_ratio = -0.004 * std::log(vector_count) + 0.0751;\n scan_ratio = std::max(scan_ratio, 0.0001f);\n return scan_ratio;\n }\n\n //! Transpose the vectors in row major order to column major order\n static inline void Transpose(size_t align_size, const void *src, size_t m,\n size_t dim, void *dst);\n\n //! Transpose the vectors in column major order to row major order\n static inline void ReverseTranspose(size_t align_size, const void *src,\n size_t m, size_t dim, void *dst);\n\n //! Aligned size of a block vectors buffer\n static inline size_t AlignedSize(size_t fnum, size_t element_size);\n\n //! Aligned size of one vector buffer\n static inline size_t AlignedSize(size_t element_size);\n\n //! Sort arr with size in ascending order, and keep the index postion\n //! n2o keep the mapping: new position => origin postion\n //! For example, the input arr = [5, 3, 9, 6, 7], size = 5, after sort\n // arr = [3, 5, 6, 7, 9]\n // n2o = [1, 0, 3, 4, 2]\n //! To save memory, no extra memory is allocated\n template \n static void Sort(T *arr, std::vector *n2o, size_t size) {\n std::vector o2n;\n o2n.resize(size);\n n2o->resize(size);\n\n std::iota(n2o->begin(), n2o->end(), 0U);\n std::sort(n2o->begin(), n2o->end(),\n [&](I i, I j) { return arr[i] < arr[j]; });\n for (I i = 0U; i < size; ++i) {\n o2n[(*n2o)[i]] = i;\n }\n //! reorder arr in place, according to given n2o index\n for (I i = 0; i < size; ++i) {\n if (i != (*n2o)[i]) {\n T tmp = arr[i];\n I j = i, k;\n while (i != (k = (*n2o)[j])) {\n arr[j] = arr[k];\n (*n2o)[j] = j;\n j = k;\n }\n arr[j] = tmp;\n (*n2o)[j] = j;\n }\n }\n\n for (I i = 0U; i < size; ++i) {\n (*n2o)[o2n[i]] = i;\n }\n }\n\n //! Transpose one vector in block\n template \n static inline void TransposeOne(const void *src, size_t M, size_t N,\n void *dst) {\n for (size_t i = 0; i < N; ++i) {\n reinterpret_cast(dst)[i] = reinterpret_cast(src)[i * M];\n }\n }\n};\n\nvoid IVFUtility::Transpose(size_t align_size, const void *src, size_t m,\n size_t dim, void *dst) {\n switch (align_size) {\n case 2:\n ailego::MatrixHelper::Transpose(src, m, dim, dst);\n break;\n case 4:\n ailego::MatrixHelper::Transpose(src, m, dim, dst);\n break;\n case 8:\n ailego::MatrixHelper::Transpose(src, m, dim, dst);\n break;\n }\n}\n\nvoid IVFUtility::ReverseTranspose(size_t align_size, const void *src, size_t m,\n size_t dim, void *dst) {\n switch (align_size) {\n case 2:\n ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst);\n break;\n case 4:\n ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst);\n break;\n case 8:\n ailego::MatrixHelper::ReverseTranspose(src, m, dim, dst);\n break;\n }\n}\n\nsize_t IVFUtility::AlignedSize(size_t fnum, size_t element_size) {\n return ailego_align(fnum * element_size, 32);\n}\n\nsize_t IVFUtility::AlignedSize(size_t element_size) {\n return ailego_align(element_size, 32);\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\ncc_library(\n NAME core_framework \n STATIC STRICT ALWAYS_LINK\n SRCS *.cc\n LIBS zvec_ailego\n INCS . ${PROJECT_ROOT_DIR}/src/core\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )\n" }, { "path": "src/core/framework/index_cluster.cc", "content": "// namespace aitheta2\n// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nstatic const std::string CLUSTER_CENTROIDS_FEATURES_NAME =\n \"IndexCluster.Centroids.Features\";\nstatic const std::string CLUSTER_CENTROIDS_INDEXES_NAME =\n \"IndexCluster.Centroids.Indexes\";\n\n/*! Item Centroid Format\n */\nstruct ItemCentroidFormat {\n uint32_t parent;\n uint32_t reserved0_;\n uint64_t follows;\n double score;\n uint64_t reserved1_;\n};\n\nstatic inline bool GatherSubitemsCount(const ItemCentroidFormat *format,\n size_t count,\n std::vector *out) {\n out->resize(count + 1);\n\n for (const ItemCentroidFormat *it = format, *end = format + count; it != end;\n ++it) {\n uint32_t parent = it->parent + 1;\n if (parent > count) {\n return false;\n }\n (*out)[parent] += 1;\n }\n return (out->front() != 0);\n}\n\nint IndexCluster::Deserialize(const IndexMeta &meta,\n IndexBundle::Pointer bundle,\n CentroidList *cents) {\n if (!bundle || !cents) {\n return IndexError_InvalidArgument;\n }\n\n ailego::BlobWrap features = bundle->get(CLUSTER_CENTROIDS_FEATURES_NAME);\n ailego::BlobWrap indexes = bundle->get(CLUSTER_CENTROIDS_INDEXES_NAME);\n\n if (!features.is_valid() || !indexes.is_valid()) {\n return IndexError_InvalidArgument;\n }\n\n if (features.size() % meta.element_size() != 0 ||\n indexes.size() % sizeof(ItemCentroidFormat) != 0) {\n return IndexError_InvalidLength;\n }\n\n size_t count = features.size() / meta.element_size();\n if (indexes.size() / sizeof(ItemCentroidFormat) != count) {\n return IndexError_InvalidLength;\n }\n\n const ItemCentroidFormat *format =\n reinterpret_cast(indexes.buffer());\n std::vector subitems;\n\n if (!GatherSubitemsCount(format, count, &subitems)) {\n return IndexError_InvalidFormat;\n }\n\n std::vector items;\n items.reserve(count);\n cents->clear();\n cents->reserve(subitems.front());\n\n const uint8_t *feat = reinterpret_cast(features.buffer());\n size_t feat_size = meta.element_size();\n\n for (size_t i = 0; i < count; ++i, ++format, feat += feat_size) {\n CentroidList *current = cents;\n\n if (format->parent != static_cast(-1)) {\n if (format->parent >= items.size()) {\n return IndexError_InvalidFormat;\n }\n current = items[format->parent]->mutable_subitems();\n }\n current->emplace_back(feat, feat_size);\n\n // Update information\n Centroid *last_one = &(current->back());\n last_one->set_follows(static_cast(format->follows));\n last_one->set_score(format->score);\n last_one->mutable_subitems()->reserve(subitems[i + 1]);\n items.push_back(last_one);\n }\n return 0;\n}\n\nstatic void SerializeToBuffers(const IndexCluster::CentroidList ¢s,\n std::string *features, std::string *indexes) {\n uint32_t parent =\n static_cast(indexes->size() / sizeof(ItemCentroidFormat)) - 1;\n\n for (const auto &it : cents) {\n ItemCentroidFormat format{parent, 0, it.follows(), it.score(), 0};\n indexes->append(reinterpret_cast(&format), sizeof(format));\n features->append(reinterpret_cast(it.feature()), it.size());\n\n if (!it.subitems().empty()) {\n SerializeToBuffers(it.subitems(), features, indexes);\n }\n }\n}\n\nint IndexCluster::Serialize(const IndexMeta &meta, const CentroidList ¢s,\n IndexBundle::Pointer *out) {\n size_t cents_total = cents.size();\n\n // Check the centroids\n for (const auto &it : cents) {\n if (!it.is_matched(meta)) {\n return IndexError_Mismatch;\n }\n cents_total += it.subcount();\n }\n\n std::string features, indexes;\n features.reserve(cents_total * meta.element_size());\n indexes.reserve(cents_total * sizeof(ItemCentroidFormat));\n SerializeToBuffers(cents, &features, &indexes);\n\n std::shared_ptr bundle =\n std::make_shared();\n\n bundle->set(CLUSTER_CENTROIDS_FEATURES_NAME, std::move(features));\n bundle->set(CLUSTER_CENTROIDS_INDEXES_NAME, std::move(indexes));\n *out = std::move(bundle);\n\n return 0;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_context.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nuint32_t IndexContext::GenerateMagic(void) {\n static std::atomic_uint32_t magic_number{std::random_device()()};\n return magic_number.fetch_add(1);\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_converter.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nint IndexConverter::TrainAndTransform(const IndexConverter::Pointer &converter,\n IndexHolder::Pointer holder) {\n auto two_pass_holder = IndexHelper::MakeTwoPassHolder(std::move(holder));\n int ret = converter->train(two_pass_holder);\n if (ret == 0) {\n ret = converter->transform(std::move(two_pass_holder));\n }\n return ret;\n}\n\nint IndexConverter::TrainTransformAndDump(\n const IndexConverter::Pointer &converter, IndexHolder::Pointer holder,\n const IndexDumper::Pointer &dumper) {\n int ret = IndexConverter::TrainAndTransform(converter, std::move(holder));\n if (ret == 0) {\n ret = converter->dump(dumper);\n }\n return ret;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_error.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n\nnamespace zvec {\nnamespace core {\n\nINDEX_ERROR_CODE_DEFINE(Success, 0, \"Success\");\nINDEX_ERROR_CODE_DEFINE(Runtime, 1, \"Runtime error\");\nINDEX_ERROR_CODE_DEFINE(Logic, 2, \"Logic error\");\nINDEX_ERROR_CODE_DEFINE(Type, 3, \"Type error\");\nINDEX_ERROR_CODE_DEFINE(System, 4, \"System call error\");\nINDEX_ERROR_CODE_DEFINE(Cast, 5, \"Cast error\");\nINDEX_ERROR_CODE_DEFINE(IO, 6, \"IO error\");\nINDEX_ERROR_CODE_DEFINE(AuthExpired, 7, \"Auth expired error\");\n\nINDEX_ERROR_CODE_DEFINE(NotImplemented, 11, \"Not implemented\");\nINDEX_ERROR_CODE_DEFINE(Unsupported, 12, \"Unsupported\");\nINDEX_ERROR_CODE_DEFINE(Denied, 13, \"Permission denied\");\nINDEX_ERROR_CODE_DEFINE(Canceled, 14, \"Operation canceled\");\nINDEX_ERROR_CODE_DEFINE(Overflow, 15, \"Overflow\");\nINDEX_ERROR_CODE_DEFINE(Underflow, 16, \"Underflow\");\nINDEX_ERROR_CODE_DEFINE(OutOfRange, 17, \"Out of range\");\nINDEX_ERROR_CODE_DEFINE(NoBuffer, 18, \"No buffer space available\");\nINDEX_ERROR_CODE_DEFINE(NoMemory, 19, \"Not enough space\");\nINDEX_ERROR_CODE_DEFINE(NoParamFound, 20, \"No parameter found\");\nINDEX_ERROR_CODE_DEFINE(NoReady, 21, \"No ready\");\nINDEX_ERROR_CODE_DEFINE(NoExist, 22, \"No exist\");\nINDEX_ERROR_CODE_DEFINE(Exist, 23, \"Already exist\");\nINDEX_ERROR_CODE_DEFINE(Mismatch, 24, \"Mismatch\");\nINDEX_ERROR_CODE_DEFINE(Duplicate, 25, \"Duplicate\");\nINDEX_ERROR_CODE_DEFINE(Uninitialized, 26, \"Uninitialized\");\n\nINDEX_ERROR_CODE_DEFINE(InvalidArgument, 31, \"Invalid argument\");\nINDEX_ERROR_CODE_DEFINE(InvalidFormat, 32, \"Invalid format\");\nINDEX_ERROR_CODE_DEFINE(InvalidLength, 33, \"Invalid length\");\nINDEX_ERROR_CODE_DEFINE(InvalidChecksum, 34, \"Invalid checksum\");\nINDEX_ERROR_CODE_DEFINE(InvalidValue, 35, \"Invalid value\");\n\nINDEX_ERROR_CODE_DEFINE(CreateDirectory, 101, \"Create directory error\");\nINDEX_ERROR_CODE_DEFINE(OpenDirectory, 102, \"Open directory error\");\nINDEX_ERROR_CODE_DEFINE(Serialize, 105, \"Serialize error\");\nINDEX_ERROR_CODE_DEFINE(Deserialize, 106, \"Deserialize error\");\nINDEX_ERROR_CODE_DEFINE(CreateFile, 111, \"Create file error\");\nINDEX_ERROR_CODE_DEFINE(OpenFile, 112, \"Open file error\");\nINDEX_ERROR_CODE_DEFINE(SeekFile, 113, \"Seek file error\");\nINDEX_ERROR_CODE_DEFINE(CloseFile, 114, \"Close file error\");\nINDEX_ERROR_CODE_DEFINE(TruncateFile, 115, \"TruncateFile file error\");\nINDEX_ERROR_CODE_DEFINE(MMapFile, 116, \"MMap file error\");\nINDEX_ERROR_CODE_DEFINE(FlushFile, 117, \"Flush file error\");\nINDEX_ERROR_CODE_DEFINE(WriteData, 121, \"Write data error\");\nINDEX_ERROR_CODE_DEFINE(ReadData, 122, \"Read data error\");\n\nINDEX_ERROR_CODE_DEFINE(PackIndex, 201, \"Read data error\");\nINDEX_ERROR_CODE_DEFINE(UnpackIndex, 202, \"Read data error\");\nINDEX_ERROR_CODE_DEFINE(IndexLoaded, 203, \"Index loaded\");\nINDEX_ERROR_CODE_DEFINE(NoIndexLoaded, 204, \"No index loaded\");\nINDEX_ERROR_CODE_DEFINE(NoTrained, 205, \"No trained\");\nINDEX_ERROR_CODE_DEFINE(IndexFull, 206, \"Index full\");\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_factory.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n\nnamespace zvec {\nnamespace core {\n\nIndexMetric::Pointer IndexFactory::CreateMetric(const std::string &name) {\n IndexMetric::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasMetric(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllMetrics(void) {\n return ailego::Factory::Classes();\n}\n\nIndexLogger::Pointer IndexFactory::CreateLogger(const std::string &name) {\n IndexLogger::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasLogger(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllLoggers(void) {\n return ailego::Factory::Classes();\n}\n\nIndexDumper::Pointer IndexFactory::CreateDumper(const std::string &name) {\n IndexDumper::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasDumper(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllDumpers(void) {\n return ailego::Factory::Classes();\n}\n\nIndexStorage::Pointer IndexFactory::CreateStorage(const std::string &name) {\n IndexStorage::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasStorage(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllStorages(void) {\n return ailego::Factory::Classes();\n}\n\nIndexConverter::Pointer IndexFactory::CreateConverter(const std::string &name) {\n IndexConverter::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasConverter(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllConverters(void) {\n return ailego::Factory::Classes();\n}\n\nIndexReformer::Pointer IndexFactory::CreateReformer(const std::string &name) {\n IndexReformer::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasReformer(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllReformers(void) {\n return ailego::Factory::Classes();\n}\n\nIndexTrainer::Pointer IndexFactory::CreateTrainer(const std::string &name) {\n IndexTrainer::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasTrainer(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllTrainers(void) {\n return ailego::Factory::Classes();\n}\n\nIndexBuilder::Pointer IndexFactory::CreateBuilder(const std::string &name) {\n IndexBuilder::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasBuilder(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllBuilders(void) {\n return ailego::Factory::Classes();\n}\n\nIndexSearcher::Pointer IndexFactory::CreateSearcher(const std::string &name) {\n IndexSearcher::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasSearcher(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllSearchers(void) {\n return ailego::Factory::Classes();\n}\n\nIndexStreamer::Pointer IndexFactory::CreateStreamer(const std::string &name) {\n IndexStreamer::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasStreamer(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllStreamers(void) {\n return ailego::Factory::Classes();\n}\n\nIndexReducer::Pointer IndexFactory::CreateReducer(const std::string &name) {\n IndexReducer::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasReducer(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllReducers(void) {\n return ailego::Factory::Classes();\n}\n\n\nIndexCluster::Pointer IndexFactory::CreateCluster(const std::string &name) {\n IndexCluster::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasCluster(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllClusters(void) {\n return ailego::Factory::Classes();\n}\n\nIndexStreamerReducer::Pointer IndexFactory::CreateStreamerReducer(\n const std::string &name) {\n IndexStreamerReducer::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasStreamerReducer(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllStreamerReducers(void) {\n return ailego::Factory::Classes();\n}\n\nIndexRefiner::Pointer IndexFactory::CreateRefiner(const std::string &name) {\n IndexRefiner::Pointer obj =\n ailego::Factory::MakeShared(name.c_str());\n if (obj) {\n obj->set_name(name);\n }\n return obj;\n}\n\nbool IndexFactory::HasRefiner(const std::string &name) {\n return ailego::Factory::Has(name.c_str());\n}\n\nstd::vector IndexFactory::AllRefiners(void) {\n return ailego::Factory::Classes();\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_flow.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n\n//! Default storage\n#define INDEX_FLOW_STORAGE_DEFAULT \"MMapFileReadStorage\"\n\nnamespace zvec {\nnamespace core {\n\n// Index Flow\nint IndexFlow::set_storage(const std::string &name,\n const ailego::Params ¶ms) {\n storage_ = IndexFactory::CreateStorage(name);\n if (!storage_) {\n LOG_ERROR(\"Failed to create a index storage with name: %s\", name.c_str());\n return IndexError_NoExist;\n }\n int ret = storage_->init(params);\n if (ret < 0) {\n storage_ = nullptr;\n LOG_ERROR(\"Failed to initialize index storage %s\", name.c_str());\n return ret;\n }\n return 0;\n}\n\nint IndexFlow::set_searcher(IndexSearcher::Pointer searcher) {\n user_searcher_ = searcher;\n\n return 0;\n}\n\nint IndexFlow::set_searcher(const std::string &name,\n const ailego::Params ¶ms) {\n user_searcher_ = IndexFactory::CreateSearcher(name);\n if (!user_searcher_) {\n LOG_ERROR(\"Failed to create a index searcher with name: %s\", name.c_str());\n return IndexError_NoExist;\n }\n int ret = user_searcher_->init(params);\n if (ret < 0) {\n user_searcher_ = nullptr;\n LOG_ERROR(\"Failed to initialize index searcher %s\", name.c_str());\n return ret;\n }\n return 0;\n}\n\nint IndexFlow::set_reformer(const std::string &name,\n const ailego::Params ¶ms) {\n user_reformer_ = IndexFactory::CreateReformer(name);\n if (!user_reformer_) {\n LOG_ERROR(\"Failed to create a index reformer with name: %s\", name.c_str());\n return IndexError_NoExist;\n }\n int ret = user_reformer_->init(params);\n if (ret < 0) {\n user_reformer_ = nullptr;\n LOG_ERROR(\"Failed to initialize index reformer %s\", name.c_str());\n return ret;\n }\n return 0;\n}\n\nint IndexFlow::set_metric(const std::string &name,\n const ailego::Params ¶ms) {\n if (!IndexFactory::HasMetric(name)) {\n LOG_ERROR(\"The index metric with name %s does not exist.\", name.c_str());\n return IndexError_NoExist;\n }\n user_metric_name_ = name;\n user_metric_params_ = params;\n return 0;\n}\n\nint IndexFlow::load(const std::string &path) {\n // Prepare storage\n if (!storage_) {\n this->set_storage(INDEX_FLOW_STORAGE_DEFAULT, ailego::Params());\n }\n\n if (!storage_) {\n LOG_ERROR(\"The index storage is uninitialized.\");\n return IndexError_Uninitialized;\n }\n\n int ret = storage_->open(path, false);\n if (ret != 0) {\n LOG_ERROR(\"Failed to load index with storage %s\", storage_->name().c_str());\n return ret;\n }\n\n ret = IndexHelper::DeserializeFromStorage(storage_.get(), &meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to deserialize index meta with storage %s\",\n storage_->name().c_str());\n return ret;\n }\n\n ret = load_internal();\n if (ret != 0) {\n LOG_ERROR(\"Failed to load index with storage %s\", storage_->name().c_str());\n return ret;\n }\n\n return 0;\n}\n\nint IndexFlow::load_internal() {\n // Prepare metric\n const std::string &metric_name =\n user_metric_name_.empty() ? meta_.metric_name() : user_metric_name_;\n const ailego::Params &metric_params =\n user_metric_name_.empty() ? meta_.metric_params() : user_metric_params_;\n if (metric_name.empty()) {\n LOG_ERROR(\"The metric name from index file is empty.\");\n return IndexError_NoExist;\n }\n metric_ = IndexFactory::CreateMetric(metric_name);\n if (!metric_) {\n LOG_ERROR(\"Failed to create a index metric with name: %s\",\n metric_name.c_str());\n return IndexError_NoExist;\n }\n int ret = metric_->init(meta_, metric_params);\n if (ret < 0) {\n LOG_ERROR(\"Failed to initialize index metric %s\", metric_name.c_str());\n metric_ = nullptr;\n return ret;\n }\n if (!metric_->is_matched(meta_)) {\n LOG_ERROR(\"The index meta is unmatched for index metric %s\",\n metric_->name().c_str());\n return IndexError_Mismatch;\n }\n auto query_metric = metric_->query_metric();\n if (query_metric) {\n metric_ = query_metric;\n }\n\n // Prepare reformer\n if (!user_reformer_) {\n const std::string &reformer_name = meta_.reformer_name();\n if (!reformer_name.empty()) {\n reformer_ = IndexFactory::CreateReformer(reformer_name);\n if (!reformer_) {\n LOG_ERROR(\"Failed to create a index reformer with name: %s\",\n reformer_name.c_str());\n return IndexError_NoExist;\n }\n ret = reformer_->init(meta_.reformer_params());\n if (ret < 0) {\n LOG_ERROR(\"Failed to initialize index reformer %s\",\n reformer_name.c_str());\n reformer_ = nullptr;\n return ret;\n }\n }\n } else {\n // Using user reformer\n reformer_ = user_reformer_;\n }\n\n if (reformer_) {\n ret = reformer_->load(storage_);\n if (ret < 0) {\n LOG_ERROR(\"Failed to load index with reformer %s, storage %s\",\n reformer_->name().c_str(), storage_->name().c_str());\n return ret;\n }\n }\n\n // Prepare searcher\n if (!user_searcher_) {\n const std::string &name = meta_.searcher_name();\n if (name.empty()) {\n LOG_ERROR(\"The searcher name from index file is empty.\");\n return IndexError_NoExist;\n }\n searcher_ = IndexFactory::CreateSearcher(name);\n if (!searcher_) {\n LOG_ERROR(\"Failed to create a index searcher with name: %s\",\n name.c_str());\n return IndexError_NoExist;\n }\n ret = searcher_->init(meta_.searcher_params());\n if (ret < 0) {\n LOG_ERROR(\"Failed to initialize index searcher %s\", name.c_str());\n searcher_ = nullptr;\n return ret;\n }\n } else {\n // Using user searcher\n searcher_ = user_searcher_;\n }\n\n ret = searcher_->load(storage_, metric_);\n if (ret < 0) {\n LOG_ERROR(\"Failed to load index with searcher %s, storage %s, metric %s\",\n searcher_->name().c_str(), storage_->name().c_str(),\n metric_->name().c_str());\n return ret;\n }\n\n // searcher_->print_all_neighbour();\n\n return 0;\n}\n\nint IndexFlow::unload(void) {\n if (searcher_) {\n int ret = searcher_->unload();\n if (ret < 0) {\n LOG_WARN(\"Unload index searcher %s error, %d\", searcher_->name().c_str(),\n ret);\n }\n searcher_ = nullptr;\n }\n if (reformer_) {\n int ret = reformer_->unload();\n if (ret < 0) {\n LOG_WARN(\"Unload index reformer %s error, %d\", reformer_->name().c_str(),\n ret);\n }\n reformer_ = nullptr;\n }\n if (metric_) {\n int ret = metric_->cleanup();\n if (ret < 0) {\n LOG_WARN(\"Cleanup index metric %s error, %d\", metric_->name().c_str(),\n ret);\n }\n metric_ = nullptr;\n }\n if (storage_) {\n int ret = storage_->cleanup();\n if (ret < 0) {\n LOG_WARN(\"Unload index searcher %s error, %d\", storage_->name().c_str(),\n ret);\n }\n storage_ = nullptr;\n }\n return 0;\n}\n\nint IndexFlow::search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n if (ailego_unlikely(!query || !context)) {\n return IndexError_InvalidArgument;\n }\n\n int error_code = 0;\n if (reformer_) {\n IndexQueryMeta new_qmeta;\n error_code = reformer_->transform(query, qmeta, context->mutable_features(),\n &new_qmeta);\n if (error_code == 0) {\n if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) {\n return IndexError_Mismatch;\n }\n error_code = searcher_->search_bf_impl(\n reinterpret_cast(context->features().data()), new_qmeta,\n context->searcher_context());\n }\n } else {\n if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) {\n return IndexError_Mismatch;\n }\n error_code =\n searcher_->search_bf_impl(query, qmeta, context->searcher_context());\n }\n\n if (error_code == 0) {\n if (metric_->support_normalize()) {\n for (auto &it : const_cast(\n context->searcher_context()->result())) {\n metric_->normalize(it.mutable_score());\n }\n }\n if (reformer_) {\n error_code =\n reformer_->normalize(query, qmeta,\n const_cast(\n context->searcher_context()->result()));\n }\n }\n return error_code;\n}\n\nint IndexFlow::search_impl(const void *query, const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n if (ailego_unlikely(!query || !context)) {\n return IndexError_InvalidArgument;\n }\n\n int error_code = 0;\n if (reformer_) {\n IndexQueryMeta new_qmeta;\n error_code = reformer_->transform(query, qmeta, context->mutable_features(),\n &new_qmeta);\n if (error_code == 0) {\n if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) {\n return IndexError_Mismatch;\n }\n error_code = searcher_->search_impl(\n reinterpret_cast(context->features().data()), new_qmeta,\n context->searcher_context());\n }\n } else {\n if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) {\n return IndexError_Mismatch;\n }\n error_code =\n searcher_->search_impl(query, qmeta, context->searcher_context());\n }\n\n if (error_code == 0) {\n if (metric_->support_normalize()) {\n for (auto &it : const_cast(\n context->searcher_context()->result())) {\n metric_->normalize(it.mutable_score());\n }\n }\n if (reformer_) {\n error_code =\n reformer_->normalize(query, qmeta,\n const_cast(\n context->searcher_context()->result()));\n }\n }\n return error_code;\n}\n\nint IndexFlow::search_bf_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const {\n if (ailego_unlikely(!query || !count || !context)) {\n return IndexError_InvalidArgument;\n }\n\n int error_code = 0;\n if (reformer_) {\n IndexQueryMeta new_qmeta;\n error_code = reformer_->transform(query, qmeta, count,\n context->mutable_features(), &new_qmeta);\n if (error_code == 0) {\n if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) {\n return IndexError_Mismatch;\n }\n error_code = searcher_->search_bf_impl(\n reinterpret_cast(context->features().data()), new_qmeta,\n count, context->searcher_context());\n }\n } else {\n if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) {\n return IndexError_Mismatch;\n }\n error_code = searcher_->search_bf_impl(query, qmeta, count,\n context->searcher_context());\n }\n\n if (error_code == 0) {\n if (metric_->support_normalize()) {\n for (uint32_t i = 0; i < count; ++i) {\n IndexDocumentList &result = const_cast(\n context->searcher_context()->result(i));\n\n for (auto &it : result) {\n metric_->normalize(it.mutable_score());\n }\n }\n }\n if (reformer_) {\n size_t offset = 0;\n for (uint32_t i = 0; i < count; ++i) {\n error_code = reformer_->normalize(\n reinterpret_cast(query) + offset, qmeta,\n const_cast(\n context->searcher_context()->result(i)));\n if (error_code != 0) {\n break;\n }\n offset += qmeta.element_size();\n }\n }\n }\n return error_code;\n}\n\nint IndexFlow::search_impl(const void *query, const IndexQueryMeta &qmeta,\n uint32_t count, Context::Pointer &context) const {\n if (ailego_unlikely(!query || !count || !context)) {\n return IndexError_InvalidArgument;\n }\n\n int error_code = 0;\n if (reformer_) {\n IndexQueryMeta new_qmeta;\n error_code = reformer_->transform(query, qmeta, count,\n context->mutable_features(), &new_qmeta);\n if (error_code == 0) {\n if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) {\n return IndexError_Mismatch;\n }\n error_code = searcher_->search_impl(\n reinterpret_cast(context->features().data()), new_qmeta,\n count, context->searcher_context());\n }\n } else {\n if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) {\n return IndexError_Mismatch;\n }\n error_code = searcher_->search_impl(query, qmeta, count,\n context->searcher_context());\n }\n\n if (error_code == 0) {\n if (metric_->support_normalize()) {\n for (uint32_t i = 0; i < count; ++i) {\n IndexDocumentList &result = const_cast(\n context->searcher_context()->result(i));\n\n for (auto &it : result) {\n metric_->normalize(it.mutable_score());\n }\n }\n }\n if (reformer_) {\n size_t offset = 0;\n for (uint32_t i = 0; i < count; ++i) {\n error_code = reformer_->normalize(\n reinterpret_cast(query) + offset, qmeta,\n const_cast(\n context->searcher_context()->result(i)));\n if (error_code != 0) {\n break;\n }\n offset += qmeta.element_size();\n }\n }\n }\n return error_code;\n}\n\n// Index Sparse Flow\nint IndexSparseFlow::set_storage(const std::string &name,\n const ailego::Params ¶ms) {\n storage_ = IndexFactory::CreateStorage(name);\n if (!storage_) {\n LOG_ERROR(\"Failed to create a index storage with name: %s\", name.c_str());\n return IndexError_NoExist;\n }\n int ret = storage_->init(params);\n if (ret < 0) {\n storage_ = nullptr;\n LOG_ERROR(\"Failed to initialize index storage %s\", name.c_str());\n return ret;\n }\n return 0;\n}\n\nint IndexSparseFlow::set_searcher(IndexSearcher::Pointer searcher) {\n user_searcher_ = searcher;\n\n return 0;\n}\n\nint IndexSparseFlow::set_searcher(const std::string &name,\n const ailego::Params ¶ms) {\n user_searcher_ = IndexFactory::CreateSearcher(name);\n if (!user_searcher_) {\n LOG_ERROR(\"Failed to create a index sparse searcher with name: %s\",\n name.c_str());\n return IndexError_NoExist;\n }\n int ret = user_searcher_->init(params);\n if (ret < 0) {\n user_searcher_ = nullptr;\n LOG_ERROR(\"Failed to initialize index sparse searcher %s\", name.c_str());\n return ret;\n }\n return 0;\n}\n\nint IndexSparseFlow::set_reformer(const std::string &name,\n const ailego::Params ¶ms) {\n user_reformer_ = IndexFactory::CreateReformer(name);\n if (!user_reformer_) {\n LOG_ERROR(\"Failed to create a index sparse reformer with name: %s\",\n name.c_str());\n return IndexError_NoExist;\n }\n int ret = user_reformer_->init(params);\n if (ret < 0) {\n user_reformer_ = nullptr;\n LOG_ERROR(\"Failed to initialize index sparse reformer %s\", name.c_str());\n return ret;\n }\n return 0;\n}\n\nint IndexSparseFlow::set_metric(const std::string &name,\n const ailego::Params ¶ms) {\n if (!IndexFactory::HasMetric(name)) {\n LOG_ERROR(\"The index metric with name %s does not exist.\", name.c_str());\n return IndexError_NoExist;\n }\n user_metric_name_ = name;\n user_metric_params_ = params;\n return 0;\n}\n\nint IndexSparseFlow::load(const std::string &path) {\n // Prepare storage\n if (!storage_) {\n this->set_storage(INDEX_FLOW_STORAGE_DEFAULT, ailego::Params());\n }\n\n if (!storage_) {\n LOG_ERROR(\"The index storage is uninitialized.\");\n return IndexError_Uninitialized;\n }\n\n int ret = storage_->open(path, false);\n if (ret != 0) {\n LOG_ERROR(\"Failed to load index with storage %s\", storage_->name().c_str());\n return ret;\n }\n\n ret = IndexHelper::DeserializeFromStorage(storage_.get(), &meta_);\n if (ret != 0) {\n LOG_ERROR(\"Failed to deserialize index meta with storage %s\",\n storage_->name().c_str());\n return ret;\n }\n\n ret = load_internal();\n if (ret != 0) {\n LOG_ERROR(\"Failed to load index with storage %s\", storage_->name().c_str());\n return ret;\n }\n\n return 0;\n}\n\nint IndexSparseFlow::load_internal() {\n // Prepare metric\n const std::string &metric_name =\n user_metric_name_.empty() ? meta_.metric_name() : user_metric_name_;\n const ailego::Params &metric_params =\n user_metric_name_.empty() ? meta_.metric_params() : user_metric_params_;\n if (metric_name.empty()) {\n LOG_ERROR(\"The metric name from index file is empty.\");\n return IndexError_NoExist;\n }\n metric_ = IndexFactory::CreateMetric(metric_name);\n if (!metric_) {\n LOG_ERROR(\"Failed to create a index metric with name: %s\",\n metric_name.c_str());\n return IndexError_NoExist;\n }\n int ret = metric_->init(meta_, metric_params);\n if (ret < 0) {\n LOG_ERROR(\"Failed to initialize index metric %s\", metric_name.c_str());\n metric_ = nullptr;\n return ret;\n }\n\n auto query_metric = metric_->query_metric();\n if (query_metric) {\n metric_ = query_metric;\n }\n\n // Prepare reformer\n if (!user_reformer_) {\n const std::string &reformer_name = meta_.reformer_name();\n if (!reformer_name.empty()) {\n reformer_ = IndexFactory::CreateReformer(reformer_name);\n if (!reformer_) {\n LOG_ERROR(\"Failed to create a index sparse reformer with name: %s\",\n reformer_name.c_str());\n return IndexError_NoExist;\n }\n ret = reformer_->init(meta_.reformer_params());\n if (ret < 0) {\n LOG_ERROR(\"Failed to initialize index reformer %s\",\n reformer_name.c_str());\n reformer_ = nullptr;\n return ret;\n }\n }\n } else {\n // Using user reformer\n reformer_ = user_reformer_;\n }\n\n if (reformer_) {\n ret = reformer_->load(storage_);\n if (ret < 0) {\n LOG_ERROR(\"Failed to load index with reformer %s, storage %s\",\n reformer_->name().c_str(), storage_->name().c_str());\n return ret;\n }\n }\n\n // Prepare searcher\n if (!user_searcher_) {\n const std::string &name = meta_.searcher_name();\n if (name.empty()) {\n LOG_ERROR(\"The searcher name from index file is empty.\");\n return IndexError_NoExist;\n }\n searcher_ = IndexFactory::CreateSearcher(name);\n if (!searcher_) {\n LOG_ERROR(\"Failed to create a index searcher with name: %s\",\n name.c_str());\n return IndexError_NoExist;\n }\n ret = searcher_->init(meta_.searcher_params());\n if (ret < 0) {\n LOG_ERROR(\"Failed to initialize index searcher %s\", name.c_str());\n searcher_ = nullptr;\n return ret;\n }\n } else {\n // Using user searcher\n searcher_ = user_searcher_;\n }\n\n ret = searcher_->load(storage_, metric_);\n if (ret < 0) {\n LOG_ERROR(\"Failed to load index with searcher %s, storage %s, metric %s\",\n searcher_->name().c_str(), storage_->name().c_str(),\n metric_->name().c_str());\n return ret;\n }\n\n // searcher_->print_all_neighbour();\n\n return 0;\n}\n\nint IndexSparseFlow::unload(void) {\n if (searcher_) {\n int ret = searcher_->unload();\n if (ret < 0) {\n LOG_WARN(\"Unload index searcher %s error, %d\", searcher_->name().c_str(),\n ret);\n }\n searcher_ = nullptr;\n }\n if (reformer_) {\n int ret = reformer_->unload();\n if (ret < 0) {\n LOG_WARN(\"Unload index reformer %s error, %d\", reformer_->name().c_str(),\n ret);\n }\n reformer_ = nullptr;\n }\n if (metric_) {\n int ret = metric_->cleanup();\n if (ret < 0) {\n LOG_WARN(\"Cleanup index metric %s error, %d\", metric_->name().c_str(),\n ret);\n }\n metric_ = nullptr;\n }\n if (storage_) {\n int ret = storage_->cleanup();\n if (ret < 0) {\n LOG_WARN(\"Unload index searcher %s error, %d\", storage_->name().c_str(),\n ret);\n }\n storage_ = nullptr;\n }\n return 0;\n}\n\nint IndexSparseFlow::search_bf_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n if (ailego_unlikely(!context)) {\n return IndexError_InvalidArgument;\n }\n\n int error_code = 0;\n if (reformer_) {\n std::string ovec;\n IndexQueryMeta new_qmeta;\n error_code = reformer_->transform(sparse_count, sparse_indices,\n sparse_query, qmeta, &ovec, &new_qmeta);\n if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) {\n return IndexError_Mismatch;\n }\n if (error_code == 0) {\n error_code =\n searcher_->search_bf_impl(sparse_count, sparse_indices, ovec.data(),\n new_qmeta, context->searcher_context());\n }\n } else {\n if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) {\n return IndexError_Mismatch;\n }\n error_code =\n searcher_->search_bf_impl(sparse_count, sparse_indices, sparse_query,\n qmeta, context->searcher_context());\n }\n\n if (error_code == 0) {\n if (metric_->support_normalize()) {\n for (auto &it : const_cast(\n context->searcher_context()->result())) {\n metric_->normalize(it.mutable_score());\n }\n }\n }\n return error_code;\n}\n\nint IndexSparseFlow::search_impl(const uint32_t sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta,\n Context::Pointer &context) const {\n if (ailego_unlikely(!context)) {\n return IndexError_InvalidArgument;\n }\n\n int error_code = 0;\n if (reformer_) {\n std::string ovec;\n IndexQueryMeta new_qmeta;\n error_code = reformer_->transform(sparse_count, sparse_indices,\n sparse_query, qmeta, &ovec, &new_qmeta);\n if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) {\n return IndexError_Mismatch;\n }\n if (error_code == 0) {\n error_code =\n searcher_->search_impl(sparse_count, sparse_indices, ovec.data(),\n new_qmeta, context->searcher_context());\n }\n } else {\n if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) {\n return IndexError_Mismatch;\n }\n error_code =\n searcher_->search_impl(sparse_count, sparse_indices, sparse_query,\n qmeta, context->searcher_context());\n }\n\n if (error_code == 0) {\n if (metric_->support_normalize()) {\n for (auto &it : const_cast(\n context->searcher_context()->result())) {\n metric_->normalize(it.mutable_score());\n }\n }\n }\n return error_code;\n}\n\nint IndexSparseFlow::search_bf_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n if (ailego_unlikely(!count || !context)) {\n return IndexError_InvalidArgument;\n }\n\n int error_code = 0;\n if (reformer_) {\n std::string ovec;\n IndexQueryMeta new_qmeta;\n error_code =\n reformer_->transform(sparse_count, sparse_indices, sparse_query, qmeta,\n count, &ovec, &new_qmeta);\n\n if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) {\n return IndexError_Mismatch;\n }\n\n if (error_code == 0) {\n error_code = searcher_->search_bf_impl(sparse_count, sparse_indices,\n ovec.data(), new_qmeta, count,\n context->searcher_context());\n }\n } else {\n if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) {\n return IndexError_Mismatch;\n }\n error_code =\n searcher_->search_bf_impl(sparse_count, sparse_indices, sparse_query,\n qmeta, count, context->searcher_context());\n }\n\n if (error_code == 0) {\n if (metric_->support_normalize()) {\n for (uint32_t i = 0; i < count; ++i) {\n IndexDocumentList &result = const_cast(\n context->searcher_context()->result(i));\n\n for (auto &it : result) {\n metric_->normalize(it.mutable_score());\n }\n }\n }\n }\n return error_code;\n}\n\nint IndexSparseFlow::search_impl(const uint32_t *sparse_count,\n const uint32_t *sparse_indices,\n const void *sparse_query,\n const IndexQueryMeta &qmeta, uint32_t count,\n Context::Pointer &context) const {\n if (ailego_unlikely(!count || !context)) {\n return IndexError_InvalidArgument;\n }\n\n int error_code = 0;\n if (reformer_) {\n std::string ovec;\n IndexQueryMeta new_qmeta;\n error_code =\n reformer_->transform(sparse_count, sparse_indices, sparse_query, qmeta,\n count, &ovec, &new_qmeta);\n\n if (ailego_unlikely(!metric_->is_matched(meta_, new_qmeta))) {\n return IndexError_Mismatch;\n }\n\n if (error_code == 0) {\n error_code =\n searcher_->search_impl(sparse_count, sparse_indices, ovec.data(),\n new_qmeta, count, context->searcher_context());\n }\n } else {\n if (ailego_unlikely(!metric_->is_matched(meta_, qmeta))) {\n return IndexError_Mismatch;\n }\n error_code =\n searcher_->search_impl(sparse_count, sparse_indices, sparse_query,\n qmeta, count, context->searcher_context());\n }\n\n if (error_code == 0) {\n if (metric_->support_normalize()) {\n for (uint32_t i = 0; i < count; ++i) {\n IndexDocumentList &result = const_cast(\n context->searcher_context()->result(i));\n\n for (auto &it : result) {\n metric_->normalize(it.mutable_score());\n }\n }\n }\n }\n return error_code;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_helper.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nint IndexHelper::SerializeToDumper(const IndexMeta &mt, IndexDumper *dumper,\n const std::string &key) {\n std::string buffer;\n mt.serialize(&buffer);\n\n size_t data_size = buffer.size();\n uint32_t data_crc = ailego::Crc32c::Hash(buffer.data(), buffer.size(), 0);\n buffer.resize((data_size + 31u) & ~31u);\n\n if (dumper->write(buffer.data(), buffer.size()) != buffer.size()) {\n return IndexError_WriteData;\n }\n if (dumper->append(key, data_size, buffer.size() - data_size, data_crc) !=\n 0) {\n return IndexError_WriteData;\n }\n return IndexError_Success;\n}\n\nint IndexHelper::SerializeToStorage(const IndexMeta &mt, IndexStorage *storage,\n const std::string &key) {\n std::string buffer;\n mt.serialize(&buffer);\n\n auto segment = storage->get(key);\n if (!segment) {\n const size_t align_size = 4096 * 4;\n size_t meta_size =\n (buffer.size() + align_size - 1) / align_size * align_size;\n\n if (storage->append(key, meta_size) != 0) {\n return IndexError_WriteData;\n }\n\n segment = storage->get(key);\n if (!segment) {\n return IndexError_NoExist;\n }\n }\n\n if (segment->write(0, buffer.data(), buffer.size()) != buffer.size()) {\n return IndexError_WriteData;\n }\n segment->resize(buffer.size());\n segment->update_data_crc(\n ailego::Crc32c::Hash(buffer.data(), buffer.size(), 0));\n return IndexError_Success;\n}\n\nint IndexHelper::DeserializeFromStorage(IndexStorage *storage,\n const std::string &key,\n IndexMeta *out) {\n auto segment = storage->get(key);\n if (!segment) {\n return IndexError_NoExist;\n }\n\n uint32_t crc = segment->data_crc();\n size_t len = segment->data_size();\n const void *data = nullptr;\n\n if (segment->read(0, &data, len) != len) {\n return IndexError_ReadData;\n }\n if (crc != 0u && ailego::Crc32c::Hash(data, len, 0u) != crc) {\n return IndexError_InvalidChecksum;\n }\n if (!out->deserialize(data, len)) {\n return IndexError_Deserialize;\n }\n return IndexError_Success;\n}\n\n/*! Two Pass Index Holder\n */\nclass TwoPassIndexHolder : public IndexHolder {\n private:\n /*! First Pass Iterator\n * store elements during iterating for second iterating.\n */\n class FirstPassIterator : public IndexHolder::Iterator {\n public:\n //! Index Holder Iterator Pointer\n typedef std::unique_ptr Pointer;\n\n //! Constructor\n FirstPassIterator(TwoPassIndexHolder *owner,\n IndexHolder::Iterator::Pointer &&iter)\n : holder_(owner), front_iter_(std::move(iter)) {}\n\n //! Destructor\n virtual ~FirstPassIterator(void) {}\n\n //! Retrieve pointer of data\n const void *data(void) const override {\n return front_iter_->data();\n }\n\n //! Test if the iterator is valid\n bool is_valid(void) const override {\n return front_iter_->is_valid();\n }\n\n //! Retrieve primary key\n uint64_t key(void) const override {\n return front_iter_->key();\n }\n\n //! Next iterator\n void next(void) override {\n holder_->features_.emplace_back(\n front_iter_->key(), std::string((const char *)front_iter_->data(),\n holder_->front_->element_size()));\n front_iter_->next();\n }\n\n private:\n TwoPassIndexHolder *holder_{nullptr};\n IndexHolder::Iterator::Pointer front_iter_{};\n };\n\n class SecondPassIterator : public IndexHolder::Iterator {\n public:\n //! Second Pass Iterator Pointer\n typedef std::unique_ptr Pointer;\n\n //! Constructor\n SecondPassIterator(TwoPassIndexHolder *owner) : holder_(owner) {\n features_iter_ = holder_->features_.begin();\n }\n\n //! Destructor\n virtual ~SecondPassIterator(void) {}\n\n //! Retrieve pointer of data\n const void *data(void) const override {\n return features_iter_->second.data();\n }\n\n //! Test if the iterator is valid\n bool is_valid(void) const override {\n return (features_iter_ != holder_->features_.end());\n }\n\n //! Retrieve primary key\n uint64_t key(void) const override {\n return features_iter_->first;\n }\n\n //! Next iterator\n void next(void) override {\n holder_->features_.erase(features_iter_++);\n }\n\n private:\n TwoPassIndexHolder *holder_{nullptr};\n typename std::list>::iterator\n features_iter_{};\n };\n\n public:\n //! Constructor\n TwoPassIndexHolder(IndexHolder::Pointer &&front)\n : front_(std::move(front)),\n data_type_(front_->data_type()),\n dimension_(front_->dimension()),\n element_size_(front_->element_size()),\n count_(front_->count()) {}\n\n //! Retrieve count of elements in holder (-1 indicates unknown)\n size_t count(void) const override {\n return count_;\n }\n\n //! Retrieve dimension\n size_t dimension(void) const override {\n return dimension_;\n }\n\n //! Retrieve type information\n IndexMeta::DataType data_type(void) const override {\n return data_type_;\n }\n\n //! Retrieve element size in bytes\n size_t element_size(void) const override {\n return element_size_;\n }\n\n //! Retrieve if it can multi-pass\n bool multipass(void) const override {\n return false;\n }\n\n //! Create a new iterator\n IndexHolder::Iterator::Pointer create_iterator(void) override {\n ++pass_;\n if (pass_ == 1) {\n IndexHolder::Iterator::Pointer iter = front_->create_iterator();\n return iter ? IndexHolder::Iterator::Pointer(\n new TwoPassIndexHolder::FirstPassIterator(\n this, std::move(iter)))\n : IndexHolder::Iterator::Pointer();\n } else if (pass_ == 2) {\n return IndexHolder::Iterator::Pointer(\n new TwoPassIndexHolder::SecondPassIterator(this));\n }\n return nullptr;\n }\n\n private:\n //! Disable them\n TwoPassIndexHolder(void) = delete;\n\n //! Members\n IndexHolder::Pointer front_{};\n std::list> features_{};\n size_t pass_{0};\n IndexMeta::DataType data_type_{IndexMeta::DataType::DT_UNDEFINED};\n size_t dimension_;\n size_t element_size_;\n size_t count_;\n};\n\nIndexHolder::Pointer IndexHelper::MakeTwoPassHolder(\n IndexHolder::Pointer holder) {\n if (holder->multipass()) {\n return holder;\n }\n return IndexHolder::Pointer(new TwoPassIndexHolder(std::move(holder)));\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_logger.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nconst int IndexLogger::LEVEL_DEBUG = 0;\nconst int IndexLogger::LEVEL_INFO = 1;\nconst int IndexLogger::LEVEL_WARN = 2;\nconst int IndexLogger::LEVEL_ERROR = 3;\nconst int IndexLogger::LEVEL_FATAL = 4;\n\n/*! Console Logger\n */\nstruct ConsoleLogger : public IndexLogger {\n //! Initialize Logger\n int init(const zvec::ailego::Params &) override {\n return 0;\n }\n\n //! Cleanup Logger\n int cleanup(void) override {\n return 0;\n }\n\n //! Log Message\n void log(int level, const char *file, int line, const char *format,\n va_list args) override {\n char buffer[8192];\n std::ostringstream stream;\n\n ailego::Realtime::Localtime(buffer, sizeof(buffer));\n stream << '[' << LevelString(level) << ' ' << buffer << ' '\n << std::this_thread::get_id() << ' ' << ailego::File::BaseName(file)\n << ':' << line << \"] \";\n\n vsnprintf(buffer, sizeof(buffer), format, args);\n stream << buffer << '\\n';\n\n if (level <= LEVEL_INFO) {\n std::cout << stream.str() << std::flush;\n } else {\n std::cerr << stream.str() << std::flush;\n }\n }\n};\n\n//! Logger Level\nint IndexLoggerBroker::logger_level_ = 0;\n\n//! Logger\nIndexLogger::Pointer IndexLoggerBroker::logger_(new ConsoleLogger);\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/framework/index_mapping.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \"ailego/utility/memory_helper.h\"\n\n#ifdef __linux__\n#include \n#include \n#ifndef HUGETLBFS_MAGIC\n#define HUGETLBFS_MAGIC 0x958458f6\n#endif\n#endif\n\nnamespace zvec {\nnamespace core {\n\nstatic inline size_t CalcPageAlignedSize(size_t size, bool huge_size) {\n size_t page_size = ailego::MemoryHelper::PageSize();\n if (huge_size) {\n page_size = ailego::MemoryHelper::HugePageSize();\n }\n return (size + page_size - 1) / page_size * page_size;\n}\n\nstatic inline bool WritePadding(ailego::File &file, size_t size) {\n std::string padding(ailego::MemoryHelper::PageSize(), 0);\n for (size_t i = 0, count = size / padding.size(); i < count; ++i) {\n if (file.write(padding.data(), padding.size()) != padding.size()) {\n return false;\n }\n }\n padding.resize(size % padding.size());\n if (padding.size()) {\n if (file.write(padding.data(), padding.size()) != padding.size()) {\n return false;\n }\n }\n return true;\n}\n\nstatic inline int UnpackMappingSize(ailego::File &file, size_t *len) {\n IndexFormat::MetaHeader header;\n if (file.read(&header, sizeof(header)) != sizeof(header)) {\n LOG_ERROR(\"Failed to read file, errno %d, %s\", errno, std::strerror(errno));\n return IndexError_ReadData;\n }\n\n if (header.meta_header_size != sizeof(IndexFormat::MetaHeader) ||\n header.meta_footer_size != sizeof(IndexFormat::MetaFooter)) {\n return IndexError_InvalidValue;\n }\n\n if (ailego::Crc32c::Hash(&header, sizeof(header), header.header_crc) !=\n header.header_crc) {\n return IndexError_InvalidChecksum;\n }\n\n if ((int32_t)header.meta_footer_offset < 0) {\n return IndexError_Unsupported;\n }\n\n *len = header.meta_footer_offset + header.meta_footer_size;\n if (*len > file.size()) {\n return IndexError_InvalidLength;\n }\n return 0;\n}\n\nint IndexMapping::open(const std::string &path, bool cow, bool full_mode) {\n path_ = path;\n full_mode_ = full_mode;\n copy_on_write_ = cow;\n huge_page_ = Ishugetlbfs(path);\n\n bool read_only = copy_on_write_ && !full_mode_;\n if (!file_.open(path.c_str(), read_only, false)) {\n LOG_ERROR(\"Failed to open file %s, errno %d, %s\", path.c_str(), errno,\n std::strerror(errno));\n return IndexError_OpenFile;\n }\n\n size_t mapping_size = 0u;\n int error_code = UnpackMappingSize(file_, &mapping_size);\n if (error_code != 0) {\n file_.close();\n return error_code;\n }\n\n if (!file_.seek(0, ailego::File::Origin::End)) {\n LOG_ERROR(\"Failed to seek file %s, errno %d, %s\", path.c_str(), errno,\n std::strerror(errno));\n return IndexError_SeekFile;\n }\n return this->init_index_mapping(mapping_size);\n}\n\nint IndexMapping::create(const std::string &path, size_t seg_meta_capacity) {\n path_ = path;\n seg_meta_capacity_ = seg_meta_capacity;\n current_header_start_offset_ = 0;\n\n // write() & copying to mmap() will auto extend the file size\n if (!file_.create(path.c_str(), 0)) {\n LOG_ERROR(\"Failed to create file %s, errno %d, %s\", path.c_str(), errno,\n std::strerror(errno));\n return IndexError_CreateFile;\n }\n huge_page_ = Ishugetlbfs(path);\n if (huge_page_) {\n return init_hugepage_meta_section();\n }\n return init_meta_section();\n}\n\nint IndexMapping::init_meta_section() {\n if (current_header_start_offset_ % ailego::MemoryHelper::PageSize() != 0) {\n LOG_ERROR(\"File offset %zu is not a multiple of the page size: %zu\",\n current_header_start_offset_, ailego::MemoryHelper::PageSize());\n return IndexError_InvalidValue;\n }\n\n auto &path = path_;\n size_t len =\n CalcPageAlignedSize(seg_meta_capacity_ + sizeof(IndexFormat::MetaHeader) +\n sizeof(IndexFormat::MetaFooter),\n false);\n\n IndexFormat::MetaHeader meta_header;\n IndexFormat::MetaFooter meta_footer;\n\n // Write index header\n IndexFormat::SetupMetaHeader(&meta_header, len - sizeof(meta_footer), len);\n if (!file_.seek(current_header_start_offset_, ailego::File::Origin::Begin)) {\n LOG_ERROR(\"Failed to seek file %s, errno %d, %s\", path.c_str(), errno,\n std::strerror(errno));\n return IndexError_SeekFile;\n }\n if (file_.write(&meta_header, sizeof(meta_header)) != sizeof(meta_header)) {\n LOG_ERROR(\"Failed to write file: %s, errno %d, %s\", path.c_str(), errno,\n std::strerror(errno));\n return IndexError_WriteData;\n }\n\n // Write padding data\n uint32_t segments_meta_size =\n static_cast(len - (sizeof(meta_header) + sizeof(meta_footer)));\n if (!WritePadding(file_, segments_meta_size)) {\n LOG_ERROR(\"Failed to write file: %s, errno %d, %s\", path.c_str(), errno,\n std::strerror(errno));\n return IndexError_WriteData;\n }\n\n // Write index footer\n IndexFormat::SetupMetaFooter(&meta_footer);\n meta_footer.segments_meta_size = segments_meta_size;\n meta_footer.total_size = len;\n IndexFormat::UpdateMetaFooter(&meta_footer, 0);\n if (file_.write(&meta_footer, sizeof(meta_footer)) != sizeof(meta_footer)) {\n LOG_ERROR(\"Failed to write file: %s, errno %d, %s\", path.c_str(), errno,\n std::strerror(errno));\n return IndexError_WriteData;\n }\n return this->init_index_mapping(len);\n}\n\nint IndexMapping::init_hugepage_meta_section() {\n ssize_t file_offset = (ssize_t)current_header_start_offset_;\n if (file_offset % ailego::MemoryHelper::HugePageSize() != 0) {\n LOG_ERROR(\"File offset %zu is not a multiple of the page size: %zu\",\n file_offset, ailego::MemoryHelper::HugePageSize());\n return IndexError_InvalidValue;\n }\n\n size_t len =\n CalcPageAlignedSize(seg_meta_capacity_ + sizeof(IndexFormat::MetaHeader) +\n sizeof(IndexFormat::MetaFooter),\n true);\n int opts = ailego::File::MMAP_SHARED | ailego::File::MMAP_HUGE_PAGE;\n void *addr =\n ailego::File::MemoryMap(file_.native_handle(), file_offset, len, opts);\n\n IndexFormat::MetaHeader meta_header;\n IndexFormat::MetaFooter meta_footer;\n\n // Write index header\n IndexFormat::SetupMetaHeader(&meta_header, len - sizeof(meta_footer), len);\n memcpy((char *)addr + file_offset, &meta_header, sizeof(meta_header));\n file_offset += sizeof(meta_header);\n\n // Write padding data\n uint32_t segments_meta_size =\n static_cast(len - (sizeof(meta_header) + sizeof(meta_footer)));\n std::string padding(ailego::MemoryHelper::HugePageSize(), 0);\n for (size_t i = 0, count = segments_meta_size / padding.size(); i < count;\n ++i) {\n memcpy((char *)addr + file_offset, padding.data(), padding.size());\n file_offset += padding.size();\n }\n padding.resize(segments_meta_size % padding.size());\n if (padding.size()) {\n memcpy((char *)addr + file_offset, padding.data(), padding.size());\n file_offset += padding.size();\n }\n\n // Write index footer\n IndexFormat::SetupMetaFooter(&meta_footer);\n meta_footer.segments_meta_size = segments_meta_size;\n meta_footer.total_size = len;\n IndexFormat::UpdateMetaFooter(&meta_footer, 0);\n memcpy((char *)addr + file_offset, &meta_footer, sizeof(meta_footer));\n file_offset += sizeof(meta_footer);\n\n return this->init_index_mapping(len);\n}\n\nvoid IndexMapping::close(void) {\n // Unmap all memory\n this->unmap_all();\n if (header_) {\n for (auto item : header_addr_map_) {\n auto header = item.second;\n ailego::File::MemoryUnmap(header, header->content_offset);\n }\n }\n // Reset members\n segment_ids_offset_ = 0;\n segment_start_ = nullptr;\n header_ = nullptr;\n header_addr_map_.clear();\n footer_ = nullptr;\n index_size_ = 0u;\n segments_.clear();\n file_.close();\n copy_on_write_ = false;\n full_mode_ = false;\n header_dirty_ = false;\n huge_page_ = false;\n}\n\nvoid IndexMapping::refresh(uint64_t check_point) {\n // support add_with_id\n for (auto item : header_addr_map_) {\n auto header_start_offset = item.first;\n auto header = item.second;\n auto footer = reinterpret_cast(\n reinterpret_cast(header) + header->meta_footer_offset);\n auto segment_start = reinterpret_cast(\n reinterpret_cast(header) +\n (header->meta_footer_offset - footer->segments_meta_size));\n footer->segments_meta_crc =\n ailego::Crc32c::Hash(segment_start, footer->segments_meta_size, 0);\n IndexFormat::UpdateMetaFooter(footer, check_point);\n }\n header_dirty_ = true;\n}\n\nint IndexMapping::append(const std::string &id, size_t size) {\n size = CalcPageAlignedSize(size, huge_page_);\n if (size == 0) {\n return IndexError_InvalidArgument;\n }\n\n if (segments_.find(id) != segments_.end()) {\n return IndexError_Duplicate;\n }\n\n size_t id_size = std::strlen(id.c_str()) + 1;\n size_t need_size = sizeof(IndexFormat::SegmentMeta) + id_size;\n if (sizeof(IndexFormat::SegmentMeta) * footer_->segment_count + need_size >\n segment_ids_offset_) {\n LOG_DEBUG(\"segment meta section expanded: %s\", path_.c_str());\n footer_->next_meta_header_offset = index_size_;\n refresh(0);\n flush();\n // mmap file storage write() will update segment's meta\n // ailego::File::MemoryUnmap(header_, header_->content_offset);\n header_ = nullptr;\n footer_ = nullptr;\n\n current_header_start_offset_ = index_size_;\n const int ret =\n huge_page_ ? init_hugepage_meta_section() : init_meta_section();\n if (ret != 0) {\n return ret;\n }\n }\n\n if (!copy_on_write_ && !file_.truncate(index_size_ + size)) {\n LOG_ERROR(\"Failed to truncate file, errno %d, %s\", errno,\n std::strerror(errno));\n return IndexError_TruncateFile;\n }\n\n // Update segment table\n segment_ids_offset_ -= static_cast(id_size);\n IndexFormat::SegmentMeta *segment = segment_start_ + footer_->segment_count;\n segment->segment_id_offset = segment_ids_offset_;\n segment->data_index =\n index_size_ - header_->content_offset - current_header_start_offset_;\n segment->data_size = 0;\n segment->data_crc = 0;\n segment->padding_size = size;\n memcpy((uint8_t *)segment_start_ + segment_ids_offset_, id.c_str(), id_size);\n index_size_ += size;\n\n // Update index footer\n footer_->segments_meta_crc =\n ailego::Crc32c::Hash(segment_start_, footer_->segments_meta_size, 0);\n footer_->segment_count += 1;\n footer_->content_size += size;\n footer_->total_size += size;\n IndexFormat::UpdateMetaFooter(footer_, 0);\n segments_.emplace(\n id, SegmentInfo{Segment{segment}, current_header_start_offset_, header_});\n header_dirty_ = true;\n return 0;\n}\n\nIndexMapping::Segment *IndexMapping::map(const std::string &id, bool warmup,\n bool locked) {\n auto iter = segments_.find(id);\n if (iter == segments_.end()) {\n return nullptr;\n }\n SegmentInfo &segment_info = iter->second;\n Segment *item = &segment_info.segment;\n if (!item->data()) {\n auto meta = item->meta();\n size_t mapping_size = meta->data_size + meta->padding_size;\n size_t offset = segment_info.segment_header_start_offset +\n segment_info.segment_header->content_offset +\n meta->data_index;\n\n void *addr = nullptr;\n if (!copy_on_write_) {\n int opts = ailego::File::MMAP_SHARED;\n if (huge_page_) {\n opts |= ailego::File::MMAP_HUGE_PAGE;\n }\n addr = ailego::File::MemoryMap(file_.native_handle(), offset,\n mapping_size, opts);\n } else {\n size_t file_size = file_.size();\n int opts = ailego::File::MMAP_POPULATE;\n if (huge_page_) {\n opts |= ailego::File::MMAP_HUGE_PAGE;\n }\n if (offset < file_size) {\n ailego_assert(offset + mapping_size <= file_size);\n addr = ailego::File::MemoryMap(file_.native_handle(), offset,\n mapping_size, opts);\n } else {\n addr = ailego::File::MemoryMap(mapping_size, opts);\n }\n }\n\n if (!addr) {\n LOG_ERROR(\"Map segment failed, segment id %s\", id.c_str());\n return nullptr;\n }\n item->set_data(addr);\n\n // Lock memory\n if (locked) {\n ailego::File::MemoryLock(item->data(), mapping_size);\n }\n // Warmup memory\n if (warmup && meta->data_size) {\n ailego::File::MemoryWarmup(item->data(), meta->data_size);\n }\n }\n return item;\n}\n\nvoid IndexMapping::unmap(const std::string &id) {\n auto iter = segments_.find(id);\n if (iter != segments_.end()) {\n SegmentInfo &segment_info = iter->second;\n Segment *item = &segment_info.segment;\n\n if (item->data()) {\n ailego::File::MemoryUnmap(\n item->data(), item->meta()->data_size + item->meta()->padding_size);\n item->set_data(nullptr);\n }\n }\n}\n\nvoid IndexMapping::unmap_all(void) {\n for (auto iter = segments_.begin(); iter != segments_.end(); ++iter) {\n SegmentInfo &segment_info = iter->second;\n Segment *item = &segment_info.segment;\n\n if (item->data()) {\n ailego::File::MemoryUnmap(\n item->data(), item->meta()->data_size + item->meta()->padding_size);\n item->set_data(nullptr);\n }\n }\n}\n\nint IndexMapping::flush(void) {\n if ((file_.size() < index_size_) && !file_.truncate(index_size_)) {\n LOG_ERROR(\"Failed to truncate file size %zu, errno %d, %s\", index_size_,\n errno, std::strerror(errno));\n return IndexError_TruncateFile;\n }\n\n for (auto iter = segments_.begin(); iter != segments_.end(); ++iter) {\n SegmentInfo &segment_info = iter->second;\n Segment *item = &segment_info.segment;\n if (!item->data() || !item->dirty()) {\n continue;\n }\n\n size_t segment_size = item->meta()->data_size + item->meta()->padding_size;\n if (full_mode_ && copy_on_write_) {\n size_t off = segment_info.segment_header_start_offset +\n segment_info.segment_header->content_offset +\n item->meta()->data_index;\n if (file_.write(off, item->data(), segment_size) != segment_size) {\n LOG_ERROR(\"Failed to write segment, size %zu, errno %d, %s\",\n segment_size, errno, std::strerror(errno));\n return IndexError_WriteData;\n }\n } else {\n ailego::File::MemoryFlush(item->data(), segment_size);\n }\n item->reset_dirty();\n }\n\n if (!header_dirty_) {\n return 0;\n }\n\n header_dirty_ = false;\n if (full_mode_ && copy_on_write_) {\n for (auto item : header_addr_map_) {\n auto header_start_offset = item.first;\n auto header = item.second;\n if (file_.write(header_start_offset, header, header->content_offset) !=\n header->content_offset) {\n LOG_ERROR(\"Failed to write segment, size %lu, errno %d, %s\",\n header->content_offset, errno, std::strerror(errno));\n return IndexError_WriteData;\n }\n }\n } else {\n for (auto item : header_addr_map_) {\n auto header = item.second;\n ailego::File::MemoryFlush(header, header->content_offset);\n }\n }\n return 0;\n}\n\nint IndexMapping::init_index_mapping(size_t len) {\n int opts =\n copy_on_write_ ? ailego::File::MMAP_POPULATE : ailego::File::MMAP_SHARED;\n if (huge_page_) {\n opts |= ailego::File::MMAP_HUGE_PAGE;\n }\n uint8_t *start = reinterpret_cast(ailego::File::MemoryMap(\n file_.native_handle(), current_header_start_offset_, len, opts));\n if (!start) {\n LOG_ERROR(\"Failed to map file, errno %d, %s\", errno, std::strerror(errno));\n return IndexError_MMapFile;\n }\n\n // Unpack header\n header_ = reinterpret_cast(start);\n header_addr_map_.insert({current_header_start_offset_, header_});\n if (header_->meta_header_size != sizeof(IndexFormat::MetaHeader)) {\n return IndexError_InvalidLength;\n }\n if (ailego::Crc32c::Hash(header_, sizeof(*header_), header_->header_crc) !=\n header_->header_crc) {\n return IndexError_InvalidChecksum;\n }\n\n switch (header_->version) {\n case IndexFormat::FORMAT_VERSION:\n break;\n default:\n LOG_ERROR(\"Unsupported index version: %u\", header_->version);\n return IndexError_Unsupported;\n }\n\n // Unpack footer\n if (header_->meta_footer_size != sizeof(IndexFormat::MetaFooter)) {\n return IndexError_InvalidLength;\n }\n if ((int32_t)header_->meta_footer_offset < 0) {\n return IndexError_Unsupported;\n }\n size_t footer_offset = header_->meta_footer_offset;\n if (footer_offset + header_->meta_footer_size > len) {\n return IndexError_InvalidLength;\n }\n\n footer_ = reinterpret_cast(start + footer_offset);\n if (footer_offset < footer_->segments_meta_size) {\n return IndexError_InvalidLength;\n }\n\n index_size_ = file_.size();\n if ((footer_->total_size > index_size_) ||\n (footer_->content_size + footer_->content_padding_size +\n header_->content_offset >\n index_size_)) {\n return IndexError_InvalidLength;\n }\n if (ailego::Crc32c::Hash(footer_, sizeof(*footer_), footer_->footer_crc) !=\n footer_->footer_crc) {\n return IndexError_InvalidChecksum;\n }\n\n // Unpack segment table\n if (sizeof(IndexFormat::SegmentMeta) * footer_->segment_count >\n footer_->segments_meta_size) {\n return IndexError_InvalidLength;\n }\n\n segment_start_ = reinterpret_cast(\n start + (footer_offset - footer_->segments_meta_size));\n if (ailego::Crc32c::Hash(segment_start_, footer_->segments_meta_size, 0u) !=\n footer_->segments_meta_crc) {\n LOG_ERROR(\"Index segments meta checksum is invalid.\");\n return IndexError_InvalidChecksum;\n }\n\n segment_ids_offset_ = footer_->segments_meta_size;\n for (IndexFormat::SegmentMeta *iter = segment_start_,\n *end = segment_start_ + footer_->segment_count;\n iter != end; ++iter) {\n if (iter->segment_id_offset > footer_->segments_meta_size) {\n return IndexError_InvalidValue;\n }\n if (iter->data_index > footer_->content_size) {\n return IndexError_InvalidValue;\n }\n if (iter->data_index + iter->data_size > footer_->content_size) {\n return IndexError_InvalidLength;\n }\n\n if (iter->segment_id_offset < segment_ids_offset_) {\n segment_ids_offset_ = iter->segment_id_offset;\n }\n segments_.emplace(\n std::string(reinterpret_cast(segment_start_) +\n iter->segment_id_offset),\n SegmentInfo{Segment{iter}, current_header_start_offset_, header_});\n }\n if (sizeof(IndexFormat::SegmentMeta) * footer_->segment_count >\n segment_ids_offset_) {\n return IndexError_InvalidLength;\n }\n\n // if (header_->version == IndexFormat::COMPATIBLE_FORMAT_VERSION_0X0002) {\n // header_->version = IndexFormat::CURRENT_FORMAT_VERSION;\n // LOG_INFO(\"Index file format upgraded\");\n // IndexFormat::UpdateMetaHeader(header_);\n // footer_->segments_meta_crc =\n // ailego::Crc32c::Hash(segment_start_, footer_->segments_meta_size, 0);\n // IndexFormat::UpdateMetaFooter(footer_, 0);\n // header_dirty_ = true;\n // }\n\n if (footer_->next_meta_header_offset > 0) {\n current_header_start_offset_ = footer_->next_meta_header_offset;\n // Meta sections have all the same size, so we can use the same size to map\n // the next meta section\n return this->init_index_mapping(len);\n }\n\n return 0;\n}\n\nbool IndexMapping::Ishugetlbfs(const std::string &path) const {\n#ifdef __linux__\n struct statfs buf;\n if (statfs(path.c_str(), &buf) != 0) {\n perror(\"statfs\");\n return false;\n }\n return static_cast(buf.f_type) == HUGETLBFS_MAGIC;\n#else\n static_cast(path);\n return false;\n#endif\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_meta.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\n/*! Index Meta Buffer Format\n */\nstruct IndexMetaFormatHeader {\n uint32_t header_size;\n uint32_t meta_type;\n uint32_t major_order;\n uint32_t data_type;\n uint32_t dimension;\n uint32_t unit_size;\n uint32_t space_id;\n uint32_t attachment_offset;\n uint32_t attachment_size;\n uint8_t reserved_[4092];\n};\n\nstatic_assert(sizeof(IndexMetaFormatHeader) % 32 == 0,\n \"IndexMetaBufferFormat must be aligned with 32 bytes\");\n\nvoid IndexMeta::serialize(std::string *out) const {\n ailego::Params attachment;\n IndexMetaFormatHeader format;\n memset(&format, 0, sizeof(format));\n format.header_size = sizeof(format);\n format.meta_type = static_cast(meta_type_);\n format.major_order = static_cast(major_order_);\n format.data_type = static_cast(data_type_);\n format.dimension = dimension_;\n format.unit_size = unit_size_;\n format.space_id = space_id_;\n\n if (!metric_name_.empty()) {\n ailego::Params item;\n item.set(\"name\", metric_name_);\n item.set(\"revision\", metric_revision_);\n item.set(\"params\", metric_params_);\n attachment.set(\"metric\", std::move(item));\n }\n\n if (!converter_name_.empty()) {\n ailego::Params item;\n item.set(\"name\", converter_name_);\n item.set(\"revision\", converter_revision_);\n item.set(\"params\", converter_params_);\n attachment.set(\"converter\", std::move(item));\n }\n if (!reformer_name_.empty()) {\n ailego::Params item;\n item.set(\"name\", reformer_name_);\n item.set(\"revision\", reformer_revision_);\n item.set(\"params\", reformer_params_);\n attachment.set(\"reformer\", std::move(item));\n }\n if (!trainer_name_.empty()) {\n ailego::Params item;\n item.set(\"name\", trainer_name_);\n item.set(\"revision\", trainer_revision_);\n item.set(\"params\", trainer_params_);\n attachment.set(\"trainer\", std::move(item));\n }\n if (!builder_name_.empty()) {\n ailego::Params item;\n item.set(\"name\", builder_name_);\n item.set(\"revision\", builder_revision_);\n item.set(\"params\", builder_params_);\n attachment.set(\"builder\", std::move(item));\n }\n if (!reducer_name_.empty()) {\n ailego::Params item;\n item.set(\"name\", reducer_name_);\n item.set(\"revision\", reducer_revision_);\n item.set(\"params\", reducer_params_);\n attachment.set(\"reducer\", std::move(item));\n }\n if (!searcher_name_.empty()) {\n ailego::Params item;\n item.set(\"name\", searcher_name_);\n item.set(\"revision\", searcher_revision_);\n item.set(\"params\", searcher_params_);\n attachment.set(\"searcher\", std::move(item));\n }\n if (!streamer_name_.empty()) {\n ailego::Params item;\n item.set(\"name\", streamer_name_);\n item.set(\"revision\", streamer_revision_);\n item.set(\"params\", streamer_params_);\n attachment.set(\"streamer\", std::move(item));\n }\n\n if (!attributes_.empty()) {\n attachment.set(\"attributes\", attributes_);\n }\n\n out->assign(reinterpret_cast(&format), sizeof(format));\n size_t offset = static_cast(out->size());\n\n if (!attachment.empty()) {\n std::string buf;\n ailego::Params::SerializeToBuffer(attachment, &buf);\n out->append(buf.data(), buf.size());\n IndexMetaFormatHeader *header = (IndexMetaFormatHeader *)out->data();\n header->attachment_offset = static_cast(offset);\n header->attachment_size = static_cast(buf.size());\n offset += buf.size();\n }\n}\n\nbool IndexMeta::deserialize(const void *data, size_t len) {\n const IndexMetaFormatHeader *format =\n reinterpret_cast(data);\n\n this->clear();\n if (sizeof(IndexMetaFormatHeader) > len) {\n return false;\n }\n if (sizeof(IndexMetaFormatHeader) > format->header_size) {\n return false;\n }\n\n meta_type_ = static_cast(format->meta_type);\n major_order_ = static_cast(format->major_order);\n data_type_ = static_cast(format->data_type);\n dimension_ = format->dimension;\n unit_size_ = format->unit_size;\n element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_);\n space_id_ = format->space_id;\n\n // Read attachment\n ailego::Params attachment;\n if (format->attachment_size) {\n if (format->attachment_offset + format->attachment_size > len) {\n return false;\n }\n std::string str(\n reinterpret_cast(data) + format->attachment_offset,\n format->attachment_size);\n if (!ailego::Params::ParseFromBuffer(str, &attachment)) {\n return false;\n }\n }\n\n ailego::Params item;\n if (attachment.get(\"metric\", &item)) {\n item.get(\"name\", &metric_name_);\n item.get(\"revision\", &metric_revision_);\n item.get(\"params\", &metric_params_);\n }\n if (attachment.get(\"converter\", &item)) {\n item.get(\"name\", &converter_name_);\n item.get(\"revision\", &converter_revision_);\n item.get(\"params\", &converter_params_);\n }\n if (attachment.get(\"reformer\", &item)) {\n item.get(\"name\", &reformer_name_);\n item.get(\"revision\", &reformer_revision_);\n item.get(\"params\", &reformer_params_);\n }\n if (attachment.get(\"trainer\", &item)) {\n item.get(\"name\", &trainer_name_);\n item.get(\"revision\", &trainer_revision_);\n item.get(\"params\", &trainer_params_);\n }\n if (attachment.get(\"builder\", &item)) {\n item.get(\"name\", &builder_name_);\n item.get(\"revision\", &builder_revision_);\n item.get(\"params\", &builder_params_);\n }\n if (attachment.get(\"reducer\", &item)) {\n item.get(\"name\", &reducer_name_);\n item.get(\"revision\", &reducer_revision_);\n item.get(\"params\", &reducer_params_);\n }\n if (attachment.get(\"searcher\", &item)) {\n item.get(\"name\", &searcher_name_);\n item.get(\"revision\", &searcher_revision_);\n item.get(\"params\", &searcher_params_);\n }\n if (attachment.get(\"streamer\", &item)) {\n item.get(\"name\", &streamer_name_);\n item.get(\"revision\", &streamer_revision_);\n item.get(\"params\", &streamer_params_);\n }\n attachment.get(\"attributes\", &attributes_);\n\n return true;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_plugin.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\nbool IndexPlugin::load(const std::string &path) {\n if (handle_) {\n return false;\n }\n handle_ = ailego::DLHelper::Load(path, nullptr);\n return (!!handle_);\n}\n\nbool IndexPlugin::load(const std::string &path, std::string *err) {\n if (handle_) {\n *err = \"plugin loaded\";\n return false;\n }\n handle_ = ailego::DLHelper::Load(path, err);\n return !!handle_;\n}\n\nvoid IndexPlugin::unload(void) {\n if (handle_) {\n ailego::DLHelper::Unload(handle_);\n handle_ = nullptr;\n }\n}\n\nbool IndexPluginBroker::emplace(IndexPlugin &&plugin) {\n if (!plugin.is_valid()) {\n return false;\n }\n for (auto iter = plugins_.begin(); iter != plugins_.end(); ++iter) {\n if (iter->handle() == plugin.handle()) {\n plugin.unload();\n return true;\n }\n }\n plugins_.push_back(std::move(plugin));\n return true;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/framework/index_version.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n\n\nnamespace zvec {\nnamespace core {\n\nstatic const char AITHETA_VERSION_DETAILS[] =\n AILEGO_VERSION_COMPILE_DETAILS(\"All rights reserved.\\n\");\n\nconst char *IndexVersion::String(void) {\n return AITHETA_VERSION_DETAILS;\n}\n\nconst char *IndexVersion::Details(void) {\n return AITHETA_VERSION_DETAILS;\n}\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/interface/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\ncc_library(\n NAME core_interface STATIC STRICT ALWAYS_LINK\n SRCS *.cc indexes/*.cc\n INCS . ${PROJECT_ROOT_DIR}/src/ ${PROJECT_ROOT_DIR}/src/core\n LIBS zvec_ailego core_framework sparsehash magic_enum rabitqlib\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n)\n" }, { "path": "src/core/interface/index.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \"mixed_reducer/mixed_reducer_params.h\"\n\nnamespace zvec::core_interface {\n\n// eliminate the pre-alloc of the context pool\nthread_local static std::array() - 1) * 2>\n _context_list;\n\n\nbool Index::init_context() {\n context_index_ = (magic_enum::enum_integer(param_.index_type) - 1) * 2 +\n static_cast(is_sparse_);\n if (_context_list[context_index_] == nullptr) {\n if ((_context_list[context_index_] = streamer_->create_context()) ==\n nullptr) {\n LOG_ERROR(\"Failed to create context\");\n return false;\n }\n }\n return true;\n}\n\ncore::IndexContext::Pointer &Index::acquire_context() {\n init_context();\n return _context_list[context_index_];\n}\n\nint Index::ParseMetricName(const BaseIndexParam ¶m) {\n std::string metric_name;\n if (is_sparse_) {\n // only inner product is supported for sparse index\n switch (param.metric_type) {\n case MetricType::kInnerProduct:\n metric_name = \"InnerProductSparse\";\n break;\n case MetricType::kMIPSL2sq:\n metric_name = \"MipsSquaredEuclideanSparse\";\n break;\n default:\n LOG_ERROR(\"Unsupported metric type\");\n return core::IndexError_Runtime;\n }\n } else {\n switch (param.metric_type) {\n case MetricType::kL2sq:\n metric_name = \"SquaredEuclidean\";\n break;\n case MetricType::kInnerProduct:\n metric_name = \"InnerProduct\";\n break;\n case MetricType::kCosine:\n metric_name = \"Cosine\"; // This is already the normalizedCosine\n break;\n case MetricType::kMIPSL2sq:\n metric_name = \"MipsSquaredEuclidean\";\n break;\n default:\n LOG_ERROR(\"Unsupported metric type\");\n return core::IndexError_Runtime;\n }\n }\n // TODO: MIPS need to set some param\n // for streamer open()\n proxima_index_meta_.set_metric(metric_name, 0, ailego::Params());\n return 0;\n}\n\nint Index::CreateAndInitMetric(const BaseIndexParam & /*param*/) {\n auto &metric_name = proxima_index_meta_.metric_name();\n\n metric_ = core::IndexFactory::CreateMetric(metric_name);\n if (!metric_) {\n LOG_ERROR(\"Failed to create metric, name %s\", metric_name.c_str());\n return core::IndexError_Runtime;\n }\n if (const auto ret = metric_->init(proxima_index_meta_,\n proxima_index_meta_.metric_params());\n ret != 0) {\n LOG_ERROR(\"Failed to create and init metric, name %s, code %d, desc: %s\",\n metric_name.c_str(), ret, core::IndexError::What(ret));\n return core::IndexError_Runtime;\n }\n if (metric_->query_metric()) {\n metric_ = metric_->query_metric();\n }\n\n return core::IndexError_Success;\n}\n\nint Index::CreateAndInitConverterReformer(const QuantizerParam ¶m,\n const BaseIndexParam &index_param) {\n ailego::Params converter_params;\n std::string converter_name;\n if (is_sparse_) {\n switch (param.type) {\n case QuantizerType::kNone:\n return core::IndexError_Success;\n case QuantizerType::kFP16:\n converter_name = \"HalfFloatSparseConverter\";\n break;\n default:\n LOG_ERROR(\"Unsupported quantizer type: \");\n return core::IndexError_Unsupported;\n }\n } else {\n if (index_param.metric_type == MetricType::kCosine) {\n switch (param.type) {\n case QuantizerType::kNone:\n if (index_param.data_type == DataType::DT_FP16) {\n converter_name = \"CosineHalfFloatConverter\";\n } else if (index_param.data_type == DataType::DT_FP32) {\n converter_name = \"CosineNormalizeConverter\";\n } else {\n LOG_ERROR(\"Unsupported data type: \");\n return core::IndexError_Unsupported;\n }\n break;\n case QuantizerType::kRabitq:\n if (index_param.data_type == DataType::DT_FP32) {\n converter_name = \"CosineNormalizeConverter\";\n } else {\n LOG_ERROR(\"Unsupported data type: \");\n return core::IndexError_Unsupported;\n }\n break;\n case QuantizerType::kFP16:\n converter_name = \"CosineFp16Converter\";\n break;\n case QuantizerType::kInt8:\n converter_name = \"CosineInt8Converter\";\n break;\n case QuantizerType::kInt4:\n converter_name = \"CosineInt4Converter\";\n break;\n default:\n LOG_ERROR(\"Unsupported quantizer type: \");\n return core::IndexError_Unsupported;\n }\n } else {\n switch (param.type) {\n case QuantizerType::kNone:\n return core::IndexError_Success;\n case QuantizerType::kFP16:\n converter_name = \"HalfFloatConverter\";\n break;\n case QuantizerType::kInt8:\n converter_name = \"Int8StreamingConverter\";\n break;\n case QuantizerType::kInt4:\n converter_name = \"Int4StreamingConverter\";\n break;\n case QuantizerType::kRabitq:\n // no converter here\n return 0;\n default:\n LOG_ERROR(\"Unsupported quantizer type: \");\n return core::IndexError_Unsupported;\n }\n }\n }\n\n proxima_index_meta_.set_converter(converter_name, 0, converter_params);\n converter_ = core::IndexFactory::CreateConverter(converter_name);\n if (converter_ == nullptr ||\n converter_->init(proxima_index_meta_, converter_params) != 0) {\n LOG_ERROR(\"Failed to create and init converter\");\n return core::IndexError_Runtime;\n }\n\n proxima_index_meta_ = converter_->meta();\n reformer_ =\n core::IndexFactory::CreateReformer(proxima_index_meta_.reformer_name());\n if (reformer_ == nullptr ||\n reformer_->init(proxima_index_meta_.reformer_params()) != 0) {\n LOG_ERROR(\"Failed to create and init reformer\");\n return core::IndexError_Runtime;\n }\n streamer_vector_meta_.set_meta(proxima_index_meta_.data_type(),\n proxima_index_meta_.dimension());\n streamer_vector_meta_.set_meta_type(proxima_index_meta_.meta_type());\n\n return core::IndexError_Success;\n}\n\nint Index::Init(const BaseIndexParam ¶m) {\n param_ = param; // will lose the original type info\n\n is_sparse_ = param.is_sparse;\n is_huge_page_ = param.is_huge_page;\n\n proxima_index_meta_.set_meta(param.data_type, param.dimension);\n proxima_index_meta_.set_meta_type(is_sparse_ ? IndexMeta::MetaType::MT_SPARSE\n : IndexMeta::MetaType::MT_DENSE);\n\n input_vector_meta_.set_meta(proxima_index_meta_.data_type(),\n proxima_index_meta_.dimension());\n input_vector_meta_.set_meta_type(proxima_index_meta_.meta_type());\n streamer_vector_meta_ = input_vector_meta_;\n\n\n // when quantizer=int8/int4, the converter.init() will change the metric to\n // QuantizedInteger with params\n if (ParseMetricName(param) != 0) {\n LOG_ERROR(\"Failed to parse metric name\");\n return core::IndexError_Runtime;\n }\n\n if (CreateAndInitConverterReformer(param.quantizer_param, param) != 0) {\n LOG_ERROR(\"Failed to create and init converter\");\n return core::IndexError_Runtime;\n }\n\n // must after quantizer handled. e.g., cosine doesn't support int8 quantizer\n if (CreateAndInitMetric(param) != 0) {\n LOG_ERROR(\"Failed to create and init metric\");\n return core::IndexError_Runtime;\n }\n\n if (CreateAndInitStreamer(param) != 0) {\n LOG_ERROR(\"Failed to create and init streamer\");\n return core::IndexError_Runtime;\n }\n return 0;\n}\n\n\nint Index::Open(const std::string &file_path, StorageOptions storage_options) {\n ailego::Params storage_params;\n // storage_params.set(\"proxima.mmap_file.storage.memory_warmup\", true);\n // storage_params.set(\"proxima.mmap_file.storage.segment_meta_capacity\",\n // 1024);\n switch (storage_options.type) {\n case StorageOptions::StorageType::kMMAP: {\n storage_ = core::IndexFactory::CreateStorage(\"MMapFileStorage\");\n if (storage_ == nullptr) {\n LOG_ERROR(\"Failed to create MMapFileStorage\");\n return core::IndexError_Runtime;\n }\n int ret = storage_->init(storage_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to init MMapFileStorage, path: %s, err: %s\",\n file_path.c_str(), core::IndexError::What(ret));\n return ret;\n }\n break;\n }\n case StorageOptions::StorageType::kBufferPool: {\n storage_ = core::IndexFactory::CreateStorage(\"BufferStorage\");\n if (storage_ == nullptr) {\n LOG_ERROR(\"Failed to create BufferStorage\");\n return core::IndexError_Runtime;\n }\n int ret = storage_->init(storage_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to init BufferStorage, path: %s, err: %s\",\n file_path.c_str(), core::IndexError::What(ret));\n return ret;\n }\n break;\n }\n default:\n LOG_ERROR(\"Unsupported storage type\");\n return core::IndexError_Unsupported;\n }\n\n // read_options.create_new\n int ret = storage_->open(file_path, storage_options.create_new);\n if (ret != 0) {\n LOG_ERROR(\"Failed to open storage, path: %s, err: %s\", file_path.c_str(),\n core::IndexError::What(ret));\n return core::IndexError_Runtime;\n }\n if (streamer_ == nullptr || streamer_->open(storage_) != 0) {\n LOG_ERROR(\"Failed to open streamer, path: %s\", file_path.c_str());\n return core::IndexError_Runtime;\n }\n\n // converter/reformer/metric are created in IndexFactory::CreateIndex\n // TODO: init\n\n // TODO: context pool\n if (!init_context()) { // to validate if any error, will be overwritten\n LOG_ERROR(\"Failed to init context\");\n return core::IndexError_Runtime;\n }\n\n is_open_ = true;\n is_read_only_ = storage_options.read_only;\n return 0;\n}\n\nint Index::Close() {\n if (!is_open_) {\n LOG_ERROR(\"Index is not open\");\n return core::IndexError_Runtime;\n }\n\n if (!is_read_only_) {\n if (ailego_unlikely(Flush() != 0)) {\n LOG_ERROR(\"Failed to cleanup streamer\");\n return core::IndexError_Runtime;\n }\n }\n if (ailego_unlikely(streamer_->cleanup() != 0)) {\n LOG_ERROR(\"Failed to cleanup streamer\");\n return core::IndexError_Runtime;\n }\n if (ailego_unlikely(storage_->close() != 0)) {\n LOG_ERROR(\"Failed to close storage\");\n return core::IndexError_Runtime;\n }\n is_open_ = false;\n return 0;\n}\n\nint Index::Flush() {\n if (!is_open_) {\n LOG_ERROR(\"Index is not open\");\n return core::IndexError_Runtime;\n }\n\n if (is_read_only_) {\n LOG_ERROR(\"Cannot flush read-only index\");\n return core::IndexError_Runtime;\n }\n if (ailego_unlikely(streamer_->flush(0) != 0)) {\n LOG_ERROR(\"Failed to flush streamer\");\n return core::IndexError_Runtime;\n }\n if (ailego_unlikely(storage_->flush() != 0)) {\n LOG_ERROR(\"Failed to flush storage\");\n return core::IndexError_Runtime;\n }\n return 0;\n}\n\nint Index::Fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) {\n if (!is_open_) {\n LOG_ERROR(\"Index is not open\");\n return core::IndexError_Runtime;\n }\n if (is_sparse_) {\n return _sparse_fetch(doc_id, vector_data_buffer);\n }\n return _dense_fetch(doc_id, vector_data_buffer);\n}\n\nint Index::Add(const VectorData &vector_data, const uint32_t doc_id) {\n if (!is_open_) {\n LOG_ERROR(\"Index is not open\");\n return core::IndexError_Runtime;\n }\n\n if (is_read_only_) {\n LOG_ERROR(\"Cannot add to read-only index\");\n return core::IndexError_Runtime;\n }\n\n auto &context = acquire_context();\n if (!context) {\n LOG_ERROR(\"Failed to acquire context\");\n return core::IndexError_Runtime;\n }\n\n int ret = 0;\n if (is_sparse_) {\n ret = _sparse_add(vector_data, doc_id, context);\n } else {\n ret = _dense_add(vector_data, doc_id, context);\n }\n context->reset();\n return ret;\n}\n\n\nint Index::Search(const VectorData &vector_data,\n const BaseIndexQueryParam::Pointer &search_param,\n SearchResult *result) {\n if (!is_open_) {\n LOG_ERROR(\"Index is not open\");\n return core::IndexError_Runtime;\n }\n\n if (!is_trained_ && this->Train() != 0) {\n LOG_ERROR(\"Failed to train index\");\n return core::IndexError_Runtime;\n }\n\n auto &context = acquire_context();\n if (!context) {\n LOG_ERROR(\"Failed to acquire context\");\n return core::IndexError_Runtime;\n }\n\n if (_prepare_for_search(vector_data, search_param, context) != 0) {\n LOG_ERROR(\"Failed to prepare for search\");\n context->reset();\n return core::IndexError_Runtime;\n }\n\n if (is_sparse_) {\n int ret = _sparse_search(vector_data, search_param, result, context);\n context->reset();\n return ret;\n }\n\n // dense support refiner, but sparse doesn't\n int ret = 0;\n if (search_param->refiner_param == nullptr) {\n ret = _dense_search(vector_data, search_param, result, context);\n context->reset();\n } else {\n auto &reference_index = search_param->refiner_param->reference_index;\n if (reference_index == nullptr) {\n LOG_ERROR(\"Reference index is not set\");\n context->reset();\n return core::IndexError_Runtime;\n }\n // TODO: tackle query_param's type info loss to loosen the constraint\n if (reference_index->param_.index_type != IndexType::kFlat) {\n LOG_ERROR(\"Reference index is not flat\");\n context->reset();\n return core::IndexError_Runtime;\n }\n\n context->set_topk(_get_coarse_search_topk(search_param));\n context->set_fetch_vector(false); // no need to fetch vector\n if (_dense_search(vector_data, search_param, result, context) != 0) {\n LOG_ERROR(\"Failed to search\");\n context->reset();\n return core::IndexError_Runtime;\n }\n\n auto &base_result = context->result();\n std::vector keys(base_result.size());\n for (size_t i = 0; i < base_result.size(); ++i) {\n keys[i] = base_result[i].key();\n }\n\n FlatQueryParam::Pointer flat_search_param =\n std::make_shared();\n flat_search_param->topk = search_param->topk;\n flat_search_param->fetch_vector = search_param->fetch_vector;\n flat_search_param->filter = search_param->filter;\n // TODO: should copy other params?\n flat_search_param->bf_pks = std::make_shared>(keys);\n\n ret = reference_index->Search(vector_data, flat_search_param, result);\n }\n context->reset();\n return ret;\n}\n\n\nint Index::_dense_fetch(const uint32_t doc_id,\n VectorDataBuffer *vector_data_buffer) {\n core::IndexStorage::MemoryBlock vector_block;\n int ret = streamer_->get_vector_by_id(doc_id, vector_block);\n if (ret != 0) {\n LOG_ERROR(\"Failed to fetch vector, doc_id: %u\", doc_id);\n return core::IndexError_Runtime;\n }\n const void *vector = vector_block.data();\n\n DenseVectorBuffer dense_vector_buffer;\n std::string &out_vector_buffer = dense_vector_buffer.data;\n // for int4, unit_size * dim != element_size\n out_vector_buffer.resize(input_vector_meta_.element_size());\n\n if (reformer_ != nullptr) {\n if (reformer_->revert(vector, streamer_vector_meta_, &out_vector_buffer) !=\n 0) {\n LOG_ERROR(\"Failed to convert vector\");\n return core::IndexError_Runtime;\n }\n } else {\n out_vector_buffer = std::string(\n static_cast(vector),\n input_vector_meta_.dimension() * input_vector_meta_.unit_size());\n }\n vector_data_buffer->vector_buffer = std::move(dense_vector_buffer);\n return 0;\n}\n\n\nint Index::_sparse_fetch(const uint32_t doc_id,\n VectorDataBuffer *vector_data_buffer) {\n SparseVectorBuffer sparse_vector_buffer;\n\n if (0 != streamer_->get_sparse_vector_by_id(\n doc_id, &sparse_vector_buffer.count,\n &sparse_vector_buffer.indices, &sparse_vector_buffer.values)) {\n LOG_ERROR(\"Failed to fetch vector\");\n return core::IndexError_Runtime;\n }\n\n if (reformer_ != nullptr) {\n std::string reverted_sparse_values_buffer;\n if (reformer_->revert(\n sparse_vector_buffer.count, sparse_vector_buffer.get_indices(),\n sparse_vector_buffer.get_values(), streamer_vector_meta_,\n &reverted_sparse_values_buffer) != 0) {\n LOG_ERROR(\"Failed to convert vector\");\n return core::IndexError_Runtime;\n }\n sparse_vector_buffer.values = std::move(reverted_sparse_values_buffer);\n }\n vector_data_buffer->vector_buffer = std::move(sparse_vector_buffer);\n return 0;\n}\n\nint Index::_dense_add(const VectorData &vector_data, const uint32_t doc_id,\n core::IndexContext::Pointer &context) {\n if (!std::holds_alternative(vector_data.vector)) {\n LOG_ERROR(\"Invalid vector data\");\n return core::IndexError_Runtime;\n }\n const DenseVector &dense_vector = std::get(vector_data.vector);\n if (reformer_ != nullptr) {\n core::IndexQueryMeta new_meta;\n std::string new_vector;\n int ret;\n ret = reformer_->convert(dense_vector.data, input_vector_meta_, &new_vector,\n &new_meta);\n if (ret != 0) {\n LOG_ERROR(\"Failed to convert vector\");\n return core::IndexError_Runtime;\n }\n ret = streamer_->add_with_id_impl(doc_id, new_vector.data(), new_meta,\n context);\n if (ret != 0) {\n LOG_ERROR(\"Failed to add vector\");\n return core::IndexError_Runtime;\n }\n } else {\n int ret = streamer_->add_with_id_impl(doc_id, dense_vector.data,\n input_vector_meta_, context);\n if (ret != 0) {\n LOG_ERROR(\"Failed to add vector\");\n return core::IndexError_Runtime;\n }\n }\n return 0;\n}\n\n\nint Index::_sparse_add(const VectorData &vector_data, const uint32_t doc_id,\n core::IndexContext::Pointer &context) {\n if (!std::holds_alternative(vector_data.vector)) {\n LOG_ERROR(\"Invalid vector data\");\n return core::IndexError_Runtime;\n }\n const SparseVector &sparse_vector =\n std::get(vector_data.vector);\n\n if (reformer_ != nullptr) {\n std::string converted_sparse_values_buffer;\n core::IndexQueryMeta new_meta;\n int ret;\n ret = reformer_->convert(sparse_vector.count, sparse_vector.get_indices(),\n sparse_vector.get_values(), input_vector_meta_,\n &converted_sparse_values_buffer, &new_meta);\n if (ret != 0) {\n LOG_ERROR(\"Failed to convert vector\");\n return core::IndexError_Runtime;\n }\n ret = streamer_->add_with_id_impl(\n doc_id, sparse_vector.count, sparse_vector.get_indices(),\n converted_sparse_values_buffer.data(), new_meta, context);\n if (ret != 0) {\n LOG_ERROR(\"Failed to add vector\");\n return core::IndexError_Runtime;\n }\n } else {\n int ret = streamer_->add_with_id_impl(\n doc_id, sparse_vector.count, sparse_vector.get_indices(),\n sparse_vector.get_values(), input_vector_meta_, context);\n if (ret != 0) {\n LOG_ERROR(\"Failed to add vector\");\n return core::IndexError_Runtime;\n }\n }\n return 0;\n}\n\n\nint Index::_dense_search(const VectorData &vector_data,\n const BaseIndexQueryParam::Pointer &search_param,\n SearchResult *result,\n core::IndexContext::Pointer &context) {\n if (!std::holds_alternative(vector_data.vector)) {\n LOG_ERROR(\"Invalid vector data\");\n return core::IndexError_Runtime;\n }\n const DenseVector &dense_vector = std::get(vector_data.vector);\n auto vector = dense_vector.data;\n // Check if need to transform feature\n std::string new_vector;\n core::IndexQueryMeta new_meta = input_vector_meta_;\n if (reformer_ != nullptr) {\n if (reformer_->transform(dense_vector.data, input_vector_meta_, &new_vector,\n &new_meta) != 0) {\n LOG_ERROR(\"Failed to transform vector\");\n return core::IndexError_Runtime;\n }\n vector = new_vector.data();\n }\n // TODO: group by\n if (search_param->bf_pks != nullptr) {\n // should we eliminate the copy of bf_pks?\n if (streamer_->search_bf_by_p_keys_impl(\n vector, std::vector>{*search_param->bf_pks},\n new_meta, 1, context) != 0) {\n LOG_ERROR(\"Failed to search_bf_by_p_keys_impl vector\");\n return core::IndexError_Runtime;\n }\n } else if (search_param->is_linear) {\n if (streamer_->search_bf_impl(vector, new_meta, 1, context) != 0) {\n LOG_ERROR(\"Failed to search vector\");\n return core::IndexError_Runtime;\n }\n } else {\n if (streamer_->search_impl(vector, new_meta, 1, context) != 0) {\n LOG_ERROR(\"Failed to search vector\");\n return core::IndexError_Runtime;\n }\n }\n result->doc_list_ = std::move(context->result());\n\n if (metric_->support_normalize()) {\n for (uint32_t i = 0; i < result->doc_list_.size(); ++i) {\n metric_->normalize(result->doc_list_[i].mutable_score());\n }\n }\n if (reformer_) {\n if (reformer_->normalize(dense_vector.data, input_vector_meta_,\n result->doc_list_) != 0) {\n LOG_ERROR(\"Failed to normalize vector\");\n return core::IndexError_Runtime;\n }\n if (context->fetch_vector() && reformer_->need_revert()) {\n // TODO: use std::pmr to optimize memory allocation\n result->reverted_vector_list_.resize(context->result().size());\n for (uint32_t i = 0; i < context->result().size(); ++i) {\n std::string &reverted_vector = result->reverted_vector_list_[i];\n reverted_vector.resize(input_vector_meta_.dimension() *\n input_vector_meta_.unit_size());\n if (reformer_->revert(context->result()[i].vector(), new_meta,\n &reverted_vector) != 0) {\n LOG_ERROR(\"Failed to revert vector\");\n return core::IndexError_Runtime;\n }\n }\n }\n }\n\n return 0;\n}\n\n\nint Index::_sparse_search(const VectorData &vector_data,\n const BaseIndexQueryParam::Pointer &search_param,\n SearchResult *result,\n core::IndexContext::Pointer &context) {\n if (!std::holds_alternative(vector_data.vector)) {\n LOG_ERROR(\"Invalid vector data\");\n return core::IndexError_Runtime;\n }\n const SparseVector &sparse_vector =\n std::get(vector_data.vector);\n auto indices = sparse_vector.get_indices();\n auto values = sparse_vector.get_values();\n\n std::string converted_sparse_values_buffer;\n core::IndexQueryMeta new_meta = input_vector_meta_;\n if (reformer_ != nullptr) {\n if (reformer_->transform(sparse_vector.count, indices, values,\n input_vector_meta_,\n &converted_sparse_values_buffer, &new_meta) != 0) {\n LOG_ERROR(\"Failed to transform vector\");\n return core::IndexError_Runtime;\n }\n values = converted_sparse_values_buffer.data();\n }\n\n if (search_param->bf_pks != nullptr) {\n if (streamer_->search_bf_by_p_keys_impl(\n sparse_vector.count, indices, values,\n std::vector>{*search_param->bf_pks}, new_meta,\n context) != 0) {\n LOG_ERROR(\"Failed to search_bf_by_p_keys_impl vector\");\n return core::IndexError_Runtime;\n }\n } else if (search_param->is_linear) {\n if (streamer_->search_bf_impl(sparse_vector.count, indices, values,\n new_meta, context) != 0) {\n LOG_ERROR(\"Failed to search vector\");\n return core::IndexError_Runtime;\n }\n } else {\n if (streamer_->search_impl(sparse_vector.count, indices, values, new_meta,\n context) != 0) {\n LOG_ERROR(\"Failed to search vector\");\n return core::IndexError_Runtime;\n }\n }\n result->doc_list_ = std::move(context->result());\n\n if (metric_->support_normalize()) {\n for (uint32_t i = 0; i < result->doc_list_.size(); ++i) {\n metric_->normalize(result->doc_list_[i].mutable_score());\n }\n }\n if (reformer_) {\n // TODO: no need to call reformer_->normalize() when sparse?\n if (context->fetch_vector() && reformer_->need_revert()) {\n // TODO: use std::pmr to optimize memory allocation\n auto &result_doc_list = context->result();\n result->reverted_sparse_values_list_.resize(result_doc_list.size());\n for (uint32_t i = 0; i < result_doc_list.size(); ++i) {\n auto &result_doc = result_doc_list[i].sparse_doc();\n std::string &reverted_sparse_values =\n result->reverted_sparse_values_list_[i];\n reverted_sparse_values.resize(result_doc.sparse_count() *\n input_vector_meta_.unit_size());\n if (reformer_->revert(result_doc.sparse_count(),\n reinterpret_cast(\n result_doc.sparse_indices().data()),\n reinterpret_cast(\n result_doc.sparse_values().data()),\n new_meta, &reverted_sparse_values) != 0) {\n LOG_ERROR(\"Failed to revert sparse vector\");\n return core::IndexError_Runtime;\n }\n }\n }\n }\n return 0;\n}\n\n\nint Index::Merge(const std::vector &indexes,\n const IndexFilter &filter, const MergeOptions &options) {\n if (indexes.empty()) {\n return core::IndexError_Success;\n }\n // ivf need builder\n auto reducer =\n core::IndexFactory::CreateStreamerReducer(\"MixedStreamerReducer\");\n if (reducer == nullptr) {\n LOG_ERROR(\"Failed to create reducer\");\n return core::IndexError_Runtime;\n }\n\n if (options.write_concurrency == 0) {\n LOG_ERROR(\"Write concurrency must be greater than 0\");\n return core::IndexError_InvalidArgument;\n }\n // must declare here to ensure its lifespan can cover reducer->reduce()\n std::unique_ptr local_thread_pool = nullptr;\n if (options.pool != nullptr) {\n reducer->set_thread_pool(options.pool);\n } else {\n local_thread_pool =\n std::make_unique(options.write_concurrency);\n reducer->set_thread_pool(local_thread_pool.get());\n }\n\n ailego::Params reducer_params;\n reducer_params.set(core::PARAM_MIXED_STREAMER_REDUCER_ENABLE_PK_REWRITE,\n true);\n reducer_params.set(core::PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS,\n options.write_concurrency);\n if (reducer->init(reducer_params) != 0) {\n LOG_ERROR(\"Failed to init reducer\");\n return core::IndexError_Runtime;\n }\n if (reducer->set_target_streamer_wiht_info(builder_, streamer_, converter_,\n reformer_,\n input_vector_meta_) != 0) {\n LOG_ERROR(\"Failed to set target streamer\");\n return core::IndexError_Runtime;\n }\n\n for (const auto &index : indexes) {\n if (reducer->feed_streamer_with_reformer(index->streamer_,\n index->reformer_) != 0) {\n LOG_ERROR(\"Failed to feed streamer\");\n return core::IndexError_Runtime;\n }\n }\n if (reducer->reduce(filter) != 0) {\n LOG_ERROR(\"Failed to reduce\");\n return core::IndexError_Runtime;\n }\n is_trained_ = true;\n return 0;\n}\n\nint Index::_get_coarse_search_topk(\n const BaseIndexQueryParam::Pointer &search_param) {\n float scale_factor = search_param->refiner_param->scale_factor_;\n if (scale_factor == 0) {\n scale_factor = 1;\n }\n return floor(search_param->topk * scale_factor);\n}\n\nstd::string Index::get_metric_name(MetricType metric_type, bool is_sparse) {\n if (is_sparse) {\n switch (metric_type) {\n case MetricType::kInnerProduct:\n return \"InnerProductSparse\";\n case MetricType::kMIPSL2sq:\n return \"MipsSquaredEuclideanSparse\";\n default:\n return \"\";\n }\n } else {\n switch (metric_type) {\n case MetricType::kL2sq:\n return \"SquaredEuclidean\";\n case MetricType::kInnerProduct:\n return \"InnerProduct\";\n case MetricType::kCosine:\n return \"Cosine\";\n case MetricType::kMIPSL2sq:\n return \"MipsSquaredEuclidean\";\n default:\n return \"\";\n }\n }\n}\n\n} // namespace zvec::core_interface\n" }, { "path": "src/core/interface/index_factory.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n#include \"core/interface/utils/utils.h\"\n\nnamespace zvec::core_interface {\n\n\nIndex::Pointer IndexFactory::CreateAndInitIndex(const BaseIndexParam ¶m) {\n Index::Pointer ptr = nullptr;\n // if (param.index_type == IndexType::kIVF) {\n // const IVFIndexParam *_param = dynamic_cast(¶m); ptr = std::make_shared(param);\n\n // if (_param->l1Index) {\n // // TODO: create l1 index\n // }\n // if (_param->l2Index) {\n // // TODO: create l2 index\n // }\n // }\n // if (param.index_type == IndexType::kHNSW) {\n // ptr = std::make_shared(param);\n // }\n if (param.index_type == IndexType::kFlat) {\n // ptr = std::make_shared(param);\n ptr = std::make_shared();\n } else if (param.index_type == IndexType::kHNSW) {\n ptr = std::make_shared();\n } else if (param.index_type == IndexType::kIVF) {\n ptr = std::make_shared();\n } else if (param.index_type == IndexType::kHNSWRabitq) {\n ptr = std::make_shared();\n } else {\n LOG_ERROR(\"Unsupported index type: \");\n return nullptr;\n }\n\n if (!ptr) {\n LOG_ERROR(\"Failed to create index\");\n return nullptr;\n }\n if (0 != ptr->Init(param)) {\n LOG_ERROR(\"Failed to init index\");\n return nullptr;\n }\n return ptr;\n}\n\nBaseIndexParam::Pointer IndexFactory::DeserializeIndexParamFromJson(\n const std::string &json_str) {\n ailego::JsonValue json_value;\n if (!json_value.parse(json_str)) {\n LOG_ERROR(\"Failed to parse json string: %s\", json_str.c_str());\n return nullptr;\n }\n ailego::JsonObject json_obj = json_value.as_object();\n ailego::JsonValue tmp_json_value;\n\n IndexType index_type;\n\n if (!extract_enum_from_json(json_obj, \"index_type\", index_type,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize index type\");\n return nullptr;\n }\n\n switch (index_type) {\n case IndexType::kFlat: {\n FlatIndexParam::Pointer param = std::make_shared();\n if (!param->DeserializeFromJson(json_str)) {\n LOG_ERROR(\"Failed to deserialize flat index param\");\n return nullptr;\n }\n return param;\n }\n case IndexType::kHNSW: {\n HNSWIndexParam::Pointer param = std::make_shared();\n if (!param->DeserializeFromJson(json_str)) {\n LOG_ERROR(\"Failed to deserialize hnsw index param\");\n return nullptr;\n }\n return param;\n }\n case IndexType::kIVF: {\n IVFIndexParam::Pointer param = std::make_shared();\n if (!param->DeserializeFromJson(json_str)) {\n LOG_ERROR(\"Failed to deserialize hnsw index param\");\n return nullptr;\n }\n return param;\n }\n case IndexType::kHNSWRabitq: {\n HNSWRabitqIndexParam::Pointer param =\n std::make_shared();\n if (!param->DeserializeFromJson(json_str)) {\n LOG_ERROR(\"Failed to deserialize hnsqrabitq index param\");\n return nullptr;\n }\n return param;\n }\n default:\n LOG_ERROR(\"Unsupported index type: %s\",\n magic_enum::enum_name(index_type).data());\n return nullptr;\n }\n}\n\ntemplate , bool> >\nstd::string IndexFactory::QueryParamSerializeToJson(const QueryParamType ¶m,\n bool omit_empty_value) {\n ailego::JsonObject json_obj;\n\n // BaseIndexQueryParam\n // omit filter & bf_pks\n if (!omit_empty_value || param.topk != 0) {\n json_obj.set(\"topk\", ailego::JsonValue(param.topk));\n }\n if (!omit_empty_value || param.fetch_vector) {\n json_obj.set(\"fetch_vector\", ailego::JsonValue(param.fetch_vector));\n }\n if (!omit_empty_value || param.radius != 0.0f) {\n json_obj.set(\"radius\", ailego::JsonValue(param.radius));\n }\n if (!omit_empty_value || param.is_linear) {\n json_obj.set(\"is_linear\", ailego::JsonValue(param.is_linear));\n }\n\n IndexType index_type{IndexType::kNone};\n if constexpr (std::is_same_v) {\n // index_type\n index_type = IndexType::kFlat;\n } else if constexpr (std::is_same_v) {\n if (!omit_empty_value || param.ef_search != 0) {\n json_obj.set(\"ef_search\", ailego::JsonValue(param.ef_search));\n }\n index_type = IndexType::kHNSW;\n } else if constexpr (std::is_same_v) {\n if (!omit_empty_value || param.nprobe != 0) {\n json_obj.set(\"nprobe\", ailego::JsonValue(param.nprobe));\n }\n index_type = IndexType::kIVF;\n // json_obj.set(\"l1QueryParam\",\n // ailego::JsonValue(QueryParamSerializeToJson(param.l1QueryParam)));\n // json_obj.set(\"l2QueryParam\",\n // ailego::JsonValue(QueryParamSerializeToJson(param.l2QueryParam)));\n } else if constexpr (std::is_same_v) {\n if (!omit_empty_value || param.ef_search != 0) {\n json_obj.set(\"ef_search\", ailego::JsonValue(param.ef_search));\n }\n index_type = IndexType::kHNSWRabitq;\n }\n\n json_obj.set(\"index_type\",\n ailego::JsonValue(magic_enum::enum_name(index_type).data()));\n\n return ailego::JsonValue(json_obj).as_json_string().as_stl_string();\n}\n\ntemplate std::string\nIndexFactory::QueryParamSerializeToJson(\n const BaseIndexQueryParam ¶m, bool omit_empty_value);\ntemplate std::string IndexFactory::QueryParamSerializeToJson(\n const FlatQueryParam ¶m, bool omit_empty_value);\ntemplate std::string IndexFactory::QueryParamSerializeToJson(\n const HNSWQueryParam ¶m, bool omit_empty_value);\ntemplate std::string IndexFactory::QueryParamSerializeToJson(\n const IVFQueryParam ¶m, bool omit_empty_value);\n\ntemplate , bool> >\ntypename QueryParamType::Pointer IndexFactory::QueryParamDeserializeFromJson(\n const std::string &json_str) {\n ailego::JsonValue tmp_json_value;\n if (!tmp_json_value.parse(json_str)) {\n LOG_ERROR(\"Failed to parse json string: %s\", json_str.c_str());\n return nullptr;\n }\n ailego::JsonObject json_obj = tmp_json_value.as_object();\n\n auto parse_common_fields = [&](auto ¶m) -> bool {\n if (!extract_value_from_json(json_obj, \"topk\", param->topk,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize topk\");\n return false;\n }\n\n if (!extract_value_from_json(json_obj, \"fetch_vector\", param->fetch_vector,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize fetch_vector\");\n return false;\n }\n\n if (!extract_value_from_json(json_obj, \"radius\", param->radius,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize radius\");\n return false;\n }\n\n if (!extract_value_from_json(json_obj, \"is_linear\", param->is_linear,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize is_linear\");\n return false;\n }\n return true;\n };\n\n IndexType index_type;\n\n if (!extract_enum_from_json(json_obj, \"index_type\", index_type,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize index type\");\n return nullptr;\n }\n\n if constexpr (std::is_same_v) {\n if (index_type == IndexType::kFlat) {\n auto param = std::make_shared();\n if (!parse_common_fields(param)) {\n return nullptr;\n }\n return param;\n } else if (index_type == IndexType::kHNSW) {\n auto param = std::make_shared();\n if (!parse_common_fields(param)) {\n return nullptr;\n }\n if (!extract_value_from_json(json_obj, \"ef_search\", param->ef_search,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize ef_search\");\n return nullptr;\n }\n return param;\n } else if (index_type == IndexType::kIVF) {\n auto param = std::make_shared();\n if (!parse_common_fields(param)) {\n return nullptr;\n }\n if (!extract_value_from_json(json_obj, \"nprobe\", param->nprobe,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize nprobe\");\n return nullptr;\n }\n return param;\n } else if (index_type == IndexType::kHNSWRabitq) {\n auto param = std::make_shared();\n if (!parse_common_fields(param)) {\n return nullptr;\n }\n if (!extract_value_from_json(json_obj, \"ef_search\", param->ef_search,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize ef_search\");\n return nullptr;\n }\n return param;\n } else {\n LOG_ERROR(\"Unsupported index type: %s\",\n magic_enum::enum_name(index_type).data());\n return nullptr;\n }\n } else {\n auto param = std::make_shared();\n if (!parse_common_fields(param)) {\n return nullptr;\n }\n if constexpr (std::is_same_v) {\n } else if constexpr (std::is_same_v) {\n if (!extract_value_from_json(json_obj, \"ef_search\", param->ef_search,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize ef_search\");\n return nullptr;\n }\n } else if constexpr (std::is_same_v) {\n if (!extract_value_from_json(json_obj, \"nprobe\", param->nprobe,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize nprobe\");\n return nullptr;\n }\n } else if constexpr (std::is_same_v) {\n if (!extract_value_from_json(json_obj, \"ef_search\", param->ef_search,\n tmp_json_value)) {\n LOG_ERROR(\"Failed to deserialize ef_search\");\n return nullptr;\n }\n } else {\n LOG_ERROR(\"Unsupported index type: %s\",\n magic_enum::enum_name(index_type).data());\n return nullptr;\n }\n return param;\n }\n}\n\ntemplate BaseIndexQueryParam::Pointer\nIndexFactory::QueryParamDeserializeFromJson(\n const std::string &json_str);\ntemplate FlatQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson<\n FlatQueryParam>(const std::string &json_str);\ntemplate HNSWQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson<\n HNSWQueryParam>(const std::string &json_str);\ntemplate IVFQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson<\n IVFQueryParam>(const std::string &json_str);\n\n} // namespace zvec::core_interface\n" }, { "path": "src/core/interface/index_param.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \"core/interface/utils/utils.h\"\n\nnamespace zvec {\nnamespace core_interface {\nailego::JsonObject BaseIndexParam::SerializeToJsonObject(\n bool omit_empty_value) const {\n ailego::JsonObject json_obj;\n\n if (!omit_empty_value || index_type != IndexType::kNone) {\n json_obj.set(\"index_type\",\n ailego::JsonValue(magic_enum::enum_name(index_type).data()));\n }\n if (!omit_empty_value || metric_type != MetricType::kNone) {\n json_obj.set(\"metric_type\",\n ailego::JsonValue(magic_enum::enum_name(metric_type).data()));\n }\n if (!omit_empty_value || dimension != 0) {\n json_obj.set(\"dimension\", ailego::JsonValue(dimension));\n }\n if (!omit_empty_value || version != 0) {\n json_obj.set(\"version\", ailego::JsonValue(version));\n }\n if (!omit_empty_value || is_sparse) {\n json_obj.set(\"is_sparse\", ailego::JsonValue(is_sparse));\n }\n if (!omit_empty_value || data_type != DataType::DT_UNDEFINED) {\n json_obj.set(\"data_type\",\n ailego::JsonValue(magic_enum::enum_name(data_type).data()));\n }\n if (!omit_empty_value || use_id_map) {\n json_obj.set(\"use_id_map\", ailego::JsonValue(use_id_map));\n }\n if (!omit_empty_value || is_huge_page) {\n json_obj.set(\"is_huge_page\", ailego::JsonValue(is_huge_page));\n }\n\n // if (preprocess_param) {\n // json.set(\"preprocess_param\", preprocess_param->SerializeToJson());\n // }\n if (!omit_empty_value || quantizer_param.type != QuantizerType::kNone) {\n json_obj.set(\"quantizer_param\",\n quantizer_param.SerializeToJsonObject(omit_empty_value));\n }\n // if (refiner_param) {\n // json.set(\"refiner_param\", refiner_param->SerializeToJson());\n // }\n // if (default_query_param) {\n // json.set(\"default_query_param\",\n // default_query_param->SerializeToJson());\n // }\n return json_obj;\n}\n\n\nailego::JsonObject FlatIndexParam::SerializeToJsonObject(\n bool omit_empty_value) const {\n auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value);\n if (!omit_empty_value || major_order != IndexMeta::MajorOrder::MO_UNDEFINED) {\n json_obj.set(\"major_order\",\n ailego::JsonValue(magic_enum::enum_name(major_order).data()));\n }\n return json_obj;\n}\n\nailego::JsonObject HNSWIndexParam::SerializeToJsonObject(\n bool omit_empty_value) const {\n auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value);\n json_obj.set(\"m\", ailego::JsonValue(m));\n json_obj.set(\"ef_construction\", ailego::JsonValue(ef_construction));\n return json_obj;\n}\n\nbool BaseIndexParam::DeserializeFromJsonObject(\n const ailego::JsonObject &json_obj) {\n DESERIALIZE_ENUM_FIELD(json_obj, index_type, IndexType);\n DESERIALIZE_ENUM_FIELD(json_obj, metric_type, MetricType);\n DESERIALIZE_ENUM_FIELD(json_obj, data_type, DataType);\n\n DESERIALIZE_VALUE_FIELD(json_obj, dimension);\n DESERIALIZE_VALUE_FIELD(json_obj, version);\n DESERIALIZE_VALUE_FIELD(json_obj, is_sparse);\n DESERIALIZE_VALUE_FIELD(json_obj, use_id_map);\n DESERIALIZE_VALUE_FIELD(json_obj, is_huge_page);\n\n ailego::JsonValue tmp_json_value;\n if (json_obj.has(\"quantizer_param\")) {\n if (json_obj.get(\"quantizer_param\", &tmp_json_value);\n tmp_json_value.is_object()) {\n quantizer_param.DeserializeFromJsonObject(tmp_json_value.as_object());\n }\n }\n\n return true;\n}\n\nbool FlatIndexParam::DeserializeFromJsonObject(\n const ailego::JsonObject &json_obj) {\n if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) {\n return false;\n }\n\n if (index_type != IndexType::kFlat) {\n LOG_ERROR(\"index_type is not kFlat\");\n return false;\n }\n\n DESERIALIZE_ENUM_FIELD(json_obj, major_order, IndexMeta::MajorOrder);\n return true;\n}\n\nbool HNSWIndexParam::DeserializeFromJsonObject(\n const ailego::JsonObject &json_obj) {\n if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) {\n return false;\n }\n\n if (index_type != IndexType::kHNSW) {\n LOG_ERROR(\"index_type is not kHNSW\");\n return false;\n }\n\n DESERIALIZE_VALUE_FIELD(json_obj, m);\n DESERIALIZE_VALUE_FIELD(json_obj, ef_construction);\n\n return true;\n}\n\nbool HNSWRabitqIndexParam::DeserializeFromJsonObject(\n const ailego::JsonObject &json_obj) {\n if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) {\n return false;\n }\n\n if (index_type != IndexType::kHNSWRabitq) {\n LOG_ERROR(\"index_type is not kHNSWRabitq\");\n return false;\n }\n\n DESERIALIZE_VALUE_FIELD(json_obj, m);\n DESERIALIZE_VALUE_FIELD(json_obj, ef_construction);\n DESERIALIZE_VALUE_FIELD(json_obj, total_bits);\n DESERIALIZE_VALUE_FIELD(json_obj, num_clusters);\n DESERIALIZE_VALUE_FIELD(json_obj, sample_count);\n\n return true;\n}\n\nailego::JsonObject HNSWRabitqIndexParam::SerializeToJsonObject(\n bool omit_empty_value) const {\n auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value);\n json_obj.set(\"m\", ailego::JsonValue(m));\n json_obj.set(\"ef_construction\", ailego::JsonValue(ef_construction));\n json_obj.set(\"total_bits\", ailego::JsonValue(total_bits));\n json_obj.set(\"num_clusters\", ailego::JsonValue(num_clusters));\n if (!omit_empty_value || sample_count != 0) {\n json_obj.set(\"sample_count\", ailego::JsonValue(sample_count));\n }\n return json_obj;\n}\n\nailego::JsonObject QuantizerParam::SerializeToJsonObject(\n bool omit_empty_value) const {\n ailego::JsonObject json_obj;\n if (!omit_empty_value || type != QuantizerType::kNone) {\n json_obj.set(\"type\",\n zvec::ailego::JsonValue(magic_enum::enum_name(type).data()));\n }\n return json_obj;\n}\n\nbool QuantizerParam::DeserializeFromJsonObject(\n const ailego::JsonObject &json_obj) {\n DESERIALIZE_ENUM_FIELD(json_obj, type, QuantizerType);\n return true;\n}\n\n// bool BaseIndexQueryParam::DeserializeFromJsonObject(\n// const ailego::JsonObject &json_obj) {\n// DESERIALIZE_ENUM_FIELD(json_obj, index_type, IndexType);\n// DESERIALIZE_VALUE_FIELD(json_obj, topk);\n// DESERIALIZE_VALUE_FIELD(json_obj, fetch_vector);\n// DESERIALIZE_VALUE_FIELD(json_obj, radius);\n// DESERIALIZE_VALUE_FIELD(json_obj, is_linear);\n// return true;\n// }\n\n// ailego::JsonObject BaseIndexQueryParam::SerializeToJsonObject(\n// bool omit_empty_value) const {\n// ailego::JsonObject json_obj;\n// if (!omit_empty_value || index_type != IndexType::kNone) {\n// json_obj.set(\"index_type\",\n// ailego::JsonValue(magic_enum::enum_name(index_type).data()));\n// }\n// if (!omit_empty_value || topk != 0) {\n// json_obj.set(\"topk\", ailego::JsonValue(topk));\n// }\n// if (!omit_empty_value || fetch_vector) {\n// json_obj.set(\"fetch_vector\", ailego::JsonValue(fetch_vector));\n// }\n// if (!omit_empty_value || radius != 0.0f) {\n// json_obj.set(\"radius\", ailego::JsonValue(radius));\n// }\n// if (!omit_empty_value || is_linear) {\n// json_obj.set(\"is_linear\", ailego::JsonValue(is_linear));\n// }\n// return json_obj;\n// }\n\n// bool FlatQueryParam::DeserializeFromJsonObject(\n// const ailego::JsonObject &json_obj) {\n// if (!BaseIndexQueryParam::DeserializeFromJsonObject(json_obj)) {\n// return false;\n// }\n// if (index_type != IndexType::kFlat) {\n// LOG_ERROR(\"index_type is not kFlat\");\n// return false;\n// }\n// return true;\n// }\n\n// ailego::JsonObject FlatQueryParam::SerializeToJsonObject(\n// bool omit_empty_value) const {\n// auto json_obj =\n// BaseIndexQueryParam::SerializeToJsonObject(omit_empty_value);\n// return json_obj;\n// }\n\n// bool HNSWQueryParam::DeserializeFromJsonObject(\n// const ailego::JsonObject &json_obj) {\n// if (!BaseIndexQueryParam::DeserializeFromJsonObject(json_obj)) {\n// return false;\n// }\n// if (index_type != IndexType::kHNSW) {\n// LOG_ERROR(\"index_type is not kHNSW\");\n// return false;\n// }\n// DESERIALIZE_VALUE_FIELD(json_obj, ef_search);\n// return true;\n// }\n\n// ailego::JsonObject HNSWQueryParam::SerializeToJsonObject(\n// bool omit_empty_value) const {\n// auto json_obj =\n// BaseIndexQueryParam::SerializeToJsonObject(omit_empty_value);\n// if (!omit_empty_value || ef_search != 0) {\n// json_obj.set(\"ef_search\", ailego::JsonValue(ef_search));\n// }\n// return json_obj;\n// }\n\n\n} // namespace core_interface\n} // namespace zvec" }, { "path": "src/core/interface/indexes/flat_index.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \"algorithm/flat/flat_utility.h\"\n\nnamespace zvec::core_interface {\n\nint FlatIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) {\n param_ = dynamic_cast(param);\n\n proxima_index_params_.set(core::PARAM_FLAT_COLUMN_MAJOR_ORDER,\n param_.major_order == IndexMeta::MO_COLUMN);\n proxima_index_params_.set(core::PARAM_FLAT_USE_ID_MAP, param_.use_id_map);\n if (is_sparse_) {\n streamer_ = core::IndexFactory::CreateStreamer(\"FlatSparseStreamer\");\n } else {\n streamer_ = core::IndexFactory::CreateStreamer(\"FlatStreamer\");\n }\n\n if (ailego_unlikely(!streamer_)) {\n LOG_ERROR(\"Failed to create streamer\");\n return core::IndexError_Runtime;\n }\n if (ailego_unlikely(\n streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) {\n LOG_ERROR(\"Failed to init streamer\");\n return core::IndexError_Runtime;\n }\n return 0;\n}\n\nint FlatIndex::_prepare_for_search(\n const VectorData & /*vector_data*/,\n const BaseIndexQueryParam::Pointer &search_param,\n core::IndexContext::Pointer &context) {\n auto flat_search_param =\n std::dynamic_pointer_cast(search_param);\n\n if (ailego_unlikely(!flat_search_param)) {\n LOG_ERROR(\"Invalid search param type, expected FlatQueryParam\");\n return core::IndexError_Runtime;\n }\n\n context->set_topk(flat_search_param->topk);\n context->set_fetch_vector(flat_search_param->fetch_vector);\n if (flat_search_param->filter) {\n context->set_filter(std::move(*flat_search_param->filter));\n }\n if (flat_search_param->radius > 0.0f) {\n context->set_threshold(flat_search_param->radius);\n }\n\n return 0;\n}\n\n\n} // namespace zvec::core_interface" }, { "path": "src/core/interface/indexes/hnsw_index.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \"algorithm/hnsw/hnsw_params.h\"\n#include \"algorithm/hnsw_sparse/hnsw_sparse_params.h\"\n\nnamespace zvec::core_interface {\n\nint HNSWIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) {\n param_ = dynamic_cast(param);\n\n // valid\n param_.ef_construction = std::max(1, std::min(2048, param_.ef_construction));\n param_.m = std::max(5, std::min(1024, param_.m));\n\n if (is_sparse_) {\n proxima_index_params_.set(core::PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION,\n param_.ef_construction);\n proxima_index_params_.set(\n core::PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT, param_.m);\n\n // TODO: add_vector_with_id & fetch_by_id don't rely on this param\n proxima_index_params_.set(\n core::PARAM_HNSW_SPARSE_STREAMER_GET_VECTOR_ENABLE, true);\n\n // TODO: use index params' default query param here\n proxima_index_params_.set(core::PARAM_HNSW_SPARSE_STREAMER_EF,\n kDefaultHnswEfSearch);\n streamer_ = core::IndexFactory::CreateStreamer(\"HnswSparseStreamer\");\n\n } else {\n proxima_index_params_.set(core::PARAM_HNSW_STREAMER_EFCONSTRUCTION,\n param_.ef_construction);\n proxima_index_params_.set(core::PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT,\n param_.m);\n\n // TODO: add_vector_with_id & fetch_by_id don't rely on this param\n proxima_index_params_.set(core::PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE,\n true);\n\n // TODO: use index params' default query param here\n proxima_index_params_.set(core::PARAM_HNSW_STREAMER_EF,\n kDefaultHnswEfSearch);\n proxima_index_params_.set(core::PARAM_HNSW_STREAMER_USE_ID_MAP,\n param_.use_id_map);\n streamer_ = core::IndexFactory::CreateStreamer(\"HnswStreamer\");\n }\n\n if (ailego_unlikely(!streamer_)) {\n LOG_ERROR(\"Failed to create streamer\");\n return core::IndexError_Runtime;\n }\n if (ailego_unlikely(\n streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) {\n LOG_ERROR(\"Failed to init streamer\");\n return core::IndexError_Runtime;\n }\n return 0;\n}\n\n\nint HNSWIndex::_prepare_for_search(\n const VectorData & /*vector_data*/,\n const BaseIndexQueryParam::Pointer &search_param,\n core::IndexContext::Pointer &context) {\n const auto &hnsw_search_param =\n std::dynamic_pointer_cast(search_param);\n\n if (ailego_unlikely(!hnsw_search_param)) {\n LOG_ERROR(\"Invalid search param type, expected HNSWQueryParam\");\n return core::IndexError_Runtime;\n }\n\n if (0 >= hnsw_search_param->ef_search ||\n hnsw_search_param->ef_search > 2048) {\n LOG_ERROR(\n \"ef_search must be greater than 0 and less than or equal to 2048.\");\n return core::IndexError_Runtime;\n }\n\n context->set_topk(hnsw_search_param->topk);\n context->set_fetch_vector(hnsw_search_param->fetch_vector);\n if (hnsw_search_param->filter) {\n context->set_filter(std::move(*hnsw_search_param->filter));\n }\n if (hnsw_search_param->radius > 0.0f) {\n context->set_threshold(hnsw_search_param->radius);\n }\n ailego::Params params;\n const int real_search_ef =\n std::max(1u, std::min(2048u, hnsw_search_param->ef_search));\n params.set(core::PARAM_HNSW_STREAMER_EF, real_search_ef);\n context->update(params);\n return 0;\n}\n\nint HNSWIndex::_get_coarse_search_topk(\n const BaseIndexQueryParam::Pointer &search_param) {\n const auto &hnsw_search_param =\n std::dynamic_pointer_cast(search_param);\n\n // scale_factor doesn't take effect for hnsw.\n auto ret = std::max(search_param->topk, hnsw_search_param->ef_search);\n return ret;\n}\n\n} // namespace zvec::core_interface" }, { "path": "src/core/interface/indexes/hnsw_rabitq_index.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \"zvec/core/framework/index_error.h\"\n\n#if RABITQ_SUPPORTED\n#include \"algorithm/hnsw_rabitq/hnsw_rabitq_params.h\"\n#include \"algorithm/hnsw_rabitq/hnsw_rabitq_streamer.h\"\n#include \"algorithm/hnsw_rabitq/rabitq_params.h\"\n#endif\n\nnamespace zvec::core_interface {\n\nint HNSWRabitqIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) {\n#if !RABITQ_SUPPORTED\n LOG_ERROR(\"RaBitQ is not supported on this platform (Linux x86_64 only)\");\n return core::IndexError_Unsupported;\n#else\n param_ = dynamic_cast(param);\n\n if (is_sparse_) {\n LOG_ERROR(\"Sparse index is not supported\");\n return core::IndexError_Runtime;\n }\n\n if (param.dimension < core::kMinRabitqDimSize ||\n param.dimension > core::kMaxRabitqDimSize) {\n LOG_ERROR(\"Unsupported dimension: %d\", param.dimension);\n return core::IndexError_Unsupported;\n }\n\n // validate parameters\n param_.ef_construction = std::max(1, std::min(2048, param_.ef_construction));\n param_.m = std::max(5, std::min(1024, param_.m));\n\n proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_EFCONSTRUCTION,\n param_.ef_construction);\n proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_MAX_NEIGHBOR_COUNT,\n param_.m);\n proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_GET_VECTOR_ENABLE,\n true);\n proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_EF,\n kDefaultHnswEfSearch);\n proxima_index_params_.set(core::PARAM_HNSW_RABITQ_STREAMER_USE_ID_MAP,\n param_.use_id_map);\n proxima_index_params_.set(core::PARAM_HNSW_RABITQ_GENERAL_DIMENSION,\n input_vector_meta_.dimension());\n proxima_index_params_.set(core::PARAM_RABITQ_TOTAL_BITS, param_.total_bits);\n // num_clusters, sample_count are parameters for rabitq converter\n // proxima_index_params_.set(core::PARAM_RABITQ_NUM_CLUSTERS,\n // param_.num_clusters);\n\n auto streamer = std::make_shared();\n streamer->set_provider(param_.provider);\n streamer->set_reformer(param_.reformer);\n streamer_ = streamer;\n\n if (ailego_unlikely(!streamer_)) {\n LOG_ERROR(\"Failed to create HnswRabitqStreamer\");\n return core::IndexError_Runtime;\n }\n if (ailego_unlikely(\n streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) {\n LOG_ERROR(\"Failed to init HnswRabitqStreamer\");\n return core::IndexError_Runtime;\n }\n return 0;\n#endif // RABITQ_SUPPORTED\n}\n\nint HNSWRabitqIndex::_prepare_for_search(\n const VectorData & /*vector_data*/,\n const BaseIndexQueryParam::Pointer &search_param,\n core::IndexContext::Pointer &context) {\n#if !RABITQ_SUPPORTED\n LOG_ERROR(\"RaBitQ is not supported on this platform (Linux x86_64 only)\");\n return core::IndexError_Unsupported;\n#else\n const auto &hnsw_search_param =\n std::dynamic_pointer_cast(search_param);\n\n if (ailego_unlikely(!hnsw_search_param)) {\n LOG_ERROR(\"Invalid search param type, expected HNSWRabitqQueryParam\");\n return core::IndexError_Runtime;\n }\n\n if (0 >= hnsw_search_param->ef_search ||\n hnsw_search_param->ef_search > 2048) {\n LOG_ERROR(\n \"ef_search must be greater than 0 and less than or equal to 2048.\");\n return core::IndexError_Runtime;\n }\n\n context->set_topk(hnsw_search_param->topk);\n context->set_fetch_vector(hnsw_search_param->fetch_vector);\n if (hnsw_search_param->filter) {\n context->set_filter(std::move(*hnsw_search_param->filter));\n }\n if (hnsw_search_param->radius > 0.0f) {\n context->set_threshold(hnsw_search_param->radius);\n }\n ailego::Params params;\n const int real_search_ef =\n std::max(1u, std::min(2048u, hnsw_search_param->ef_search));\n params.set(core::PARAM_HNSW_RABITQ_STREAMER_EF, real_search_ef);\n context->update(params);\n return 0;\n#endif // RABITQ_SUPPORTED\n}\n\nint HNSWRabitqIndex::_get_coarse_search_topk(\n const BaseIndexQueryParam::Pointer &search_param) {\n#if !RABITQ_SUPPORTED\n LOG_ERROR(\"RaBitQ is not supported on this platform (Linux x86_64 only)\");\n return -1;\n#else\n const auto &hnsw_search_param =\n std::dynamic_pointer_cast(search_param);\n\n auto ret = std::max(search_param->topk, hnsw_search_param->ef_search);\n return ret;\n#endif // RABITQ_SUPPORTED\n}\n\n\n} // namespace zvec::core_interface\n" }, { "path": "src/core/interface/indexes/ivf_index.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \"algorithm/ivf/ivf_params.h\"\n\nnamespace zvec::core_interface {\n\nstatic constexpr uint64_t kInvalidKey = std::numeric_limits::max();\n\nint IVFIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) {\n if (is_sparse_) {\n LOG_ERROR(\"IVF Index not support sparse vector\");\n return core::IndexError_InvalidArgument;\n }\n\n param_ = dynamic_cast(param);\n param_.nlist = std::max(1, std::min(1024, param_.nlist));\n param_.niters = std::max(1, std::min(1024, param_.niters));\n\n proxima_index_params_.set(core::PARAM_IVF_BUILDER_CENTROID_COUNT,\n param_.nlist);\n\n // TODO: add_vector_with_id & fetch_by_id don't rely on this param\n builder_ = core::IndexFactory::CreateBuilder(\"IVFBuilder\");\n streamer_ = core::IndexFactory::CreateStreamer(\"IVFStreamer\");\n\n if (ailego_unlikely(!builder_)) {\n LOG_ERROR(\"Failed to create builder\");\n return core::IndexError_Runtime;\n }\n if (ailego_unlikely(!streamer_)) {\n LOG_ERROR(\"Failed to create streamer\");\n return core::IndexError_Runtime;\n }\n IndexMeta real_meta;\n if (converter_) {\n real_meta = converter_->meta();\n } else {\n real_meta = proxima_index_meta_;\n }\n if (ailego_unlikely(builder_->init(real_meta, proxima_index_params_) != 0)) {\n LOG_ERROR(\"Failed to init builder\");\n return core::IndexError_Runtime;\n }\n if (ailego_unlikely(streamer_->init(real_meta, proxima_index_params_) != 0)) {\n LOG_ERROR(\"Failed to init streamer\");\n return core::IndexError_Runtime;\n }\n return 0;\n}\n\nint IVFIndex::Open(const std::string &file_path,\n StorageOptions storage_options) {\n ailego::Params storage_params;\n file_path_ = file_path;\n is_read_only_ = storage_options.read_only;\n switch (storage_options.type) {\n case StorageOptions::StorageType::kMMAP: {\n storage_ = core::IndexFactory::CreateStorage(\"MMapFileReadStorage\");\n if (storage_ == nullptr) {\n LOG_ERROR(\"Failed to create MMapFileStorage\");\n return core::IndexError_Runtime;\n }\n int ret = storage_->init(storage_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to init MMapFileStorage, path: %s, err: %s\",\n file_path_.c_str(), core::IndexError::What(ret));\n return ret;\n }\n break;\n }\n case StorageOptions::StorageType::kBufferPool: {\n storage_ = core::IndexFactory::CreateStorage(\"BufferStorage\");\n if (storage_ == nullptr) {\n LOG_ERROR(\"Failed to create BufferStorage\");\n return core::IndexError_Runtime;\n }\n int ret = storage_->init(storage_params);\n if (ret != 0) {\n LOG_ERROR(\"Failed to init BufferStorage, path: %s, err: %s\",\n file_path_.c_str(), core::IndexError::What(ret));\n return ret;\n }\n break;\n }\n default: {\n LOG_ERROR(\"Unsupported storage type\");\n return core::IndexError_Unsupported;\n }\n }\n\n if (is_read_only_ || !storage_options.create_new) {\n // read_options.create_new\n int ret = storage_->open(file_path_, false);\n if (ret != 0) {\n LOG_ERROR(\"Failed to open storage, path: %s, err: %s\", file_path_.c_str(),\n core::IndexError::What(ret));\n return core::IndexError_Runtime;\n }\n if (streamer_ == nullptr || streamer_->open(storage_) != 0) {\n LOG_ERROR(\"Failed to open streamer, path: %s\", file_path_.c_str());\n return core::IndexError_Runtime;\n }\n is_trained_ = true;\n }\n is_open_ = true;\n return 0;\n}\n\nint IVFIndex::GenerateHolder() {\n if (param_.data_type == DataType::DT_FP16) {\n auto holder =\n std::make_shared>(\n param_.dimension);\n for (auto doc : doc_cache_) {\n ailego::NumericalVector vec(doc.second);\n if (doc.first == kInvalidKey) {\n continue;\n }\n if (!holder->emplace(doc.first, vec)) {\n LOG_ERROR(\"Failed to add vector\");\n return core::IndexError_Runtime;\n }\n }\n holder_ = holder;\n } else if (param_.data_type == DataType::DT_FP32) {\n auto holder =\n std::make_shared>(\n param_.dimension);\n for (auto doc : doc_cache_) {\n ailego::NumericalVector vec(doc.second);\n if (doc.first == kInvalidKey) {\n continue;\n }\n if (!holder->emplace(doc.first, vec)) {\n LOG_ERROR(\"Failed to add vector\");\n return core::IndexError_Runtime;\n }\n }\n holder_ = holder;\n } else if (param_.data_type == DataType::DT_INT8) {\n auto holder =\n std::make_shared>(\n param_.dimension);\n for (auto doc : doc_cache_) {\n ailego::NumericalVector vec(doc.second);\n if (doc.first == kInvalidKey) {\n continue;\n }\n if (!holder->emplace(doc.first, vec)) {\n LOG_ERROR(\"Failed to add vector\");\n return core::IndexError_Runtime;\n }\n }\n holder_ = holder;\n } else {\n LOG_ERROR(\"data_type is not support\");\n return core::IndexError_Runtime;\n }\n if (converter_) {\n core::IndexConverter::TrainAndTransform(converter_, holder_);\n holder_ = converter_->result();\n }\n return 0;\n}\n\nint IVFIndex::Add(const VectorData &vector, uint32_t doc_id) {\n if (is_trained_) {\n LOG_ERROR(\"this IVF index is trained\");\n return core::IndexError_Runtime;\n }\n if (!std::holds_alternative(vector.vector)) {\n LOG_ERROR(\"Invalid vector data\");\n return core::IndexError_Runtime;\n }\n const DenseVector &dense_vector = std::get(vector.vector);\n std::string out_vector_buffer = std::string(\n static_cast(dense_vector.data),\n input_vector_meta_.dimension() * input_vector_meta_.unit_size());\n\n std::lock_guard lock(mutex_);\n while (doc_cache_.size() <= doc_id) {\n std::string fake_data(\n input_vector_meta_.dimension() * input_vector_meta_.unit_size(), 0);\n doc_cache_.push_back(std::make_pair(kInvalidKey, fake_data));\n }\n doc_cache_[doc_id] = std::make_pair(doc_id, out_vector_buffer);\n return 0;\n}\n\nint IVFIndex::Train() {\n GenerateHolder();\n builder_->train(holder_);\n builder_->build(holder_);\n auto dumper = core::IndexFactory::CreateDumper(\"FileDumper\");\n\n dumper->create(file_path_);\n builder_->dump(dumper);\n dumper->close();\n int ret = storage_->open(file_path_, false);\n if (ret != 0) {\n LOG_ERROR(\"Failed to open storage, path: %s, err: %s\", file_path_.c_str(),\n core::IndexError::What(ret));\n return core::IndexError_Runtime;\n }\n if (streamer_ == nullptr || streamer_->open(storage_) != 0) {\n LOG_ERROR(\"Failed to open streamer, path: %s\", file_path_.c_str());\n return core::IndexError_Runtime;\n }\n is_trained_ = true;\n return 0;\n}\n\nint IVFIndex::_dense_fetch(const uint32_t doc_id,\n VectorDataBuffer *vector_data_buffer) {\n if (is_trained_) {\n return Index::_dense_fetch(doc_id, vector_data_buffer);\n } else {\n DenseVectorBuffer dense_vector_buffer;\n std::string &out_vector_buffer = dense_vector_buffer.data;\n out_vector_buffer = doc_cache_[doc_id].second;\n vector_data_buffer->vector_buffer = std::move(dense_vector_buffer);\n return 0;\n }\n}\n\nint IVFIndex::_prepare_for_search(\n const VectorData & /*query*/,\n const BaseIndexQueryParam::Pointer &search_param,\n core::IndexContext::Pointer &context) {\n const auto &ivf_search_param =\n std::dynamic_pointer_cast(search_param);\n\n context->set_topk(ivf_search_param->topk);\n context->set_fetch_vector(ivf_search_param->fetch_vector);\n if (ivf_search_param->filter) {\n context->set_filter(std::move(*ivf_search_param->filter));\n }\n if (ivf_search_param->radius > 0.0f) {\n context->set_threshold(ivf_search_param->radius);\n }\n\n if (ivf_search_param->nprobe > 0) {\n // TODO: 1. sparse; 2. default ef\n ailego::Params params;\n // need fix\n params.set(core::PARAM_IVF_BUILDER_CENTROID_COUNT,\n ivf_search_param->nprobe);\n context->update(params);\n }\n return 0;\n}\n\nint IVFIndex::Merge(const std::vector &indexes,\n const IndexFilter &filter, const MergeOptions &options) {\n int pre_ret = Index::Merge(indexes, filter, options);\n if (pre_ret != 0) {\n return pre_ret;\n }\n auto dumper = core::IndexFactory::CreateDumper(\"FileDumper\");\n\n dumper->create(file_path_);\n builder_->dump(dumper);\n dumper->close();\n int ret = storage_->open(file_path_, false);\n if (ret != 0) {\n LOG_ERROR(\"Failed to open storage, path: %s, err: %s\", file_path_.c_str(),\n core::IndexError::What(ret));\n return core::IndexError_Runtime;\n }\n if (streamer_ == nullptr || streamer_->open(storage_) != 0) {\n LOG_ERROR(\"Failed to open streamer, path: %s\", file_path_.c_str());\n return core::IndexError_Runtime;\n }\n is_trained_ = true;\n return 0;\n}\n} // namespace zvec::core_interface" }, { "path": "src/core/interface/utils/utils.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#pragma once\n\n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core_interface {\n\ntemplate \nconstexpr bool extract_enum_from_json(const ailego::JsonObject &json_obj,\n const char *key, EnumType &enum_value,\n ailego::JsonValue &tmp_json_value) {\n if (json_obj.has(key)) {\n if (json_obj.get(key, &tmp_json_value); tmp_json_value.is_string()) {\n auto optional_enum_value =\n magic_enum::enum_cast(tmp_json_value.as_stl_string());\n if (optional_enum_value.has_value()) {\n enum_value = optional_enum_value.value();\n } else {\n LOG_ERROR(\"Invalid enum value for key: %s, value: %s\", key,\n tmp_json_value.as_c_string());\n return false;\n }\n } else {\n LOG_ERROR(\"Invalid json field type for key: %s\", key);\n return false;\n }\n }\n return true;\n}\n\ntemplate \nconstexpr bool extract_value_from_json(const ailego::JsonObject &json_obj,\n const char *key, T &value,\n ailego::JsonValue &tmp_json_value) {\n if (json_obj.has(key)) {\n json_obj.get(key, &tmp_json_value);\n if constexpr (std::is_same_v) {\n if (tmp_json_value.is_boolean()) {\n value = tmp_json_value.as_bool();\n } else {\n LOG_ERROR(\"Invalid json field type for key: %s; expected: boolean\",\n key);\n return false;\n }\n } else if constexpr (std::is_floating_point_v) {\n if (tmp_json_value.is_float() || tmp_json_value.is_integer()) {\n value = static_cast(tmp_json_value.as_float());\n } else {\n LOG_ERROR(\"Invalid json field type for key: %s; expected: float\", key);\n return false;\n }\n } else if constexpr (std::is_integral_v) {\n if (tmp_json_value.is_integer()) {\n value = static_cast(tmp_json_value.as_integer());\n } else {\n LOG_ERROR(\"Invalid json field type for key: %s; expected: integer\",\n key);\n return false;\n }\n } else {\n abort();\n }\n }\n return true;\n}\n\n#define DESERIALIZE_ENUM_FIELD(json_obj, field_name, EnumType) \\\n { \\\n ailego::JsonValue tmp_json_value; \\\n if (!extract_enum_from_json(json_obj, #field_name, field_name, \\\n tmp_json_value)) { \\\n LOG_ERROR(\"Error when deserialize json - field:%s\", #field_name); \\\n return false; \\\n } \\\n }\n\n\n#define DESERIALIZE_VALUE_FIELD(json_obj, field_name) \\\n { \\\n ailego::JsonValue tmp_json_value; \\\n if (!extract_value_from_json(json_obj, #field_name, field_name, \\\n tmp_json_value)) { \\\n LOG_ERROR(\"Error when deserialize json - field:%s\", #field_name); \\\n return false; \\\n } \\\n }\n} // namespace core_interface\n} // namespace zvec" }, { "path": "src/core/metric/CMakeLists.txt", "content": "include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)\ninclude(${PROJECT_ROOT_DIR}/cmake/option.cmake)\n\ncc_library(\n NAME core_metric \n STATIC SHARED STRICT ALWAYS_LINK\n SRCS *.cc\n LIBS zvec_ailego zvec_turbo core_framework \n INCS . ${PROJECT_ROOT_DIR}/src/core\n VERSION \"${PROXIMA_ZVEC_VERSION}\"\n )\n" }, { "path": "src/core/metric/cosine_metric.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\n//! Retrieve distance function for index features\ninline IndexMetric::MatrixDistanceHandle CosineDistanceMatrixFp32(size_t m,\n size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n//! Retrieve distance function for index features\ninline IndexMetric::MatrixDistanceHandle CosineDistanceMatrixFp16(size_t m,\n size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n/*! Cosine Metric\n */\nclass CosineMetric : public IndexMetric {\n public:\n //! Initialize Metric\n int init(const IndexMeta &meta, const ailego::Params &index_params) override {\n IndexMeta::DataType ft = meta.data_type();\n if (ft != IndexMeta::DataType::DT_FP16 &&\n ft != IndexMeta::DataType::DT_FP32) {\n return IndexError_Unsupported;\n }\n if (IndexMeta::UnitSizeof(ft) != meta.unit_size()) {\n return IndexError_Unsupported;\n }\n data_type_ = ft;\n params_ = index_params;\n\n return 0;\n }\n\n //! Cleanup Metric\n int cleanup(void) override {\n return 0;\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta) const override {\n return (meta.data_type() == data_type_ &&\n meta.unit_size() == IndexMeta::UnitSizeof(data_type_));\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta,\n const IndexQueryMeta &qmeta) const override {\n return (qmeta.data_type() == data_type_ &&\n qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) &&\n qmeta.dimension() == meta.dimension());\n }\n\n //! Retrieve distance function for query\n MatrixDistance distance(void) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP16:\n return reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute);\n case IndexMeta::DataType::DT_FP32:\n return reinterpret_cast(\n ailego::CosineDistanceMatrix::Compute);\n\n default:\n return nullptr;\n }\n }\n\n //! Retrieve distance function for index features\n MatrixDistance distance_matrix(size_t m, size_t n) const override {\n if (m != 1 || n != 1) {\n return nullptr;\n }\n return distance();\n }\n\n //! Retrieve distance function for query\n MatrixBatchDistance batch_distance(void) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n case IndexMeta::DataType::DT_FP16:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n default:\n return nullptr;\n }\n }\n\n //! Retrieve params of Metric\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n //! Retrieve query metric object of this index metric\n Pointer query_metric(void) const override {\n return nullptr;\n }\n\n private:\n IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32};\n ailego::Params params_{};\n};\n\nINDEX_FACTORY_REGISTER_METRIC_ALIAS(Cosine, CosineMetric);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/metric/euclidean_metric.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\n//! Retrieve distance function for index features\nstatic inline IndexMetric::MatrixDistanceHandle\nSquaredEuclideanDistanceMatrixFp32(size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n//! Retrieve distance function for index features\nstatic inline IndexMetric::MatrixDistanceHandle\nSquaredEuclideanDistanceMatrixFp16(size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\nstatic inline IndexMetric::MatrixDistanceHandle\nSquaredEuclideanDistanceMatrixInt8(size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n//! Retrieve distance function for index features in Int4\nstatic inline IndexMetric::MatrixDistanceHandle\nSquaredEuclideanDistanceMatrixInt4(size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n//! Retrieve distance function for index features\nstatic inline IndexMetric::MatrixDistanceHandle EuclideanDistanceMatrixFp32(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n//! Retrieve distance function for index features\nstatic inline IndexMetric::MatrixDistanceHandle EuclideanDistanceMatrixFp16(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\nstatic inline IndexMetric::MatrixDistanceHandle EuclideanDistanceMatrixInt8(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n//! Retrieve distance function for index features in Int4\nstatic inline IndexMetric::MatrixDistanceHandle EuclideanDistanceMatrixInt4(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n//! Retrieve distance function for index features\nstatic inline IndexMetric::MatrixDistanceHandle HammingDistanceMatrix32(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n#if defined(AILEGO_M64)\nstatic inline IndexMetric::MatrixDistanceHandle HammingDistanceMatrix64(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n#endif // AILEGO_M64\n\n//! Retrieve distance function for index features\nstatic inline IndexMetric::MatrixDistanceHandle\nHammingSquareRootDistanceMatrix32(size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n#if defined(AILEGO_M64)\nstatic inline IndexMetric::MatrixDistanceHandle\nHammingSquareRootDistanceMatrix64(size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n#endif // AILEGO_M64\n\n/*! Squared Euclidean Distance Metric\n */\nclass SquaredEuclideanMetric : public IndexMetric {\n public:\n //! Initialize Metric\n int init(const IndexMeta &meta, const ailego::Params &index_params) override {\n IndexMeta::DataType dt = meta.data_type();\n if (dt != IndexMeta::DataType::DT_FP16 &&\n dt != IndexMeta::DataType::DT_FP32 &&\n dt != IndexMeta::DataType::DT_INT8 &&\n dt != IndexMeta::DataType::DT_INT4 &&\n dt != IndexMeta::DataType::DT_BINARY32 &&\n dt != IndexMeta::DataType::DT_BINARY64) {\n return IndexError_Unsupported;\n }\n if (IndexMeta::UnitSizeof(dt) != meta.unit_size()) {\n return IndexError_Unsupported;\n }\n data_type_ = dt;\n params_ = index_params;\n\n return 0;\n }\n\n //! Cleanup Metric\n int cleanup(void) override {\n return 0;\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta) const override {\n return (meta.data_type() == data_type_ &&\n meta.unit_size() == IndexMeta::UnitSizeof(data_type_));\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta,\n const IndexQueryMeta &qmeta) const override {\n return (qmeta.data_type() == data_type_ &&\n qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) &&\n qmeta.dimension() == meta.dimension());\n }\n\n //! Retrieve distance function for query\n MatrixDistance distance(void) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_BINARY32:\n return reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute);\n\n#if defined(AILEGO_M64)\n case IndexMeta::DataType::DT_BINARY64:\n return reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute);\n#endif // AILEGO_M64\n\n case IndexMeta::DataType::DT_FP16:\n return reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute);\n\n case IndexMeta::DataType::DT_FP32:\n return reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute);\n\n case IndexMeta::DataType::DT_INT8:\n return reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute);\n\n case IndexMeta::DataType::DT_INT4:\n return reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute);\n\n default:\n return nullptr;\n }\n }\n\n //! Retrieve sparse distance function for query\n MatrixSparseDistance sparse_distance(void) const override {\n return reinterpret_cast(\n ailego::SquaredEuclideanSparseDistanceMatrix::Compute);\n }\n\n //! Retrieve distance function for index features\n MatrixDistance distance_matrix(size_t m, size_t n) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_BINARY32:\n return HammingDistanceMatrix32(m, n);\n\n#if defined(AILEGO_M64)\n case IndexMeta::DataType::DT_BINARY64:\n return HammingDistanceMatrix64(m, n);\n#endif // AILEGO_M64\n\n case IndexMeta::DataType::DT_FP16:\n return SquaredEuclideanDistanceMatrixFp16(m, n);\n\n case IndexMeta::DataType::DT_FP32:\n return SquaredEuclideanDistanceMatrixFp32(m, n);\n\n case IndexMeta::DataType::DT_INT8:\n return SquaredEuclideanDistanceMatrixInt8(m, n);\n\n case IndexMeta::DataType::DT_INT4:\n return SquaredEuclideanDistanceMatrixInt4(m, n);\n\n default:\n return nullptr;\n }\n }\n\n //! Retrieve distance function for query\n MatrixBatchDistance batch_distance(void) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_BINARY32:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n\n#if defined(AILEGO_M64)\n case IndexMeta::DataType::DT_BINARY64:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n#endif // AILEGO_M64\n\n case IndexMeta::DataType::DT_FP16:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n\n case IndexMeta::DataType::DT_FP32:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n\n case IndexMeta::DataType::DT_INT8:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n\n case IndexMeta::DataType::DT_INT4:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n\n default:\n return nullptr;\n }\n }\n\n //! Retrieve params of Metric\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n //! Retrieve query metric object of this index metric\n Pointer query_metric(void) const override {\n return nullptr;\n }\n\n private:\n IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32};\n ailego::Params params_{};\n};\n\n/*! Euclidean Distance Metric\n */\nclass EuclideanMetric : public IndexMetric {\n public:\n //! Initialize Metric\n int init(const IndexMeta &meta, const ailego::Params &index_params) override {\n IndexMeta::DataType dt = meta.data_type();\n if (dt != IndexMeta::DataType::DT_FP16 &&\n dt != IndexMeta::DataType::DT_FP32 &&\n dt != IndexMeta::DataType::DT_INT8 &&\n dt != IndexMeta::DataType::DT_INT4 &&\n dt != IndexMeta::DataType::DT_BINARY32 &&\n dt != IndexMeta::DataType::DT_BINARY64) {\n return IndexError_Unsupported;\n }\n if (IndexMeta::UnitSizeof(dt) != meta.unit_size()) {\n return IndexError_Unsupported;\n }\n data_type_ = dt;\n params_ = index_params;\n return 0;\n }\n\n //! Cleanup Metric\n int cleanup(void) override {\n return 0;\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta) const override {\n return (meta.data_type() == data_type_ &&\n meta.unit_size() == IndexMeta::UnitSizeof(data_type_));\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta,\n const IndexQueryMeta &qmeta) const override {\n return (qmeta.data_type() == data_type_ &&\n qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) &&\n qmeta.dimension() == meta.dimension());\n }\n\n //! Retrieve distance function for query\n MatrixDistance distance(void) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_BINARY32:\n return reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute);\n\n#if defined(AILEGO_M64)\n case IndexMeta::DataType::DT_BINARY64:\n return reinterpret_cast(\n ailego::HammingSquareRootDistanceMatrix::Compute);\n#endif // AILEGO_M64\n\n case IndexMeta::DataType::DT_FP16:\n return reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute);\n\n case IndexMeta::DataType::DT_FP32:\n return reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute);\n\n case IndexMeta::DataType::DT_INT8:\n return reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute);\n\n case IndexMeta::DataType::DT_INT4:\n return reinterpret_cast(\n ailego::EuclideanDistanceMatrix::Compute);\n\n default:\n return nullptr;\n }\n }\n\n //! Retrieve distance function for index features\n MatrixDistance distance_matrix(size_t m, size_t n) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_BINARY32:\n return HammingSquareRootDistanceMatrix32(m, n);\n\n#if defined(AILEGO_M64)\n case IndexMeta::DataType::DT_BINARY64:\n return HammingSquareRootDistanceMatrix64(m, n);\n#endif // AILEGO_M64\n\n case IndexMeta::DataType::DT_FP16:\n return EuclideanDistanceMatrixFp16(m, n);\n\n case IndexMeta::DataType::DT_FP32:\n return EuclideanDistanceMatrixFp32(m, n);\n\n case IndexMeta::DataType::DT_INT8:\n return EuclideanDistanceMatrixInt8(m, n);\n\n case IndexMeta::DataType::DT_INT4:\n return EuclideanDistanceMatrixInt4(m, n);\n\n default:\n return nullptr;\n }\n }\n\n //! Retrieve params of Metric\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n //! Retrieve query metric object of this index metric\n Pointer query_metric(void) const override {\n return nullptr;\n }\n\n private:\n IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32};\n ailego::Params params_{};\n};\n\n/*! Squared Euclidean Sparse Metric\n */\nclass SquaredEuclideanSparseMetric : public IndexMetric {\n public:\n //! Initialize Metric\n int init(const IndexMeta &meta, const ailego::Params &index_params) override {\n IndexMeta::DataType data_type = meta.data_type();\n if (data_type != IndexMeta::DataType::DT_FP16 &&\n data_type != IndexMeta::DataType::DT_FP32) {\n return IndexError_Unsupported;\n }\n\n if (IndexMeta::UnitSizeof(data_type) != meta.unit_size()) {\n return IndexError_Unsupported;\n }\n\n data_type_ = data_type;\n params_ = index_params;\n\n return 0;\n }\n\n //! Cleanup Metric\n int cleanup(void) override {\n return 0;\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta) const override {\n return (meta.data_type() == data_type_ &&\n meta.unit_size() == IndexMeta::UnitSizeof(data_type_));\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta,\n const IndexQueryMeta &qmeta) const override {\n return (qmeta.data_type() == data_type_ &&\n qmeta.data_type() == meta.data_type() &&\n qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) &&\n qmeta.unit_size() == meta.unit_size());\n }\n\n //! Retrieve sparse distance function for query\n MatrixSparseDistance sparse_distance(void) const override {\n return reinterpret_cast(\n ailego::SquaredEuclideanSparseDistanceMatrix::Compute);\n }\n\n //! Retrieve params of Metric\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n //! Retrieve query metric object of this index metric\n Pointer query_metric(void) const override {\n return nullptr;\n }\n\n private:\n IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32};\n\n ailego::Params params_{};\n};\n\nINDEX_FACTORY_REGISTER_METRIC_ALIAS(SquaredEuclidean, SquaredEuclideanMetric);\nINDEX_FACTORY_REGISTER_METRIC_ALIAS(Euclidean, EuclideanMetric);\n\nINDEX_FACTORY_REGISTER_METRIC_ALIAS(SquaredEuclideanSparse,\n SquaredEuclideanSparseMetric);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/metric/hamming_metric.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \"ailego/math_batch/distance_batch.h\"\n\nnamespace zvec {\nnamespace core {\n\n//! Retrieve distance function for index features\nstatic inline IndexMetric::MatrixDistanceHandle HammingDistanceMatrix32(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n#if defined(AILEGO_M64)\nstatic inline IndexMetric::MatrixDistanceHandle HammingDistanceMatrix64(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute),\n reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n#endif // AILEOG_M64\n\n/*! Hamming Metric\n */\nclass HammingMetric : public IndexMetric {\n public:\n //! Initialize Metric\n int init(const IndexMeta &meta, const ailego::Params &index_params) override {\n if (meta.data_type() != IndexMeta::DataType::DT_BINARY32 &&\n meta.data_type() != IndexMeta::DataType::DT_BINARY64) {\n return IndexError_Unsupported;\n }\n if (IndexMeta::UnitSizeof(meta.data_type()) != meta.unit_size()) {\n return IndexError_Unsupported;\n }\n feature_type_ = meta.data_type();\n params_ = index_params;\n return 0;\n }\n\n //! Cleanup Metric\n int cleanup(void) override {\n return 0;\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta) const override {\n return (meta.data_type() == feature_type_ &&\n meta.unit_size() == IndexMeta::UnitSizeof(feature_type_));\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta,\n const IndexQueryMeta &qmeta) const override {\n return (qmeta.data_type() == feature_type_ &&\n qmeta.unit_size() == IndexMeta::UnitSizeof(feature_type_) &&\n qmeta.dimension() == meta.dimension());\n }\n\n //! Retrieve distance function for query\n MatrixDistance distance(void) const override {\n#if defined(AILEGO_M64)\n if (feature_type_ == IndexMeta::DataType::DT_BINARY64) {\n return reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute);\n }\n#endif\n if (feature_type_ == IndexMeta::DataType::DT_BINARY32) {\n return reinterpret_cast(\n ailego::HammingDistanceMatrix::Compute);\n }\n return nullptr;\n }\n\n MatrixBatchDistance batch_distance(void) const override {\n#if defined(AILEGO_M64)\n if (feature_type_ == IndexMeta::DataType::DT_BINARY64) {\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n }\n#endif\n if (feature_type_ == IndexMeta::DataType::DT_BINARY32) {\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n }\n return nullptr;\n }\n\n //! Retrieve distance function for index features\n MatrixDistance distance_matrix(size_t m, size_t n) const override {\n#if defined(AILEGO_M64)\n if (feature_type_ == IndexMeta::DataType::DT_BINARY64) {\n return HammingDistanceMatrix64(m, n);\n }\n#endif\n if (feature_type_ == IndexMeta::DataType::DT_BINARY32) {\n return HammingDistanceMatrix32(m, n);\n }\n return nullptr;\n }\n\n //! Retrieve params of Metric\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n //! Retrieve query metric object of this index metric\n Pointer query_metric(void) const override {\n return nullptr;\n }\n\n private:\n IndexMeta::DataType feature_type_{IndexMeta::DataType::DT_BINARY32};\n ailego::Params params_{};\n};\n\nINDEX_FACTORY_REGISTER_METRIC_ALIAS(Hamming, HammingMetric);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/metric/inner_product_metric.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n#include \n#include \n#include \n#include \n\nnamespace zvec {\nnamespace core {\n\n//! Retrieve distance function for index features\nstatic inline IndexMetric::MatrixDistanceHandle MinusInnerProductMatrixFp32(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n//! Retrieve distance function for index features\nstatic inline IndexMetric::MatrixDistanceHandle MinusInnerProductMatrixFp16(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\nstatic inline IndexMetric::MatrixDistanceHandle MinusInnerProductMatrixInt8(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n//! Retrieve distance function for index features in Int4\nstatic inline IndexMetric::MatrixDistanceHandle MinusInnerProductMatrixInt4(\n size_t m, size_t n) {\n static const IndexMetric::MatrixDistanceHandle distance_table[6][6] = {\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr, nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n nullptr},\n {reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute),\n reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute)},\n };\n if (m > 32 || n > 32 || ailego_popcount(m) != 1 || ailego_popcount(n) != 1) {\n return nullptr;\n }\n return distance_table[ailego_ctz(m)][ailego_ctz(n)];\n}\n\n/*! Inner Product Metric\n */\nclass InnerProductMetric : public IndexMetric {\n public:\n //! Initialize Metric\n int init(const IndexMeta &meta, const ailego::Params &index_params) override {\n IndexMeta::MetaType mt = meta.meta_type();\n if (mt != IndexMeta::MetaType::MT_DENSE) {\n return IndexError_Unsupported;\n }\n\n IndexMeta::DataType dt = meta.data_type();\n if (dt != IndexMeta::DataType::DT_FP16 &&\n dt != IndexMeta::DataType::DT_FP32 &&\n dt != IndexMeta::DataType::DT_INT8 &&\n dt != IndexMeta::DataType::DT_INT4) {\n return IndexError_Unsupported;\n }\n if (IndexMeta::UnitSizeof(dt) != meta.unit_size()) {\n return IndexError_Unsupported;\n }\n\n meta_type_ = mt;\n data_type_ = dt;\n params_ = index_params;\n\n return 0;\n }\n\n //! Cleanup Metric\n int cleanup(void) override {\n return 0;\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta) const override {\n return (meta.data_type() == data_type_ &&\n meta.unit_size() == IndexMeta::UnitSizeof(data_type_));\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta,\n const IndexQueryMeta &qmeta) const override {\n return (qmeta.data_type() == data_type_ &&\n qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) &&\n qmeta.dimension() == meta.dimension());\n }\n\n //! Retrieve distance function for query\n MatrixDistance distance(void) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP16:\n return reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute);\n\n case IndexMeta::DataType::DT_FP32:\n return reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute);\n\n case IndexMeta::DataType::DT_INT8:\n return reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute);\n\n case IndexMeta::DataType::DT_INT4:\n return reinterpret_cast(\n ailego::MinusInnerProductMatrix::Compute);\n\n default:\n return nullptr;\n }\n }\n\n //! Retrieve sparse distance function for query\n MatrixSparseDistance sparse_distance(void) const override {\n return reinterpret_cast(\n ailego::MinusInnerProductSparseMatrix::Compute);\n }\n\n //! Retrieve distance function for index features\n MatrixDistance distance_matrix(size_t m, size_t n) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP16:\n return MinusInnerProductMatrixFp16(m, n);\n\n case IndexMeta::DataType::DT_FP32:\n return MinusInnerProductMatrixFp32(m, n);\n\n case IndexMeta::DataType::DT_INT8:\n return MinusInnerProductMatrixInt8(m, n);\n\n case IndexMeta::DataType::DT_INT4:\n return MinusInnerProductMatrixInt4(m, n);\n\n default:\n return nullptr;\n }\n }\n\n //! Retrieve distance function for query\n MatrixBatchDistance batch_distance(void) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n case IndexMeta::DataType::DT_FP16:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n case IndexMeta::DataType::DT_INT8:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n case IndexMeta::DataType::DT_INT4:\n return reinterpret_cast(\n ailego::BaseDistance::ComputeBatch);\n default:\n return nullptr;\n }\n }\n\n //! Normalize result\n void normalize(float *score) const override {\n *score = -(*score);\n }\n\n //! Denormalize threshold\n void denormalize(float *score) const override {\n *score = -(*score);\n }\n\n //! Retrieve if it supports normalization\n bool support_normalize(void) const override {\n return true;\n }\n\n //! Retrieve params of Metric\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n //! Retrieve query measure object of this index measure\n Pointer query_metric(void) const override {\n return nullptr;\n }\n\n private:\n IndexMeta::MetaType meta_type_{IndexMeta::MetaType::MT_DENSE};\n IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32};\n ailego::Params params_{};\n};\n\n/*! Normalized Cosine Metric\n */\nclass NormalizedCosineMetric : public InnerProductMetric {\n public:\n //! Initialize Metric\n int init(const IndexMeta &meta, const ailego::Params &index_params) override {\n IndexMeta::DataType dt = meta.data_type();\n if (dt != IndexMeta::DataType::DT_FP16 &&\n dt != IndexMeta::DataType::DT_FP32) {\n return IndexError_Unsupported;\n }\n\n InnerProductMetric::init(meta, index_params);\n\n return 0;\n }\n\n //! Normalize result\n void normalize(float *score) const override {\n *score = 1 + (*score);\n }\n\n //! Denormalize threshold\n\n void denormalize(float *score) const override {\n *score -= 1;\n }\n};\n\n/*! Inner Product Sparse Metric\n */\nclass InnerProductSparseMetric : public IndexMetric {\n public:\n //! Initialize Metric\n int init(const IndexMeta &meta, const ailego::Params &index_params) override {\n IndexMeta::DataType dt = meta.data_type();\n if (dt != IndexMeta::DataType::DT_FP16 &&\n dt != IndexMeta::DataType::DT_FP32) {\n return IndexError_Unsupported;\n }\n\n if (IndexMeta::UnitSizeof(dt) != meta.unit_size()) {\n return IndexError_Unsupported;\n }\n\n data_type_ = dt;\n params_ = index_params;\n\n return 0;\n }\n\n //! Cleanup Metric\n int cleanup(void) override {\n return 0;\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta) const override {\n return (meta.data_type() == data_type_ &&\n meta.unit_size() == IndexMeta::UnitSizeof(data_type_));\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta,\n const IndexQueryMeta &qmeta) const override {\n return (qmeta.data_type() == data_type_ &&\n qmeta.data_type() == meta.data_type() &&\n qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) &&\n qmeta.unit_size() == meta.unit_size());\n }\n\n //! Retrieve distance function for query\n MatrixDistance distance(void) const override {\n return nullptr;\n }\n\n //! Retrieve sparse distance function for query\n MatrixSparseDistance sparse_distance(void) const override {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP16:\n return reinterpret_cast(\n ailego::MinusInnerProductSparseMatrix::Compute);\n case IndexMeta::DataType::DT_FP32:\n return reinterpret_cast(\n ailego::MinusInnerProductSparseMatrix::Compute);\n default:\n return nullptr;\n }\n }\n\n //! Normalize result\n void normalize(float *score) const override {\n *score = -(*score);\n }\n\n //! Denormalize threshold\n void denormalize(float *score) const override {\n *score = -(*score);\n }\n\n //! Retrieve if it supports normalization\n bool support_normalize(void) const override {\n return true;\n }\n\n //! Retrieve params of Metric\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n //! Retrieve query measure object of this index measure\n Pointer query_metric(void) const override {\n return nullptr;\n }\n\n private:\n IndexMeta::DataType data_type_{IndexMeta::DataType::DT_FP32};\n ailego::Params params_{};\n};\n\nINDEX_FACTORY_REGISTER_METRIC_ALIAS(InnerProduct, InnerProductMetric);\nINDEX_FACTORY_REGISTER_METRIC_ALIAS(NormalizedCosine, NormalizedCosineMetric);\n\nINDEX_FACTORY_REGISTER_METRIC_ALIAS(InnerProductSparse,\n InnerProductSparseMetric);\n\n} // namespace core\n} // namespace zvec\n" }, { "path": "src/core/metric/metric_params.h", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#pragma once\n\n#include \n\nnamespace zvec {\nnamespace core {\n\n//! MipsEuclideanMetric\nstatic const std::string MIPS_EUCLIDEAN_METRIC_M_VALUE =\n \"mips_euclidean.metric.m_value\";\nstatic const std::string MIPS_EUCLIDEAN_METRIC_U_VALUE =\n \"mips_euclidean.metric.u_value\";\nstatic const std::string MIPS_EUCLIDEAN_METRIC_MAX_L2_NORM =\n \"mips_euclidean.metric.max_l2_norm\";\nstatic const std::string MIPS_EUCLIDEAN_METRIC_INJECTION_TYPE =\n \"mips_euclidean.metric.injection_type\";\n\n//! QuantizedInteger Metric\nstatic const std::string QUANTIZED_INTEGER_METRIC_ORIGIN_METRIC_NAME =\n \"proxima.quantized_integer.metric.origin_metric_name\";\nstatic const std::string QUANTIZED_INTEGER_METRIC_ORIGIN_METRIC_PARAMS =\n \"proxima.quantized_integer.metric.origin_metric_params\";\n\n} // namespace core\n} // namespace zvec" }, { "path": "src/core/metric/mips_euclidean_metric.cc", "content": "// Copyright 2025-present the zvec project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n#include \n#include \n#include \n#include \n#include \n#include \n#include \"metric_params.h\"\n\nnamespace zvec {\nnamespace core {\n\n/*! Mips Squared Euclidean Metric\n */\ntemplate \nclass MipsSquaredEuclideanMetric : public IndexMetric {\n public:\n //! Initialize Metric\n int init(const IndexMeta &meta, const ailego::Params &index_params) override {\n data_type_ = meta.data_type();\n dimension_ = meta.dimension();\n\n int injection_type = static_cast(kDefaultInjectionType);\n index_params.get(MIPS_EUCLIDEAN_METRIC_INJECTION_TYPE, &injection_type);\n if (injection_type >= static_cast(Injection::kNumInjections)) {\n LOG_WARN(\"Unsupported injection_type %u, using '%s' instead\",\n injection_type, InjectionName(0));\n injection_type = static_cast(Injection::kLocalizedSpherical);\n }\n injection_ = static_cast(injection_type);\n LOG_DEBUG(\n \"Initializing MipsSquaredEuclideanMetric with injection %s\"\n \" type %d dimension %d\",\n InjectionName(injection_), data_type_, dimension_);\n\n float max_l2_norm = 0.0f;\n float u_value = 0.0f;\n index_params.get(MIPS_EUCLIDEAN_METRIC_M_VALUE, &m_value_);\n index_params.get(MIPS_EUCLIDEAN_METRIC_U_VALUE, &u_value);\n index_params.get(MIPS_EUCLIDEAN_METRIC_MAX_L2_NORM, &max_l2_norm);\n CheckAndFixM(injection_, &m_value_);\n CheckAndFixU(injection_, m_value_, &u_value);\n\n squared_u_value_ = u_value * u_value;\n max_squared_l2_norm_ = max_l2_norm * max_l2_norm;\n if (injection_ == Injection::kIdentity ||\n injection_ == Injection::kLocalizedSpherical) {\n eta_ = 0.0f;\n } else if (max_squared_l2_norm_ < std::numeric_limits::epsilon()) {\n eta_ = kDefaultEta;\n } else {\n eta_ = squared_u_value_ / max_squared_l2_norm_;\n }\n\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n squared_norm2_handle_ = reinterpret_cast(\n ailego::SquaredNorm2Matrix::Compute);\n break;\n\n case IndexMeta::DataType::DT_FP16:\n squared_norm2_handle_ = reinterpret_cast(\n ailego::SquaredNorm2Matrix::Compute);\n break;\n\n case IndexMeta::DataType::DT_INT8:\n squared_norm2_handle_ = reinterpret_cast(\n ailego::SquaredNorm2Matrix::Compute);\n break;\n\n case IndexMeta::DataType::DT_INT4:\n squared_norm2_handle_ = reinterpret_cast(\n ailego::SquaredNorm2Matrix::Compute);\n break;\n\n default:\n return IndexError_Unsupported;\n }\n\n query_metric_ = IndexFactory::CreateMetric(kQueryMetric);\n if (!query_metric_) {\n LOG_ERROR(\"Failed to create metric %s\", kQueryMetric);\n return IndexError_NoExist;\n }\n int ret = query_metric_->init(meta, ailego::Params());\n if (ret != 0) {\n LOG_ERROR(\"Failed to initialize metric %s\", kQueryMetric);\n return ret;\n }\n params_ = index_params;\n return 0;\n }\n\n //! Cleanup Metric\n int cleanup(void) override {\n eta_ = 0.0f;\n m_value_ = 0;\n squared_u_value_ = 0.0f;\n max_squared_l2_norm_ = 0.0f;\n query_metric_.reset();\n return 0;\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta) const override {\n return (meta.data_type() == data_type_ &&\n meta.unit_size() == IndexMeta::UnitSizeof(data_type_));\n }\n\n //! Retrieve if it matched\n bool is_matched(const IndexMeta &meta,\n const IndexQueryMeta &qmeta) const override {\n return (qmeta.data_type() == data_type_ &&\n qmeta.unit_size() == IndexMeta::UnitSizeof(data_type_) &&\n qmeta.dimension() == meta.dimension());\n }\n\n //! Retrieve distance function for query\n MatrixBatchDistance batch_distance() const override {\n MatrixDistance dist_func = distance();\n\n return\n [=](const void **m, const void *q, size_t num, size_t dim, float *out) {\n for (size_t i = 0; i < num; ++i) {\n dist_func(m[i], q, dim, out + i);\n }\n };\n }\n\n\n //! Retrieve distance function for query\n MatrixDistance distance(void) const override {\n if (injection_ == Injection::kLocalizedSpherical) {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, 0.0f, out);\n };\n\n case IndexMeta::DataType::DT_FP16:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::\n Compute(reinterpret_cast(m),\n reinterpret_cast(q), dim, 0.0f,\n out);\n };\n\n case IndexMeta::DataType::DT_INT8:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, 0.0f, out);\n };\n\n case IndexMeta::DataType::DT_INT4:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, 0.0f, out);\n };\n\n default:\n return nullptr;\n }\n }\n\n if (injection_ == Injection::kRepeatedQuadratic) {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, m_value_, eta_, out);\n };\n\n case IndexMeta::DataType::DT_FP16:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::\n Compute(reinterpret_cast(m),\n reinterpret_cast(q), dim,\n m_value_, eta_, out);\n };\n\n case IndexMeta::DataType::DT_INT8:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, m_value_, eta_, out);\n };\n\n case IndexMeta::DataType::DT_INT4:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, m_value_, eta_, out);\n };\n\n default:\n return nullptr;\n }\n }\n\n if (injection_ == Injection::kSpherical) {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, eta_, out);\n };\n\n case IndexMeta::DataType::DT_FP16:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::\n Compute(reinterpret_cast(m),\n reinterpret_cast(q), dim, eta_,\n out);\n };\n\n case IndexMeta::DataType::DT_INT8:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, eta_, out);\n };\n\n case IndexMeta::DataType::DT_INT4:\n return [&](const void *m, const void *q, size_t dim, float *out) {\n ailego::MipsSquaredEuclideanDistanceMatrix::Compute(\n reinterpret_cast(m),\n reinterpret_cast(q), dim, eta_, out);\n };\n\n default:\n return nullptr;\n }\n }\n\n if (injection_ == Injection::kIdentity) {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n return reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute);\n\n case IndexMeta::DataType::DT_FP16:\n return reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute);\n\n case IndexMeta::DataType::DT_INT8:\n return reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute);\n\n case IndexMeta::DataType::DT_INT4:\n return reinterpret_cast(\n ailego::SquaredEuclideanDistanceMatrix::Compute);\n\n default:\n return nullptr;\n }\n }\n return nullptr;\n }\n\n //! Retrieve distance function for query\n MatrixSparseDistance sparse_distance(void) const override {\n if (injection_ == Injection::kLocalizedSpherical) {\n return [&](const void *m_sparse, const void *q_sparse, float *out) {\n ailego::MipsSquaredEuclideanSparseDistanceMatrix::Compute(\n m_sparse, q_sparse, out);\n };\n }\n\n if (injection_ == Injection::kRepeatedQuadratic) {\n LOG_ERROR(\n \"Repeated Quadratic is not supported in MipsEuclideanMetric for \"\n \"Hybrid Vector!\");\n\n return nullptr;\n }\n\n if (injection_ == Injection::kSpherical) {\n LOG_ERROR(\n \"Spherical is not supported in MipsEuclideanMetric for Hybrid \"\n \"Vector!\");\n\n return nullptr;\n }\n\n if (injection_ == Injection::kIdentity) {\n LOG_ERROR(\n \"Identity is not supported in MipsEuclideanMetric for Hybrid \"\n \"Vector!\");\n\n return nullptr;\n }\n\n return nullptr;\n }\n\n //! Retrieve matrix distance function for index features\n MatrixDistance distance_matrix(size_t m, size_t n) const override {\n if (injection_ == Injection::kLocalizedSpherical) {\n SphericalHandle compute;\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n case IndexMeta::DataType::DT_FP16:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n case IndexMeta::DataType::DT_INT8:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n case IndexMeta::DataType::DT_INT4:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n default:\n return nullptr;\n }\n return [=](const void *d, const void *q, size_t dim, float *out) {\n compute(d, q, dim, 0.0f, out);\n };\n }\n\n if (injection_ == Injection::kRepeatedQuadratic) {\n RepeatedQuadraticHandle compute;\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n case IndexMeta::DataType::DT_FP16:\n compute =\n DistanceMatrixCompute(\n m, n);\n break;\n case IndexMeta::DataType::DT_INT8:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n case IndexMeta::DataType::DT_INT4:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n default:\n return nullptr;\n }\n return [=](const void *d, const void *q, size_t dim, float *out) {\n compute(d, q, dim, m_value_, eta_, out);\n };\n }\n\n if (injection_ == Injection::kSpherical) {\n SphericalHandle compute;\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n case IndexMeta::DataType::DT_FP16:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n case IndexMeta::DataType::DT_INT8:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n case IndexMeta::DataType::DT_INT4:\n compute =\n DistanceMatrixCompute(m, n);\n break;\n default:\n return nullptr;\n }\n return [=](const void *d, const void *q, size_t dim, float *out) {\n compute(d, q, dim, eta_, out);\n };\n }\n\n if (injection_ == Injection::kIdentity) {\n switch (data_type_) {\n case IndexMeta::DataType::DT_FP32:\n return DistanceMatrixCompute(m, n);\n case IndexMeta::DataType::DT_FP16:\n return DistanceMatrixCompute(m,\n n);\n case IndexMeta::DataType::DT_INT8:\n return DistanceMatrixCompute(m, n);\n case IndexMeta::DataType::DT_INT4:\n return DistanceMatrixCompute(m, n);\n default:\n return nullptr;\n }\n }\n return nullptr;\n }\n\n //! Normalize result\n void normalize(float *score) const override {\n query_metric_->normalize(score);\n }\n\n //! Denormalize threshold\n void denormalize(float *score) const override {\n query_metric_->denormalize(score);\n }\n\n //! Retrieve if it supports normalization\n bool support_normalize(void) const override {\n return query_metric_->support_normalize();\n }\n\n //! Retrieve params of Metric\n const ailego::Params ¶ms(void) const override {\n return params_;\n }\n\n //! Train the metric\n int train(const void *vec, size_t dim) override {\n if (eta_ == 0.0f) { // No global norm scaling => no training.\n return 0;\n }\n if (!squared_norm2_handle_) {\n return IndexError_Unsupported;\n }\n\n float score;\n squared_norm2_handle_(vec, dim, &score);\n if (score > max_squared_l2_norm_) {\n max_squared_l2_norm_ = score;\n const float max_l2_norm = std::sqrt(score);\n params_.set(MIPS_EUCLIDEAN_METRIC_MAX_L2_NORM, max_l2_norm);\n if (max_squared_l2_norm_ < 1.0 &&\n max_squared_l2_norm_ > squared_u_value_) {\n squared_u_value_ = max_squared_l2_norm_;\n params_.set(MIPS_EUCLIDEAN_METRIC_U_VALUE, max_l2_norm);\n }\n eta_ = squared_u_value_ / max_squared_l2_norm_;\n }\n return 0;\n }\n\n //! Retrieve if it supports training\n bool support_train(void) const override {\n // No global norm scaling => eta_ == 0 => no training.\n return eta_ != 0.0f;\n }\n\n //! Retrieve query metric object of this index metric\n Pointer query_metric(void) const override {\n return query_metric_;\n }\n\n private:\n //! Type of MipsSquaredEuclideanDistanceMatrix::Compute overloaded for\n // Spherical injection and LocalizedSpherical nonmetric.\n template \n using SphericalHandle = void (*)(const T *m, const T *q, size_t dim,\n float eta, float *out);\n\n //! Type of MipsSquaredEuclideanDistanceMatrix::Compute overloaded for\n // RepeatedQuadratic injection.\n template \n using RepeatedQuadraticHandle = void (*)(const T *m, const T *q, size_t dim,\n size_t m_value, float eta,\n float *out);\n\n //! Type of squared L2 norm function.\n using SquaredNorm2Handle = void (*)(const void *m, size_t dim, float *out);\n\n enum struct Injection { // Type of injective mapping into Euclidean space.\n kLocalizedSpherical = 0, // spherical with pair-only max-norm\n kSpherical = 1, // require global scaling/training\n kRepeatedQuadratic = 2, // require global scaling/training\n kIdentity = 3, // plain Euclidean distance\n kNumInjections = 4\n };\n\n static const char *InjectionName(int injection) {\n static const char *injection_names[] = {\"LocalizedSpherical\", \"Spherical\",\n \"RepeatedQuadratic\", \"Identity\"};\n if (injection >= 0 &&\n injection < static_cast(Injection::kNumInjections)) {\n return injection_names[injection];\n }\n return \"Invalid\";\n }\n\n static const char *InjectionName(Injection injection) {\n return InjectionName(static_cast(injection));\n }\n\n // Checks (and fixes) `*m_value`, no. additional dimensions for injection.\n // `dim` is the original dimension, used ONLY by RepeatedQuadratic\n // injection, where dim = 1 induces the default *m_value = 3. It's\n // positioned last to allow other injections to skip it.\n // Returns true if `*m_value` is modified.\n static bool CheckAndFixM(Injection injection, uint32_t *m_value) {\n if (injection == Injection::kRepeatedQuadratic) {\n if (*m_value == 0) {\n *m_value = 3u; // Recommend value in paper (3.5 Practical\n // Recommendation of Parameters)\n return true;\n }\n } else if (injection == Injection::kSpherical) {\n if (*m_value != 1) {\n if (*m_value != 0) {\n LOG_WARN(\"M value (%u) set to 1 for Spherical injection\", *m_value);\n }\n *m_value = 1;\n return true;\n }\n } else { // kLocalizedSpherical, kIdentity, or kInvalid\n if (*m_value != 0) {\n LOG_WARN(\"M value (%u) set to 0 for %s injections\", *m_value,\n InjectionName(injection));\n *m_value = 0;\n return true;\n }\n }\n return false;\n }\n\n // Checks and fixes `*u_value`, global L2 norm scalar.\n // `m_value` is no. additional dimensions, used ONLY by RepeatedQuadratic\n // injection. It's positioned last to allow other injections to skip it.\n // Returns true if `*u_value` is set to a new value.\n static bool CheckAndFixU(Injection injection, uint32_t m_value,\n float *u_value) {\n if (injection == Injection::kRepeatedQuadratic) {\n if (*u_value <= std::numeric_limits::epsilon() ||\n *u_value >= 1.0) {\n // Try computing a default U value\n constexpr float kLogError = -5.0; // log_10(distance_error)\n float new_u_value = std::pow(10, kLogError / (1 << (m_value + 1)));\n if (*u_value != 0) {\n LOG_WARN(\"U value (%f) set to %f for RepeatedQuadratic injection\",\n *u_value, new_u_value);\n }\n *u_value = new_u_value;\n return true;\n } else if (std::pow(*u_value, (1 << m_value)) <\n std::numeric_limits::epsilon()) {\n LOG_WARN(\n \"U value %f is too small, may cause loss of distance precision\",\n *u_value);\n }\n } else if (injection == Injection::kSpherical) {\n // Spherical injection requires ||x'|| <= 1.0 for computing\n // std::sqrt(1 - ||x'||^2), x' = u_value * x / max_norm. Set u_value\n // to slightly < 1.0 in case of precision loss in float computation.\n if (*u_value <= std::numeric_limits::epsilon() ||\n *u_value >= 1.0) {\n static constexpr float kSphericalUValue = 1.0f - 1e-3;\n if (*u_value != 0.0f) {\n LOG_WARN(\"U value (%f) set to %f for Spherical injection\", *u_value,\n kSphericalUValue);\n }\n *u_value = kSphericalUValue;\n return true;\n }\n } else { // kLocalizedSpherical, kIdentity, or kInvalid\n if (*u_value != 1.0) {\n if (*u_value != 0) {\n LOG_WARN(\"U value (%f) set to 1.0 for %s injection\", *u_value,\n InjectionName(injection));\n }\n *u_value = 1.0;\n return true;\n }\n }\n return false;\n }\n\n private:\n //! Type of basic DistanceMatrix::Compute function with typed parameter.\n template \n using TypedDistanceHandle = void (*)(const T *m, const T *q, size_t dim,\n float *out);\n\n //! Returns m x n distance matrix compute function.\n // Handle is used to resolve potential DistanceMatrix::Compute overload.\n template