[mlir][linalg] Remove ConstraintsSet class

[springerm]

Summary

The same functionality can be implemented with FlatAffineValueConstraints.

Depends On D108185

[springerm]

rebase

[nicolasvasilache]

Great, thanks for the refactoring!

mlir/include/mlir/Dialect/Linalg/Analysis/ConstraintsSet.h

-//===- ConstraintsSet.h - Ext. for FlatAffineValueConstraints ---*- C++ -*-===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-//
-// Linalg-specific constraints set extensions.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_DIALECT_LINALG_ANALYSIS_CONSTRAINTS_SET_H_
-#define MLIR_DIALECT_LINALG_ANALYSIS_CONSTRAINTS_SET_H_
-
-#include "mlir/Analysis/AffineStructures.h"
-#include "mlir/IR/AffineMap.h"
-
-namespace mlir {
-class ValueRange;
-
-/// Linalg-specific constraints set extensions.
-class ConstraintsSet : public FlatAffineValueConstraints {
-public:
-  ConstraintsSet() : FlatAffineValueConstraints() {}
-
-  /// Assuming `val` is defined by `val = affine.min map (operands)`, introduce
-  /// all the constraints `val <= expr_i(operands)`, where expr_i are all the
-  /// results of `map`.
-  // This API avoids taking a dependence on the AffineMinOp definition.
-  LogicalResult composeMin(Value val, AffineMap map, ValueRange operands) {
-    return composeMinOrMaxMapAndOperands(val, map, operands, /*min=*/true);
-  }
-
-  /// Assuming `val` is defined by `val = affine.max map (operands)`, introduce
-  /// all the constraints `val >= expr_i(operands)`, where expr_i are all the
-  /// results of `map`.
-  // This API avoids taking a dependence on the AffineMaxOp definition.
-  LogicalResult composeMax(Value val, AffineMap map, ValueRange operands) {
-    return composeMinOrMaxMapAndOperands(val, map, operands, /*min=*/false);
-  }
-
-  /// Assuming `val` is defined by `val = affine.apply map (operands)`, call
-  /// composeMap.
-  // This API avoids taking a dependence on the AffineMApplyOp definition.
-  LogicalResult composeAffineApply(Value val, AffineMap map,
-                                   ValueRange operands);
-
-  /// Asserts the identifier `id` is in the constraints set and returns it.
-  unsigned lookupPos(Value id) const;
-
-  /// If v is not in the constraint set, insert it as a dim or symbol depending
-  /// on `asDim`.
-  /// Return success if v is of dim id type when `asDim` is true and of symbol
-  /// id type when `asDim` is false.
-  /// Return failure otherwise.
-  LogicalResult ensureIdOfType(Value v, bool asDim);
-
-private:
-  /// Implementation detail for composeMin/Max.
-  LogicalResult composeMinOrMaxMapAndOperands(Value val, AffineMap map,
-                                              ValueRange operands, bool min);
-};
-
-} // namespace mlir
-
-#endif // MLIR_DIALECT_LINALG_ANALYSIS_CONSTRAINTS_SET_H_

mlir/lib/Dialect/Linalg/Analysis/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgAnalysis
-  ConstraintsSet.cpp
   DependenceAnalysis.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
   MLIRIR
   MLIRLinalg
   MLIRLoopAnalysis
   MLIRMemRef
   MLIRStandard
   )

mlir/lib/Dialect/Linalg/Analysis/ConstraintsSet.cpp

-//===- ConstraintsSet.cpp - Extensions for FlatAffineConstraints ----------===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-//
-// Linalg-specific constraints set extensions.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Dialect/Linalg/Analysis/ConstraintsSet.h"
-#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
-#include "mlir/IR/AffineMap.h"
-
-using namespace mlir;
-
-unsigned ConstraintsSet::lookupPos(Value id) const {
-  unsigned pos;
-  if (!findId(id, &pos)) {
-    llvm::errs() << "Lookup failed: " << id << "\n";
-    llvm_unreachable("Lookup failed");
-  }
-  return pos;
-}
-
-LogicalResult ConstraintsSet::ensureIdOfType(Value v, bool asDim) {
-  if (!containsId(v)) {
-    if (asDim)
-      addDimId(getNumDimIds(), v);
-    else
-      addSymbolId(getNumSymbolIds(), v);
-    return success();
-  }
-  unsigned pos = lookupPos(v);
-  return success((asDim && pos < getNumDimIds()) ||
-                 (!asDim && getNumDimIds() <= pos &&
-                  pos < getNumDimIds() + getNumSymbolIds()));
-}
-
-LogicalResult ConstraintsSet::composeAffineApply(Value val, AffineMap map,
-                                                 ValueRange operands) {
-  AffineValueMap avm(map, operands, val);
-  return composeMap(&avm);
-}
-
-LogicalResult ConstraintsSet::composeMinOrMaxMapAndOperands(Value val,
-                                                            AffineMap map,
-                                                            ValueRange operands,
-                                                            bool min) {
-  ConstraintsSet localCst;
-  std::vector<SmallVector<int64_t, 8>> flatExprs;
-  if (failed(getFlattenedAffineExprs(map, &flatExprs, &localCst)))
-    return failure();
-  assert(flatExprs.size() == map.getNumResults() &&
-         "incorrect number of flattened expressiosn");
-
-  // Local vars on a per-need basis.
-  if (localCst.getNumLocalIds() != 0)
-    return failure();
-
-  // Add one inequality for each result connecting `val` to the other ids in
-  // `operands`. For instance, uf the expression is:
-  //   `16 * i0 + i1` and
-  //   `min` is true
-  // add:
-  //  -d_val + 16 * i0 + i1 >= 0.
-  for (const auto &flatExpr : flatExprs) {
-    assert(flatExpr.size() >= operands.size() + 1);
-    SmallVector<int64_t, 8> ineq(getNumCols(), 0);
-    for (unsigned i = 0, e = operands.size(); i < e; i++)
-      ineq[lookupPos(operands[i])] = min ? flatExpr[i] : -flatExpr[i];
-
-    // Set the coefficient for `d_val`.
-    ineq[lookupPos(val)] = min ? -1 : 1;
-
-    // Set the constant term (upper bound in flatExpr is exclusive).
-    ineq[getNumCols() - 1] = min ? flatExpr[flatExpr.size() - 1] - 1
-                                 : -flatExpr[flatExpr.size() - 1];
-
-    // Add the inequality connecting the result of the map to the rest.
-    addInequality(ineq);
-  }
-
-  return success();
-}

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

 //===- Hoisting.cpp - Linalg hoisting transformations ---------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements functions concerned with hoisting invariant operations
 // in the context of Linalg transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Affine/Utils.h"
-#include "mlir/Dialect/Linalg/Analysis/ConstraintsSet.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SCF/Utils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/Vector/VectorUtils.h"
static Value buildLoopIterationCount(OpBuilder &b, scf::ForOp outer,
   auto loc = forOp->getLoc();
   return b.createOrFold<AffineApplyOp>(loc, (iv - lb).ceilDiv(step),
                                        ValueRange{ivVal, lbVal, stepVal});
 }
 
 /// 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.
-static ConstraintsSet initLoopIvsAndBounds(ArrayRef<Operation *> loops) {
-  ConstraintsSet constraints;
+static FlatAffineValueConstraints
+initLoopIvsAndBounds(ArrayRef<Operation *> loops) {
+  FlatAffineValueConstraints constraints;
   for (Operation *op : loops)
     constraints.addDimId(constraints.getNumDimIds(),
                          cast<scf::ForOp>(op).getInductionVar());
   for (Operation *op : loops)
     constraints.addDimId(constraints.getNumDimIds(),
                          cast<scf::ForOp>(op).lowerBound());
   for (Operation *op : loops)
     constraints.addDimId(constraints.getNumDimIds(),

 /// constraints set.
 /// The following operations are supported:
 ///   - scf::ForOp are simply skipped.
 ///   - AffineApplyOp are composed to replace the result by an equality.
 ///   - AffineMinOp are composed by adding each entry as an upper bound.
 /// If any other operation is met, return failure.
 // TODO: extend on a per-need basis.
 static LogicalResult
-foldUpperBoundsIntoConstraintsSet(ConstraintsSet &constraints,
+foldUpperBoundsIntoConstraintsSet(FlatAffineValueConstraints &constraints,
                                   scf::ForOp outerLimit,
                                   ArrayRef<Operation *> loops) {
   SetVector<Value> toProjectOut;
   for (Operation *loop : loops) {
     auto ub = cast<scf::ForOp>(loop).upperBound();
     if (isDefinedOutsideOrConstant(outerLimit, ub))
       continue;
 
     // Compute a backward slice up to, but not including, `outerLimit`.
     SetVector<Operation *> backwardSlice;
     getBackwardSlice(ub, &backwardSlice, [&](Operation *op) {
       return outerLimit->isProperAncestor(op);
     });
     backwardSlice.insert(ub.getDefiningOp());
 
     // Iterate over all ops in the slice and compose them in the constraints.
     for (Operation *op : llvm::reverse(backwardSlice)) {
       if (!isa<scf::ForOp, AffineApplyOp, AffineMinOp>(op))
         return failure();
       if (isa<scf::ForOp>(op))
         continue;
       // Ensure there is a
       auto ensureIdFailed = [&](Value v) {
-        return failed(constraints.ensureIdOfType(v, /*asDim=*/true));
+        if (constraints.containsId(v)) {
+          unsigned pos;
+          constraints.findId(v, &pos);
+          return pos >= constraints.getNumDimIds();
+        }
+        constraints.addDimId(constraints.getNumDimIds(), v);
+        return false;
       };
 
       // Ensure all ids exist and add results for later projection.
       if (llvm::any_of(op->getResults(), ensureIdFailed) ||
           llvm::any_of(op->getOperands(), ensureIdFailed))
         return failure();
 
       // All supported ops have 1 result.
       // TODO: extend when needed.
       toProjectOut.insert(op->getResult(0));
 
       // Compose supported ops.
       if (auto affineApplyOp = dyn_cast<AffineApplyOp>(op)) {
-        if (failed(constraints.composeAffineApply(affineApplyOp.getResult(),
-                                                  affineApplyOp.getAffineMap(),
-                                                  affineApplyOp.getOperands())))
+        AffineValueMap avm(affineApplyOp.getAffineMap(),
+                           affineApplyOp.getOperands(),
+                           affineApplyOp.getResult());
+        if (failed(constraints.composeMap(&avm)))
           return failure();
         continue;
       }
       auto affineMinOp = cast<AffineMinOp>(op);
-      if (failed(constraints.composeMin(affineMinOp.getResult(),
-                                        affineMinOp.getAffineMap(),
-                                        affineMinOp.operands())))
+      unsigned pos;
+      bool foundMinOp = constraints.findId(affineMinOp.getResult(), &pos);
+      (void)foundMinOp;
+      assert(foundMinOp);
+      AffineMap alignedMap = constraints.computeAlignedMap(
+          affineMinOp.getAffineMap(), affineMinOp.getOperands());
+      if (failed(
+              constraints.addBound(FlatAffineConstraints::UB, pos, alignedMap)))
         return failure();
     }
   }
   for (Value v : toProjectOut)
     constraints.projectOut(v);
   return success();
 }
 
-/// Compute dynamic tensor sizes, independent of any value defined inside
-/// `outer` and such that every n-D iteration of the packingLoops has its own
-/// space (so that each packed buffer has a storage location). This is achieved
-/// by computing the extent for each of the packing loops.
-static LogicalResult computeBounds(scf::ForOp outer,
-                                   ArrayRef<Operation *> packingLoops,
-                                   SmallVector<AffineMap> &lbs,
-                                   SmallVector<AffineMap> &ubs) {
-  // Packing loop IVs are introduced as the first positions.
-  ConstraintsSet constraints = initLoopIvsAndBounds(packingLoops);
-  if (failed(
-          foldUpperBoundsIntoConstraintsSet(constraints, outer, packingLoops)))
-    return failure();
-  // Compute the bounds of the first positions, assuming the others are fixed.
-  constraints.getSliceBounds(/*pos=*/0, /*num=*/packingLoops.size(),
-                             outer->getContext(), &lbs, &ubs);
-  return success();
-}
-
 /// Ensure prerequisites that guarantee pad op hoisting can occur.
 /// Return failure in the cases when we cannot perform hoisting; i.e. if either:
 ///   1. There exists a use of `padTensorOp` that is not a linalg input operand.
 ///   2. There isn't an enclosing `outermostEnclosingForOp` loop.
 ///   3. There exists an op with a region that is dominated by
 ///   `outermostEnclosingForOp` and that isn't a LoopLikeInterface or a
 ///    LinalgOp.
 ///   4. There exists an op with side effects that is dominated by
hoistPaddingOnTensorsPrerequisites(linalg::PadTensorOp padTensorOp, int nLevels,
   }
 
   // Backward slice is a topologically sorted list of ops starting at
   // `outermostEnclosingForOp`.
   assert(outermostEnclosingForOp == backwardSlice.front());
 
   scf::ForOp outer = cast<scf::ForOp>(outermostEnclosingForOp);
 
-  ConstraintsSet constraints = initLoopIvsAndBounds(packingLoops.getArrayRef());
+  FlatAffineValueConstraints constraints =
+      initLoopIvsAndBounds(packingLoops.getArrayRef());
   if (failed(foldUpperBoundsIntoConstraintsSet(constraints, outer,
                                                packingLoops.getArrayRef())))
     return failure();
 
   unsigned numLoops = packingLoops.size();
   SmallVector<AffineMap> lbs(numLoops), ubs(numLoops);
-  if (failed(computeBounds(outer, packingLoops.getArrayRef(), lbs, ubs)))
-    return failure();
+  // Compute the bounds of the first positions, assuming the others are fixed.
+  constraints.getSliceBounds(/*pos=*/0, /*num=*/packingLoops.size(),
+                             outer->getContext(), &lbs, &ubs);
 
   SmallVector<Value> allValues;
   constraints.getAllValues(&allValues);
   SmallVector<Value> allNonLoopValues(allValues.begin() + numLoops,
                                       allValues.end());
 
   // For each packingLoop, create the extent by (ub - lb).ceilDiv(step).
   // IP just before the outermost loop considered that we hoist above.

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 )
 
 cc_library(
     name = "LinalgTransforms",
     srcs = glob([
         "lib/Dialect/Linalg/Transforms/*.cpp",
         "lib/Dialect/Linalg/Transforms/*.h",
     ]) + [
-        "lib/Dialect/Linalg/Analysis/ConstraintsSet.cpp",
         "lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp",
         "lib/Dialect/Linalg/Utils/Utils.cpp",
     ],
     hdrs = [
-        "include/mlir/Dialect/Linalg/Analysis/ConstraintsSet.h",
         "include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h",
         "include/mlir/Dialect/Linalg/Passes.h",
         "include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h",
         "include/mlir/Dialect/Linalg/Transforms/Hoisting.h",
         "include/mlir/Dialect/Linalg/Transforms/Transforms.h",
         "include/mlir/Dialect/Linalg/Utils/Utils.h",
     ],
     includes = ["include"],