diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h index 0ddacad04792..b307f0d4d2e3 100644 --- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h @@ -431,16 +431,11 @@ public: return {values.data(), values.size()}; } - inline ArrayRef> getMaybeDimValues() const { - return {values.data(), getNumDimIds()}; - } - - inline ArrayRef> getMaybeSymbolValues() const { - return {values.data() + getNumDimIds(), getNumSymbolIds()}; - } - - inline ArrayRef> getMaybeDimAndSymbolValues() const { - return {values.data(), getNumDimIds() + getNumSymbolIds()}; + inline ArrayRef> + getMaybeValues(presburger::IdKind kind) const { + assert(kind != IdKind::Local && + "Local identifiers do not have any value attached to them."); + return {values.data() + getIdKindOffset(kind), getNumIdKind(kind)}; } /// Sets the Value associated with the pos^th identifier. diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp index b3bf543f93dc..2c8e3acacb4d 100644 --- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp @@ -1666,8 +1666,8 @@ void FlatAffineRelation::compose(const FlatAffineRelation &other) { convertToLocal(IdKind::SetDim, getNumDomainDims() - removeDims, getNumDomainDims()); - auto thisMaybeValues = getMaybeDimValues(); - auto relMaybeValues = rel.getMaybeDimValues(); + auto thisMaybeValues = getMaybeValues(IdKind::SetDim); + auto relMaybeValues = rel.getMaybeValues(IdKind::SetDim); // Add and match domain of `rel` to domain of `this`. for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i) diff --git a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp index 78ad8532e0f3..5abc780e22d8 100644 --- a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp +++ b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp @@ -43,8 +43,8 @@ static LogicalResult alignAndAddBound(FlatAffineValueConstraints &constraints, unsigned pos, AffineMap map, ValueRange operands) { SmallVector dims, syms, newSyms; - unpackOptionalValues(constraints.getMaybeDimValues(), dims); - unpackOptionalValues(constraints.getMaybeSymbolValues(), syms); + unpackOptionalValues(constraints.getMaybeValues(IdKind::SetDim), dims); + unpackOptionalValues(constraints.getMaybeValues(IdKind::Symbol), syms); AffineMap alignedMap = alignAffineMapWithValues(map, operands, dims, syms, &newSyms); @@ -182,7 +182,7 @@ canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap map, // Lower and upper bound of `op` are equal. Replace `minOp` with its bound. AffineMap newMap = alignedBoundMap; SmallVector newOperands; - unpackOptionalValues(constraints.getMaybeDimAndSymbolValues(), newOperands); + unpackOptionalValues(constraints.getMaybeValues(), newOperands); // If dims/symbols have known constant values, use those in order to simplify // the affine map further. for (int64_t i = 0, e = constraints.getNumIds(); i < e; ++i) {