Skip to content

[do-not-merge] providing cyclic layout support to enable Transformer evaluation - #3388

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_974067900
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_974067900

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Aug 31, 2026

Copy link
Copy Markdown
Contributor

[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:

    • 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.

@copybara-service copybara-service Bot changed the title wip: add closed-form constant packing and compile-time controls for 32k slots wip: [do-not-merge] add closed-form constant packing and compile-time controls for 32k slots Sep 4, 2026
@copybara-service copybara-service Bot changed the title wip: [do-not-merge] add closed-form constant packing and compile-time controls for 32k slots [do-not-merge] providing cyclic layout support to enable Transformer evaluation Sep 4, 2026
@copybara-service
copybara-service Bot force-pushed the test_974067900 branch 2 times, most recently from a326849 to 2b47503 Compare September 5, 2026 00:47
…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 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
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.

1 participant