[mlir] Factor out constraint set creation from hoist padding.

[nicolasvasilache]

Summary

This revision adds a

FlatAffineValueConstraints(ValueRange ivs, ValueRange lbs, ValueRange ubs)

method and use it in hoist padding.

[ftynse]

mlir/include/mlir/Analysis/AffineStructures.h

586

I wouldn't drop these, somebody may be using them.

604

I'd consider a named static constructor instead, something like FlatAffineValueConstraints::getHyperrectangular that makes the intent clear.

635

Is this reflow necessary?

[springerm]

mlir/lib/Analysis/AffineStructures.cpp

213–215

can use appendDimId instead of insertDimId

214

maybe write: unsigned lbsStart = appendDimId(lbs)

Then you can reference it in the loop via AffineMap lb = ... getAffineDimExpr(lbsStart + ivIdx, ctx) ...

221

nit: /*dimCount=*/3 * nIvs, /*symbolCount=*/0

[nicolasvasilache]

Address comments.

[nicolasvasilache]

Fix stale commit.

mlir/include/mlir/Analysis/AffineStructures.h

FlatAffineValueConstraints(unsigned numDims = 0, unsigned numSymbols = 0,
       : FlatAffineValueConstraints(/*numReservedInequalities=*/0,
                                    /*numReservedEqualities=*/0,
                                    /*numReservedCols=*/numDims + numSymbols +
                                        numLocals + 1,
                                    numDims, numSymbols, numLocals, valArgs) {}
 
   /// Create a flat affine constraint system from an AffineValueMap or a list of
   /// these. The constructed system will only include equalities.
   explicit FlatAffineValueConstraints(const AffineValueMap &avm);
   explicit FlatAffineValueConstraints(ArrayRef<const AffineValueMap *> avmRef);
 
   /// Creates an affine constraint system from an IntegerSet.
   explicit FlatAffineValueConstraints(IntegerSet set);
 
   FlatAffineValueConstraints(ArrayRef<const AffineValueMap *> avmRef,
                              IntegerSet set);
 
+  // Construct a hyperrectangular constraint set from ValueRanges that represent
+  // induction variables, lower and upper bounds. `ivs`, `lbs` and `ubs` are
+  // expected to match one to one. The order of variables and constraints is:
+  //
+  // ivs | lbs | ubs | eq/ineq
+  // ----+-----+-----+---------
+  //   1   -1     0      >= 0
+  // ----+-----+-----+---------
+  //  -1    0     1      >= 0
+  //
+  // All dimensions as set as DimId.
+  static FlatAffineValueConstraints
+  getHyperrectangular(ValueRange ivs, ValueRange lbs, ValueRange ubs);
+
   /// Return the kind of this FlatAffineConstraints.
   Kind getKind() const override { return Kind::FlatAffineValueConstraints; }
 
   static bool classof(const FlatAffineConstraints *cst) {
     return cst->getKind() == Kind::FlatAffineValueConstraints;
   }
 
   /// Clears any existing data and reserves memory for the specified
public:
 
   /// Clones this object.
   std::unique_ptr<FlatAffineValueConstraints> clone() const;
 
   /// Adds constraints (lower and upper bounds) for the specified 'affine.for'
   /// operation's Value using IR information stored in its bound maps. The
   /// right identifier is first looked up using `forOp`'s Value. Asserts if the
   /// Value corresponding to the 'affine.for' operation isn't found in the
   /// constraint system. Returns failure for the yet unimplemented/unsupported
   /// cases.  Any new identifiers that are found in the bound operands of the
   /// 'affine.for' operation are added as trailing identifiers (either
   /// dimensional or symbolic depending on whether the operand is a valid
   /// symbol).
   //  TODO: add support for non-unit strides.
   LogicalResult addAffineForOpDomain(AffineForOp forOp);
 
   /// Adds constraints (lower and upper bounds) for each loop in the loop nest

mlir/lib/Analysis/AffineStructures.cpp

 }
 
 // Construct from an IntegerSet.
 FlatAffineValueConstraints::FlatAffineValueConstraints(IntegerSet set)
     : FlatAffineConstraints(set) {
   values.resize(numIds, None);
 }
 
+// Construct a hyperrectangular constraint set from ValueRanges that represent
+// induction variables, lower and upper bounds. `ivs`, `lbs` and `ubs` are
+// expected to match one to one. The order of variables and constraints is:
+//
+// ivs | lbs | ubs | eq/ineq
+// ----+-----+-----+---------
+//   1   -1     0      >= 0
+// ----+-----+-----+---------
+//  -1    0     1      >= 0
+//
+// All dimensions as set as DimId.
+FlatAffineValueConstraints
+FlatAffineValueConstraints::getHyperrectangular(ValueRange ivs, ValueRange lbs,
+                                                ValueRange ubs) {
+  FlatAffineValueConstraints res;
+  unsigned nIvs = ivs.size();
+  assert(nIvs == lbs.size() && "expected as many lower bounds as ivs");
+  assert(nIvs == ubs.size() && "expected as many upper bounds as ivs");
+
+  if (nIvs == 0)
+    return res;
+
+  res.appendDimId(ivs);
+  unsigned lbsStart = res.appendDimId(lbs);
+  unsigned ubsStart = res.appendDimId(ubs);
+
+  MLIRContext *ctx = ivs.front().getContext();
+  for (int ivIdx = 0, e = nIvs; ivIdx < e; ++ivIdx) {
+    // iv - lb >= 0
+    AffineMap lb = AffineMap::get(/*dimCount=*/3 * nIvs, /*symbolCount=*/0,
+                                  getAffineDimExpr(lbsStart + ivIdx, ctx));
+    if (failed(res.addBound(BoundType::LB, ivIdx, lb)))
+      llvm_unreachable("Unexpected FlatAffineValueConstraints creation error");
+    // -iv + ub >= 0
+    AffineMap ub = AffineMap::get(/*dimCount=*/3 * nIvs, /*symbolCount=*/0,
+                                  getAffineDimExpr(ubsStart + ivIdx, ctx));
+    if (failed(res.addBound(BoundType::UB, ivIdx, ub)))
+      llvm_unreachable("Unexpected FlatAffineValueConstraints creation error");
+  }
+  return res;
+}
+
 void FlatAffineConstraints::reset(unsigned numReservedInequalities,
                                   unsigned numReservedEqualities,
                                   unsigned newNumReservedCols,
                                   unsigned newNumDims, unsigned newNumSymbols,
                                   unsigned newNumLocals) {
   assert(newNumReservedCols >= newNumDims + newNumSymbols + newNumLocals + 1 &&
          "minimum 1 column");
   *this = FlatAffineConstraints(numReservedInequalities, numReservedEqualities,

mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp

for (Operation *user : forOp.getInductionVar().getUsers()) {
       }
     }
   }
 
   // The analysis is valid and hoisting can occur.
   valid = true;
 }
 
-/// Given a set of loops, assumed to be scf::ForOp, create a constraint set
-/// containing the inequalities `iv - lb >= 0` and `-iv + ub - 1 >= 0` for each
-/// loop. The order of the constraints follows:
-///
-///  ivs | lbs | ubs | eq/ineq
-///  ----+-----+-----+---------
-///   1    -1     0      >= 0
-///  ----+-----+-----+---------
-///  -1    0      1      >= 0
-///
-static FlatAffineValueConstraints
-initLoopIvsAndBounds(ArrayRef<scf::ForOp> loops) {
-  FlatAffineValueConstraints constraints;
-  // Append dims for all ivs, lbs, ubs: the order is important.
-  for (scf::ForOp op : loops)
-    constraints.appendDimId(op.getInductionVar());
-  for (scf::ForOp op : loops)
-    constraints.appendDimId(op.lowerBound());
-  for (scf::ForOp op : loops)
-    constraints.appendDimId(op.upperBound());
-  int numLoops = loops.size();
-  for (int ivIdx = 0, e = numLoops; ivIdx < e; ++ivIdx) {
-    // iv - lb >= 0
-    SmallVector<int64_t, 8> ineqLb(constraints.getNumCols(), 0);
-    ineqLb[ivIdx] = 1;
-    ineqLb[ivIdx + numLoops] = -1;
-    // -iv + ub >= 0
-    SmallVector<int64_t, 8> ineqUb(constraints.getNumCols(), 0);
-    ineqUb[ivIdx] = -1;
-    ineqUb[ivIdx + 2 * numLoops] = 1;
-    ineqUb[constraints.getNumCols() - 1] = -1;
-    constraints.addInequality(ineqLb);
-    constraints.addInequality(ineqUb);
-  }
-  return constraints;
-}
-
 static bool isDefinedOutsideOrConstant(scf::ForOp outer, Value v) {
   return outer.isDefinedOutsideOfLoop(v) || v.getDefiningOp<ConstantOp>();
 }
 
 /// For each loop in `loops`, determine the ops involved in the construction of
 /// its upper bound---up to the outerLimit loop--- and fold them as new
 /// inequalities in the constraint set.
 /// This is achieved by computing the backwardSlice of the loop's upper bound
for (Value v : toProjectOut)
     constraints.projectOut(v);
   return success();
 }
 
 // Footprint of the packedTensor, computed from the packingLoops and
 // `backwardSlice`.
 FailureOr<SmallVector<Value>>
 HoistingAnalysis::getPackedTensorSizes(ImplicitLocOpBuilder &b) {
-  FlatAffineValueConstraints constraints =
-      initLoopIvsAndBounds(packingLoops.getArrayRef());
+  // Create the base affine constaints for the packedLoops.
+  auto constraints = FlatAffineValueConstraints::getHyperrectangular(
+      llvm::to_vector<8>(llvm::map_range(
+          packingLoops, [](scf::ForOp op) { return op.getInductionVar(); })),
+      llvm::to_vector<8>(llvm::map_range(
+          packingLoops, [](scf::ForOp op) { return op.lowerBound(); })),
+      llvm::to_vector<8>(llvm::map_range(
+          packingLoops, [](scf::ForOp op) { return op.upperBound(); })));
+
+  // Iteratively try to fold the upper bounds into the constraints set.
   if (failed(foldUpperBoundsIntoConstraintsSet(
           constraints, outermostEnclosingForOp, packingLoops.getArrayRef())))
     return failure();
 
   int nPackedLoops = packingLoops.size();
   SmallVector<AffineMap> lbs(nPackedLoops), ubs(nPackedLoops);
   // Compute the bounds of the first positions, assuming the others are fixed.
   constraints.getSliceBounds(/*pos=*/0, /*num=*/nPackedLoops,