[MLIR][Presburger] Cleanup getMaybeValues in FACV

This patch cleans up multiple getMaybeValue functions to take an IdKind instead
of special functions.

Reviewed By: arjunp

Differential Revision: https://reviews.llvm.org/D125617

3 files changed, 10 insertions, 15 deletions

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<Optional<Value>> getMaybeDimValues() const {
-    return {values.data(), getNumDimIds()};
-  }
-
-  inline ArrayRef<Optional<Value>> getMaybeSymbolValues() const {
-    return {values.data() + getNumDimIds(), getNumSymbolIds()};
-  }
-
-  inline ArrayRef<Optional<Value>> getMaybeDimAndSymbolValues() const {
-    return {values.data(), getNumDimIds() + getNumSymbolIds()};
+  inline ArrayRef<Optional<Value>>
+  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<Value> 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<Value> 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) {