diff --git a/lib/Dialect/Polynomial/Transforms/BUILD b/lib/Dialect/Polynomial/Transforms/BUILD index 01430a9f90..8e437347b4 100644 --- a/lib/Dialect/Polynomial/Transforms/BUILD +++ b/lib/Dialect/Polynomial/Transforms/BUILD @@ -59,8 +59,10 @@ cc_library( "@heir//lib/Dialect/ModArith/IR:TypeInterfaces", "@heir//lib/Dialect/Polynomial/IR:Dialect", "@llvm-project//llvm:Support", + "@llvm-project//mlir:ControlFlowInterfaces", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", + "@llvm-project//mlir:LoopLikeInterface", "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", "@llvm-project//mlir:TensorDialect", diff --git a/lib/Dialect/Polynomial/Transforms/NTTSolver.cpp b/lib/Dialect/Polynomial/Transforms/NTTSolver.cpp index 8da5250ede..f401817b4b 100644 --- a/lib/Dialect/Polynomial/Transforms/NTTSolver.cpp +++ b/lib/Dialect/Polynomial/Transforms/NTTSolver.cpp @@ -68,25 +68,29 @@ NTTSolver::RepVars& NTTSolver::getOrCreateVars(const Value& v) { if (it != vars.end()) { return it->second; } - int convCost = getConversionCost(v); + int64_t convCost = getConversionCost(v); + auto multiplierIt = conversionCostMultipliers.find(v); + if (multiplierIt != conversionCostMultipliers.end()) { + convCost *= multiplierIt->second; + } RepVars repVars{/*c=*/model.NewBoolVar(), /*e=*/model.NewBoolVar(), /*conv=*/model.NewBoolVar(), /*mode=*/BoolVar()}; - objective += repVars.conv; - // add a new conversion variable equal to the "representative" conversion cost - // and force equality between them. We never need references to these other - // variables though; we just use repVars.conv as their proxy. - for (int i = 1; i < convCost; i++) { - BoolVar b = model.NewBoolVar(); - model.AddEquality(repVars.conv, b); - objective += b; - } + objective += LinearExpr::Term(repVars.conv, convCost); vars[v] = repVars; return vars[v]; } +void NTTSolver::setConversionCostMultiplier(const Value& v, + int64_t multiplier) { + assert(multiplier >= 1 && "conversion cost multiplier must be positive"); + assert(!vars.contains(v) && + "conversion cost multiplier must be set before value is modeled"); + conversionCostMultipliers[v] = multiplier; +} + const BoolVar& NTTSolver::RepVars::getVarForm(Form form) const { return form == Form::COEFF ? c : e; } @@ -112,12 +116,28 @@ void NTTSolver::prohibitBothForms(const Value& v) { {vs.getVarForm(Form::COEFF).Not(), vs.getVarForm(Form::EVAL).Not()}); } +void NTTSolver::equateNativeForm(const Value& a, const Value& b) { + RepVars& as = getOrCreateVars(a); + RepVars& bs = getOrCreateVars(b); + model.AddEquality(as.c, bs.c); +} + void NTTSolver::implyUse(const Value& out, const Value& in, Form form) { RepVars& outs = getOrCreateVars(out); RepVars& ins = getOrCreateVars(in); model.AddImplication(outs.getVarForm(form), ins.getVarForm(form)); } +void NTTSolver::requireSourceMatchesNativeForm(const Value& target, + const Value& source) { + RepVars& ts = getOrCreateVars(target); + RepVars& ss = getOrCreateVars(source); + // native(target) == COEFF (target.c) => source must supply COEFF. + model.AddImplication(ts.c, ss.c); + // native(target) == EVAL (i.e. !target.c) => source must supply EVAL. + model.AddBoolOr({ts.c, ss.e}); +} + void NTTSolver::implyMode(const Value& out, const Value& in) { RepVars& outs = getOrCreateVars(out); RepVars& ins = getOrCreateVars(in); diff --git a/lib/Dialect/Polynomial/Transforms/NTTSolver.h b/lib/Dialect/Polynomial/Transforms/NTTSolver.h index 42a754cf7d..32b579509d 100644 --- a/lib/Dialect/Polynomial/Transforms/NTTSolver.h +++ b/lib/Dialect/Polynomial/Transforms/NTTSolver.h @@ -1,6 +1,8 @@ #ifndef LIB_DIALECT_POLYNOMIAL_TRANSFORMS_NTT_SOLVER_H_ #define LIB_DIALECT_POLYNOMIAL_TRANSFORMS_NTT_SOLVER_H_ +#include + #include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h" #include "llvm/include/llvm/ADT/DenseMap.h" // from @llvm-project #include "mlir/include/mlir/IR/Value.h" // from @llvm-project @@ -31,15 +33,51 @@ class NTTSolver { RepVars& getOrCreateVars(const Value& v); operations_research::sat::CpModelBuilder model; llvm::DenseMap vars; + llvm::DenseMap conversionCostMultipliers; operations_research::sat::LinearExpr objective; public: + // Scales v's conversion cost in the objective by `multiplier` (e.g. the + // number of times a loop containing v's conversion site will execute). + // Must be called before any other solver method touches v. + void setConversionCostMultiplier(const Value& v, int64_t multiplier); void forceDemandEitherForm(const Value& v); void forceDemandFixedForm(const Value& v, Form form); void implyForm(const Value& v, Form a, Form b); void implyUse(const Value& out, const Value& in, Form form); + // Requires `source` to supply whichever form `target` is materialized in + // natively -- COEFF if target's coeff-demand bit is set, EVAL otherwise + // (mirroring the "needsForm(COEFF) ? COEFF : EVAL" tie-break PolyMulToNTT.cpp + // uses when actually materializing a value). This is weaker than requiring + // `source` to supply *every* form `target` needs: if target additionally + // needs the other form too, that is satisfied by a separate, locally + // materialized conversion at target's own definition site, which does not + // require `source` to supply it. Forwarding edges into a region-branch + // successor input (e.g. a loop's entry operand into its iter_arg) use this + // instead of two implyUse calls to avoid over-constraining `source` into + // needing a form nothing actually consumes. + void requireSourceMatchesNativeForm(const Value& target, const Value& source); void implyMode(const Value& out, const Value& in); void prohibitBothForms(const Value& v); + // Forces two values to resolve to the same *native* (materialized-in-the-IR) + // form. This is for region-branch successor inputs (e.g. a loop iter_arg) + // that share a single physical operand: MLIR's RegionBranchOpInterface can + // forward one operand to several successor inputs at once (e.g. scf.for's + // scf.yield operand doubles as both the next iteration's iter_arg and the + // loop's own result), so those targets have no independent operand slot to + // diverge on and must end up with the same materialized type. + // + // This only needs to tie the coeff-demand bit, not the full demand pattern: + // the native form of any value in this pass is chosen as + // "needsForm(COEFF) ? COEFF : EVAL" (see PolyMulToNTT.cpp), a function of + // the coeff-demand bit alone. Tying just that bit is therefore sufficient to + // guarantee the two values resolve to the same materialized type, while + // leaving each value free to independently need (or not need) the other + // form as a separate, locally materialized conversion -- e.g. a loop + // iter_arg used in eval form only inside the loop body shouldn't force the + // loop's result to also be materialized in eval form if nothing outside the + // loop needs it. + void equateNativeForm(const Value& a, const Value& b); void addConversionCostForForm(const Value& v, Form form); void addConversionCostIfBothForms(const Value& v); void setZeroConversionCost(const Value& v); diff --git a/lib/Dialect/Polynomial/Transforms/PolyMulToNTT.cpp b/lib/Dialect/Polynomial/Transforms/PolyMulToNTT.cpp index 48ec017406..607f14a187 100644 --- a/lib/Dialect/Polynomial/Transforms/PolyMulToNTT.cpp +++ b/lib/Dialect/Polynomial/Transforms/PolyMulToNTT.cpp @@ -1,11 +1,15 @@ #include "lib/Dialect/Polynomial/Transforms/PolyMulToNTT.h" +#include +#include + #include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h" #include "lib/Dialect/Polynomial/IR/PolynomialOps.h" #include "lib/Dialect/Polynomial/IR/PolynomialTypes.h" #include "lib/Dialect/Polynomial/Transforms/NTTSolver.h" #include "llvm/include/llvm/ADT/DenseMap.h" // from @llvm-project #include "llvm/include/llvm/ADT/STLExtras.h" // from @llvm-project +#include "llvm/include/llvm/ADT/SetVector.h" // from @llvm-project #include "llvm/include/llvm/ADT/SmallVector.h" // from @llvm-project #include "llvm/include/llvm/ADT/TypeSwitch.h" // from @llvm-project #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project @@ -13,13 +17,16 @@ #include "mlir/include/mlir/IR/Builders.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/include/mlir/IR/Matchers.h" // from @llvm-project #include "mlir/include/mlir/IR/Operation.h" // from @llvm-project #include "mlir/include/mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/include/mlir/IR/Types.h" // from @llvm-project #include "mlir/include/mlir/IR/Value.h" // from @llvm-project #include "mlir/include/mlir/IR/ValueRange.h" // from @llvm-project -#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project -#include "mlir/include/mlir/Support/WalkResult.h" // from @llvm-project +#include "mlir/include/mlir/Interfaces/ControlFlowInterfaces.h" // from @llvm-project +#include "mlir/include/mlir/Interfaces/LoopLikeInterface.h" // from @llvm-project +#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project +#include "mlir/include/mlir/Support/WalkResult.h" // from @llvm-project #include "mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project // IWYU pragma: begin_keep @@ -108,6 +115,118 @@ static llvm::SmallVector filterPolynomialOps(ValueRange values) { return result; } +// For loops whose iteration count is not statically-known, we set the cost of +// conversions inside the loop to this large value so that the solver +// strongly prefers hoisting a conversion out of it over leaving it inside. +static constexpr int64_t kUnknownLoopIterations = 1000; + +// Returns the cost of a conversion inside a specific single region. +// E.g., for loops with a statically-known iteration count, returns the +// iteration count via LoopLikeOpInterface::getStaticTripCount. +// Falls back to RegionBranchOpInterface::getRegionInvocationBounds, +// which is more general (e.g., it also covers multi-region ops like scf.while). +// For loops whose iteration count is not statically-known, returns +// kUnknownLoopIterations. +static int64_t getRegionIterationCount(RegionBranchOpInterface loopOp, + Region* region) { + if (auto loopLike = dyn_cast(loopOp.getOperation())) { + if (std::optional tripCount = loopLike.getStaticTripCount()) { + return static_cast( + tripCount->getLimitedValue(std::numeric_limits::max())); + } + } + + // getRegionInvocationBounds takes *attributes* rather than raw operands. + // `matchPattern` and `m_Constant` are built-ins that put either a + // constant-like op or the null attribute into `operandConstants`. + SmallVector operandConstants(loopOp->getNumOperands()); + for (auto [i, operand] : llvm::enumerate(loopOp->getOperands())) { + matchPattern(operand, m_Constant(&operandConstants[i])); + } + SmallVector bounds; + loopOp.getRegionInvocationBounds(operandConstants, bounds); + for (auto [candidate, bound] : llvm::zip(loopOp->getRegions(), bounds)) { + if (&candidate != region) continue; + // the "bound" here is a pair with a lower-bound and upper-bound on the + // number of times the region will execute. We choose the upper-bound as + // a worst-case cost. Note that for non-static bounds, this API sets the + // lower-bound to zero, so it's not useful as an estimate if the + // upper-bound isn't available. + if (std::optional upper = bound.getUpperBound()) { + return *upper; + } + } + return kUnknownLoopIterations; +} + +// Return how many times a conversion at v's materialization site actually runs: +// once per iteration of every loop it's nested in (multiplied together for +// nested loops), or 1 if it's not in a loop at all. +static int64_t getConversionCostMultiplier(Value v) { + int64_t weight = 1; + // This loop starts in the op's defining region and iterates upward to + // capture nested regions + for (Region* region = getEnclosingRepetitiveRegion(v); region; + region = getEnclosingRepetitiveRegion(region->getParentOp())) { + auto loopOp = cast(region->getParentOp()); + weight *= getRegionIterationCount(loopOp, region); + } + return weight; +} + +// A note on terminology, using scf.for as an example: +// %result = scf.for %i = %lb to %ub step %step iter_args(%acc = %init) -> +// (!poly_ty) { +// %next = polynomial.mul %acc, %acc : !poly_ty +// scf.yield %next : !poly_ty +// } +// +// %init and %next are operands forwarded across a region boundary. %acc and +// %result are the "successor inputs" that receive them -- %acc is the body +// region's successor input (a block argument), %result is the parent's +// successor input (an op result). A "successor" is the *destination* +// itself (the region, or "the parent"), not a value. +// +// getSuccessorOperandInputMapping returns a map from each such operand to +// the successor input(s) it feeds, e.g. {%init: [%acc, %result], %next: [%acc, +// %result]} (both fan out to both here, since the trip count isn't +// statically known in this example). +// +// Adds solver constraints for a RegionBranchOpInterface op (scf.for, scf.if, +// scf.while, ...): every successor input can be fed by more than one +// operand (as above), so each gets its own coeff/eval demand, independent +// of the op(s) that produced the operands flowing into it. This one call +// covers loop entry, backedges, and results, without needing to +// special-case any individual op. +static void addRegionBranchConstraints(NTTSolver& solver, + RegionBranchOpInterface regionBranchOp) { + RegionBranchSuccessorMapping operandToInputs; + regionBranchOp.getSuccessorOperandInputMapping(operandToInputs); + + for (auto& [operand, inputs] : operandToInputs) { + Value source = operand->get(); + if (!isPolyValue(source)) continue; + + for (Value target : inputs) { + solver.addConversionCostIfBothForms(target); + // Whichever form target (a successor input) needs natively, source + // (the operand feeding it) must supply -- e.g. if %result needs COEFF, + // the yielded value %next must too. + solver.requireSourceMatchesNativeForm(target, source); + } + + // One operand can feed more than one target -- e.g. scf.for's yielded + // value feeds both the next iter_arg and the loop's result. Since both + // are rewritten from that same operand, they must end up the same type. + // In short, the previous loop tied the yielded value to iterArgs.front(), + // this loop ties iterArgs.front() to loop.getResults(0) so that all inputs + // are forced to the same form + for (Value other : llvm::drop_begin(inputs)) { + solver.equateNativeForm(inputs.front(), other); + } + } +} + void PolyMulToNTT::runOnOperation() { func::FuncOp func = getOperation(); MLIRContext* context = &getContext(); @@ -168,8 +287,22 @@ void PolyMulToNTT::runOnOperation() { // fix up the inputs // 5. Fix the function signature and arguments // - // TODO(#2685): This pass only handles polynomial and tensor ops in functions; - // it does *not* support ops with regions (e.g., loops). + // Ops implementing RegionBranchOpInterface (scf.for, scf.if, scf.while, + // ...) are supported via a separate mechanism layered on top of the five + // steps above; see addRegionBranchConstraints and Steps 3b/4b for details. + // A conversion's cost is weighted by how many times it actually runs + // (getConversionCostMultiplier), using RegionBranchOpInterface's own + // invocation-bounds query: a loop with a statically known bound (e.g. + // scf.for with constant lower/upper bound and step) contributes its exact + // iteration count, and any loop we can't size statically is assumed to + // run many times, so the solver strongly prefers hoisting a conversion + // out of it. + // + // A block argument or region-branch result that isn't fed by any operand + // (e.g. a loop's induction variable) is not a forwarding edge and is not + // supported; this isn't a real limitation in practice since such + // "produced" successor inputs are never polynomial-typed for the ops this + // pass deals with. // Steps 1, 3, and 4 above involve walking the AST. Since we're going to be // doing multiple walks and adding some nodes on the way, we first memoize @@ -177,8 +310,50 @@ void PolyMulToNTT::runOnOperation() { // prune it to remove ops that don't involve polynomials. This doesn't remove // ops from the AST, it just means that we don't walk over them later. llvm::SmallVector rewriteOrder; + // RegionBranchOpInterface ops (e.g. scf.for, scf.if, scf.while) collected + // here have their polynomial operands/results handled entirely by the + // dedicated region-branch forwarding logic below (see + // addRegionBranchConstraints), not by the generic per-op walks over + // rewriteOrder. + llvm::SmallVector regionBranchOps; WalkResult wr = func.walk([&](Operation* op) -> WalkResult { + // RegionBranchTerminatorOpInterface ops that actually terminate a + // region of a RegionBranchOpInterface parent (e.g. scf.yield inside + // scf.for/scf.if, scf.condition inside scf.while) never need their + // own entry in rewriteOrder: every polynomial operand they have is a + // forwarding edge into some successor input, captured via the + // parent's successor mapping and rewritten directly through that + // mapping's OpOperand pointers. + // + // The interface check alone isn't enough to identify these: many + // ReturnLike ops (e.g. func.return, polynomial.yield) satisfy + // RegionBranchTerminatorOpInterface "for free" regardless of what + // their parent op is, even when that parent has nothing to do with + // region-branching (e.g. func.func, polynomial.apply_coefficientwise) + // and must still be handled by the ordinary per-op walks below. So we + // additionally require the parent op to actually implement + // RegionBranchOpInterface. + Operation* parentOp = op->getParentOp(); + if (isa(op) && parentOp && + isa(parentOp)) { + return WalkResult::advance(); + } + if (auto regionBranchOp = dyn_cast(op)) { + // Likewise, a region branch op's own operands/results (e.g. + // scf.for's initial iter_arg operand, or a value returned to the + // parent) are forwarding edges, not ordinary op inputs/outputs, so + // this op is excluded from the single-poly-result restriction + // below: each of its successor inputs gets its own, independently + // solved form, so there's no bound on how many polynomial + // loop-carried values it may have. + if (!filterPolynomialOps(op->getOperands()).empty() || + !filterPolynomialOps(op->getResults()).empty()) { + regionBranchOps.push_back(regionBranchOp); + } + return WalkResult::advance(); + } + auto polyResults = filterPolynomialOps(op->getResults()); auto polyOperands = filterPolynomialOps(op->getOperands()); @@ -192,6 +367,15 @@ void PolyMulToNTT::runOnOperation() { signalPassFailure(); return WalkResult::interrupt(); } + // A conversion on this op's result, if one is needed, is inserted + // right after the op, so it costs however many times that site + // actually runs. + for (Value result : polyResults) { + int64_t multiplier = getConversionCostMultiplier(result); + if (multiplier != 1) { + solver.setConversionCostMultiplier(result, multiplier); + } + } } return WalkResult::advance(); }); @@ -207,6 +391,41 @@ void PolyMulToNTT::runOnOperation() { } } + // Region-branch successor inputs (loop iter_args, scf.if results, etc.) + // are collected here so that Step 3 can materialize them once solving is + // done; see addRegionBranchConstraints for how they're constrained. Cost + // multipliers for all of them must be registered + // before any of them are constrained + // (addRegionBranchConstraints) -- see that function's comment for why. + llvm::SetVector polySuccessorInputs; + for (RegionBranchOpInterface regionBranchOp : regionBranchOps) { + // setConversionCostMultiplier must run on a value before anything else + // touches it -- but a nested loop's successor input can be fed by an + // enclosing loop's own successor input (e.g. the outer iter_arg feeding the + // inner loop's entry operand), so we can't set multipliers while we build + // constraints: an inner loop could reference the outer loop's block + // argument as a source (via requireSourceMatchesNativeForm below) before + // the outer loop's own turn to register it. So this runs, for every + // RegionBranchOpInterface op, before addRegionBranchConstraints runs for + // any of them. + RegionBranchSuccessorMapping operandToInputs; + regionBranchOp.getSuccessorOperandInputMapping(operandToInputs); + for (auto& [operand, inputs] : operandToInputs) { + if (!isPolyValue(operand->get())) continue; + for (Value target : inputs) { + if (polySuccessorInputs.insert(target)) { + int64_t multiplier = getConversionCostMultiplier(target); + if (multiplier != 1) { + solver.setConversionCostMultiplier(target, multiplier); + } + } + } + } + } + for (RegionBranchOpInterface regionBranchOp : regionBranchOps) { + addRegionBranchConstraints(solver, regionBranchOp); + } + for (Operation* op : rewriteOrder) { auto polyResults = filterPolynomialOps(op->getResults()); auto polyOperands = filterPolynomialOps(op->getOperands()); @@ -593,6 +812,49 @@ void PolyMulToNTT::runOnOperation() { } } + /**************************************************************** + ***** Step 3b: Materialize region-branch successor inputs ****** + *****************************************************************/ + // Block arguments and region-branch results are materialized exactly like + // function arguments (see the loop over func.getArguments() above): we + // arbitrarily fix one native form into the IR -- preferring coeff form + // when both are needed -- and, if the other form is also required by some + // use, insert a single conversion right where the value comes into + // existence: at the start of the owning block for a block argument, or + // right after the op for a value returned to the parent. We remember which + // form each successor input was fixed to so Step 4b can rewrite the + // forwarding operands that feed it to match. + llvm::DenseMap successorNativeForm; + for (Value target : polySuccessorInputs) { + Form f = soln.needsForm(target, Form::COEFF) ? Form::COEFF : Form::EVAL; + successorNativeForm[target] = f; + + Type newTy = typeToForm(target.getType(), f); + if (!newTy) { + signalPassFailure(); + return; + } + target.setType(newTy); + + if (auto blockArg = dyn_cast(target)) { + b.setInsertionPointToStart(blockArg.getOwner()); + } else { + b.setInsertionPointAfter(target.getDefiningOp()); + } + + if (f == Form::COEFF) { + coeffFormCache[target] = target; + if (soln.needsForm(target, Form::EVAL)) { + evalFormCache[target] = addConversion(target, Form::EVAL); + } + } else { + evalFormCache[target] = target; + if (soln.needsForm(target, Form::COEFF)) { + coeffFormCache[target] = addConversion(target, Form::COEFF); + } + } + } + /************************************************ *********** Step 4: Fix up AST inputs ********** ***********************************************/ @@ -683,6 +945,26 @@ void PolyMulToNTT::runOnOperation() { } } + /********************************************************** + ***** Step 4b: Rewrite region-branch forwarding edges ***** + **********************************************************/ + // Every successor input now has a resolved, materialized native form + // (Step 3b). Point each operand that forwards into one -- the region + // branch op's own entry operands, and every operand yielded/forwarded + // inside its regions -- at the cached value in that form. + for (RegionBranchOpInterface regionBranchOp : regionBranchOps) { + RegionBranchSuccessorMapping operandToInputs; + regionBranchOp.getSuccessorOperandInputMapping(operandToInputs); + for (auto& [operand, inputs] : operandToInputs) { + if (!isPolyValue(operand->get())) continue; + // All of this operand's targets were tied to the same form in + // addRegionBranchConstraints (equateForms), so any one of them tells + // us the form this operand must be rewritten to. + Form form = successorNativeForm.at(inputs.front()); + operand->set(formToValue(operand->get(), form)); + } + } + /************************************************ ********* Step 5: Fix function signature ******* ***********************************************/ diff --git a/tests/Dialect/Polynomial/Transforms/poly_mul_to_ntt_region_branch.mlir b/tests/Dialect/Polynomial/Transforms/poly_mul_to_ntt_region_branch.mlir new file mode 100644 index 0000000000..0989e9afbb --- /dev/null +++ b/tests/Dialect/Polynomial/Transforms/poly_mul_to_ntt_region_branch.mlir @@ -0,0 +1,138 @@ +// RUN: heir-opt --convert-polynomial-mul-to-ntt %s | FileCheck %s + +!Zq0 = !mod_arith.int<1095233372161 : i64> +#ring_1 = #polynomial.ring, polynomialModulus = <1 + x**1024>> +!poly_ty_1 = !polynomial.polynomial +!ntt_poly_ty_1 = !polynomial.polynomial + +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 + // to exactly one occurrence per iteration, inside the loop body around the + // mul, with no conversions needed on the function argument or the loop + // result themselves. + // CHECK: func.func @for_iter_arg_needs_both_forms([[x:%.+]]: [[poly_ty_1:![^ ]+]], + // CHECK-SAME: -> tensor<1024x[[RNS:![^ ]+]]> { + // CHECK: [[r:%.+]] = scf.for {{.*}} iter_args([[acc:%.+]] = [[x]]) -> ([[poly_ty_1]]) { + // CHECK: [[acce:%.+]] = polynomial.ntt [[acc]] : [[poly_ty_1]] + // CHECK: [[sq:%.+]] = polynomial.mul [[acce]], [[acce]] : [[ntt_poly_ty_1:![^ ]+]] + // CHECK: [[sqc:%.+]] = polynomial.intt [[sq]] : [[ntt_poly_ty_1]] + // CHECK: scf.yield [[sqc]] : [[poly_ty_1]] + // CHECK: [[t:%.+]] = polynomial.to_tensor [[r]] : [[poly_ty_1]] -> tensor<1024x[[RNS]]> + // CHECK: return [[t]] : tensor<1024x[[RNS]]> + func.func @for_iter_arg_needs_both_forms(%x: !poly_ty_1, %lb: index, %ub: index, %step: index) -> tensor<1024x!rns.rns> { + %r = scf.for %i = %lb to %ub step %step iter_args(%acc = %x) -> !poly_ty_1 { + %sq = polynomial.mul %acc, %acc : !poly_ty_1 + scf.yield %sq : !poly_ty_1 + } + %t = polynomial.to_tensor %r : !poly_ty_1 -> tensor<1024x!rns.rns> + return %t : tensor<1024x!rns.rns> + } + + // Covers: two independent polynomial iter_args in the same scf.for, each + // getting its own, independently solved form. %acc0 needs coeff (feeds + // to_tensor after the loop) and gets NTT/INTT'd locally around the + // eval-only mul, exactly as above. %acc1 merely passes through unused, so + // it's free to settle on eval form throughout with no conversions at all. + // CHECK: func.func @for_multiple_independent_iter_args([[a:%.+]]: [[poly_ty_1]], [[b:%.+]]: [[ntt_poly_ty_1]], + // CHECK: [[loop:%.+]]:2 = scf.for {{.*}} iter_args([[acc0:%.+]] = [[a]], [[acc1:%.+]] = [[b]]) -> ([[poly_ty_1]], [[ntt_poly_ty_1]]) { + // CHECK: [[acc0e:%.+]] = polynomial.ntt [[acc0]] : [[poly_ty_1]] + // CHECK: [[sq0:%.+]] = polynomial.mul [[acc0e]], [[acc0e]] : [[ntt_poly_ty_1]] + // CHECK: [[sq0c:%.+]] = polynomial.intt [[sq0]] : [[ntt_poly_ty_1]] + // CHECK: scf.yield [[sq0c]], [[acc1]] : [[poly_ty_1]], [[ntt_poly_ty_1]] + // CHECK: [[t:%.+]] = polynomial.to_tensor [[loop]]#0 : [[poly_ty_1]] -> tensor<1024x[[RNS]]> + // CHECK: return [[t]], [[loop]]#1 : tensor<1024x[[RNS]]>, [[ntt_poly_ty_1]] + func.func @for_multiple_independent_iter_args(%a: !poly_ty_1, %b: !poly_ty_1, %lb: index, %ub: index, %step: index) -> (tensor<1024x!rns.rns>, !poly_ty_1) { + %r0, %r1 = scf.for %i = %lb to %ub step %step iter_args(%acc0 = %a, %acc1 = %b) -> (!poly_ty_1, !poly_ty_1) { + %sq0 = polynomial.mul %acc0, %acc0 : !poly_ty_1 + scf.yield %sq0, %acc1 : !poly_ty_1, !poly_ty_1 + } + %t = polynomial.to_tensor %r0 : !poly_ty_1 -> tensor<1024x!rns.rns> + return %t, %r1 : tensor<1024x!rns.rns>, !poly_ty_1 + } + + // Covers: a tensor loop-carried value. Since the result is returned + // directly with no coeff-only consumer, the whole loop (argument, iter_arg, + // and result) settles on eval form for free, with no conversions inserted. + // CHECK: func.func @for_tensor_iter_arg([[xt:%.+]]: tensor<2x[[ntt_poly_ty_1]]>, + // CHECK: [[rt:%.+]] = scf.for {{.*}} iter_args([[acct:%.+]] = [[xt]]) -> (tensor<2x[[ntt_poly_ty_1]]>) { + // CHECK-NOT: polynomial.ntt + // CHECK-NOT: polynomial.intt + // CHECK: [[sqt:%.+]] = polynomial.mul [[acct]], [[acct]] : tensor<2x[[ntt_poly_ty_1]]> + // CHECK: scf.yield [[sqt]] : tensor<2x[[ntt_poly_ty_1]]> + // CHECK: return [[rt]] : tensor<2x[[ntt_poly_ty_1]]> + func.func @for_tensor_iter_arg(%x: tensor<2x!poly_ty_1>, %lb: index, %ub: index, %step: index) -> tensor<2x!poly_ty_1> { + %r = scf.for %i = %lb to %ub step %step iter_args(%acc = %x) -> tensor<2x!poly_ty_1> { + %sq = polynomial.mul %acc, %acc : tensor<2x!poly_ty_1> + scf.yield %sq : tensor<2x!poly_ty_1> + } + return %r : tensor<2x!poly_ty_1> + } + + // Covers: scf.if yielding the same polynomial from both branches, with a + // coeff-only consumer of the original value and an eval-only consumer of + // the if's result. Since scf.if's regions take no block arguments, this + // exercises only the "value returned to the parent" successor-input case + // (not entry/backedge edges), forwarded from two different scf.yield ops + // (one per branch) to the same result. x needs both forms (coeff for + // to_tensor, eval to feed the if); the if's own result only ever needs + // eval, so the one unavoidable conversion lands on x rather than on it. + // CHECK: func.func @if_yields_both_branches([[cond:%.+]]: i1, [[x2:%.+]]: [[poly_ty_1]]) + // CHECK: [[x2e:%.+]] = polynomial.ntt [[x2]] : [[poly_ty_1]] + // CHECK: [[r2:%.+]] = scf.if [[cond]] -> ([[ntt_poly_ty_1]]) { + // CHECK: scf.yield [[x2e]] : [[ntt_poly_ty_1]] + // CHECK: } else { + // CHECK: scf.yield [[x2e]] : [[ntt_poly_ty_1]] + // CHECK: } + // 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>, !poly_ty_1) { + %r = scf.if %cond -> !poly_ty_1 { + scf.yield %x : !poly_ty_1 + } else { + scf.yield %x : !poly_ty_1 + } + %t = polynomial.to_tensor %x : !poly_ty_1 -> tensor<1024x!rns.rns> + %m = polynomial.mul %r, %r : !poly_ty_1 + return %t, %m : tensor<1024x!rns.rns>, !poly_ty_1 + } + + // Covers: scf.while, whose scf.condition op forwards the very same + // operand to two different successor inputs at once -- the "after" + // region's block argument (if the loop continues) and the scf.while op's + // own result (if it exits) -- exercising the case where a single physical + // operand fans out to multiple successor inputs that must share one + // resolved form (see NTTSolver::equateNativeForm). to_tensor forces the + // while's result (and, by the fan-out above, the "after" region's + // argument) to coeff. The "before" region's argument has no such + // constraint, so it's free to settle on eval directly to feed the + // eval-only mul with no local conversion; the fan-out's coeff requirement + // is instead satisfied by a single intt before scf.condition, and the + // "after" region converts back to eval before looping around. + // CHECK: func.func @while_loop([[x3:%.+]]: [[ntt_poly_ty_1]], [[cond3:%.+]]: i1) -> tensor<1024x[[RNS]]> { + // CHECK: [[r3:%.+]] = scf.while ([[acc3:%.+]] = [[x3]]) : ([[ntt_poly_ty_1]]) -> [[poly_ty_1]] { + // CHECK: [[sq3:%.+]] = polynomial.mul [[acc3]], [[acc3]] : [[ntt_poly_ty_1]] + // CHECK: [[sq3c:%.+]] = polynomial.intt [[sq3]] : [[ntt_poly_ty_1]] + // CHECK: scf.condition({{.*}}) [[sq3c]] : [[poly_ty_1]] + // CHECK: } do { + // CHECK: ^bb0([[after3:%.+]]: [[poly_ty_1]]): + // CHECK: [[after3e:%.+]] = polynomial.ntt [[after3]] : [[poly_ty_1]] + // CHECK: scf.yield [[after3e]] : [[ntt_poly_ty_1]] + // CHECK: } + // CHECK: [[t3:%.+]] = polynomial.to_tensor [[r3]] : [[poly_ty_1]] -> tensor<1024x[[RNS]]> + // CHECK: return [[t3]] : tensor<1024x[[RNS]]> + func.func @while_loop(%x: !poly_ty_1, %cond0: i1) -> tensor<1024x!rns.rns> { + %r = scf.while (%acc = %x) : (!poly_ty_1) -> !poly_ty_1 { + %sq = polynomial.mul %acc, %acc : !poly_ty_1 + %c = arith.constant true + scf.condition(%c) %sq : !poly_ty_1 + } do { + ^bb0(%arg: !poly_ty_1): + scf.yield %arg : !poly_ty_1 + } + %t = polynomial.to_tensor %r : !poly_ty_1 -> tensor<1024x!rns.rns> + return %t : tensor<1024x!rns.rns> + } +}