-
Notifications
You must be signed in to change notification settings - Fork 61
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
Introduce NoiseAnalysis Framework #1343
base: main
Are you sure you want to change the base?
Introduce NoiseAnalysis Framework #1343
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added some comments on the structure/organization.
Having a noise analysis that assumes known scheme parameters is the right way to start. Even if we have a sophisticated parameter selection optimization, we will want to regularly verify noise bounds are still respected after various program transformations.
To that end, let's test this by creating a utility function validateNoise
that accepts as input a top-level op, known parameters, and a noise model, and returns a LogicalResult for success/failure if the noise stays within correctness bounds.
Then the right way to test this PR would be to wrap the validateNoise
in a trivial pass validate-noise
and add test IRs that pass and fail the noise bounds.
// for BGV. Should observe the result using --debug-only=NoiseAnalysis | ||
// and --debug-only=secret-insert-mgmt-bgv |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
// for BGV. Should observe the result using --debug-only=NoiseAnalysis | |
// and --debug-only=secret-insert-mgmt-bgv | |
// for BGV. Should observe the result using --debug-only=NoiseAnalysis,secret-insert-mgmt-bgv |
FYI these flags accept a comma-separated list of debug-type ids.
// the N in Z[X]/(X^N+1) | ||
int ringDim; | ||
|
||
// the plaintext modulud for BGV |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
// the plaintext modulud for BGV | |
// the plaintext modulus for BGV |
namespace mlir { | ||
namespace heir { | ||
|
||
class SchemeParam { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think these classes are fine, but I would prefer if we made two small tweaks:
- Name them specific to the scheme they are operating on (if not the scheme + variant).
- Put them in a new directory, maybe
lib/Parameters
.
The reason for (2) is that I suspect we will wrap these classes in IR attributes with custom parser/printers so we can attach them to the IR later. So having them in an isolated place allows us to avoid cyclic dependencies more easily.
LogicalResult visitOperation( | ||
Operation *op, ArrayRef<const NoiseLattice<Noise> *> operands, | ||
ArrayRef<NoiseLattice<Noise> *> results) override; | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You may also want to override visitExternalCall if you are passing things to a debug function, otherwise the default behavior is to reset the lattice to uninitialized.
namespace bgv { | ||
|
||
// is worst case or not | ||
template <bool W = false> |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think a template with a boolean option is not the right modeling construct here. First, most of the formulas above don't depend on it, while the one that does switches on a constexpr (presumably because you can't add the bool as member data to the Noise
constructor without mucking up the lattice). Second, more advanced noise models would likely have completely different APIs; e.g., they might track a lower and upper bound for a given probability. We don't want to leak that throughout the rest of the program, an in this case it's leaking through the template instantiation.
I can think of a few possible improvements.
(1) Add a base class to represent the common parts of the two noise models given here, with virtual methods for the parts that differ. Since the parts that differ are only used internally, you can further simplify the analysis pass below to accept a NoiseModel<NoiseBase>
and avoid the templating.
(2) Separate the noise model from the noise data. Since the lattice itself only needs "join", the visit
helpers could instead live in a NoiseModel
class with a more well-defined interface (e.g., propagateThroughOp
as the front-facing function, which takes as input the operand noises and produces result noises), and then the NoiseAnalysis
can avoid hard-coding anything about how the noise model works internally (in particular, which ops need type switches) while two noise models that both use "single scalar with a max join" can reuse the same Noise
struct. In this case you might still need a Noise
base class or template, if two different model classes use two different kinds of noise, but after this refactoring that would be relatively straightforward and could be deferred until the point where we need it.
I like (2) better because I can see how it will be more future-proof: if we have a noise model that uses an upper bound on error probability, that upper bound would be data given to the noise model, and not live in the noise struct, and would only be visible during the noise model construction time.
Open to any further ideas, or I can sketch out the class structure for (2) if you would like more details.
namespace heir { | ||
namespace bgv { | ||
|
||
template <bool W> |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please add a code comment linking to the source for these formulas. Either a paper or code reference.
Before incorporating parameter selection procedures, it is vital to support various forms of noise analyse in the HE circuit, and the analysis itself worths design discussion.
Although there are multiple papers discussing noise analysis (i.e. worst-case v.s. average-case, coefficient embedding v.s. canonical embedding, implementation-agnostic v.s. implementation-specific), I think we can model them using MLIR Dataflow framework with different
LatticeState
and re-use theAnalysis
. The focus of this PR is not the detailed function inNoise.cpp/NoiseAnalysis.cpp
but the design ofNoiseAnalysis.h/Noise.h/Params.h
.The tricky part of parameter selection and noise analysis is that, noise analysis requires concrete parameter, yet parameter should be selected according to the analysis result, forming a circular dependency.
To break this dependency, we can use an iterative approach by first giving a conservative param and later passes can give optimized one. Previous papers would use optimizer to solve it at once, but for us that would be left for future PR.
The
BGV/Noise.cpp
uses formulas from KPZ21, and there are various assumptions made to the circuit, but that should not be critical and we could revisit it when we incorporate the param selection procedure and have tested various programs.I could not think of a test for such analysis. The test should be considered when param selection is in.
Example
Analysed result
Real noise dumped