[MLIR][Presburger] Cleanup getMaybeValues in FACV

[Groverkss]

Summary

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

[arjunp]

Not sure how to think of the tradeoff between having this runtime Idkind::Local check and the minor duplication. On the other hand, with more IdKinds like range, etc., it probably becomes worth it to deduplicate.

[Groverkss]

Rebase

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

public:
   inline void getAllValues(SmallVectorImpl<Value> *values) const {
     getValues(0, getNumDimAndSymbolIds(), values);
   }
 
   inline ArrayRef<Optional<Value>> getMaybeValues() const {
     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.
   inline void setValue(unsigned pos, Value val) {
     assert(pos < getNumDimAndSymbolIds() && "invalid id position");
     values[pos] = val;
   }
 

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

void FlatAffineRelation::compose(const FlatAffineRelation &other) {
   rel.convertToLocal(IdKind::SetDim, rel.getNumDomainDims(),
                      rel.getNumDomainDims() + removeDims);
   // Convert `this` from [otherDomain thisDomain] -> [thisRange] to
   // [otherDomain] -> [thisRange] by converting last thisDomain domain ids
   // to local ids.
   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)
     if (relMaybeValues[i].hasValue())
       setValue(i, relMaybeValues[i].getValue());
   // Add and match range of `this` to range of `rel`.
   for (unsigned i = 0, e = getNumRangeDims(); i < e; ++i) {
     unsigned rangeIdx = rel.getNumDomainDims() + i;

mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp

 /// `operands` are the SSA Value operands used with the affine map.
 /// Note: This function adds a new symbol column to the `constraints` for each
 /// dimension/symbol that exists in the affine map but not in `constraints`.
 static LogicalResult alignAndAddBound(FlatAffineValueConstraints &constraints,
                                       IntegerPolyhedron::BoundType type,
                                       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);
   for (unsigned i = syms.size(); i < newSyms.size(); ++i)
     constraints.appendSymbolId(newSyms[i]);
   return constraints.addBound(type, pos, alignedMap);
 }
 
for (unsigned i = resultDimStart; i < resultDimStart + numResults; ++i) {
     newConstr.addInequality(ineq);
     if (!newConstr.isEmpty())
       return failure();
   }
 
   // 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) {
     // Skip unused operands and operands that are already constants.
     if (!newOperands[i] || getConstantIntValue(newOperands[i]))
       continue;
     if (auto bound = constraints.getConstantBound(IntegerPolyhedron::EQ, i))
       newOperands[i] =