[Scheduling] Set up infrastructure for using OR-Tools' solvers. (#2465)

CMakeLists.txt

@@ -295,6 +295,26 @@ else()
   endif()
 endif()
 
+#-------------------------------------------------------------------------------
+# OR-Tools Configuration
+#-------------------------------------------------------------------------------
+
+option(OR_TOOLS_DISABLE "Disable OR-Tools.")
+if (OR_TOOLS_DISABLE)
+  message(STATUS "Disabling OR-Tools.")
+else()
+  if(DEFINED OR_TOOLS_PATH)
+    list(APPEND CMAKE_PREFIX_PATH ${OR_TOOLS_PATH}/lib/cmake)
+  endif()
+
+  find_package(ortools CONFIG)
+
+  if (ortools_FOUND)
+    list(APPEND CMAKE_INSTALL_RPATH ${ortools_LIBDIR})
+    set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE)
+  endif()
+endif()
+
 #-------------------------------------------------------------------------------
 # llhd-sim Configuration
 #-------------------------------------------------------------------------------

include/circt/Scheduling/Algorithms.h

@@ -50,6 +50,11 @@ LogicalResult scheduleSimplex(SharedOperatorsProblem &prob, Operation *lastOp);
 /// that do not include at least one edge with a non-zero distance.
 LogicalResult scheduleSimplex(ModuloProblem &prob, Operation *lastOp);
 
+/// Solve the basic problem using linear programming and an external LP solver.
+/// The objective is to minimize the start time of the given \p lastOp. Fails if
+/// the dependence graph contains cycles.
+LogicalResult scheduleLP(Problem &prob, Operation *lastOp);
+
 } // namespace scheduling
 } // namespace circt
 

lib/Scheduling/CMakeLists.txt

@@ -1,21 +1,39 @@
 set(LLVM_OPTIONAL_SOURCES
   ASAPScheduler.cpp
+  LPSchedulers.cpp
   Problems.cpp
   SimplexSchedulers.cpp
   TestPasses.cpp
   )
 
-add_circt_library(CIRCTScheduling
+set(SCHEDULING_SOURCES
   ASAPScheduler.cpp
   Problems.cpp
   SimplexSchedulers.cpp
+  )
 
-  LINK_LIBS PUBLIC
+set(SCHEDULING_LIBS
   MLIRIR
   MLIRStandard
   MLIRSupport
   )
 
+if(ortools_FOUND)
+  option(SCHEDULING_OR_TOOLS "Enable schedulers relying on an external solver from OR-Tools" ON)
+endif()
+
+if(SCHEDULING_OR_TOOLS)
+  list(APPEND SCHEDULING_SOURCES LPSchedulers.cpp)
+  list(APPEND SCHEDULING_LIBS ortools::ortools)
+endif()
+
+add_circt_library(CIRCTScheduling
+  ${SCHEDULING_SOURCES}
+
+  LINK_LIBS PUBLIC
+  ${SCHEDULING_LIBS}
+  )
+
 add_circt_library(CIRCTSchedulingTestPasses
   TestPasses.cpp
 
@@ -23,3 +41,7 @@ add_circt_library(CIRCTSchedulingTestPasses
   CIRCTScheduling
   MLIRPass
   )
+
+if(SCHEDULING_OR_TOOLS)
+  target_compile_definitions(obj.CIRCTSchedulingTestPasses PRIVATE SCHEDULING_OR_TOOLS)
+endif()

lib/Scheduling/LPSchedulers.cpp

@@ -0,0 +1,77 @@
+//===- LPSchedulers.cpp - Schedulers using external LP solvers ------------===//
+//
+// 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 using external solvers
+// via OR-Tools.
+//
+//===----------------------------------------------------------------------===//
+
+#include "circt/Scheduling/Algorithms.h"
+
+#include "mlir/IR/Operation.h"
+
+#include "ortools/linear_solver/linear_solver.h"
+
+using namespace circt;
+using namespace circt::scheduling;
+using namespace operations_research;
+
+LogicalResult scheduling::scheduleLP(Problem &prob, Operation *lastOp) {
+  Operation *containingOp = prob.getContainingOp();
+
+  MPSolver::OptimizationProblemType problemType;
+  if (!MPSolver::ParseSolverType("GLOP_LINEAR_PROGRAMMING", &problemType) ||
+      !MPSolver::SupportsProblemType(problemType))
+    return containingOp->emitError("Solver is unvailable");
+
+  MPSolver solver("Problem", problemType);
+  double infinity = solver.infinity();
+
+  // Create start time variables.
+  DenseMap<Operation *, MPVariable *> vars;
+  unsigned i = 0;
+  for (auto *op : prob.getOperations()) {
+    vars[op] = solver.MakeNumVar(0, infinity, (Twine("t_") + Twine(i)).str());
+    ++i;
+  }
+
+  // The objective is to minimize the start time of the last operation.
+  MPObjective *objective = solver.MutableObjective();
+  objective->SetCoefficient(vars[lastOp], 1);
+  objective->SetMinimization();
+
+  // Construct a linear constraint for each dependence.
+  for (auto *op : prob.getOperations())
+    for (auto dep : prob.getDependences(op)) {
+      Operation *src = dep.getSource();
+      Operation *dst = dep.getDestination();
+      if (src == dst)
+        return containingOp->emitError() << "dependence cycle detected";
+
+      //     t_src + t.linkedOperatorType.latency <= t_dst
+      // <=> 1 * t_src + -1 * t_dst <= -latency
+      unsigned latency = *prob.getLatency(*prob.getLinkedOperatorType(src));
+      MPConstraint *constraint =
+          solver.MakeRowConstraint(-infinity, -((double)latency));
+      constraint->SetCoefficient(vars[src], 1);
+      constraint->SetCoefficient(vars[dst], -1);
+    }
+
+  // Invoke solver. The LP is infeasible if the scheduling problem contained
+  // dependence cycles. Otherwise, we expect the result to be optimal.
+  MPSolver::ResultStatus result = solver.Solve();
+  if (result == MPSolver::INFEASIBLE)
+    return containingOp->emitError() << "dependence cycle detected";
+  assert(result == MPSolver::OPTIMAL);
+
+  // Retrieve start times.
+  for (auto *op : prob.getOperations())
+    prob.setStartTime(op, std::round(vars[op]->solution_value()));
+
+  return success();
+}

lib/Scheduling/TestPasses.cpp

@@ -497,6 +497,55 @@ void TestSimplexSchedulerPass::runOnFunction() {
   llvm_unreachable("Unsupported scheduling problem");
 }
 
+//===----------------------------------------------------------------------===//
+// LPScheduler
+//===----------------------------------------------------------------------===//
+
+#ifdef SCHEDULING_OR_TOOLS
+
+namespace {
+struct TestLPSchedulerPass
+    : public PassWrapper<TestLPSchedulerPass, FunctionPass> {
+  TestLPSchedulerPass() = default;
+  TestLPSchedulerPass(const TestLPSchedulerPass &) {}
+  Option<std::string> problemToTest{*this, "with", llvm::cl::init("Problem")};
+  void runOnFunction() override;
+  StringRef getArgument() const override { return "test-lp-scheduler"; }
+  StringRef getDescription() const override {
+    return "Emit an LP scheduler's solution as attributes";
+  }
+};
+} // anonymous namespace
+
+void TestLPSchedulerPass::runOnFunction() {
+  auto func = getFunction();
+  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(scheduleLP(prob, lastOp))) {
+      func->emitError("scheduling failed");
+      return signalPassFailure();
+    }
+
+    if (failed(prob.verify())) {
+      func->emitError("schedule verification failed");
+      return signalPassFailure();
+    }
+
+    emitSchedule(prob, "lpStartTime", builder);
+    return;
+  }
+
+  llvm_unreachable("Unsupported scheduling problem");
+}
+
+#endif // SCHEDULING_OR_TOOLS
+
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
@@ -525,6 +574,11 @@ void registerSchedulingTestPasses() {
   mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
     return std::make_unique<TestSimplexSchedulerPass>();
   });
+#ifdef SCHEDULING_OR_TOOLS
+  mlir::registerPass([]() -> std::unique_ptr<::mlir::Pass> {
+    return std::make_unique<TestLPSchedulerPass>();
+  });
+#endif
 }
 } // namespace test
 } // namespace circt

test/Scheduling/or-tools/problems.mlir

@@ -0,0 +1,50 @@
+// REQUIRES: or-tools
+// RUN: circt-opt %s -test-lp-scheduler=with=Problem -allow-unregistered-dialect | FileCheck %s -check-prefix=LP
+
+// LP-LABEL: unit_latencies
+func @unit_latencies(%a1 : i32, %a2 : i32, %a3 : i32, %a4 : i32) -> i32 {
+  %0 = arith.addi %a1, %a2 : i32
+  %1 = arith.addi %0, %a3 : i32
+  %2:3 = "more.results"(%0, %1) : (i32, i32) -> (i32, i32, i32)
+  %3 = arith.addi %a4, %2#1 : i32
+  %4 = arith.addi %2#0, %2#2 : i32
+  %5 = arith.addi %3, %3 : i32
+  %6 = "more.operands"(%3, %4, %5) : (i32, i32, i32) -> i32
+  // LP: return
+  // LP-SAME: lpStartTime = 6
+  return %6 : i32
+}
+
+// LP-LABEL: arbitrary_latencies
+func @arbitrary_latencies(%v : complex<f32>) -> f32 attributes {
+  operatortypes = [
+    { name = "extr", latency = 0 },
+    { name = "add", latency = 3 },
+    { name = "mult", latency = 6 },
+    { name = "sqrt", latency = 10 }
+  ] } {
+  %0 = "complex.re"(%v) { opr = "extr" } : (complex<f32>) -> f32
+  %1 = "complex.im"(%v) { opr = "extr" } : (complex<f32>) -> f32
+  %2 = arith.mulf %0, %0 { opr = "mult" } : f32
+  %3 = arith.mulf %1, %1 { opr = "mult" } : f32
+  %4 = arith.addf %2, %3 { opr = "add" } : f32
+  %5 = "math.sqrt"(%4) { opr = "sqrt" } : (f32) -> f32
+  // LP: return
+  // LP-SAME: lpStartTime = 19
+  return %5 : f32
+}
+
+// LP-LABEL: auxiliary_dependences
+func @auxiliary_dependences() attributes { auxdeps = [
+    [0,1], [0,2], [2,3], [3,4], [3,6], [4,5], [5,6]
+  ] } {
+  %0 = arith.constant 0 : i32
+  %1 = arith.constant 1 : i32
+  %2 = arith.constant 2 : i32
+  %3 = arith.constant 3 : i32
+  %4 = arith.constant 4 : i32
+  %5 = arith.constant 5 : i32
+  // LP: return
+  // LP-SAME: lpStartTime = 5
+  return { problemStartTime = 6 }
+}

test/lit.cfg.py

@@ -70,6 +70,10 @@ if config.verilator_path != "":
 if config.esi_capnp != "":
   config.available_features.add('capnp')
 
+# Enable tests for schedulers relying on an external solver from OR-Tools.
+if config.scheduling_or_tools != "":
+  config.available_features.add('or-tools')
+
 # Add llhd-sim if it is built.
 if config.llhd_sim_enabled:
   config.available_features.add('llhd-sim')

test/lit.site.cfg.py.in

@@ -37,6 +37,7 @@ config.circt_tools_dir = "@CIRCT_TOOLS_DIR@"
 config.circt_shlib_dir = "@LLVM_LIBRARY_OUTPUT_INTDIR@"
 config.verilator_path = "@VERILATOR_PATH@"
 config.esi_capnp = "@ESI_CAPNP@"
+config.scheduling_or_tools = "@SCHEDULING_OR_TOOLS@"
 config.llhd_sim_enabled = @CIRCT_LLHD_SIM_ENABLED@
 
 # Support substitution of the tools_dir with user parameters. This is