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

feat: create a tf2fhe pipeline #267

Merged
merged 1 commit into from
Feb 1, 2024

Conversation

asraa
Copy link
Collaborator

@asraa asraa commented Nov 17, 2023

Blocked by #268

@j2kun
Copy link
Collaborator

j2kun commented Dec 3, 2023

Putting some notes here as I think about them:

  • Start with a tosa model, where the model is a single function and all inputs are private.
  • Run secretize to annotate the function args to be converted to secret
  • Run wrap-generic to wrap the body of the annotated function in a secret.generic
  • Run distribute-generic with option distribute-through=affine.for to distribute generic down to the bodies of loop nests.
  • Run yosys-optimizer to optimize the individual generic bodies, resulting in an IR composed of secret.generic blocks (possibly inside loop nests) with booleanized comb ops inside and some secret.cast ops around each generic.
  • Run comb-to-cggi which converts the comb ops and also handles the remaining removal of secret.generic, secret.cast, and secret.secret types.
  • Run cggi-to-tfhe-rust and emit-tfhe-rust

@asraa
Copy link
Collaborator Author

asraa commented Dec 4, 2023

Run yosys-optimizer to optimize the individual generic bodies, resulting in an IR composed of secret.generic blocks (possibly inside loop nests) with booleanized comb ops inside and some secret.cast ops around each generic.

After #308 the for loop bodies will contain affine.load / affine.store - which translateToVerilog will be able to handle - then after YosysOptimizer booleanizes the inputs, we'll be converting secret<memref<?xi8>> multi-bit types to secret<memref<tensor<?xi8>>> types.

Question: At what point is more accurate to call these memref<?xsecret<i8>>...

@j2kun
Copy link
Collaborator

j2kun commented Dec 4, 2023

Question: At what point is more accurate to call these memref<?xsecret<i8>>...

I think the types will be, in order of passes

  • secret<memref<?xi8>>
  • secret<memref<?xtensor<8xi1>>

It sounds like we'd want another bufferization pass... but that seems like it won't work if it starts from a mixed tensor/memref IR. I wonder if it makes sense to put bufferization later in the pipeline, after Yosys optimization? Then the types are tensor<?xtensor<8xi1>> and it should be easier to collapse that into a single multidimensional tensor.

tools/heir-opt.cpp Outdated Show resolved Hide resolved
tools/heir-opt.cpp Outdated Show resolved Hide resolved
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.

Some nits, otherwise LGTM

tools/heir-opt.cpp Outdated Show resolved Hide resolved
@asraa asraa force-pushed the add-tf2fhe-pipeline branch 2 times, most recently from b7ae74d to d274d1c Compare January 24, 2024 19:13
@asraa asraa force-pushed the add-tf2fhe-pipeline branch 2 times, most recently from 44ab842 to 985baf5 Compare January 31, 2024 16:58
@asraa asraa marked this pull request as ready for review January 31, 2024 16:59
@asraa asraa force-pushed the add-tf2fhe-pipeline branch from 985baf5 to 24fc2a7 Compare January 31, 2024 17:00
@asraa asraa changed the title draft: create a tf2fhe pipeline feat: create a tf2fhe pipeline Jan 31, 2024
@asraa
Copy link
Collaborator Author

asraa commented Jan 31, 2024

@j2kun pipeline cleaned up!

@asraa asraa force-pushed the add-tf2fhe-pipeline branch from 24fc2a7 to f1d4524 Compare January 31, 2024 18:08
@asraa asraa force-pushed the add-tf2fhe-pipeline branch from f1d4524 to ae59267 Compare January 31, 2024 19:43
@asraa asraa added the pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing label Jan 31, 2024
Signed-off-by: Asra <asraa@google.com>
@asraa asraa force-pushed the add-tf2fhe-pipeline branch from ae59267 to ef46fc8 Compare January 31, 2024 19:52
@copybara-service copybara-service bot merged commit 812ff19 into google:main Feb 1, 2024
6 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants