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

Introduce NoiseAnalysis Framework #1343

Open
wants to merge 6 commits into
base: main
Choose a base branch
from

Conversation

ZenithalHourlyRate
Copy link
Collaborator

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 the Analysis. The focus of this PR is not the detailed function in Noise.cpp/NoiseAnalysis.cpp but the design of NoiseAnalysis.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

$ --mlir-to-secret-arithmetic --secret-insert-mgmt-bgv $PWD/tests/Examples/openfhe/dot_product_8.mlir --debug-only=secret-insert-mgmt-bgv --debug-only=NoiseAnalysis --debug-only=BGVNoise

ringDim: 16384
plaintextModulus: 4295294977
level: 2
logqi: 60 60 60 
dnum: 2
logpi: 60 
Propagating 45.27 to <block argument> of type 'tensor<8xi16>' at index: 0
Propagating 45.27 to <block argument> of type 'tensor<8xi16>' at index: 1
Propagating 98.53 to %1 = arith.muli %input0, %input1 {mgmt.mgmt = #mgmt.mgmt<level = 2, dimension = 3>} : tensor<8xi16>
Propagating 98.53 to %2 = mgmt.relinearize %1 {mgmt.mgmt = #mgmt.mgmt<level = 2>} : tensor<8xi16>
Propagating 98.53 to %3 = tensor_ext.rotate %2, %c4 {mgmt.mgmt = #mgmt.mgmt<level = 2>} : tensor<8xi16>, index
Propagating 99.53 to %4 = arith.addi %2, %3 {mgmt.mgmt = #mgmt.mgmt<level = 2>} : tensor<8xi16>
Propagating 99.53 to %5 = tensor_ext.rotate %4, %c2 {mgmt.mgmt = #mgmt.mgmt<level = 2>} : tensor<8xi16>, index
Propagating 100.53 to %6 = arith.addi %4, %5 {mgmt.mgmt = #mgmt.mgmt<level = 2>} : tensor<8xi16>
Propagating 100.53 to %7 = tensor_ext.rotate %6, %c1 {mgmt.mgmt = #mgmt.mgmt<level = 2>} : tensor<8xi16>, index
Propagating 101.53 to %8 = arith.addi %6, %7 {mgmt.mgmt = #mgmt.mgmt<level = 2>} : tensor<8xi16>
Propagating 41.76 to %9 = mgmt.modreduce %8 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<8xi16>
Propagating 113.76 to %extracted = tensor.extract %9[%c7] {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<8xi16>
Propagating 53.76 to %10 = mgmt.modreduce %extracted {mgmt.mgmt = #mgmt.mgmt<level = 0>} : i16

Real noise dumped

$ ./bazel-out/k8-dbg/bin/tests/Examples/openfhe/dot_product_8_debug_test

( 1 2 3 4 5 6 7 8 ... )
cv 2 Ql 4 logQ: 167 logqi: [ 47 52 52 16 ] budget 123.579 noise: 42.4211
( 2 3 4 5 6 7 8 9 ... )
cv 2 Ql 4 logQ: 167 logqi: [ 47 52 52 16 ] budget 123.69 noise: 42.3096
( 2 6 12 20 30 42 56 72 ... )
cv 3 Ql 3 logQ: 151 logqi: [ 47 52 52 ] budget 66.4859 noise: 83.5141
( 2 6 12 20 30 42 56 72 ... )
cv 2 Ql 3 logQ: 151 logqi: [ 47 52 52 ] budget 66.4859 noise: 83.5141
( 30 42 56 72 2 6 12 20 ... )
cv 2 Ql 3 logQ: 151 logqi: [ 47 52 52 ] budget 66.4859 noise: 83.5141
( 32 48 68 92 32 48 68 92 ... )
cv 2 Ql 3 logQ: 151 logqi: [ 47 52 52 ] budget 66.096 noise: 83.904
( 68 92 32 48 68 92 32 48 ... )
cv 2 Ql 3 logQ: 151 logqi: [ 47 52 52 ] budget 66.096 noise: 83.904
( 100 140 100 140 100 140 100 140 ... )
cv 2 Ql 3 logQ: 151 logqi: [ 47 52 52 ] budget 65.6047 noise: 84.3953
( 140 100 140 100 140 100 140 100 ... )
cv 2 Ql 3 logQ: 151 logqi: [ 47 52 52 ] budget 65.6047 noise: 84.3953
( 240 240 240 240 240 240 240 240 ... )
cv 2 Ql 3 logQ: 151 logqi: [ 47 52 52 ] budget 64.8355 noise: 85.1645
( 240 240 240 240 240 240 240 240 ... )
cv 2 Ql 3 logQ: 151 logqi: [ 47 52 52 ] budget 64.8355 noise: 85.1645
( 0 0 0 0 0 0 0 240 ... )
cv 2 Ql 2 logQ: 99 logqi: [ 47 52 ] budget 26.5106 noise: 71.4894
( 240 ... )
cv 2 Ql 2 logQ: 99 logqi: [ 47 52 ] budget 26.5106 noise: 71.4894
( 240 ... )
cv 2 Ql 2 logQ: 99 logqi: [ 47 52 ] budget 26.5106 noise: 71.4894
( 240 ... )
cv 2 Ql 2 logQ: 99 logqi: [ 47 52 ] budget 26.5106 noise: 71.4894

@j2kun j2kun self-requested a review January 31, 2025 18:15
Copy link
Collaborator

@j2kun j2kun left a 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.

Comment on lines +97 to +98
// for BGV. Should observe the result using --debug-only=NoiseAnalysis
// and --debug-only=secret-insert-mgmt-bgv
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
// 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
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
// the plaintext modulud for BGV
// the plaintext modulus for BGV

namespace mlir {
namespace heir {

class SchemeParam {
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 these classes are fine, but I would prefer if we made two small tweaks:

  1. Name them specific to the scheme they are operating on (if not the scheme + variant).
  2. 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;

Copy link
Collaborator

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>
Copy link
Collaborator

@j2kun j2kun Jan 31, 2025

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>
Copy link
Collaborator

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.

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.

2 participants