diff --git a/mlir/include/mlir/Analysis/AffineStructures.h b/mlir/include/mlir/Analysis/AffineStructures.h --- a/mlir/include/mlir/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Analysis/AffineStructures.h @@ -59,7 +59,11 @@ 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 }; @@ -556,7 +560,8 @@ 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 @@ -574,49 +579,13 @@ /// Clones this object. std::unique_ptr 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 lbMaps, - ArrayRef ubMaps, - ArrayRef 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); @@ -672,14 +641,6 @@ 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; @@ -700,9 +661,6 @@ 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 @@ -833,6 +791,76 @@ SmallVector, 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 lbMaps, + ArrayRef ubMaps, + ArrayRef 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 diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp --- a/mlir/lib/Analysis/AffineStructures.cpp +++ b/mlir/lib/Analysis/AffineStructures.cpp @@ -548,7 +548,7 @@ } // 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 loopIVs; for (unsigned i = getNumDimIds(), e = getNumDimAndSymbolIds(); i < e; i++) { @@ -561,7 +561,7 @@ } } -void FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) { +void FlatAffineOpValueConstraints::addInductionVarOrTerminalSymbol(Value val) { if (containsId(val)) return; @@ -584,7 +584,7 @@ } LogicalResult -FlatAffineValueConstraints::addAffineForOpDomain(AffineForOp forOp) { +FlatAffineOpValueConstraints::addAffineForOpDomain(AffineForOp forOp) { unsigned pos; // Pre-condition for this method. if (!findId(forOp.getInductionVar(), &pos)) { @@ -636,9 +636,9 @@ } LogicalResult -FlatAffineValueConstraints::addDomainFromSliceMaps(ArrayRef lbMaps, - ArrayRef ubMaps, - ArrayRef operands) { +FlatAffineOpValueConstraints::addDomainFromSliceMaps(ArrayRef lbMaps, + ArrayRef ubMaps, + ArrayRef operands) { assert(lbMaps.size() == ubMaps.size()); assert(lbMaps.size() <= getNumDimIds()); @@ -685,7 +685,7 @@ 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()); @@ -2043,6 +2043,26 @@ 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()) + 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 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)); @@ -2092,9 +2112,13 @@ if (failed(addBound(BoundType::UB, pos, ubMap, operands))) return failure(); } else { - auto loop = getForInductionVarOwner(values[i]); - if (failed(this->addAffineForOpDomain(loop))) - return failure(); + if (auto *constr = dyn_cast(this)) { + auto loop = getForInductionVarOwner(values[i]); + if (failed(constr->addAffineForOpDomain(loop))) + return failure(); + } else { + assert(false && "unexpected empty bound"); + } } } return success();