Skip to content

Add guards to prevent lowering eval-form polynomial constants#3056

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

Add guards to prevent lowering eval-form polynomial constants#3056
copybara-service[bot] merged 1 commit into
google:mainfrom
crockeea:safepolymultontt

Conversation

@crockeea

@crockeea crockeea commented Jun 11, 2026

Copy link
Copy Markdown
Collaborator

PolyMulToNTT can create EVAL-form constants, but the PolynomialToModArith lowering didn't look at the form when lowering the constants, resulting in bogus lowerings (e.g., the constant polynomial f(x)=5 with the EVAL tag would be lowered to the tensor [5, 0, 0, ..., 0]). This PR adds guards to prevent such lowerings. It also adds support for lowering eval-form constants in one special case, namely when the constant is a degree-zero polynomial, because the NTT of a degree-zero polynomial with constant-term c is [c, c, c, c, ...., c].

Note that there are discussions about changing the behavior of PolyMulToNTT so that constants are forced into COEFF form, which forces PolyMulToNTT to insert explicit NTTs, which would subsequently be removed by an sccf pass. Until we get that working, I'm keeping the behavior of PolyMulToNTT on constants unchanged.

I realized that this is more or less what #2902 was trying to do, so I incorporated those changes here. Here's the rundown:

  • PolyMulToNTT allows constants to be emitted in either form, since it is efficient materialize EVAL-form constants rather than performing a runtime NTT. However, we don't have the code to actually do the compile-time NTT, so this resulted in invalid lowerings. That's what the initial PR guarded against.
  • For non-constant Polynomial ops, PolyMulToNTT already enforces the expected input/output forms. For example, the pass requires a COEFF-form polynomial as input to LeadingTermOp. add defensive guards for eval/coeff forms #2902 added guards for those ops as well. This is the safe thing to do, although the current code never emitted EVAL-form inputs to these ops.
  • ToTensorOp and FromTensorOp are a bit special: if you know what you're doing, there's nothing wrong with extracting the NTT coefficients of a polynomial. In fact, we [need to] do so in the NTT tests. The PolyMulToNTT pass requires COEFF-form inputs/outputs for these ops, but I do not enforce that in the PolynomialToModArith lowering since we need EVAL-form inputs/outputs for the tests.

The implementation of the degree-zero NTT was done by AI.

@crockeea crockeea requested a review from j2kun June 11, 2026 16:57
@crockeea

Copy link
Copy Markdown
Collaborator Author

Actually, I see that this is more or less precisely what #2902 was trying to solve. That PR appears abandoned, so I will add the extra guards from that PR to this one.

@crockeea crockeea marked this pull request as draft June 11, 2026 16:59
@crockeea crockeea marked this pull request as ready for review June 11, 2026 17:21
@crockeea

Copy link
Copy Markdown
Collaborator Author

Ready for review.

Comment thread lib/Dialect/Polynomial/Conversions/PolynomialToModArith/PolynomialToModArith.cpp Outdated
@crockeea

Copy link
Copy Markdown
Collaborator Author

Applied suggested changes and squashed.

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