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 @@ -243,19 +243,23 @@ /// Swap the posA^th identifier with the posB^th identifier. virtual void swapId(unsigned posA, unsigned posB); - /// Add identifiers of the specified kind - specified positions are relative - /// to the kind of identifier. The coefficient column corresponding to the - /// added identifier is initialized to zero. - void addDimId(unsigned pos); - void addSymbolId(unsigned pos); - void addLocalId(unsigned pos); - virtual unsigned addId(IdKind kind, unsigned pos); - /// 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. - unsigned addDimId(); - unsigned addSymbolId(); - unsigned addLocalId(); + /// Insert `num` identifiers of the specified kind at position `pos`. + /// Positions are relative to the kind of identifier. The coefficient columns + /// corresponding to the added identifiers are initialized to zero. Return the + /// absolute column position (i.e., not relative to the kind of identifier) + /// of the first added identifier. + unsigned insertDimId(unsigned pos, unsigned num = 1); + unsigned insertSymbolId(unsigned pos, unsigned num = 1); + unsigned insertLocalId(unsigned pos, unsigned num = 1); + virtual unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1); + + /// Append `num` identifiers of the specified kind after the last identifier. + /// of that kind. Return the position of the first appended column. The + /// coefficient columns corresponding to the added identifiers are initialized + /// to zero. + unsigned appendDimId(unsigned num = 1); + unsigned appendSymbolId(unsigned num = 1); + unsigned appendLocalId(unsigned num = 1); /// Composes an affine map whose dimensions and symbols match one to one with /// the dimensions and symbols of this FlatAffineConstraints. The results of @@ -655,22 +659,31 @@ /// Swap the posA^th identifier with the posB^th identifier. void swapId(unsigned posA, unsigned posB) override; - /// Add identifiers of the specified kind - specified positions are relative - /// to the kind of identifier. 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. - void addDimId(unsigned pos, Value val); - using FlatAffineConstraints::addDimId; - void addSymbolId(unsigned pos, Value val); - using FlatAffineConstraints::addSymbolId; - unsigned addId(IdKind kind, unsigned pos) override; - 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); + /// Insert identifiers of the specified kind at position `pos`. Positions are + /// relative to the kind of identifier. The coefficient columns corresponding + /// to the added identifiers are initialized to zero. `vals` are the Values + /// corresponding to the identifiers. Return the absolute column position + /// (i.e., not relative to the kind of identifier) of the first added + /// identifier. + /// + /// Note: Empty Values are allowed in `vals`. + unsigned insertDimId(unsigned pos, ValueRange vals); + using FlatAffineConstraints::insertDimId; + unsigned insertSymbolId(unsigned pos, ValueRange vals); + using FlatAffineConstraints::insertSymbolId; + unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override; + unsigned insertId(IdKind kind, unsigned pos, ValueRange vals); + + /// Append identifiers of the specified kind after the last identifier of that + /// kind. The coefficient columns corresponding to the added identifiers are + /// initialized to zero. `vals` are the Values corresponding to the + /// identifiers. Return the position of the first added column. + /// + /// Note: Empty Values are allowed in `vals`. + unsigned appendDimId(ValueRange vals); + using FlatAffineConstraints::appendDimId; + unsigned appendSymbolId(ValueRange vals); + using FlatAffineConstraints::appendSymbolId; /// 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 diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp --- a/mlir/lib/Analysis/AffineAnalysis.cpp +++ b/mlir/lib/Analysis/AffineAnalysis.cpp @@ -614,9 +614,7 @@ unsigned srcNumLocalIds = srcLocalVarCst.getNumLocalIds(); unsigned dstNumLocalIds = destLocalVarCst.getNumLocalIds(); unsigned numLocalIdsToAdd = srcNumLocalIds + dstNumLocalIds; - for (unsigned i = 0; i < numLocalIdsToAdd; i++) { - dependenceDomain->addLocalId(dependenceDomain->getNumLocalIds()); - } + dependenceDomain->appendLocalId(numLocalIdsToAdd); unsigned numDims = dependenceDomain->getNumDimIds(); unsigned numSymbols = dependenceDomain->getNumSymbolIds(); @@ -859,9 +857,7 @@ unsigned numIdsToEliminate = dependenceDomain->getNumIds(); // Add new variables to 'dependenceDomain' to represent the direction // constraints for each shared loop. - for (unsigned j = 0; j < numCommonLoops; ++j) { - dependenceDomain->addDimId(j); - } + dependenceDomain->insertDimId(/*pos=*/0, /*num=*/numCommonLoops); // Add equality constraints for each common loop, setting newly introduced // variable at column 'j' to the 'dst' IV minus the 'src IV. 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 @@ -168,9 +168,7 @@ return; } assert(flatExprs.size() == set.getNumConstraints()); - for (unsigned l = 0, e = localVarCst.getNumLocalIds(); l < e; l++) { - addLocalId(getNumLocalIds()); - } + appendLocalId(/*num=*/localVarCst.getNumLocalIds()); for (unsigned i = 0, e = flatExprs.size(); i < e; ++i) { const auto &flatExpr = flatExprs[i]; @@ -259,57 +257,60 @@ } } -void FlatAffineConstraints::addLocalId(unsigned pos) { - addId(IdKind::Local, pos); +unsigned FlatAffineConstraints::appendDimId(unsigned num) { + unsigned pos = getNumDimIds(); + insertId(IdKind::Dimension, pos, num); + return pos; } -unsigned FlatAffineConstraints::addLocalId() { - unsigned pos = getNumLocalIds(); - addId(IdKind::Local, pos); +unsigned FlatAffineValueConstraints::appendDimId(ValueRange vals) { + unsigned pos = getNumDimIds(); + insertId(IdKind::Dimension, pos, vals); return pos; } -void FlatAffineConstraints::addDimId(unsigned pos) { - addId(IdKind::Dimension, pos); +unsigned FlatAffineConstraints::appendSymbolId(unsigned num) { + unsigned pos = getNumSymbolIds(); + insertId(IdKind::Symbol, pos, num); + return pos; } -unsigned FlatAffineConstraints::addDimId() { - unsigned pos = getNumDimIds(); - addId(IdKind::Dimension, pos); +unsigned FlatAffineValueConstraints::appendSymbolId(ValueRange vals) { + unsigned pos = getNumSymbolIds(); + insertId(IdKind::Symbol, pos, vals); return pos; } -void FlatAffineValueConstraints::addDimId(unsigned pos, Value val) { - addId(IdKind::Dimension, pos, val); +unsigned FlatAffineConstraints::appendLocalId(unsigned num) { + unsigned pos = getNumLocalIds(); + insertId(IdKind::Local, pos, num); + return pos; } -unsigned FlatAffineValueConstraints::addDimId(Value val) { - unsigned pos = getNumDimIds(); - addId(IdKind::Dimension, pos, val); - return pos; +unsigned FlatAffineConstraints::insertDimId(unsigned pos, unsigned num) { + return insertId(IdKind::Dimension, pos, num); } -void FlatAffineConstraints::addSymbolId(unsigned pos) { - addId(IdKind::Symbol, pos); +unsigned FlatAffineValueConstraints::insertDimId(unsigned pos, + ValueRange vals) { + return insertId(IdKind::Dimension, pos, vals); } -unsigned FlatAffineConstraints::addSymbolId() { - unsigned pos = getNumSymbolIds(); - addId(IdKind::Symbol, pos); - return pos; +unsigned FlatAffineConstraints::insertSymbolId(unsigned pos, unsigned num) { + return insertId(IdKind::Symbol, pos, num); } -void FlatAffineValueConstraints::addSymbolId(unsigned pos, Value val) { - addId(IdKind::Symbol, pos, val); +unsigned FlatAffineValueConstraints::insertSymbolId(unsigned pos, + ValueRange vals) { + return insertId(IdKind::Symbol, pos, vals); } -unsigned FlatAffineValueConstraints::addSymbolId(Value val) { - unsigned pos = getNumSymbolIds(); - addId(IdKind::Symbol, pos, val); - return pos; +unsigned FlatAffineConstraints::insertLocalId(unsigned pos, unsigned num) { + return insertId(IdKind::Local, pos, num); } -unsigned FlatAffineConstraints::addId(IdKind kind, unsigned pos) { +unsigned FlatAffineConstraints::insertId(IdKind kind, unsigned pos, + unsigned num) { if (kind == IdKind::Dimension) assert(pos <= getNumDimIds()); else if (kind == IdKind::Symbol) @@ -320,36 +321,41 @@ unsigned absolutePos; if (kind == IdKind::Dimension) { absolutePos = pos; - numDims++; + numDims += num; } else if (kind == IdKind::Symbol) { absolutePos = pos + getNumDimIds(); - numSymbols++; + numSymbols += num; } else { absolutePos = pos + getNumDimIds() + getNumSymbolIds(); } - numIds++; + numIds += num; - inequalities.insertColumn(absolutePos); - equalities.insertColumn(absolutePos); + inequalities.insertColumns(absolutePos, num); + equalities.insertColumns(absolutePos, num); return absolutePos; } -unsigned FlatAffineValueConstraints::addId(IdKind kind, unsigned pos) { - return addId(kind, pos, /*val=*/{}); +unsigned FlatAffineValueConstraints::insertId(IdKind kind, unsigned pos, + unsigned num) { + unsigned absolutePos = FlatAffineConstraints::insertId(kind, pos, num); + values.insert(values.begin() + absolutePos, num, None); + assert(values.size() == getNumIds()); + return absolutePos; } -unsigned FlatAffineValueConstraints::addId(IdKind kind, unsigned pos, - Value val) { - unsigned absolutePos = FlatAffineConstraints::addId(kind, pos); +unsigned FlatAffineValueConstraints::insertId(IdKind kind, unsigned pos, + ValueRange vals) { + assert(!vals.empty() && "expected ValueRange with Values"); + unsigned num = vals.size(); + unsigned absolutePos = FlatAffineConstraints::insertId(kind, pos, num); - // If an 'id' is provided, insert it; otherwise use None. - if (val) - values.insert(values.begin() + absolutePos, val); - else - values.insert(values.begin() + absolutePos, None); - assert(values.size() == getNumIds()); + // If a Value is provided, insert it; otherwise use None. + for (unsigned i = 0; i < num; ++i) + values.insert(values.begin() + absolutePos + i, + vals[i] ? Optional(vals[i]) : None); + assert(values.size() == getNumIds()); return absolutePos; } @@ -412,13 +418,8 @@ [](Optional id) { return id.hasValue(); })); // Place local id's of A after local id's of B. - for (unsigned l = 0, e = a->getNumLocalIds(); l < e; l++) { - b->addLocalId(0); - } - for (unsigned t = 0, e = b->getNumLocalIds() - a->getNumLocalIds(); t < e; - t++) { - a->addLocalId(a->getNumLocalIds()); - } + b->insertLocalId(/*pos=*/0, /*num=*/a->getNumLocalIds()); + a->appendLocalId(/*num=*/b->getNumLocalIds() - a->getNumLocalIds()); SmallVector aDimValues, aSymValues; a->getValues(offset, a->getNumDimIds(), &aDimValues); @@ -434,16 +435,14 @@ "A's dim appears in B's non-dim position"); b->swapId(d, loc); } else { - b->addDimId(d); - b->setValue(d, aDimValue); + b->insertDimId(d, aDimValue); } d++; } // Dimensions that are in B, but not in A, are added at the end. for (unsigned t = a->getNumDimIds(), e = b->getNumDimIds(); t < e; t++) { - a->addDimId(a->getNumDimIds()); - a->setValue(a->getNumDimIds() - 1, b->getValue(t)); + a->appendDimId(b->getValue(t)); } } { @@ -456,7 +455,8 @@ "A's symbol appears in B's non-symbol position"); b->swapId(s, loc); } else { - b->addSymbolId(s - b->getNumDimIds()); + b->insertSymbolId(s - b->getNumDimIds()); + // TODO: Is this correct? b->setValue(s, aSymValue); } s++; @@ -465,8 +465,7 @@ for (unsigned t = a->getNumDimAndSymbolIds(), e = b->getNumDimAndSymbolIds(); t < e; t++) { - a->addSymbolId(a->getNumSymbolIds()); - a->setValue(a->getNumDimAndSymbolIds() - 1, b->getValue(t)); + a->appendSymbolId(b->getValue(t)); } } assert(areIdsAligned(*a, *b) && "IDs expected to be aligned"); @@ -497,9 +496,7 @@ assert(flatExprs.size() == other.getNumResults()); // Add dimensions corresponding to the map's results. - for (unsigned t = 0, e = other.getNumResults(); t < e; t++) { - addDimId(0); - } + insertDimId(/*pos=*/0, /*num=*/other.getNumResults()); // We add one equality for each result connecting the result dim of the map to // the other identifiers. @@ -570,14 +567,14 @@ "non-terminal symbol / loop IV expected"); // Outer loop IVs could be used in forOp's bounds. if (auto loop = getForInductionVarOwner(val)) { - addDimId(getNumDimIds(), val); + appendDimId(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); + appendSymbolId(val); // Check if the symbol is a constant. if (auto constOp = val.getDefiningOp()) addBound(BoundType::EQ, val, constOp.getValue()); @@ -1944,11 +1941,9 @@ if (localCst.getNumLocalIds() > 0) { unsigned numLocalIds = getNumLocalIds(); // Insert local dims of localCst at the beginning. - for (unsigned l = 0, e = localCst.getNumLocalIds(); l < e; ++l) - addLocalId(0); + insertLocalId(/*pos=*/0, /*num=*/localCst.getNumLocalIds()); // Insert local dims of `this` at the end of localCst. - for (unsigned l = 0; l < numLocalIds; ++l) - localCst.addLocalId(localCst.getNumLocalIds()); + localCst.appendLocalId(/*num=*/numLocalIds); // Dimensions of localCst and this constraint set match. Append localCst to // this constraint set. append(localCst); @@ -2150,7 +2145,7 @@ assert(dividend.size() == getNumCols() && "incorrect dividend size"); assert(divisor > 0 && "positive divisor expected"); - addLocalId(getNumLocalIds()); + appendLocalId(); // Add two constraints for this new identifier 'q'. SmallVector bound(dividend.size() + 1); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp @@ -559,14 +559,11 @@ initLoopIvsAndBounds(ArrayRef loops) { FlatAffineValueConstraints constraints; for (Operation *op : loops) - constraints.addDimId(constraints.getNumDimIds(), - cast(op).getInductionVar()); + constraints.appendDimId(cast(op).getInductionVar()); for (Operation *op : loops) - constraints.addDimId(constraints.getNumDimIds(), - cast(op).lowerBound()); + constraints.appendDimId(cast(op).lowerBound()); for (Operation *op : loops) - constraints.addDimId(constraints.getNumDimIds(), - cast(op).upperBound()); + constraints.appendDimId(cast(op).upperBound()); unsigned numLoops = loops.size(); for (unsigned ivIdx = 0, e = numLoops; ivIdx < e; ++ivIdx) { // iv - lb >= 0 @@ -628,7 +625,7 @@ constraints.findId(v, &pos); return pos >= constraints.getNumDimIds(); } - constraints.addDimId(constraints.getNumDimIds(), v); + constraints.appendDimId(v); return false; }; diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp --- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp +++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp @@ -186,7 +186,7 @@ AffineMap alignedMap = alignAffineMapWithValues(map, operands, dims, syms, &newSyms); for (unsigned i = syms.size(); i < newSyms.size(); ++i) - constraints.addSymbolId(constraints.getNumSymbolIds(), newSyms[i]); + constraints.appendSymbolId(newSyms[i]); return constraints.addBound(type, pos, alignedMap); } @@ -236,11 +236,9 @@ unsigned numResults = map.getNumResults(); // Add a few extra dimensions. - unsigned dimOp = constraints.addDimId(); // `op` - unsigned dimOpBound = constraints.addDimId(); // `op` lower/upper bound - unsigned resultDimStart = constraints.getNumDimIds(); - for (unsigned i = 0; i < numResults; ++i) - constraints.addDimId(); + unsigned dimOp = constraints.appendDimId(); // `op` + unsigned dimOpBound = constraints.appendDimId(); // `op` lower/upper bound + unsigned resultDimStart = constraints.appendDimId(/*num=*/numResults); // Add an inequality for each result expr_i of map: // isMin: op <= expr_i, !isMin: op >= expr_i @@ -346,9 +344,7 @@ Value iv, Value ub, Value step, bool insideLoop) { FlatAffineValueConstraints constraints; - constraints.addDimId(0, iv); - constraints.addDimId(1, ub); - constraints.addDimId(2, step); + constraints.appendDimId({iv, ub, step}); if (auto constUb = getConstantIntValue(ub)) constraints.addBound(FlatAffineConstraints::EQ, 1, *constUb); if (auto constStep = getConstantIntValue(step)) @@ -443,9 +439,9 @@ if (!stepInt) continue; - unsigned dimIv = constraints.addDimId(iv); - unsigned dimLb = constraints.addDimId(lb); - unsigned dimUb = constraints.addDimId(ub); + unsigned dimIv = constraints.appendDimId(iv); + unsigned dimLb = constraints.appendDimId(lb); + unsigned dimUb = constraints.appendDimId(ub); // If loop lower/upper bounds are constant: Add EQ constraint. Optional lbInt = getConstantIntValue(lb);