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.

Sign up for GitHub

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

Jump to bottom

Re-implement multi-parameter parametrization pass using type inference #674

Merged
merged 10 commits into from
Feb 2, 2024
Merged

Re-implement multi-parameter parametrization pass using type inference #674

merged 10 commits into from
Feb 2, 2024

Conversation

andidr
Copy link
Contributor

@andidr andidr commented Jan 22, 2024

No description provided.

@cla-bot cla-bot bot added the cla-signed label Jan 22, 2024
@andidr andidr force-pushed the andi/type-inference branch 7 times, most recently from 8b7e103 to 55e534a Compare January 24, 2024 14:21
aPere3
aPere3 approved these changes Jan 25, 2024
View reviewed changes
Copy link
Contributor

@aPere3 aPere3 left a comment

Choose a reason for hiding this comment

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

Thanks for the good work Andi. I really like how you implemented the TypeResolver, and I learned quite a few things along the way. I made a few comments here and there, but there is nothing major blocking. Best.

compilers/concrete-compiler/compiler/lib/Support/TFHECircuitKeys.cpp Show resolved Hide resolved
...oncrete-compiler/compiler/lib/Dialect/TFHE/Transforms/TFHECircuitSolutionParametrization.cpp Outdated Show resolved Hide resolved
...oncrete-compiler/compiler/lib/Dialect/TFHE/Transforms/TFHECircuitSolutionParametrization.cpp Show resolved Hide resolved
compilers/concrete-compiler/compiler/include/concretelang/Analysis/TypeInferenceAnalysis.h Show resolved Hide resolved
...oncrete-compiler/compiler/lib/Dialect/TFHE/Transforms/TFHECircuitSolutionParametrization.cpp Outdated Show resolved Hide resolved
compilers/concrete-compiler/compiler/include/concretelang/Analysis/TypeInferenceAnalysis.h Outdated Show resolved Hide resolved
compilers/concrete-compiler/compiler/include/concretelang/Analysis/TypeInferenceAnalysis.h Outdated Show resolved Hide resolved
compilers/concrete-compiler/compiler/include/concretelang/Analysis/TypeInferenceAnalysis.h Show resolved Hide resolved
compilers/concrete-compiler/compiler/include/concretelang/Transforms/TypeInferenceRewriter.h Outdated Show resolved Hide resolved
compilers/concrete-compiler/compiler/include/concretelang/Transforms/TypeInferenceRewriter.h Show resolved Hide resolved
aPere3
aPere3 approved these changes Jan 25, 2024
View reviewed changes
Copy link
Contributor

@aPere3 aPere3 left a comment

Choose a reason for hiding this comment

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

Thanks for the good work Andi. I really like how you implemented the TypeResolver, and I learned quite a few things along the way. I made a few comments here and there, but there is nothing major blocking. Best.

@andidr andidr force-pushed the andi/type-inference branch 2 times, most recently from f6c9b05 to e2faffb Compare January 26, 2024 16:19
@andidr
Copy link
Contributor Author

andidr commented Jan 26, 2024

Thanks a lot @aPere3 and @antoniupop for the in-depth review of the PR.

@aPere3
Copy link
Contributor

aPere3 commented Jan 29, 2024

Agreed. However, I fixed the minimal amount of things in that pass to make all tests pass. Maybe this should be a followup commit by @aPere3?

I can look into it if needed yes ! I'll have to focus on composition for a few weeks I guess, but when I am done, I can add the remaining patterns. Can you create an issue and tag me ?

@andidr
Copy link
Contributor Author

andidr commented Jan 29, 2024 •
edited

Can you create an issue and tag me ?

Done: zama-ai/concrete-internal#607 and zama-ai/concrete-internal#606 (I have also created an issue for the key gathering routines).

@andidr andidr force-pushed the andi/type-inference branch 10 times, most recently from bebb6ed to 09ebf06 Compare February 1, 2024 05:53
@andidr
fix(compiler): Fix typos in the summaries of batched TFHE ops
8b1c6ab
@andidr
chore(compiler): Bump MLIR fork to version with fix for backward data…
9ca9369 
…flow analysis
@andidr
feat(compiler): Add a dialect for type inference debugging
e78883c 
The `TypeInference` dialect provides three operations.

The operation `TypeInference.propagate_downwards` respresents a type
barrier, which is supposed to forward the type of its operand as its
result type during type inference.

The operation `TypeInference.propagate_upwards` also respresents a
type barrier, but is supposed to forward the type of its result as the
type for its operand during type inference.

The operation `TypeInference.unresolved_conflict` can be used as a
marker when two different types have beed inferred for a value (e.g.,
one type during forward dataflow analysis and the other during
backward dataflow analysis)
@andidr
feat(compiler): Add infrastructure for type inference
43180a7 
Type inference is implemented through the two classes
`ForwardTypeInferenceAnalysis` and `BackwardTypeInferenceAnalysis`,
which can be used as forward and backward dataflow analyses with the
MLIR sparse dataflow analysis framework.

Both classes rely on a type resolver, which must be a class inheriting
`TypeResolver` and that specifies which types are to be considered as
unresolved and that resolves the actual types for the values related
to an operation based on the previous state of type inference.

The type inference state for an operation is represented by an
instance of the class `LocalInferenceState`, which maps the values
related to an operation to instances of `InferredType` (either
indicating the inferred type as an `mlir::Type` or indicating that no
type has been inferred, yet).

The local type inference by a type resolver can be implemented with
type constraints (instances of sub-classes of `TypeConstraint`), which
can be combined into a `TypeConstraintSet`. The latter provides a
function that attempts to apply the constraints until the resulting
type inference state converges.

There are multiple, predefined type constraint classes for common
constraints (e.g., if two values must have the same type or the same
element type). These exist both as static constraints and as dynamic
constraints. Some pre-defined type constraints depend on a class that
yields a pair of values for which the contraints shall be applied
(e.g., yielding two operands or an operand and a result, etc.).
@andidr
feat(compiler): Add rewriter applying the results from type inference
1faeb2d 
This adds the `TypeInferenceRewriter` class, which applies the results
for type inference obtained from the state of a `DataFlowSolver` and
from a final invocation of a type resolver to a module.
@andidr
feat(compiler): Add support for batched keyswitch ops in TFHE key nor…
79d8337 
…malization
@andidr
test(compiler): Run canonicalizer in TFHE parametrization tests
bc171dd
@andidr
fix(compiler): Run canonicalizer after TFHE circuit parametrization
775171c 
In order to avoid leftover TFHE operations to be lowered further down
the pipeline after parametrization, run the canonicalizer, which
includes dead code eliminiation.
@andidr
refactor(compiler): Re-implement TFHE multi-parameter parametrization…
21a7eea 
… with type inference

The current pass applying the parameters determined by the optimizer
to the IR propagates the parametrized TFHE types to operations not
directly tagged with an optimizer ID only under certain conditions. In
particular, it does not always properly propagate types into nested
regions (e.g., of `scf.for` loops).

This burdens preceding transformations that are applied in between the
invocation of the optimizer and the parametrization pass with
data-flow analysis and book-keeping in order to tag newly inserted
operations with the right optimizer IDs that ensure proper
parametrization.

This commit replaces the current parametrization pass with a new pass
that propagates parametrized TFHE types up and down def-use chains
using type inference and a proper rewriter. The pass is limited to the
operations supported by `TFHEParametrizationTypeResolver::resolve`.
@andidr
test(compiler): Add check tests for type inference
972eb96 
Add various check tests for type inference testing forward and
backward propagation, including propagation into nested regions.
@andidr andidr merged commit 972eb96 into main Feb 2, 2024
27 checks passed
@andidr andidr deleted the andi/type-inference branch February 2, 2024 05:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@aPere3 aPere3 aPere3 approved these changes

@antoniupop antoniupop antoniupop approved these changes

@BourgerieQuentin BourgerieQuentin Awaiting requested review from BourgerieQuentin

Assignees
No one assigned
Labels
cla-signed
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
3 participants
@andidr @aPere3 @antoniupop