Implementation of loop support in PolyMulToNTT

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",

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);

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 <cstdint>
+
 #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<Value, RepVars> vars;
+  llvm::DenseMap<Value, int64_t> 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);