[do-not-merge] providing cyclic layout support to enable Transformer evaluation - #3388
Open
copybara-service[bot] wants to merge 1 commit into
Open
copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
copybara-service
Bot
force-pushed
the
test_974067900
branch
from
September 4, 2026 07:35
11fbdd6 to
30ff4d2
Compare
copybara-service
Bot
force-pushed
the
test_974067900
branch
from
September 4, 2026 21:16
30ff4d2 to
87abf55
Compare
copybara-service
Bot
force-pushed
the
test_974067900
branch
2 times, most recently
from
September 5, 2026 00:47
a326849 to
2b47503
Compare
…evaluation
Introduce end-to-end cyclic (bicyclic and tricyclic) layout support and
compiler optimizations required for evaluating Transformer models (e.g.
BERT-tiny SST-2) at 32k slots:
- In-layout strided reduce:
- Support in-layout strided reductions along CRT-based cyclic layouts in
ConvertToCiphertextSemantics without requiring physical layout conversion.
- Lower reductions via binary span-doubling rotate-and-reduce operations
followed by periodic replication of the valid prefix across tail slots.
- Extend LayoutPropagation for linalg.reduce on cyclic layouts and upgrade
partial target relations to canonical total congruences in elementwise
rectification.
- Closed-form constant layout packing:
- Add closed-form diagonal packing routines for rank-2 and rank-3
cleartext weights (packBicyclicDiagonalClosedForm,
packTricyclicDiagonalClosedForm, packDiagonalWeightClosedForm).
- Bypass expensive polyhedral integer relation evaluation (Presburger /
ISL) in AssignLayout, enabling instant constant packing for 32k/64k
slot sizes.
- Add fast paths in AssignLayout for splat constants and zero-filled pad
buffers, and raise constant folding threshold to 65536.
- Shift network & rotation group optimizations:
- Optimize conflict graph construction in ImplementShiftNetwork by grouping
sources by slot position to eliminate redundant quadratic conflict edges.
- Skip graph coloring entirely when no position collisions occur, safely
defaulting unconflicted vertices to rotation group 0.
- Pipeline controls for large ciphertext sizes:
- Introduce CommonPipelineOptions with --skip-layout-optimization,
--layout-optimization-vve-tries, and --codegen-strategy flags wired
through MlirToRLWE and PlaintextBackend pipelines.
- Bypass expensive hour-scale Vos-Vos-Erkin (VVE) layout optimization cost
evaluations at 32,768 slots where CRT-closed layouts are already optimal.
- Testing:
- Add UtilsTest suites validating mathematical equivalence between
closed-form diagonal packing formulas and Presburger relation evaluation.
- Add lit tests for layout propagation and ciphertext semantics lowering of
multidimensional and strided reductions.
PiperOrigin-RevId: 974067900
copybara-service
Bot
force-pushed
the
test_974067900
branch
from
September 5, 2026 01:36
2b47503 to
bd2802e
Compare
copybara-service Bot
pushed a commit
that referenced
this pull request
Sep 5, 2026
…ipeline Introduce the LowerPadToIdentityMatmul pass and pipeline controls to lower zero-padding operations on secret tensors into cleartext identity matrix multiplications, avoiding irregular shift networks during layout conversion. This change is stacked on top of all experimental pending commits: - cl/974067900 (or GitHub PR #3388): cyclic layout support for Transformers - cl/976505644 (or GitHub PR #3400): Lattigo backend support and CKKS fixes - cl/976542552 (or GitHub PR #3402): approximate non-linearities and management - LowerPadToIdentityMatmul pass: - Rewrite static zero tensor.pad ops on secret tensors to explicit linalg.matmul or linalg.batch_matmul with cleartext identity diagonal matrices (P_L * X for row padding, X * P_R for column padding). - Handle 2D tensors ([M, K] -> [M', K']) and 3D batched tensors ([H, M, K] -> [H, M', K']) across row padding, column padding, and simultaneous 2D padding via sequential Left and Right PCMM. - Leverage structured BSGS rotate-and-reduce schedules (2*sqrt(n) Galois keys) of existing Left and Right PCMM kernels, completely avoiding irregular shift networks and eliminating hundreds of Galois rotation keys. - CKKS pipeline integration: - Introduce --enable-identity-pcmm-padding in ArithmeticPipelineRegistration to wire LowerPadToIdentityMatmul into the mlirToRLWE CKKS pipeline prior to layout propagation. - Follow identity PCMM lowering with canonicalization and CSE to fold intermediate structures. - Register createLowerPadToIdentityMatmul pass in heir-opt. - Waterline analysis & forced bootstrapping: - Support mgmt.force_bootstrap attribute in BootstrapWaterlineAnalysis to force bootstrap placement on critical join operations where level consumption must be reset. PiperOrigin-RevId: 976603963
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[do-not-merge] providing cyclic layout support to enable Transformer evaluation
Introduce end-to-end cyclic (bicyclic and tricyclic) layout support and
compiler optimizations required for evaluating Transformer models (e.g.
BERT-tiny SST-2) at 32k slots:
In-layout strided reduce:
ConvertToCiphertextSemantics without requiring physical layout conversion.
followed by periodic replication of the valid prefix across tail slots.
partial target relations to canonical total congruences in elementwise
rectification.
Closed-form constant layout packing:
cleartext weights (packBicyclicDiagonalClosedForm,
packTricyclicDiagonalClosedForm, packDiagonalWeightClosedForm).
ISL) in AssignLayout, enabling instant constant packing for 32k/64k
slot sizes.
buffers, and raise constant folding threshold to 65536.
Shift network & rotation group optimizations:
sources by slot position to eliminate redundant quadratic conflict edges.
defaulting unconflicted vertices to rotation group 0.
Pipeline controls for large ciphertext sizes:
--layout-optimization-vve-tries, and --codegen-strategy flags wired
through MlirToRLWE and PlaintextBackend pipelines.
evaluations at 32,768 slots where CRT-closed layouts are already optimal.
Testing:
closed-form diagonal packing formulas and Presburger relation evaluation.
multidimensional and strided reductions.