Implementation of loop support in PolyMulToNTT - #3295
Conversation
ae099e6 to
04feed0
Compare
43b5865 to
37803c5
Compare
| module { | ||
| // Covers: an scf.for iter_arg that needs both forms -- coeff to match its | ||
| // entry operand/loop result (consumed by to_tensor, coeff-only) and eval to | ||
| // feed the eval-only MulOp in the loop body. The NTT/INTT pair is hoisted |
There was a problem hiding this comment.
Wouldn't it be more optimal to apply NTT once before the loop even begins, and then INTT once after the loop ends, rather than in each loop iteration?
| // CHECK: [[t2:%.+]] = polynomial.to_tensor [[x2]] : [[poly_ty_1]] -> tensor<1024x[[RNS]]> | ||
| // CHECK: [[m2:%.+]] = polynomial.mul [[r2]], [[r2]] : [[ntt_poly_ty_1]] | ||
| // CHECK: return [[t2]], [[m2]] : tensor<1024x[[RNS]]>, [[ntt_poly_ty_1]] | ||
| func.func @if_yields_both_branches(%cond: i1, %x: !poly_ty_1) -> (tensor<1024x!rns.rns<!Zq0>>, !poly_ty_1) { |
There was a problem hiding this comment.
I think it would be useful to have a test case where an if statement requires two different forms in the two different branches.
| // forwarding operands that feed it to match. | ||
| llvm::DenseMap<Value, Form> successorNativeForm; | ||
| for (Value target : polySuccessorInputs) { | ||
| Form f = soln.needsForm(target, Form::COEFF) ? Form::COEFF : Form::EVAL; |
There was a problem hiding this comment.
I feel like this is what's making the for_iter_arg_needs_both_forms test feel incorrect to me.
In my view, it's not that the successor input must have COEFF form if it has uses of both in different branches, but that all region successor operands and all corresponding successor inputs must have the same form (you can't change form during a control flow branch), and then allow the solver to choose which form it is based on the implied cost of changing form before/after the branch.
In the for_iter_arg_needs_both_forms test, it would require x, acc, sq and r to all have the same form. If they're all coeff, the ntt/intt would be in the loop and have cost given by the lb/ub (which I guess is the default 1k iterations in that example?), and if they're all eval, then it would have to insert conversions for x and for r.
|
Your review exposed some major gaps (thank you!) which will take some time to fix. I'll let you know when I'm ready again. |
|
Given that the gaps are optimizations, we could merge what you have and follow up with the improvements in a second PR. Up to you. |
2 participants
Fixes #2685
This PR adds support for
RegionBranchOpInterfaceto PolyMulToNTT, enabling NTT optimization inside things likeforloops,ifstatements, andwhileloops.Assisted by AI.