Skip to content

Removed mod_arith::ExtractOp#3094

Merged
copybara-service[bot] merged 1 commit into
google:mainfrom
crockeea:liftedBasisConversion
Jun 17, 2026
Merged

Removed mod_arith::ExtractOp#3094
copybara-service[bot] merged 1 commit into
google:mainfrom
crockeea:liftedBasisConversion

Conversation

@crockeea

Copy link
Copy Markdown
Collaborator

Removed mod_arith::ExtractOp in favor of the new mod_arith::LiftOp. This turned out to be more work than anticipated.

The expected changes:

  • Remove references in ModArithOps.td/.cpp
  • Change basis conversion to use LiftOp instead of ExtractOp
  • Remove ConvertExtract from ModArithToArith
  • Remove other occurrences of ExtractOp:
    • in SecretToModArith, I preserved functionality by lifting to the standard representative
  • Many tests were updated to use lift instead of extract
  • Updated lower_convert_basis tests, in particular the runner now expects centered representatives (converted back to standard since that's the internal representation used in ModArith).

Tangentially related changes:

  • Make inner loop of computeMixedRadixCoeffs and lowering for ConvertBasisOp use scf:ForOp loops instead of C++ to reduce IR size/compile time. Note that in both cases, types change in the outer loop, so the outer loop cannot be implemented as an MLIR loop.
  • I did a bit of cleanup and refactoring in the ConvertBasis lowering
  • I realized the lowering for ConvertMonomial for EVAL form polynomials was unnecessarily convoluted (and used mod_arith::ExtractOp); it is now much simpler.
  • There is no need to explicitly reduce before calling LiftOp; it is implicit.

Unexpected changes:

  • ExtractOp is not exactly a subset of LiftOp. For ModArith, you could view it that way (except that ExtractOp might not have output a canonical representative, while LiftOp always does). For RNS, they have different meanings: ExtractOp would output the tensor of extracted ModArith coeffs. LiftOp doesn't support RNS for now, but if it did, it would do a CRT reconstruction on the coefficients and output a scalar. This affected a couple of tests, but mostly the lowering of polynomial.mul_scalar, which does some fancy unpacking and repacking of the coefficients. Since LiftOp doesn't support RNS, I added a special path when multiplying a polynomial by an RNS scalar that avoids it. For the non-RNS case, I changed mod_arith::ExtractOp to mod_arith::LiftOp standard; the representative doesn't matter since we immediately re-encapsulate. The representative only matters in "mathematical" code; there the "lifting" is for the compiler's benefit. The same holds in a few other places where LiftOp replaced ExtractOp.
  • Convertion utils: Written by AI, but reviewed by me.
    • ConvertFromElements was tightened a bit: the main change is that the pattern fails if the input's element type doesn't lower to a tensor. (Note that with the new scf loops in basis conversion, we have to make tensors of ints and ModAriths; I don't think those had ever been tested before).
    • ConvertSplat: Add conversion for splatting a type that lowers to a tensor (e.g., RNS). The use of splat in mul_scalar necessitated this change.

Misc changes:

  • Remove deprecated syntax in LinalgCanonicalizations.cpp

Tests:

  • added runner tests for liftOp in the process of debugging basis conversion. It tests even and odd moduli with both standard and centered representatives.
  • Removed test for invalid mod_arith.extract on RNS, since there is no analog anymore

@j2kun j2kun left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Looks great!

@j2kun j2kun added the pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing label Jun 16, 2026
@copybara-service copybara-service Bot merged commit 689e074 into google:main Jun 17, 2026
19 checks passed
@crockeea crockeea deleted the liftedBasisConversion branch June 17, 2026 22:11
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