Implemented support for loops via RegionBranchOpInterface

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

@@ -1,16 +1,22 @@
 #include "lib/Dialect/Polynomial/Transforms/NTTSolver.h"
 
 #include <cassert>
+#include <cstdint>
+#include <limits>
+#include <optional>
 
 #include "lib/Dialect/ModArith/IR/ModArithTypeInterfaces.h"
 #include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
 #include "lib/Dialect/Polynomial/IR/PolynomialTypes.h"
 #include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/Operation.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/Support/LLVM.h"              // from @llvm-project
-#include "ortools/sat/cp_model.h"         // from @com_google_ortools
-#include "ortools/sat/cp_model_solver.h"  // from @com_google_ortools
+#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 "ortools/sat/cp_model.h"            // from @com_google_ortools
+#include "ortools/sat/cp_model_solver.h"     // from @com_google_ortools
 
 namespace mlir {
 namespace heir {
@@ -39,10 +45,17 @@ Form CPSATSolution::getMode(const Value& v) const {
 }
 
 bool CPSATSolution::isValid() const {
+  // Both a proven optimum and a merely FEASIBLE solution are accepted. A
+  // FEASIBLE solution may use more conversions than necessary, so callers
+  // should check isOptimal() and warn.
   return soln.status() == CpSolverStatus::OPTIMAL ||
          soln.status() == CpSolverStatus::FEASIBLE;
 }
 
+bool CPSATSolution::isOptimal() const {
+  return soln.status() == CpSolverStatus::OPTIMAL;
+}
+
 int getConversionCost(const Value& v) {
   Type t = v.getType();
   int multiplicity = 1;
@@ -63,26 +76,51 @@ int getConversionCost(const Value& v) {
   return coeffTy.getNumResidues() * multiplicity;
 }
 
+// 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 number of times the body of `op` runs: its static trip count if
+// `op` is a loop with one, and kUnknownLoopIterations otherwise.
+static int64_t getLoopIterationCount(Operation* op) {
+  if (auto loopLike = dyn_cast<LoopLikeOpInterface>(op)) {
+    if (std::optional<APInt> tripCount = loopLike.getStaticTripCount()) {
+      return static_cast<int64_t>(
+          tripCount->getLimitedValue(std::numeric_limits<int64_t>::max()));
+    }
+  }
+  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())) {
+    weight *= getLoopIterationCount(region->getParentOp());
+  }
+  return weight;
+}
+
 NTTSolver::RepVars& NTTSolver::getOrCreateVars(const Value& v) {
   auto it = vars.find(v);
   if (it != vars.end()) {
     return it->second;
   }
-  int convCost = getConversionCost(v);
+  // A conversion on v, if one is needed, is inserted right where v comes into
+  // existence, so it costs however many times that site actually runs.
+  int64_t convCost = getConversionCost(v) * getConversionCostMultiplier(v);
   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];
 }
@@ -139,17 +177,25 @@ void NTTSolver::addConversionCostForForm(const Value& v, Form form) {
   model.AddEquality(vs.conv, form == Form::COEFF ? vs.c : vs.e);
 }
 
-void NTTSolver::addConversionCostIfBothForms(const Value& v) {
-  RepVars& vs = getOrCreateVars(v);
-  // vs.conv <=> vs.c and vs.e (otherwise we just run the op in the mode
-  // required)
+BoolVar NTTSolver::addFloatingMode(llvm::ArrayRef<Value> members) {
+  BoolVar mode = model.NewBoolVar();
+  modeVars.push_back(mode);
+  for (const Value& v : members) {
+    getOrCreateVars(v).mode = mode;
+  }
+  return mode;
+}
 
-  // The first two give vs.conv => vs.c and vs.conv => vs.e, hence vs.conv =>
-  // vs.c ^ vs.e
-  model.AddImplication(vs.conv, vs.c);
-  model.AddImplication(vs.conv, vs.e);
-  // The last is vs.c ^ vs.e => vs.conv
-  model.AddBoolOr({vs.c.Not(), vs.e.Not(), vs.conv});
+void NTTSolver::addModeConversionCost(const Value& member,
+                                      const BoolVar& mode) {
+  RepVars& vs = getOrCreateVars(member);
+  // conv == (mode ^ c) v (!mode ^ e), i.e. a conversion is needed at
+  // member's own materialization site iff member's own demand disagrees
+  // with the group's shared mode.
+  model.AddBoolOr({mode.Not(), vs.c.Not(), vs.conv});  // mode ^ c => conv
+  model.AddBoolOr({mode, vs.e.Not(), vs.conv});        // !mode ^ e => conv
+  model.AddBoolOr({vs.conv.Not(), mode.Not(), vs.c});  // conv ^ mode => c
+  model.AddBoolOr({vs.conv.Not(), mode, vs.e});        // conv ^ !mode => e
 }
 
 void NTTSolver::setZeroConversionCost(const Value& v) {
@@ -168,20 +214,46 @@ void NTTSolver::setZeroConversionCost(const Value& v) {
   objective += vs.e;
 }
 
-void NTTSolver::addOpMode(const Value& v) {
-  RepVars& vs = getOrCreateVars(v);
-  vs.mode = model.NewBoolVar();
-  // 1 = E-mode, 0 = C-mode
-  // If v_e is needed but not v_c, set v_mode = 1
-  // v_e ^ ¬v_c => v_mode
-  model.AddBoolOr({vs.e.Not(), vs.c, vs.mode});
-  // If v_c is needed but not v_e, set v_mode = 0
-  // v_c ^ ¬v_e => v_mode
-  model.AddBoolOr({vs.c.Not(), vs.e, vs.mode.Not()});
-}
-
 CPSATSolution NTTSolver::solve() {
-  model.Minimize(objective);
+  // `objective` is the total number of per-limb NTTs+iNTTs for a fixed
+  // assignment to the set of variables. Multiple assignments can result in the
+  // same objective, making deterministic tests fragile. We *partially* address
+  // this problem by adding a secondary objective `tieBreak` which is the total
+  // number of modes set to EVAL. This ensures that, among all possible
+  // solutions that minimize the primary objective, we pick one which
+  // (arbitrarily) minimizes the number of EVAL modes. `primaryScale` is a
+  // constant (i.e., independent of the mode assignments) which is always larger
+  // than `tieBreak`, therefore `objective * primaryScale + tieBreak` results in
+  // a two-tiered optimization. The tie-break term is at most modeVars.size(),
+  // so scaling the primary objective by modeVars.size() + 1 means saving one
+  // unit of conversion cost always outweighs any change in the tie-break term.
+  //
+  // The current tests all happen to have a unique solution under this
+  // condition, but in general, multiple solutions may still exist. For example:
+  //  - f(a, b) has the same objective cost when exactly one of a/b is EVAL.
+  //      Both solutions also clearly have the same *secondary* objective cost,
+  //      so a non-deterministic solver (either from randomness or from
+  //      multi-threading) or different solver version is free to choose a
+  //      solution different from the one chosen by the current solver.
+  //  - Regions that run zero times (e.g., a loop with a static trip count of 0)
+  //      can include conversions (or not) that don't change the objective,
+  //      since the region never runs, but which changes the IR.
+  //  - Returned constants, which have no conversion cost, can be output in
+  //  either form
+  //  - Zero-length tensors have arbitrary form
+  //
+  // Currently, non-determinism isn't hurting anything, so finding a robust,
+  // cheap way to produce a globally-unique solution to the CPSAT instance is
+  // left as an open problem. But if you update the solver version, run this on
+  // a very big machine with lots of cores, or make a seemingly innocuous change
+  // in PolyMulToNTT that shouldn't affect the optimal solution, be aware that
+  // the concrete IR output by this pass can change!
+  LinearExpr tieBreak;
+  for (const BoolVar& mode : modeVars) {
+    tieBreak += mode;
+  }
+  const int64_t primaryScale = static_cast<int64_t>(modeVars.size()) + 1;
+  model.Minimize(objective * primaryScale + tieBreak);
   const CpSolverResponse resp = Solve(model.Build());
   return CPSATSolution(*this, resp);
 }

lib/Dialect/Polynomial/Transforms/NTTSolver.h

@@ -1,10 +1,14 @@
 #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
-#include "ortools/sat/cp_model.h"            // from @com_google_ortools
+#include "llvm/include/llvm/ADT/ArrayRef.h"     // from @llvm-project
+#include "llvm/include/llvm/ADT/DenseMap.h"     // from @llvm-project
+#include "llvm/include/llvm/ADT/SmallVector.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/Value.h"         // from @llvm-project
+#include "ortools/sat/cp_model.h"               // from @com_google_ortools
 
 namespace mlir {
 namespace heir {
@@ -28,22 +32,53 @@ class NTTSolver {
 
     const operations_research::sat::BoolVar& getVarForm(Form form) const;
   };
+  // v's conversion cost in the objective is scaled by the number of times
+  // v's conversion site executes (i.e., the product of the trip counts of all
+  // loops enclosing it).
   RepVars& getOrCreateVars(const Value& v);
   operations_research::sat::CpModelBuilder model;
   llvm::DenseMap<Value, RepVars> vars;
+  // Every mode variable created by addFloatingMode. solve() uses these to
+  // break ties between equal-cost solutions in favor of COEFF mode.
+  llvm::SmallVector<operations_research::sat::BoolVar> modeVars;
+  // Total conversion cost. solve() minimizes this first, then the tie-break.
   operations_research::sat::LinearExpr objective;
 
  public:
   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);
+  // `out`'s mode (0 = COEFF, 1 = EVAL) determines which form `in` must supply,
+  // gated on `out` being used in some form at all (if `out` is dead, nothing
+  // is forced on `in`). Also used to tie a region-branch successor input's
+  // shared mode (see addFloatingMode) to the operand forwarding into it.
   void implyMode(const Value& out, const Value& in);
   void prohibitBothForms(const Value& v);
   void addConversionCostForForm(const Value& v, Form form);
-  void addConversionCostIfBothForms(const Value& v);
+  // Allocates one fresh mode variable (0 = COEFF, 1 = EVAL) shared by every
+  // value in `members`, and returns it. The mode is left fully free -- driven
+  // only by the objective via addModeConversionCost -- since any value in
+  // `members` may have its own, independently-forced demand that the shared
+  // mode has to reconcile rather than dictate. For a single-output "either
+  // form" op, call this with a one-element list; for a region-branch
+  // successor-input group (e.g. a loop's iter_arg and its result, which MLIR
+  // requires to share one materialized type), call it with every member of
+  // that group.
+  operations_research::sat::BoolVar addFloatingMode(
+      llvm::ArrayRef<Value> members);
+  // Charges `member`'s conversion variable iff `member`'s own demand
+  // disagrees with the shared `mode`: conv == (mode ∧ c) ∨ (¬mode ∧ e). This
+  // generalizes "conversion needed iff both forms independently demanded" to
+  // the case where `member` shares its native form with other values via
+  // addFloatingMode and may need a conversion even when its own demand is
+  // one-sided, because a different member of the group pulled the shared mode
+  // the other way.
+  void addModeConversionCost(const Value& member,
+                             const operations_research::sat::BoolVar& mode);
   void setZeroConversionCost(const Value& v);
-  void addOpMode(const Value& v);
+  // Returns a minimum-conversion-cost solution. Solutions are not guaranteed
+  // to be unique or repeatable.
   CPSATSolution solve();
   friend class CPSATSolution;
 };
@@ -59,7 +94,11 @@ class CPSATSolution {
   bool needsForm(const Value& v, Form form) const;
   bool needsConversion(const Value& v) const;
   Form getMode(const Value& v) const;
+  // True if CP-SAT proved the solution optimal or feasible; anything else is an
+  // error.
   bool isValid() const;
+  // True iff CP-SAT proved the solution optimal.
+  bool isOptimal() const;
 
  private:
   const NTTSolver& solver;