diff --git a/.bazelrc b/.bazelrc index 7f49c964..56efa6ae 100644 --- a/.bazelrc +++ b/.bazelrc @@ -46,10 +46,17 @@ build:linux-release --action_env=BAZEL_LINKLIBS=-l%:libstdc++.a:-l%:libgcc.a # platform specific config # Bazel will automatic pick platform config since we have enable_platform_specific_config set build:macos --copt=-Wno-unused-command-line-argument +build:macos --copt=-Wno-vla-cxx-extension build:macos --features=-supports_dynamic_linker build:macos --macos_minimum_os=12.0 build:macos --host_macos_minimum_os=12.0 +# Fix yacl VLA warning on macOS Clang (Xcode 16.4+) +build --per_file_copt=external/yacl.*openssl.*@-Wno-vla-cxx-extension +# Fix brpc deprecated protobuf API warning +build --per_file_copt=external/yacl.*brpc.*@-Wno-deprecated-declarations +build --per_file_copt=external/brpc.*@-Wno-deprecated-declarations + # NOTE: fix build on linux_aarch64, since `build:linux_aarch64` works only if `--config linux_aarch64` option included. build:linux --cxxopt=-Wno-mismatched-new-delete diff --git a/.circleci/test.yml b/.circleci/test.yml index 6e963faa..c9ac788f 100644 --- a/.circleci/test.yml +++ b/.circleci/test.yml @@ -37,6 +37,12 @@ commands: - run: name: Setup BuildBuddy Cache command: ../devtools/bazel_cache_setup.py + - run: + name: Install Intel HEXL + command: | + export HEXL_ROOT="${HOME}/.local/hexl" + ./bazel/install_hexl.sh "${HEXL_ROOT}" + echo "export HEXL_ROOT=${HEXL_ROOT}" >> "${BASH_ENV}" - run: name: "build" command: bazel build //... -c opt --ui_event_filters=-info,-debug,-warning --jobs 20 @@ -76,8 +82,8 @@ jobs: - build_and_test macos_ut: macos: - xcode: 16.2.0 - resource_class: macos.m1.large.gen1 + xcode: 16.4.0 + resource_class: m4pro.medium steps: - checkout - run: diff --git a/.codex b/.codex new file mode 100644 index 00000000..e69de29b diff --git a/.gitignore b/.gitignore index 1bdfef08..96682349 100644 --- a/.gitignore +++ b/.gitignore @@ -48,3 +48,4 @@ ossutil_output/ psi/apps/pir/docker/linux/ psi/apps/pir/docker/version.txt +.vscode/ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 02d6def6..918b70cc 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -13,6 +13,8 @@ part of the project. In general, please use clang-format to format code, and follow clang-tidy tips. +For Bazel files (`MODULE.bazel`, `*.bzl`, `BUILD`, `BUILD.bazel`), please use `buildifier`. + Most of the code style is derived from the [Google C++ style guidelines](https://google.github.io/styleguide/cppguide.html), except: @@ -57,6 +59,10 @@ docker exec -it psi-dev-$(whoami) bash ```sh Install gcc>=11.2, cmake>=3.26, ninja, nasm>=2.15, python>=3.10, bazelisk, xxd, lld + +# Install Intel HEXL, or set HEXL_ROOT to an existing install prefix. +./bazel/install_hexl.sh "$HOME/.local/hexl" +export HEXL_ROOT="$HOME/.local/hexl" ``` #### macOS @@ -80,6 +86,10 @@ brew install bazelisk cmake ninja libomp wget # For Intel mac only brew install nasm + +# Install Intel HEXL, or set HEXL_ROOT to an existing install prefix. +./bazel/install_hexl.sh "$HOME/.local/hexl" +export HEXL_ROOT="$HOME/.local/hexl" ``` ### Build & UnitTest diff --git a/MODULE.bazel b/MODULE.bazel index 78147e81..84e97d5d 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -58,6 +58,8 @@ use_repo( new_local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl", "new_local_repository") +hexl_repository = use_repo_rule("//bazel:hexl_repo.bzl", "hexl_repository") + new_local_repository( name = "macos_omp_x64", build_file = "@yacl//bazel:local_openmp_macos.BUILD", @@ -70,6 +72,8 @@ new_local_repository( path = "/opt/homebrew/opt/libomp/", ) +hexl_repository(name = "hexl") + bazel_dep(name = "kuscia", version = "0.14.0b0") bazel_dep(name = "perfetto", version = "41.0") diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 6196ea06..2b1e424d 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -910,128 +910,6 @@ }, "selectedYankedVersions": {}, "moduleExtensions": { - "//bazel:defs.bzl%non_module_dependencies": { - "general": { - "bzlTransitiveDigest": "NAC0+n0skvguFGiSWQSZMhbrQQVV+nNsrNIVDde/xQE=", - "usagesDigest": "89PlExpoF4wVcL+M7Z1YGe8ZyC3E1M5lGgLzpZSTyvE=", - "recordedFileInputs": {}, - "recordedDirentsInputs": {}, - "envVariables": {}, - "generatedRepoSpecs": { - "apsi": { - "bzlFile": "@@bazel_tools//tools/build_defs/repo:http.bzl", - "ruleClassName": "http_archive", - "attributes": { - "sha256": "82c0f9329c79222675109d4a3682d204acd3ea9a724bcd98fa58eabe53851333", - "strip_prefix": "APSI-0.11.0", - "urls": [ - "https://github.com/microsoft/APSI/archive/refs/tags/v0.11.0.tar.gz" - ], - "build_file": "@@//bazel:microsoft_apsi.BUILD", - "patch_args": [ - "-p1" - ], - "patches": [ - "@@//bazel/patches:apsi.patch", - "@@//bazel/patches:apsi-fourq.patch" - ], - "patch_cmds": [ - "rm -rf common/apsi/fourq" - ] - } - }, - "kuku": { - "bzlFile": "@@bazel_tools//tools/build_defs/repo:http.bzl", - "ruleClassName": "http_archive", - "attributes": { - "sha256": "96ed5fad82ea8c8a8bb82f6eaf0b5dce744c0c2566b4baa11d8f5443ad1f83b7", - "strip_prefix": "Kuku-2.1.0", - "type": "tar.gz", - "urls": [ - "https://github.com/microsoft/Kuku/archive/refs/tags/v2.1.0.tar.gz" - ], - "build_file": "@@//bazel:microsoft_kuku.BUILD" - } - }, - "com_google_flatbuffers": { - "bzlFile": "@@bazel_tools//tools/build_defs/repo:http.bzl", - "ruleClassName": "http_archive", - "attributes": { - "sha256": "4157c5cacdb59737c5d627e47ac26b140e9ee28b1102f812b36068aab728c1ed", - "strip_prefix": "flatbuffers-24.3.25", - "urls": [ - "https://github.com/google/flatbuffers/archive/refs/tags/v24.3.25.tar.gz" - ], - "patch_cmds": [ - "rm grpc/BUILD.bazel", - "rm grpc/src/compiler/BUILD.bazel", - "rm src/BUILD.bazel" - ], - "build_file": "@@//bazel:flatbuffers.BUILD" - } - }, - "curve25519-donna": { - "bzlFile": "@@bazel_tools//tools/build_defs/repo:http.bzl", - "ruleClassName": "http_archive", - "attributes": { - "strip_prefix": "curve25519-donna-2fe66b65ea1acb788024f40a3373b8b3e6f4bbb2", - "sha256": "ba57d538c241ad30ff85f49102ab2c8dd996148456ed238a8c319f263b7b149a", - "type": "tar.gz", - "build_file": "@@//bazel:curve25519-donna.BUILD", - "urls": [ - "https://github.com/floodyberry/curve25519-donna/archive/2fe66b65ea1acb788024f40a3373b8b3e6f4bbb2.tar.gz" - ] - } - }, - "com_github_zeromq_cppzmq": { - "bzlFile": "@@bazel_tools//tools/build_defs/repo:http.bzl", - "ruleClassName": "http_archive", - "attributes": { - "build_file": "@@//bazel:cppzmq.BUILD", - "strip_prefix": "cppzmq-4.10.0", - "sha256": "c81c81bba8a7644c84932225f018b5088743a22999c6d82a2b5f5cd1e6942b74", - "type": ".tar.gz", - "urls": [ - "https://github.com/zeromq/cppzmq/archive/refs/tags/v4.10.0.tar.gz" - ] - } - }, - "com_github_zeromq_libzmq": { - "bzlFile": "@@bazel_tools//tools/build_defs/repo:http.bzl", - "ruleClassName": "http_archive", - "attributes": { - "build_file": "@@//bazel:libzmq.BUILD", - "strip_prefix": "libzmq-4.3.5", - "sha256": "6c972d1e6a91a0ecd79c3236f04cf0126f2f4dfbbad407d72b4606a7ba93f9c6", - "type": ".tar.gz", - "urls": [ - "https://github.com/zeromq/libzmq/archive/refs/tags/v4.3.5.tar.gz" - ] - } - }, - "com_github_open_source_parsers_jsoncpp": { - "bzlFile": "@@bazel_tools//tools/build_defs/repo:http.bzl", - "ruleClassName": "http_archive", - "attributes": { - "build_file": "@@//bazel:jsoncpp.BUILD", - "strip_prefix": "jsoncpp-1.9.6", - "sha256": "f93b6dd7ce796b13d02c108bc9f79812245a82e577581c4c9aabe57075c90ea2", - "type": ".tar.gz", - "urls": [ - "https://github.com/open-source-parsers/jsoncpp/archive/refs/tags/1.9.6.tar.gz" - ] - } - } - }, - "recordedRepoMappingEntries": [ - [ - "", - "bazel_tools", - "bazel_tools" - ] - ] - } - }, "@@apple_support~//crosstool:setup.bzl%apple_cc_configure_extension": { "general": { "bzlTransitiveDigest": "7ii+gFxWSxHhQPrBxfMEHhtrGvHmBTvsh+KOyGunP/s=", diff --git a/README.md b/README.md index 3e78deae..18c0faa0 100644 --- a/README.md +++ b/README.md @@ -218,3 +218,4 @@ Please refer to [PSI V2 Benchmark](docs/user_guide/psi_v2_benchmark.md) ## APSI Benchmark Please refer to [APSI Benchmark](docs/user_guide/apsi_benchmark.md) + diff --git a/bazel/hexl.BUILD b/bazel/hexl.BUILD new file mode 100644 index 00000000..faf6aa63 --- /dev/null +++ b/bazel/hexl.BUILD @@ -0,0 +1,15 @@ +load("@rules_cc//cc:defs.bzl", "cc_import", "cc_library") + +package(default_visibility = ["//visibility:public"]) + +cc_import( + name = "hexl_lib", + static_library = "lib/libhexl.a", +) + +cc_library( + name = "hexl", + hdrs = glob(["include/hexl/**/*.h", "include/hexl/**/*.hpp"]), + includes = ["include"], + deps = [":hexl_lib"], +) diff --git a/bazel/hexl_repo.bzl b/bazel/hexl_repo.bzl new file mode 100644 index 00000000..b41edc5e --- /dev/null +++ b/bazel/hexl_repo.bzl @@ -0,0 +1,65 @@ +_HEXL_BUILD = """load("@rules_cc//cc:defs.bzl", "cc_import", "cc_library") + +package(default_visibility = ["//visibility:public"]) + +cc_import( + name = "hexl_lib", + static_library = "lib/libhexl.a", +) + +cc_library( + name = "hexl", + hdrs = glob(["include/hexl/**/*.h", "include/hexl/**/*.hpp"]), + includes = ["include"], + deps = [":hexl_lib"], +) +""" + +_HEXL_STUB_BUILD = """load("@rules_cc//cc:defs.bzl", "cc_library") + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "hexl", +) +""" + +def _hexl_repository_impl(repository_ctx): + roots = [] + env_root = repository_ctx.os.environ.get("HEXL_ROOT", "") + if env_root: + roots.append(env_root) + + roots.extend([ + "/usr/local/opt/intel-hexl", + "/usr/local", + "/opt/homebrew/opt/intel-hexl", + "/opt/homebrew", + "/usr", + ]) + + for root in roots: + include_dir = repository_ctx.path(root + "/include") + header_dir = repository_ctx.path(root + "/include/hexl") + for lib_dir_name in ["lib", "lib64"]: + lib_dir = repository_ctx.path(root + "/" + lib_dir_name) + static_lib = repository_ctx.path(root + "/" + lib_dir_name + "/libhexl.a") + if header_dir.exists and static_lib.exists: + repository_ctx.symlink(include_dir, "include") + repository_ctx.symlink(lib_dir, "lib") + repository_ctx.file("BUILD.bazel", _HEXL_BUILD) + repository_ctx.file("WORKSPACE.bazel", "workspace(name = \"hexl\")\n") + return + + # HEXL not found — create a stub so the workspace resolves on non-x86 + # platforms. Targets that actually depend on @hexl are gated by + # target_compatible_with = ["@platforms//cpu:x86_64"] and will never be + # built on ARM / macOS Apple Silicon. + repository_ctx.file("BUILD.bazel", _HEXL_STUB_BUILD) + repository_ctx.file("WORKSPACE.bazel", "workspace(name = \"hexl\")\n") + +hexl_repository = repository_rule( + implementation = _hexl_repository_impl, + environ = ["HEXL_ROOT"], + local = True, +) diff --git a/bazel/install_hexl.sh b/bazel/install_hexl.sh new file mode 100755 index 00000000..05d2a686 --- /dev/null +++ b/bazel/install_hexl.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash + +set -euo pipefail + +hexl_version="${HEXL_VERSION:-v1.2.6}" +install_prefix="${1:-${HEXL_ROOT:-${HOME}/.local/hexl}}" +arch="$(uname -m)" + +case "${arch}" in + x86_64|amd64) + ;; + *) + printf 'Skipping Intel HEXL installation on unsupported architecture %s\n' \ + "${arch}" + exit 0 + ;; +esac + +if [[ -f "${install_prefix}/lib/libhexl.a" && -d "${install_prefix}/include/hexl" ]]; then + printf 'Intel HEXL already installed at %s\n' "${install_prefix}" + exit 0 +fi + +parallelism="$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)" +workdir="$(mktemp -d)" +repo_dir="${workdir}/hexl" +build_dir="${workdir}/build" + +cleanup() { + rm -rf "${workdir}" +} + +trap cleanup EXIT + +printf 'Installing Intel HEXL %s into %s\n' "${hexl_version}" "${install_prefix}" + +git clone --depth 1 --branch "${hexl_version}" https://github.com/IntelLabs/hexl.git "${repo_dir}" + +cmake -S "${repo_dir}" -B "${build_dir}" -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX="${install_prefix}" -DHEXL_BENCHMARK=OFF -DHEXL_COVERAGE=OFF -DHEXL_DOCS=OFF -DHEXL_SHARED_LIB=OFF -DHEXL_TESTING=OFF -DHEXL_TREAT_WARNING_AS_ERROR=OFF + +cmake --build "${build_dir}" --parallel "${parallelism}" +cmake --install "${build_dir}" diff --git a/bazel/psi.bzl b/bazel/psi.bzl index 2ebadb70..adf7b1f8 100644 --- a/bazel/psi.bzl +++ b/bazel/psi.bzl @@ -26,7 +26,7 @@ WARNING_FLAGS = [ ] DEBUG_FLAGS = ["-O0", "-g", "-DSPDLOG_ACTIVE_LEVEL=1"] RELEASE_FLAGS = ["-O2"] -FAST_FLAGS = ["-O1"] +FAST_FLAGS = ["-O3"] def _psi_copts(): return select({ diff --git a/docs/user_guide/pir.rst b/docs/user_guide/pir.rst index e7923919..26aba993 100644 --- a/docs/user_guide/pir.rst +++ b/docs/user_guide/pir.rst @@ -13,6 +13,8 @@ Supported Protocols +----------------+-------------+---------------+ | APSI | Keyword PIR | Single Server | +----------------+-------------+---------------+ +| YPIR | Index PIR | Single Server | ++----------------+-------------+---------------+ At this moment, SealPIR is under development. diff --git a/psi/algorithm/pir_interface/pir_type.proto b/psi/algorithm/pir_interface/pir_type.proto index fb8cde26..37d13812 100644 --- a/psi/algorithm/pir_interface/pir_type.proto +++ b/psi/algorithm/pir_interface/pir_type.proto @@ -21,6 +21,7 @@ enum PirType { INVALID = 0; SEAL_PIR = 1; SPIRAL_PIR = 2; + YPIR_PIR = 3; } message PirTypeProto { diff --git a/psi/algorithm/spiral/arith/ntt.cc b/psi/algorithm/spiral/arith/ntt.cc index eb2e693e..95f04efb 100644 --- a/psi/algorithm/spiral/arith/ntt.cc +++ b/psi/algorithm/spiral/arith/ntt.cc @@ -77,15 +77,14 @@ void NttForward(const Params& params, absl::Span operand_overall) { curr_x + (static_cast(two_times_modulus_small) - q_new); } } - - // Update the operand with modulus constraints - for (std::size_t i = 0; i < n; ++i) { - operand[i] -= - static_cast(operand[i] >= two_times_modulus_small) * - two_times_modulus_small; - operand[i] -= static_cast(operand[i] >= modulus_small) * - modulus_small; - } + } + // Update the operand with modulus constraints + for (std::size_t i = 0; i < n; ++i) { + operand[i] -= + static_cast(operand[i] >= two_times_modulus_small) * + two_times_modulus_small; + operand[i] -= static_cast(operand[i] >= modulus_small) * + modulus_small; } } } @@ -99,6 +98,7 @@ void NttForward(const Params& params, absl::Span operand_overall) { std::size_t log_n = params.PolyLenLog2(); std::size_t n = static_cast(1) << log_n; + YACL_ENFORCE(n >= 4, "PolyLen must be >= 4 for AVX2 impl"); YACL_ENFORCE(operand_overall.size() >= params.CrtCount() * n); for (std::size_t coeff_mod = 0; coeff_mod < params.CrtCount(); ++coeff_mod) { @@ -116,11 +116,8 @@ void NttForward(const Params& params, absl::Span operand_overall) { for (std::size_t i = 0; i < m; ++i) { uint64_t w = forward_table[m + i]; uint64_t w_prime = forward_table_prime[m + i]; - auto op = operand.subspan(i * (2 * t), 2 * t); - SPDLOG_DEBUG("Processing coeff_mod: {}, m: {}, i: {}", coeff_mod, m, i); - if (t < 4) { for (std::size_t j = 0; j < t; ++j) { uint32_t x = static_cast(op[j]); @@ -129,6 +126,7 @@ void NttForward(const Params& params, absl::Span operand_overall) { std::uint32_t curr_x = x - (two_times_modulus_small * static_cast(x >= two_times_modulus_small)); + std::uint64_t q_tmp = (static_cast(y) * static_cast(w_prime)) >> 32; @@ -143,46 +141,36 @@ void NttForward(const Params& params, absl::Span operand_overall) { } } else { for (std::size_t j = 0; j < t; j += 4) { - if (j + 4 > t) break; // Ensure we do not exceed bounds - __m256i* p_x = reinterpret_cast<__m256i*>(&op[j]); __m256i* p_y = reinterpret_cast<__m256i*>(&op[j + t]); __m256i x = _mm256_loadu_si256(p_x); __m256i y = _mm256_loadu_si256(p_y); - __m256i cmp_val = _mm256_set1_epi64x( + __m256i threshold_2q = _mm256_set1_epi64x( + static_cast(two_times_modulus_small - 1)); + __m256i val_2q = _mm256_set1_epi64x( static_cast(two_times_modulus_small)); - // reuse this variable to reduce variable num - // gt_mask - __m256i tmp1 = _mm256_cmpgt_epi64(x, cmp_val); - // __m256i to_subtract = _mm256_and_si256(gt_mask_reused, cmp_val); - tmp1 = _mm256_and_si256(tmp1, cmp_val); - __m256i curr_x = _mm256_sub_epi64(x, tmp1); + __m256i gt_mask = _mm256_cmpgt_epi64(x, threshold_2q); + __m256i to_subtract = _mm256_and_si256(gt_mask, val_2q); + __m256i curr_x = _mm256_sub_epi64(x, to_subtract); - // __m256i w_prime_vec = - // _mm256_set1_epi64x(static_cast(w_prime)); - tmp1 = _mm256_set1_epi64x(static_cast(w_prime)); - // __m256i product = _mm256_mul_epu32(y, tmp1); - tmp1 = _mm256_mul_epu32(y, tmp1); - // __m256i q_val = _mm256_srli_epi64(tmp1, 32); - tmp1 = _mm256_srli_epi64(tmp1, 32); + __m256i w_prime_vec = + _mm256_set1_epi64x(static_cast(w_prime)); + __m256i product = _mm256_mul_epu32(y, w_prime_vec); + __m256i q_val = _mm256_srli_epi64(product, 32); - // __m256i w_vec = _mm256_set1_epi64x(static_cast(w)); - __m256i tmp2 = _mm256_set1_epi64x(static_cast(w)); - // __m256i w_times_y = _mm256_mul_epu32(y, w_vec); - // __m256i w_times_y = _mm256_mul_epu32(y, tmp2); - tmp2 = _mm256_mul_epu32(y, tmp2); + __m256i w_vec = _mm256_set1_epi64x(static_cast(w)); + __m256i w_times_y = _mm256_mul_epu32(y, w_vec); __m256i modulus_small_vec = _mm256_set1_epi64x(static_cast(modulus_small)); - // __m256i q_scaled = _mm256_mul_epu32(q_val, modulus_small_vec); - __m256i q_scaled = _mm256_mul_epu32(tmp1, modulus_small_vec); - __m256i q_final = _mm256_sub_epi64(tmp2, q_scaled); + __m256i q_scaled = _mm256_mul_epu32(q_val, modulus_small_vec); + __m256i q_final = _mm256_sub_epi64(w_times_y, q_scaled); __m256i new_x = _mm256_add_epi64(curr_x, q_final); - __m256i q_final_inverted = _mm256_sub_epi64(cmp_val, q_final); + __m256i q_final_inverted = _mm256_sub_epi64(val_2q, q_final); __m256i new_y = _mm256_add_epi64(curr_x, q_final_inverted); _mm256_storeu_si256(p_x, new_x); @@ -193,21 +181,28 @@ void NttForward(const Params& params, absl::Span operand_overall) { } for (std::size_t i = 0; i < n; i += 4) { - if (i + 4 > n) break; // Ensure we do not exceed bounds __m256i* p_x = reinterpret_cast<__m256i*>(&operand[i]); + __m256i x = _mm256_loadu_si256(p_x); - __m256i cmp_val1 = + // Check >= 2q + __m256i threshold_2q = + _mm256_set1_epi64x(static_cast(two_times_modulus_small - 1)); + __m256i val_2q = _mm256_set1_epi64x(static_cast(two_times_modulus_small)); - __m256i x = _mm256_loadu_si256(p_x); - __m256i gt_mask = _mm256_cmpgt_epi64(x, cmp_val1); - __m256i to_subtract = _mm256_and_si256(gt_mask, cmp_val1); + + __m256i gt_mask = _mm256_cmpgt_epi64(x, threshold_2q); + __m256i to_subtract = _mm256_and_si256(gt_mask, val_2q); x = _mm256_sub_epi64(x, to_subtract); - __m256i cmp_val2 = - _mm256_set1_epi64x(static_cast(modulus_small)); - gt_mask = _mm256_cmpgt_epi64(x, cmp_val2); - to_subtract = _mm256_and_si256(gt_mask, cmp_val2); + // Check >= q + __m256i threshold_q = + _mm256_set1_epi64x(static_cast(modulus_small - 1)); + __m256i val_q = _mm256_set1_epi64x(static_cast(modulus_small)); + + gt_mask = _mm256_cmpgt_epi64(x, threshold_q); + to_subtract = _mm256_and_si256(gt_mask, val_q); x = _mm256_sub_epi64(x, to_subtract); + _mm256_storeu_si256(p_x, x); } } @@ -265,7 +260,7 @@ void NttInverse(const Params& params, absl::Span operand_overall) { #else void NttInverse(const Params& params, absl::Span operand_overall) { - SPDLOG_DEBUG("use AVX2 NttInverse"); + SPDLOG_DEBUG("using AVX2 NttInverse"); for (size_t coeff_mod = 0; coeff_mod < params.CrtCount(); ++coeff_mod) { size_t n = params.PolyLen(); @@ -295,7 +290,9 @@ void NttInverse(const Params& params, absl::Span operand_overall) { uint64_t y = op[t + j]; uint64_t t_tmp = two_times_modulus - y + x; - uint64_t curr_x = x + y - (two_times_modulus * ((x << 1) >= t_tmp)); + uint64_t curr_x = + x + y - + (two_times_modulus * static_cast((x << 1) >= t_tmp)); uint64_t h_tmp = (t_tmp * w_prime) >> 32; uint64_t res_x = (curr_x + (modulus * (t_tmp & 1))) >> 1; diff --git a/psi/algorithm/spiral/arith/ntt.h b/psi/algorithm/spiral/arith/ntt.h index 65109b64..ac682d16 100644 --- a/psi/algorithm/spiral/arith/ntt.h +++ b/psi/algorithm/spiral/arith/ntt.h @@ -23,6 +23,7 @@ namespace psi::spiral::arith { void NttForward(const Params& params, absl::Span operand_overall); + void NttInverse(const Params& params, absl::Span operand_overall); } // namespace psi::spiral::arith diff --git a/psi/algorithm/spiral/params.cc b/psi/algorithm/spiral/params.cc index ce417e68..dfd51810 100644 --- a/psi/algorithm/spiral/params.cc +++ b/psi/algorithm/spiral/params.cc @@ -163,6 +163,7 @@ std::uint64_t Params::CrtCompose2(std::uint64_t x, std::uint64_t y) const { crt_params_.modulus_.value()); } +// TODO: optimize memory copy span? std::uint64_t Params::CrtCompose(const std::vector& a, std::size_t idx) const { if (crt_params_.crt_count_ == 1) { diff --git a/psi/algorithm/spiral/params.h b/psi/algorithm/spiral/params.h index 2ab57bd8..80ed4538 100644 --- a/psi/algorithm/spiral/params.h +++ b/psi/algorithm/spiral/params.h @@ -276,9 +276,9 @@ class Params { [[nodiscard]] size_t CrtCount() const { return crt_params_.crt_count_; } - [[nodiscard]] uint64_t DbDim2() const { return query_params_.db_dim2_; } + [[nodiscard]] size_t DbDim2() const { return query_params_.db_dim2_; } - [[nodiscard]] uint64_t DbDim1() const { return query_params_.db_dim1_; } + [[nodiscard]] size_t DbDim1() const { return query_params_.db_dim1_; } [[nodiscard]] uint64_t Modulus() const { return crt_params_.modulus_.value(); @@ -310,7 +310,15 @@ class Params { std::uint64_t CrtCompose(const std::vector& a, std::size_t idx) const; + uint64_t GetQPrime1() const { return 1 << 20; } + uint64_t GetQPrime2() const { + if (Q2Bits() == ModulusLog2()) { + return Modulus(); + } else { + return kQ2Values[Q2Bits()]; + } + } // other util methods [[nodiscard]] std::string ToString(); @@ -385,11 +393,16 @@ class Params { ParamsId Id() const { return id_; } - private: - void ComputeId(); + size_t DbRowsPadded() const { + return 1 << (query_params_.db_dim1_ + poly_len_log2_); + } + void SetDbDim1(size_t v1) { query_params_.db_dim1_ = v1; } void SetDbDim2(size_t v2) { query_params_.db_dim2_ = v2; } + private: + void ComputeId(); + // d in R = Z[x]/x^d + 1 std::size_t poly_len_ = 0; diff --git a/psi/algorithm/spiral/poly_matrix.cc b/psi/algorithm/spiral/poly_matrix.cc index 07079af6..28f67817 100644 --- a/psi/algorithm/spiral/poly_matrix.cc +++ b/psi/algorithm/spiral/poly_matrix.cc @@ -238,7 +238,39 @@ PolyMatrixRaw PolyMatrixRaw::RandomPrg(const Params& params, size_t rows, return out; } -//--------------- PolyMatrixNtt +PolyMatrixRaw PolyMatrixRaw::Recover(const Params& params, uint64_t q_1, + uint64_t q_2, + const std::vector& ciphertext) { + const size_t q_1_bits = + static_cast(std::ceil(std::log2(static_cast(q_2)))); + const size_t q_2_bits = + static_cast(std::ceil(std::log2(static_cast(q_1)))); + + const size_t total_sz_bits = (q_1_bits + q_2_bits) * params.PolyLen(); + const size_t total_sz_bytes = (total_sz_bits + 7) / 8; + YACL_ENFORCE_EQ(ciphertext.size(), total_sz_bytes, + "Ciphertext size mismatch"); + + PolyMatrixRaw res = PolyMatrixRaw::Zero(params.PolyLen(), 2, 1); + size_t bit_offs = 0; + + uint64_t* row_0_ptr = res.Data().data(); + uint64_t* row_1_ptr = res.Data().data() + params.PolyLen(); + + for (size_t z = 0; z < params.PolyLen(); ++z) { + uint64_t val = util::ReadArbitraryBits(ciphertext, bit_offs, q_1_bits); + row_0_ptr[z] = arith::Rescale(val, q_2, params.Modulus()); + bit_offs += q_1_bits; + } + + for (size_t z = 0; z < params.PolyLen(); ++z) { + uint64_t val = util::ReadArbitraryBits(ciphertext, bit_offs, q_2_bits); + row_1_ptr[z] = arith::Rescale(val, q_1, params.Modulus()); + bit_offs += q_2_bits; + } + + return res; +} void PolyMatrixNtt::CopyInto(const PolyMatrixNtt& p, size_t target_row, size_t target_col) { diff --git a/psi/algorithm/spiral/poly_matrix.h b/psi/algorithm/spiral/poly_matrix.h index 1314616e..114f79c9 100644 --- a/psi/algorithm/spiral/poly_matrix.h +++ b/psi/algorithm/spiral/poly_matrix.h @@ -170,6 +170,8 @@ class PolyMatrixRaw { static PolyMatrixRaw RandomPrg(const Params& params, size_t rows, size_t cols, yacl::crypto::Prg& prg); + static PolyMatrixRaw Recover(const Params& params, uint64_t q_1, uint64_t q_2, + const std::vector& ciphertext); private: size_t poly_len_ = 0; @@ -339,7 +341,7 @@ void AddPoly(const Params& params, absl::Span res, void AddPolyInto(const Params& params, absl::Span res, absl::Span a); // res = -a -void InvertPoly(const Params& params, absl::Span res, +void NegatePoly(const Params& params, absl::Span res, absl::Span a); void AutomotphPoly(const Params& params, absl::Span res, @@ -369,6 +371,8 @@ void Multiply(const Params& params, PolyMatrixNtt& res, const PolyMatrixNtt& a, PolyMatrixNtt Multiply(const Params& params, const PolyMatrixNtt& a, const PolyMatrixNtt& b); +void MultiplyNoReduce(PolyMatrixNtt& res, const PolyMatrixNtt& a, + const PolyMatrixNtt& b, size_t start_inner_dim); void Automorphism(const Params& params, PolyMatrixRaw& res, const PolyMatrixRaw& a, size_t t); @@ -385,11 +389,13 @@ void AddInto(const Params& params, PolyMatrixNtt& res, const PolyMatrixNtt& a); void AddIntoAt(const Params& params, PolyMatrixNtt& res, const PolyMatrixNtt& a, size_t t_row, size_t t_col); -void Invert(const Params& params, PolyMatrixRaw& res, const PolyMatrixRaw& a); -PolyMatrixRaw Invert(const Params& params, const PolyMatrixRaw& a); +void Negate(const Params& params, PolyMatrixRaw& res, const PolyMatrixRaw& a); +PolyMatrixRaw Negate(const Params& params, const PolyMatrixRaw& a); void FromNtt(const Params& params, PolyMatrixRaw& out, const PolyMatrixNtt& in); PolyMatrixRaw FromNtt(const Params& params, const PolyMatrixNtt& in); +void FromNttScratch(const Params& params, PolyMatrixRaw& out, + absl::Span scratch, const PolyMatrixNtt& in); void ToNtt(const Params& params, PolyMatrixNtt& out, const PolyMatrixRaw& in); PolyMatrixNtt ToNtt(const Params& params, const PolyMatrixRaw& in); diff --git a/psi/algorithm/spiral/poly_matrix_utils.cc b/psi/algorithm/spiral/poly_matrix_utils.cc index 46154b6f..ffb785b7 100644 --- a/psi/algorithm/spiral/poly_matrix_utils.cc +++ b/psi/algorithm/spiral/poly_matrix_utils.cc @@ -121,15 +121,15 @@ void AddPolyInto(const Params& params, absl::Span res, } } -void InvertPoly(const Params& params, absl::Span res, +void NegatePoly(const Params& params, absl::Span res, absl::Span a) { for (size_t i = 0; i < params.PolyLen(); ++i) { res[i] = params.Modulus() - a[i]; } } -void AutomotphPoly(const Params& params, absl::Span res, - absl::Span a, size_t t) { +void AutomorphismPoly(const Params& params, absl::Span res, + absl::Span a, size_t t) { auto poly_len = params.PolyLen(); for (size_t i = 0; i < poly_len; ++i) { uint64_t num = (i * t) / poly_len; @@ -142,6 +142,32 @@ void AutomotphPoly(const Params& params, absl::Span res, } } +void AutomorphismPolyUncrtd(const Params& params, absl::Span res, + absl::Span a, size_t t) { + size_t poly_len = params.PolyLen(); + size_t crt_count = params.CrtCount(); + for (size_t m = 0; m < crt_count; ++m) { + size_t offset = m * poly_len; + auto a_chunk = a.subspan(offset, poly_len); + auto res_chunk = res.subspan(offset, poly_len); + uint64_t current_modulus = params.Moduli(m); + for (size_t i = 0; i < poly_len; ++i) { + size_t num = (i * t) / poly_len; + size_t rem = (i * t) % poly_len; + + if (num % 2 == 0) { + res_chunk[rem] = a_chunk[i]; + } else { + if (a_chunk[i] == 0) { + res_chunk[rem] = 0; + } else { + res_chunk[rem] = current_modulus - a_chunk[i]; + } + } + } + } +} + void ReduceCopy(const Params& params, absl::Span res, absl::Span in) { for (size_t i = 0; i < params.CrtCount(); ++i) { @@ -250,7 +276,7 @@ void Automorphism(const Params& params, PolyMatrixRaw& res, for (size_t j = 0; j < a.Cols(); ++j) { auto poly_a = a.Poly(i, j); auto poly_res = res.Poly(i, j); - AutomotphPoly(params, poly_res, poly_a, t); + AutomorphismPoly(params, poly_res, poly_a, t); } } } @@ -320,7 +346,7 @@ void AddIntoAt(const Params& params, PolyMatrixNtt& res, const PolyMatrixNtt& a, } } -void Invert(const Params& params, PolyMatrixRaw& res, const PolyMatrixRaw& a) { +void Negate(const Params& params, PolyMatrixRaw& res, const PolyMatrixRaw& a) { WEAK_ENFORCE(res.Rows() == a.Rows()); WEAK_ENFORCE(res.Cols() == a.Cols()); @@ -328,18 +354,18 @@ void Invert(const Params& params, PolyMatrixRaw& res, const PolyMatrixRaw& a) { for (size_t j = 0; j < a.Cols(); ++j) { auto res_poly = res.Poly(i, j); auto a_poly = a.Poly(i, j); - InvertPoly(params, res_poly, a_poly); + NegatePoly(params, res_poly, a_poly); } } } -PolyMatrixRaw Invert(const Params& params, const PolyMatrixRaw& a) { +PolyMatrixRaw Negate(const Params& params, const PolyMatrixRaw& a) { PolyMatrixRaw res(params.PolyLen(), a.Rows(), a.Cols()); for (size_t i = 0; i < a.Rows(); ++i) { for (size_t j = 0; j < a.Cols(); ++j) { auto res_poly = res.Poly(i, j); auto a_poly = a.Poly(i, j); - InvertPoly(params, res_poly, a_poly); + NegatePoly(params, res_poly, a_poly); } } return res; @@ -375,6 +401,34 @@ PolyMatrixRaw FromNtt(const Params& params, const PolyMatrixNtt& in) { return res; } +void FromNttScratch(const Params& params, PolyMatrixRaw& out, + absl::Span scratch, const PolyMatrixNtt& in) { + YACL_ENFORCE(in.Rows() == 2, "in rows must be 2, got {}", in.Rows()); + YACL_ENFORCE(in.Cols() == 1, "in cols must be 1, got {}", in.Cols()); + + size_t poly_size = in.Poly(0, 0).size(); + YACL_ENFORCE(scratch.size() >= poly_size, "scratch size too small: {} < {}", + scratch.size(), poly_size); + + for (size_t r = 0; r < in.Rows(); ++r) { + size_t c = 0; + + auto in_poly = in.Poly(r, c); + auto scratch_sub = scratch.subspan(0, poly_size); + std::copy(in_poly.begin(), in_poly.end(), scratch_sub.begin()); + + arith::NttInverse(params, scratch_sub); + + if (r == 0) { + size_t raw_poly_idx = out.PolyStartIndex(r, c); + std::vector temp_vec(scratch_sub.begin(), scratch_sub.end()); + for (size_t i = 0; i < params.PolyLen(); ++i) { + out.Data()[raw_poly_idx + i] = params.CrtCompose(temp_vec, i); + } + } + } +} + void ToNtt(const Params& params, PolyMatrixNtt& out, const PolyMatrixRaw& in) { for (size_t r = 0; r < out.Rows(); ++r) { for (size_t c = 0; c < out.Cols(); ++c) { @@ -456,6 +510,32 @@ PolyMatrixNtt Multiply(const Params& params, const PolyMatrixNtt& a, return res; } +void MultiplyNoReduce(PolyMatrixNtt& res, const PolyMatrixNtt& a, + const PolyMatrixNtt& b, size_t start_inner_dim) { + WEAK_ENFORCE(res.Rows() == a.Rows(), "Row mismatch: res={}, a={}", res.Rows(), + a.Rows()); + WEAK_ENFORCE(res.Cols() == b.Cols(), "Col mismatch: res={}, b={}", res.Cols(), + b.Cols()); + WEAK_ENFORCE(a.Cols() == b.Rows(), "Inner dim mismatch: a.cols={}, b.rows={}", + a.Cols(), b.Rows()); + + for (size_t i = 0; i < a.Rows(); ++i) { + for (size_t j = 0; j < b.Cols(); ++j) { + auto res_poly = res.Poly(i, j); + for (size_t k = start_inner_dim; k < a.Cols(); ++k) { + auto a_poly = a.Poly(i, k); + auto b_poly = b.Poly(k, j); + + size_t total_len = res_poly.size(); + + for (size_t z = 0; z < total_len; ++z) { + res_poly[z] += a_poly[z] * b_poly[z]; + } + } + } + } +} + PolyMatrixRaw MatrixWithIdentity(const PolyMatrixRaw& p) { WEAK_ENFORCE(p.Cols() == 1U); diff --git a/psi/algorithm/spiral/poly_matrix_utils.h b/psi/algorithm/spiral/poly_matrix_utils.h index 5910c59a..4f16103e 100644 --- a/psi/algorithm/spiral/poly_matrix_utils.h +++ b/psi/algorithm/spiral/poly_matrix_utils.h @@ -45,11 +45,13 @@ void AddPoly(const Params& params, absl::Span res, void AddPolyInto(const Params& params, absl::Span res, absl::Span a); // res = -a -void InvertPoly(const Params& params, absl::Span res, +void NegatePoly(const Params& params, absl::Span res, absl::Span a); -void AutomotphPoly(const Params& params, absl::Span res, - absl::Span a, size_t t); +void AutomorphismPoly(const Params& params, absl::Span res, + absl::Span a, size_t t); +void AutomorphismPolyUncrtd(const Params& params, absl::Span res, + absl::Span a, size_t t); // in is a PolyRaw, res is a RNS void ReduceCopy(const Params& params, absl::Span res, @@ -83,6 +85,9 @@ void Multiply(const Params& params, PolyMatrixNtt& res, const PolyMatrixNtt& a, PolyMatrixNtt Multiply(const Params& params, const PolyMatrixNtt& a, const PolyMatrixNtt& b); +void MultiplyNoReduce(PolyMatrixNtt& res, const PolyMatrixNtt& a, + const PolyMatrixNtt& b, size_t start_inner_dim); + void Automorphism(const Params& params, PolyMatrixRaw& res, const PolyMatrixRaw& a, size_t t); PolyMatrixRaw Automorphism(const Params& params, const PolyMatrixRaw& a, @@ -98,11 +103,13 @@ void AddInto(const Params& params, PolyMatrixNtt& res, const PolyMatrixNtt& a); void AddIntoAt(const Params& params, PolyMatrixNtt& res, const PolyMatrixNtt& a, size_t t_row, size_t t_col); -void Invert(const Params& params, PolyMatrixRaw& res, const PolyMatrixRaw& a); -PolyMatrixRaw Invert(const Params& params, const PolyMatrixRaw& a); +void Negate(const Params& params, PolyMatrixRaw& res, const PolyMatrixRaw& a); +PolyMatrixRaw Negate(const Params& params, const PolyMatrixRaw& a); void FromNtt(const Params& params, PolyMatrixRaw& out, const PolyMatrixNtt& in); PolyMatrixRaw FromNtt(const Params& params, const PolyMatrixNtt& in); +void FromNttScratch(const Params& params, PolyMatrixRaw& out, + absl::Span scratch, const PolyMatrixNtt& in); void ToNtt(const Params& params, PolyMatrixNtt& out, const PolyMatrixRaw& in); PolyMatrixNtt ToNtt(const Params& params, const PolyMatrixRaw& in); diff --git a/psi/algorithm/spiral/spiral_client.cc b/psi/algorithm/spiral/spiral_client.cc index 383ed9f2..d66761c9 100644 --- a/psi/algorithm/spiral/spiral_client.cc +++ b/psi/algorithm/spiral/spiral_client.cc @@ -20,6 +20,7 @@ #include "yacl/utils/parallel.h" #include "psi/algorithm/pir_interface/pir_db.h" +#include "psi/algorithm/spiral/arith/arith.h" #include "psi/algorithm/spiral/arith/ntt.h" #include "psi/algorithm/spiral/common.h" #include "psi/algorithm/spiral/gadget.h" @@ -115,8 +116,8 @@ void SpiralClient::Init() { PolyMatrixRaw::Zero(params_.PolyLen(), sk_regev_rows, sk_regev_cols); dg_ = DiscreteGaussian(params_.NoiseWidth()); - // Gen secret key for SpiralClient - GenSecretKeys(); + // Note: Secret keys are NOT generated here anymore. + // Constructors must call GenSecretKeys() explicitly. } PolyMatrixRaw SpiralClient::GetFreshGswPublicKey( @@ -128,7 +129,7 @@ PolyMatrixRaw SpiralClient::GetFreshGswPublicKey( auto a_ntt = ToNtt(params_, a); auto e = Noise(params_, n, m, dg_, rng); auto e_ntt = ToNtt(params_, e); - auto a_inv = Invert(params_, a); + auto a_inv = Negate(params_, a); auto sk_gsw_ntt = ToNtt(params_, sk_gsw_); auto b_p = Multiply(params_, sk_gsw_ntt, a_ntt); auto b = Add(params_, e_ntt, b_p); @@ -142,7 +143,7 @@ PolyMatrixNtt SpiralClient::GetRegevSample( yacl::crypto::Prg& rng_pub) const { auto a = PolyMatrixRaw::RandomPrg(params_, 1, 1, rng_pub); auto a_ntt = ToNtt(params_, a); - auto a_inv = ToNtt(params_, Invert(params_, a)); + auto a_inv = ToNtt(params_, Negate(params_, a)); auto e = Noise(params_, 1, 1, dg_, rng); auto e_ntt = ToNtt(params_, e); @@ -156,6 +157,30 @@ PolyMatrixNtt SpiralClient::GetRegevSample( return p; } +PolyMatrixNtt SpiralClient::GetScaledRegevSample( + yacl::crypto::Prg& rng, yacl::crypto::Prg& rng_pub, + uint64_t scale) const { + auto a = PolyMatrixRaw::RandomPrg(params_, 1, 1, rng_pub); + auto a_ntt = ToNtt(params_, a); + auto a_inv = ToNtt(params_, Negate(params_, a)); + auto e = Noise(params_, 1, 1, dg_, rng); + + for (size_t i = 0; i < params_.PolyLen(); ++i) { + e.Data()[i] = arith::MultiplyUintMod(e.Data()[i], scale, params_.Modulus()); + } + + auto e_ntt = ToNtt(params_, e); + auto sk_reg_ntt = ToNtt(params_, sk_reg_); + auto b_p = Multiply(params_, sk_reg_ntt, a_ntt); + auto b = Add(params_, e_ntt, b_p); + + auto p = PolyMatrixNtt::Zero(params_.CrtCount(), params_.PolyLen(), 2, 1); + p.CopyInto(a_inv, 0, 0); + p.CopyInto(b, 1, 0); + + return p; +} + PolyMatrixNtt SpiralClient::GetFreshRegevPublicKey( size_t m, yacl::crypto::Prg& rng, yacl::crypto::Prg& rng_pub) const { @@ -166,6 +191,16 @@ PolyMatrixNtt SpiralClient::GetFreshRegevPublicKey( return p; } +PolyMatrixNtt SpiralClient::GetFreshScaledRegevPublicKey( + size_t m, yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub, uint64_t scale) const { + auto p = PolyMatrixNtt::Zero(params_.CrtCount(), params_.PolyLen(), 2, m); + for (size_t i = 0; i < m; ++i) { + p.CopyInto(GetScaledRegevSample(rng, rng_pub, scale), 0, i); + } + return p; +} + PolyMatrixNtt SpiralClient::EncryptMatrixGsw( PolyMatrixNtt& ag, yacl::crypto::Prg& rng, yacl::crypto::Prg& rng_pub) const { @@ -187,6 +222,15 @@ PolyMatrixNtt SpiralClient::EncryptMatrixRegev( return Add(params_, p, a.PadTop(1)); } +PolyMatrixNtt SpiralClient::EncryptMatrixScaledRegev( + PolyMatrixNtt& a, yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub, uint64_t scale) const { + YACL_ENFORCE(sk_inited_, "Secret Key must be inited"); + auto m = a.Cols(); + auto p = GetFreshScaledRegevPublicKey(m, rng, rng_pub, scale); + return Add(params_, p, a.PadTop(1)); +} + void SpiralClient::GenSecretKeys(yacl::crypto::Prg& rng) { GenTernaryMatrix(params_, sk_gsw_, kHammingWeight, rng); GenTernaryMatrix(params_, sk_reg_, kHammingWeight, rng); diff --git a/psi/algorithm/spiral/spiral_client.h b/psi/algorithm/spiral/spiral_client.h index 00d0e58c..79f3f475 100644 --- a/psi/algorithm/spiral/spiral_client.h +++ b/psi/algorithm/spiral/spiral_client.h @@ -60,7 +60,15 @@ struct SpiralQuery { class SpiralClient : public psi::pir::IndexPirClient { public: - explicit SpiralClient(Params params) : params_(std::move(params)) { Init(); } + explicit SpiralClient(Params params) : params_(std::move(params)) { + Init(); + GenSecretKeys(); // Use random seed + } + explicit SpiralClient(Params params, uint128_t seed) + : params_(std::move(params)) { + Init(); + GenSecretKeys(seed); + } SpiralClient(Params params, DatabaseMetaInfo database_info) : params_(std::move(params)), database_info_(database_info) { @@ -74,6 +82,7 @@ class SpiralClient : public psi::pir::IndexPirClient { params_.MaxByteLenOfPt(); element_size_of_pt_ = params_.ElementSizeOfPt(element_byte_len); + GenSecretKeys(); // Use random seed } PublicKeys GenPublicKeys() const; @@ -133,7 +142,7 @@ class SpiralClient : public psi::pir::IndexPirClient { pir::PirType GetPirType() const override { return pir::PirType::SPIRAL_PIR; } - protected: + // protected: SpiralQuery GenQueryInternal(size_t pt_idx_target) const; PolyMatrixRaw DecodeResponseInternal( @@ -148,9 +157,18 @@ class SpiralClient : public psi::pir::IndexPirClient { PolyMatrixNtt GetRegevSample(yacl::crypto::Prg& rng, yacl::crypto::Prg& rng_pub) const; + PolyMatrixNtt GetScaledRegevSample(yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub, + uint64_t scale) const; + PolyMatrixNtt GetFreshRegevPublicKey( size_t m, yacl::crypto::Prg& rng, yacl::crypto::Prg& rng_pub) const; + + PolyMatrixNtt GetFreshScaledRegevPublicKey( + size_t m, yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub, uint64_t scale) const; + PolyMatrixNtt DecryptMatrixRegev(const PolyMatrixNtt& a) const { auto sk_reg_full_ntt = ToNtt(GetParams(), sk_reg_full_); return Multiply(GetParams(), sk_reg_full_ntt, a); @@ -169,6 +187,11 @@ class SpiralClient : public psi::pir::IndexPirClient { yacl::crypto::Prg& rng, yacl::crypto::Prg& rng_pub) const; + PolyMatrixNtt EncryptMatrixScaledRegev(PolyMatrixNtt& a, + yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub, + uint64_t scale) const; + private: // core implementation void GenSecretKeys(yacl::crypto::Prg& rng); diff --git a/psi/algorithm/spiral/spiral_server.cc b/psi/algorithm/spiral/spiral_server.cc index 49a4b233..0dc5e1f4 100644 --- a/psi/algorithm/spiral/spiral_server.cc +++ b/psi/algorithm/spiral/spiral_server.cc @@ -618,7 +618,7 @@ std::vector SpiralServer::GetVFoldingNeg( // -C auto ct_gsw_inv = PolyMatrixRaw::Zero(params_.PolyLen(), 2, 2 * params_.TGsw()); - Invert(params_, ct_gsw_inv, FromNtt(params_, v_folding[i])); + Negate(params_, ct_gsw_inv, FromNtt(params_, v_folding[i])); // G_{n+1, z} - C auto ct_gsw_neg = PolyMatrixNtt::Zero(params_.CrtCount(), params_.PolyLen(), 2, 2 * params_.TGsw()); @@ -963,7 +963,7 @@ std::vector GetVneg1(const Params& params) { auto idx = params.PolyLen() - (static_cast(1) << j); auto ng1 = PolyMatrixRaw::Zero(params.PolyLen(), 1, 1); ng1.Data()[idx] = 1ULL; - v_neg1.push_back(ToNtt(params, Invert(params, ng1))); + v_neg1.push_back(ToNtt(params, Negate(params, ng1))); } return v_neg1; } diff --git a/psi/algorithm/spiral/spiral_server_test.cc b/psi/algorithm/spiral/spiral_server_test.cc index 850ad90d..cfe876ba 100644 --- a/psi/algorithm/spiral/spiral_server_test.cc +++ b/psi/algorithm/spiral/spiral_server_test.cc @@ -382,7 +382,7 @@ TEST(SpiralServer, FoldCiphertexts) { auto ct_gsw_inv = PolyMatrixRaw::Zero(params.PolyLen(), 2, 2 * params.TGsw()); for (size_t i = 0; i < params.DbDim2(); ++i) { - Invert(params, ct_gsw_inv, FromNtt(params, v_folding[i])); + Negate(params, ct_gsw_inv, FromNtt(params, v_folding[i])); auto ct_gsw_neg = PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 2, 2 * params.TGsw()); diff --git a/psi/algorithm/spiral/util.cc b/psi/algorithm/spiral/util.cc index a8202bb1..520b6689 100644 --- a/psi/algorithm/spiral/util.cc +++ b/psi/algorithm/spiral/util.cc @@ -210,4 +210,66 @@ std::vector ConvertBytesToU8Coeffs( return coeff_array; } +uint64_t ReadArbitraryBits(const std::vector& buffer, + size_t bit_offset, size_t num_bits) { + assert(num_bits > 0 && num_bits <= 64 && "num_bits must be between 1 and 64"); + + const size_t word_offset = bit_offset / 64; + const size_t bit_offset_within_word = bit_offset % 64; + const size_t byte_offset = word_offset * 8; + + if (bit_offset_within_word + num_bits <= 64) { + assert(byte_offset + 8 <= buffer.size() && "Buffer read-out-of-bounds"); + + uint64_t val; + memcpy(&val, buffer.data() + byte_offset, sizeof(uint64_t)); + const uint64_t mask = (num_bits == 64) ? ~0ULL : (1ULL << num_bits) - 1; + return (val >> bit_offset_within_word) & mask; + + } else { + assert(byte_offset + 16 <= buffer.size() && "Buffer read-out-of-bounds"); + + uint128_t val; + memcpy(&val, buffer.data() + byte_offset, sizeof(uint128_t)); + + const uint128_t mask = (static_cast(1) << num_bits) - 1; + + return static_cast((val >> bit_offset_within_word) & mask); + } +} + +void WriteArbitraryBits(std::vector& buffer, uint64_t val, + size_t bit_offset, size_t num_bits) { + assert(num_bits > 0 && num_bits <= 64 && "num_bits must be between 1 and 64"); + + const size_t word_offset = bit_offset / 64; + const size_t bit_offset_within_word = bit_offset % 64; + const size_t byte_offset = word_offset * 8; + + const uint64_t val_mask = (num_bits == 64) ? ~0ULL : (1ULL << num_bits) - 1; + val &= val_mask; + + if (bit_offset_within_word + num_bits <= 64) { + assert(byte_offset + 8 <= buffer.size() && "Buffer write-out-of-bounds"); + + uint64_t current_val; + std::memcpy(¤t_val, buffer.data() + byte_offset, sizeof(uint64_t)); + const uint64_t write_mask = val_mask << bit_offset_within_word; + current_val &= ~write_mask; + current_val |= (val << bit_offset_within_word); + std::memcpy(buffer.data() + byte_offset, ¤t_val, sizeof(uint64_t)); + + } else { + assert(byte_offset + 16 <= buffer.size() && "Buffer write-out-of-bounds"); + + uint128_t current_val; + std::memcpy(¤t_val, buffer.data() + byte_offset, sizeof(uint128_t)); + + const uint128_t wide_val_mask = (static_cast(1) << num_bits) - 1; + const uint128_t write_mask = wide_val_mask << bit_offset_within_word; + current_val &= ~write_mask; + current_val |= (static_cast(val) << bit_offset_within_word); + std::memcpy(buffer.data() + byte_offset, ¤t_val, sizeof(uint128_t)); + } +} } // namespace psi::spiral::util diff --git a/psi/algorithm/spiral/util.h b/psi/algorithm/spiral/util.h index 252d6d05..50411b41 100644 --- a/psi/algorithm/spiral/util.h +++ b/psi/algorithm/spiral/util.h @@ -49,5 +49,8 @@ std::vector ConvertCoeffsToBytes( std::vector ConvertCoeffsToBytes(absl::Span coeff_array, size_t logt); - +uint64_t ReadArbitraryBits(const std::vector& buffer, + size_t bit_offset, size_t num_bits); +void WriteArbitraryBits(std::vector& buffer, uint64_t val, + size_t bit_offset, size_t num_bits); } // namespace psi::spiral::util diff --git a/psi/algorithm/ypir/BUILD.bazel b/psi/algorithm/ypir/BUILD.bazel new file mode 100644 index 00000000..81426b39 --- /dev/null +++ b/psi/algorithm/ypir/BUILD.bazel @@ -0,0 +1,275 @@ +# Copyright 2026 The secretflow authors. + +load("//bazel:psi.bzl", "psi_cc_library", "psi_cc_test") + +package(default_visibility = ["//visibility:public"]) + +X86_64_COMPATIBLE = ["@platforms//cpu:x86_64"] + +psi_cc_library( + name = "legacy_ypir_impl", + srcs = [ + "legacy/aes_prng.cc", + "legacy/client.cc", + "legacy/hexl.cc", + "legacy/server.cc", + "legacy/ypir_params.cc", + "legacy/ypir_util.cc", + ], + hdrs = [ + "legacy/aes_prng.h", + "legacy/client.h", + "legacy/hexl.h", + "legacy/server.h", + "legacy/ypir_params.h", + "legacy/ypir_util.h", + ], + copts = ["-maes"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + "@hexl", + "@yacl//yacl/base:exception", + ], +) + +psi_cc_library( + name = "legacy_simplepir", + srcs = [ + "legacy/matmul.cpp", + "legacy/util.cc", + ], + hdrs = [ + "legacy/types.h", + "legacy/util.h", + ], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + "//psi/algorithm/spiral:discrete_gaussian", + "//psi/algorithm/spiral:gadget", + "//psi/algorithm/spiral:params", + "//psi/algorithm/spiral:poly_matrix", + "//psi/algorithm/spiral:poly_matrix_utils", + "//psi/algorithm/spiral:util", + "@abseil-cpp//absl/types:span", + "@yacl//yacl/base:int128", + "@yacl//yacl/crypto/rand", + "@yacl//yacl/crypto/tools:prg", + "@yacl//yacl/utils:parallel", + ], +) + +psi_cc_library( + name = "types", + hdrs = ["types.h"], + target_compatible_with = X86_64_COMPATIBLE, +) + +psi_cc_library( + name = "params", + srcs = ["params.cc"], + hdrs = ["params.h"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":types", + "//psi/algorithm/spiral:params", + "@yacl//yacl/base:exception", + ], +) + +psi_cc_library( + name = "util", + srcs = ["util.cc"], + hdrs = ["util.h"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":params", + "//psi/algorithm/spiral/arith", + "@yacl//yacl/base:exception", + ], +) + +psi_cc_library( + name = "serialize", + srcs = ["serialize.cc"], + hdrs = ["serialize.h"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":types", + "@yacl//yacl/base:buffer", + "@yacl//yacl/base:byte_container_view", + "@yacl//yacl/base:exception", + ], +) + +psi_cc_library( + name = "ypir_internal_params", + srcs = ["ypir_internal_params.cc"], + hdrs = ["ypir_internal_params.h"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":legacy_ypir_impl", + ":params", + "@yacl//yacl/base:exception", + ], +) + +psi_cc_library( + name = "ypir_internal_client", + srcs = ["ypir_internal_client.cc"], + hdrs = ["ypir_internal_client.h"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":legacy_ypir_impl", + ":types", + ":util", + ":ypir_internal_params", + "@yacl//yacl/base:exception", + ], +) + +psi_cc_library( + name = "ypir_internal_server", + srcs = ["ypir_internal_server.cc"], + hdrs = ["ypir_internal_server.h"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":legacy_ypir_impl", + ":types", + ":ypir_internal_params", + "@yacl//yacl/base:exception", + ], +) + +psi_cc_library( + name = "client", + srcs = ["client.cc"], + hdrs = ["client.h"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":params", + ":serialize", + ":util", + ":ypir_internal_client", + "//psi/algorithm/pir_interface:index_pir", + "//psi/algorithm/pir_interface:pir_db", + "@yacl//yacl/base:buffer", + "@yacl//yacl/base:byte_container_view", + "@yacl//yacl/base:exception", + ], +) + +psi_cc_library( + name = "server", + srcs = ["server.cc"], + hdrs = ["server.h"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":legacy_simplepir", + ":params", + ":serialize", + ":ypir_internal_server", + "//psi/algorithm/pir_interface:pir_db", + "@yacl//yacl/base:buffer", + "@yacl//yacl/base:byte_container_view", + "@yacl//yacl/base:exception", + ], +) + +psi_cc_library( + name = "entry", + srcs = ["entry.cc"], + hdrs = ["entry.h"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":client", + ":params", + ":server", + "//psi/algorithm/pir_interface:pir_db", + "@abseil-cpp//absl/strings", + "@yacl//yacl/link:context", + ], +) + +psi_cc_test( + name = "smoke_test", + srcs = ["smoke_test.cc"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":client", + ":server", + ], +) + +psi_cc_test( + name = "params_test", + srcs = ["params_test.cc"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [":params"], +) + +psi_cc_test( + name = "serialize_test", + srcs = ["serialize_test.cc"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [":serialize"], +) + +psi_cc_test( + name = "client_test", + srcs = ["client_test.cc"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":client", + ":serialize", + ":server", + ], +) + +psi_cc_test( + name = "server_test", + srcs = ["server_test.cc"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [":server"], +) + +psi_cc_test( + name = "pir_interface_test", + srcs = ["pir_interface_test.cc"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":client", + ":server", + "//psi/algorithm/pir_interface:index_pir", + ], +) + +psi_cc_test( + name = "simplepir_flow_test", + srcs = ["simplepir_flow_test.cc"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":client", + ":server", + ], +) + +psi_cc_test( + name = "ypir_flow_test", + srcs = ["ypir_flow_test.cc"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":client", + ":server", + ], +) + +psi_cc_test( + name = "entry_test", + srcs = ["entry_test.cc"], + target_compatible_with = X86_64_COMPATIBLE, + deps = [ + ":entry", + "@abseil-cpp//absl/strings", + "@yacl//yacl/link:test_util", + ], +) diff --git a/psi/algorithm/ypir/client.cc b/psi/algorithm/ypir/client.cc new file mode 100644 index 00000000..e4180731 --- /dev/null +++ b/psi/algorithm/ypir/client.cc @@ -0,0 +1,80 @@ +#include "psi/algorithm/ypir/client.h" + +#include +#include + +#include "yacl/base/exception.h" + +#include "psi/algorithm/ypir/serialize.h" +#include "psi/algorithm/ypir/util.h" + +namespace psi::ypir { + +YpirClient::YpirClient(YpirParameters params) : params_(std::move(params)) { + if (params_.mode == YpirMode::kDoublepir) { + ypir_context_ = std::make_unique( + internal::ypir::CreateContext(params_)); + } +} + +yacl::Buffer YpirClient::GeneratePksBuffer() const { return yacl::Buffer(); } + +std::string YpirClient::GeneratePksString() const { return {}; } + +YpirQuery YpirClient::GenerateQuery(uint64_t raw_idx) const { + YACL_ENFORCE_LT(raw_idx, params_.NumItems()); + + if (params_.mode == YpirMode::kSimplepir) { + YpirQuery query; + query.mode = YpirMode::kSimplepir; + query.packed_query_row = BuildPackedSimplepirQuery(params_, raw_idx); + return query; + } + + YACL_ENFORCE(ypir_context_ != nullptr); + return internal::ypir::GenerateQuery(raw_idx, params_, client_secrets_, + *ypir_context_); +} + +yacl::Buffer YpirClient::GenerateQueryBuffer(uint64_t raw_idx) const { + return SerializeQuery(GenerateQuery(raw_idx)); +} + +yacl::Buffer YpirClient::GenerateIndexQuery(uint64_t raw_idx) const { + return GenerateQueryBuffer(raw_idx); +} + +std::string YpirClient::GenerateIndexQueryStr(uint64_t raw_idx) const { + auto buffer = GenerateQueryBuffer(raw_idx); + return std::string(static_cast(buffer)); +} + +std::vector YpirClient::DecodeResponse(const YpirResponse& response, + uint64_t raw_idx) const { + YACL_ENFORCE_LT(raw_idx, params_.NumItems()); + YACL_ENFORCE(response.mode == params_.mode); + + if (params_.mode == YpirMode::kSimplepir) { + YACL_ENFORCE_EQ(response.simplepir_response.size(), params_.db_cols); + const uint64_t col_idx = raw_idx % params_.db_cols; + const uint64_t decoded = + DecodeSimplepirValue(params_, response.simplepir_response[col_idx]); + return EncodeIntegerValue(decoded, params_.value_bytes); + } + + YACL_ENFORCE(ypir_context_ != nullptr); + return internal::ypir::RecoverResponse(response, params_, client_secrets_, + *ypir_context_); +} + +std::vector YpirClient::DecodeResponseBuffer( + const yacl::ByteContainerView& response_buffer, uint64_t raw_idx) const { + return DecodeResponse(DeserializeResponse(response_buffer), raw_idx); +} + +std::vector YpirClient::DecodeIndexResponse( + const yacl::ByteContainerView& response_buffer, uint64_t raw_idx) const { + return DecodeResponseBuffer(response_buffer, raw_idx); +} + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/client.h b/psi/algorithm/ypir/client.h new file mode 100644 index 00000000..0be8ba3c --- /dev/null +++ b/psi/algorithm/ypir/client.h @@ -0,0 +1,47 @@ +#pragma once + +#include +#include +#include + +#include "yacl/base/buffer.h" +#include "yacl/base/byte_container_view.h" + +#include "psi/algorithm/pir_interface/index_pir.h" +#include "psi/algorithm/ypir/params.h" +#include "psi/algorithm/ypir/types.h" +#include "psi/algorithm/ypir/ypir_internal_client.h" + +namespace psi::ypir { + +class YpirClient : public psi::pir::IndexPirClient { + public: + explicit YpirClient(YpirParameters params); + + const YpirParameters& GetParameters() const { return params_; } + + pir::PirType GetPirType() const override { return pir::PirType::YPIR_PIR; } + + yacl::Buffer GeneratePksBuffer() const override; + std::string GeneratePksString() const override; + + YpirQuery GenerateQuery(uint64_t raw_idx) const; + yacl::Buffer GenerateQueryBuffer(uint64_t raw_idx) const; + yacl::Buffer GenerateIndexQuery(uint64_t raw_idx) const override; + std::string GenerateIndexQueryStr(uint64_t raw_idx) const override; + + std::vector DecodeResponse(const YpirResponse& response, + uint64_t raw_idx) const; + std::vector DecodeResponseBuffer( + const yacl::ByteContainerView& response_buffer, uint64_t raw_idx) const; + std::vector DecodeIndexResponse( + const yacl::ByteContainerView& response_buffer, + uint64_t raw_idx) const override; + + private: + YpirParameters params_; + mutable std::unique_ptr ypir_context_; + mutable internal::ypir::ClientSecrets client_secrets_; +}; + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/client_test.cc b/psi/algorithm/ypir/client_test.cc new file mode 100644 index 00000000..47ff5e0e --- /dev/null +++ b/psi/algorithm/ypir/client_test.cc @@ -0,0 +1,49 @@ +#include "psi/algorithm/ypir/client.h" + +#include + +#include "gtest/gtest.h" + +#include "psi/algorithm/ypir/serialize.h" + +namespace psi::ypir { +namespace { + +TEST(YpirClientTest, SimplepirQueryUsesPackedRowShape) { + auto params = CreateSmallTestParamsSimplePIR(); + YpirClient client(params); + + const auto query = client.GenerateQuery(123ULL * params.db_cols + 456ULL); + EXPECT_EQ(query.mode, YpirMode::kSimplepir); + EXPECT_EQ(query.packed_query_row.size(), params.db_rows); + EXPECT_TRUE(query.qu0.empty()); + EXPECT_TRUE(query.qu1.empty()); + EXPECT_TRUE(query.ksk_b.empty()); + + const auto buffer = client.GenerateQueryBuffer(7); + const auto decoded = DeserializeQuery(buffer); + EXPECT_EQ(decoded.mode, YpirMode::kSimplepir); + EXPECT_EQ(decoded.packed_query_row.size(), params.db_rows); +} + +TEST(YpirClientTest, DoublepirQueryContainsLegacyInternalPayload) { + auto params = CreateSmallTestParamsDoublePIR(); + YpirClient client(params); + + const auto query = client.GenerateQuery(111ULL * params.db_cols + 222ULL); + EXPECT_EQ(query.mode, YpirMode::kDoublepir); + EXPECT_TRUE(query.packed_query_row.empty()); + EXPECT_FALSE(query.qu0.empty()); + EXPECT_FALSE(query.qu1.empty()); + EXPECT_FALSE(query.ksk_b.empty()); + + const auto buffer = client.GenerateQueryBuffer(3); + const auto decoded = DeserializeQuery(buffer); + EXPECT_EQ(decoded.mode, YpirMode::kDoublepir); + EXPECT_FALSE(decoded.qu0.empty()); + EXPECT_FALSE(decoded.qu1.empty()); + EXPECT_FALSE(decoded.ksk_b.empty()); +} + +} // namespace +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/entry.cc b/psi/algorithm/ypir/entry.cc new file mode 100644 index 00000000..916644e3 --- /dev/null +++ b/psi/algorithm/ypir/entry.cc @@ -0,0 +1,243 @@ +#include "psi/algorithm/ypir/entry.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/strings/escaping.h" +#include "fmt/format.h" +#include "spdlog/spdlog.h" +#include "yacl/base/byte_container_view.h" +#include "yacl/base/exception.h" + +#include "psi/algorithm/pir_interface/pir_db.h" +#include "psi/algorithm/ypir/client.h" +#include "psi/algorithm/ypir/params.h" +#include "psi/algorithm/ypir/server.h" + +namespace psi::ypir { +namespace { + +constexpr char kQueryCountTag[] = "ypir/query_count"; + +std::string TrimAscii(std::string_view text) { + size_t begin = 0; + size_t end = text.size(); + while (begin < end && + std::isspace(static_cast(text[begin])) != 0) { + ++begin; + } + while (end > begin && + std::isspace(static_cast(text[end - 1])) != 0) { + --end; + } + return std::string(text.substr(begin, end - begin)); +} + +uint8_t ParseHexNibble(char ch) { + if (ch >= '0' && ch <= '9') { + return static_cast(ch - '0'); + } + if (ch >= 'a' && ch <= 'f') { + return static_cast(10 + ch - 'a'); + } + if (ch >= 'A' && ch <= 'F') { + return static_cast(10 + ch - 'A'); + } + YACL_THROW("invalid hex character: {}", ch); +} + +std::vector ParseHexBytes(std::string text, size_t value_bytes) { + if (text.size() >= 2 && text[0] == '0' && + (text[1] == 'x' || text[1] == 'X')) { + text = text.substr(2); + } + YACL_ENFORCE_EQ(text.size(), value_bytes * 2, + "hex value width mismatch, expect {} bytes", value_bytes); + std::vector out(value_bytes, 0); + for (size_t i = 0; i < value_bytes; ++i) { + out[i] = static_cast((ParseHexNibble(text[2 * i]) << 4) | + ParseHexNibble(text[2 * i + 1])); + } + return out; +} + +std::vector ReadLines(const std::string& path) { + std::ifstream input(path); + YACL_ENFORCE(input.is_open(), "failed to open file: {}", path); + + std::vector lines; + std::string line; + while (std::getline(input, line)) { + auto trimmed = TrimAscii(line); + if (!trimmed.empty()) { + lines.push_back(std::move(trimmed)); + } + } + return lines; +} + +void WriteLines(const std::string& path, + const std::vector& lines) { + std::ofstream output(path, std::ios::out | std::ios::trunc); + YACL_ENFORCE(output.is_open(), "failed to open output file: {}", path); + for (const auto& line : lines) { + output << line << '\n'; + } +} + +template +yacl::Buffer SerializeScalar(T value) { + yacl::Buffer out(sizeof(T)); + std::memcpy(out.data(), &value, sizeof(T)); + return out; +} + +template +T DeserializeScalar(const yacl::ByteContainerView& buffer) { + YACL_ENFORCE_EQ(buffer.size(), sizeof(T)); + T out = 0; + std::memcpy(&out, buffer.data(), sizeof(T)); + return out; +} + +std::string QueryTag(uint64_t idx) { return fmt::format("ypir/query/{}", idx); } + +std::string ResponseTag(uint64_t idx) { + return fmt::format("ypir/response/{}", idx); +} + +YpirParameters BuildParameters(YpirMode mode, uint64_t db_rows, + uint64_t db_cols, uint64_t item_size_bits) { + YACL_ENFORCE_GT(db_rows, 0U); + YACL_ENFORCE_GT(db_cols, 0U); + YACL_ENFORCE_GT(item_size_bits, 0U); + + if (mode == YpirMode::kSimplepir) { + return CreateParamsForShapeSimplePIR(db_rows, db_cols, item_size_bits); + } + return CreateParamsForShapeDoublePIR(db_rows, db_cols, item_size_bits); +} + +std::vector> LoadDatabaseRows(const std::string& db_file, + size_t value_bytes, + uint64_t num_items) { + const auto lines = ReadLines(db_file); + YACL_ENFORCE_EQ(lines.size(), num_items, + "db_file must contain exactly {} values", num_items); + + std::vector> rows; + rows.reserve(lines.size()); + for (const auto& line : lines) { + rows.push_back(ParseHexBytes(line, value_bytes)); + } + return rows; +} + +std::vector LoadQueryIndices(const std::string& query_file, + uint64_t num_items) { + const auto lines = ReadLines(query_file); + std::vector indices; + indices.reserve(lines.size()); + for (const auto& line : lines) { + uint64_t raw_idx = std::stoull(line); + YACL_ENFORCE_LT(raw_idx, num_items, "query index out of range"); + indices.push_back(raw_idx); + } + return indices; +} + +std::vector EncodeOutputLines( + const std::vector>& values) { + std::vector out; + out.reserve(values.size()); + for (const auto& value : values) { + out.push_back(absl::BytesToHexString(absl::string_view( + reinterpret_cast(value.data()), value.size()))); + } + return out; +} + +template +int RunSenderWithServer(const YpirParameters& params, + const YpirSenderOptions& options, + std::shared_ptr lctx) { + YpirServer server(params); + server.GenerateFromRawData(psi::pir::RawDatabase(LoadDatabaseRows( + options.db_file, params.value_bytes, params.NumItems()))); + + lctx->ConnectToMesh(); + const uint64_t query_count = + DeserializeScalar(lctx->Recv(lctx->NextRank(), kQueryCountTag)); + for (uint64_t idx = 0; idx < query_count; ++idx) { + auto query = lctx->Recv(lctx->NextRank(), QueryTag(idx)); + auto response = server.Response(query, yacl::Buffer()); + lctx->Send(lctx->NextRank(), response, ResponseTag(idx)); + } + return 0; +} + +} // namespace + +int SenderOnline(const YpirSenderOptions& options, + std::shared_ptr lctx) { + YACL_ENFORCE(lctx != nullptr, "link context is required for YPIR sender"); + const auto params = BuildParameters(options.mode, options.db_rows, + options.db_cols, options.item_size_bits); + + SPDLOG_INFO("Starting YPIR sender, mode={}, db_rows={}, db_cols={}", + params.mode == YpirMode::kSimplepir ? "simplepir" : "doublepir", + params.db_rows, params.db_cols); + + if (params.mode == YpirMode::kSimplepir) { + YACL_ENFORCE_LE( + params.value_bytes, sizeof(uint16_t), + "SimplePIR launcher currently supports up to 16-bit values"); + if (params.value_bytes == 1) { + return RunSenderWithServer(params, options, std::move(lctx)); + } + return RunSenderWithServer(params, options, std::move(lctx)); + } + + YACL_ENFORCE_EQ(params.value_bytes, 1U, + "DoublePIR launcher currently supports 8-bit values"); + return RunSenderWithServer(params, options, std::move(lctx)); +} + +int ReceiverOnline(const YpirReceiverOptions& options, + std::shared_ptr lctx) { + YACL_ENFORCE(lctx != nullptr, "link context is required for YPIR receiver"); + const auto params = BuildParameters(options.mode, options.db_rows, + options.db_cols, options.item_size_bits); + const auto query_indices = + LoadQueryIndices(options.query_file, params.NumItems()); + + SPDLOG_INFO("Starting YPIR receiver, mode={}, query_count={}", + params.mode == YpirMode::kSimplepir ? "simplepir" : "doublepir", + query_indices.size()); + + YpirClient client(params); + lctx->ConnectToMesh(); + lctx->Send(lctx->NextRank(), SerializeScalar(query_indices.size()), + kQueryCountTag); + + std::vector> decoded_values; + decoded_values.reserve(query_indices.size()); + for (size_t idx = 0; idx < query_indices.size(); ++idx) { + const auto raw_idx = query_indices[idx]; + auto query = client.GenerateIndexQuery(raw_idx); + lctx->Send(lctx->NextRank(), query, QueryTag(idx)); + auto response = lctx->Recv(lctx->NextRank(), ResponseTag(idx)); + decoded_values.push_back(client.DecodeIndexResponse(response, raw_idx)); + } + + WriteLines(options.output_file, EncodeOutputLines(decoded_values)); + return 0; +} + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/entry.h b/psi/algorithm/ypir/entry.h new file mode 100644 index 00000000..6c273573 --- /dev/null +++ b/psi/algorithm/ypir/entry.h @@ -0,0 +1,36 @@ +#pragma once + +#include +#include +#include + +#include "yacl/link/context.h" + +#include "psi/algorithm/ypir/types.h" + +namespace psi::ypir { + +struct YpirSenderOptions { + YpirMode mode = YpirMode::kSimplepir; + uint64_t db_rows = 0; + uint64_t db_cols = 0; + uint64_t item_size_bits = 0; + std::string db_file; +}; + +struct YpirReceiverOptions { + YpirMode mode = YpirMode::kSimplepir; + uint64_t db_rows = 0; + uint64_t db_cols = 0; + uint64_t item_size_bits = 0; + std::string query_file; + std::string output_file; +}; + +int SenderOnline(const YpirSenderOptions& options, + std::shared_ptr lctx); + +int ReceiverOnline(const YpirReceiverOptions& options, + std::shared_ptr lctx); + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/entry_test.cc b/psi/algorithm/ypir/entry_test.cc new file mode 100644 index 00000000..d210e3c7 --- /dev/null +++ b/psi/algorithm/ypir/entry_test.cc @@ -0,0 +1,165 @@ +#include "psi/algorithm/ypir/entry.h" + +#include +#include +#include +#include +#include + +#include "absl/strings/escaping.h" +#include "gtest/gtest.h" +#include "yacl/link/test_util.h" + +#include "psi/algorithm/ypir/params.h" + +namespace psi::ypir { +namespace { + +std::string HexByte(uint64_t value) { + std::vector bytes = {static_cast(value)}; + return absl::BytesToHexString(absl::string_view( + reinterpret_cast(bytes.data()), bytes.size())); +} + +std::string HexUint16(uint64_t value) { + const uint16_t narrowed = static_cast(value); + return absl::BytesToHexString(absl::string_view( + reinterpret_cast(&narrowed), sizeof(narrowed))); +} + +void WriteLines(const std::filesystem::path& path, + const std::vector& lines) { + std::ofstream output(path, std::ios::out | std::ios::trunc); + ASSERT_TRUE(output.is_open()); + for (const auto& line : lines) { + output << line << '\n'; + } +} + +std::vector ReadLines(const std::filesystem::path& path) { + std::ifstream input(path); + EXPECT_TRUE(input.is_open()); + std::vector lines; + std::string line; + while (std::getline(input, line)) { + if (!line.empty()) { + lines.push_back(line); + } + } + return lines; +} + +TEST(YpirEntryTest, SimplepirOnlineFlowWorksWithHexFiles) { + const auto params = CreateSmallTestParamsSimplePIR(); + const auto tmp_dir = + std::filesystem::temp_directory_path() / "ypir_entry_simplepir"; + std::filesystem::create_directories(tmp_dir); + + const auto db_path = tmp_dir / "db.hex"; + const auto query_path = tmp_dir / "query.txt"; + const auto output_path = tmp_dir / "result.hex"; + + std::vector db_lines; + db_lines.reserve(params.NumItems()); + for (uint64_t raw_idx = 0; raw_idx < params.NumItems(); ++raw_idx) { + const uint64_t row = raw_idx / params.db_cols; + const uint64_t col = raw_idx % params.db_cols; + db_lines.push_back( + HexUint16((row * 17 + col * 3) % params.spiral_params.PtModulus())); + } + WriteLines(db_path, db_lines); + WriteLines(query_path, {"0", "123456", "1048575"}); + + auto lctxs = yacl::link::test::SetupWorld(2); + YpirSenderOptions sender_options; + sender_options.mode = YpirMode::kSimplepir; + sender_options.db_rows = params.db_rows; + sender_options.db_cols = params.db_cols; + sender_options.item_size_bits = 16; + sender_options.db_file = db_path; + + YpirReceiverOptions receiver_options; + receiver_options.mode = YpirMode::kSimplepir; + receiver_options.db_rows = params.db_rows; + receiver_options.db_cols = params.db_cols; + receiver_options.item_size_bits = 16; + receiver_options.query_file = query_path; + receiver_options.output_file = output_path; + + auto sender = std::async(std::launch::async, [&] { + return SenderOnline(sender_options, lctxs[0]); + }); + auto receiver = std::async(std::launch::async, [&] { + return ReceiverOnline(receiver_options, lctxs[1]); + }); + + EXPECT_EQ(sender.get(), 0); + EXPECT_EQ(receiver.get(), 0); + + EXPECT_EQ( + ReadLines(output_path), + (std::vector{HexUint16(0), + HexUint16((120ULL * 17 + 576ULL * 3) % + params.spiral_params.PtModulus()), + HexUint16((1023ULL * 17 + 1023ULL * 3) % + params.spiral_params.PtModulus())})); + + std::filesystem::remove_all(tmp_dir); +} + +TEST(YpirEntryTest, DoublepirOnlineFlowWorksWithHexFiles) { + const auto params = CreateSmallTestParamsDoublePIR(); + const auto tmp_dir = + std::filesystem::temp_directory_path() / "ypir_entry_doublepir"; + std::filesystem::create_directories(tmp_dir); + + const auto db_path = tmp_dir / "db.hex"; + const auto query_path = tmp_dir / "query.txt"; + const auto output_path = tmp_dir / "result.hex"; + + std::vector db_lines; + db_lines.reserve(params.NumItems()); + for (uint64_t raw_idx = 0; raw_idx < params.NumItems(); ++raw_idx) { + const uint64_t row = raw_idx / params.db_cols; + const uint64_t col = raw_idx % params.db_cols; + db_lines.push_back(HexByte((row + col) % 251)); + } + WriteLines(db_path, db_lines); + WriteLines(query_path, {"0", "113886", "1048575"}); + + auto lctxs = yacl::link::test::SetupWorld(2); + YpirSenderOptions sender_options; + sender_options.mode = YpirMode::kDoublepir; + sender_options.db_rows = params.db_rows; + sender_options.db_cols = params.db_cols; + sender_options.item_size_bits = 8; + sender_options.db_file = db_path; + + YpirReceiverOptions receiver_options; + receiver_options.mode = YpirMode::kDoublepir; + receiver_options.db_rows = params.db_rows; + receiver_options.db_cols = params.db_cols; + receiver_options.item_size_bits = 8; + receiver_options.query_file = query_path; + receiver_options.output_file = output_path; + + auto sender = std::async(std::launch::async, [&] { + return SenderOnline(sender_options, lctxs[0]); + }); + auto receiver = std::async(std::launch::async, [&] { + return ReceiverOnline(receiver_options, lctxs[1]); + }); + + EXPECT_EQ(sender.get(), 0); + EXPECT_EQ(receiver.get(), 0); + + EXPECT_EQ( + ReadLines(output_path), + (std::vector{HexByte(0), HexByte((111ULL + 222ULL) % 251), + HexByte((1023ULL + 1023ULL) % 251)})); + + std::filesystem::remove_all(tmp_dir); +} + +} // namespace +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/legacy/aes_prng.cc b/psi/algorithm/ypir/legacy/aes_prng.cc new file mode 100644 index 00000000..43cfef8b --- /dev/null +++ b/psi/algorithm/ypir/legacy/aes_prng.cc @@ -0,0 +1,98 @@ +// Copyright 2026 The secretflow authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "psi/algorithm/ypir/legacy/aes_prng.h" + +namespace psi::ypir::ypir_internal { + +AESCTR_PRNG::AESCTR_PRNG() { + const uint8_t key[16] = {}; + aes128_key_expansion(key); + ctr_lo_ = 0; + ctr_hi_ = 0; +} + +AESCTR_PRNG::AESCTR_PRNG(const uint8_t key[16], uint64_t seed) { + aes128_key_expansion(key); + ctr_lo_ = seed; + ctr_hi_ = 0; +} + +void AESCTR_PRNG::refresh(uint64_t seed) { + ctr_lo_ = seed; + ctr_hi_ = 0; +} + +void AESCTR_PRNG::aes128_key_expansion(const uint8_t key[16]) { + __m128i tmp1, tmp2; + tmp1 = _mm_loadu_si128(reinterpret_cast(key)); + round_keys_[0] = tmp1; + + AES_128_ASSIST(tmp1, tmp2, 0x01); + round_keys_[1] = tmp1; + AES_128_ASSIST(tmp1, tmp2, 0x02); + round_keys_[2] = tmp1; + AES_128_ASSIST(tmp1, tmp2, 0x04); + round_keys_[3] = tmp1; + AES_128_ASSIST(tmp1, tmp2, 0x08); + round_keys_[4] = tmp1; + AES_128_ASSIST(tmp1, tmp2, 0x10); + round_keys_[5] = tmp1; + AES_128_ASSIST(tmp1, tmp2, 0x20); + round_keys_[6] = tmp1; + AES_128_ASSIST(tmp1, tmp2, 0x40); + round_keys_[7] = tmp1; + AES_128_ASSIST(tmp1, tmp2, 0x80); + round_keys_[8] = tmp1; + AES_128_ASSIST(tmp1, tmp2, 0x1B); + round_keys_[9] = tmp1; + AES_128_ASSIST(tmp1, tmp2, 0x36); + round_keys_[10] = tmp1; +} + +__m128i AESCTR_PRNG::aes128_encrypt_block(__m128i block, + const __m128i round_keys[11]) { + block = _mm_xor_si128(block, round_keys[0]); + for (int i = 1; i < 10; ++i) { + block = _mm_aesenc_si128(block, round_keys[i]); + } + block = _mm_aesenclast_si128(block, round_keys[10]); + return block; +} + +__m128i AESCTR_PRNG::next_block() { + __m128i ctr_block = _mm_set_epi64x(ctr_hi_, ctr_lo_); + __m128i out = aes128_encrypt_block(ctr_block, round_keys_); + if (++ctr_lo_ == 0) { + ++ctr_hi_; + } + return out; +} + +void AESCTR_PRNG::fill_bytes(uint8_t* dst, size_t len) { + while (len >= 16) { + __m128i b = next_block(); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst), b); + dst += 16; + len -= 16; + } + if (len > 0) { + __m128i b = next_block(); + alignas(16) uint8_t tmp[16]; + _mm_store_si128(reinterpret_cast<__m128i*>(tmp), b); + std::memcpy(dst, tmp, len); + } +} + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/aes_prng.h b/psi/algorithm/ypir/legacy/aes_prng.h new file mode 100644 index 00000000..b7e7bc00 --- /dev/null +++ b/psi/algorithm/ypir/legacy/aes_prng.h @@ -0,0 +1,51 @@ +// Copyright 2026 The secretflow authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +#include +#include +#include + +namespace psi::ypir::ypir_internal { + +#define AES_128_ASSIST(t1, t2, rc) \ + t2 = _mm_aeskeygenassist_si128(t1, rc); \ + t2 = _mm_shuffle_epi32(t2, _MM_SHUFFLE(3, 3, 3, 3)); \ + t1 = _mm_xor_si128(t1, _mm_slli_si128(t1, 4)); \ + t1 = _mm_xor_si128(t1, _mm_slli_si128(t1, 4)); \ + t1 = _mm_xor_si128(t1, _mm_slli_si128(t1, 4)); \ + t1 = _mm_xor_si128(t1, t2); + +class AESCTR_PRNG { + public: + AESCTR_PRNG(); + AESCTR_PRNG(const uint8_t key[16], uint64_t seed = 0); + + void refresh(uint64_t seed); + void fill_bytes(uint8_t* dst, size_t len); + + private: + void aes128_key_expansion(const uint8_t key[16]); + __m128i aes128_encrypt_block(__m128i block, const __m128i round_keys[11]); + __m128i next_block(); + + __m128i round_keys_[11]; + uint64_t ctr_lo_ = 0; + uint64_t ctr_hi_ = 0; +}; + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/client.cc b/psi/algorithm/ypir/legacy/client.cc new file mode 100644 index 00000000..9beb4fe0 --- /dev/null +++ b/psi/algorithm/ypir/legacy/client.cc @@ -0,0 +1,275 @@ +#include "psi/algorithm/ypir/legacy/client.h" + +#include +#include +#include +#include + +namespace psi::ypir::ypir_internal { +namespace { + +uint64_t MakeSeedMaterial() { + const uint64_t time_seed = static_cast( + std::chrono::high_resolution_clock::now().time_since_epoch().count()); + std::random_device rd; + return time_seed ^ (static_cast(rd()) << 32) ^ rd(); +} + +std::mt19937_64& GlobalPrng() { + thread_local std::mt19937_64 rng(MakeSeedMaterial()); + return rng; +} + +uint64_t GetBase(uint64_t b, uint64_t z, uint64_t ti) { + return 1ULL << (b + ti * z); +} + +void LweEncrypt(Secret& sk, const std::vector& a, uint64_t message, + uint64_t& b, uint64_t pmod, double sig) { + const uint64_t lwe_dimension = sk.get_len(); + const uint64_t cmod = sk.get_mod(); + const long double delta = + static_cast(cmod) / static_cast(pmod); + + const uint64_t e = SampleGauss(sig, cmod, GlobalPrng()); + + std::vector tmp(lwe_dimension, 0); + EltwiseMultMod(tmp.data(), a.data(), sk.data.data(), lwe_dimension, cmod); + b = 0; + for (uint64_t i = 0; i < lwe_dimension; ++i) { + b = (b + tmp[i]) % cmod; + } + b = (b + e) % cmod; + b = (b + static_cast(message * delta)) % cmod; +} + +void LweDecrypt(Secret& sk, std::vector& a, uint64_t& message, + uint64_t& b, uint64_t pmod) { + const uint64_t lwe_dimension = sk.get_len(); + const uint64_t cmod = sk.get_mod(); + const long double delta = + static_cast(cmod) / static_cast(pmod); + + std::vector tmp(lwe_dimension, 0); + EltwiseMultMod(tmp.data(), a.data(), sk.data.data(), lwe_dimension, cmod); + for (uint64_t i = 0; i < lwe_dimension; ++i) { + b = (cmod + b - tmp[i]) % cmod; + } + + if (b > (cmod >> 1)) { + b = static_cast(std::llround( + (static_cast(b) - static_cast(cmod)) / + delta)) % + pmod; + } else { + b = static_cast( + std::llround(static_cast(b) / delta)) % + pmod; + } + message = b; +} + +void PowerOfBase(const uint64_t* vec, uint64_t& num, uint64_t b, uint64_t z, + uint64_t t) { + num = 0; + for (uint64_t i = 0; i < t; ++i) { + num = (num << z) | vec[t - 1 - i]; + } + num <<= b; +} + +void RlweEncode(Secret& sk, const std::vector& a, + std::vector& message, std::vector& b, + const FheParams& fparm) { + const uint64_t poly_degree = fparm.get_poly_degree(); + const uint64_t rlwe_cmod = fparm.get_rlwe_cmod(); + const double sig = fparm.get_sig_ring(); + + YpirHexlNtt& ntt = fparm.get_ntt(); + if (!sk.get_ntt_form()) { + ntt.Forward(sk.data.data(), poly_degree); + sk.switch_ntt_format(); + } + + std::vector err(poly_degree, 0); + SampleGauss(err, sig, rlwe_cmod, GlobalPrng()); + EltwiseFMAMod(err.data(), message.data(), 1, err.data(), poly_degree, + rlwe_cmod); + ntt.Forward(err.data(), poly_degree); + + EltwiseMultMod(b.data(), sk.data.data(), a.data(), poly_degree, rlwe_cmod); + EltwiseAddMod(b.data(), err.data(), b.data(), poly_degree, rlwe_cmod); +} + +void RlweDecrypt(Secret& sk, std::vector& a, + std::vector& message, std::vector& b, + const FheParams& fparm) { + const uint64_t poly_degree = fparm.get_poly_degree(); + const uint64_t rlwe_cmod = fparm.get_rlwe_cmod(); + const uint64_t rlwe_pmod = fparm.get_rlwe_pmod(); + const uint64_t delta = rlwe_cmod / rlwe_pmod; + + YpirHexlNtt& ntt = fparm.get_ntt(); + if (!sk.get_ntt_form()) { + ntt.Forward(sk.data.data(), poly_degree); + sk.switch_ntt_format(); + } + + std::vector tmp(poly_degree, 0); + EltwiseMultMod(tmp.data(), sk.data.data(), a.data(), poly_degree, rlwe_cmod); + EltwiseSubMod(b.data(), b.data(), tmp.data(), poly_degree, rlwe_cmod); + ntt.Inverse(b.data(), poly_degree); + + message.assign(poly_degree, 0); + for (uint64_t i = 0; i < poly_degree; ++i) { + if (b[i] > (rlwe_cmod >> 1)) { + message[i] = static_cast( + std::llround((static_cast(b[i]) - + static_cast(rlwe_cmod)) / + static_cast(delta))) % + rlwe_pmod; + } else { + message[i] = + static_cast(std::llround(static_cast(b[i]) / + static_cast(delta))) % + rlwe_pmod; + } + } +} + +void MatrixReconstruct(const std::vector& vec, + std::vector>& mat, uint64_t row, + uint64_t col) { + mat.assign(row, std::vector(col, 0)); + uint64_t ptr = 0; + for (uint64_t i = 0; i < row; ++i) { + for (uint64_t j = 0; j < col; ++j, ++ptr) { + mat[i][j] = vec[ptr]; + } + } +} + +void GadgetEncrypt(Secret& sk, const std::vector>& a, + std::vector& message, + std::vector>& b, + const FheParams& fparm) { + const uint64_t b_auto = fparm.get_b_auto(); + const uint64_t z_auto = fparm.get_z_auto(); + const uint64_t t_auto = fparm.get_t_auto(); + const uint64_t poly_degree = fparm.get_poly_degree(); + const uint64_t rlwe_cmod = fparm.get_rlwe_cmod(); + + std::vector message_tmp(poly_degree, 0); + for (uint64_t i = 0; i < t_auto; ++i) { + const uint64_t base = GetBase(b_auto, z_auto, i) % rlwe_cmod; + EltwiseFMAMod(message_tmp.data(), message.data(), base, nullptr, + poly_degree, rlwe_cmod); + RlweEncode(sk, a[i], message_tmp, b[i], fparm); + } +} + +void GenerateAutokey( + Secret& sk, const std::vector>>& ksk_a, + std::vector>>& ksk_b, + const FheParams& fparm) { + const uint64_t degree = fparm.get_poly_degree(); + const uint64_t cmod = fparm.get_rlwe_cmod(); + const uint64_t expo = GetLog2(degree); + + YpirHexlNtt& ntt = fparm.get_ntt(); + if (sk.get_ntt_form()) { + ntt.Inverse(sk.data.data(), degree); + sk.switch_ntt_format(); + } + + std::vector newkey(degree, 0); + const std::vector copykey = sk.data; + for (uint64_t i = 0; i < expo; ++i) { + const uint64_t idx = (1ULL << (expo - i)) + 1; + ApplyAutoCoefForm(newkey, copykey, static_cast(idx), cmod); + GadgetEncrypt(sk, ksk_a[i], newkey, ksk_b[i], fparm); + } +} + +} // namespace + +void YpirRecover(Secret& simple_sk, Secret& double_sk, + std::vector>& res, uint64_t& message, + const FheParams& fparm, const PirParams&) { + const uint64_t degree = fparm.get_poly_degree(); + const uint64_t lwe_pmod = fparm.get_lwe_pmod(); + const uint64_t lwe_dimension = fparm.get_lwe_dimension(); + const uint64_t b = fparm.get_b_decomp(); + const uint64_t z = fparm.get_z_decomp(); + const uint64_t t = fparm.get_t_decomp(); + + std::vector pack_a_message(degree, 0); + std::vector pack_b_message(degree, 0); + RlweDecrypt(double_sk, res[0], pack_a_message, res[1], fparm); + RlweDecrypt(double_sk, res[2], pack_b_message, res[3], fparm); + + std::vector> matrix_pack_a_message; + MatrixReconstruct(pack_a_message, matrix_pack_a_message, lwe_dimension, t); + + std::vector a(lwe_dimension, 0); + for (uint64_t i = 0; i < lwe_dimension; ++i) { + PowerOfBase(matrix_pack_a_message[i].data(), a[i], b, z, t); + } + + uint64_t b_value = 0; + PowerOfBase(pack_b_message.data(), b_value, b, z, t); + LweDecrypt(simple_sk, a, message, b_value, lwe_pmod); +} + +YpirQuery Generate_query_ypir(uint64_t c_idx, uint64_t r_idx, Secret& lwe_sk, + Secret& rlwe_sk, AESCTR_PRNG& prng, + const FheParams& fparm, const PirParams& pparm) { + YpirQuery out; + + const uint64_t cols = pparm.get_col(); + std::vector query_vec_col(cols, 0); + query_vec_col[c_idx] = 1; + + out.qu0.resize(cols); + const double sig = fparm.get_sig(); + const uint64_t lwe_pmod = fparm.get_lwe_pmod(); + const auto& matrix0 = fparm.get_persudo_matrix_simplepir(); + for (uint64_t i = 0; i < cols; ++i) { + LweEncrypt(lwe_sk, matrix0[i], query_vec_col[i], out.qu0[i], lwe_pmod, sig); + } + + const uint64_t rows = pparm.get_row(); + const uint64_t degree = fparm.get_poly_degree(); + std::vector query_vec_row(rows, 0); + query_vec_row[r_idx] = 1; + + out.qu1.resize(rows); + const double sig_ring = fparm.get_sig_ring(); + const uint64_t rlwe_pmod = fparm.get_rlwe_pmod(); + std::vector> matrix1(rows, + std::vector(degree, 0)); + prng.refresh(kSecondDimensionSeed); + PseudorandomMatrixGenerate(matrix1, fparm.get_rlwe_cmod(), prng); + for (uint64_t i = 0; i < rows; ++i) { + LweEncrypt(rlwe_sk, matrix1[i], query_vec_row[i], out.qu1[i], rlwe_pmod, + sig_ring); + } + + const uint64_t t_auto = fparm.get_t_auto(); + const uint64_t expo = GetLog2(degree); + out.ksk_b.assign(expo, std::vector>( + t_auto, std::vector(degree, 0))); + + std::vector>> ksk_a( + expo, std::vector>( + t_auto, std::vector(degree, 0))); + prng.refresh(kThirdDimensionSeed); + for (uint64_t i = 0; i < expo; ++i) { + PseudorandomMatrixGenerate(ksk_a[i], fparm.get_rlwe_cmod(), prng); + } + GenerateAutokey(rlwe_sk, ksk_a, out.ksk_b, fparm); + + return out; +} + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/client.h b/psi/algorithm/ypir/legacy/client.h new file mode 100644 index 00000000..3fa963d2 --- /dev/null +++ b/psi/algorithm/ypir/legacy/client.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include + +#include "psi/algorithm/ypir/legacy/ypir_params.h" + +namespace psi::ypir::ypir_internal { + +struct YpirQuery { + std::vector qu0; + std::vector qu1; + std::vector>> ksk_b; +}; + +void YpirRecover(Secret& simple_sk, Secret& double_sk, + std::vector>& res, uint64_t& message, + const FheParams& fparm, const PirParams& pparm); + +YpirQuery Generate_query_ypir(uint64_t c_idx, uint64_t r_idx, Secret& lwe_sk, + Secret& rlwe_sk, AESCTR_PRNG& prng, + const FheParams& fparm, const PirParams& pparm); + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/hexl.cc b/psi/algorithm/ypir/legacy/hexl.cc new file mode 100644 index 00000000..86d056cb --- /dev/null +++ b/psi/algorithm/ypir/legacy/hexl.cc @@ -0,0 +1,60 @@ +// Copyright 2026 The secretflow authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "psi/algorithm/ypir/legacy/hexl.h" + +#include "yacl/base/exception.h" + +namespace psi::ypir::ypir_internal { + +YpirHexlNtt::YpirHexlNtt(uint64_t degree, uint64_t modulus, + uint64_t root_of_unity) + : ntt_(degree, modulus, root_of_unity), degree_(degree), modulus_(modulus) { + YACL_ENFORCE(degree_ > 0); + YACL_ENFORCE(modulus_ > 0); +} + +void YpirHexlNtt::Forward(uint64_t* data, size_t len) { + YACL_ENFORCE(len == degree_, "NTT Forward expects len == degree ({} != {})", + len, degree_); + ntt_.ComputeForward(data, data, 1, 1); +} + +void YpirHexlNtt::Inverse(uint64_t* data, size_t len) { + YACL_ENFORCE(len == degree_, "NTT Inverse expects len == degree ({} != {})", + len, degree_); + ntt_.ComputeInverse(data, data, 1, 1); +} + +void EltwiseMultMod(uint64_t* out, const uint64_t* a, const uint64_t* b, + size_t len, uint64_t modulus) { + intel::hexl::EltwiseMultMod(out, a, b, len, modulus, 1); +} + +void EltwiseAddMod(uint64_t* out, const uint64_t* a, const uint64_t* b, + size_t len, uint64_t modulus) { + intel::hexl::EltwiseAddMod(out, a, b, len, modulus); +} + +void EltwiseSubMod(uint64_t* out, const uint64_t* a, const uint64_t* b, + size_t len, uint64_t modulus) { + intel::hexl::EltwiseSubMod(out, a, b, len, modulus); +} + +void EltwiseFMAMod(uint64_t* out, const uint64_t* a, uint64_t scalar, + const uint64_t* add, size_t len, uint64_t modulus) { + intel::hexl::EltwiseFMAMod(out, a, scalar, add, len, modulus, 1); +} + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/hexl.h b/psi/algorithm/ypir/legacy/hexl.h new file mode 100644 index 00000000..cd0edfe3 --- /dev/null +++ b/psi/algorithm/ypir/legacy/hexl.h @@ -0,0 +1,53 @@ +// Copyright 2026 The secretflow authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +#include "hexl/hexl.hpp" + +namespace psi::ypir::ypir_internal { + +constexpr uint64_t kCrtQ1 = 268369921ULL; +constexpr uint64_t kCrtQ2 = 249561089ULL; +constexpr uint64_t kCrtMod = kCrtQ1 * kCrtQ2; +constexpr uint64_t kRootOfUnityCrt = 38878761190133527ULL; + +class YpirHexlNtt { + public: + YpirHexlNtt(uint64_t degree, uint64_t modulus, uint64_t root_of_unity); + + void Forward(uint64_t* data, size_t len); + void Inverse(uint64_t* data, size_t len); + + intel::hexl::NTT& Raw() { return ntt_; } + + private: + intel::hexl::NTT ntt_; + uint64_t degree_ = 0; + uint64_t modulus_ = 0; +}; + +void EltwiseMultMod(uint64_t* out, const uint64_t* a, const uint64_t* b, + size_t len, uint64_t modulus); +void EltwiseAddMod(uint64_t* out, const uint64_t* a, const uint64_t* b, + size_t len, uint64_t modulus); +void EltwiseSubMod(uint64_t* out, const uint64_t* a, const uint64_t* b, + size_t len, uint64_t modulus); +void EltwiseFMAMod(uint64_t* out, const uint64_t* a, uint64_t scalar, + const uint64_t* add, size_t len, uint64_t modulus); + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/matmul.cpp b/psi/algorithm/ypir/legacy/matmul.cpp new file mode 100644 index 00000000..b8114ef4 --- /dev/null +++ b/psi/algorithm/ypir/legacy/matmul.cpp @@ -0,0 +1,1556 @@ +/* +See https://github.com/ahenzinger/simplepir for original. + +MIT License + +Copyright (c) 2022, Alexandra Henzinger + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +// #define PROFILE + +#include +#include +#include + +#ifdef PROFILE +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +struct System { + static void profile(const std::string &name, std::function body) { + std::string filename = + name.find(".data") == std::string::npos ? (name + ".data") : name; + + // Launch profiler + pid_t pid; + std::stringstream s; + s << getpid(); + pid = fork(); + if (pid == 0) { + exit(execl("/usr/bin/perf", "perf", "stat", "-e", + "cache-references,cache-misses,cycles,instructions,branches," + "faults,migrations,l1d.replacement,l2_rqsts.all_demand_miss," + "cycle_activity.stalls_l3_miss", + "-o", filename.c_str(), "-p", s.str().c_str(), nullptr)); + } + + // Run body + body(); + + // Kill profiler + kill(pid, SIGINT); + waitpid(pid, nullptr, 0); + } + + static void profile(std::function body) { + profile("perf.data", body); + } +}; +#endif + +typedef uint32_t Elem; + +extern "C" { +void matMulVecPacked(uint32_t *out, const uint32_t *a, const uint32_t *b, + size_t aRows, size_t aCols); + +void matMulVecPacked2(Elem *out, const Elem *a, const Elem *b_full, + size_t aRows, size_t aCols); + +void matMulVecPacked4(Elem *out, const Elem *a, const Elem *b_full, + size_t aRows, size_t aCols); + +void matMulVecPacked6(Elem *out, const Elem *a, const Elem *b_full, + size_t aRows, size_t aCols); + +void matMulVecPacked8(Elem *out, const Elem *a, const Elem *b_full, + size_t aRows, size_t aCols); + +void matMulVecPacked8Alt(Elem *out, const Elem *a, const Elem *b_full, + size_t aRows, size_t aCols); + +void matMulVecPacked8Orig(Elem *out, const Elem *a, const Elem *b_full, + size_t aRows, size_t aCols); +} + +// Hard-coded, to allow for compiler optimizations: +#define COMPRESSION 4 +#define BASIS 8 +#define BASIS2 16 +#define BASIS3 24 +#define MASK 0xff + +void matMulVecPacked(uint32_t *out, const uint32_t *a, const uint32_t *b, + size_t aRows, size_t aCols) { + uint32_t db, db2, db3, db4, db5, db6, db7, db8; + uint32_t val, val2, val3, val4, val5, val6, val7, val8; + uint32_t tmp, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, tmp8; + size_t index = 0; + size_t index2; + + for (size_t i = 0; i < aRows; i += 8) { + tmp = 0; + tmp2 = 0; + tmp3 = 0; + tmp4 = 0; + tmp5 = 0; + tmp6 = 0; + tmp7 = 0; + tmp8 = 0; + + index2 = 0; + for (size_t j = 0; j < aCols; j++) { + db = a[index]; + db2 = a[index + 1 * aCols]; + db3 = a[index + 2 * aCols]; + db4 = a[index + 3 * aCols]; + db5 = a[index + 4 * aCols]; + db6 = a[index + 5 * aCols]; + db7 = a[index + 6 * aCols]; + db8 = a[index + 7 * aCols]; + + val = db & MASK; + val2 = db2 & MASK; + val3 = db3 & MASK; + val4 = db4 & MASK; + val5 = db5 & MASK; + val6 = db6 & MASK; + val7 = db7 & MASK; + val8 = db8 & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + index2 += 1; + + val = (db >> BASIS) & MASK; + val2 = (db2 >> BASIS) & MASK; + val3 = (db3 >> BASIS) & MASK; + val4 = (db4 >> BASIS) & MASK; + val5 = (db5 >> BASIS) & MASK; + val6 = (db6 >> BASIS) & MASK; + val7 = (db7 >> BASIS) & MASK; + val8 = (db8 >> BASIS) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + index2 += 1; + + val = (db >> BASIS2) & MASK; + val2 = (db2 >> BASIS2) & MASK; + val3 = (db3 >> BASIS2) & MASK; + val4 = (db4 >> BASIS2) & MASK; + val5 = (db5 >> BASIS2) & MASK; + val6 = (db6 >> BASIS2) & MASK; + val7 = (db7 >> BASIS2) & MASK; + val8 = (db8 >> BASIS2) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + index2 += 1; + + val = (db >> BASIS3) & MASK; + val2 = (db2 >> BASIS3) & MASK; + val3 = (db3 >> BASIS3) & MASK; + val4 = (db4 >> BASIS3) & MASK; + val5 = (db5 >> BASIS3) & MASK; + val6 = (db6 >> BASIS3) & MASK; + val7 = (db7 >> BASIS3) & MASK; + val8 = (db8 >> BASIS3) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + index2 += 1; + index += 1; + } + out[i] += tmp; + out[i + 1] += tmp2; + out[i + 2] += tmp3; + out[i + 3] += tmp4; + out[i + 4] += tmp5; + out[i + 5] += tmp6; + out[i + 6] += tmp7; + out[i + 7] += tmp8; + index += aCols * 7; + } +} + +void matMulVecPacked2(Elem *out, const Elem *a, const Elem *b_full, + size_t aRows, size_t aCols) { + // in this variant, b has 2 columns + + Elem db, db2, db3, db4, db5, db6, db7, db8; + Elem val, val2, val3, val4, val5, val6, val7, val8; + Elem tmp, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, tmp8; + Elem x_tmp, x_tmp2, x_tmp3, x_tmp4, x_tmp5, x_tmp6, x_tmp7, x_tmp8; + size_t index = 0; + size_t index2; + + const Elem *b = b_full; + const Elem *b2 = b_full + aCols * COMPRESSION; + + for (size_t i = 0; i < aRows; i += 8) { + tmp = 0; + tmp2 = 0; + tmp3 = 0; + tmp4 = 0; + tmp5 = 0; + tmp6 = 0; + tmp7 = 0; + tmp8 = 0; + + x_tmp = 0; + x_tmp2 = 0; + x_tmp3 = 0; + x_tmp4 = 0; + x_tmp5 = 0; + x_tmp6 = 0; + x_tmp7 = 0; + x_tmp8 = 0; + + index2 = 0; + for (size_t j = 0; j < aCols; j++) { + db = a[index]; + db2 = a[index + 1 * aCols]; + db3 = a[index + 2 * aCols]; + db4 = a[index + 3 * aCols]; + db5 = a[index + 4 * aCols]; + db6 = a[index + 5 * aCols]; + db7 = a[index + 6 * aCols]; + db8 = a[index + 7 * aCols]; + + val = db & MASK; + val2 = db2 & MASK; + val3 = db3 & MASK; + val4 = db4 & MASK; + val5 = db5 & MASK; + val6 = db6 & MASK; + val7 = db7 & MASK; + val8 = db8 & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + index2 += 1; + + val = (db >> BASIS) & MASK; + val2 = (db2 >> BASIS) & MASK; + val3 = (db3 >> BASIS) & MASK; + val4 = (db4 >> BASIS) & MASK; + val5 = (db5 >> BASIS) & MASK; + val6 = (db6 >> BASIS) & MASK; + val7 = (db7 >> BASIS) & MASK; + val8 = (db8 >> BASIS) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + index2 += 1; + + val = (db >> BASIS2) & MASK; + val2 = (db2 >> BASIS2) & MASK; + val3 = (db3 >> BASIS2) & MASK; + val4 = (db4 >> BASIS2) & MASK; + val5 = (db5 >> BASIS2) & MASK; + val6 = (db6 >> BASIS2) & MASK; + val7 = (db7 >> BASIS2) & MASK; + val8 = (db8 >> BASIS2) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + index2 += 1; + + val = (db >> BASIS3) & MASK; + val2 = (db2 >> BASIS3) & MASK; + val3 = (db3 >> BASIS3) & MASK; + val4 = (db4 >> BASIS3) & MASK; + val5 = (db5 >> BASIS3) & MASK; + val6 = (db6 >> BASIS3) & MASK; + val7 = (db7 >> BASIS3) & MASK; + val8 = (db8 >> BASIS3) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + index2 += 1; + + index += 1; + } + out[2 * i] += tmp; + out[2 * i + 1] += x_tmp; + out[2 * i + 2] += tmp2; + out[2 * i + 3] += x_tmp2; + out[2 * i + 4] += tmp3; + out[2 * i + 5] += x_tmp3; + out[2 * i + 6] += tmp4; + out[2 * i + 7] += x_tmp4; + out[2 * i + 8] += tmp5; + out[2 * i + 9] += x_tmp5; + out[2 * i + 10] += tmp6; + out[2 * i + 11] += x_tmp6; + out[2 * i + 12] += tmp7; + out[2 * i + 13] += x_tmp7; + out[2 * i + 14] += tmp8; + out[2 * i + 15] += x_tmp8; + index += aCols * 7; + } +} + +void trueMatMulVecPacked4(Elem *out, const Elem *a, const Elem *b_full, + size_t aRows, size_t aCols) { + // in this variant, b has 4 columns + + Elem db, db2, db3, db4, db5, db6, db7, db8; + Elem val, val2, val3, val4, val5, val6, val7, val8; + Elem tmp, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, tmp8; + Elem x_tmp, x_tmp2, x_tmp3, x_tmp4, x_tmp5, x_tmp6, x_tmp7, x_tmp8; + Elem y_tmp, y_tmp2, y_tmp3, y_tmp4, y_tmp5, y_tmp6, y_tmp7, y_tmp8; + Elem z_tmp, z_tmp2, z_tmp3, z_tmp4, z_tmp5, z_tmp6, z_tmp7, z_tmp8; + size_t index = 0; + size_t index2; + + const Elem *b = b_full; + const Elem *b2 = b_full + aCols * COMPRESSION; + const Elem *b3 = b_full + 2 * aCols * COMPRESSION; + const Elem *b4 = b_full + 3 * aCols * COMPRESSION; + + for (size_t i = 0; i < aRows; i += 8) { + tmp = 0; + tmp2 = 0; + tmp3 = 0; + tmp4 = 0; + tmp5 = 0; + tmp6 = 0; + tmp7 = 0; + tmp8 = 0; + + x_tmp = 0; + x_tmp2 = 0; + x_tmp3 = 0; + x_tmp4 = 0; + x_tmp5 = 0; + x_tmp6 = 0; + x_tmp7 = 0; + x_tmp8 = 0; + + y_tmp = 0; + y_tmp2 = 0; + y_tmp3 = 0; + y_tmp4 = 0; + y_tmp5 = 0; + y_tmp6 = 0; + y_tmp7 = 0; + y_tmp8 = 0; + + z_tmp = 0; + z_tmp2 = 0; + z_tmp3 = 0; + z_tmp4 = 0; + z_tmp5 = 0; + z_tmp6 = 0; + z_tmp7 = 0; + z_tmp8 = 0; + + index2 = 0; + for (size_t j = 0; j < aCols; j++) { + db = a[index]; + db2 = a[index + 1 * aCols]; + db3 = a[index + 2 * aCols]; + db4 = a[index + 3 * aCols]; + db5 = a[index + 4 * aCols]; + db6 = a[index + 5 * aCols]; + db7 = a[index + 6 * aCols]; + db8 = a[index + 7 * aCols]; + + val = db & MASK; + val2 = db2 & MASK; + val3 = db3 & MASK; + val4 = db4 & MASK; + val5 = db5 & MASK; + val6 = db6 & MASK; + val7 = db7 & MASK; + val8 = db8 & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + + index2 += 1; + + val = (db >> BASIS) & MASK; + val2 = (db2 >> BASIS) & MASK; + val3 = (db3 >> BASIS) & MASK; + val4 = (db4 >> BASIS) & MASK; + val5 = (db5 >> BASIS) & MASK; + val6 = (db6 >> BASIS) & MASK; + val7 = (db7 >> BASIS) & MASK; + val8 = (db8 >> BASIS) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + + index2 += 1; + + val = (db >> BASIS2) & MASK; + val2 = (db2 >> BASIS2) & MASK; + val3 = (db3 >> BASIS2) & MASK; + val4 = (db4 >> BASIS2) & MASK; + val5 = (db5 >> BASIS2) & MASK; + val6 = (db6 >> BASIS2) & MASK; + val7 = (db7 >> BASIS2) & MASK; + val8 = (db8 >> BASIS2) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + + index2 += 1; + + val = (db >> BASIS3) & MASK; + val2 = (db2 >> BASIS3) & MASK; + val3 = (db3 >> BASIS3) & MASK; + val4 = (db4 >> BASIS3) & MASK; + val5 = (db5 >> BASIS3) & MASK; + val6 = (db6 >> BASIS3) & MASK; + val7 = (db7 >> BASIS3) & MASK; + val8 = (db8 >> BASIS3) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + + index2 += 1; + + index += 1; + } + out[4 * i] += tmp; + out[4 * i + 1] += x_tmp; + out[4 * i + 2] += y_tmp; + out[4 * i + 3] += z_tmp; + out[4 * i + 4] += tmp2; + out[4 * i + 5] += x_tmp2; + out[4 * i + 6] += y_tmp2; + out[4 * i + 7] += z_tmp2; + out[4 * i + 8] += tmp3; + out[4 * i + 9] += x_tmp3; + out[4 * i + 10] += y_tmp3; + out[4 * i + 11] += z_tmp3; + out[4 * i + 12] += tmp4; + out[4 * i + 13] += x_tmp4; + out[4 * i + 14] += y_tmp4; + out[4 * i + 15] += z_tmp4; + out[4 * i + 16] += tmp5; + out[4 * i + 17] += x_tmp5; + out[4 * i + 18] += y_tmp5; + out[4 * i + 19] += z_tmp5; + out[4 * i + 20] += tmp6; + out[4 * i + 21] += x_tmp6; + out[4 * i + 22] += y_tmp6; + out[4 * i + 23] += z_tmp6; + out[4 * i + 24] += tmp7; + out[4 * i + 25] += x_tmp7; + out[4 * i + 26] += y_tmp7; + out[4 * i + 27] += z_tmp7; + out[4 * i + 28] += tmp8; + out[4 * i + 29] += x_tmp8; + out[4 * i + 30] += y_tmp8; + out[4 * i + 31] += z_tmp8; + index += aCols * 7; + } +} + +void matMulVecPacked4(Elem *out, const Elem *a, const Elem *b_full, + size_t aRows, size_t aCols) { +#ifdef PROFILE + System::profile("matmul4", [&]() { +#endif + trueMatMulVecPacked4(out, a, b_full, aRows, aCols); + +#ifdef PROFILE + }); +#endif +} + +void matMulVecPacked6(Elem *__restrict__ out, const Elem *__restrict__ a, + const Elem *__restrict__ b_full, size_t aRows, + size_t aCols) { + // in this variant, b has 6 columns + + Elem db, db2, db3, db4, db5, db6, db7, db8; + Elem val, val2, val3, val4, val5, val6, val7, val8; + Elem tmp, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, tmp8; + Elem x_tmp, x_tmp2, x_tmp3, x_tmp4, x_tmp5, x_tmp6, x_tmp7, x_tmp8; + Elem y_tmp, y_tmp2, y_tmp3, y_tmp4, y_tmp5, y_tmp6, y_tmp7, y_tmp8; + Elem z_tmp, z_tmp2, z_tmp3, z_tmp4, z_tmp5, z_tmp6, z_tmp7, z_tmp8; + Elem r_tmp, r_tmp2, r_tmp3, r_tmp4, r_tmp5, r_tmp6, r_tmp7, r_tmp8; + Elem s_tmp, s_tmp2, s_tmp3, s_tmp4, s_tmp5, s_tmp6, s_tmp7, s_tmp8; + size_t index = 0; + size_t index2; + + const Elem *b = b_full; + const Elem *b2 = b_full + aCols * COMPRESSION; + const Elem *b3 = b_full + 2 * aCols * COMPRESSION; + const Elem *b4 = b_full + 3 * aCols * COMPRESSION; + const Elem *b5 = b_full + 4 * aCols * COMPRESSION; + const Elem *b6 = b_full + 5 * aCols * COMPRESSION; + + for (size_t i = 0; i < aRows; i += 8) { + tmp = 0; + tmp2 = 0; + tmp3 = 0; + tmp4 = 0; + tmp5 = 0; + tmp6 = 0; + tmp7 = 0; + tmp8 = 0; + + x_tmp = 0; + x_tmp2 = 0; + x_tmp3 = 0; + x_tmp4 = 0; + x_tmp5 = 0; + x_tmp6 = 0; + x_tmp7 = 0; + x_tmp8 = 0; + + y_tmp = 0; + y_tmp2 = 0; + y_tmp3 = 0; + y_tmp4 = 0; + y_tmp5 = 0; + y_tmp6 = 0; + y_tmp7 = 0; + y_tmp8 = 0; + + z_tmp = 0; + z_tmp2 = 0; + z_tmp3 = 0; + z_tmp4 = 0; + z_tmp5 = 0; + z_tmp6 = 0; + z_tmp7 = 0; + z_tmp8 = 0; + + r_tmp = 0; + r_tmp2 = 0; + r_tmp3 = 0; + r_tmp4 = 0; + r_tmp5 = 0; + r_tmp6 = 0; + r_tmp7 = 0; + r_tmp8 = 0; + + s_tmp = 0; + s_tmp2 = 0; + s_tmp3 = 0; + s_tmp4 = 0; + s_tmp5 = 0; + s_tmp6 = 0; + s_tmp7 = 0; + s_tmp8 = 0; + + index2 = 0; + for (size_t j = 0; j < aCols; j++) { + db = a[index]; + db2 = a[index + 1 * aCols]; + db3 = a[index + 2 * aCols]; + db4 = a[index + 3 * aCols]; + db5 = a[index + 4 * aCols]; + db6 = a[index + 5 * aCols]; + db7 = a[index + 6 * aCols]; + db8 = a[index + 7 * aCols]; + + val = db & MASK; + val2 = db2 & MASK; + val3 = db3 & MASK; + val4 = db4 & MASK; + val5 = db5 & MASK; + val6 = db6 & MASK; + val7 = db7 & MASK; + val8 = db8 & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + r_tmp += val * b5[index2]; + r_tmp2 += val2 * b5[index2]; + r_tmp3 += val3 * b5[index2]; + r_tmp4 += val4 * b5[index2]; + r_tmp5 += val5 * b5[index2]; + r_tmp6 += val6 * b5[index2]; + r_tmp7 += val7 * b5[index2]; + r_tmp8 += val8 * b5[index2]; + s_tmp += val * b6[index2]; + s_tmp2 += val2 * b6[index2]; + s_tmp3 += val3 * b6[index2]; + s_tmp4 += val4 * b6[index2]; + s_tmp5 += val5 * b6[index2]; + s_tmp6 += val6 * b6[index2]; + s_tmp7 += val7 * b6[index2]; + s_tmp8 += val8 * b6[index2]; + index2 += 1; + + val = (db >> BASIS) & MASK; + val2 = (db2 >> BASIS) & MASK; + val3 = (db3 >> BASIS) & MASK; + val4 = (db4 >> BASIS) & MASK; + val5 = (db5 >> BASIS) & MASK; + val6 = (db6 >> BASIS) & MASK; + val7 = (db7 >> BASIS) & MASK; + val8 = (db8 >> BASIS) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + r_tmp += val * b5[index2]; + r_tmp2 += val2 * b5[index2]; + r_tmp3 += val3 * b5[index2]; + r_tmp4 += val4 * b5[index2]; + r_tmp5 += val5 * b5[index2]; + r_tmp6 += val6 * b5[index2]; + r_tmp7 += val7 * b5[index2]; + r_tmp8 += val8 * b5[index2]; + s_tmp += val * b6[index2]; + s_tmp2 += val2 * b6[index2]; + s_tmp3 += val3 * b6[index2]; + s_tmp4 += val4 * b6[index2]; + s_tmp5 += val5 * b6[index2]; + s_tmp6 += val6 * b6[index2]; + s_tmp7 += val7 * b6[index2]; + s_tmp8 += val8 * b6[index2]; + index2 += 1; + + val = (db >> BASIS2) & MASK; + val2 = (db2 >> BASIS2) & MASK; + val3 = (db3 >> BASIS2) & MASK; + val4 = (db4 >> BASIS2) & MASK; + val5 = (db5 >> BASIS2) & MASK; + val6 = (db6 >> BASIS2) & MASK; + val7 = (db7 >> BASIS2) & MASK; + val8 = (db8 >> BASIS2) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + r_tmp += val * b5[index2]; + r_tmp2 += val2 * b5[index2]; + r_tmp3 += val3 * b5[index2]; + r_tmp4 += val4 * b5[index2]; + r_tmp5 += val5 * b5[index2]; + r_tmp6 += val6 * b5[index2]; + r_tmp7 += val7 * b5[index2]; + r_tmp8 += val8 * b5[index2]; + s_tmp += val * b6[index2]; + s_tmp2 += val2 * b6[index2]; + s_tmp3 += val3 * b6[index2]; + s_tmp4 += val4 * b6[index2]; + s_tmp5 += val5 * b6[index2]; + s_tmp6 += val6 * b6[index2]; + s_tmp7 += val7 * b6[index2]; + s_tmp8 += val8 * b6[index2]; + index2 += 1; + + val = (db >> BASIS3) & MASK; + val2 = (db2 >> BASIS3) & MASK; + val3 = (db3 >> BASIS3) & MASK; + val4 = (db4 >> BASIS3) & MASK; + val5 = (db5 >> BASIS3) & MASK; + val6 = (db6 >> BASIS3) & MASK; + val7 = (db7 >> BASIS3) & MASK; + val8 = (db8 >> BASIS3) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + r_tmp += val * b5[index2]; + r_tmp2 += val2 * b5[index2]; + r_tmp3 += val3 * b5[index2]; + r_tmp4 += val4 * b5[index2]; + r_tmp5 += val5 * b5[index2]; + r_tmp6 += val6 * b5[index2]; + r_tmp7 += val7 * b5[index2]; + r_tmp8 += val8 * b5[index2]; + s_tmp += val * b6[index2]; + s_tmp2 += val2 * b6[index2]; + s_tmp3 += val3 * b6[index2]; + s_tmp4 += val4 * b6[index2]; + s_tmp5 += val5 * b6[index2]; + s_tmp6 += val6 * b6[index2]; + s_tmp7 += val7 * b6[index2]; + s_tmp8 += val8 * b6[index2]; + index2 += 1; + + index += 1; + } + out[6 * i + 0] += tmp; + out[6 * i + 1] += x_tmp; + out[6 * i + 2] += y_tmp; + out[6 * i + 3] += z_tmp; + out[6 * i + 4] += r_tmp; + out[6 * i + 5] += s_tmp; + out[6 * i + 6] += tmp2; + out[6 * i + 7] += x_tmp2; + out[6 * i + 8] += y_tmp2; + out[6 * i + 9] += z_tmp2; + out[6 * i + 10] += r_tmp2; + out[6 * i + 11] += s_tmp2; + out[6 * i + 12] += tmp3; + out[6 * i + 13] += x_tmp3; + out[6 * i + 14] += y_tmp3; + out[6 * i + 15] += z_tmp3; + out[6 * i + 16] += r_tmp3; + out[6 * i + 17] += s_tmp3; + out[6 * i + 18] += tmp4; + out[6 * i + 19] += x_tmp4; + out[6 * i + 20] += y_tmp4; + out[6 * i + 21] += z_tmp4; + out[6 * i + 22] += r_tmp4; + out[6 * i + 23] += s_tmp4; + out[6 * i + 24] += tmp5; + out[6 * i + 25] += x_tmp5; + out[6 * i + 26] += y_tmp5; + out[6 * i + 27] += z_tmp5; + out[6 * i + 28] += r_tmp5; + out[6 * i + 29] += s_tmp5; + out[6 * i + 30] += tmp6; + out[6 * i + 31] += x_tmp6; + out[6 * i + 32] += y_tmp6; + out[6 * i + 33] += z_tmp6; + out[6 * i + 34] += r_tmp6; + out[6 * i + 35] += s_tmp6; + out[6 * i + 36] += tmp7; + out[6 * i + 37] += x_tmp7; + out[6 * i + 38] += y_tmp7; + out[6 * i + 39] += z_tmp7; + out[6 * i + 40] += r_tmp7; + out[6 * i + 41] += s_tmp7; + out[6 * i + 42] += tmp8; + out[6 * i + 43] += x_tmp8; + out[6 * i + 44] += y_tmp8; + out[6 * i + 45] += z_tmp8; + out[6 * i + 46] += r_tmp8; + out[6 * i + 47] += s_tmp8; + index += aCols * 7; + } +} + +void trueMatMulVecPacked8( + Elem *__restrict__ out, const Elem *__restrict__ a, + const Elem *__restrict__ b, const Elem *__restrict__ b2, + const Elem *__restrict__ b3, const Elem *__restrict__ b4, + const Elem *__restrict__ b5, const Elem *__restrict__ b6, + const Elem *__restrict__ b7, const Elem *__restrict__ b8, size_t aRows, + size_t aCols) { + // in this variant, b has 8 columns + + Elem db, db2, db3, db4, db5, db6, db7, db8; + Elem val, val2, val3, val4, val5, val6, val7, val8; + Elem tmp, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, tmp8; + Elem x_tmp, x_tmp2, x_tmp3, x_tmp4, x_tmp5, x_tmp6, x_tmp7, x_tmp8; + Elem y_tmp, y_tmp2, y_tmp3, y_tmp4, y_tmp5, y_tmp6, y_tmp7, y_tmp8; + Elem z_tmp, z_tmp2, z_tmp3, z_tmp4, z_tmp5, z_tmp6, z_tmp7, z_tmp8; + Elem r_tmp, r_tmp2, r_tmp3, r_tmp4, r_tmp5, r_tmp6, r_tmp7, r_tmp8; + Elem s_tmp, s_tmp2, s_tmp3, s_tmp4, s_tmp5, s_tmp6, s_tmp7, s_tmp8; + Elem t_tmp, t_tmp2, t_tmp3, t_tmp4, t_tmp5, t_tmp6, t_tmp7, t_tmp8; + Elem u_tmp, u_tmp2, u_tmp3, u_tmp4, u_tmp5, u_tmp6, u_tmp7, u_tmp8; + size_t index = 0; + size_t index2; + + for (size_t i = 0; i < aRows; i += 8) { + tmp = 0; + tmp2 = 0; + tmp3 = 0; + tmp4 = 0; + tmp5 = 0; + tmp6 = 0; + tmp7 = 0; + tmp8 = 0; + + x_tmp = 0; + x_tmp2 = 0; + x_tmp3 = 0; + x_tmp4 = 0; + x_tmp5 = 0; + x_tmp6 = 0; + x_tmp7 = 0; + x_tmp8 = 0; + + y_tmp = 0; + y_tmp2 = 0; + y_tmp3 = 0; + y_tmp4 = 0; + y_tmp5 = 0; + y_tmp6 = 0; + y_tmp7 = 0; + y_tmp8 = 0; + + z_tmp = 0; + z_tmp2 = 0; + z_tmp3 = 0; + z_tmp4 = 0; + z_tmp5 = 0; + z_tmp6 = 0; + z_tmp7 = 0; + z_tmp8 = 0; + + r_tmp = 0; + r_tmp2 = 0; + r_tmp3 = 0; + r_tmp4 = 0; + r_tmp5 = 0; + r_tmp6 = 0; + r_tmp7 = 0; + r_tmp8 = 0; + + s_tmp = 0; + s_tmp2 = 0; + s_tmp3 = 0; + s_tmp4 = 0; + s_tmp5 = 0; + s_tmp6 = 0; + s_tmp7 = 0; + s_tmp8 = 0; + + t_tmp = 0; + t_tmp2 = 0; + t_tmp3 = 0; + t_tmp4 = 0; + t_tmp5 = 0; + t_tmp6 = 0; + t_tmp7 = 0; + t_tmp8 = 0; + + u_tmp = 0; + u_tmp2 = 0; + u_tmp3 = 0; + u_tmp4 = 0; + u_tmp5 = 0; + u_tmp6 = 0; + u_tmp7 = 0; + u_tmp8 = 0; + + index2 = 0; + for (size_t j = 0; j < aCols; j++) { + db = a[index]; + db2 = a[index + 1 * aCols]; + db3 = a[index + 2 * aCols]; + db4 = a[index + 3 * aCols]; + db5 = a[index + 4 * aCols]; + db6 = a[index + 5 * aCols]; + db7 = a[index + 6 * aCols]; + db8 = a[index + 7 * aCols]; + + val = db & MASK; + val2 = db2 & MASK; + val3 = db3 & MASK; + val4 = db4 & MASK; + val5 = db5 & MASK; + val6 = db6 & MASK; + val7 = db7 & MASK; + val8 = db8 & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + r_tmp += val * b5[index2]; + r_tmp2 += val2 * b5[index2]; + r_tmp3 += val3 * b5[index2]; + r_tmp4 += val4 * b5[index2]; + r_tmp5 += val5 * b5[index2]; + r_tmp6 += val6 * b5[index2]; + r_tmp7 += val7 * b5[index2]; + r_tmp8 += val8 * b5[index2]; + s_tmp += val * b6[index2]; + s_tmp2 += val2 * b6[index2]; + s_tmp3 += val3 * b6[index2]; + s_tmp4 += val4 * b6[index2]; + s_tmp5 += val5 * b6[index2]; + s_tmp6 += val6 * b6[index2]; + s_tmp7 += val7 * b6[index2]; + s_tmp8 += val8 * b6[index2]; + t_tmp += val * b7[index2]; + t_tmp2 += val2 * b7[index2]; + t_tmp3 += val3 * b7[index2]; + t_tmp4 += val4 * b7[index2]; + t_tmp5 += val5 * b7[index2]; + t_tmp6 += val6 * b7[index2]; + t_tmp7 += val7 * b7[index2]; + t_tmp8 += val8 * b7[index2]; + u_tmp += val * b8[index2]; + u_tmp2 += val2 * b8[index2]; + u_tmp3 += val3 * b8[index2]; + u_tmp4 += val4 * b8[index2]; + u_tmp5 += val5 * b8[index2]; + u_tmp6 += val6 * b8[index2]; + u_tmp7 += val7 * b8[index2]; + u_tmp8 += val8 * b8[index2]; + + index2 += 1; + + val = (db >> BASIS) & MASK; + val2 = (db2 >> BASIS) & MASK; + val3 = (db3 >> BASIS) & MASK; + val4 = (db4 >> BASIS) & MASK; + val5 = (db5 >> BASIS) & MASK; + val6 = (db6 >> BASIS) & MASK; + val7 = (db7 >> BASIS) & MASK; + val8 = (db8 >> BASIS) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + r_tmp += val * b5[index2]; + r_tmp2 += val2 * b5[index2]; + r_tmp3 += val3 * b5[index2]; + r_tmp4 += val4 * b5[index2]; + r_tmp5 += val5 * b5[index2]; + r_tmp6 += val6 * b5[index2]; + r_tmp7 += val7 * b5[index2]; + r_tmp8 += val8 * b5[index2]; + s_tmp += val * b6[index2]; + s_tmp2 += val2 * b6[index2]; + s_tmp3 += val3 * b6[index2]; + s_tmp4 += val4 * b6[index2]; + s_tmp5 += val5 * b6[index2]; + s_tmp6 += val6 * b6[index2]; + s_tmp7 += val7 * b6[index2]; + s_tmp8 += val8 * b6[index2]; + t_tmp += val * b7[index2]; + t_tmp2 += val2 * b7[index2]; + t_tmp3 += val3 * b7[index2]; + t_tmp4 += val4 * b7[index2]; + t_tmp5 += val5 * b7[index2]; + t_tmp6 += val6 * b7[index2]; + t_tmp7 += val7 * b7[index2]; + t_tmp8 += val8 * b7[index2]; + u_tmp += val * b8[index2]; + u_tmp2 += val2 * b8[index2]; + u_tmp3 += val3 * b8[index2]; + u_tmp4 += val4 * b8[index2]; + u_tmp5 += val5 * b8[index2]; + u_tmp6 += val6 * b8[index2]; + u_tmp7 += val7 * b8[index2]; + u_tmp8 += val8 * b8[index2]; + + index2 += 1; + + val = (db >> BASIS2) & MASK; + val2 = (db2 >> BASIS2) & MASK; + val3 = (db3 >> BASIS2) & MASK; + val4 = (db4 >> BASIS2) & MASK; + val5 = (db5 >> BASIS2) & MASK; + val6 = (db6 >> BASIS2) & MASK; + val7 = (db7 >> BASIS2) & MASK; + val8 = (db8 >> BASIS2) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + r_tmp += val * b5[index2]; + r_tmp2 += val2 * b5[index2]; + r_tmp3 += val3 * b5[index2]; + r_tmp4 += val4 * b5[index2]; + r_tmp5 += val5 * b5[index2]; + r_tmp6 += val6 * b5[index2]; + r_tmp7 += val7 * b5[index2]; + r_tmp8 += val8 * b5[index2]; + s_tmp += val * b6[index2]; + s_tmp2 += val2 * b6[index2]; + s_tmp3 += val3 * b6[index2]; + s_tmp4 += val4 * b6[index2]; + s_tmp5 += val5 * b6[index2]; + s_tmp6 += val6 * b6[index2]; + s_tmp7 += val7 * b6[index2]; + s_tmp8 += val8 * b6[index2]; + t_tmp += val * b7[index2]; + t_tmp2 += val2 * b7[index2]; + t_tmp3 += val3 * b7[index2]; + t_tmp4 += val4 * b7[index2]; + t_tmp5 += val5 * b7[index2]; + t_tmp6 += val6 * b7[index2]; + t_tmp7 += val7 * b7[index2]; + t_tmp8 += val8 * b7[index2]; + u_tmp += val * b8[index2]; + u_tmp2 += val2 * b8[index2]; + u_tmp3 += val3 * b8[index2]; + u_tmp4 += val4 * b8[index2]; + u_tmp5 += val5 * b8[index2]; + u_tmp6 += val6 * b8[index2]; + u_tmp7 += val7 * b8[index2]; + u_tmp8 += val8 * b8[index2]; + + index2 += 1; + + val = (db >> BASIS3) & MASK; + val2 = (db2 >> BASIS3) & MASK; + val3 = (db3 >> BASIS3) & MASK; + val4 = (db4 >> BASIS3) & MASK; + val5 = (db5 >> BASIS3) & MASK; + val6 = (db6 >> BASIS3) & MASK; + val7 = (db7 >> BASIS3) & MASK; + val8 = (db8 >> BASIS3) & MASK; + tmp += val * b[index2]; + tmp2 += val2 * b[index2]; + tmp3 += val3 * b[index2]; + tmp4 += val4 * b[index2]; + tmp5 += val5 * b[index2]; + tmp6 += val6 * b[index2]; + tmp7 += val7 * b[index2]; + tmp8 += val8 * b[index2]; + x_tmp += val * b2[index2]; + x_tmp2 += val2 * b2[index2]; + x_tmp3 += val3 * b2[index2]; + x_tmp4 += val4 * b2[index2]; + x_tmp5 += val5 * b2[index2]; + x_tmp6 += val6 * b2[index2]; + x_tmp7 += val7 * b2[index2]; + x_tmp8 += val8 * b2[index2]; + y_tmp += val * b3[index2]; + y_tmp2 += val2 * b3[index2]; + y_tmp3 += val3 * b3[index2]; + y_tmp4 += val4 * b3[index2]; + y_tmp5 += val5 * b3[index2]; + y_tmp6 += val6 * b3[index2]; + y_tmp7 += val7 * b3[index2]; + y_tmp8 += val8 * b3[index2]; + z_tmp += val * b4[index2]; + z_tmp2 += val2 * b4[index2]; + z_tmp3 += val3 * b4[index2]; + z_tmp4 += val4 * b4[index2]; + z_tmp5 += val5 * b4[index2]; + z_tmp6 += val6 * b4[index2]; + z_tmp7 += val7 * b4[index2]; + z_tmp8 += val8 * b4[index2]; + r_tmp += val * b5[index2]; + r_tmp2 += val2 * b5[index2]; + r_tmp3 += val3 * b5[index2]; + r_tmp4 += val4 * b5[index2]; + r_tmp5 += val5 * b5[index2]; + r_tmp6 += val6 * b5[index2]; + r_tmp7 += val7 * b5[index2]; + r_tmp8 += val8 * b5[index2]; + s_tmp += val * b6[index2]; + s_tmp2 += val2 * b6[index2]; + s_tmp3 += val3 * b6[index2]; + s_tmp4 += val4 * b6[index2]; + s_tmp5 += val5 * b6[index2]; + s_tmp6 += val6 * b6[index2]; + s_tmp7 += val7 * b6[index2]; + s_tmp8 += val8 * b6[index2]; + t_tmp += val * b7[index2]; + t_tmp2 += val2 * b7[index2]; + t_tmp3 += val3 * b7[index2]; + t_tmp4 += val4 * b7[index2]; + t_tmp5 += val5 * b7[index2]; + t_tmp6 += val6 * b7[index2]; + t_tmp7 += val7 * b7[index2]; + t_tmp8 += val8 * b7[index2]; + u_tmp += val * b8[index2]; + u_tmp2 += val2 * b8[index2]; + u_tmp3 += val3 * b8[index2]; + u_tmp4 += val4 * b8[index2]; + u_tmp5 += val5 * b8[index2]; + u_tmp6 += val6 * b8[index2]; + u_tmp7 += val7 * b8[index2]; + u_tmp8 += val8 * b8[index2]; + + index2 += 1; + + index += 1; + } + out[8 * i + 0] += tmp; + out[8 * i + 1] += x_tmp; + out[8 * i + 2] += y_tmp; + out[8 * i + 3] += z_tmp; + out[8 * i + 4] += r_tmp; + out[8 * i + 5] += s_tmp; + out[8 * i + 6] += t_tmp; + out[8 * i + 7] += u_tmp; + out[8 * i + 8] += tmp2; + out[8 * i + 9] += x_tmp2; + out[8 * i + 10] += y_tmp2; + out[8 * i + 11] += z_tmp2; + out[8 * i + 12] += r_tmp2; + out[8 * i + 13] += s_tmp2; + out[8 * i + 14] += t_tmp2; + out[8 * i + 15] += u_tmp2; + out[8 * i + 16] += tmp3; + out[8 * i + 17] += x_tmp3; + out[8 * i + 18] += y_tmp3; + out[8 * i + 19] += z_tmp3; + out[8 * i + 20] += r_tmp3; + out[8 * i + 21] += s_tmp3; + out[8 * i + 22] += t_tmp3; + out[8 * i + 23] += u_tmp3; + out[8 * i + 24] += tmp4; + out[8 * i + 25] += x_tmp4; + out[8 * i + 26] += y_tmp4; + out[8 * i + 27] += z_tmp4; + out[8 * i + 28] += r_tmp4; + out[8 * i + 29] += s_tmp4; + out[8 * i + 30] += t_tmp4; + out[8 * i + 31] += u_tmp4; + out[8 * i + 32] += tmp5; + out[8 * i + 33] += x_tmp5; + out[8 * i + 34] += y_tmp5; + out[8 * i + 35] += z_tmp5; + out[8 * i + 36] += r_tmp5; + out[8 * i + 37] += s_tmp5; + out[8 * i + 38] += t_tmp5; + out[8 * i + 39] += u_tmp5; + out[8 * i + 40] += tmp6; + out[8 * i + 41] += x_tmp6; + out[8 * i + 42] += y_tmp6; + out[8 * i + 43] += z_tmp6; + out[8 * i + 44] += r_tmp6; + out[8 * i + 45] += s_tmp6; + out[8 * i + 46] += t_tmp6; + out[8 * i + 47] += u_tmp6; + out[8 * i + 48] += tmp7; + out[8 * i + 49] += x_tmp7; + out[8 * i + 50] += y_tmp7; + out[8 * i + 51] += z_tmp7; + out[8 * i + 52] += r_tmp7; + out[8 * i + 53] += s_tmp7; + out[8 * i + 54] += t_tmp7; + out[8 * i + 55] += u_tmp7; + out[8 * i + 56] += tmp8; + out[8 * i + 57] += x_tmp8; + out[8 * i + 58] += y_tmp8; + out[8 * i + 59] += z_tmp8; + out[8 * i + 60] += r_tmp8; + out[8 * i + 61] += s_tmp8; + out[8 * i + 62] += t_tmp8; + out[8 * i + 63] += u_tmp8; + index += aCols * 7; + } +} + +void matMulVecPacked8(Elem *__restrict__ out, const Elem *__restrict__ a, + const Elem *__restrict__ b_full, size_t aRows, + size_t aCols) { +#ifdef PROFILE + System::profile("matmul4", [&]() { +#endif + trueMatMulVecPacked8( + out, a, b_full, b_full + (1 * aCols * COMPRESSION), + b_full + (2 * aCols * COMPRESSION), b_full + (3 * aCols * COMPRESSION), + b_full + (4 * aCols * COMPRESSION), b_full + (5 * aCols * COMPRESSION), + b_full + (6 * aCols * COMPRESSION), b_full + (7 * aCols * COMPRESSION), + aRows, aCols); + +#ifdef PROFILE + }); +#endif +} \ No newline at end of file diff --git a/psi/algorithm/ypir/legacy/server.cc b/psi/algorithm/ypir/legacy/server.cc new file mode 100644 index 00000000..ce6e4ef3 --- /dev/null +++ b/psi/algorithm/ypir/legacy/server.cc @@ -0,0 +1,157 @@ +#include "psi/algorithm/ypir/legacy/server.h" + +#include + +namespace psi::ypir::ypir_internal { +namespace { + +void DoublepirAnswer(const uint8_t* db, uint32_t* qu0, + std::vector& qu1, + const std::vector>& server_hint, + std::vector& h2, + std::vector>& h3, + std::vector& h4, const FheParams& fparm, + const PirParams& pparm) { + const uint64_t row = pparm.get_row(); + const uint64_t col = pparm.get_col(); + const uint64_t rlwe_cmod = fparm.get_rlwe_cmod(); + const uint64_t b = fparm.get_b_decomp(); + const uint64_t z = fparm.get_z_decomp(); + const uint64_t t = fparm.get_t_decomp(); + + std::vector simple_res(row, 0); + MatVecU8U32Mod2p32(db, qu0, simple_res.data(), row, col); + + std::vector> trans_simple_res_decomp; + VectorColDecompose(simple_res, trans_simple_res_decomp, b, z, t); + + const uint64_t* matrix_d2_flat = fparm.get_persudo_matrix_doublepir_flat(); + const uint64_t poly_degree = fparm.get_poly_degree(); + + MatrixVectorFirstDimension(h2, server_hint, qu1, rlwe_cmod); + + MatrixMultiplicationFlatU16(h3, trans_simple_res_decomp, matrix_d2_flat, + poly_degree, rlwe_cmod); + + MatrixVectorMultiplicationU16(h4, trans_simple_res_decomp, qu1, rlwe_cmod); +} + +} // namespace + +void YpirHintGenerate( + std::vector>& db, std::vector& h0, + std::vector>& double_server_hint, + std::vector>>& decomp_buf, + AESCTR_PRNG& prng, const FheParams& fparm, const PirParams& pparm) { + const uint64_t lwe_dimension = fparm.get_lwe_dimension(); + const uint64_t poly_degree = fparm.get_poly_degree(); + const uint64_t rlwe_cmod = fparm.get_rlwe_cmod(); + const uint64_t t_auto = fparm.get_t_auto(); + const uint64_t t_decomp = fparm.get_t_decomp(); + const uint64_t expo = GetLog2(poly_degree); + const uint64_t pack_num = lwe_dimension * t_decomp; + + std::vector> double_client_hint( + lwe_dimension * t_decomp, std::vector(poly_degree, 0)); + + const uint64_t* matrix_flat = fparm.get_persudo_matrix_simplepir_flat(); + std::vector> simple_hint( + pparm.get_row(), std::vector(lwe_dimension, 0)); + MatrixMultiplicationFlat(simple_hint, db, matrix_flat, lwe_dimension, + fparm.get_lwe_cmod()); + + std::vector> matrix_decomp; + MatrixRowDecompose(simple_hint, matrix_decomp, fparm.get_b_decomp(), + fparm.get_z_decomp(), fparm.get_t_decomp()); + MatrixTranspose(matrix_decomp, double_server_hint); + + const uint64_t* matrix_d2_flat = fparm.get_persudo_matrix_doublepir_flat(); + MatrixMultiplicationFlat(double_client_hint, double_server_hint, + matrix_d2_flat, poly_degree, rlwe_cmod); + + std::vector>> ksk_a( + expo, std::vector>( + t_auto, std::vector(poly_degree))); + prng.refresh(kThirdDimensionSeed); + for (uint64_t i = 0; i < expo; ++i) { + PseudorandomMatrixGenerate(ksk_a[i], rlwe_cmod, prng); + } + + const uint64_t mod_inv = ModInverse(static_cast(pack_num), + static_cast(rlwe_cmod)); + YpirHexlNtt& ntt = fparm.get_ntt(); + for (uint64_t i = 0; i < pack_num; ++i) { + Cdks21Lwe2RlweInplace(double_client_hint[i].data(), poly_degree, rlwe_cmod, + ntt); + EltwiseFMAMod(double_client_hint[i].data(), double_client_hint[i].data(), + mod_inv, nullptr, poly_degree, rlwe_cmod); + } + + h0 = PackrlwePreprocess(double_client_hint, GetLog2(pack_num), + GetLog2(poly_degree) - GetLog2(pack_num), ksk_a, + decomp_buf, fparm); +} + +void YpirAnswer(const uint8_t* db, uint32_t* qu0, std::vector& qu1, + const std::vector>>& ksk_b, + std::vector>>& decomp_buf, + const std::vector>& server_hint, + std::vector>& res, const FheParams& fparm, + const PirParams& pparm) { + const uint64_t lwe_dimension = fparm.get_lwe_dimension(); + const uint64_t poly_degree = fparm.get_poly_degree(); + const uint64_t rlwe_cmod = fparm.get_rlwe_cmod(); + const uint64_t t_decomp = fparm.get_t_decomp(); + const uint64_t pack_num = lwe_dimension * t_decomp; + const uint64_t mod_inv = ModInverse(static_cast(pack_num), + static_cast(rlwe_cmod)); + const uint64_t pack_num_log2 = GetLog2(pack_num); + const uint64_t t_decomp_log2 = GetLog2(t_decomp); + const uint64_t poly_degree_log2 = GetLog2(poly_degree); + + res.reserve(res.size() + 3); + + std::vector> h3(t_decomp, + std::vector(poly_degree, 0)); + std::vector h2(lwe_dimension * t_decomp, 0); + std::vector h4(t_decomp, 0); + + DoublepirAnswer(db, qu0, qu1, server_hint, h2, h3, h4, fparm, pparm); + + for (uint64_t i = 0; i < pack_num; ++i) { + h2[i] = (static_cast(h2[i]) * mod_inv) % rlwe_cmod; + } + + uint64_t ptr = 0; + res.push_back(PackrlweOnlineConstantRows(h2, pack_num_log2, + poly_degree_log2 - pack_num_log2, + ksk_b, decomp_buf, ptr, fparm)); + + std::vector> h4_matrix( + t_decomp, std::vector(poly_degree, 0)); + const uint64_t mod_inv_2 = ModInverse(static_cast(t_decomp), + static_cast(rlwe_cmod)); + YpirHexlNtt& ntt = fparm.get_ntt(); + + for (uint64_t i = 0; i < t_decomp; ++i) { + Cdks21Lwe2RlweInplace(h3[i].data(), poly_degree, rlwe_cmod, ntt); + h4_matrix[i][0] = h4[i]; + ntt.Forward(h4_matrix[i].data(), poly_degree); + EltwiseFMAMod(h3[i].data(), h3[i].data(), mod_inv_2, nullptr, poly_degree, + rlwe_cmod); + EltwiseFMAMod(h4_matrix[i].data(), h4_matrix[i].data(), mod_inv_2, nullptr, + poly_degree, rlwe_cmod); + } + + const auto& ksk_a = fparm.get_persudo_hcube_ypir(); + std::vector>> decomp_b_buf; + ptr = 0; + res.push_back(PackrlwePreprocess(h3, t_decomp_log2, + poly_degree_log2 - t_decomp_log2, ksk_a, + decomp_b_buf, fparm)); + res.push_back(PackrlweOnline(h4_matrix, t_decomp_log2, + poly_degree_log2 - t_decomp_log2, ksk_b, + decomp_b_buf, ptr, fparm)); +} + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/server.h b/psi/algorithm/ypir/legacy/server.h new file mode 100644 index 00000000..56e01fa2 --- /dev/null +++ b/psi/algorithm/ypir/legacy/server.h @@ -0,0 +1,23 @@ +#pragma once + +#include +#include + +#include "psi/algorithm/ypir/legacy/ypir_params.h" + +namespace psi::ypir::ypir_internal { + +void YpirAnswer(const uint8_t* db, uint32_t* qu0, std::vector& qu1, + const std::vector>>& ksk_b, + std::vector>>& decomp_buf, + const std::vector>& server_hint, + std::vector>& res, const FheParams& fparm, + const PirParams& pparm); + +void YpirHintGenerate( + std::vector>& db, std::vector& h0, + std::vector>& double_server_hint, + std::vector>>& decomp_buf, + AESCTR_PRNG& prng, const FheParams& fparm, const PirParams& pparm); + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/types.h b/psi/algorithm/ypir/legacy/types.h new file mode 100644 index 00000000..2f94aaad --- /dev/null +++ b/psi/algorithm/ypir/legacy/types.h @@ -0,0 +1,57 @@ +// Copyright 2024 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include + +#include "psi/algorithm/spiral/poly_matrix.h" + +namespace psi::ypir { + +using namespace psi::spiral; + +// Precomputation tuple type for packing operations +using Precomp = + std::vector, + std::vector>>>; + +// Forward declaration for YPirServer template class +template +class YPirServer; + +// Convenience alias for the most common type +using YServer = YPirServer; + +// Offline precomputation values for SimplePIR and DoublePIR +struct OfflinePrecomputedValues { + // SimplePIR fields + std::vector hint_0; // First hint + // DoublePIR fields + std::vector hint_1; // Second dimension hint + std::vector + pseudorandom_query_1; // Pseudorandom query for second dimension + std::pair, std::vector> y_constants; + std::shared_ptr smaller_server; + // Packing-related fields + std::vector> prepacked_lwe; + std::vector fake_pack_pub_params; + Precomp precomp; +}; + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/legacy/util.cc b/psi/algorithm/ypir/legacy/util.cc new file mode 100644 index 00000000..c9c2bc1c --- /dev/null +++ b/psi/algorithm/ypir/legacy/util.cc @@ -0,0 +1,1125 @@ +#include "psi/algorithm/ypir/legacy/util.h" + +#include +#include +#include + +#include "yacl/base/int128.h" + +#include "psi/algorithm/spiral/arith/arith_params.h" +#include "psi/algorithm/spiral/arith/ntt.h" +#include "psi/algorithm/spiral/discrete_gaussian.h" +#include "psi/algorithm/spiral/gadget.h" +#include "psi/algorithm/spiral/params.h" +#include "psi/algorithm/spiral/poly_matrix.h" +#include "psi/algorithm/spiral/poly_matrix_utils.h" +#include "psi/algorithm/spiral/util.h" + +namespace psi::ypir { + +using namespace psi::spiral; + +std::vector ConcatHorizontal( + const std::vector>& v_a, size_t a_rows, + size_t a_cols) { + if (v_a.empty()) { + return {}; + } + + size_t num_vecs = v_a.size(); + + std::vector out(a_rows * a_cols * num_vecs); + size_t out_stride = num_vecs * a_cols; + + for (size_t i = 0; i < a_rows; ++i) { + for (size_t k = 0; k < num_vecs; ++k) { + auto src_begin = v_a[k].begin() + (i * a_cols); + auto src_end = src_begin + a_cols; + size_t dest_offset = (i * out_stride) + (k * a_cols); + + std::copy(src_begin, src_end, out.begin() + dest_offset); + } + } + + return out; +} + +extern "C" { +void matMulVecPacked(uint32_t* out, const uint32_t* a, const uint32_t* b, + size_t aRows, size_t aCols); + +void matMulVecPacked2(uint32_t* out, const uint32_t* a, const uint32_t* b_full, + size_t aRows, size_t aCols); + +void matMulVecPacked4(uint32_t* out, const uint32_t* a, const uint32_t* b_full, + size_t aRows, size_t aCols); + +void matMulVecPacked8(uint32_t* out, const uint32_t* a, const uint32_t* b_full, + size_t aRows, size_t aCols); +} + +// Wrapper function that matches the Rust signature +void MatMulVecPacked(uint32_t* out, const uint32_t* a, const uint32_t* b, + size_t a_rows, size_t a_cols, + [[maybe_unused]] size_t b_rows, size_t b_cols) { + assert(a_cols * 4 == b_rows); + // Note: out.len() >= a_rows + 8 should be checked by caller + + // Dispatch based on b_cols + if (b_cols == 1) { + matMulVecPacked(out, a, b, a_rows, a_cols); + } else if (b_cols == 2) { + matMulVecPacked2(out, a, b, a_rows, a_cols); + } else if (b_cols == 4) { + matMulVecPacked4(out, a, b, a_rows, a_cols); + } else if (b_cols == 8) { + matMulVecPacked8(out, a, b, a_rows, a_cols); + } else { + fprintf(stderr, "Error: b_cols must be 1, 2, 4, or 8, got %zu\n", b_cols); + assert(false && "b_cols must be 1, 2, 4, or 8"); + } +} + +PolyMatrixNtt HomomorphicAutomorph(const Params& params, size_t t, size_t t_exp, + const PolyMatrixNtt& ct, + const PolyMatrixNtt& pub_param) { + YACL_ENFORCE(ct.Rows() == static_cast(2)); + YACL_ENFORCE(ct.Cols() == static_cast(1)); + + auto ct_raw = PolyMatrixRaw::Zero(params.PolyLen(), 2, 1); + FromNtt(params, ct_raw, ct); + auto ct_auto = Automorphism(params, ct_raw, t); + + auto ginv_ct = PolyMatrixRaw::Zero(params.PolyLen(), t_exp, 1); + psi::spiral::util::GadgetInvertRdim(params, ginv_ct, ct_auto, 1); + + auto ginv_ct_ntt = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), t_exp, 1); + for (size_t i = 1; i < t_exp; ++i) { + auto pol_src = ginv_ct.Poly(i, 0); + auto pol_dst = ginv_ct_ntt.Poly(i, 0); + ReduceCopy(params, pol_dst, pol_src); + arith::NttForward(params, pol_dst); + } + + auto w_times_ginv_ct = Multiply(params, pub_param, ginv_ct_ntt); + + auto ct_auto_1 = PolyMatrixRaw::Zero(params.PolyLen(), 1, 1); + + std::memcpy(ct_auto_1.Data().data(), + ct_auto.Data().data() + ct_auto.PolyStartIndex(1, 0), + sizeof(uint64_t) * ct_auto.NumWords()); + auto ct_auto_1_ntt = ToNtt(params, ct_auto_1); + + auto res = Add(params, ct_auto_1_ntt.PadTop(1), w_times_ginv_ct); + return res; +} + +PolyMatrixNtt RingPackLwesInner( + const Params& params, size_t ell, size_t start_idx, + const std::vector& rlwe_cts, + const std::vector& pub_params, + const std::pair, std::vector>& + y_constants) { + YACL_ENFORCE_EQ(pub_params.size(), params.PolyLenLog2()); + + if (ell == 0) { + return rlwe_cts[start_idx]; + } + + size_t step = 1ULL << (params.PolyLenLog2() - ell); + size_t even = start_idx; + size_t odd = start_idx + step; + + auto ct_even = RingPackLwesInner(params, ell - 1, even, rlwe_cts, pub_params, + y_constants); + auto ct_odd = RingPackLwesInner(params, ell - 1, odd, rlwe_cts, pub_params, + y_constants); + + const auto& y = y_constants.first[ell - 1]; + const auto& neg_y = y_constants.second[ell - 1]; + + auto y_times_ct_odd = ScalarMultiply(params, y, ct_odd); + auto neg_y_times_ct_odd = ScalarMultiply(params, neg_y, ct_odd); + + auto ct_sum_1 = ct_even; + AddInto(params, ct_sum_1, neg_y_times_ct_odd); + AddInto(params, ct_even, y_times_ct_odd); + + size_t t = (1ULL << ell) + 1; + const auto& pub_param = pub_params[params.PolyLenLog2() - 1 - (ell - 1)]; + auto ct_sum_1_automorphed = + HomomorphicAutomorph(params, t, params.TExpLeft(), ct_sum_1, pub_param); + + return Add(params, ct_even, ct_sum_1_automorphed); +} + +std::pair, std::vector> GenYConstants( + const Params& params) { + std::vector y_constants; + std::vector neg_y_constants; + + for (size_t num_cts_log2 = 1; num_cts_log2 <= params.PolyLenLog2(); + ++num_cts_log2) { + size_t num_cts = 1ULL << num_cts_log2; + + auto y_raw = PolyMatrixRaw::Zero(params.PolyLen(), 1, 1); + size_t idx = params.PolyLen() / num_cts; + y_raw.Data()[idx] = 1ULL; + auto y = ToNtt(params, y_raw); + + auto neg_y_raw = PolyMatrixRaw::Zero(params.PolyLen(), 1, 1); + neg_y_raw.Data()[idx] = params.Modulus() - 1; + auto neg_y = ToNtt(params, neg_y_raw); + + y_constants.push_back(y); + neg_y_constants.push_back(neg_y); + } + + return std::make_pair(y_constants, neg_y_constants); +} + +PolyMatrixNtt RingPackLwes( + const Params& params, const std::vector& b_values, + const std::vector& rlwe_cts, size_t num_cts, + const std::vector& pub_params, + const std::pair, std::vector>& + y_constants) { + YACL_ENFORCE_EQ(b_values.size(), num_cts); + YACL_ENFORCE_EQ(rlwe_cts.size(), num_cts); + + size_t ell = params.PolyLenLog2(); + auto out = + RingPackLwesInner(params, ell, 0, rlwe_cts, pub_params, y_constants); + + auto out_raw = FromNtt(params, out); + for (size_t z = 0; z < params.PolyLen(); ++z) { + uint128_t b_val_u128 = static_cast(b_values[z]); + uint128_t poly_len_u128 = static_cast(params.PolyLen()); + uint128_t prod = b_val_u128 * poly_len_u128; + uint64_t val = arith::BarrettReductionU128(params, prod); + + size_t idx = out_raw.PolyStartIndex(1, 0) + z; + uint64_t sum = out_raw.Data()[idx] + val; + out_raw.Data()[idx] = arith::BarrettU64(params, sum); + } + + return ToNtt(params, out_raw); +} + +std::vector NegacyclicPerm(absl::Span a, size_t shift, + uint64_t modulus) { + size_t n = a.size(); + std::vector out(n); + + for (size_t i = 0; i <= shift; ++i) { + out[i] = a[shift - i]; + } + + for (size_t i = shift + 1; i < n; ++i) { + size_t src_idx = n - (i - shift); + uint64_t val = a[src_idx] % modulus; + if (val == 0) { + out[i] = 0; + } else { + out[i] = modulus - val; + } + } + + return out; +} + +std::vector NegacyclicMatrix(absl::Span a, + uint64_t modulus) { + size_t n = a.size(); + std::vector out(n * n, 0); + + for (size_t i = 0; i < n; ++i) { + std::vector perm = NegacyclicPerm(a, i, modulus); + for (size_t j = 0; j < n; ++j) { + out[j * n + i] = perm[j]; + } + } + + return out; +} + +PolyMatrixNtt CondenseMatrix(const Params& params, const PolyMatrixNtt& a) { + PolyMatrixNtt res = PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), + a.Rows(), a.Cols()); + for (size_t i = 0; i < a.Rows(); ++i) { + for (size_t j = 0; j < a.Cols(); ++j) { + auto res_poly = res.Poly(i, j); + const auto a_poly = a.Poly(i, j); + if (a_poly.size() < 2 * params.PolyLen()) { + continue; + } + for (size_t z = 0; z < params.PolyLen(); ++z) { + res_poly[z] = a_poly[z] | (a_poly[z + params.PolyLen()] << 32); + } + } + } + + return res; +} + +// from ypir/packing.rs +PolyMatrixNtt RotationPoly(const Params& params, size_t amount) { + PolyMatrixRaw res = PolyMatrixRaw::Zero(params.PolyLen(), 1, 1); + + if (amount < res.Data().size()) { + res.Data()[amount] = 1; + } else { + YACL_THROW("Rotation amount exceeds polynomial length"); + } + + return ToNtt(params, res); +} +PolyMatrixNtt PackSingleLwe(const Params& params, + const std::vector& pub_params, + const PolyMatrixNtt& lwe_ct) { + PolyMatrixNtt cur_r = lwe_ct; + + size_t log_n = params.PolyLenLog2(); + + for (size_t i = 0; i < log_n; ++i) { + size_t t = (params.PolyLen() / (1ULL << i)) + 1; + + const auto& pub_param = pub_params[i]; + + auto tau_of_r = + HomomorphicAutomorph(params, t, params.TExpLeft(), cur_r, pub_param); + AddInto(params, cur_r, tau_of_r); + } + + return cur_r; +} + +PolyMatrixNtt PackUsingSingleWithOffset( + const Params& params, const std::vector& pub_params, + const std::vector& cts, size_t offset) { + PolyMatrixNtt res = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 2, 1); + + SPDLOG_INFO( + "PackUsingSingleWithOffset: offset={}, cts.size()={}, poly_len={}", + offset, cts.size(), params.PolyLen()); + + for (size_t i = 0; i < cts.size(); ++i) { + PolyMatrixNtt packed_single = PackSingleLwe(params, pub_params, cts[i]); + + size_t rotation_amount = offset + i; + if (rotation_amount >= params.PolyLen()) { + SPDLOG_ERROR("ERROR: rotation_amount {} >= poly_len {} (offset={}, i={})", + rotation_amount, params.PolyLen(), offset, i); + YACL_THROW( + "Rotation amount exceeds polynomial length in " + "PackUsingSingleWithOffset"); + } + + PolyMatrixNtt rotation = RotationPoly(params, rotation_amount); + PolyMatrixNtt rotated = ScalarMultiply(params, rotation, packed_single); + + AddInto(params, res, rotated); + } + + return res; +} + +std::vector PrepPackLwes(const Params& params, + absl::Span lwe_cts, + size_t cols_to_do) { + size_t poly_len = params.PolyLen(); + size_t expected_size = poly_len * (poly_len + 1); + + if (lwe_cts.size() != expected_size) { + throw std::runtime_error("PrepPackLwes: lwe_cts size mismatch"); + } + if (cols_to_do != poly_len) { + throw std::runtime_error("PrepPackLwes: cols_to_do must equal poly_len"); + } + + std::vector rlwe_cts; + rlwe_cts.reserve(cols_to_do); + + for (size_t i = 0; i < cols_to_do; ++i) { + auto rlwe_ct = PolyMatrixRaw::Zero(params.PolyLen(), 2, 1); + std::vector poly; + poly.reserve(poly_len); + for (size_t j = 0; j < poly_len; ++j) { + poly.push_back(lwe_cts[j * poly_len + i]); + } + + auto nega = NegacyclicPerm(poly, 0, params.Modulus()); + + // Use Poly(0, 0) to get the first polynomial + auto first_poly = rlwe_ct.Poly(0, 0); + for (size_t j = 0; j < poly_len; ++j) { + first_poly[j] = nega[j]; + } + + rlwe_cts.push_back(ToNtt(params, rlwe_ct)); + } + + return rlwe_cts; +} + +std::vector> PrepPackManyLwes( + const Params& params, absl::Span lwe_cts, + size_t num_rlwe_outputs) { + size_t poly_len = params.PolyLen(); + + size_t expected_total_size = (poly_len + 1) * (num_rlwe_outputs * poly_len); + + if (lwe_cts.size() != expected_total_size) { + throw std::runtime_error("PrepPackManyLwes: Total size mismatch"); + } + + std::vector> vecs; + vecs.reserve(num_rlwe_outputs); + + for (size_t i = 0; i < num_rlwe_outputs; ++i) { + std::vector v; + // 每个 chunk 的大小应该是 poly_len * (poly_len + 1) + v.reserve(poly_len * (poly_len + 1)); + for (size_t j = 0; j < poly_len + 1; ++j) { + size_t start_idx = j * (num_rlwe_outputs * poly_len) + i * poly_len; + for (size_t k = 0; k < poly_len; ++k) { + v.push_back(lwe_cts[start_idx + k]); + } + } + vecs.push_back(std::move(v)); + } + std::vector> res; + res.reserve(num_rlwe_outputs); + + for (size_t i = 0; i < num_rlwe_outputs; ++i) { + res.push_back(PrepPackLwes(params, vecs[i], poly_len)); + } + + return res; +} + +std::vector> GenerateAutomorphTablesBruteForce( + const Params& params) { + std::vector> tables; + + tables.reserve(params.PolyLenLog2()); + + for (size_t i = params.PolyLenLog2(); i >= 1; --i) { + size_t poly_len = params.PolyLen(); + std::vector table_candidate(poly_len); + while (true) { + uint64_t t = (1ULL << i) + 1; + PolyMatrixRaw poly = PolyMatrixRaw::Random(params, 1, 1); + PolyMatrixNtt poly_ntt = ToNtt(params, poly); + PolyMatrixRaw poly_auto(params.PolyLen(), 1, 1); + Automorphism(params, poly_auto, poly, t); + PolyMatrixNtt poly_auto_ntt = ToNtt(params, poly_auto); + auto pol_orig = poly_ntt.Poly(0, 0); + auto pol_auto = poly_auto_ntt.Poly(0, 0); + bool must_redo = false; + for (size_t src_idx = 0; src_idx < poly_len; ++src_idx) { + uint64_t target_val = pol_orig[src_idx]; + int found_count = 0; + size_t found_pos = 0; + for (size_t dest_idx = 0; dest_idx < poly_len; ++dest_idx) { + if (pol_auto[dest_idx] == target_val) { + found_count++; + found_pos = dest_idx; + } + } + if (found_count != 1) { + must_redo = true; + break; + } + table_candidate[found_pos] = src_idx; + } + + if (!must_redo) { + break; + } + } + tables.push_back(std::move(table_candidate)); + } + + return tables; +} + +std::tuple, + std::vector>> +PrecomputePack(const Params& params, size_t poly_len_log2, + const std::vector& prepacked, + const std::vector& pub_params, + const std::pair, + std::vector>& y_constants) { + YACL_ENFORCE_EQ(pub_params.size(), params.PolyLenLog2(), + "pub_params size must match poly_len_log2"); + YACL_ENFORCE_EQ(params.CrtCount(), 2UL, "CRT count must be 2"); + + std::vector working_set = prepacked; + PolyMatrixNtt y_times_ct_odd = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 2, 1); + + PolyMatrixNtt neg_y_times_ct_odd = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 2, 1); + PolyMatrixNtt ct_sum_1 = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 2, 1); + + PolyMatrixRaw ct_raw = PolyMatrixRaw::Zero(params.PolyLen(), 1, 1); + PolyMatrixRaw ct_auto = PolyMatrixRaw::Zero(params.PolyLen(), 1, 1); + + size_t t_exp = params.TExpLeft(); + PolyMatrixRaw ginv_ct = + PolyMatrixRaw::Zero(params.PolyLen(), t_exp, 1); // Fixed: poly_len first + PolyMatrixNtt ginv_ct_ntt = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), t_exp, 1); + + PolyMatrixNtt ct_auto_1_ntt = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 1, 1); + PolyMatrixNtt w_times_ginv_ct = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 2, 1); + PolyMatrixNtt scratch = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 2, 1); + std::vector v_trgsw; + v_trgsw.reserve(poly_len_log2); + + for (size_t cur_ell = 1; cur_ell <= poly_len_log2; ++cur_ell) { + size_t num_in = 1ULL << (poly_len_log2 - cur_ell + 1); + size_t num_out = num_in >> 1; + + const auto& y = y_constants.first[cur_ell - 1]; + const auto& neg_y = y_constants.second[cur_ell - 1]; + + for (size_t i = 0; i < num_out; ++i) { + auto& ct_even = working_set[i]; + const auto& ct_odd = working_set[i + num_out]; + + // y_times_ct_odd = y * ct_odd + ScalarMultiply(params, y_times_ct_odd, y, ct_odd); + // neg_y_times_ct_odd = -y * ct_odd + ScalarMultiply(params, neg_y_times_ct_odd, neg_y, ct_odd); + + // ct_sum_1 = ct_even + neg_y * ct_odd + ct_sum_1.CopyInto(ct_even, 0, 0); + AddInto(params, ct_sum_1, neg_y_times_ct_odd); + + AddInto(params, ct_even, y_times_ct_odd); + + uint64_t t = (1ULL << cur_ell) + 1; + FromNttScratch(params, ct_raw, scratch.Data(), ct_sum_1); + + Automorphism(params, ct_auto, ct_raw, t); + util::GadgetInvertRdim(params, ginv_ct, ct_auto, 1); + ToNtt(params, ginv_ct_ntt, ginv_ct); + + auto scratch_poly = scratch.Poly(0, 0); // Get the stored 2nd row data + AutomorphismPolyUncrtd(params, ct_auto_1_ntt.Data(), scratch_poly, t); + arith::NttForward(params, ct_auto_1_ntt.Data()); + + v_trgsw.push_back(CondenseMatrix(params, ginv_ct_ntt)); + + size_t param_idx = poly_len_log2 - cur_ell; + const auto& pub_param = pub_params[param_idx]; + + // W * Ginv + std::fill(w_times_ginv_ct.Data().begin(), w_times_ginv_ct.Data().end(), + 0); + MultiplyNoReduce(w_times_ginv_ct, pub_param, ginv_ct_ntt, 0); + + AddIntoAt(params, ct_even, ct_auto_1_ntt, 1, 0); + + AddInto(params, ct_even, w_times_ginv_ct); + } + } + + auto combined_hints = GenerateAutomorphTablesBruteForce(params); + return std::make_tuple(working_set[0], v_trgsw, combined_hints); +} + +inline uint64_t FastBarrettRawU64(uint64_t input, uint64_t const_ratio_1, + uint64_t modulus) { + uint128_t mul = static_cast(input) * const_ratio_1; + uint64_t tmp = static_cast(mul >> 64); + uint64_t res = input - tmp * modulus; + return res; +} + +void FastAddInto(const Params& params, PolyMatrixNtt& res, + const PolyMatrixNtt& a) { + if (res.Rows() != a.Rows() || res.Cols() != a.Cols()) { + throw std::runtime_error("FastAddInto: Dimension mismatch"); + } + + size_t poly_len = params.PolyLen(); + size_t crt_count = params.CrtCount(); + + uint64_t* res_ptr = const_cast(res.Data().data()); + const uint64_t* a_ptr = a.Data().data(); + size_t num_polys = res.Rows() * res.Cols(); + + for (size_t p = 0; p < num_polys; ++p) { + for (size_t c = 0; c < crt_count; ++c) { + size_t poly_offset = p * (crt_count * poly_len); + size_t crt_offset = poly_offset + c * poly_len; + + for (size_t i = 0; i < poly_len; ++i) { + size_t idx = crt_offset + i; + uint64_t val = res_ptr[idx] + a_ptr[idx]; + res_ptr[idx] = + FastBarrettRawU64(val, params.BarrettCr1(c), params.Moduli(c)); + } + } + } +} + +void FastAddIntoNoReduce(PolyMatrixNtt& res, const PolyMatrixNtt& a) { + if (res.Rows() != a.Rows() || res.Cols() != a.Cols()) { + throw std::runtime_error("FastAddIntoNoReduce: Dimension mismatch"); + } + size_t len = res.Data().size(); + for (size_t i = 0; i < len; ++i) { + res.Data()[i] += a.Data()[i]; + } +} + +void MultiplyPolyAvx(uint64_t* res, const uint64_t* a, const uint64_t* b, + size_t len) { + for (size_t i = 0; i < len; ++i) { + uint64_t x = static_cast(static_cast(a[i])); + uint64_t y = static_cast(static_cast(b[i])); + res[i] = x * y; + } +} + +void ScalarMultiplyAvx(const Params& params, PolyMatrixNtt& res, + const PolyMatrixNtt& a, const PolyMatrixNtt& b) { + if (a.Rows() != 1 || a.Cols() != 1) { + throw std::runtime_error( + "ScalarMultiplyAvx: 'a' must be a 1x1 matrix (scalar)"); + } + + if (res.Rows() != b.Rows() || res.Cols() != b.Cols()) { + throw std::runtime_error( + "ScalarMultiplyAvx: 'res' and 'b' dimension mismatch"); + } + + // poly_size = poly_len * crt_count + size_t poly_size = params.PolyLen() * params.CrtCount(); + + const uint64_t* pol2_ptr = a.Data().data(); + + const uint64_t* b_ptr = b.Data().data(); + + uint64_t* res_ptr = const_cast(res.Data().data()); + + size_t total_polys = b.Rows() * b.Cols(); + + for (size_t i = 0; i < total_polys; ++i) { + size_t offset = i * poly_size; + + const uint64_t* pol1_ptr = b_ptr + offset; + + uint64_t* res_poly_ptr = res_ptr + offset; + + MultiplyPolyAvx(res_poly_ptr, pol1_ptr, pol2_ptr, poly_size); + } +} + +void FastMultiplyNoReduce(const Params& params, PolyMatrixNtt& res, + const PolyMatrixNtt& a, const PolyMatrixNtt& b, + size_t /*start_inner_dim*/) { + if (res.Rows() != 1 || res.Cols() != 1) { + throw std::runtime_error("FastMultiplyNoReduce: Result must be 1x1"); + } + if (res.Rows() != a.Rows() || res.Cols() != b.Cols()) { + throw std::runtime_error( + "FastMultiplyNoReduce: Dimension mismatch (Result)"); + } + if (a.Cols() != b.Rows()) { + throw std::runtime_error( + "FastMultiplyNoReduce: Dimension mismatch (Inner)"); + } + + if (params.CrtCount() != 2) { + throw std::runtime_error("FastMultiplyNoReduce: Requires CRT count = 2"); + } + + uint64_t* res_ptr = const_cast(res.Data().data()); + const uint64_t* a_ptr = a.Data().data(); + const uint64_t* b_ptr = b.Data().data(); + + size_t pol_sz = params.PolyLen(); + size_t k_dim = a.Cols(); + + for (size_t idx = 0; idx < pol_sz; ++idx) { + uint64_t sum_lo = 0; + uint64_t sum_hi = 0; + + for (size_t k = 0; k < k_dim; ++k) { + size_t offset = k * 2 * pol_sz + idx; + + uint64_t x = a_ptr[offset]; + uint64_t y = b_ptr[offset]; + + uint64_t x_lo = x & 0xFFFFFFFF; + uint64_t x_hi = x >> 32; + + uint64_t y_lo = y & 0xFFFFFFFF; + uint64_t y_hi = y >> 32; + + sum_lo += x_lo * y_lo; + sum_hi += x_hi * y_hi; + } + + res_ptr[idx] = sum_lo; + res_ptr[pol_sz + idx] = sum_hi; + } +} + +void ApplyAutomorphNttRaw(const Params& params, const uint64_t* poly, + uint64_t* out, size_t t, + const std::vector>& tables) { + size_t poly_len = params.PolyLen(); + if (t <= 1) throw std::runtime_error("ApplyAutomorphNttRaw: t must be > 1"); + + size_t val = poly_len / (t - 1); + // C++20 + // size_t table_idx = std::countr_zero(val); + size_t table_idx = __builtin_ctzll(val); + + if (table_idx >= tables.size()) { + throw std::runtime_error("ApplyAutomorphNttRaw: Table index out of bounds"); + } + + const std::vector& table = tables[table_idx]; + + for (size_t i = 0; i < poly_len; ++i) { + out[i] += poly[table[i]]; + } +} + +void ApplyAutomorphNtt(const Params& params, + const std::vector>& tables, + const PolyMatrixNtt& mat, PolyMatrixNtt& res, size_t t) { + if (mat.Rows() != res.Rows() || mat.Cols() != res.Cols()) { + throw std::runtime_error("ApplyAutomorphNtt: Dimension mismatch"); + } + + const uint64_t* mat_ptr = mat.Data().data(); + + uint64_t* res_ptr = const_cast(res.Data().data()); + + size_t poly_len = params.PolyLen(); + size_t total_elements = mat.Data().size(); + + for (size_t offset = 0; offset < total_elements; offset += poly_len) { + ApplyAutomorphNttRaw(params, mat_ptr + offset, res_ptr + offset, t, tables); + } +} + +void FastReduce(const Params& params, PolyMatrixNtt& res) { + size_t crt_count = params.CrtCount(); + size_t poly_len = params.PolyLen(); + + for (size_t m = 0; m < crt_count; ++m) { + for (size_t i = 0; i < poly_len; ++i) { + size_t idx = m * poly_len + i; + res.Data()[idx] = arith::BarrettCoeffU64(params, res.Data()[idx], m); + } + } +} + +PolyMatrixNtt PackUsingPrecompVals( + const Params& params, size_t ell, + absl::Span pub_params, + absl::Span b_values, const PolyMatrixNtt& precomp_res, + absl::Span precomp_vals, + const std::vector>& precomp_tables, + const std::pair, std::vector>& + y_constants) { + size_t initial_capacity = 1ULL << (ell - 1); + std::vector working_set; + working_set.reserve(initial_capacity); + + for (size_t i = 0; i < initial_capacity; ++i) { + working_set.push_back( + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 1, 1)); + } + + auto y_times_ct_odd = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 1, 1); + auto neg_y_times_ct_odd = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 1, 1); + auto ct_sum_1 = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 1, 1); + auto w_times_ginv_ct = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 1, 1); + + size_t idx_precomp = 0; + + for (size_t cur_ell = 1; cur_ell <= ell; ++cur_ell) { + size_t num_in = 1ULL << (ell - cur_ell + 1); + size_t num_out = num_in >> 1; + + if (num_in == params.PolyLen()) { + num_in = num_out; + } + + for (size_t i = 0; i < num_out; ++i) { + PolyMatrixNtt& ct_even = working_set[i]; + + const auto& y = y_constants.first[cur_ell - 1]; + const auto& neg_y = y_constants.second[cur_ell - 1]; + + if (cur_ell > 1) { + PolyMatrixNtt& ct_odd = working_set[num_out + i]; + ScalarMultiplyAvx(params, y_times_ct_odd, y, ct_odd); + ScalarMultiplyAvx(params, neg_y_times_ct_odd, neg_y, ct_odd); + } + + if (cur_ell > 1) { + std::copy(ct_even.Data().begin(), ct_even.Data().end(), + const_cast(ct_sum_1.Data().data())); + + FastAddIntoNoReduce(ct_sum_1, neg_y_times_ct_odd); + FastAddIntoNoReduce(ct_even, y_times_ct_odd); + } + + const PolyMatrixNtt* ct_ptr = &ct_sum_1; + size_t t = (1ULL << cur_ell) + 1; + + const auto& cur_ginv_ct_ntt = precomp_vals[idx_precomp]; + idx_precomp++; + + size_t w_idx = params.PolyLenLog2() - cur_ell; + const auto& w = pub_params[w_idx]; + + FastMultiplyNoReduce(params, w_times_ginv_ct, w, cur_ginv_ct_ntt, 0); + + if (cur_ell > 1) { + ApplyAutomorphNtt(params, precomp_tables, *ct_ptr, ct_even, t); + + if (i < num_out / 2 && ((cur_ell - 1) % 5 != 0)) { + FastAddIntoNoReduce(ct_even, w_times_ginv_ct); + } else { + FastAddInto(params, ct_even, w_times_ginv_ct); + } + } else { + if (i < num_out / 2) { + FastAddIntoNoReduce(ct_even, w_times_ginv_ct); + } else { + FastAddInto(params, ct_even, w_times_ginv_ct); + } + } + } + } + + if (idx_precomp != precomp_vals.size()) { + throw std::runtime_error("PackUsingPrecompVals: idx_precomp mismatch"); + } + + PolyMatrixNtt resulting_row_1 = working_set[0]; + FastReduce(params, resulting_row_1); + + PolyMatrixNtt res = precomp_res; + + size_t poly_len = params.PolyLen(); + size_t crt_count = params.CrtCount(); + size_t row_1_offset_ntt = crt_count * poly_len; + + uint64_t* res_data = const_cast(res.Data().data()); + const uint64_t* row_1_src = resulting_row_1.Data().data(); + + std::copy(row_1_src, row_1_src + (crt_count * poly_len), + res_data + row_1_offset_ntt); + + PolyMatrixRaw out_raw = FromNtt(params, res); + uint64_t modulus = params.Modulus(); + + for (size_t z = 0; z < poly_len; ++z) { + uint128_t val_u128 = static_cast(b_values[z]) * poly_len; + uint64_t val = arith::BarrettReductionU128(params, val_u128); + + size_t idx = poly_len + z; + + out_raw.Data()[idx] += val; + if (out_raw.Data()[idx] >= modulus) { + out_raw.Data()[idx] -= modulus; + } + } + + return ToNtt(params, out_raw); +} + +std::vector PackManyLwes( + const Params& params, + const std::vector>& prep_rlwe_cts, + const Precomp& precomp, absl::Span b_values, + size_t num_rlwe_outputs, + const std::vector& pack_pub_params_row_1s, + const std::pair, std::vector>& + y_constants) { + if (prep_rlwe_cts.size() != num_rlwe_outputs) { + throw std::runtime_error("PackManyLwes: prep_rlwe_cts size mismatch"); + } + + if (!prep_rlwe_cts.empty() && prep_rlwe_cts[0].size() != params.PolyLen()) { + throw std::runtime_error("PackManyLwes: prep_rlwe_cts inner size mismatch"); + } + + if (b_values.size() != num_rlwe_outputs * params.PolyLen()) { + throw std::runtime_error("PackManyLwes: b_values size mismatch"); + } + + if (precomp.size() != num_rlwe_outputs) { + throw std::runtime_error("PackManyLwes: precomp size mismatch"); + } + + std::vector res; + res.reserve(num_rlwe_outputs); + + size_t poly_len = params.PolyLen(); + size_t poly_len_log2 = params.PolyLenLog2(); + + for (size_t i = 0; i < num_rlwe_outputs; ++i) { + const auto& [precomp_res, precomp_vals, precomp_tables] = precomp[i]; + + auto b_values_slice = b_values.subspan(i * poly_len, poly_len); + + PolyMatrixNtt packed = PackUsingPrecompVals( + params, poly_len_log2, + absl::MakeConstSpan(pack_pub_params_row_1s), // Vector -> Span + b_values_slice, precomp_res, + absl::MakeConstSpan(precomp_vals), // Vector -> Span + precomp_tables, y_constants); + + res.push_back(std::move(packed)); + } + + return res; +} + +PolyMatrixNtt GetRegevSample(const Params& params, const PolyMatrixRaw& sk_reg, + yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub) { + auto a = PolyMatrixRaw::RandomPrg(params, 1, 1, rng_pub); + auto a_ntt = ToNtt(params, a); + auto a_inv = ToNtt(params, Negate(params, a)); + DiscreteGaussian dg(params.NoiseWidth()); + auto e = Noise(params, 1, 1, dg, rng); + + auto e_ntt = ToNtt(params, e); + auto sk_reg_ntt = ToNtt(params, sk_reg); + auto b_p = Multiply(params, sk_reg_ntt, a_ntt); + auto b = Add(params, e_ntt, b_p); + auto p = PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 2, 1); + p.CopyInto(a_inv, 0, 0); + p.CopyInto(b, 1, 0); + + return p; +} + +PolyMatrixNtt GetFreshRegevPublicKey(const Params& params, + const PolyMatrixRaw& sk_reg, size_t m, + yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub) { + auto p = PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), 2, m); + for (size_t i = 0; i < m; ++i) { + p.CopyInto(GetRegevSample(params, sk_reg, rng, rng_pub), 0, i); + } + return p; +} + +std::vector RawGenerateExpansionParams( + const Params& params, const PolyMatrixRaw& sk_reg, size_t num_exp, + size_t m_exp, yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub) { + auto g_exp = util::BuildGadget(params, 1, m_exp); + auto g_exp_ntt = ToNtt(params, g_exp); + + std::vector res; + res.reserve(num_exp); + + for (size_t i = 0; i < num_exp; ++i) { + size_t t = (params.PolyLen() / (1ULL << i)) + 1; + auto tau_sk_reg = Automorphism(params, sk_reg, t); + auto tau_sk_reg_ntt = ToNtt(params, tau_sk_reg); + auto prod = Multiply(params, tau_sk_reg_ntt, g_exp_ntt); + auto sample = GetFreshRegevPublicKey(params, sk_reg, m_exp, rng, rng_pub); + + PolyMatrixNtt padded_prod = prod.PadTop(1); + auto w_exp_i = Add(params, sample, padded_prod); + res.push_back(std::move(w_exp_i)); + } + return res; +} + +// from kernal.rs + +template +void FastBatchedDotProduct(const Params& params, uint64_t* c, const uint64_t* a, + size_t a_elems, const T* b_t, size_t b_rows, + size_t b_cols) { + if (a_elems != b_rows) { + throw std::runtime_error("FastBatchedDotProduct: Dimension mismatch"); + } + + const size_t chunk_size = std::min(65536, a_elems); + const size_t num_chunks = a_elems / chunk_size; + + for (size_t k_outer = 0; k_outer < num_chunks; ++k_outer) { + const uint64_t* a_chunk_ptr = a + (k_outer * chunk_size); + for (size_t j = 0; j < b_cols; ++j) { + const T* b_ptr = b_t + (j * b_rows) + (k_outer * chunk_size); + uint64_t total_sum_lo = 0; + uint64_t total_sum_hi = 0; + for (size_t k_inner = 0; k_inner < chunk_size; ++k_inner) { + uint64_t a_val = a_chunk_ptr[k_inner]; + uint64_t b_val = static_cast(b_ptr[k_inner]); + uint64_t a_lo = a_val & 0xFFFFFFFF; + uint64_t a_hi = a_val >> 32; + total_sum_lo += a_lo * b_val; + total_sum_hi += a_hi * b_val; + } + + uint64_t lo = arith::BarrettCoeffU64(params, total_sum_lo, 0); + uint64_t hi = arith::BarrettCoeffU64(params, total_sum_hi, 1); + + uint64_t res = params.CrtCompose2(lo, hi); + + uint64_t sum = c[j] + res; + c[j] = arith::BarrettU64(params, sum); + } + } +} + +// from serialize.rs +std::vector UnpackVecPm(const Params& params, size_t rows, + size_t cols, + absl::Span data) { + if (params.CrtCount() != 2) { + throw std::runtime_error("UnpackVecPm: Params CRT count must be 2"); + } + + size_t poly_len = params.PolyLen(); + size_t chunk_size = rows * cols * poly_len; + + if (chunk_size == 0) { + return {}; + } + + if (data.size() % chunk_size != 0) { + throw std::runtime_error( + "UnpackVecPm: Data size not aligned with matrix dimensions"); + } + + size_t num_matrices = data.size() / chunk_size; + std::vector v_cts; + v_cts.reserve(num_matrices); + + for (size_t i = 0; i < num_matrices; ++i) { + const uint64_t* in_data_ptr = data.data() + (i * chunk_size); + + auto ct = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), rows, cols); + + uint64_t* out_data_ptr = const_cast(ct.Data().data()); + + for (size_t row = 0; row < rows; ++row) { + for (size_t col = 0; col < cols; ++col) { + size_t in_offs = (row * cols + col) * poly_len; + + size_t out_offs = (row * cols + col) * 2 * poly_len; + + for (size_t z = 0; z < poly_len; ++z) { + out_data_ptr[out_offs + z] = in_data_ptr[in_offs + z]; + } + } + } + v_cts.push_back(std::move(ct)); + } + + return v_cts; +} +PolyMatrixNtt UncondenseMatrix(const Params& params, const PolyMatrixNtt& a) { + size_t rows = a.Rows(); + size_t cols = a.Cols(); + size_t poly_len = params.PolyLen(); + + PolyMatrixNtt res = + PolyMatrixNtt::Zero(params.CrtCount(), params.PolyLen(), rows, cols); + + for (size_t i = 0; i < rows; ++i) { + for (size_t j = 0; j < cols; ++j) { + auto res_poly = res.Poly(i, j); + const auto a_poly = a.Poly(i, j); + + for (size_t z = 0; z < poly_len; ++z) { + uint64_t val = a_poly[z]; + res_poly[z] = val & 0xFFFFFFFFULL; + res_poly[z + poly_len] = val >> 32; + } + } + } + + return res; +} +// from modulus switch.rs +std::vector ModulusSwitch(const Params& params, + const PolyMatrixRaw& poly_matrix, + uint64_t q_prime_1, uint64_t q_prime_2) { + if (poly_matrix.Rows() != 2 || poly_matrix.Cols() != 1) { + throw std::runtime_error("ModulusSwitch: Matrix must be 2x1"); + } + size_t poly_len = params.PolyLen(); + uint64_t current_modulus = params.Modulus(); + + size_t bits_for_row0 = static_cast(std::ceil(std::log2(q_prime_2))); + size_t bits_for_row1 = static_cast(std::ceil(std::log2(q_prime_1))); + + size_t total_sz_bits = (bits_for_row0 + bits_for_row1) * poly_len; + size_t total_sz_bytes = (total_sz_bits + 7) / 8; + + std::vector res(total_sz_bytes, 0); + + size_t bit_offs = 0; + + const uint64_t* row_0_ptr = poly_matrix.Data().data(); + + for (size_t z = 0; z < poly_len; ++z) { + uint64_t val = row_0_ptr[z]; + uint64_t val_rescaled = arith::Rescale(val, current_modulus, q_prime_2); + + util::WriteArbitraryBits(res, val_rescaled, bit_offs, bits_for_row0); + bit_offs += bits_for_row0; + } + + const uint64_t* row_1_ptr = poly_matrix.Data().data() + poly_len; + + for (size_t z = 0; z < poly_len; ++z) { + uint64_t val = row_1_ptr[z]; + uint64_t val_rescaled = arith::Rescale(val, current_modulus, q_prime_1); + + util::WriteArbitraryBits(res, val_rescaled, bit_offs, bits_for_row1); + bit_offs += bits_for_row1; + } + + return res; +} + +// Explicit instantiation of template functions +template void FastBatchedDotProduct(const Params& params, uint64_t* c, + const uint64_t* a, size_t a_elems, + const uint8_t* b_t, size_t b_rows, + size_t b_cols); +template void FastBatchedDotProduct(const Params& params, uint64_t* c, + const uint64_t* a, size_t a_elems, + const uint16_t* b_t, + size_t b_rows, size_t b_cols); +template void FastBatchedDotProduct(const Params& params, uint64_t* c, + const uint64_t* a, size_t a_elems, + const uint32_t* b_t, + size_t b_rows, size_t b_cols); + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/legacy/util.h b/psi/algorithm/ypir/legacy/util.h new file mode 100644 index 00000000..eff77dc9 --- /dev/null +++ b/psi/algorithm/ypir/legacy/util.h @@ -0,0 +1,129 @@ +#pragma once + +#include + +#include "absl/types/span.h" + +#include "psi/algorithm/spiral/params.h" +#include "psi/algorithm/spiral/poly_matrix.h" +#include "psi/algorithm/ypir/legacy/types.h" +namespace psi::ypir { + +using namespace psi::spiral; +std::vector ConcatHorizontal( + const std::vector>& v_a, size_t a_rows, + size_t a_cols); + +void MatMulVecPacked(uint32_t* out, const uint32_t* a, const uint32_t* b, + size_t a_rows, size_t a_cols, size_t b_rows, + size_t b_cols); + +PolyMatrixNtt HomomorphicAutomorph(const Params& params, size_t t, size_t t_exp, + const PolyMatrixNtt& ct, + const PolyMatrixNtt& pub_param); + +PolyMatrixNtt RingPackLwesInner( + const Params& params, size_t ell, size_t start_idx, + const std::vector& rlwe_cts, + const std::vector& pub_params, + const std::pair, std::vector>& + y_constants); + +std::pair, std::vector> GenYConstants( + const Params& params); + +PolyMatrixNtt RingPackLwes( + const Params& params, const std::vector& b_values, + const std::vector& rlwe_cts, size_t num_cts, + const std::vector& pub_params, + const std::pair, std::vector>& + y_constants); + +std::vector NegacyclicPerm(absl::Span a, size_t shift, + uint64_t modulus); +std::vector NegacyclicMatrix(absl::Span a, + uint64_t modulus); +PolyMatrixNtt CondenseMatrix(const Params& params, const PolyMatrixNtt& a); +PolyMatrixNtt UncondenseMatrix(const Params& params, const PolyMatrixNtt& a); +PolyMatrixNtt PackUsingSingleWithOffset( + const Params& params, const std::vector& pub_params, + const std::vector& cts, size_t offset); + +std::vector PrepPackLwes(const Params& params, + absl::Span lwe_cts, + size_t cols_to_do); + +std::vector> PrepPackManyLwes( + const Params& params, absl::Span lwe_cts, + size_t num_rlwe_outputs); + +std::vector> GenerateAutomorphTablesBruteForce( + const Params& params); + +std::tuple, + std::vector>> +PrecomputePack(const Params& params, size_t poly_len_log2, + const std::vector& prepacked, + const std::vector& pub_params, + const std::pair, + std::vector>& y_constants); +void FastAddIntoNoReduce(PolyMatrixNtt& res, const PolyMatrixNtt& a); +void MultiplyPolyAvx(uint64_t* res, const uint64_t* a, const uint64_t* b, + size_t len); +void ScalarMultiplyAvx(const Params& params, PolyMatrixNtt& res, + const PolyMatrixNtt& a, const PolyMatrixNtt& b); + +void FastMultiplyNoReduce(const Params& params, PolyMatrixNtt& res, + const PolyMatrixNtt& a, const PolyMatrixNtt& b, + size_t start_inner_dim = 0); +void ApplyAutomorphNttRaw(const Params& params, const uint64_t* poly, + uint64_t* out, size_t t, + const std::vector>& tables); +void ApplyAutomorphNtt(const Params& params, + const std::vector>& tables, + const PolyMatrixNtt& mat, PolyMatrixNtt& res, size_t t); +void FastReduce(const Params& params, PolyMatrixNtt& res); +PolyMatrixNtt PackUsingPrecompVals( + const Params& params, + size_t ell, // params.poly_len_log2 + absl::Span pub_params, + absl::Span b_values, const PolyMatrixNtt& precomp_res, + absl::Span precomp_vals, + const std::vector>& precomp_tables, + const std::pair, std::vector>& + y_constants); + +std::vector PackManyLwes( + const Params& params, + const std::vector>& prep_rlwe_cts, + const Precomp& precomp, absl::Span b_values, + size_t num_rlwe_outputs, + const std::vector& pack_pub_params_row_1s, + const std::pair, std::vector>& + y_constants); + +PolyMatrixNtt GetRegevSample(const Params& params, const PolyMatrixRaw& sk_reg, + yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub); +PolyMatrixNtt GetFreshRegevPublicKey(const Params& params, + const PolyMatrixRaw& sk_reg, size_t m, + yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub); +std::vector RawGenerateExpansionParams( + const Params& params, const PolyMatrixRaw& sk_reg, size_t num_exp, + size_t m_exp, yacl::crypto::Prg& rng, + yacl::crypto::Prg& rng_pub); + +template +void FastBatchedDotProduct(const Params& params, uint64_t* c, const uint64_t* a, + size_t a_elems, const T* b_t, size_t b_rows, + size_t b_cols); + +std::vector UnpackVecPm(const Params& params, size_t rows, + size_t cols, + absl::Span data); + +std::vector ModulusSwitch(const Params& params, + const PolyMatrixRaw& poly_matrix, + uint64_t q_prime_1, uint64_t q_prime_2); +} // namespace psi::ypir \ No newline at end of file diff --git a/psi/algorithm/ypir/legacy/ypir_params.cc b/psi/algorithm/ypir/legacy/ypir_params.cc new file mode 100644 index 00000000..5d5f1226 --- /dev/null +++ b/psi/algorithm/ypir/legacy/ypir_params.cc @@ -0,0 +1,441 @@ +// Copyright 2026 The secretflow authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "psi/algorithm/ypir/legacy/ypir_params.h" + +#include +#include +#include +#include + +#include "yacl/base/exception.h" + +namespace psi::ypir::ypir_internal { +namespace { + +uint64_t Log2Exact(uint64_t n) { + uint64_t r = 0; + while (n > 1) { + n >>= 1; + ++r; + } + return r; +} + +// Thread-local buffer pool for PackrlweOnline* recursive functions. +// Each recursion depth gets 6 dedicated scratch vectors: +// 0=tmp, 1=tmp1, 2=tmp2, 3=result, 4=b_e (left child result), 5=b_o (right +// child result) +// Slots 4/5 eliminate the ~1022 × 16 KB heap copies for b_e/b_o in +// PackrlweOnlineConstantRowsImpl. +struct PackrlwePool { + static constexpr size_t kMaxDepth = 16; + static constexpr size_t kSlots = 6; + size_t len = 0; + std::array, kSlots>, kMaxDepth> v; + + void Init(size_t length) { + if (len == length) return; + len = length; + for (auto& row : v) + for (auto& buf : row) buf.resize(length); + } + + std::vector& Get(size_t depth, size_t slot) { + return v[depth][slot]; + } +}; + +thread_local PackrlwePool g_packrlwe_pool; + +} // namespace + +Secret::Secret(uint64_t length, uint64_t cmod) : len_(length), mod_(cmod) { + std::mt19937_64 rng(static_cast( + std::chrono::high_resolution_clock::now().time_since_epoch().count())); + std::uniform_int_distribution dist(-1, 1); + data.reserve(length); + for (uint64_t i = 0; i < length; ++i) { + int tmp = dist(rng); + if (tmp == -1) { + data.push_back(cmod - 1); + } else { + data.push_back(static_cast(tmp)); + } + } + is_ntt_ = false; +} + +FheParams::FheParams(uint64_t rlwe_degree, uint64_t rlwe_ct_modulus, + uint64_t rlwe_pt_modulus, uint64_t lwe_dimension, + uint64_t lwe_ct_modulus, uint64_t lwe_pt_modulus, + double sigma, double sigma_ring, AutoParams auto_params, + DecompParams decomp_params) + : rlwe_degree_(rlwe_degree), + rlwe_ct_modulus_(rlwe_ct_modulus), + rlwe_pt_modulus_(rlwe_pt_modulus), + lwe_dimension_(lwe_dimension), + lwe_ct_modulus_(lwe_ct_modulus), + lwe_pt_modulus_(lwe_pt_modulus), + sigma_(sigma), + sigma_ring_(sigma_ring), + auto_params_(auto_params), + decomp_params_(decomp_params), + ntt_(rlwe_degree, rlwe_ct_modulus, kRootOfUnityCrt) { + YACL_ENFORCE(rlwe_degree_ > 0); + YACL_ENFORCE(lwe_dimension_ > 0); + YACL_ENFORCE(rlwe_ct_modulus_ > 0); + YACL_ENFORCE(lwe_ct_modulus_ > 0); + YACL_ENFORCE(rlwe_pt_modulus_ > 0); + YACL_ENFORCE(lwe_pt_modulus_ > 0); + YACL_ENFORCE(IsPowerOfTwo(rlwe_degree_), "rlwe_degree must be power of two"); + + rlwe_degree_log2_ = Log2Exact(rlwe_degree_); +} + +void FheParams::set_persudo_matrix_simplepir(uint64_t row) { + SetPseudorandomMatrixSimplepir(persudo_matrix_simplepir_, + persudo_matrix_simplepir_flat_, row, + lwe_dimension_, lwe_ct_modulus_, prg_); +} + +void FheParams::set_persudo_matrix_doublepir(uint64_t row) { + SetPseudorandomMatrixDoublepir(persudo_matrix_doublepir_, + persudo_matrix_doublepir_flat_, row, + rlwe_degree_, rlwe_ct_modulus_, prg_); +} + +void FheParams::set_persudo_hcube_ypir() { + uint64_t expo = rlwe_degree_log2_; + SetPseudorandomHypercubeYpir(persudo_hcube_ypir_, expo, auto_params_.t, + rlwe_degree_, rlwe_ct_modulus_, prg_); +} + +void FheParams::set_automap(std::vector& idx) { + SetAutomap(automap_, rlwe_degree_, idx); +} + +void FheParams::set_precomputed_pt(uint64_t max_lh) { + SetPrecomputedPt(precomputed_pt_, rlwe_degree_, max_lh, ntt_forward_); +} + +PirParams::PirParams(uint64_t rows, uint64_t cols) : rows_(rows), cols_(cols) { + YACL_ENFORCE(rows_ > 0); + YACL_ENFORCE(cols_ > 0); +} + +// ====================================================================== +// YPIR utilities that depend on FheParams. +// ====================================================================== + +void ApplyAutoNttForm(const std::vector& vec, + std::vector& result, uint64_t idx, + const FheParams& fparm) { + const uint64_t length = vec.size(); + idx = idx % (2 * length); + const auto& automap = fparm.get_automap(GetLog2(idx - 1) - 1); + + if (vec.data() == result.data()) { + static thread_local std::vector scratch; + scratch.resize(length); + std::copy(vec.begin(), vec.end(), scratch.begin()); + for (uint64_t i = 0; i < length; ++i) { + result[i] = scratch[automap[i]]; + } + return; + } + + if (result.size() != length) { + result.resize(length); + } + for (uint64_t i = 0; i < length; ++i) { + result[i] = vec[automap[i]]; + } +} + +void KeyswitchPreprocess( + const std::vector>& ksk_a, + std::vector& a_in, std::vector& a_out, + std::vector>>& decomp_buf, + const FheParams& fparm) { + const uint64_t z = fparm.get_z_auto(); + const uint64_t t = fparm.get_t_auto(); + const uint64_t b = fparm.get_b_auto(); + const uint64_t length = fparm.get_poly_degree(); + const uint64_t modulus = fparm.get_rlwe_cmod(); + + YpirHexlNtt& ntt = fparm.get_ntt(); + std::vector> decomp_a(t, + std::vector(length, 0)); + ntt.Inverse(a_in.data(), length); + for (uint64_t i = 0; i < length; ++i) { + uint64_t val = a_in[i] >> b; + for (uint64_t j = 0; j < t; ++j) { + decomp_a[j][i] = val & ((1ULL << z) - 1); + val >>= z; + } + } + for (uint64_t i = 0; i < t; ++i) { + ntt.Forward(decomp_a[i].data(), length); + } + + decomp_buf.push_back(std::move(decomp_a)); + if (a_out.size() != length) { + a_out.resize(length); + } + std::fill(a_out.begin(), a_out.end(), 0); + std::vector tmp(length, 0); + for (uint64_t i = 0; i < t; ++i) { + EltwiseMultMod(tmp.data(), ksk_a[i].data(), decomp_buf.back()[i].data(), + length, modulus); + EltwiseSubMod(a_out.data(), a_out.data(), tmp.data(), length, modulus); + } +} + +void KeyswitchOnline(const std::vector>& ksk_b, + std::vector& b_in, std::vector& b_out, + std::vector>& decomp_buf, + const FheParams& fparm) { + const uint64_t t = fparm.get_t_auto(); + const uint64_t length = fparm.get_poly_degree(); + const uint64_t modulus = fparm.get_rlwe_cmod(); + // Skip copy if b_in and b_out are already the same object (in-place call). + if (b_out.data() != b_in.data()) { + b_out = b_in; + } + // Reuse a thread-local buffer to avoid repeated heap allocation. + // EltwiseMultMod fully overwrites tmp, so no zero-init needed. + static thread_local std::vector tmp; + if (tmp.size() < length) tmp.resize(length); + for (uint64_t i = 0; i < t; ++i) { + EltwiseMultMod(tmp.data(), ksk_b[i].data(), decomp_buf[i].data(), length, + modulus); + EltwiseSubMod(b_out.data(), b_out.data(), tmp.data(), length, modulus); + } +} + +namespace { +void EvalAutoPreprocess( + std::vector& a_in, std::vector& a_out, + const std::vector>& ksk_a, uint64_t idx, + std::vector>>& decomp_buf, + const FheParams& fparm) { + ApplyAutoNttForm(a_in, a_in, idx, fparm); + KeyswitchPreprocess(ksk_a, a_in, a_out, decomp_buf, fparm); +} + +void EvalAutoOnline(std::vector& b_in, std::vector& b_out, + const std::vector>& ksk_b, + uint64_t idx, + std::vector>& decomp_buf, + const FheParams& fparm) { + // Permute b_in into b_out (out-of-place: no scratch copy needed). + // Then key-switch b_out in place, avoiding the extra copy in KeyswitchOnline. + ApplyAutoNttForm(b_in, b_out, idx, fparm); + KeyswitchOnline(ksk_b, b_out, b_out, decomp_buf, fparm); +} + +std::vector PackrlwePreprocessImpl( + const std::vector>& a, uint64_t start, + uint64_t stride, uint64_t l, uint64_t h, + const std::vector>>& ksk_a, + std::vector>>& decomp_buf, + const FheParams& fparm) { + if (l == 0) { + return a[start]; + } + + const uint64_t length = fparm.get_poly_degree(); + const uint64_t modulus = fparm.get_rlwe_cmod(); + const uint64_t expo = GetLog2(length); + + std::vector a_e = PackrlwePreprocessImpl( + a, start, stride << 1, l - 1, h, ksk_a, decomp_buf, fparm); + std::vector a_o = PackrlwePreprocessImpl( + a, start + stride, stride << 1, l - 1, h, ksk_a, decomp_buf, fparm); + + const std::vector& pt = fparm.get_precomputed_pt(l + h); + std::vector tmp(length, 0); + std::vector tmp1(length, 0); + std::vector tmp2(length, 0); + std::vector result(length, 0); + + EltwiseMultMod(tmp.data(), pt.data(), a_o.data(), length, modulus); + EltwiseAddMod(result.data(), a_e.data(), tmp.data(), length, modulus); + EltwiseSubMod(tmp1.data(), a_e.data(), tmp.data(), length, modulus); + + EvalAutoPreprocess(tmp1, tmp2, + ksk_a[GetAutokeyIdx((1ULL << (l + h)) + 1, expo)], + (1ULL << (l + h)) + 1, decomp_buf, fparm); + EltwiseAddMod(result.data(), result.data(), tmp2.data(), length, modulus); + return result; +} + +std::vector PackrlweOnlineImpl( + const std::vector>& b, uint64_t start, + uint64_t stride, uint64_t l, uint64_t h, + const std::vector>>& ksk_b, + std::vector>>& decomp_buf, uint64_t& ptr, + const FheParams& fparm, uint64_t depth = 0) { + if (l == 0) { + return b[start]; + } + + const uint64_t length = fparm.get_poly_degree(); + const uint64_t modulus = fparm.get_rlwe_cmod(); + const uint64_t expo = GetLog2(length); + + std::vector b_e = + PackrlweOnlineImpl(b, start, stride << 1, l - 1, h, ksk_b, decomp_buf, + ptr, fparm, depth + 1); + std::vector b_o = + PackrlweOnlineImpl(b, start + stride, stride << 1, l - 1, h, ksk_b, + decomp_buf, ptr, fparm, depth + 1); + + const std::vector& pt = fparm.get_precomputed_pt(l + h); + + g_packrlwe_pool.Init(length); + auto& tmp = g_packrlwe_pool.Get(depth, 0); + auto& tmp1 = g_packrlwe_pool.Get(depth, 1); + auto& tmp2 = g_packrlwe_pool.Get(depth, 2); + auto& result = g_packrlwe_pool.Get(depth, 3); + + EltwiseMultMod(tmp.data(), pt.data(), b_o.data(), length, modulus); + EltwiseAddMod(result.data(), b_e.data(), tmp.data(), length, modulus); + EltwiseSubMod(tmp1.data(), b_e.data(), tmp.data(), length, modulus); + + EvalAutoOnline(tmp1, tmp2, ksk_b[GetAutokeyIdx((1ULL << (l + h)) + 1, expo)], + (1ULL << (l + h)) + 1, decomp_buf[ptr++], fparm); + EltwiseAddMod(result.data(), result.data(), tmp2.data(), length, modulus); + return result; +} + +// Output-parameter variant: writes result directly into `output` (a pool slot +// owned by the parent frame), eliminating the ~16 KB copy-on-return that the +// previous return-by-value version produced at each of 1023 recursive calls. +void PackrlweOnlineConstantRowsImpl( + std::vector& output, const std::vector& b, + uint64_t start, uint64_t stride, uint64_t l, uint64_t h, + const std::vector>>& ksk_b, + std::vector>>& decomp_buf, uint64_t& ptr, + const FheParams& fparm, uint64_t depth = 0) { + const uint64_t length = fparm.get_poly_degree(); + g_packrlwe_pool.Init(length); + + if (l == 0) { + std::fill(output.begin(), output.end(), b[start]); + return; + } + + const uint64_t modulus = fparm.get_rlwe_cmod(); + const uint64_t expo = GetLog2(length); + + // Pool slots at this depth: + // 0=tmp, 1=tmp1, 2=tmp2 (scratch) + // 4=b_e (left child writes here), 5=b_o (right child writes here) + // `output` is provided by the parent (or the top-level wrapper). + auto& tmp = g_packrlwe_pool.Get(depth, 0); + auto& tmp1 = g_packrlwe_pool.Get(depth, 1); + auto& tmp2 = g_packrlwe_pool.Get(depth, 2); + + if (l == 1) { + const std::vector& pt = fparm.get_precomputed_pt(l + h); + EltwiseFMAMod(tmp.data(), pt.data(), b[start + stride], nullptr, length, + modulus); + const uint64_t scalar = b[start]; + for (uint64_t i = 0; i < length; ++i) { + output[i] = scalar + tmp[i]; + if (output[i] >= modulus) output[i] -= modulus; + tmp1[i] = scalar >= tmp[i] ? scalar - tmp[i] : scalar + modulus - tmp[i]; + } + EvalAutoOnline(tmp1, tmp2, + ksk_b[GetAutokeyIdx((1ULL << (l + h)) + 1, expo)], + (1ULL << (l + h)) + 1, decomp_buf[ptr++], fparm); + EltwiseAddMod(output.data(), output.data(), tmp2.data(), length, modulus); + return; + } + + // Recurse: children write directly into pool slots 4 and 5 at this depth, + // avoiding any heap allocation or copy for the intermediate results. + auto& b_e = g_packrlwe_pool.Get(depth, 4); + auto& b_o = g_packrlwe_pool.Get(depth, 5); + + PackrlweOnlineConstantRowsImpl(b_e, b, start, stride << 1, l - 1, h, ksk_b, + decomp_buf, ptr, fparm, depth + 1); + PackrlweOnlineConstantRowsImpl(b_o, b, start + stride, stride << 1, l - 1, h, + ksk_b, decomp_buf, ptr, fparm, depth + 1); + + const std::vector& pt = fparm.get_precomputed_pt(l + h); + + EltwiseMultMod(tmp.data(), pt.data(), b_o.data(), length, modulus); + EltwiseAddMod(output.data(), b_e.data(), tmp.data(), length, modulus); + EltwiseSubMod(tmp1.data(), b_e.data(), tmp.data(), length, modulus); + + EvalAutoOnline(tmp1, tmp2, ksk_b[GetAutokeyIdx((1ULL << (l + h)) + 1, expo)], + (1ULL << (l + h)) + 1, decomp_buf[ptr++], fparm); + EltwiseAddMod(output.data(), output.data(), tmp2.data(), length, modulus); +} +} // namespace + +std::vector PackrlwePreprocess( + std::vector>& a, uint64_t l, uint64_t h, + const std::vector>>& ksk_a, + std::vector>>& decomp_buf, + const FheParams& fparm) { + return PackrlwePreprocessImpl(a, 0, 1, l, h, ksk_a, decomp_buf, fparm); +} + +std::vector PackrlweOnline( + std::vector>& b, uint64_t l, uint64_t h, + const std::vector>>& ksk_b, + std::vector>>& decomp_buf, uint64_t& ptr, + const FheParams& fparm) { + return PackrlweOnlineImpl(b, 0, 1, l, h, ksk_b, decomp_buf, ptr, fparm); +} + +std::vector PackrlweOnlineConstantRows( + const std::vector& b, uint64_t l, uint64_t h, + const std::vector>>& ksk_b, + std::vector>>& decomp_buf, uint64_t& ptr, + const FheParams& fparm) { + g_packrlwe_pool.Init(fparm.get_poly_degree()); + // Use the depth-0 result slot as output; one final copy on return. + auto& result = g_packrlwe_pool.Get(0, 3); + PackrlweOnlineConstantRowsImpl(result, b, 0, 1, l, h, ksk_b, decomp_buf, ptr, + fparm, 0); + return result; +} + +void Cdks21Lwe2RlweInplace(uint64_t* lwe_a, uint64_t degree, uint64_t cmod, + YpirHexlNtt& ntt) { + uint64_t a0 = lwe_a[0]; + + for (uint64_t i = 1; i < (degree + 1) / 2; ++i) { + uint64_t left_val = lwe_a[i]; + uint64_t right_val = lwe_a[degree - i]; + lwe_a[i] = right_val ? cmod - right_val : 0; + lwe_a[degree - i] = left_val ? cmod - left_val : 0; + } + + if (degree % 2 == 0) { + uint64_t mid = degree / 2; + uint64_t val = lwe_a[mid]; + lwe_a[mid] = val ? cmod - val : 0; + } + + lwe_a[0] = a0; + ntt.Forward(lwe_a, degree); +} + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/ypir_params.h b/psi/algorithm/ypir/legacy/ypir_params.h new file mode 100644 index 00000000..15ec7a87 --- /dev/null +++ b/psi/algorithm/ypir/legacy/ypir_params.h @@ -0,0 +1,155 @@ +// Copyright 2026 The secretflow authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +#include "psi/algorithm/ypir/legacy/ypir_util.h" + +namespace psi::ypir::ypir_internal { + +class Secret { + public: + Secret() = default; + Secret(uint64_t length, uint64_t cmod); + + uint64_t get_data_i(uint64_t idx) const { return data[idx]; } + uint64_t get_len() const { return len_; } + uint64_t get_mod() const { return mod_; } + bool get_ntt_form() const { return is_ntt_; } + + void switch_ntt_format() { is_ntt_ = !is_ntt_; } + + std::vector data; + + private: + uint64_t len_ = 0; + uint64_t mod_ = 0; + bool is_ntt_ = false; +}; + +struct AutoParams { + uint64_t b = 0; + uint64_t z = 0; + uint64_t t = 0; +}; + +struct DecompParams { + uint64_t b = 0; + uint64_t z = 0; + uint64_t t = 0; +}; + +class FheParams { + public: + FheParams(uint64_t rlwe_degree, uint64_t rlwe_ct_modulus, + uint64_t rlwe_pt_modulus, uint64_t lwe_dimension, + uint64_t lwe_ct_modulus, uint64_t lwe_pt_modulus, double sigma, + double sigma_ring, AutoParams auto_params, + DecompParams decomp_params); + + // ypir-internal getters + uint64_t get_poly_degree() const { return rlwe_degree_; } + uint64_t get_rlwe_cmod() const { return rlwe_ct_modulus_; } + uint64_t get_rlwe_pmod() const { return rlwe_pt_modulus_; } + uint64_t get_lwe_dimension() const { return lwe_dimension_; } + uint64_t get_lwe_cmod() const { return lwe_ct_modulus_; } + uint64_t get_lwe_pmod() const { return lwe_pt_modulus_; } + uint64_t get_b_auto() const { return auto_params_.b; } + uint64_t get_z_auto() const { return auto_params_.z; } + uint64_t get_t_auto() const { return auto_params_.t; } + uint64_t get_b_decomp() const { return decomp_params_.b; } + uint64_t get_z_decomp() const { return decomp_params_.z; } + uint64_t get_t_decomp() const { return decomp_params_.t; } + double get_sig() const { return sigma_; } + double get_sig_ring() const { return sigma_ring_; } + YpirHexlNtt& get_ntt() const noexcept { return ntt_; } + + const std::vector>& get_persudo_matrix_simplepir() + const noexcept { + return persudo_matrix_simplepir_; + } + const std::vector>& get_persudo_matrix_doublepir() + const noexcept { + return persudo_matrix_doublepir_; + } + const std::vector>>& + get_persudo_hcube_ypir() const noexcept { + return persudo_hcube_ypir_; + } + const uint64_t* get_persudo_matrix_simplepir_flat() const noexcept { + return persudo_matrix_simplepir_flat_.data(); + } + const uint64_t* get_persudo_matrix_doublepir_flat() const noexcept { + return persudo_matrix_doublepir_flat_.data(); + } + + void set_persudo_matrix_simplepir(uint64_t row); + void set_persudo_matrix_doublepir(uint64_t row); + void set_persudo_hcube_ypir(); + void set_automap(std::vector& idx); + const std::vector& get_automap(uint64_t idx) const noexcept { + return automap_[idx]; + } + void set_precomputed_pt(uint64_t max_lh); + const std::vector& get_precomputed_pt(uint64_t lh) const noexcept { + return precomputed_pt_[lh]; + } + + void SetNttForward(NttForwardFn fn) { ntt_forward_ = fn; } + + private: + uint64_t rlwe_degree_ = 0; + uint64_t rlwe_ct_modulus_ = 0; + uint64_t rlwe_pt_modulus_ = 0; + + uint64_t lwe_dimension_ = 0; + uint64_t lwe_ct_modulus_ = 0; + uint64_t lwe_pt_modulus_ = 0; + + double sigma_ = 0.0; + double sigma_ring_ = 0.0; + + AutoParams auto_params_; + DecompParams decomp_params_; + + uint64_t rlwe_degree_log2_ = 0; + mutable YpirHexlNtt ntt_; + + std::vector> persudo_matrix_simplepir_; + std::vector> persudo_matrix_doublepir_; + std::vector>> persudo_hcube_ypir_; + std::vector persudo_matrix_simplepir_flat_; + std::vector persudo_matrix_doublepir_flat_; + std::vector> automap_; + std::vector> precomputed_pt_; + AESCTR_PRNG prg_; + NttForwardFn ntt_forward_ = nullptr; +}; + +class PirParams { + public: + PirParams(uint64_t rows, uint64_t cols); + + uint64_t get_row() const { return rows_; } + uint64_t get_col() const { return cols_; } + + private: + uint64_t rows_ = 0; + uint64_t cols_ = 0; +}; + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/ypir_util.cc b/psi/algorithm/ypir/legacy/ypir_util.cc new file mode 100644 index 00000000..0f98fcf2 --- /dev/null +++ b/psi/algorithm/ypir/legacy/ypir_util.cc @@ -0,0 +1,521 @@ +// Copyright 2026 The secretflow authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "psi/algorithm/ypir/legacy/ypir_util.h" + +#include +#include +#include + +#include "yacl/base/exception.h" + +#if defined(__x86_64__) || defined(_M_X64) +#include +#endif + +namespace psi::ypir::ypir_internal { +namespace { + +int64_t ExGcd(int64_t a, int64_t b, int64_t& x, int64_t& y) { + if (b == 0) { + x = 1; + y = 0; + return a; + } + int64_t x1 = 0; + int64_t y1 = 0; + int64_t d = ExGcd(b, a % b, x1, y1); + x = y1; + y = x1 - (a / b) * y1; + return d; +} + +unsigned char ReverseByte(unsigned char x) { + static const unsigned char table[] = { + 0x00, 0x80, 0x40, 0xc0, 0x20, 0xa0, 0x60, 0xe0, 0x10, 0x90, 0x50, 0xd0, + 0x30, 0xb0, 0x70, 0xf0, 0x08, 0x88, 0x48, 0xc8, 0x28, 0xa8, 0x68, 0xe8, + 0x18, 0x98, 0x58, 0xd8, 0x38, 0xb8, 0x78, 0xf8, 0x04, 0x84, 0x44, 0xc4, + 0x24, 0xa4, 0x64, 0xe4, 0x14, 0x94, 0x54, 0xd4, 0x34, 0xb4, 0x74, 0xf4, + 0x0c, 0x8c, 0x4c, 0xcc, 0x2c, 0xac, 0x6c, 0xec, 0x1c, 0x9c, 0x5c, 0xdc, + 0x3c, 0xbc, 0x7c, 0xfc, 0x02, 0x82, 0x42, 0xc2, 0x22, 0xa2, 0x62, 0xe2, + 0x12, 0x92, 0x52, 0xd2, 0x32, 0xb2, 0x72, 0xf2, 0x0a, 0x8a, 0x4a, 0xca, + 0x2a, 0xaa, 0x6a, 0xea, 0x1a, 0x9a, 0x5a, 0xda, 0x3a, 0xba, 0x7a, 0xfa, + 0x06, 0x86, 0x46, 0xc6, 0x26, 0xa6, 0x66, 0xe6, 0x16, 0x96, 0x56, 0xd6, + 0x36, 0xb6, 0x76, 0xf6, 0x0e, 0x8e, 0x4e, 0xce, 0x2e, 0xae, 0x6e, 0xee, + 0x1e, 0x9e, 0x5e, 0xde, 0x3e, 0xbe, 0x7e, 0xfe, 0x01, 0x81, 0x41, 0xc1, + 0x21, 0xa1, 0x61, 0xe1, 0x11, 0x91, 0x51, 0xd1, 0x31, 0xb1, 0x71, 0xf1, + 0x09, 0x89, 0x49, 0xc9, 0x29, 0xa9, 0x69, 0xe9, 0x19, 0x99, 0x59, 0xd9, + 0x39, 0xb9, 0x79, 0xf9, 0x05, 0x85, 0x45, 0xc5, 0x25, 0xa5, 0x65, 0xe5, + 0x15, 0x95, 0x55, 0xd5, 0x35, 0xb5, 0x75, 0xf5, 0x0d, 0x8d, 0x4d, 0xcd, + 0x2d, 0xad, 0x6d, 0xed, 0x1d, 0x9d, 0x5d, 0xdd, 0x3d, 0xbd, 0x7d, 0xfd, + 0x03, 0x83, 0x43, 0xc3, 0x23, 0xa3, 0x63, 0xe3, 0x13, 0x93, 0x53, 0xd3, + 0x33, 0xb3, 0x73, 0xf3, 0x0b, 0x8b, 0x4b, 0xcb, 0x2b, 0xab, 0x6b, 0xeb, + 0x1b, 0x9b, 0x5b, 0xdb, 0x3b, 0xbb, 0x7b, 0xfb, 0x07, 0x87, 0x47, 0xc7, + 0x27, 0xa7, 0x67, 0xe7, 0x17, 0x97, 0x57, 0xd7, 0x37, 0xb7, 0x77, 0xf7, + 0x0f, 0x8f, 0x4f, 0xcf, 0x2f, 0xaf, 0x6f, 0xef, 0x1f, 0x9f, 0x5f, 0xdf, + 0x3f, 0xbf, 0x7f, 0xff, + }; + return table[x]; +} + +uint32_t ReverseBits(uint32_t num, uint32_t msb) { + static const int shift_trick[] = {0, 7, 6, 5, 4, 3, 2, 1}; + uint32_t msbb = (msb >> 3) + ((msb & 0x7) ? 1 : 0); + switch (msbb) { + case 1: + return (ReverseByte(num & 0xff) >> shift_trick[msb & 0x7]); + case 2: + return (ReverseByte(num & 0xff) << 8 | ReverseByte((num >> 8) & 0xff)) >> + shift_trick[msb & 0x7]; + case 3: + return (ReverseByte(num & 0xff) << 16 | + ReverseByte((num >> 8) & 0xff) << 8 | + ReverseByte((num >> 16) & 0xff)) >> + shift_trick[msb & 0x7]; + case 4: + return (ReverseByte(num & 0xff) << 24 | + ReverseByte((num >> 8) & 0xff) << 16 | + ReverseByte((num >> 16) & 0xff) << 8 | + ReverseByte((num >> 24) & 0xff)) >> + shift_trick[msb & 0x7]; + default: + YACL_ENFORCE(false, "ReverseBits msb out of range"); + } +} + +#if defined(__x86_64__) || defined(_M_X64) +#if defined(__GNUC__) || defined(__clang__) +__attribute__((target("avx512f,avx512bw,avx512vl"))) void +MatVecU8U32Mod2p32Avx512(const uint8_t* A, const uint32_t* x, uint32_t* y, + size_t rows, size_t cols) { + constexpr size_t kLanes = 16; + + for (size_t r = 0; r < rows; ++r) { + const uint8_t* row = A + r * cols; + __m512i acc0 = _mm512_setzero_si512(); + __m512i acc1 = _mm512_setzero_si512(); + + size_t c = 0; + for (; c + 2 * kLanes <= cols; c += 2 * kLanes) { + const __m128i bytes0 = + _mm_loadu_si128(reinterpret_cast(row + c)); + const __m128i bytes1 = + _mm_loadu_si128(reinterpret_cast(row + c + kLanes)); + const __m512i a0 = _mm512_cvtepu8_epi32(bytes0); + const __m512i a1 = _mm512_cvtepu8_epi32(bytes1); + const __m512i x0 = + _mm512_loadu_si512(reinterpret_cast(x + c)); + const __m512i x1 = + _mm512_loadu_si512(reinterpret_cast(x + c + kLanes)); + acc0 = _mm512_add_epi32(acc0, _mm512_mullo_epi32(a0, x0)); + acc1 = _mm512_add_epi32(acc1, _mm512_mullo_epi32(a1, x1)); + } + + __m512i acc = _mm512_add_epi32(acc0, acc1); + for (; c + kLanes <= cols; c += kLanes) { + const __m128i bytes = + _mm_loadu_si128(reinterpret_cast(row + c)); + const __m512i a = _mm512_cvtepu8_epi32(bytes); + const __m512i xv = + _mm512_loadu_si512(reinterpret_cast(x + c)); + acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(a, xv)); + } + + alignas(64) uint32_t lanes[kLanes]; + _mm512_storeu_si512(reinterpret_cast(lanes), acc); + uint32_t sum = 0; + for (size_t i = 0; i < kLanes; ++i) { + sum += lanes[i]; + } + for (; c < cols; ++c) { + sum += static_cast(row[c]) * x[c]; + } + y[r] = sum; + } +} + +__attribute__((target("avx512f,avx512bw,avx512vl"))) void +MatrixVectorMultiplicationU16Avx512( + std::vector& result, + const std::vector>& matrix, + const std::vector& vec, uint64_t mod) { + const uint64_t mat_row = matrix.size(); + const uint64_t mat_col = mat_row == 0 ? 0 : matrix[0].size(); + result.assign(mat_row, 0); + if (mat_row == 0 || mat_col == 0) { + return; + } + + std::vector row64(mat_col, 0); + std::vector tmp(mat_col, 0); + constexpr uint64_t kBlock = 128; + + for (uint64_t i = 0; i < mat_row; ++i) { + for (uint64_t j = 0; j < mat_col; ++j) { + row64[j] = static_cast(matrix[i][j]); + } + + EltwiseMultMod(tmp.data(), row64.data(), vec.data(), mat_col, mod); + + uint64_t sum_mod = 0; + for (uint64_t base = 0; base < mat_col; base += kBlock) { + const uint64_t end = std::min(base + kBlock, mat_col); + __m512i acc0 = _mm512_setzero_si512(); + __m512i acc1 = _mm512_setzero_si512(); + uint64_t j = base; + for (; j + 16 <= end; j += 16) { + const __m512i v0 = _mm512_loadu_si512( + reinterpret_cast(tmp.data() + j)); + const __m512i v1 = _mm512_loadu_si512( + reinterpret_cast(tmp.data() + j + 8)); + acc0 = _mm512_add_epi64(acc0, v0); + acc1 = _mm512_add_epi64(acc1, v1); + } + + const __m512i acc = _mm512_add_epi64(acc0, acc1); + alignas(64) uint64_t buf[8]; + _mm512_store_si512(reinterpret_cast<__m512i*>(buf), acc); + + unsigned __int128 acc128 = 0; + for (uint64_t k = 0; k < 8; ++k) { + acc128 += buf[k]; + } + for (; j < end; ++j) { + acc128 += tmp[j]; + } + sum_mod = (sum_mod + static_cast(acc128 % mod)) % mod; + } + result[i] = sum_mod; + } +} +#endif +#endif + +void MatVecU8U32Mod2p32Scalar(const uint8_t* A, const uint32_t* x, uint32_t* y, + size_t rows, size_t cols) { + for (size_t r = 0; r < rows; ++r) { + uint64_t acc = 0; + const uint8_t* row = A + r * cols; + for (size_t c = 0; c < cols; ++c) { + acc += static_cast(row[c]) * x[c]; + } + y[r] = static_cast(acc); + } +} + +} // namespace + +uint64_t GetLog2(uint64_t num) { + uint64_t result = 0; + while (num > 1) { + num >>= 1; + ++result; + } + return result; +} + +bool IsPowerOfTwo(uint64_t num) { return num != 0 && (num & (num - 1)) == 0; } + +uint64_t SampleGauss(double st_dev, uint64_t modulus, std::mt19937_64& rng) { + std::normal_distribution gaussian_sampler(0.0, st_dev); + int64_t tmp = static_cast(std::llround(gaussian_sampler(rng))); + if (tmp < 0) { + return static_cast(modulus + tmp); + } + return static_cast(tmp); +} + +void SampleGauss(std::vector& err, double st_dev, uint64_t modulus, + std::mt19937_64& rng) { + std::normal_distribution gaussian_sampler(0.0, st_dev); + for (auto& v : err) { + int64_t tmp = static_cast(std::llround(gaussian_sampler(rng))); + if (tmp < 0) { + v = static_cast(modulus + tmp); + } else { + v = static_cast(tmp); + } + } +} + +int64_t ModInverse(int64_t a, int64_t modulus) { + int64_t x = 0; + int64_t y = 0; + int64_t g = ExGcd(a, modulus, x, y); + if (g != 1) { + return -1; + } + int64_t res = (x % modulus + modulus) % modulus; + return res; +} + +void ApplyAutoCoefForm(std::vector& result, + const std::vector& input, int32_t index, + uint64_t modulus) { + const uint64_t length = input.size(); + result.assign(length, 0); + for (size_t i = 0; i < length; ++i) { + uint64_t destination = (static_cast(i) * index) % (2 * length); + if (destination >= length) { + result[destination - length] = (modulus - input[i]) % modulus; + } else { + result[destination] = input[i]; + } + } +} + +void MatrixTranspose(const std::vector>& matrix, + std::vector>& trans_matrix) { + const uint64_t rows = matrix.size(); + const uint64_t cols = rows == 0 ? 0 : matrix[0].size(); + trans_matrix.assign(cols, std::vector(rows, 0)); + constexpr uint64_t kBlock = 64; + for (uint64_t ii = 0; ii < cols; ii += kBlock) { + for (uint64_t jj = 0; jj < rows; jj += kBlock) { + uint64_t i_end = std::min(ii + kBlock, cols); + uint64_t j_end = std::min(jj + kBlock, rows); + for (uint64_t i = ii; i < i_end; ++i) { + for (uint64_t j = jj; j < j_end; ++j) { + trans_matrix[i][j] = matrix[j][i]; + } + } + } + } +} + +void BitDecomp(uint64_t num, std::vector& vec, uint64_t b, uint64_t z, + uint64_t t) { + num >>= b; + vec.assign(t, 0); + const uint64_t mask = (1ULL << z) - 1; + for (uint64_t i = 0; i < t; ++i) { + vec[i] = num & mask; + num >>= z; + } +} + +void PrecomputeAutomap(uint32_t length, uint32_t idx, + std::vector& automap) { + const uint32_t m = length << 1; + const uint32_t logm = static_cast(std::llround(std::log2(m))); + const uint32_t logn = static_cast(std::llround(std::log2(length))); + automap.assign(length, 0); + for (uint32_t j = 0; j < length; ++j) { + uint32_t j_tmp = ((j << 1) + 1); + uint32_t index = ((j_tmp * idx) - (((j_tmp * idx) >> logm) << logm)) >> 1; + uint32_t j_rev = ReverseBits(j, logn); + uint32_t idx_rev = ReverseBits(index, logn); + automap[j_rev] = idx_rev; + } +} + +void PseudorandomMatrixGenerate(std::vector>& matrix, + uint64_t modulus, AESCTR_PRNG& prng) { + const uint64_t rows = matrix.size(); + const uint64_t cols = rows == 0 ? 0 : matrix[0].size(); + for (uint64_t i = 0; i < rows; ++i) { + uint8_t* ptr = reinterpret_cast(matrix[i].data()); + prng.fill_bytes(ptr, sizeof(uint64_t) * cols); + } + for (uint64_t i = 0; i < rows; ++i) + for (uint64_t j = 0; j < cols; ++j) matrix[i][j] %= modulus; +} + +void SetPseudorandomMatrixSimplepir(std::vector>& matrix, + std::vector& matrix_flat, + uint64_t rows, uint64_t cols, + uint64_t modulus, AESCTR_PRNG& prng) { + matrix.assign(rows, std::vector(cols, 0)); + prng.refresh(kFirstDimensionSeed); + PseudorandomMatrixGenerate(matrix, modulus, prng); + matrix_flat.resize(rows * cols); + for (uint64_t i = 0; i < rows; ++i) { + std::memcpy(&matrix_flat[i * cols], matrix[i].data(), + cols * sizeof(uint64_t)); + } +} + +void SetPseudorandomMatrixDoublepir(std::vector>& matrix, + std::vector& matrix_flat, + uint64_t rows, uint64_t cols, + uint64_t modulus, AESCTR_PRNG& prng) { + matrix.assign(rows, std::vector(cols, 0)); + prng.refresh(kSecondDimensionSeed); + PseudorandomMatrixGenerate(matrix, modulus, prng); + matrix_flat.resize(rows * cols); + for (uint64_t i = 0; i < rows; ++i) { + std::memcpy(&matrix_flat[i * cols], matrix[i].data(), + cols * sizeof(uint64_t)); + } +} + +void SetPseudorandomHypercubeYpir( + std::vector>>& cube, uint64_t layers, + uint64_t rows, uint64_t cols, uint64_t modulus, AESCTR_PRNG& prng) { + cube.assign(layers, std::vector>( + rows, std::vector(cols))); + prng.refresh(kThirdDimensionSeed); + for (uint64_t i = 0; i < layers; ++i) { + PseudorandomMatrixGenerate(cube[i], modulus, prng); + } +} + +void SetAutomap(std::vector>& automap, uint64_t degree, + const std::vector& idx) { + automap.assign(idx.size(), std::vector(degree, 0)); + for (size_t i = 0; i < idx.size(); ++i) { + PrecomputeAutomap(static_cast(degree), + static_cast(idx[i]), automap[i]); + } +} + +void SetPrecomputedPt(std::vector>& precomputed_pt, + uint64_t poly_degree, uint64_t max_lh, + NttForwardFn ntt_forward) { + precomputed_pt.assign(max_lh + 1, std::vector(poly_degree, 0)); + for (uint64_t lh = 2; lh <= max_lh; ++lh) { + precomputed_pt[lh][poly_degree / (1ULL << lh)] = 1; + if (ntt_forward != nullptr) { + ntt_forward(precomputed_pt[lh].data(), precomputed_pt[lh].size()); + } + } +} + +uint64_t GetAutokeyIdx(uint64_t idx, uint64_t expo) { + return expo - GetLog2(idx - 1); +} + +void VectorColDecompose(const std::vector& vec, + std::vector>& decomp_matrix, + uint64_t b, uint64_t z, uint64_t t) { + const uint64_t row = vec.size(); + const uint64_t mask = (1ULL << z) - 1; + decomp_matrix.assign(t, std::vector(row, 0)); + for (uint64_t i = 0; i < t; ++i) { + for (uint64_t j = 0; j < row; ++j) { + decomp_matrix[i][j] = + static_cast((vec[j] >> (b + i * z)) & mask); + } + } +} + +void MatrixRowDecompose(const std::vector>& matrix, + std::vector>& decomp_matrix, + uint64_t b, uint64_t z, uint64_t t) { + const uint64_t row = matrix.size(); + const uint64_t col = row == 0 ? 0 : matrix[0].size(); + decomp_matrix.assign(row, std::vector(col * t, 0)); + for (uint64_t i = 0; i < row; ++i) { + std::vector tmp(t, 0); + uint64_t ptr = 0; + for (uint64_t j = 0; j < col; ++j) { + BitDecomp(matrix[i][j], tmp, b, z, t); + for (uint64_t k = 0; k < t; ++k, ++ptr) { + decomp_matrix[i][ptr] = tmp[k]; + } + } + } +} + +void MatrixVectorFirstDimension( + std::vector& result, + const std::vector>& matrix, + const std::vector& vec, uint64_t mod) { + const uint64_t rows = matrix.size(); + const uint64_t cols = rows == 0 ? 0 : matrix[0].size(); + result.assign(cols, 0); + for (uint64_t r = 0; r < rows; ++r) { + EltwiseFMAMod(result.data(), matrix[r].data(), vec[r], result.data(), cols, + mod); + } +} + +void MatrixMultiplicationFlat(std::vector>& db_mul_matrix, + const std::vector>& db, + const uint64_t* matrix_flat, uint64_t matrix_col, + uint64_t mod) { + const uint64_t db_row = db.size(); + const uint64_t db_col = db_row == 0 ? 0 : db[0].size(); + db_mul_matrix.assign(db_row, std::vector(matrix_col, 0)); + for (uint64_t i = 0; i < db_row; ++i) { + uint64_t* res = db_mul_matrix[i].data(); + for (uint64_t k = 0; k < db_col; ++k) + EltwiseFMAMod(res, matrix_flat + k * matrix_col, db[i][k], res, + matrix_col, mod); + } +} + +void MatrixMultiplicationFlatU16( + std::vector>& db_mul_matrix, + const std::vector>& db, const uint64_t* matrix_flat, + uint64_t matrix_col, uint64_t mod) { + const uint64_t db_row = db.size(); + const uint64_t db_col = db_row == 0 ? 0 : db[0].size(); + db_mul_matrix.assign(db_row, std::vector(matrix_col, 0)); + for (uint64_t i = 0; i < db_row; ++i) { + uint64_t* res = db_mul_matrix[i].data(); + const uint16_t* row = db[i].data(); + for (uint64_t k = 0; k < db_col; ++k) { + const uint64_t s = row[k]; + if (s == 0) continue; + EltwiseFMAMod(res, matrix_flat + k * matrix_col, s, res, matrix_col, mod); + } + } +} + +void MatrixVectorMultiplicationU16( + std::vector& result, + const std::vector>& matrix, + const std::vector& vec, uint64_t mod) { +#if defined(__x86_64__) || defined(_M_X64) +#if defined(__GNUC__) || defined(__clang__) + if (__builtin_cpu_supports("avx512f") && + __builtin_cpu_supports("avx512bw") && + __builtin_cpu_supports("avx512vl")) { + MatrixVectorMultiplicationU16Avx512(result, matrix, vec, mod); + return; + } +#endif +#endif + const uint64_t mat_row = matrix.size(); + const uint64_t mat_col = mat_row == 0 ? 0 : matrix[0].size(); + result.assign(mat_row, 0); + for (uint64_t i = 0; i < mat_row; ++i) { + unsigned __int128 sum = 0; + for (uint64_t j = 0; j < mat_col; ++j) { + sum += static_cast(matrix[i][j]) * vec[j]; + if ((j & 127U) == 127U) { + sum %= mod; + } + } + result[i] = static_cast(sum % mod); + } +} + +void MatVecU8U32Mod2p32(const uint8_t* A, const uint32_t* x, uint32_t* y, + size_t rows, size_t cols) { +#if defined(__x86_64__) || defined(_M_X64) +#if defined(__GNUC__) || defined(__clang__) + if (__builtin_cpu_supports("avx512f") && + __builtin_cpu_supports("avx512bw") && + __builtin_cpu_supports("avx512vl")) { + MatVecU8U32Mod2p32Avx512(A, x, y, rows, cols); + return; + } +#endif +#endif + MatVecU8U32Mod2p32Scalar(A, x, y, rows, cols); +} + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/legacy/ypir_util.h b/psi/algorithm/ypir/legacy/ypir_util.h new file mode 100644 index 00000000..679afa44 --- /dev/null +++ b/psi/algorithm/ypir/legacy/ypir_util.h @@ -0,0 +1,134 @@ +// Copyright 2026 The secretflow authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include + +#include "psi/algorithm/ypir/legacy/aes_prng.h" +#include "psi/algorithm/ypir/legacy/hexl.h" + +namespace psi::ypir::ypir_internal { + +class FheParams; + +constexpr uint64_t kFirstDimensionSeed = 1; +constexpr uint64_t kSecondDimensionSeed = 1ULL << 16; +constexpr uint64_t kThirdDimensionSeed = 1ULL << 16; + +uint64_t GetLog2(uint64_t num); +bool IsPowerOfTwo(uint64_t num); +uint64_t GetAutokeyIdx(uint64_t idx, uint64_t expo); + +uint64_t SampleGauss(double st_dev, uint64_t modulus, std::mt19937_64& rng); +void SampleGauss(std::vector& err, double st_dev, uint64_t modulus, + std::mt19937_64& rng); + +int64_t ModInverse(int64_t a, int64_t modulus); + +void ApplyAutoCoefForm(std::vector& result, + const std::vector& input, int32_t index, + uint64_t modulus); + +void MatrixTranspose(const std::vector>& matrix, + std::vector>& trans_matrix); + +void PrecomputeAutomap(uint32_t length, uint32_t idx, + std::vector& automap); + +void PseudorandomMatrixGenerate(std::vector>& matrix, + uint64_t modulus, AESCTR_PRNG& prng); + +// Helpers for ypir-internal precomputation buffers. +void SetPseudorandomMatrixSimplepir(std::vector>& matrix, + std::vector& matrix_flat, + uint64_t rows, uint64_t cols, + uint64_t modulus, AESCTR_PRNG& prng); + +void SetPseudorandomMatrixDoublepir(std::vector>& matrix, + std::vector& matrix_flat, + uint64_t rows, uint64_t cols, + uint64_t modulus, AESCTR_PRNG& prng); + +void SetPseudorandomHypercubeYpir( + std::vector>>& cube, uint64_t layers, + uint64_t rows, uint64_t cols, uint64_t modulus, AESCTR_PRNG& prng); + +void SetAutomap(std::vector>& automap, uint64_t degree, + const std::vector& idx); + +using NttForwardFn = void (*)(uint64_t* data, size_t len); + +// If ntt_forward is nullptr, precomputed polynomials stay in coefficient form. +void SetPrecomputedPt(std::vector>& precomputed_pt, + uint64_t poly_degree, uint64_t max_lh, + NttForwardFn ntt_forward); + +// ypir server-side helpers +void VectorColDecompose(const std::vector& vec, + std::vector>& decomp_matrix, + uint64_t b, uint64_t z, uint64_t t); + +void MatrixRowDecompose(const std::vector>& matrix, + std::vector>& decomp_matrix, + uint64_t b, uint64_t z, uint64_t t); + +void MatrixVectorFirstDimension( + std::vector& result, + const std::vector>& matrix, + const std::vector& vec, uint64_t mod); + +void MatrixMultiplicationFlat(std::vector>& db_mul_matrix, + const std::vector>& db, + const uint64_t* matrix_flat, uint64_t matrix_col, + uint64_t mod); + +void MatrixMultiplicationFlatU16( + std::vector>& db_mul_matrix, + const std::vector>& db, const uint64_t* matrix_flat, + uint64_t matrix_col, uint64_t mod); + +void MatrixVectorMultiplicationU16( + std::vector& result, + const std::vector>& matrix, + const std::vector& vec, uint64_t mod); + +void MatVecU8U32Mod2p32(const uint8_t* A, const uint32_t* x, uint32_t* y, + size_t rows, size_t cols); + +std::vector PackrlwePreprocess( + std::vector>& a, uint64_t l, uint64_t h, + const std::vector>>& ksk_a, + std::vector>>& decomp_buf, + const FheParams& fparm); + +std::vector PackrlweOnline( + std::vector>& b, uint64_t l, uint64_t h, + const std::vector>>& ksk_b, + std::vector>>& decomp_buf, uint64_t& ptr, + const FheParams& fparm); + +std::vector PackrlweOnlineConstantRows( + const std::vector& b, uint64_t l, uint64_t h, + const std::vector>>& ksk_b, + std::vector>>& decomp_buf, uint64_t& ptr, + const FheParams& fparm); + +void Cdks21Lwe2RlweInplace(uint64_t* lwe_a, uint64_t degree, uint64_t cmod, + YpirHexlNtt& ntt); + +} // namespace psi::ypir::ypir_internal diff --git a/psi/algorithm/ypir/params.cc b/psi/algorithm/ypir/params.cc new file mode 100644 index 00000000..a37bfe39 --- /dev/null +++ b/psi/algorithm/ypir/params.cc @@ -0,0 +1,132 @@ +#include "psi/algorithm/ypir/params.h" + +#include +#include +#include +#include + +#include "yacl/base/exception.h" + +namespace psi::ypir { +namespace { + +size_t Log2Exact(uint64_t v) { + YACL_ENFORCE_GT(v, 0U); + YACL_ENFORCE_EQ(v & (v - 1), 0U, "value must be a power of two"); + size_t out = 0; + while (v > 1) { + v >>= 1; + ++out; + } + return out; +} + +struct MatrixShape { + uint64_t rows = 0; + uint64_t cols = 0; +}; + +MatrixShape ChooseShape(uint64_t num_items, uint64_t base) { + YACL_ENFORCE_GT(base, 0U); + uint64_t rows = base; + uint64_t cols = base; + while (rows * cols < num_items) { + if (rows <= cols) { + rows <<= 1; + } else { + cols <<= 1; + } + } + return {rows, cols}; +} + +YpirParameters BuildSimplepirParamsForShape(uint64_t db_rows, uint64_t db_cols, + uint64_t item_size_bits) { + const uint64_t poly_len = 1024; + const uint64_t pt_modulus = item_size_bits <= 8 ? (1ULL << 8) : (1ULL << 14); + YACL_ENFORCE_GE(db_rows, poly_len); + YACL_ENFORCE_EQ(db_cols % poly_len, 0U, + "simplepir db_cols must be a multiple of {}", poly_len); + const size_t db_dim_1 = Log2Exact(db_rows) - Log2Exact(poly_len); + const size_t instances = db_cols / poly_len; + + std::vector moduli = {268369921ULL, 249561089ULL}; + double noise_width = 16.042421; + psi::spiral::PolyMatrixParams poly_matrix_params(2, pt_modulus, 21, 4, 8, 8, + instances); + psi::spiral::QueryParams query_params(db_dim_1, 0, instances); + psi::spiral::Params spiral_params(poly_len, std::move(moduli), noise_width, + std::move(poly_matrix_params), + std::move(query_params)); + + YpirParameters out; + out.mode = YpirMode::kSimplepir; + out.spiral_params = std::move(spiral_params); + out.db_rows = + 1ULL << (out.spiral_params.DbDim1() + out.spiral_params.PolyLenLog2()); + out.db_cols = out.spiral_params.Instances() * out.spiral_params.PolyLen(); + out.value_bytes = std::max(1, (item_size_bits + 7) / 8); + return out; +} + +YpirParameters BuildDoublepirParamsForShape(uint64_t db_rows, uint64_t db_cols, + uint64_t item_size_bits) { + const uint64_t poly_len = 1024; + YACL_ENFORCE_GE(db_rows, poly_len); + YACL_ENFORCE_GE(db_cols, poly_len); + const size_t db_dim_1 = Log2Exact(db_rows) - Log2Exact(poly_len); + const size_t db_dim_2 = Log2Exact(db_cols) - Log2Exact(poly_len); + + std::vector moduli = {268369921ULL, 249561089ULL}; + double noise_width = 6.4; + psi::spiral::PolyMatrixParams poly_matrix_params(2, 1ULL << 15, 21, 4, 8, 8, + 1); + psi::spiral::QueryParams query_params(db_dim_1, db_dim_2, 1); + psi::spiral::Params spiral_params(poly_len, std::move(moduli), noise_width, + std::move(poly_matrix_params), + std::move(query_params)); + + YpirParameters out; + out.mode = YpirMode::kDoublepir; + out.spiral_params = std::move(spiral_params); + out.db_rows = + 1ULL << (out.spiral_params.DbDim1() + out.spiral_params.PolyLenLog2()); + out.db_cols = + 1ULL << (out.spiral_params.DbDim2() + out.spiral_params.PolyLenLog2()); + out.value_bytes = std::max(1, (item_size_bits + 7) / 8); + return out; +} + +} // namespace + +YpirParameters CreateParamsForScenarioSimplePIR(uint64_t num_items, + uint64_t item_size_bits) { + const auto shape = ChooseShape(num_items, 1024); + return BuildSimplepirParamsForShape(shape.rows, shape.cols, item_size_bits); +} + +YpirParameters CreateParamsForScenarioDoublePIR(uint64_t num_items, + uint64_t item_size_bits) { + const auto shape = ChooseShape(num_items, 1024); + return BuildDoublepirParamsForShape(shape.rows, shape.cols, item_size_bits); +} + +YpirParameters CreateParamsForShapeSimplePIR(uint64_t db_rows, uint64_t db_cols, + uint64_t item_size_bits) { + return BuildSimplepirParamsForShape(db_rows, db_cols, item_size_bits); +} + +YpirParameters CreateParamsForShapeDoublePIR(uint64_t db_rows, uint64_t db_cols, + uint64_t item_size_bits) { + return BuildDoublepirParamsForShape(db_rows, db_cols, item_size_bits); +} + +YpirParameters CreateSmallTestParamsSimplePIR() { + return BuildSimplepirParamsForShape(1ULL << 10, 1ULL << 10, 16); +} + +YpirParameters CreateSmallTestParamsDoublePIR() { + return BuildDoublepirParamsForShape(1ULL << 10, 1ULL << 10, 8); +} + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/params.h b/psi/algorithm/ypir/params.h new file mode 100644 index 00000000..c7d74324 --- /dev/null +++ b/psi/algorithm/ypir/params.h @@ -0,0 +1,32 @@ +#pragma once + +#include +#include + +#include "psi/algorithm/spiral/params.h" +#include "psi/algorithm/ypir/types.h" + +namespace psi::ypir { + +struct YpirParameters { + YpirMode mode = YpirMode::kSimplepir; + psi::spiral::Params spiral_params; + uint64_t db_rows = 0; + uint64_t db_cols = 0; + size_t value_bytes = 0; + + [[nodiscard]] uint64_t NumItems() const { return db_rows * db_cols; } +}; + +YpirParameters CreateParamsForScenarioSimplePIR(uint64_t num_items, + uint64_t item_size_bits); +YpirParameters CreateParamsForScenarioDoublePIR(uint64_t num_items, + uint64_t item_size_bits); +YpirParameters CreateParamsForShapeSimplePIR(uint64_t db_rows, uint64_t db_cols, + uint64_t item_size_bits); +YpirParameters CreateParamsForShapeDoublePIR(uint64_t db_rows, uint64_t db_cols, + uint64_t item_size_bits); +YpirParameters CreateSmallTestParamsSimplePIR(); +YpirParameters CreateSmallTestParamsDoublePIR(); + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/params_test.cc b/psi/algorithm/ypir/params_test.cc new file mode 100644 index 00000000..3350ab5b --- /dev/null +++ b/psi/algorithm/ypir/params_test.cc @@ -0,0 +1,51 @@ +#include "psi/algorithm/ypir/params.h" + +#include "gtest/gtest.h" + +namespace psi::ypir { +namespace { + +TEST(YpirParamsTest, SmallSimplepirParamsMatchFastRegressionShape) { + const auto params = CreateSmallTestParamsSimplePIR(); + EXPECT_EQ(params.mode, YpirMode::kSimplepir); + EXPECT_EQ(params.db_rows, 1ULL << 10); + EXPECT_EQ(params.db_cols, 1ULL << 10); + EXPECT_EQ(params.value_bytes, 2U); +} + +TEST(YpirParamsTest, SmallDoublepirParamsMatchFastRegressionShape) { + const auto params = CreateSmallTestParamsDoublePIR(); + EXPECT_EQ(params.mode, YpirMode::kDoublepir); + EXPECT_EQ(params.db_rows, 1ULL << 10); + EXPECT_EQ(params.db_cols, 1ULL << 10); + EXPECT_EQ(params.value_bytes, 1U); +} + +TEST(YpirParamsTest, ScenarioHelpersRespectRequestedValueSize) { + const auto simple_params = + CreateParamsForScenarioSimplePIR((1ULL << 18) + 7, 8); + const auto double_params = + CreateParamsForScenarioDoublePIR((1ULL << 19) + 9, 24); + + EXPECT_GE(simple_params.NumItems(), (1ULL << 18) + 7); + EXPECT_EQ(simple_params.value_bytes, 1U); + EXPECT_GE(double_params.NumItems(), (1ULL << 19) + 9); + EXPECT_EQ(double_params.value_bytes, 3U); +} + +TEST(YpirParamsTest, ShapeHelpersPreserveRequestedShape) { + const auto simple_params = + CreateParamsForShapeSimplePIR(1ULL << 11, 3ULL << 10, 16); + const auto double_params = + CreateParamsForShapeDoublePIR(1ULL << 11, 1ULL << 12, 8); + + EXPECT_EQ(simple_params.db_rows, 1ULL << 11); + EXPECT_EQ(simple_params.db_cols, 3ULL << 10); + EXPECT_EQ(simple_params.value_bytes, 2U); + EXPECT_EQ(double_params.db_rows, 1ULL << 11); + EXPECT_EQ(double_params.db_cols, 1ULL << 12); + EXPECT_EQ(double_params.value_bytes, 1U); +} + +} // namespace +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/pir_interface_test.cc b/psi/algorithm/ypir/pir_interface_test.cc new file mode 100644 index 00000000..8e2f279f --- /dev/null +++ b/psi/algorithm/ypir/pir_interface_test.cc @@ -0,0 +1,94 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "psi/algorithm/ypir/client.h" +#include "psi/algorithm/ypir/server.h" + +namespace psi::ypir { +namespace { + +uint64_t BytesToU64(const std::vector& bytes) { + uint64_t value = 0; + for (size_t i = 0; i < bytes.size(); ++i) { + value |= static_cast(bytes[i]) << (8 * i); + } + return value; +} + +template +psi::pir::RawDatabase BuildItemDatabase(const YpirParameters& params, + Fn&& make_value) { + std::vector> rows( + params.NumItems(), std::vector(params.value_bytes, 0)); + for (uint64_t raw_idx = 0; raw_idx < params.NumItems(); ++raw_idx) { + const T value = make_value(raw_idx); + std::memcpy(rows[raw_idx].data(), &value, params.value_bytes); + } + return psi::pir::RawDatabase(std::move(rows)); +} + +TEST(YpirPirInterfaceTest, SimplepirImplementsIndexPirInterfaces) { + auto params = CreateSmallTestParamsSimplePIR(); + YpirClient concrete_client(params); + YpirServer concrete_server(params); + + psi::pir::IndexPirClient& client = concrete_client; + psi::pir::IndexPirDataBase& server = concrete_server; + + const auto raw_db = + BuildItemDatabase(params, [&](uint64_t raw_idx) { + const uint64_t row = raw_idx / params.db_cols; + const uint64_t col = raw_idx % params.db_cols; + return static_cast((row * 17 + col * 3) % + params.spiral_params.PtModulus()); + }); + server.GenerateFromRawData(raw_db); + + const uint64_t raw_idx = 123ULL * params.db_cols + 456ULL; + const auto pks = client.GeneratePksBuffer(); + const auto query = client.GenerateIndexQuery(raw_idx); + const auto response = server.Response(query, pks); + const auto decoded = client.DecodeIndexResponse(response, raw_idx); + + EXPECT_EQ(client.GetPirType(), psi::pir::PirType::YPIR_PIR); + EXPECT_EQ(server.GetPirType(), psi::pir::PirType::YPIR_PIR); + EXPECT_EQ(server.MaxElementsOfOnePt(), 1U); + EXPECT_TRUE(server.DbSeted()); + EXPECT_EQ(BytesToU64(decoded), + (123ULL * 17 + 456ULL * 3) % params.spiral_params.PtModulus()); +} + +TEST(YpirPirInterfaceTest, DoublepirImplementsIndexPirInterfaces) { + auto params = CreateSmallTestParamsDoublePIR(); + YpirClient concrete_client(params); + YpirServer concrete_server(params); + + psi::pir::IndexPirClient& client = concrete_client; + psi::pir::IndexPirDataBase& server = concrete_server; + + const auto raw_db = BuildItemDatabase(params, [&](uint64_t raw_idx) { + const uint64_t row = raw_idx / params.db_cols; + const uint64_t col = raw_idx % params.db_cols; + return static_cast((row + col) % 251); + }); + server.GenerateFromRawData(raw_db); + + const uint64_t raw_idx = 111ULL * params.db_cols + 222ULL; + const auto pks = client.GeneratePksBuffer(); + const auto query = client.GenerateIndexQuery(raw_idx); + const auto response = server.Response(query, pks); + const auto decoded = client.DecodeIndexResponse(response, raw_idx); + + EXPECT_EQ(client.GetPirType(), psi::pir::PirType::YPIR_PIR); + EXPECT_EQ(server.GetPirType(), psi::pir::PirType::YPIR_PIR); + EXPECT_EQ(server.MaxElementsOfOnePt(), 1U); + EXPECT_TRUE(server.DbSeted()); + ASSERT_EQ(decoded.size(), 1U); + EXPECT_EQ(decoded[0], static_cast((111 + 222) % 251)); +} + +} // namespace +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/serialize.cc b/psi/algorithm/ypir/serialize.cc new file mode 100644 index 00000000..ef9b441f --- /dev/null +++ b/psi/algorithm/ypir/serialize.cc @@ -0,0 +1,140 @@ +#include "psi/algorithm/ypir/serialize.h" + +#include +#include +#include +#include +#include + +#include "yacl/base/exception.h" + +namespace psi::ypir { +namespace { + +template +void AppendPod(std::string& out, const T& value) { + out.append(reinterpret_cast(&value), sizeof(T)); +} + +template +T ReadPod(const char*& ptr, const char* end) { + YACL_ENFORCE_GE(end - ptr, static_cast(sizeof(T))); + T value; + std::memcpy(&value, ptr, sizeof(T)); + ptr += sizeof(T); + return value; +} + +template +void AppendVector(std::string& out, const std::vector& values) { + AppendPod(out, values.size()); + if (!values.empty()) { + out.append(reinterpret_cast(values.data()), + values.size() * sizeof(T)); + } +} + +template +std::vector ReadVector(const char*& ptr, const char* end) { + const auto size = ReadPod(ptr, end); + YACL_ENFORCE_GE(end - ptr, static_cast(size * sizeof(T))); + std::vector out(size); + if (size > 0) { + std::memcpy(out.data(), ptr, size * sizeof(T)); + ptr += size * sizeof(T); + } + return out; +} + +template +void AppendVector2D(std::string& out, + const std::vector>& values) { + AppendPod(out, values.size()); + for (const auto& inner : values) { + AppendVector(out, inner); + } +} + +template +std::vector> ReadVector2D(const char*& ptr, const char* end) { + const auto outer = ReadPod(ptr, end); + std::vector> out; + out.reserve(outer); + for (uint64_t i = 0; i < outer; ++i) { + out.push_back(ReadVector(ptr, end)); + } + return out; +} + +template +void AppendVector3D(std::string& out, + const std::vector>>& values) { + AppendPod(out, values.size()); + for (const auto& inner : values) { + AppendVector2D(out, inner); + } +} + +template +std::vector>> ReadVector3D(const char*& ptr, + const char* end) { + const auto outer = ReadPod(ptr, end); + std::vector>> out; + out.reserve(outer); + for (uint64_t i = 0; i < outer; ++i) { + out.push_back(ReadVector2D(ptr, end)); + } + return out; +} + +yacl::Buffer ToBuffer(const std::string& bytes) { + return yacl::Buffer(bytes.data(), bytes.size()); +} + +} // namespace + +yacl::Buffer SerializeQuery(const YpirQuery& query) { + std::string bytes; + AppendPod(bytes, static_cast(query.mode)); + AppendVector(bytes, query.packed_query_row); + AppendVector(bytes, query.qu0); + AppendVector(bytes, query.qu1); + AppendVector3D(bytes, query.ksk_b); + return ToBuffer(bytes); +} + +YpirQuery DeserializeQuery(const yacl::ByteContainerView& buffer) { + const char* ptr = reinterpret_cast(buffer.data()); + const char* end = ptr + buffer.size(); + + YpirQuery query; + query.mode = static_cast(ReadPod(ptr, end)); + query.packed_query_row = ReadVector(ptr, end); + query.qu0 = ReadVector(ptr, end); + query.qu1 = ReadVector(ptr, end); + query.ksk_b = ReadVector3D(ptr, end); + YACL_ENFORCE(ptr == end, "unexpected trailing bytes in YpirQuery"); + return query; +} + +yacl::Buffer SerializeResponse(const YpirResponse& response) { + std::string bytes; + AppendPod(bytes, static_cast(response.mode)); + AppendVector(bytes, response.simplepir_response); + AppendVector2D(bytes, response.doublepir_response); + return ToBuffer(bytes); +} + +YpirResponse DeserializeResponse(const yacl::ByteContainerView& buffer) { + const char* ptr = reinterpret_cast(buffer.data()); + const char* end = ptr + buffer.size(); + + YpirResponse response; + response.mode = static_cast(ReadPod(ptr, end)); + response.simplepir_response = ReadVector(ptr, end); + response.doublepir_response = ReadVector2D(ptr, end); + YACL_ENFORCE(ptr == end, "unexpected trailing bytes in YpirResponse"); + return response; +} + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/serialize.h b/psi/algorithm/ypir/serialize.h new file mode 100644 index 00000000..60198dab --- /dev/null +++ b/psi/algorithm/ypir/serialize.h @@ -0,0 +1,16 @@ +#pragma once + +#include "yacl/base/buffer.h" +#include "yacl/base/byte_container_view.h" + +#include "psi/algorithm/ypir/types.h" + +namespace psi::ypir { + +yacl::Buffer SerializeQuery(const YpirQuery& query); +YpirQuery DeserializeQuery(const yacl::ByteContainerView& buffer); + +yacl::Buffer SerializeResponse(const YpirResponse& response); +YpirResponse DeserializeResponse(const yacl::ByteContainerView& buffer); + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/serialize_test.cc b/psi/algorithm/ypir/serialize_test.cc new file mode 100644 index 00000000..2a61c212 --- /dev/null +++ b/psi/algorithm/ypir/serialize_test.cc @@ -0,0 +1,43 @@ +#include "psi/algorithm/ypir/serialize.h" + +#include + +#include "gtest/gtest.h" + +namespace psi::ypir { +namespace { + +TEST(YpirSerializeTest, QueryRoundTrip) { + YpirQuery query; + query.mode = YpirMode::kDoublepir; + query.packed_query_row = {1, 2, 3}; + query.qu0 = {4, 5}; + query.qu1 = {6, 7, 8}; + query.ksk_b = {{{9, 10}, {11}}, {{12, 13, 14}}}; + + const auto buffer = SerializeQuery(query); + const auto decoded = DeserializeQuery(buffer); + + EXPECT_EQ(decoded.mode, query.mode); + EXPECT_EQ(decoded.packed_query_row, query.packed_query_row); + EXPECT_EQ(decoded.qu0, query.qu0); + EXPECT_EQ(decoded.qu1, query.qu1); + EXPECT_EQ(decoded.ksk_b, query.ksk_b); +} + +TEST(YpirSerializeTest, ResponseRoundTrip) { + YpirResponse response; + response.mode = YpirMode::kDoublepir; + response.simplepir_response = {15, 16, 17}; + response.doublepir_response = {{18, 19}, {20}}; + + const auto buffer = SerializeResponse(response); + const auto decoded = DeserializeResponse(buffer); + + EXPECT_EQ(decoded.mode, response.mode); + EXPECT_EQ(decoded.simplepir_response, response.simplepir_response); + EXPECT_EQ(decoded.doublepir_response, response.doublepir_response); +} + +} // namespace +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/server.cc b/psi/algorithm/ypir/server.cc new file mode 100644 index 00000000..bfb4feb7 --- /dev/null +++ b/psi/algorithm/ypir/server.cc @@ -0,0 +1,186 @@ +#include "psi/algorithm/ypir/server.h" + +#include +#include +#include +#include +#include +#include + +#include "yacl/base/exception.h" + +#include "psi/algorithm/ypir/legacy/util.h" +#include "psi/algorithm/ypir/serialize.h" + +namespace psi::ypir { +namespace { + +template +std::vector RowBytesToValues(const std::vector& bytes) { + YACL_ENFORCE_EQ(bytes.size() % sizeof(T), 0U); + std::vector out(bytes.size() / sizeof(T)); + std::memcpy(out.data(), bytes.data(), bytes.size()); + return out; +} + +template +T ReadScalarValue(const std::vector& bytes) { + YACL_ENFORCE_EQ(bytes.size(), sizeof(T)); + T out = 0; + std::memcpy(&out, bytes.data(), sizeof(T)); + return out; +} + +} // namespace + +template +YpirServer::YpirServer(YpirParameters params) + : psi::pir::IndexPirDataBase(psi::pir::PirType::YPIR_PIR), + params_(std::move(params)) { + if (params_.mode == YpirMode::kDoublepir) { + ypir_context_ = std::make_unique( + internal::ypir::CreateContext(params_)); + } +} + +template +void YpirServer::GenerateFromRawData( + const psi::pir::RawDatabase& raw_database) { + static_assert(std::is_same_v || std::is_same_v, + "YpirServer only supports uint8_t and uint16_t values"); + + const bool item_layout = raw_database.Rows() <= params_.NumItems() && + raw_database.RowByteLen() == params_.value_bytes; + const bool matrix_layout = + raw_database.Rows() == params_.db_rows && + raw_database.RowByteLen() == params_.db_cols * sizeof(T); + YACL_ENFORCE(item_layout || matrix_layout, + "raw database shape does not match YPIR parameters"); + + if (params_.mode == YpirMode::kSimplepir) { + simplepir_db_column_major_.assign(params_.db_rows * params_.db_cols, 0); + if (item_layout) { + YACL_ENFORCE_EQ(params_.value_bytes, sizeof(T)); + for (uint64_t raw_idx = 0; raw_idx < raw_database.Rows(); ++raw_idx) { + const uint64_t row = raw_idx / params_.db_cols; + const uint64_t col = raw_idx % params_.db_cols; + simplepir_db_column_major_[col * params_.db_rows + row] = + ReadScalarValue(raw_database.At(raw_idx)); + } + } else { + for (uint64_t row = 0; row < params_.db_rows; ++row) { + auto row_values = RowBytesToValues(raw_database.At(row)); + for (uint64_t col = 0; col < params_.db_cols; ++col) { + simplepir_db_column_major_[col * params_.db_rows + row] = + row_values[col]; + } + } + } + } else { + YACL_ENFORCE(sizeof(T) == 1, "DoublePIR currently expects uint8_t rows"); + doublepir_db_row_major_.assign(params_.db_rows * params_.db_cols, 0); + if (item_layout) { + YACL_ENFORCE_EQ(params_.value_bytes, sizeof(uint8_t)); + for (uint64_t raw_idx = 0; raw_idx < raw_database.Rows(); ++raw_idx) { + doublepir_db_row_major_[raw_idx] = + ReadScalarValue(raw_database.At(raw_idx)); + } + } else { + for (uint64_t row = 0; row < params_.db_rows; ++row) { + const auto& row_bytes = raw_database.At(row); + std::memcpy(doublepir_db_row_major_.data() + row * params_.db_cols, + row_bytes.data(), row_bytes.size()); + } + } + } + + db_set_ = true; +} + +template +void YpirServer::GenerateFromSimpleHashTable( + const psi::pir::RawDatabase& raw_database) { + GenerateFromRawData(raw_database); +} + +template +void YpirServer::Dump(std::ostream& out_stream) const { + out_stream << "YpirServer{mode=" + << (params_.mode == YpirMode::kSimplepir ? "simplepir" + : "doublepir") + << ", db_rows=" << params_.db_rows + << ", db_cols=" << params_.db_cols + << ", value_bytes=" << params_.value_bytes + << ", db_set=" << db_set_ << "}"; +} + +template +YpirPrecomputedState YpirServer::PerformOfflinePrecomputation() const { + YACL_ENFORCE(db_set_, "database must be loaded before precomputation"); + + if (params_.mode == YpirMode::kSimplepir) { + YpirPrecomputedState state; + state.mode = YpirMode::kSimplepir; + return state; + } + + YACL_ENFORCE(ypir_context_ != nullptr); + return internal::ypir::PrepareOfflineState(doublepir_db_row_major_, params_, + *ypir_context_); +} + +template +YpirResponse YpirServer::ProcessQuery(const YpirQuery& query) const { + return ProcessQuery(query, PerformOfflinePrecomputation()); +} + +template +YpirResponse YpirServer::ProcessQuery( + const YpirQuery& query, const YpirPrecomputedState& state) const { + YACL_ENFORCE(db_set_, "database must be loaded before query processing"); + YACL_ENFORCE(query.mode == params_.mode); + + if (params_.mode == YpirMode::kSimplepir) { + YACL_ENFORCE_EQ(query.packed_query_row.size(), params_.db_rows); + + std::vector result(params_.db_cols, 0); + psi::ypir::FastBatchedDotProduct( + params_.spiral_params, result.data(), query.packed_query_row.data(), + params_.db_rows, simplepir_db_column_major_.data(), params_.db_rows, + params_.db_cols); + + YpirResponse response; + response.mode = YpirMode::kSimplepir; + response.simplepir_response = std::move(result); + return response; + } + + YACL_ENFORCE(ypir_context_ != nullptr); + return internal::ypir::ProcessQuery(doublepir_db_row_major_, query, state, + params_, *ypir_context_); +} + +template +yacl::Buffer YpirServer::Response( + const yacl::ByteContainerView& query_buffer) const { + return SerializeResponse(ProcessQuery(DeserializeQuery(query_buffer))); +} + +template +yacl::Buffer YpirServer::Response( + const yacl::ByteContainerView& query_buffer, + const yacl::Buffer& /*pks_buffer*/) const { + return Response(query_buffer); +} + +template +std::string YpirServer::Response(const yacl::ByteContainerView& query_buffer, + const std::string& /*pks_buffer*/) const { + auto buffer = Response(query_buffer); + return std::string(static_cast(buffer)); +} + +template class YpirServer; +template class YpirServer; + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/server.h b/psi/algorithm/ypir/server.h new file mode 100644 index 00000000..9f6c1e46 --- /dev/null +++ b/psi/algorithm/ypir/server.h @@ -0,0 +1,55 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "yacl/base/buffer.h" +#include "yacl/base/byte_container_view.h" + +#include "psi/algorithm/pir_interface/pir_db.h" +#include "psi/algorithm/ypir/params.h" +#include "psi/algorithm/ypir/types.h" +#include "psi/algorithm/ypir/ypir_internal_server.h" + +namespace psi::ypir { + +template +class YpirServer : public psi::pir::IndexPirDataBase { + public: + explicit YpirServer(YpirParameters params); + + void GenerateFromRawData(const psi::pir::RawDatabase& raw_database) override; + void GenerateFromSimpleHashTable( + const psi::pir::RawDatabase& raw_database) override; + void Dump(std::ostream& out_stream) const override; + std::size_t MaxElementsOfOnePt() const override { return 1; } + bool DbSeted() const override { return db_set_; } + + [[nodiscard]] bool DbSet() const { return db_set_; } + [[nodiscard]] const YpirParameters& GetParameters() const { return params_; } + + YpirPrecomputedState PerformOfflinePrecomputation() const; + YpirResponse ProcessQuery(const YpirQuery& query) const; + YpirResponse ProcessQuery(const YpirQuery& query, + const YpirPrecomputedState& state) const; + yacl::Buffer Response(const yacl::ByteContainerView& query_buffer) const; + yacl::Buffer Response(const yacl::ByteContainerView& query_buffer, + const yacl::Buffer& pks_buffer) const override; + std::string Response(const yacl::ByteContainerView& query_buffer, + const std::string& pks_buffer) const override; + + private: + YpirParameters params_; + bool db_set_ = false; + std::vector simplepir_db_column_major_; + std::vector doublepir_db_row_major_; + std::unique_ptr ypir_context_; +}; + +extern template class YpirServer; +extern template class YpirServer; + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/server_test.cc b/psi/algorithm/ypir/server_test.cc new file mode 100644 index 00000000..f1530d59 --- /dev/null +++ b/psi/algorithm/ypir/server_test.cc @@ -0,0 +1,52 @@ +#include "psi/algorithm/ypir/server.h" + +#include +#include +#include + +#include "gtest/gtest.h" + +namespace psi::ypir { +namespace { + +template +psi::pir::RawDatabase BuildDatabase(const YpirParameters& params) { + std::vector> rows( + params.db_rows, std::vector(params.db_cols * sizeof(T), 0)); + for (uint64_t row = 0; row < params.db_rows; ++row) { + for (uint64_t col = 0; col < params.db_cols; ++col) { + const T value = static_cast((row * 17 + col * 3) % 251); + std::memcpy(rows[row].data() + col * sizeof(T), &value, sizeof(T)); + } + } + return psi::pir::RawDatabase(std::move(rows)); +} + +TEST(YpirServerTest, SimplepirOfflineStateIsEmpty) { + auto params = CreateSmallTestParamsSimplePIR(); + YpirServer server(params); + server.GenerateFromRawData(BuildDatabase(params)); + + const auto state = server.PerformOfflinePrecomputation(); + EXPECT_EQ(state.mode, YpirMode::kSimplepir); + EXPECT_TRUE(state.hint_0.empty()); + EXPECT_TRUE(state.server_hint.empty()); + EXPECT_TRUE(state.decomp_buf.empty()); +} + +TEST(YpirServerTest, DoublepirOfflineStateCanBeReused) { + auto params = CreateSmallTestParamsDoublePIR(); + YpirServer server(params); + server.GenerateFromRawData(BuildDatabase(params)); + + const auto state1 = server.PerformOfflinePrecomputation(); + const auto state2 = server.PerformOfflinePrecomputation(); + + EXPECT_EQ(state1.mode, YpirMode::kDoublepir); + EXPECT_EQ(state2.mode, YpirMode::kDoublepir); + EXPECT_EQ(state1.hint_0, state2.hint_0); + EXPECT_EQ(state1.server_hint, state2.server_hint); +} + +} // namespace +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/simplepir_flow_test.cc b/psi/algorithm/ypir/simplepir_flow_test.cc new file mode 100644 index 00000000..320adb15 --- /dev/null +++ b/psi/algorithm/ypir/simplepir_flow_test.cc @@ -0,0 +1,56 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "psi/algorithm/ypir/client.h" +#include "psi/algorithm/ypir/server.h" + +namespace psi::ypir { +namespace { + +uint64_t BytesToU64(const std::vector& bytes) { + uint64_t value = 0; + for (size_t i = 0; i < bytes.size(); ++i) { + value |= static_cast(bytes[i]) << (8 * i); + } + return value; +} + +psi::pir::RawDatabase BuildSimplepirDatabase(const YpirParameters& params) { + std::vector> rows( + params.db_rows, + std::vector(params.db_cols * sizeof(uint16_t), 0)); + for (uint64_t row = 0; row < params.db_rows; ++row) { + for (uint64_t col = 0; col < params.db_cols; ++col) { + const uint16_t value = static_cast( + (row * 17 + col * 3) % params.spiral_params.PtModulus()); + std::memcpy(rows[row].data() + col * sizeof(uint16_t), &value, + sizeof(value)); + } + } + return psi::pir::RawDatabase(std::move(rows)); +} + +TEST(YpirSimplepirFlowTest, EndToEndSmallDatabase) { + auto params = CreateSmallTestParamsSimplePIR(); + YpirClient client(params); + YpirServer server(params); + server.GenerateFromRawData(BuildSimplepirDatabase(params)); + + const uint64_t row = 123; + const uint64_t col = 456; + const uint64_t raw_idx = row * params.db_cols + col; + + const auto query_buffer = client.GenerateQueryBuffer(raw_idx); + const auto response_buffer = server.Response(query_buffer); + const auto decoded = client.DecodeResponseBuffer(response_buffer, raw_idx); + + const uint64_t expected = + (row * 17 + col * 3) % params.spiral_params.PtModulus(); + EXPECT_EQ(BytesToU64(decoded), expected); +} + +} // namespace +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/smoke_test.cc b/psi/algorithm/ypir/smoke_test.cc new file mode 100644 index 00000000..71ef48e6 --- /dev/null +++ b/psi/algorithm/ypir/smoke_test.cc @@ -0,0 +1,25 @@ +#include "gtest/gtest.h" + +#include "psi/algorithm/ypir/client.h" +#include "psi/algorithm/ypir/server.h" + +namespace psi::ypir { +namespace { + +TEST(YpirSmokeTest, ConstructPublicObjects) { + auto simple_params = CreateSmallTestParamsSimplePIR(); + auto double_params = CreateSmallTestParamsDoublePIR(); + + YpirClient simple_client(simple_params); + YpirClient double_client(double_params); + YpirServer simple_server(simple_params); + YpirServer double_server(double_params); + + EXPECT_EQ(simple_client.GetParameters().mode, YpirMode::kSimplepir); + EXPECT_EQ(double_client.GetParameters().mode, YpirMode::kDoublepir); + EXPECT_FALSE(simple_server.DbSet()); + EXPECT_FALSE(double_server.DbSet()); +} + +} // namespace +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/types.h b/psi/algorithm/ypir/types.h new file mode 100644 index 00000000..24d301c2 --- /dev/null +++ b/psi/algorithm/ypir/types.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include + +namespace psi::ypir { + +enum class YpirMode : uint8_t { + kSimplepir = 0, + kDoublepir = 1, +}; + +struct YpirPrecomputedState { + YpirMode mode = YpirMode::kSimplepir; + std::vector hint_0; + std::vector> server_hint; + std::vector>> decomp_buf; +}; + +struct YpirQuery { + YpirMode mode = YpirMode::kSimplepir; + std::vector packed_query_row; + std::vector qu0; + std::vector qu1; + std::vector>> ksk_b; +}; + +struct YpirResponse { + YpirMode mode = YpirMode::kSimplepir; + std::vector simplepir_response; + std::vector> doublepir_response; +}; + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/util.cc b/psi/algorithm/ypir/util.cc new file mode 100644 index 00000000..aa622971 --- /dev/null +++ b/psi/algorithm/ypir/util.cc @@ -0,0 +1,44 @@ +#include "psi/algorithm/ypir/util.h" + +#include +#include +#include + +#include "yacl/base/exception.h" + +#include "psi/algorithm/spiral/arith/arith.h" + +namespace psi::ypir { + +std::vector BuildPackedSimplepirQuery(const YpirParameters& params, + uint64_t raw_idx) { + YACL_ENFORCE(params.mode == YpirMode::kSimplepir); + YACL_ENFORCE_LT(raw_idx, params.NumItems()); + + const uint64_t row_idx = raw_idx / params.db_cols; + std::vector query(params.db_rows, 0); + query[row_idx] = params.spiral_params.ScaleK(); + + const uint64_t m0 = params.spiral_params.Moduli(0); + const uint64_t m1 = params.spiral_params.Moduli(1); + std::vector packed(params.db_rows, 0); + for (size_t i = 0; i < packed.size(); ++i) { + packed[i] = (query[i] % m0) | ((query[i] % m1) << 32); + } + return packed; +} + +std::vector EncodeIntegerValue(uint64_t value, size_t value_bytes) { + std::vector out(value_bytes, 0); + for (size_t i = 0; i < value_bytes; ++i) { + out[i] = static_cast((value >> (8 * i)) & 0xff); + } + return out; +} + +uint64_t DecodeSimplepirValue(const YpirParameters& params, uint64_t value) { + return psi::spiral::arith::Rescale(value, params.spiral_params.Modulus(), + params.spiral_params.PtModulus()); +} + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/util.h b/psi/algorithm/ypir/util.h new file mode 100644 index 00000000..905b4ccc --- /dev/null +++ b/psi/algorithm/ypir/util.h @@ -0,0 +1,16 @@ +#pragma once + +#include +#include +#include + +#include "psi/algorithm/ypir/params.h" + +namespace psi::ypir { + +std::vector BuildPackedSimplepirQuery(const YpirParameters& params, + uint64_t raw_idx); +std::vector EncodeIntegerValue(uint64_t value, size_t value_bytes); +uint64_t DecodeSimplepirValue(const YpirParameters& params, uint64_t value); + +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/ypir_flow_test.cc b/psi/algorithm/ypir/ypir_flow_test.cc new file mode 100644 index 00000000..c7099c4f --- /dev/null +++ b/psi/algorithm/ypir/ypir_flow_test.cc @@ -0,0 +1,45 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "psi/algorithm/ypir/client.h" +#include "psi/algorithm/ypir/serialize.h" +#include "psi/algorithm/ypir/server.h" + +namespace psi::ypir { +namespace { + +psi::pir::RawDatabase BuildDoublepirDatabase(const YpirParameters& params) { + std::vector> rows( + params.db_rows, std::vector(params.db_cols, 0)); + for (uint64_t row = 0; row < params.db_rows; ++row) { + for (uint64_t col = 0; col < params.db_cols; ++col) { + rows[row][col] = static_cast((row + col) % 251); + } + } + return psi::pir::RawDatabase(std::move(rows)); +} + +TEST(YpirDoublepirFlowTest, EndToEndSmallDatabase) { + auto params = CreateSmallTestParamsDoublePIR(); + YpirClient client(params); + YpirServer server(params); + server.GenerateFromRawData(BuildDoublepirDatabase(params)); + + const auto state = server.PerformOfflinePrecomputation(); + const uint64_t row = 111; + const uint64_t col = 222; + const uint64_t raw_idx = row * params.db_cols + col; + + const auto query_buffer = client.GenerateQueryBuffer(raw_idx); + const auto response = + server.ProcessQuery(DeserializeQuery(query_buffer), state); + const auto decoded = client.DecodeResponse(response, raw_idx); + + ASSERT_EQ(decoded.size(), 1U); + EXPECT_EQ(decoded[0], static_cast((row + col) % 251)); +} + +} // namespace +} // namespace psi::ypir diff --git a/psi/algorithm/ypir/ypir_internal_client.cc b/psi/algorithm/ypir/ypir_internal_client.cc new file mode 100644 index 00000000..4c675417 --- /dev/null +++ b/psi/algorithm/ypir/ypir_internal_client.cc @@ -0,0 +1,57 @@ +#include "psi/algorithm/ypir/ypir_internal_client.h" + +#include "yacl/base/exception.h" + +#include "psi/algorithm/ypir/util.h" + +namespace psi::ypir::internal::ypir { + +YpirQuery GenerateQuery(uint64_t raw_idx, const YpirParameters& params, + ClientSecrets& secrets, const Context& context) { + YACL_ENFORCE(params.mode == YpirMode::kDoublepir); + YACL_ENFORCE_LT(raw_idx, params.NumItems()); + + const uint64_t row_idx = raw_idx / params.db_cols; + const uint64_t col_idx = raw_idx % params.db_cols; + + secrets.simple_secret = + psi::ypir::ypir_internal::Secret(context.fhe_params->get_lwe_dimension(), + context.fhe_params->get_lwe_cmod()); + secrets.double_secret = + psi::ypir::ypir_internal::Secret(context.fhe_params->get_poly_degree(), + context.fhe_params->get_rlwe_cmod()); + secrets.initialized = true; + + auto legacy_query = psi::ypir::ypir_internal::Generate_query_ypir( + col_idx, row_idx, secrets.simple_secret, secrets.double_secret, + *context.prng, *context.fhe_params, *context.pir_params); + + YpirQuery query; + query.mode = YpirMode::kDoublepir; + query.qu0 = std::move(legacy_query.qu0); + query.qu1 = std::move(legacy_query.qu1); + query.ksk_b = std::move(legacy_query.ksk_b); + return query; +} + +std::vector RecoverResponse(const YpirResponse& response, + const YpirParameters& params, + const ClientSecrets& secrets, + const Context& context) { + YACL_ENFORCE(params.mode == YpirMode::kDoublepir); + YACL_ENFORCE(secrets.initialized, + "GenerateQuery must be called before decode"); + YACL_ENFORCE(response.mode == YpirMode::kDoublepir); + + auto simple_secret = secrets.simple_secret; + auto double_secret = secrets.double_secret; + auto result = response.doublepir_response; + uint64_t message = 0; + + psi::ypir::ypir_internal::YpirRecover(simple_secret, double_secret, result, + message, *context.fhe_params, + *context.pir_params); + return EncodeIntegerValue(message, params.value_bytes); +} + +} // namespace psi::ypir::internal::ypir diff --git a/psi/algorithm/ypir/ypir_internal_client.h b/psi/algorithm/ypir/ypir_internal_client.h new file mode 100644 index 00000000..29053ec7 --- /dev/null +++ b/psi/algorithm/ypir/ypir_internal_client.h @@ -0,0 +1,25 @@ +#pragma once + +#include + +#include "psi/algorithm/ypir/legacy/client.h" +#include "psi/algorithm/ypir/types.h" +#include "psi/algorithm/ypir/ypir_internal_params.h" + +namespace psi::ypir::internal::ypir { + +struct ClientSecrets { + psi::ypir::ypir_internal::Secret simple_secret; + psi::ypir::ypir_internal::Secret double_secret; + bool initialized = false; +}; + +YpirQuery GenerateQuery(uint64_t raw_idx, const YpirParameters& params, + ClientSecrets& secrets, const Context& context); + +std::vector RecoverResponse(const YpirResponse& response, + const YpirParameters& params, + const ClientSecrets& secrets, + const Context& context); + +} // namespace psi::ypir::internal::ypir diff --git a/psi/algorithm/ypir/ypir_internal_params.cc b/psi/algorithm/ypir/ypir_internal_params.cc new file mode 100644 index 00000000..6ab1f052 --- /dev/null +++ b/psi/algorithm/ypir/ypir_internal_params.cc @@ -0,0 +1,60 @@ +#include "psi/algorithm/ypir/ypir_internal_params.h" + +#include + +#include "yacl/base/exception.h" + +namespace psi::ypir::internal::ypir { +namespace { + +thread_local psi::ypir::ypir_internal::YpirHexlNtt* g_ntt = nullptr; + +void NttForwardWrapper(uint64_t* data, size_t len) { + g_ntt->Forward(data, len); +} + +Context BuildContext(uint64_t rows, uint64_t cols) { + using namespace psi::ypir::ypir_internal; + + const uint64_t lwe_dimension = 1024; + const uint64_t lwe_pmod = (1 << 8); + const uint64_t lwe_cmod = (1ULL << 32); + const uint64_t rlwe_degree = 2048; + const uint64_t rlwe_pmod = (1 << 16); + const uint64_t rlwe_cmod = kCrtMod; + const double sigma = 3.19; + const double sigma_ring = 3.19; + const AutoParams auto_params{20, 18, 2}; + const DecompParams decomp_params{16, 16, 1}; + + auto fhe_params = std::make_shared( + rlwe_degree, rlwe_cmod, rlwe_pmod, lwe_dimension, lwe_cmod, lwe_pmod, + sigma, sigma_ring, auto_params, decomp_params); + auto pir_params = std::make_shared(rows, cols); + auto prng = std::make_shared(); + + g_ntt = &fhe_params->get_ntt(); + fhe_params->SetNttForward(&NttForwardWrapper); + fhe_params->set_persudo_matrix_simplepir(pir_params->get_col()); + fhe_params->set_persudo_matrix_doublepir(pir_params->get_row()); + fhe_params->set_persudo_hcube_ypir(); + + const uint64_t expo = GetLog2(rlwe_degree); + std::vector auto_idx; + for (uint64_t i = 1; i <= expo; ++i) { + auto_idx.push_back((1ULL << i) + 1); + } + fhe_params->set_automap(auto_idx); + fhe_params->set_precomputed_pt(expo); + + return {std::move(fhe_params), std::move(pir_params), std::move(prng)}; +} + +} // namespace + +Context CreateContext(const YpirParameters& params) { + YACL_ENFORCE(params.mode == YpirMode::kDoublepir); + return BuildContext(params.db_rows, params.db_cols); +} + +} // namespace psi::ypir::internal::ypir diff --git a/psi/algorithm/ypir/ypir_internal_params.h b/psi/algorithm/ypir/ypir_internal_params.h new file mode 100644 index 00000000..4e665a26 --- /dev/null +++ b/psi/algorithm/ypir/ypir_internal_params.h @@ -0,0 +1,18 @@ +#pragma once + +#include + +#include "psi/algorithm/ypir/legacy/ypir_params.h" +#include "psi/algorithm/ypir/params.h" + +namespace psi::ypir::internal::ypir { + +struct Context { + std::shared_ptr fhe_params; + std::shared_ptr pir_params; + std::shared_ptr prng; +}; + +Context CreateContext(const YpirParameters& params); + +} // namespace psi::ypir::internal::ypir diff --git a/psi/algorithm/ypir/ypir_internal_server.cc b/psi/algorithm/ypir/ypir_internal_server.cc new file mode 100644 index 00000000..526885e7 --- /dev/null +++ b/psi/algorithm/ypir/ypir_internal_server.cc @@ -0,0 +1,70 @@ +#include "psi/algorithm/ypir/ypir_internal_server.h" + +#include "yacl/base/exception.h" + +#include "psi/algorithm/ypir/legacy/server.h" + +namespace psi::ypir::internal::ypir { +namespace { + +std::vector> ExpandDatabase( + const std::vector& db, uint64_t rows, uint64_t cols) { + YACL_ENFORCE_EQ(db.size(), rows * cols); + std::vector> out(rows, std::vector(cols, 0)); + for (uint64_t row = 0; row < rows; ++row) { + for (uint64_t col = 0; col < cols; ++col) { + out[row][col] = db[row * cols + col]; + } + } + return out; +} + +} // namespace + +YpirPrecomputedState PrepareOfflineState(const std::vector& db, + const YpirParameters& params, + const Context& context) { + YACL_ENFORCE(params.mode == YpirMode::kDoublepir); + auto db_matrix = ExpandDatabase(db, params.db_rows, params.db_cols); + + YpirPrecomputedState state; + state.mode = YpirMode::kDoublepir; + std::vector> server_hint; + psi::ypir::ypir_internal::YpirHintGenerate( + db_matrix, state.hint_0, server_hint, state.decomp_buf, *context.prng, + *context.fhe_params, *context.pir_params); + psi::ypir::ypir_internal::MatrixTranspose(server_hint, state.server_hint); + return state; +} + +YpirResponse ProcessQuery(const std::vector& db, + const YpirQuery& query, + const YpirPrecomputedState& state, + const YpirParameters& params, + const Context& context) { + YACL_ENFORCE(params.mode == YpirMode::kDoublepir); + YACL_ENFORCE(query.mode == YpirMode::kDoublepir); + YACL_ENFORCE(state.mode == YpirMode::kDoublepir); + + std::vector qu0(query.qu0.size(), 0); + for (size_t i = 0; i < query.qu0.size(); ++i) { + qu0[i] = static_cast(query.qu0[i]); + } + + auto qu1 = query.qu1; + auto decomp_buf = state.decomp_buf; + const auto& server_hint = state.server_hint; + + std::vector> result; + result.push_back(state.hint_0); + psi::ypir::ypir_internal::YpirAnswer( + db.data(), qu0.data(), qu1, query.ksk_b, decomp_buf, server_hint, result, + *context.fhe_params, *context.pir_params); + + YpirResponse response; + response.mode = YpirMode::kDoublepir; + response.doublepir_response = std::move(result); + return response; +} + +} // namespace psi::ypir::internal::ypir diff --git a/psi/algorithm/ypir/ypir_internal_server.h b/psi/algorithm/ypir/ypir_internal_server.h new file mode 100644 index 00000000..dd3070e5 --- /dev/null +++ b/psi/algorithm/ypir/ypir_internal_server.h @@ -0,0 +1,20 @@ +#pragma once + +#include +#include + +#include "psi/algorithm/ypir/types.h" +#include "psi/algorithm/ypir/ypir_internal_params.h" + +namespace psi::ypir::internal::ypir { + +YpirPrecomputedState PrepareOfflineState(const std::vector& db, + const YpirParameters& params, + const Context& context); + +YpirResponse ProcessQuery(const std::vector& db, + const YpirQuery& query, + const YpirPrecomputedState& state, + const YpirParameters& params, const Context& context); + +} // namespace psi::ypir::internal::ypir diff --git a/psi/apps/psi_launcher/BUILD.bazel b/psi/apps/psi_launcher/BUILD.bazel index 7723ed47..05aea38e 100644 --- a/psi/apps/psi_launcher/BUILD.bazel +++ b/psi/apps/psi_launcher/BUILD.bazel @@ -42,12 +42,35 @@ psi_cc_library( hdrs = ["report.h"], ) +psi_cc_library( + name = "launch_ypir", + srcs = select({ + "@platforms//cpu:x86_64": ["ypir_launch.cc"], + "//conditions:default": ["ypir_launch_stub.cc"], + }), + hdrs = ["ypir_launch.h"], + deps = select({ + "@platforms//cpu:x86_64": [ + "//psi/algorithm/ypir:entry", + "//psi/proto:pir_cc_proto", + "@yacl//yacl/base:exception", + "@yacl//yacl/link:context", + ], + "//conditions:default": [ + "//psi/proto:pir_cc_proto", + "@yacl//yacl/base:exception", + "@yacl//yacl/link:context", + ], + }), +) + psi_cc_library( name = "launch", srcs = ["launch.cc"], hdrs = ["launch.h"], deps = [ ":factory", + ":launch_ypir", ":report", "//psi:trace_categories", "//psi/algorithm/dkpir:entry", @@ -72,6 +95,17 @@ psi_cc_test( ], ) +psi_cc_test( + name = "ypir_test", + srcs = ["ypir_test.cc"], + target_compatible_with = ["@platforms//cpu:x86_64"], + deps = [ + ":launch_ypir", + "@abseil-cpp//absl/strings", + "@yacl//yacl/link:test_util", + ], +) + psi_cc_library( name = "kuscia_adapter", srcs = [ diff --git a/psi/apps/psi_launcher/launch.h b/psi/apps/psi_launcher/launch.h index 080f550a..0b5dad32 100644 --- a/psi/apps/psi_launcher/launch.h +++ b/psi/apps/psi_launcher/launch.h @@ -47,6 +47,12 @@ PirResultReport RunPir(const ApsiReceiverConfig& apsi_receiver_config, PirResultReport RunPir(const ApsiSenderConfig& apsi_sender_config, const std::shared_ptr& lctx); +PirResultReport RunPir(const YpirReceiverConfig& ypir_receiver_config, + const std::shared_ptr& lctx); + +PirResultReport RunPir(const YpirSenderConfig& ypir_sender_config, + const std::shared_ptr& lctx); + PirResultReport RunDkPir(const DkPirReceiverConfig& dk_pir_receiver_config, const std::shared_ptr& lctx); diff --git a/psi/apps/psi_launcher/main.cc b/psi/apps/psi_launcher/main.cc index 39650064..1d65811e 100644 --- a/psi/apps/psi_launcher/main.cc +++ b/psi/apps/psi_launcher/main.cc @@ -121,6 +121,18 @@ int main(int argc, char* argv[]) { YACL_ENFORCE(google::protobuf::util::MessageToJsonString( report, &report_json, json_print_options) .ok()); + } else if (launch_config.has_ypir_sender_config()) { + psi::PirResultReport report = + psi::RunPir(launch_config.ypir_sender_config(), lctx); + YACL_ENFORCE(google::protobuf::util::MessageToJsonString( + report, &report_json, json_print_options) + .ok()); + } else if (launch_config.has_ypir_receiver_config()) { + psi::PirResultReport report = + psi::RunPir(launch_config.ypir_receiver_config(), lctx); + YACL_ENFORCE(google::protobuf::util::MessageToJsonString( + report, &report_json, json_print_options) + .ok()); } else if (launch_config.has_dk_pir_sender_config()) { psi::PirResultReport report = psi::RunDkPir(launch_config.dk_pir_sender_config(), lctx); diff --git a/psi/apps/psi_launcher/ypir_launch.cc b/psi/apps/psi_launcher/ypir_launch.cc new file mode 100644 index 00000000..53f33255 --- /dev/null +++ b/psi/apps/psi_launcher/ypir_launch.cc @@ -0,0 +1,50 @@ +#include "psi/apps/psi_launcher/ypir_launch.h" + +#include "yacl/base/exception.h" + +#include "psi/algorithm/ypir/entry.h" + +namespace psi { +namespace { + +psi::ypir::YpirMode ToYpirMode(YpirProtocolMode mode) { + switch (mode) { + case YPIR_PROTOCOL_MODE_SIMPLEPIR: + return psi::ypir::YpirMode::kSimplepir; + case YPIR_PROTOCOL_MODE_DOUBLEPIR: + return psi::ypir::YpirMode::kDoublepir; + default: + YACL_THROW("unsupported YPIR mode"); + } +} + +} // namespace + +PirResultReport RunPir(const YpirReceiverConfig& ypir_receiver_config, + const std::shared_ptr& lctx) { + psi::ypir::YpirReceiverOptions options; + options.mode = ToYpirMode(ypir_receiver_config.mode()); + options.db_rows = ypir_receiver_config.db_rows(); + options.db_cols = ypir_receiver_config.db_cols(); + options.item_size_bits = ypir_receiver_config.item_size_bits(); + options.query_file = ypir_receiver_config.query_file(); + options.output_file = ypir_receiver_config.output_file(); + + YACL_ENFORCE_EQ(psi::ypir::ReceiverOnline(options, lctx), 0); + return PirResultReport(); +} + +PirResultReport RunPir(const YpirSenderConfig& ypir_sender_config, + const std::shared_ptr& lctx) { + psi::ypir::YpirSenderOptions options; + options.mode = ToYpirMode(ypir_sender_config.mode()); + options.db_rows = ypir_sender_config.db_rows(); + options.db_cols = ypir_sender_config.db_cols(); + options.item_size_bits = ypir_sender_config.item_size_bits(); + options.db_file = ypir_sender_config.db_file(); + + YACL_ENFORCE_EQ(psi::ypir::SenderOnline(options, lctx), 0); + return PirResultReport(); +} + +} // namespace psi diff --git a/psi/apps/psi_launcher/ypir_launch.h b/psi/apps/psi_launcher/ypir_launch.h new file mode 100644 index 00000000..86151c06 --- /dev/null +++ b/psi/apps/psi_launcher/ypir_launch.h @@ -0,0 +1,17 @@ +#pragma once + +#include + +#include "yacl/link/context.h" + +#include "psi/proto/pir.pb.h" + +namespace psi { + +PirResultReport RunPir(const YpirReceiverConfig& ypir_receiver_config, + const std::shared_ptr& lctx); + +PirResultReport RunPir(const YpirSenderConfig& ypir_sender_config, + const std::shared_ptr& lctx); + +} // namespace psi diff --git a/psi/apps/psi_launcher/ypir_launch_stub.cc b/psi/apps/psi_launcher/ypir_launch_stub.cc new file mode 100644 index 00000000..793f6acd --- /dev/null +++ b/psi/apps/psi_launcher/ypir_launch_stub.cc @@ -0,0 +1,17 @@ +#include "yacl/base/exception.h" + +#include "psi/apps/psi_launcher/ypir_launch.h" + +namespace psi { + +PirResultReport RunPir(const YpirReceiverConfig&, + const std::shared_ptr&) { + YACL_THROW("YPIR is only supported on x86_64"); +} + +PirResultReport RunPir(const YpirSenderConfig&, + const std::shared_ptr&) { + YACL_THROW("YPIR is only supported on x86_64"); +} + +} // namespace psi diff --git a/psi/apps/psi_launcher/ypir_test.cc b/psi/apps/psi_launcher/ypir_test.cc new file mode 100644 index 00000000..33bfb3c6 --- /dev/null +++ b/psi/apps/psi_launcher/ypir_test.cc @@ -0,0 +1,146 @@ +#include +#include +#include +#include +#include + +#include "absl/strings/escaping.h" +#include "gtest/gtest.h" +#include "yacl/link/test_util.h" + +#include "psi/apps/psi_launcher/ypir_launch.h" + +namespace psi { +namespace { + +std::string HexByte(uint64_t value) { + std::vector bytes = {static_cast(value)}; + return absl::BytesToHexString(absl::string_view( + reinterpret_cast(bytes.data()), bytes.size())); +} + +std::string HexUint16(uint64_t value) { + const uint16_t narrowed = static_cast(value); + return absl::BytesToHexString(absl::string_view( + reinterpret_cast(&narrowed), sizeof(narrowed))); +} + +void WriteLines(const std::filesystem::path& path, + const std::vector& lines) { + std::ofstream output(path, std::ios::out | std::ios::trunc); + ASSERT_TRUE(output.is_open()); + for (const auto& line : lines) { + output << line << '\n'; + } +} + +std::vector ReadLines(const std::filesystem::path& path) { + std::ifstream input(path); + EXPECT_TRUE(input.is_open()); + std::vector lines; + std::string line; + while (std::getline(input, line)) { + if (!line.empty()) { + lines.push_back(line); + } + } + return lines; +} + +TEST(YpirLauncherTest, RunPirSimplepirUsesLauncherEntry) { + const auto tmp_dir = + std::filesystem::temp_directory_path() / "ypir_launcher_simplepir"; + std::filesystem::create_directories(tmp_dir); + const auto db_path = tmp_dir / "db.hex"; + const auto query_path = tmp_dir / "query.txt"; + const auto output_path = tmp_dir / "result.hex"; + + std::vector db_lines; + db_lines.reserve(1ULL << 20); + for (uint64_t raw_idx = 0; raw_idx < (1ULL << 20); ++raw_idx) { + const uint64_t row = raw_idx / (1ULL << 10); + const uint64_t col = raw_idx % (1ULL << 10); + db_lines.push_back(HexUint16((row * 17 + col * 3) % (1ULL << 14))); + } + WriteLines(db_path, db_lines); + WriteLines(query_path, {"0", "123456", "1048575"}); + + YpirSenderConfig sender_config; + sender_config.set_mode(YPIR_PROTOCOL_MODE_SIMPLEPIR); + sender_config.set_db_rows(1ULL << 10); + sender_config.set_db_cols(1ULL << 10); + sender_config.set_item_size_bits(16); + sender_config.set_db_file(db_path); + + YpirReceiverConfig receiver_config; + receiver_config.set_mode(YPIR_PROTOCOL_MODE_SIMPLEPIR); + receiver_config.set_db_rows(1ULL << 10); + receiver_config.set_db_cols(1ULL << 10); + receiver_config.set_item_size_bits(16); + receiver_config.set_query_file(query_path); + receiver_config.set_output_file(output_path); + + auto lctxs = yacl::link::test::SetupWorld(2); + auto sender = std::async(std::launch::async, + [&] { return RunPir(sender_config, lctxs[0]); }); + auto receiver = std::async(std::launch::async, + [&] { return RunPir(receiver_config, lctxs[1]); }); + + EXPECT_EQ(sender.get().match_cnt(), 0); + EXPECT_EQ(receiver.get().match_cnt(), 0); + EXPECT_EQ(ReadLines(output_path), + (std::vector{HexUint16(0), HexUint16(3768), + HexUint16(4076)})); + + std::filesystem::remove_all(tmp_dir); +} + +TEST(YpirLauncherTest, RunPirDoublepirUsesLauncherEntry) { + const auto tmp_dir = + std::filesystem::temp_directory_path() / "ypir_launcher_doublepir"; + std::filesystem::create_directories(tmp_dir); + const auto db_path = tmp_dir / "db.hex"; + const auto query_path = tmp_dir / "query.txt"; + const auto output_path = tmp_dir / "result.hex"; + + std::vector db_lines; + db_lines.reserve(1ULL << 20); + for (uint64_t raw_idx = 0; raw_idx < (1ULL << 20); ++raw_idx) { + const uint64_t row = raw_idx / (1ULL << 10); + const uint64_t col = raw_idx % (1ULL << 10); + db_lines.push_back(HexByte((row + col) % 251)); + } + WriteLines(db_path, db_lines); + WriteLines(query_path, {"0", "113886", "1048575"}); + + YpirSenderConfig sender_config; + sender_config.set_mode(YPIR_PROTOCOL_MODE_DOUBLEPIR); + sender_config.set_db_rows(1ULL << 10); + sender_config.set_db_cols(1ULL << 10); + sender_config.set_item_size_bits(8); + sender_config.set_db_file(db_path); + + YpirReceiverConfig receiver_config; + receiver_config.set_mode(YPIR_PROTOCOL_MODE_DOUBLEPIR); + receiver_config.set_db_rows(1ULL << 10); + receiver_config.set_db_cols(1ULL << 10); + receiver_config.set_item_size_bits(8); + receiver_config.set_query_file(query_path); + receiver_config.set_output_file(output_path); + + auto lctxs = yacl::link::test::SetupWorld(2); + auto sender = std::async(std::launch::async, + [&] { return RunPir(sender_config, lctxs[0]); }); + auto receiver = std::async(std::launch::async, + [&] { return RunPir(receiver_config, lctxs[1]); }); + + EXPECT_EQ(sender.get().match_cnt(), 0); + EXPECT_EQ(receiver.get().match_cnt(), 0); + EXPECT_EQ(ReadLines(output_path), + (std::vector{HexByte(0), HexByte(82), HexByte(38)})); + + std::filesystem::remove_all(tmp_dir); +} + +} // namespace +} // namespace psi diff --git a/psi/proto/entry.proto b/psi/proto/entry.proto index 265ebaa5..49c65bca 100644 --- a/psi/proto/entry.proto +++ b/psi/proto/entry.proto @@ -47,5 +47,9 @@ message LaunchConfig { DkPirSenderConfig dk_pir_sender_config = 8; DkPirReceiverConfig dk_pir_receiver_config = 9; + + YpirSenderConfig ypir_sender_config = 10; + + YpirReceiverConfig ypir_receiver_config = 11; } } diff --git a/psi/proto/pir.proto b/psi/proto/pir.proto index 094ecba8..87a426ed 100644 --- a/psi/proto/pir.proto +++ b/psi/proto/pir.proto @@ -225,7 +225,50 @@ message DkPirReceiverConfig { bool skip_count_check = 11; } +enum YpirProtocolMode { + YPIR_PROTOCOL_MODE_UNSPECIFIED = 0; + YPIR_PROTOCOL_MODE_SIMPLEPIR = 1; + YPIR_PROTOCOL_MODE_DOUBLEPIR = 2; +} + +message YpirSenderConfig { + // YPIR mode. SIMPLEPIR currently supports up to 16-bit values and DOUBLEPIR + // currently supports 8-bit values. + YpirProtocolMode mode = 1; + + // Logical database shape. + uint64 db_rows = 2; + uint64 db_cols = 3; + + // Plaintext bit width of each value. + uint64 item_size_bits = 4; + + // Path to a text file that contains exactly db_rows * db_cols hex-encoded + // values, one per line, without a header. + string db_file = 5; +} + +message YpirReceiverConfig { + // YPIR mode. Must match sender config. + YpirProtocolMode mode = 1; + + // Logical database shape. Must match sender config. + uint64 db_rows = 2; + uint64 db_cols = 3; + + // Plaintext bit width of each value. Must match sender config. + uint64 item_size_bits = 4; + + // Path to a text file containing raw indexes in decimal, one per line, + // without a header. + string query_file = 5; + + // Path to a text file where hex-encoded results will be written, one per + // line, without a header. + string output_file = 6; +} + // The report of pir task. message PirResultReport { int64 match_cnt = 1; -} \ No newline at end of file +}