[Scheduling] Add Presburger Simplex scheduler

include/circt/Scheduling/Algorithms.h

@@ -80,6 +80,19 @@ LogicalResult scheduleLP(CyclicProblem &prob, Operation *lastOp);
 /// prob does not include \p lastOp.
 LogicalResult scheduleCPSAT(SharedOperatorsProblem &prob, Operation *lastOp);
 
+/// Solve the basic problem using linear programming and the Presburger solver.
+/// The objective is to minimize the start time of the given \p lastOp. Fails
+/// if the dependence graph contains cycles, or \p prob does not include \p
+/// lastOp.
+LogicalResult schedulePresburger(Problem &prob, Operation *lastOp);
+
+/// Solve the resource-free cyclic problem using linear programming and the
+/// Presburger solver. The objectives are to determine the smallest feasible
+/// initiation interval, and to minimize the start time of the given \p lastOp.
+/// Fails if the dependence graph contains cycles that do not include at least
+/// one edge with a non-zero distance, or \p prob does not include \p lastOp.
+LogicalResult schedulePresburger(CyclicProblem &prob, Operation *lastOp);
+
 } // namespace scheduling
 } // namespace circt
 

lib/Scheduling/CMakeLists.txt

@@ -3,6 +3,7 @@ set(LLVM_OPTIONAL_SOURCES
   ChainingSupport.cpp
   CPSATSchedulers.cpp
   LPSchedulers.cpp
+  PresburgerSchedulers.cpp
   Problems.cpp
   SimplexSchedulers.cpp
   TestPasses.cpp
@@ -12,6 +13,7 @@ set(LLVM_OPTIONAL_SOURCES
 set(SCHEDULING_SOURCES
   ASAPScheduler.cpp
   ChainingSupport.cpp
+  PresburgerSchedulers.cpp
   Problems.cpp
   SimplexSchedulers.cpp
   Utilities.cpp

lib/Scheduling/PresburgerSchedulers.cpp

@@ -0,0 +1,231 @@
+//===- PresburgerSchedulers.cpp -  Presburger lib based schedulers --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Implementation of linear programming-based schedulers with the Presburger
+// library Simplex.
+//
+//===----------------------------------------------------------------------===//
+
+#include "circt/Scheduling/Algorithms.h"
+#include "circt/Scheduling/Utilities.h"
+
+#include "mlir/Analysis/Presburger/Simplex.h"
+#include "mlir/Analysis/Presburger/Utils.h"
+
+#include "mlir/IR/Operation.h"
+
+using namespace circt;
+using namespace circt::scheduling;
+
+using namespace mlir::presburger;
+
+namespace {
+
+/// The Solver finds the smallest II that satisfies the constraints and
+/// minimizes the objective function. The solver also finds when a particular
+/// operation should be scheduled.
+class Solver : private LexSimplex {
+public:
+  Solver(Problem &prob, unsigned numObj)
+      : LexSimplex(1 + numObj + prob.getOperations().size()), prob(prob) {
+    // Offsets for variable types.
+    unsigned problemVarOffset = numObj + 1;
+
+    // Map each operation to a variable representing its start time and make
+    // their start time positive.
+    unsigned var = problemVarOffset;
+    for (Operation *op : prob.getOperations()) {
+      opToVar[op] = var;
+      addLowerBound(var, MPInt(0));
+      ++var;
+    }
+  }
+
+  /// Get the number of columns in the constraint system.
+  unsigned getNumCols() const { return LexSimplex::getNumVariables() + 1; }
+
+  using LexSimplex::addEquality;
+  using LexSimplex::addInequality;
+
+  /// Get the index of the operation in the solver.
+  unsigned getOpIndex(Operation *op) const { return opToVar.lookup(op); }
+
+  /// Create a latency constraint representing the given dependence.
+  /// The constraint is represented as:
+  /// dstOpStartTime >= srcOpStartTime + latency.
+  void createLatencyConstraint(MutableArrayRef<MPInt> row,
+                               Problem::Dependence dep) {
+    // Constraint: dst >= src + latency.
+    Operation *src = dep.getSource();
+    Operation *dst = dep.getDestination();
+    unsigned latency = *prob.getLatency(*prob.getLinkedOperatorType(src));
+    row.back() = -latency; // note the negation
+    if (src !=
+        dst) { // note that these coefficients just zero out in self-arcs.
+      row[opToVar[src]] = -1;
+      row[opToVar[dst]] = 1;
+    }
+  }
+
+  /// Create a cyclic latency constraint representing the given dependence and
+  /// the given dependence distance.
+  /// The constraint is represented as:
+  /// dstOpStartTime >= srcOpStartTime + latency + II * distance.
+  void createCyclicLatencyConstraint(MutableArrayRef<MPInt> row,
+                                     Problem::Dependence dep,
+                                     const MPInt &distance) {
+    createLatencyConstraint(row, dep);
+    row[0] = distance;
+  }
+
+  /// Add a constant lower bound on the variable at position `var`, representing
+  /// the constraint: `var >= bound`.
+  void addLowerBound(unsigned var, const MPInt &bound) {
+    SmallVector<MPInt, 8> row(getNumCols());
+    row[var] = 1;
+    row.back() = -bound;
+    addInequality(row);
+  }
+
+  /// Fix the variable at position `var` to a constant, representing the
+  /// constraint: `var == bound`.
+  void addEqBound(unsigned var, const MPInt &bound) {
+    SmallVector<MPInt, 8> row(getNumCols());
+    row[var] = 1;
+    row.back() = -bound;
+    addEquality(row);
+  }
+
+  // Solve the problem, keeping II integer, but allowing the solutions can be
+  // rational. We use a rational lexicographic simplex solver to do this.
+  // To keep II integer, if we obtain a rational II, we fix the II to the
+  // ceiling of the rational II. Since any II greater than the minimum II
+  // is valid, this is a valid solution.
+  MaybeOptimum<SmallVector<Fraction, 8>> solveRationally() {
+    MaybeOptimum<SmallVector<Fraction, 8>> sample =
+        LexSimplex::findRationalLexMin();
+
+    if (!sample.isBounded())
+      return sample;
+
+    ArrayRef<Fraction> res = *sample;
+
+    // If we have an integer II, we can just return the solution.
+    Fraction ii = res[0];
+    if (ii.num % ii.den == 0) {
+      return sample;
+    }
+
+    // We have a rational solution for II. We fix II to the ceiling of the
+    // given solution.
+    addEqBound(0, ceil(ii));
+
+    sample = LexSimplex::findRationalLexMin();
+    assert(sample.isBounded() && "Rounded up II should be feasible");
+
+    return sample;
+  }
+
+private:
+  // Reference to the problem.
+  Problem &prob;
+
+  /// A mapping from operation to their index in the simplex.
+  DenseMap<Operation *, unsigned> opToVar;
+};
+
+}; // namespace
+
+LogicalResult scheduling::schedulePresburger(Problem &prob, Operation *lastOp) {
+  Solver solver(prob, 1);
+
+  // II = 0 for acyclic problems.
+  solver.addEqBound(0, MPInt(0));
+
+  // There is a single objective, minimize the last operation.
+  {
+    SmallVector<MPInt, 8> row(solver.getNumCols());
+    row[1] = 1;
+    row[solver.getOpIndex(lastOp)] = -1;
+    solver.addEquality(row);
+  }
+
+  // Setup constraints for dependencies.
+  {
+    SmallVector<MPInt, 8> row(solver.getNumCols());
+    for (auto *op : prob.getOperations()) {
+      for (auto &dep : prob.getDependences(op)) {
+        solver.createLatencyConstraint(row, dep);
+        solver.addInequality(row);
+        std::fill(row.begin(), row.end(), MPInt(0));
+      }
+    }
+  }
+
+  // The constraints from dependence are built in a way that the solution always
+  // has integer rational start times. So, we can solve rationally keeping II
+  // integer.
+  auto res = solver.solveRationally();
+  if (!res.isBounded())
+    return prob.getContainingOp()->emitError() << "problem is infeasible";
+
+  auto sample = *res;
+
+  for (auto *op : prob.getOperations())
+    prob.setStartTime(op,
+                      int64_t(sample[solver.getOpIndex(op)].getAsInteger()));
+
+  return success();
+}
+
+LogicalResult scheduling::schedulePresburger(CyclicProblem &prob,
+                                             Operation *lastOp) {
+  Solver solver(prob, 1);
+
+  // II >= 1 for cyclic problems.
+  solver.addLowerBound(0, MPInt(1));
+
+  // There is a single objective, minimize the last operation.
+  {
+    SmallVector<MPInt, 8> row(solver.getNumCols());
+    row[1] = 1;
+    row[solver.getOpIndex(lastOp)] = -1;
+    solver.addEquality(row);
+  }
+
+  // Setup constraints for dependencies.
+  {
+    SmallVector<MPInt, 8> row(solver.getNumCols());
+    for (auto *op : prob.getOperations()) {
+      for (auto &dep : prob.getDependences(op)) {
+        if (auto dist = prob.getDistance(dep))
+          solver.createCyclicLatencyConstraint(row, dep, MPInt(*dist));
+        else
+          solver.createLatencyConstraint(row, dep);
+        solver.addInequality(row);
+        std::fill(row.begin(), row.end(), MPInt(0));
+      }
+    }
+  }
+
+  // The constraints from dependence are built in a way that the solution always
+  // has integer rational start times. So, we can solve rationally keeping II
+  // integer.
+  auto res = solver.solveRationally();
+  if (!res.isBounded())
+    return prob.getContainingOp()->emitError() << "problem is infeasible";
+
+  auto sample = *res;
+
+  prob.setInitiationInterval(int64_t(sample[0].getAsInteger()));
+  for (auto *op : prob.getOperations())
+    prob.setStartTime(op,
+                      int64_t(sample[solver.getOpIndex(op)].getAsInteger()));
+
+  return success();
+}

lib/Scheduling/TestPasses.cpp

@@ -558,6 +558,76 @@ void TestSimplexSchedulerPass::runOnOperation() {
   llvm_unreachable("Unsupported scheduling problem");
 }
 
+//===----------------------------------------------------------------------===//
+// PresburgerScheduler
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct TestPresburgerSchedulerPass
+    : public PassWrapper<TestPresburgerSchedulerPass,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPresburgerSchedulerPass)
+
+  TestPresburgerSchedulerPass() = default;
+  TestPresburgerSchedulerPass(const TestPresburgerSchedulerPass &) {}
+  Option<std::string> problemToTest{*this, "with", llvm::cl::init("Problem")};
+  void runOnOperation() override;
+  StringRef getArgument() const override { return "test-presburger-scheduler"; }
+  StringRef getDescription() const override {
+    return "Emit a presburger scheduler's solution as attributes";
+  }
+};
+} // anonymous namespace
+
+void TestPresburgerSchedulerPass::runOnOperation() {
+  auto func = getOperation();
+  Operation *lastOp = func.getBlocks().front().getTerminator();
+  OpBuilder builder(func.getContext());
+
+  if (problemToTest == "Problem") {
+    auto prob = Problem::get(func);
+    constructProblem(prob, func);
+    assert(succeeded(prob.check()));
+
+    if (failed(schedulePresburger(prob, lastOp))) {
+      func->emitError("scheduling failed");
+      return signalPassFailure();
+    }
+
+    if (failed(prob.verify())) {
+      func->emitError("schedule verification failed");
+      return signalPassFailure();
+    }
+
+    emitSchedule(prob, "simplexStartTime", builder);
+    return;
+  }
+
+  if (problemToTest == "CyclicProblem") {
+    auto prob = CyclicProblem::get(func);
+    constructProblem(prob, func);
+    constructCyclicProblem(prob, func);
+    assert(succeeded(prob.check()));
+
+    if (failed(schedulePresburger(prob, lastOp))) {
+      func->emitError("scheduling failed");
+      return signalPassFailure();
+    }
+
+    if (failed(prob.verify())) {
+      func->emitError("schedule verification failed");
+      return signalPassFailure();
+    }
+
+    func->setAttr("simplexInitiationInterval",
+                  builder.getI32IntegerAttr(*prob.getInitiationInterval()));
+    emitSchedule(prob, "simplexStartTime", builder);
+    return;
+  }
+
+  llvm_unreachable("Unsupported scheduling problem");
+}
+
 //===----------------------------------------------------------------------===//
 // LPScheduler
 //===----------------------------------------------------------------------===//
@@ -703,6 +773,9 @@ void registerSchedulingTestPasses() {
   mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
     return std::make_unique<TestSimplexSchedulerPass>();
   });
+  mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
+    return std::make_unique<TestPresburgerSchedulerPass>();
+  });
 #ifdef SCHEDULING_OR_TOOLS
   mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
     return std::make_unique<TestLPSchedulerPass>();

test/Scheduling/cyclic-problems.mlir

@@ -1,5 +1,6 @@
 // RUN: circt-opt %s -test-cyclic-problem
 // RUN: circt-opt %s -test-simplex-scheduler=with=CyclicProblem | FileCheck %s -check-prefix=SIMPLEX
+// RUN: circt-opt %s -test-presburger-scheduler=with=CyclicProblem | FileCheck %s -check-prefix=SIMPLEX
 
 // SIMPLEX-LABEL: cyclic
 // SIMPLEX-SAME: simplexInitiationInterval = 2

test/Scheduling/problems.mlir

@@ -1,6 +1,7 @@
 // RUN: circt-opt %s -test-scheduling-problem -allow-unregistered-dialect
 // RUN: circt-opt %s -test-asap-scheduler -allow-unregistered-dialect | FileCheck %s -check-prefix=ASAP
 // RUN: circt-opt %s -test-simplex-scheduler=with=Problem -allow-unregistered-dialect | FileCheck %s -check-prefix=SIMPLEX
+// RUN: circt-opt %s -test-presburger-scheduler=with=Problem -allow-unregistered-dialect | FileCheck %s -check-prefix=SIMPLEX
 
 // ASAP-LABEL: unit_latencies
 // SIMPLEX-LABEL: unit_latencies