Ryan Henry
0d22946a0e
Ship the TLS mesh, composer, Beaver/Yao/leaf MPC, prep/online paths, apps, and docs so the tree is pushable before elevating share_expr, security_mode, and prep resume. Co-authored-by: Cursor <cursoragent@cursor.com>
311 lines
10 KiB
C++
311 lines
10 KiB
C++
#include <gtest/gtest.h>
|
|
|
|
#include "dpf.hpp"
|
|
|
|
#include <array>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <type_traits>
|
|
#include <vector>
|
|
|
|
namespace
|
|
{
|
|
|
|
template <typename LaneT>
|
|
using storage = dpf::leaf_storage_t<LaneT>;
|
|
|
|
template <typename LaneT>
|
|
constexpr unsigned lane_bits()
|
|
{
|
|
if constexpr (dpf::utils::is_packed_subbyte_v<LaneT>)
|
|
return static_cast<unsigned>(dpf::utils::packed_lane_bits_v<LaneT>);
|
|
else
|
|
return static_cast<unsigned>(sizeof(LaneT) * 8u);
|
|
}
|
|
|
|
template <typename LaneT>
|
|
constexpr std::uint64_t lane_mask()
|
|
{
|
|
if constexpr (dpf::utils::is_packed_subbyte_v<LaneT>)
|
|
return (std::uint64_t{1} << dpf::utils::packed_lane_bits_v<LaneT>) - 1u;
|
|
else if constexpr (sizeof(LaneT) >= 8)
|
|
return ~std::uint64_t{0};
|
|
else
|
|
return (std::uint64_t{1} << (sizeof(LaneT) * 8u)) - 1u;
|
|
}
|
|
|
|
template <typename LaneT>
|
|
std::uint64_t get_lane(const storage<LaneT> * buf, std::size_t i)
|
|
{
|
|
if constexpr (dpf::utils::is_packed_subbyte_v<LaneT>)
|
|
{
|
|
constexpr unsigned w = dpf::utils::packed_lane_bits_v<LaneT>;
|
|
const std::size_t bit = i * w;
|
|
const std::size_t word = bit / 64u;
|
|
const unsigned shift = static_cast<unsigned>(bit % 64u);
|
|
std::uint64_t v = buf[word] >> shift;
|
|
if (shift + w > 64u)
|
|
v |= buf[word + 1u] << (64u - shift);
|
|
return v & lane_mask<LaneT>();
|
|
}
|
|
else if constexpr (std::is_class_v<LaneT>)
|
|
{
|
|
return static_cast<std::uint64_t>(
|
|
static_cast<typename LaneT::integral_type>(buf[i])) & lane_mask<LaneT>();
|
|
}
|
|
else
|
|
{
|
|
return static_cast<std::uint64_t>(buf[i]) & lane_mask<LaneT>();
|
|
}
|
|
}
|
|
|
|
template <typename LaneT>
|
|
void set_lane(storage<LaneT> * buf, std::size_t i, std::uint64_t val)
|
|
{
|
|
val &= lane_mask<LaneT>();
|
|
if constexpr (dpf::utils::is_packed_subbyte_v<LaneT>)
|
|
{
|
|
constexpr unsigned w = dpf::utils::packed_lane_bits_v<LaneT>;
|
|
const std::size_t bit = i * w;
|
|
const std::size_t word = bit / 64u;
|
|
const unsigned shift = static_cast<unsigned>(bit % 64u);
|
|
const std::uint64_t mask = lane_mask<LaneT>();
|
|
buf[word] = (buf[word] & ~(mask << shift)) | (val << shift);
|
|
if (shift + w > 64u)
|
|
{
|
|
const unsigned lo = 64u - shift;
|
|
buf[word + 1u] = (buf[word + 1u] & ~(mask >> lo)) | (val >> lo);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
buf[i] = static_cast<storage<LaneT>>(val);
|
|
}
|
|
}
|
|
|
|
template <typename LaneT>
|
|
std::vector<storage<LaneT>> make_key(std::size_t nleaves, std::size_t key_id)
|
|
{
|
|
std::vector<storage<LaneT>> v(dpf::leaf_storage_words<LaneT>(nleaves),
|
|
storage<LaneT>{0});
|
|
for (std::size_t i = 0; i < nleaves; ++i)
|
|
{
|
|
// Distinct pattern per (key, leaf); fits every width's mask.
|
|
const std::uint64_t val
|
|
= (key_id * 131u + i * 17u + 3u) & lane_mask<LaneT>();
|
|
set_lane<LaneT>(v.data(), i, val);
|
|
}
|
|
return v;
|
|
}
|
|
|
|
template <typename LaneT>
|
|
void check_roundtrip(std::size_t nkeys, std::size_t nleaves)
|
|
{
|
|
std::vector<std::vector<storage<LaneT>>> owned;
|
|
owned.reserve(nkeys);
|
|
std::vector<const storage<LaneT> *> in_ptrs;
|
|
std::vector<storage<LaneT> *> out_ptrs;
|
|
in_ptrs.reserve(nkeys);
|
|
out_ptrs.reserve(nkeys);
|
|
for (std::size_t k = 0; k < nkeys; ++k)
|
|
{
|
|
owned.push_back(make_key<LaneT>(nleaves, k));
|
|
in_ptrs.push_back(owned.back().data());
|
|
}
|
|
|
|
std::vector<storage<LaneT>> interleaved(
|
|
dpf::leaf_storage_words<LaneT>(nkeys * nleaves), storage<LaneT>{0});
|
|
dpf::interleave_leaves<LaneT>(interleaved.data(), in_ptrs.data(), nkeys,
|
|
nleaves);
|
|
|
|
for (std::size_t i = 0; i < nleaves; ++i)
|
|
{
|
|
for (std::size_t k = 0; k < nkeys; ++k)
|
|
{
|
|
const std::size_t g = dpf::cohort_index(i, k, nkeys);
|
|
EXPECT_EQ(get_lane<LaneT>(interleaved.data(), g),
|
|
get_lane<LaneT>(owned[k].data(), i))
|
|
<< "LaneT bits=" << lane_bits<LaneT>() << " nkeys=" << nkeys
|
|
<< " nleaves=" << nleaves << " i=" << i << " k=" << k;
|
|
}
|
|
}
|
|
|
|
std::vector<std::vector<storage<LaneT>>> round;
|
|
round.reserve(nkeys);
|
|
for (std::size_t k = 0; k < nkeys; ++k)
|
|
{
|
|
round.emplace_back(dpf::leaf_storage_words<LaneT>(nleaves),
|
|
storage<LaneT>{0});
|
|
out_ptrs.push_back(round.back().data());
|
|
}
|
|
dpf::deinterleave_leaves<LaneT>(out_ptrs.data(), interleaved.data(), nkeys,
|
|
nleaves);
|
|
|
|
for (std::size_t k = 0; k < nkeys; ++k)
|
|
{
|
|
for (std::size_t i = 0; i < nleaves; ++i)
|
|
{
|
|
EXPECT_EQ(get_lane<LaneT>(round[k].data(), i),
|
|
get_lane<LaneT>(owned[k].data(), i))
|
|
<< "deinterleave bits=" << lane_bits<LaneT>()
|
|
<< " nkeys=" << nkeys << " i=" << i << " k=" << k;
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
template <typename LaneT>
|
|
class InterleaveLeavesTest : public ::testing::Test
|
|
{
|
|
};
|
|
|
|
using InterleaveTypes = ::testing::Types<
|
|
std::uint64_t,
|
|
std::uint32_t,
|
|
std::uint16_t,
|
|
std::uint8_t,
|
|
dpf::nyble,
|
|
dpf::twobit,
|
|
dpf::bit,
|
|
dpf::gf2,
|
|
dpf::gf22,
|
|
dpf::gf24,
|
|
dpf::gf28,
|
|
dpf::gf216,
|
|
dpf::gf232,
|
|
dpf::gf264>;
|
|
|
|
TYPED_TEST_SUITE(InterleaveLeavesTest, InterleaveTypes);
|
|
|
|
TYPED_TEST(InterleaveLeavesTest, KeyCountsAndOddLength)
|
|
{
|
|
// Length not a multiple of pack width (8 for bits-in-byte, 32 for
|
|
// 2-bit-in-u64, 16 for 4-bit-in-u64, etc.): 13 is odd for all of those.
|
|
constexpr std::size_t nleaves = 13;
|
|
for (std::size_t nkeys : {std::size_t{1}, std::size_t{3}, std::size_t{5},
|
|
std::size_t{16}})
|
|
check_roundtrip<TypeParam>(nkeys, nleaves);
|
|
}
|
|
|
|
TYPED_TEST(InterleaveLeavesTest, LongerNotMultipleOfWord)
|
|
{
|
|
// 70 leaves: not a multiple of 64 (bit word) or 32 (twobit word).
|
|
check_roundtrip<TypeParam>(5, 70);
|
|
check_roundtrip<TypeParam>(8, 70);
|
|
}
|
|
|
|
TEST(InterleaveLeaves, BitPackingIsKeyMajorWithinLeaf)
|
|
{
|
|
// Three keys, five 1-bit leaves. Logical order of bits in the output
|
|
// stream is (i0,k0),(i0,k1),(i0,k2),(i1,k0),... — not a byte index of
|
|
// i*m+k into a uint8_t array.
|
|
constexpr std::size_t nkeys = 3;
|
|
constexpr std::size_t nleaves = 5;
|
|
std::array<std::uint64_t, 1> k0{0}, k1{0}, k2{0};
|
|
// key0: 1 0 1 1 0
|
|
// key1: 0 1 1 0 1
|
|
// key2: 1 1 0 0 1
|
|
k0[0] = 0b01101ull;
|
|
k1[0] = 0b10110ull;
|
|
k2[0] = 0b10011ull;
|
|
const std::uint64_t * keys[3] = {k0.data(), k1.data(), k2.data()};
|
|
std::uint64_t out = 0;
|
|
dpf::interleave_leaves<dpf::bit>(&out, keys, nkeys, nleaves);
|
|
|
|
// Expected 15 bits, low-first: for i=0..4, bits of keys 0,1,2
|
|
// i0: 1,0,1 | i1: 0,1,1 | i2: 1,1,0 | i3: 1,0,0 | i4: 0,1,1
|
|
// bit0..14 = 1,0,1,0,1,1,1,1,0,1,0,0,0,1,1
|
|
const std::uint64_t expect = 0b110001011110101ull;
|
|
EXPECT_EQ(out & ((1ull << 15) - 1ull), expect);
|
|
}
|
|
|
|
TEST(InterleaveLeaves, InnerProductModintFromBits)
|
|
{
|
|
// Same three keys as above. Leaf integers (key 0 in the low bit):
|
|
// i0 = 0b101 = 5, i1 = 0b110 = 6, i2 = 0b011 = 3, i3 = 0b001 = 1, i4 = 0b110 = 6.
|
|
constexpr std::size_t nkeys = 3;
|
|
constexpr std::size_t nleaves = 5;
|
|
std::array<std::uint64_t, 1> k0{0b01101ull}, k1{0b10110ull}, k2{0b10011ull};
|
|
const std::uint64_t * keys[3] = {k0.data(), k1.data(), k2.data()};
|
|
std::uint64_t bits = 0;
|
|
dpf::interleave_leaves<dpf::bit>(&bits, keys, nkeys, nleaves);
|
|
|
|
const std::uint64_t w[5] = {1, 2, 3, 4, 5};
|
|
const auto got = dpf::interleaved_bits_inner_product<3>(&bits, nleaves, nkeys, w);
|
|
// 5*1 + 6*2 + 3*3 + 1*4 + 6*5 = 5+12+9+4+30 = 60 ≡ 60-7*8 = 4 (mod 8)
|
|
EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<3>::integral_type>(got)), 4u);
|
|
|
|
std::array<dpf::modint<3>, 5> mw{1, 2, 3, 4, 5};
|
|
const auto got_m = dpf::interleaved_bits_inner_product<3>(&bits, nleaves, nkeys, mw);
|
|
EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<3>::integral_type>(got_m)), 4u);
|
|
}
|
|
|
|
TEST(InterleaveLeaves, InnerProductModint128)
|
|
{
|
|
constexpr std::size_t nkeys = 128;
|
|
constexpr std::size_t nleaves = 3;
|
|
std::vector<std::uint64_t> key_bits(nkeys * 2, 0);
|
|
std::vector<const std::uint64_t *> ptrs(nkeys);
|
|
for (std::size_t k = 0; k < nkeys; ++k)
|
|
{
|
|
// Leaf 0: bit k set. Leaf 1: only key 0. Leaf 2: zero.
|
|
if (k == 0)
|
|
key_bits[k * 2] = 0b011ull;
|
|
else
|
|
key_bits[k * 2] = 0b001ull;
|
|
ptrs[k] = key_bits.data() + k * 2;
|
|
}
|
|
std::vector<std::uint64_t> bits(dpf::leaf_storage_words<dpf::bit>(nkeys * nleaves));
|
|
dpf::interleave_leaves<dpf::bit>(bits.data(), ptrs.data(), nkeys, nleaves);
|
|
|
|
const std::uint64_t w[3] = {1, 3, 9};
|
|
const auto got = dpf::interleaved_bits_inner_product<128>(
|
|
bits.data(), nleaves, nkeys, w);
|
|
using limb = dpf::modint<128>::integral_type;
|
|
const limb all = ~limb{0};
|
|
const limb expect = all * limb{1} + limb{1} * limb{3};
|
|
EXPECT_EQ(static_cast<limb>(got), expect);
|
|
}
|
|
|
|
TEST(InterleaveLeaves, InnerProductLowBitsOfWideGroup)
|
|
{
|
|
// Eight 1-bit keys, read as modint<4>: only keys 0..3.
|
|
constexpr std::size_t nkeys = 8;
|
|
constexpr std::size_t nleaves = 2;
|
|
std::array<std::uint64_t, 8> raw{};
|
|
const std::uint64_t * ptrs[8];
|
|
for (std::size_t k = 0; k < nkeys; ++k)
|
|
{
|
|
raw[k] = (k < 4) ? 0b11ull : 0b01ull; // both leaves set for low keys
|
|
ptrs[k] = &raw[k];
|
|
}
|
|
std::uint64_t bits[4]{};
|
|
dpf::interleave_leaves<dpf::bit>(bits, ptrs, nkeys, nleaves);
|
|
const std::uint64_t w[2] = {1, 1};
|
|
const auto got = dpf::interleaved_bits_inner_product<4>(bits, nleaves, nkeys, w);
|
|
// Each leaf's low 4 bits are 0b1111 = 15. 15+15 = 30 ≡ 14 (mod 16).
|
|
EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<4>::integral_type>(got)), 14u);
|
|
}
|
|
|
|
TYPED_TEST(InterleaveLeavesTest, EmptyAndSingleLeaf)
|
|
{
|
|
check_roundtrip<TypeParam>(0, 0);
|
|
check_roundtrip<TypeParam>(3, 0);
|
|
check_roundtrip<TypeParam>(0, 5);
|
|
check_roundtrip<TypeParam>(1, 1);
|
|
check_roundtrip<TypeParam>(4, 1);
|
|
}
|
|
|
|
TEST(InterleaveLeaves, InnerProductEmptyIsZero)
|
|
{
|
|
const std::uint64_t w[1] = {9};
|
|
EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<3>::integral_type>(
|
|
dpf::interleaved_bits_inner_product<3>(
|
|
static_cast<const std::uint64_t *>(nullptr), 0, 3, w))),
|
|
0u);
|
|
std::uint64_t bits = 0;
|
|
EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<3>::integral_type>(
|
|
dpf::interleaved_bits_inner_product<3>(&bits, 5, 0, w))),
|
|
0u);
|
|
}
|