[WIP][mlir][Analysis] Split FlatAffineValueConstraints

[springerm]

Summary

Factor out all functionality that relies on AffineForOp and AffineIfOp. When dealing with these ops, there are additional assumptions about how SSA values are defined (e.g., addInductionVarOrTerminalSymbol).

Association of columns with SSA values proved useful in places such as "SCF for loop peeling". However, these places do not deal with Affine dialect ops (they use AffineMinOp, AffineMaxOp, AffineApplyOp; but only for computing mathematical computations, as a replacement for std.addi etc.), so extracting all Affine dialect-related functionality into a derived class would simplify the class's API and avoid accidentally calling a function such as addInductionVarOrTerminalSymbol, which may suddenly raise an assertion and crash the program.

Depends On D108270

mlir/include/mlir/Analysis/AffineStructures.h

 /// identifier can ultimately always be recovered as an affine function of
 /// dimensional and symbolic identifiers (involving floordiv's); note however
 /// that some floordiv combinations are converted to mod's by AffineExpr
 /// construction.
 ///
 class FlatAffineConstraints {
 public:
   /// All derived classes of FlatAffineConstraints.
-  enum class Kind { FlatAffineConstraints, FlatAffineValueConstraints };
+  enum class Kind {
+    FlatAffineConstraints,
+    FlatAffineValueConstraints,
+    FlatAffineOpValueConstraints
+  };
 
   /// Kind of identifier (column).
   enum IdKind { Dimension, Symbol, Local };
 
   /// Constructs a constraint system reserving memory for the specified number
   /// of constraints and identifiers.
   FlatAffineConstraints(unsigned numReservedInequalities,
                         unsigned numReservedEqualities,
public:
 
   FlatAffineValueConstraints(ArrayRef<const AffineValueMap *> avmRef,
                              IntegerSet set);
 
   /// Return the kind of this FlatAffineConstraints.
   Kind getKind() const override { return Kind::FlatAffineValueConstraints; }
 
   static bool classof(const FlatAffineConstraints *cst) {
-    return cst->getKind() == Kind::FlatAffineValueConstraints;
+    return cst->getKind() == Kind::FlatAffineValueConstraints ||
+           cst->getKind() == Kind::FlatAffineOpValueConstraints;
   }
 
   /// Clears any existing data and reserves memory for the specified
   /// constraints.
   void reset(unsigned numReservedInequalities, unsigned numReservedEqualities,
              unsigned numReservedCols, unsigned numDims, unsigned numSymbols,
              unsigned numLocals = 0) override;
   void reset(unsigned numReservedInequalities, unsigned numReservedEqualities,
              unsigned numReservedCols, unsigned numDims, unsigned numSymbols,
              unsigned numLocals, ArrayRef<Value> valArgs);
   void reset(unsigned numDims, unsigned numSymbols, unsigned numLocals,
              ArrayRef<Value> valArgs);
   using FlatAffineConstraints::reset;
 
   /// 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
-  /// described by the bound maps `lbMaps` and `ubMaps` of a computation slice.
-  /// Every pair (`lbMaps[i]`, `ubMaps[i]`) describes the bounds of a loop in
-  /// the nest, sorted outer-to-inner. `operands` contains the bound operands
-  /// for a single bound map. All the bound maps will use the same bound
-  /// operands. Note that some loops described by a computation slice might not
-  /// exist yet in the IR so the Value attached to those dimension identifiers
-  /// might be empty. For that reason, this method doesn't perform Value
-  /// look-ups to retrieve the dimension identifier positions. Instead, it
-  /// assumes the position of the dim identifiers in the constraint system is
-  /// the same as the position of the loop in the loop nest.
-  LogicalResult addDomainFromSliceMaps(ArrayRef<AffineMap> lbMaps,
-                                       ArrayRef<AffineMap> ubMaps,
-                                       ArrayRef<Value> operands);
-
-  /// Adds constraints imposed by the `affine.if` operation. These constraints
-  /// are collected from the IntegerSet attached to the given `affine.if`
-  /// instance argument (`ifOp`). It is asserted that:
-  /// 1) The IntegerSet of the given `affine.if` instance should not contain
-  /// semi-affine expressions,
-  /// 2) The columns of the constraint system created from `ifOp` should match
-  /// the columns in the current one regarding numbers and values.
-  void addAffineIfOpDomain(AffineIfOp ifOp);
-
   /// Adds a bound for the identifier at the specified position with constraints
   /// being drawn from the specified bound map and operands. In case of an
   /// EQ bound, the  bound map is expected to have exactly one result. In case
   /// of a LB/UB, the bound map may have more than one result, for each of which
   /// an inequality is added.
-  LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap,
-                         ValueRange operands);
+  virtual LogicalResult addBound(BoundType type, unsigned pos,
+                                 AffineMap boundMap, ValueRange operands);
 
   /// Adds a constant bound for the identifier associated with the given Value.
   void addBound(BoundType type, Value val, int64_t value);
 
   using FlatAffineConstraints::addBound;
 
   /// Returns the bound for the identifier at `pos` from the inequality at
   /// `ineqPos` as a 1-d affine value map (affine map + operands). The returned
public:
   unsigned addId(IdKind kind, unsigned pos, Value val);
   /// Add identifiers of the specified kind at the end of the table. Return the
   /// position of the column. The coefficient column corresponding to the
   /// added identifier is initialized to zero. `val` is the Value corresponding
   /// to the identifier that can optionally be provided.
   unsigned addDimId(Value val);
   unsigned addSymbolId(Value val);
 
-  /// Add the specified values as a dim or symbol id depending on its nature, if
-  /// it already doesn't exist in the system. `val` has to be either a terminal
-  /// symbol or a loop IV, i.e., it cannot be the result affine.apply of any
-  /// symbols or loop IVs. The identifier is added to the end of the existing
-  /// dims or symbols. Additional information on the identifier is extracted
-  /// from the IR and added to the constraint system.
-  void addInductionVarOrTerminalSymbol(Value val);
-
   /// Align `map` with this constraint system based on `operands`. Each operand
   /// must already have a corresponding dim/symbol in this constraint system.
   AffineMap computeAlignedMap(AffineMap map, ValueRange operands) const;
 
   /// Composes the affine value map with this FlatAffineValueConstrains, adding
   /// the results of the map as dimensions at the front
   /// [0, vMap->getNumResults()) and with the dimensions set to the equalities
   /// specified by the value map.
   ///
   /// Returns failure if the composition fails (when vMap is a semi-affine map).
   /// The vMap's operand Value's are used to look up the right positions in
   /// the FlatAffineConstraints with which to associate. Every operand of vMap
   /// should have a matching dim/symbol column in this constraint system (with
   /// the same associated Value).
   LogicalResult composeMap(const AffineValueMap *vMap);
 
   /// Projects out the identifier that is associate with Value.
   void projectOut(Value val);
   using FlatAffineConstraints::projectOut;
 
-  /// Changes all symbol identifiers which are loop IVs to dim identifiers.
-  void convertLoopIVSymbolsToDims();
-
   /// Updates the constraints to be the smallest bounding (enclosing) box that
   /// contains the points of `this` set and that of `other`, with the symbols
   /// being treated specially. For each of the dimensions, the min of the lower
   /// bounds (symbolic) and the max of the upper bounds (symbolic) is computed
   /// to determine such a bounding box. `other` is expected to have the same
   /// dimensional identifiers as this constraint system (in the same order).
   ///
   /// E.g.:
protected:
 
   /// Values corresponding to the (column) identifiers of this constraint
   /// system appearing in the order the identifiers correspond to columns.
   /// Temporary ones or those that aren't associated with any Value are set to
   /// None.
   SmallVector<Optional<Value>, 8> values;
 };
 
+/// An extension of FlatAffineConstraints in which dimensions and symbols can
+/// optionally be associated with an SSA value.
+class FlatAffineOpValueConstraints : public FlatAffineValueConstraints {
+public:
+  using FlatAffineValueConstraints::FlatAffineValueConstraints;
+
+  /// Return the kind of this FlatAffineConstraints.
+  Kind getKind() const override { return Kind::FlatAffineOpValueConstraints; }
+
+  static bool classof(const FlatAffineConstraints *cst) {
+    return cst->getKind() == Kind::FlatAffineOpValueConstraints;
+  }
+
+  /// 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
+  /// described by the bound maps `lbMaps` and `ubMaps` of a computation slice.
+  /// Every pair (`lbMaps[i]`, `ubMaps[i]`) describes the bounds of a loop in
+  /// the nest, sorted outer-to-inner. `operands` contains the bound operands
+  /// for a single bound map. All the bound maps will use the same bound
+  /// operands. Note that some loops described by a computation slice might not
+  /// exist yet in the IR so the Value attached to those dimension identifiers
+  /// might be empty. For that reason, this method doesn't perform Value
+  /// look-ups to retrieve the dimension identifier positions. Instead, it
+  /// assumes the position of the dim identifiers in the constraint system is
+  /// the same as the position of the loop in the loop nest.
+  LogicalResult addDomainFromSliceMaps(ArrayRef<AffineMap> lbMaps,
+                                       ArrayRef<AffineMap> ubMaps,
+                                       ArrayRef<Value> operands);
+
+  /// Adds constraints imposed by the `affine.if` operation. These constraints
+  /// are collected from the IntegerSet attached to the given `affine.if`
+  /// instance argument (`ifOp`). It is asserted that:
+  /// 1) The IntegerSet of the given `affine.if` instance should not contain
+  /// semi-affine expressions,
+  /// 2) The columns of the constraint system created from `ifOp` should match
+  /// the columns in the current one regarding numbers and values.
+  void addAffineIfOpDomain(AffineIfOp ifOp);
+
+  /// Adds a bound for the identifier at the specified position with constraints
+  /// being drawn from the specified bound map and operands. In case of an
+  /// EQ bound, the  bound map is expected to have exactly one result. In case
+  /// of a LB/UB, the bound map may have more than one result, for each of which
+  /// an inequality is added.
+  LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap,
+                         ValueRange operands) override;
+  using FlatAffineValueConstraints::addBound;
+
+  /// Add the specified values as a dim or symbol id depending on its nature, if
+  /// it already doesn't exist in the system. `val` has to be either a terminal
+  /// symbol or a loop IV, i.e., it cannot be the result affine.apply of any
+  /// symbols or loop IVs. The identifier is added to the end of the existing
+  /// dims or symbols. Additional information on the identifier is extracted
+  /// from the IR and added to the constraint system.
+  void addInductionVarOrTerminalSymbol(Value val);
+
+  /// Changes all symbol identifiers which are loop IVs to dim identifiers.
+  void convertLoopIVSymbolsToDims();
+};
+
 /// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the
 /// dimensions, symbols, and additional variables that represent floor divisions
 /// of dimensions, symbols, and in turn other floor divisions.  Returns failure
 /// if 'expr' could not be flattened (i.e., semi-affine is not yet handled).
 /// 'cst' contains constraints that connect newly introduced local identifiers
 /// to existing dimensional and symbolic identifiers. See documentation for
 /// AffineExprFlattener on how mod's and div's are flattened.
 LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,

mlir/lib/Analysis/AffineStructures.cpp

static void turnSymbolIntoDim(FlatAffineValueConstraints *cst, Value id) {
   if (cst->findId(id, &pos) && pos >= cst->getNumDimIds() &&
       pos < cst->getNumDimAndSymbolIds()) {
     cst->swapId(pos, cst->getNumDimIds());
     cst->setDimSymbolSeparation(cst->getNumSymbolIds() - 1);
   }
 }
 
 // Changes all symbol identifiers which are loop IVs to dim identifiers.
-void FlatAffineValueConstraints::convertLoopIVSymbolsToDims() {
+void FlatAffineOpValueConstraints::convertLoopIVSymbolsToDims() {
   // Gather all symbols which are loop IVs.
   SmallVector<Value, 4> loopIVs;
   for (unsigned i = getNumDimIds(), e = getNumDimAndSymbolIds(); i < e; i++) {
     if (hasValue(i) && getForInductionVarOwner(getValue(i)))
       loopIVs.push_back(getValue(i));
   }
   // Turn each symbol in 'loopIVs' into a dim identifier.
   for (auto iv : loopIVs) {
     turnSymbolIntoDim(this, iv);
   }
 }
 
-void FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
+void FlatAffineOpValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
   if (containsId(val))
     return;
 
   // Caller is expected to fully compose map/operands if necessary.
   assert((isTopLevelValue(val) || isForInductionVar(val)) &&
          "non-terminal symbol / loop IV expected");
   // Outer loop IVs could be used in forOp's bounds.
   if (auto loop = getForInductionVarOwner(val)) {
     addDimId(getNumDimIds(), val);
     if (failed(this->addAffineForOpDomain(loop)))
       LLVM_DEBUG(
           loop.emitWarning("failed to add domain info to constraint system"));
     return;
   }
   // Add top level symbol.
   addSymbolId(getNumSymbolIds(), val);
   // Check if the symbol is a constant.
   if (auto constOp = val.getDefiningOp<ConstantIndexOp>())
     addBound(BoundType::EQ, val, constOp.getValue());
 }
 
 LogicalResult
-FlatAffineValueConstraints::addAffineForOpDomain(AffineForOp forOp) {
+FlatAffineOpValueConstraints::addAffineForOpDomain(AffineForOp forOp) {
   unsigned pos;
   // Pre-condition for this method.
   if (!findId(forOp.getInductionVar(), &pos)) {
     assert(false && "Value not found");
     return failure();
   }
 
   int64_t step = forOp.getStep();
if (forOp.hasConstantUpperBound()) {
     return success();
   }
   // Non-constant upper bound case.
   return addBound(BoundType::UB, pos, forOp.getUpperBoundMap(),
                   forOp.getUpperBoundOperands());
 }
 
 LogicalResult
-FlatAffineValueConstraints::addDomainFromSliceMaps(ArrayRef<AffineMap> lbMaps,
+FlatAffineOpValueConstraints::addDomainFromSliceMaps(ArrayRef<AffineMap> lbMaps,
                                                      ArrayRef<AffineMap> ubMaps,
                                                      ArrayRef<Value> operands) {
   assert(lbMaps.size() == ubMaps.size());
   assert(lbMaps.size() <= getNumDimIds());
 
   for (unsigned i = 0, e = lbMaps.size(); i < e; ++i) {
     AffineMap lbMap = lbMaps[i];
     AffineMap ubMap = ubMaps[i];
     assert(!lbMap || lbMap.getNumInputs() == operands.size());
     assert(!ubMap || ubMap.getNumInputs() == operands.size());
for (unsigned i = 0, e = lbMaps.size(); i < e; ++i) {
     if (lbMap && failed(addBound(BoundType::LB, i, lbMap, operands)))
       return failure();
     if (ubMap && failed(addBound(BoundType::UB, i, ubMap, operands)))
       return failure();
   }
   return success();
 }
 
-void FlatAffineValueConstraints::addAffineIfOpDomain(AffineIfOp ifOp) {
+void FlatAffineOpValueConstraints::addAffineIfOpDomain(AffineIfOp ifOp) {
   // Create the base constraints from the integer set attached to ifOp.
   FlatAffineValueConstraints cst(ifOp.getIntegerSet());
 
   // Bind ids in the constraints to ifOp operands.
   SmallVector<Value, 4> operands = ifOp.getOperands();
   cst.setValues(0, cst.getNumDimAndSymbolIds(), operands);
 
   // Merge the constraints from ifOp to the current domain. We need first merge
LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos,
                                                    ValueRange boundOperands) {
   // 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) {
+    addSymbolId(getNumSymbolIds(), operand);
+    // Check if the symbol is a constant.
+    if (auto constOp = operand.getDefiningOp<ConstantIndexOp>())
+      addBound(BoundType::EQ, operand, constOp.getValue());
+  }
+  return addBound(type, pos, computeAlignedMap(map, operands));
+}
+
+LogicalResult FlatAffineOpValueConstraints::addBound(BoundType type,
+                                                     unsigned pos,
+                                                     AffineMap boundMap,
+                                                     ValueRange boundOperands) {
+  // 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);
   return addBound(type, pos, computeAlignedMap(map, operands));
 }
 
 // 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
for (unsigned i = 0, e = lbMaps.size(); i < e; ++i) {
     // part of the slice.
     if (lbMap && lbMap.getNumResults() != 0 && ubMap &&
         ubMap.getNumResults() != 0) {
       if (failed(addBound(BoundType::LB, pos, lbMap, operands)))
         return failure();
       if (failed(addBound(BoundType::UB, pos, ubMap, operands)))
         return failure();
     } else {
+      if (auto *constr = dyn_cast<FlatAffineOpValueConstraints>(this)) {
         auto loop = getForInductionVarOwner(values[i]);
-      if (failed(this->addAffineForOpDomain(loop)))
+        if (failed(constr->addAffineForOpDomain(loop)))
           return failure();
+      } else {
+        assert(false && "unexpected empty bound");
+      }
     }
   }
   return success();
 }
 
 void FlatAffineConstraints::addEquality(ArrayRef<int64_t> eq) {
   assert(eq.size() == getNumCols());
   unsigned row = equalities.appendExtraRow();