Skip to content

Add RNSPolynomial constant folding#3135

Merged
copybara-service[bot] merged 1 commit into
mainfrom
test_938805181
Jun 29, 2026
Merged

Add RNSPolynomial constant folding#3135
copybara-service[bot] merged 1 commit into
mainfrom
test_938805181

Conversation

@copybara-service

Copy link
Copy Markdown
Contributor

Add RNSPolynomial constant folding

This change supports ops:

  • add
  • sub
  • mul
  • ntt
  • intt
  • extract_slice

I explicitly fail if the NTT/iNTT ops don't have an attached primitive root. I think this is correct because if the RNSPolynomialAttrs are in evaluation form already as input to the folder, it is invalid to combine such polynomials. So applying folding here requires all ntt/intt ops that could cause two polynomials in evaluation form to be folded together to use the same primitive root.

@j2kun
j2kun requested a review from crockeea June 26, 2026 23:10
@google google deleted a comment from github-actions Bot Jun 26, 2026

@crockeea crockeea 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.

LGTM. Does this mean that I can now attempt removal of the special cases for constants (and the special case of the PolyToModArith lowering) in PolyMulToNTT?

@j2kun

j2kun commented Jun 29, 2026

Copy link
Copy Markdown
Collaborator

LGTM. Does this mean that I can now attempt removal of the special cases for constants (and the special case of the PolyToModArith lowering) in PolyMulToNTT?

I think I have lost the mental thread on what exactly the blocker was here, but if you try this, you would need to either:

  • Add a -sccp pass at some point in the pipeline, which will try to do constant folding for the whole IR
  • Replace some create calls with OpBuilder::createOrFold
  • Enable folding in the greedy rewrite driver which will try to fold before and after applying patterns.

Since we haven't done a ton of constant folding in HEIR yet, I'm not sure which will be the cleanest for your purposes.

@copybara-service
copybara-service Bot force-pushed the test_938805181 branch 3 times, most recently from ea603c5 to 71a632e Compare June 29, 2026 19:38
This change supports ops:

- add
- sub
- mul
- ntt
- intt
- extract_slice

I explicitly fail if the NTT/iNTT ops don't have an attached primitive root. I think this is correct because if the RNSPolynomialAttrs are in evaluation form already as input to the folder, it is invalid to combine such polynomials. So applying folding here requires all ntt/intt ops that could cause two polynomials in evaluation form to be folded together to use the same primitive root.

PiperOrigin-RevId: 939980273
@copybara-service
copybara-service Bot merged commit b122e4a into main Jun 29, 2026
1 check was pending
@copybara-service
copybara-service Bot deleted the test_938805181 branch June 29, 2026 19:51
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.

2 participants