[Scheduling] [2/3] Add scheduler for shared pipelined operators problem: Tableau transformations. (#1621)

This implements the tableau transformations necessary to fix an operation's start time to a certain value.

lib/Scheduling/SimplexSchedulers.cpp

@@ -108,6 +108,10 @@ protected:
   /// variables.
   SmallVector<int> startTimeLocations;
 
+  /// Non-basic variables can be "frozen", which prevents them from being
+  /// pivoted into basis again.
+  DenseSet<unsigned> frozenVariables;
+
   /// Number of rows in the tableau = 1 + |deps|.
   unsigned nRows;
   /// Number of explicitly stored columns in the tableau = 3 + |ops|.
@@ -131,12 +135,16 @@ protected:
   virtual void fillConstraintRow(SmallVector<int> &row, detail::Dependence dep);
   void buildTableau();
   Optional<unsigned> findPivotRow();
-  Optional<unsigned> findPivotColumn(unsigned pivotRow);
+  Optional<unsigned> findPivotColumn(unsigned pivotRow,
+                                     bool allowPositive = false);
   void multiplyRow(unsigned row, int factor);
   void addMultipleOfRow(unsigned sourceRow, int factor, unsigned targetRow);
   void pivot(unsigned pivotRow, unsigned pivotColumn);
   LogicalResult solveTableau();
   bool isInBasis(unsigned startTimeVariable);
+  unsigned freeze(unsigned startTimeVariable);
+  void translate(unsigned column, int factor1, int factorS, int factorT);
+  LogicalResult scheduleAt(unsigned startTimeVariable, unsigned timeStep);
   void storeStartTimes();
 
   void dumpTableau();
@@ -256,19 +264,22 @@ Optional<unsigned> SimplexSchedulerBase::findPivotRow() {
   return None;
 }
 
-Optional<unsigned> SimplexSchedulerBase::findPivotColumn(unsigned pivotRow) {
+Optional<unsigned> SimplexSchedulerBase::findPivotColumn(unsigned pivotRow,
+                                                         bool allowPositive) {
   Optional<int> maxQuot;
   Optional<unsigned> pivotCol;
   // Look for negative entries in the ~A part of the tableau. If multiple
   // candidates exist, take the one with maximum of the quotient:
   // tableau[objectiveRow][col] / pivotCand
   for (unsigned col = firstNonBasicVariableColumn; col < nColumns; ++col) {
+    if (frozenVariables.count(
+            nonBasicVariables[col - firstNonBasicVariableColumn]))
+      continue;
     int pivotCand = tableau[pivotRow][col];
-    if (pivotCand < 0) {
+    if (pivotCand < 0 || (allowPositive && pivotCand > 0)) {
       // The ~A part of the tableau has only {-1, 0, 1} entries by construction.
-      assert(pivotCand == -1);
-      int quot = -tableau[objectiveRow][col];
-      // Quotient in general: tableau[objectiveRow][col] / pivotCand
+      assert(pivotCand * pivotCand == 1);
+      int quot = tableau[objectiveRow][col] / pivotCand;
       if (!maxQuot || quot > *maxQuot) {
         maxQuot = quot;
         pivotCol = col;
@@ -308,10 +319,9 @@ void SimplexSchedulerBase::pivot(unsigned pivotRow, unsigned pivotColumn) {
 
   int pivotElem = tableau[pivotRow][pivotColumn];
   // The ~A part of the tableau has only {-1, 0, 1} entries by construction.
-  // The pivot element must be negative, so it can only be -1.
-  assert(pivotElem == -1);
+  assert(pivotElem * pivotElem == 1);
   // Make `tableau[pivotRow][pivotColumn]` := 1
-  multiplyRow(pivotRow, -1); // Factor in general: 1 / pivotElement
+  multiplyRow(pivotRow, 1 / pivotElem);
 
   for (unsigned row = 0; row < nRows; ++row) {
     if (row == pivotRow)
@@ -322,7 +332,6 @@ void SimplexSchedulerBase::pivot(unsigned pivotRow, unsigned pivotColumn) {
       continue; // nothing to do
 
     // Make `tableau[row][pivotColumn]` := 0.
-    // Factor in general: -elem / pivotElem
     addMultipleOfRow(pivotRow, -elem, row);
   }
 
@@ -396,14 +405,107 @@ bool SimplexSchedulerBase::isInBasis(unsigned startTimeVariable) {
   llvm_unreachable("Invalid variable location");
 }
 
+// The following `freeze`, `translate` and `scheduleAt` methods are implemented
+// based on a follow-up publication to the paper mentioned above:
+//
+//   B. D. de Dinechin, "Fast Modulo Scheduling Under the Simplex Scheduling
+//   Framework", PRISM 1995.01, 1995.
+
+unsigned SimplexSchedulerBase::freeze(unsigned startTimeVariable) {
+  assert(startTimeVariable < startTimeLocations.size());
+  assert(!frozenVariables.count(startTimeVariable));
+
+  // Mark variable.
+  frozenVariables.insert(startTimeVariable);
+
+  if (!isInBasis(startTimeVariable))
+    // That's all for non-basic variables.
+    return startTimeLocations[startTimeVariable];
+
+  // We need to pivot this variable one out of basis.
+  unsigned pivotRow = -startTimeLocations[startTimeVariable];
+
+  // Here, positive pivot elements can be considered as well, hence finding a
+  // suitable column should not fail.
+  auto pivotCol = findPivotColumn(pivotRow, /* allowPositive= */ true);
+  assert(pivotCol);
+
+  // Perform the exchange.
+  pivot(pivotRow, *pivotCol);
+
+  // `startTimeVariable` is now represented by `pivotCol`.
+  return *pivotCol;
+}
+
+void SimplexSchedulerBase::translate(unsigned column, int factor1, int factorS,
+                                     int factorT) {
+  for (unsigned row = 0; row < nRows; ++row) {
+    auto &rowVec = tableau[row];
+    int elem = rowVec[column];
+    if (elem == 0)
+      continue;
+
+    rowVec[parameter1Column] += -elem * factor1;
+    rowVec[parameterSColumn] += -elem * factorS;
+    rowVec[parameterTColumn] += -elem * factorT;
+  }
+}
+
+LogicalResult SimplexSchedulerBase::scheduleAt(unsigned startTimeVariable,
+                                               unsigned timeStep) {
+  assert(startTimeVariable < startTimeLocations.size());
+  assert(!frozenVariables.count(startTimeVariable));
+
+  // Freeze variable and translate its column by parameter S.
+  unsigned frozenCol = freeze(startTimeVariable);
+  translate(frozenCol, /* factor1= */ 0, /* factorS= */ 1, /* factorT= */ 0);
+
+  // Temporarily set S to the desired value, and attempt to solve.
+  parameterS = timeStep;
+  auto res = solveTableau();
+  parameterS = 0;
+
+  if (failed(res)) {
+    // The LP is infeasible with the new constraint. We could try other values
+    // for S, but for now, we just roll back and signal failure to the driver.
+    translate(frozenCol, /* factor1= */ 0, /* factorS= */ -1, /* factorT= */ 0);
+    frozenVariables.erase(startTimeVariable);
+    res = solveTableau();
+    assert(succeeded(res));
+    return failure();
+  }
+
+  // Translate S by the other parameter(s). Currently, this means setting
+  // `factor1` to `timeStep`. For cyclic problems, one would perform a modulo
+  // decomposition: S = `factor1` + `factorT` * T, with `factor1` < T.
+  //
+  // This translation does not change the value of ~B_p u (the dot product of
+  // the the first three columns with parameters), hence we do not need to solve
+  // the tableau again.
+  //
+  // Note: I added a negation of the factors here, which is not mentioned in the
+  // paper's text, but apparently used in the example. Without it, the intended
+  // effect, i.e. making the S-column all-zero again, is not achieved.
+  assert(parameterT == 0);
+  translate(parameterSColumn, -timeStep, /* factorS= */ 1, /* factorT= */ 0);
+
+  // Record the start time.
+  Problem &prob = getProblem();
+  prob.setStartTime(prob.getOperations()[startTimeVariable], timeStep);
+
+  return success();
+}
+
 void SimplexSchedulerBase::storeStartTimes() {
   auto &prob = getProblem();
 
   for (auto *op : prob.getOperations()) {
     unsigned startTimeVar = startTimeVariables[op];
     if (!isInBasis(startTimeVar)) {
-      // Non-basic variables are 0 at the end of the simplex algorithm.
-      prob.setStartTime(op, 0);
+      // Non-basic variables that are not already fixed to a specific time step
+      // are 0 at the end of the simplex algorithm.
+      if (!frozenVariables.contains(startTimeVar))
+        prob.setStartTime(op, 0);
       continue;
     }
     // For the variables currently in basis, we look up the solution in the ~B