Skip to content

Add extended basis rotation passes for OpenFHE (ConvertToExtendedBasis + HoistKeySwitchDown)#2624

Merged
copybara-service[bot] merged 1 commit intogoogle:mainfrom
toprakmurat:feat/openfhe-fast-rotation-ext
Mar 6, 2026
Merged

Add extended basis rotation passes for OpenFHE (ConvertToExtendedBasis + HoistKeySwitchDown)#2624
copybara-service[bot] merged 1 commit intogoogle:mainfrom
toprakmurat:feat/openfhe-fast-rotation-ext

Conversation

@toprakmurat
Copy link
Copy Markdown
Contributor

@toprakmurat toprakmurat commented Feb 8, 2026

Summary

Continuation of #2574, related to #2519.
Adds two new passes for optimizing rotation-heavy FHE programs by deferring expensive key-switch operations:

  1. --openfhe-convert-to-extended-basis - Converts fast_rotationfast_rotation_ext + key_switch_down
  2. --openfhe-hoist-key-switch-down

Changes

Pass 1: ConvertToExtendedBasis

  • Converts each openfhe.fast_rotation to openfhe.fast_rotation_ext + openfhe.key_switch_down
  • Semantically equivalent transformation that enables further optimization

Pass 2: HoistKeySwitchDown

  • Pushes key_switch_down operations later in the IR to reduce key-switch count.
  • Works on Add, AddInPlace, and Mul operations.
  • When at least one operand is key_switch_down (with single use), the pass swaps:
// Case 1: Both operands are key_switch_down
%a = key_switch_down(%a_ext)
%b = key_switch_down(%b_ext)
%sum = add(%a, %b)

to:

%sum_ext = add(%a_ext, %b_ext)
%sum = key_switch_down(%sum_ext)

and,

// Case 2: One operand is key_switch_down (enables chained optimization)
%a = key_switch_down(%a_ext)
%sum = add(%a, %ct)  // %ct is not key_switch_down

to:

%sum_ext = add(%a_ext, %ct)
%sum = key_switch_down(%sum_ext)

Through greedy iteration, this reduces multiple key-switches in chained operations to a single final key-switch.

@toprakmurat toprakmurat marked this pull request as ready for review February 9, 2026 17:00
@toprakmurat toprakmurat force-pushed the feat/openfhe-fast-rotation-ext branch from a9ac151 to a9a1ecf Compare February 14, 2026 15:51
Comment thread lib/Dialect/Openfhe/Transforms/ConvertToExtendedBasis.cpp Outdated
Comment thread lib/Dialect/Openfhe/Transforms/HoistKeySwitchDown.cpp
@toprakmurat
Copy link
Copy Markdown
Contributor Author

Hi @asraa, I'd appreciate your review of these changes when you have time. Should I also add interpreter benchmarks comparing performance with and without these passes?

@j2kun
Copy link
Copy Markdown
Collaborator

j2kun commented Feb 18, 2026

@AlexanderViand Asra and I are a bit strapped for time these days. Do you think you'd have time to review this PR?

@j2kun j2kun self-requested a review February 23, 2026 18:36
Copy link
Copy Markdown
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.

LGTM! Some minor changes suggested. In my view the next step would be to add these two passes to the openfhe pipeline, and run them with something like a matvec op that lowers to an IR with lots of rotations and summations, which should be able to take advantage of these patterns.

Comment thread lib/Dialect/Openfhe/Transforms/HoistKeySwitchDown.cpp Outdated
Comment thread lib/Dialect/Openfhe/Transforms/ConvertToExtendedBasis.cpp Outdated
Comment thread lib/Dialect/Openfhe/Transforms/ConvertToExtendedBasis.cpp Outdated
@j2kun
Copy link
Copy Markdown
Collaborator

j2kun commented Feb 23, 2026

Also when you're ready, please squash your commits into one, as per https://heir.dev/docs/development/#pull-request-review-flow

@toprakmurat toprakmurat force-pushed the feat/openfhe-fast-rotation-ext branch from a9a1ecf to 561d04a Compare March 2, 2026 15:47
@toprakmurat toprakmurat requested a review from j2kun March 2, 2026 15:48
@j2kun j2kun added the pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing label Mar 5, 2026
@j2kun
Copy link
Copy Markdown
Collaborator

j2kun commented Mar 5, 2026

Thank you @toprakmurat!

@copybara-service copybara-service Bot merged commit ad5751d into google:main Mar 6, 2026
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