[mlir][Analysis][NFC] Reimplement FlatAffineConstraints::addLowerOrUpperBound

[springerm]

Summary

Reimplement this function in terms of the function variant without Value semantics.

Depends On D107727

[ftynse]

mlir/lib/Analysis/AffineStructures.cpp

2097

2103–2104

Nit: please add assertion messages.

[springerm]

rebase

[nicolasvasilache]

Nice cleanup!

mlir/lib/Analysis/AffineStructures.cpp

for (unsigned r = 0, e = flatExprs.size(); r < e; r++) {
 
     // Add the equality connecting the result of the map to this constraint set.
     addEquality(eqToAdd);
   }
 
   return success();
 }
 
-// Turn a dimension into a symbol.
-static void turnDimIntoSymbol(FlatAffineValueConstraints *cst, Value id) {
-  unsigned pos;
-  if (cst->findId(id, &pos) && pos < cst->getNumDimIds()) {
-    cst->swapId(pos, cst->getNumDimIds() - 1);
-    cst->setDimSymbolSeparation(cst->getNumSymbolIds() + 1);
-  }
-}
-
 // Turn a symbol into a dimension.
 static void turnSymbolIntoDim(FlatAffineValueConstraints *cst, Value id) {
   unsigned pos;
   if (cst->findId(id, &pos) && pos >= cst->getNumDimIds() &&
       pos < cst->getNumDimAndSymbolIds()) {
     cst->swapId(pos, cst->getNumDimIds());
     cst->setDimSymbolSeparation(cst->getNumSymbolIds() - 1);
   }
LogicalResult FlatAffineConstraints::addLowerOrUpperBound(unsigned pos,
   }
 
   return success();
 }
 
 LogicalResult FlatAffineValueConstraints::addLowerOrUpperBound(
     unsigned pos, AffineMap boundMap, ValueRange boundOperands, bool eq,
     bool lower) {
-  assert(pos < getNumDimAndSymbolIds() && "invalid position");
-  // Equality follows the logic of lower bound except that we add an equality
-  // instead of an inequality.
-  assert((!eq || boundMap.getNumResults() == 1) && "single result expected");
-  if (eq)
-    lower = true;
-
   // Fully compose map and operands; canonicalize and simplify so that we
   // transitively get to terminal symbols or loop IVs.
   auto map = boundMap;
   SmallVector<Value, 4> operands(boundOperands.begin(), boundOperands.end());
   fullyComposeAffineMapAndOperands(&map, &operands);
   map = simplifyAffineMap(map);
   canonicalizeMapAndOperands(&map, &operands);
   for (auto operand : operands)
     addInductionVarOrTerminalSymbol(operand);
 
-  FlatAffineValueConstraints localVarCst;
-  std::vector<SmallVector<int64_t, 8>> flatExprs;
-  if (failed(getFlattenedAffineExprs(map, &flatExprs, &localVarCst))) {
-    LLVM_DEBUG(llvm::dbgs() << "semi-affine expressions not yet supported\n");
-    return failure();
-  }
-
-  // Merge and align with localVarCst.
-  if (localVarCst.getNumLocalIds() > 0) {
-    // Set values for localVarCst.
-    localVarCst.setIdValues(0, localVarCst.getNumDimAndSymbolIds(), operands);
-    for (auto operand : operands) {
-      unsigned pos;
-      if (findId(operand, &pos)) {
-        if (pos >= getNumDimIds() && pos < getNumDimAndSymbolIds()) {
-          // If the local var cst has this as a dim, turn it into its symbol.
-          turnDimIntoSymbol(&localVarCst, operand);
-        } else if (pos < getNumDimIds()) {
-          // Or vice versa.
-          turnSymbolIntoDim(&localVarCst, operand);
-        }
-      }
-    }
-    mergeAndAlignIds(/*offset=*/0, this, &localVarCst);
-    append(localVarCst);
-  }
-
-  // Record positions of the operands in the constraint system. Need to do
-  // this here since the constraint system changes after a bound is added.
-  SmallVector<unsigned, 8> positions;
-  unsigned numOperands = operands.size();
-  for (auto operand : operands) {
-    unsigned pos;
-    if (!findId(operand, &pos))
-      assert(0 && "expected to be found");
-    positions.push_back(pos);
-  }
-
-  for (const auto &flatExpr : flatExprs) {
-    // Invalid bound: pos appears among the operands.
-    if (llvm::find(positions, pos) != positions.end())
-      continue;
+  SmallVector<Value> dims, syms;
+#ifndef NDEBUG
+  SmallVector<Value> newSyms;
+  SmallVector<Value> *newSymsPtr = &newSyms;
+#else
+  SmallVector<Value> *newSymsPtr = nullptr;
+#endif // NDEBUG
+
+  dims.reserve(numDims);
+  syms.reserve(numSymbols);
+  for (unsigned i = 0; i < numDims; ++i)
+    dims.push_back(ids[i] ? *ids[i] : Value());
+  for (unsigned i = numDims; i < numDims + numSymbols; ++i)
+    syms.push_back(ids[i] ? *ids[i] : Value());
+
+  AffineMap alignedMap =
+      alignAffineMapWithValues(map, operands, dims, syms, newSymsPtr);
+  // All symbols are already part of this FlatAffineConstraints.
+  assert(syms.size() == newSymsPtr->size());
+  assert(std::equal(syms.begin(), syms.end(), newSymsPtr->begin()));
 
-    SmallVector<int64_t, 4> ineq(getNumCols(), 0);
-    ineq[pos] = lower ? 1 : -1;
-    // Dims and symbols.
-    for (unsigned j = 0, e = map.getNumInputs(); j < e; j++) {
-      ineq[positions[j]] = lower ? -flatExpr[j] : flatExpr[j];
-    }
-    // Copy over the local id coefficients.
-    unsigned numLocalIds = flatExpr.size() - 1 - numOperands;
-    for (unsigned jj = 0, j = getNumIds() - numLocalIds; jj < numLocalIds;
-         jj++, j++) {
-      ineq[j] =
-          lower ? -flatExpr[numOperands + jj] : flatExpr[numOperands + jj];
-    }
-    // Constant term.
-    ineq[getNumCols() - 1] =
-        lower ? -flatExpr[flatExpr.size() - 1]
-              // Upper bound in flattenedExpr is an exclusive one.
-              : flatExpr[flatExpr.size() - 1] - 1;
-    eq ? addEquality(ineq) : addInequality(ineq);
-  }
-  return success();
+  return addLowerOrUpperBound(pos, alignedMap, eq, lower);
 }
 
 // Adds slice lower bounds represented by lower bounds in 'lbMaps' and upper
 // bounds in 'ubMaps' to each value in `values' that appears in the constraint
 // system. Note that both lower/upper bounds share the same operand list
 // 'operands'.
 // This function assumes 'values.size' == 'lbMaps.size' == 'ubMaps.size', and
 // skips any null AffineMaps in 'lbMaps' or 'ubMaps'.

 void FlatAffineConstraints::print(raw_ostream &os) const {
   assert(hasConsistentState());
   os << "\nConstraints (" << getNumDimIds() << " dims, " << getNumSymbolIds()
      << " symbols, " << getNumLocalIds() << " locals), (" << getNumConstraints()
      << " constraints)\n";
   os << "(";
   for (unsigned i = 0, e = getNumIds(); i < e; i++) {
     if (auto *valueConstrSet =
             dynamic_cast<const FlatAffineValueConstraints *>(this)) {
       if (valueConstrSet->hasIdValue(i))
         os << "Value ";
       else
         os << "None ";
     } else {
       os << "None ";
     }
   }
void FlatAffineConstraints::removeTrivialRedundancy() {
 
   // TODO: consider doing this for equalities as well, but probably not worth
   // the savings.
 }
 
 void FlatAffineConstraints::clearAndCopyFrom(
     const FlatAffineConstraints &other) {
   if (auto *otherValueSet =
           dynamic_cast<const FlatAffineValueConstraints *>(&other))
     assert(!otherValueSet->hasIdValues() &&
            "cannot copy associated Values into FlatAffineConstraints");
   *this = other;
 }
 
 void FlatAffineValueConstraints::clearAndCopyFrom(
     const FlatAffineConstraints &other) {
   if (auto *otherValueSet =
           dynamic_cast<const FlatAffineValueConstraints *>(&other)) {
     *this = *otherValueSet;
   } else {
     *static_cast<FlatAffineConstraints *>(this) = other;
     ids.clear();
     ids.resize(numIds, None);
   }
 }