Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Compact step reprise #961

Merged

Conversation

martinthomson
Copy link
Member

@martinthomson martinthomson commented Feb 27, 2024

A working version of a new CompactStep and CompactGate implementation, automatically derived from annotations on XxxStep types.

This works in two phases: build.rs and main.

build.rs phase

Here, code is generated during the execution of build.rs. This is incomplete code, but only for CompactGate, which only includes stubbed implementations of AsRef<str> and From<&str>.

Step implementations generate complete implementations of CompactStep, which includes a bunch of helpful functions that allow for traversing the tree of steps. Some of these functions are not used in the main code, but they are used in a build.rs file.

The build.rs file includes (using the include!() macro, hidden by a convenience macro) all of the files that contain steps. It then runs a build process that generates a file for each CompactGate target. This will create implementations of AsRef<str> and From<&str> that the main build can pick up.

main build phase

Here, the same code is generated, but an include!() call will be made to pick up the implementations generated by build.rs.

Usage

The derivation of CompactStep is a proc-macro. This is moderately complex, but not overtly so. Here is a complete example of what an annotated Step might look like:

use ipa_step_derive::{CompactStep, CompactGate};
#[derive(CompactStep, CompactGate)]
enum MyStep {
  Empty,
  #[step(count = 10)]
  Int(u8),
  #[step(child = OtherStep)]
  OperationWithOtherSteps,
  #[step(child = OtherStep, count = 2, name = "short-name")]
  SetOfOperationsWithOtherSteps(u8),
}

This is a top-level step as well, so it requires that it is put in a file on its own, so that it can be included in build.rs.

The build.rs file looks fairly simple (I tried to compress this further, but ran into issues building a macros-by-example wrapper for this, so this tiny bit of repetition will have to do for now):

use ipa_step::{build_gate, load_steps, track_steps};
load_steps!(path::to::my_step);
fn main() {
  track_steps!(path::to::my_step);
  build_gate::<path::to::my_step::MyStep>();
}

The two macros list modules (or modules and files, load_steps!(foo @ "src/foo/mod.rs") might be necessary, though I doubt it will be common). load_steps!() ensures that the modules are loaded with a module path that works for generating files. track_steps!() ensures that build.rs is re-run if anything changes in those modules. Finally, build_gate() generates the code necessary for each gate and puts that in a file in $OUT_DIR, where the proc-macro for the corresponding CompactGate generation will be able to find and include!() it. The build process uses the complete implementation of CompactStep on the target type to guide the building of the implementation.

Example

There is an ipa-step-test crate in the branch that demonstrates how this might be used. That includes some tests that demonstrate that the the StepNarrow implementation is working and correct.

TODO

This currently duplicates the Step-related machinery in ipa-core. The goal is to migrate that code out and start to use this instead. Unfortunately, there is no easy path to transition, so I'll have to just go in and annotate and move everything around.

I'm not 100% sure that each step only has a single child step type. If we have any places where two different types are found, some tweaking might be needed.

I haven't implemented anything for struct Foo(u8), only enum Foo { Bar(u8), }. That might require a small amount of adjustment.

Copy link

codecov bot commented Feb 27, 2024

Codecov Report

Attention: Patch coverage is 92.35569% with 147 lines in your changes are missing coverage. Please review.

Project coverage is 90.81%. Comparing base (908c68e) to head (4b01187).

Files Patch % Lines
ipa-step/src/gate.rs 30.76% 63 Missing ⚠️
ipa-step-derive/src/track.rs 89.78% 19 Missing ⚠️
ipa-step-derive/src/lib.rs 98.92% 10 Missing ⚠️
ipa-step-derive/src/variant.rs 97.27% 10 Missing ⚠️
ipa-core/src/protocol/boolean/mod.rs 60.00% 6 Missing ⚠️
ipa-core/src/protocol/boolean/step.rs 71.42% 6 Missing ⚠️
ipa-core/src/protocol/context/mod.rs 68.75% 5 Missing ⚠️
ipa-core/src/protocol/step.rs 16.66% 5 Missing ⚠️
ipa-core/src/query/executor.rs 54.54% 5 Missing ⚠️
ipa-core/src/protocol/ipa_prf/aggregation/step.rs 55.55% 4 Missing ⚠️
... and 10 more
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #961      +/-   ##
==========================================
+ Coverage   89.81%   90.81%   +1.00%     
==========================================
  Files         164      184      +20     
  Lines       24460    25713    +1253     
==========================================
+ Hits        21969    23352    +1383     
+ Misses       2491     2361     -130     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Copy link
Collaborator

@akoshelev akoshelev left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks awesome, I can't wait to start using it!

integer: None,
child: step_child,
} = self
else {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VariantAttribute is a struct, so I am not sure I understand the purpose of else here

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh, didn't know you can do that

ipa-step/src/name.rs Show resolved Hide resolved
@@ -42,6 +42,11 @@ impl Step for String {}
#[cfg(any(feature = "test-fixture", debug_assertions))]
impl Step for str {}

pub trait CompactStep: Step {
const STEP_COUNT: usize;
fn step_string(i: usize) -> String;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should probably document this trait

contents.lines().map(|s| s.to_owned()).collect::<Vec<_>>()
contents
.lines()
.map(std::borrow::ToOwned::to_owned)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

probably better to import ToOwned trait as it's been used more than once in this file

ipa-step/src/gate.rs Outdated Show resolved Hide resolved
($($a:ident$(::$b:ident)* $(@ $f:expr)*),*) => {
$(println!("cargo:rerun-if-changed={f}", f = ::ipa_step::module_file!($a$(::$b)* $(@ $f)?));)*
println!("cargo:rerun-if-changed={f}", f = ::std::file!());
assert!(::std::env::var(::ipa_step::COMPACT_GATE_INCLUDE_ENV).is_err(),
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

ipa-step/src/gate.rs Outdated Show resolved Hide resolved
@@ -0,0 +1,78 @@
#![allow(clippy::module_name_repetitions)]

#[cfg(any(feature = "descriptive", debug_assertions))]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does it work when feature is defined outside of this crate?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No. We'll have to pass it through, as I do in Cargo.toml, though I realize now that I only added it in the defaults, which isn't right.

akoshelev added a commit to akoshelev/raw-ipa that referenced this pull request Mar 15, 2024
PR private-attribution#963 changed the root step name, now our script is failing.

Would be great to add it to CI, but I have high hopes that private-attribution#961 is not
too far ahead.

Tested by manually running this script
akoshelev added 4 commits May 5, 2024 21:52
Got a lot of clippy complains about trivial clones when compact gate is enabled. We need to keep the traits at pair with `DescriptiveGate`, otherwise we will have a lot of conditionals
akoshelev added a commit to akoshelev/raw-ipa that referenced this pull request May 13, 2024
We are working towards reducing the total number of steps (private-attribution#1056 is the step in that direction) and compact gate rearchitecture (private-attribution#961) in parallel. The old compact gate gets broken because of that but there is no real reason to fix it.

To avoid red CI, this temporarily disables compact gate checks until private-attribution#961 is merged
akoshelev added a commit that referenced this pull request May 13, 2024
We are working towards reducing the total number of steps (#1056 is the step in that direction) and compact gate rearchitecture (#961) in parallel. The old compact gate gets broken because of that but there is no real reason to fix it.

To avoid red CI, this temporarily disables compact gate checks until #961 is merged
akoshelev added 14 commits May 13, 2024 20:07
We don't need compact gate functionality to check for correctness
Test multiplication does not work properly with the previous implementation of compact gate. W/o narrowing to PRSS, we send one message for key exchange and then try to multiply starting from record_id = 0 which causes infra to complain.

It is not great that PRSS uses someone else's capacity to exchange messages, so this commit fixes that
Some of them were not defined properly and compact gate caught that, others did not have child specified.

Confirmed that it works by running compact gate test

```
cargo test --test compact_gate compact_gate_cap_16_no_window_semi_honest --no-default-features --features "cli web-app real-world-infra test-fixture compact-gate"
   Compiling ipa-core v0.1.0 (/Users/koshelev/workspace/raw-ipa/ipa-core)
    Finished `test` profile [unoptimized + debuginfo] target(s) in 37.25s
     Running tests/compact_gate.rs (target/debug/deps/compact_gate-ccb7a8f65d996521)

running 1 test
test compact_gate_cap_16_no_window_semi_honest ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 3 filtered out; finished in 6.41s
```
///
/// This is a temporary solution for narrowing contexts until the infra is
/// updated with a new step scheme.
pub trait BitStep: Step + From<usize> {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

moved outside of step module because compact gate module does not want anything in it except step definitions @benjaminsavage

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we can parameterize this one further... pub struct NBitStep<const N: u32>. Then we can do away with the bit depth bit...

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Happy to improve this, but let's do it as a follow up, and get this PR merged first.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this requires some changes inside the proc macro to support const generics. We can't use constants in macro attributes, so the way to make it work could be compact step reading the generic parameters and deriving the number of unique steps from it. Not sure if we can make it work, but it is worth experimenting with this

/// we've done for bit steps
#[derive(CompactStep)]
pub(crate) enum SaturatedAdditionStep {
#[step(child = crate::protocol::boolean::step::SixteenBitStep)]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that's the hacky fix for the compact gate @benjaminsavage @martinthomson. This works for now, but if we use integer_sat_add inside some other circuit that requires more than 2 bytes, this will panic

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, this seems like it is something we can work on. Issue?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Filed an issue: #1075

/// gates set, will use [`Unique`] which is the default behavior and benchmarks will use [`Fixed`]
///
/// [`TestWorld`]: super::TestWorld
pub(super) trait TestGateVendor: Send + Sync + 'static {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is done to support benchmarks with compact gate enabled. We can't put PRF IPA root step under Test because that will cause an explosion of steps. Instead, we narrow in the benchark to the correct root step and then create the test world

@@ -79,7 +79,7 @@ check() {
}

check "Benchmark compilation" \
cargo build --benches --no-default-features --features "enable-benches descriptive-gate"
cargo build --benches --no-default-features --features "enable-benches compact-gate"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure if we want to turn off descriptive gate for local checks completely, thoughts @martinthomson ?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm happy leaving this as mutually exclusive. We can run both. As long as we don't have too many #[cfg... bits, we're not risking too much.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is going on here?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good point, I also noticed a lot of empty files in main. Will follow up with a PR to remove it

///
/// This is a temporary solution for narrowing contexts until the infra is
/// updated with a new step scheme.
pub trait BitStep: Step + From<usize> {
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we can parameterize this one further... pub struct NBitStep<const N: u32>. Then we can do away with the bit depth bit...

}

#[cfg(test)]
#[derive(CompactStep)]
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess we need to look into what the codecov warnings mean here. Maybe we need to suppress them.

@@ -150,13 +140,15 @@ pub trait UpgradedContext<F: ExtendableField>: Context {
T: Send,
for<'a> UpgradeContext<'a, Self, F>: UpgradeToMalicious<'a, T, M>,
{
#[cfg(feature = "descriptive-gate")]
#[cfg(descriptive_gate)]
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess we'll have to deal with upgrades later.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we will have very few uses of upgrade. It should only be used in the shuffle protocol.

#[derive(CompactStep)]
pub enum BucketStep {
/// should be equal to `MAX_BREAKDOWNS`
#[step(count = 512, child = crate::protocol::boolean::step::BoolAndStep)]
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
#[step(count = 512, child = crate::protocol::boolean::step::BoolAndStep)]
#[step(count = 512, child = crate::protocol::boolean::step::BoolAndStep), name = "b"]

@benjaminsavage was sad that this said "bit" rather than "bucket", so "bucket" is also fine. I just prefer to use shorter names when possible.

We could also make this struct BucketStep(usize), which could be even shorter.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would like "b" better than "bit" =). But if it's easier to change this in a follow-up PR go for it.

/// gate.
///
/// A few tests and some benchmarks want to execute the actual protocol, but they also care about
/// performance they use compact gate to manage step transitions.
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
/// performance they use compact gate to manage step transitions.
/// performance, so they use compact gate to manage step transitions.

self.gate_vendor.current()
}

fn next_gate(&self) -> Gate {
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
fn next_gate(&self) -> Gate {
#[must_use]
fn next_gate(&self) -> Gate {

@@ -79,7 +79,7 @@ check() {
}

check "Benchmark compilation" \
cargo build --benches --no-default-features --features "enable-benches descriptive-gate"
cargo build --benches --no-default-features --features "enable-benches compact-gate"
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm happy leaving this as mutually exclusive. We can run both. As long as we don't have too many #[cfg... bits, we're not risking too much.


cargo test --bench oneshot_ipa --no-default-features --features "enable-benches descriptive-gate" -- -n 62 -c 16
# descriptive-gate is enabled by default
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
# descriptive-gate is enabled by default
# descriptive-gate is enabled when compact-gate is not


cargo test --bench criterion_arithmetic --no-default-features --features "enable-benches descriptive-gate"
cargo test --bench oneshot_ipa --no-default-features --features "enable-benches" -- -n 62 -c 16
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This might benefit from a comment:

Suggested change
cargo test --bench oneshot_ipa --no-default-features --features "enable-benches" -- -n 62 -c 16
# Test that our benchmark works when compact-gate is disabled
cargo test --bench oneshot_ipa --no-default-features --features "enable-benches" -- -n 62 -c 16

@martinthomson martinthomson marked this pull request as ready for review May 17, 2024 06:31
Copy link
Collaborator

@benjaminsavage benjaminsavage left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will admit I do not really understand the parsing / generating / proc_macro code. But I do understand the step names and annotations. I'm happy for this to be merged.

///
/// This is a temporary solution for narrowing contexts until the infra is
/// updated with a new step scheme.
pub trait BitStep: Step + From<usize> {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Happy to improve this, but let's do it as a follow up, and get this PR merged first.

Comment on lines +4 to +7
pub(crate) enum BoolAndStep {
#[step(count = 8)] // keep in sync with MAX_BITS defined inside and.rs
Bit(usize),
}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Feels like maybe we should just use EightBitStep and eliminate BoolAndStep... Or better yet, make bool_and support a generic BitStep type. Again, happy for this to be a follow up.

@@ -150,13 +140,15 @@ pub trait UpgradedContext<F: ExtendableField>: Context {
T: Send,
for<'a> UpgradeContext<'a, Self, F>: UpgradeToMalicious<'a, T, M>,
{
#[cfg(feature = "descriptive-gate")]
#[cfg(descriptive_gate)]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we will have very few uses of upgrade. It should only be used in the shuffle protocol.

#[derive(CompactStep)]
pub enum BucketStep {
/// should be equal to `MAX_BREAKDOWNS`
#[step(count = 512, child = crate::protocol::boolean::step::BoolAndStep)]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would like "b" better than "bit" =). But if it's easier to change this in a follow-up PR go for it.

Comment on lines +18 to +21
impl From<u32> for BucketStep {
fn from(v: u32) -> Self {
Self::Bit(usize::try_from(v).unwrap())
}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We may not even need this.

#[step(child = UserNthRowStep)]
BinaryValidator,
PrimeFieldValidator,
ModulusConvertBreakdownKeyBitsAndTriggerValues,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is dead code and can be removed.

}

#[derive(CompactStep)]
pub(crate) enum AttributionZeroTriggerStep {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh I see, this is about zeroing-out a trigger event right? Might I suggest naming it something like ZeroOutTriggerStep?

@@ -174,7 +169,7 @@ where

assert!(
K::BITS <= ThirtyTwoBitStep::max_bit_depth(),
"ThirtyTwoBitStep is not large enough to accomodate this sort"
"ThirtyTwoBitStep is not large enough to accommodate this sort"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Whoops... thanks for fixing my spelling errors =P.

Comment on lines +154 to +158
fn prss_gate() -> Gate {
use crate::protocol::step::{ProtocolGate, ProtocolStep};

ProtocolGate::default().narrow(&ProtocolStep::Prss)
}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Kind of strange that there is no test coverage for this.

/// A complex sum, comprised of multiple statements, plus a simple counter.
/// This implements `ToTokens` so that it can be dropped into `quote!()` easily.
#[derive(Default, Clone)]
pub struct ExtendedSum {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where is this used?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is used inside the compact gate module to keep track of number of steps in a tree

@benjaminsavage benjaminsavage merged commit 98fc9ff into private-attribution:main May 18, 2024
12 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants