diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h --- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h @@ -430,16 +430,11 @@ 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 --- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp @@ -1666,8 +1666,8 @@ 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 --- a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp +++ b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp @@ -43,8 +43,8 @@ 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 @@ // 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) {