[Mlir-commits] [mlir] [MLIR][Presburger] Use Identifiers outside Presburger library (PR #77316)

Bharathi Ramana Joshi llvmlistbot at llvm.org
Fri Jan 12 10:03:32 PST 2024


https://github.com/iambrj updated https://github.com/llvm/llvm-project/pull/77316

>From 2ec01d509e624dac91ac25b2e2077569d4a6400e Mon Sep 17 00:00:00 2001
From: iambrj <joshibharathiramana at gmail.com>
Date: Fri, 12 Jan 2024 10:03:30 +0530
Subject: [PATCH 1/2] [MLIR][Presburger] Implement IntegerRelation::setId

---
 .../Analysis/Presburger/IntegerRelation.h     |  4 +++
 .../Analysis/Presburger/IntegerRelation.cpp   |  8 +++++
 .../Presburger/IntegerRelationTest.cpp        | 36 +++++++++++++++++++
 3 files changed, 48 insertions(+)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index 8e2c9fca0a17cb..00dd230e853991 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -122,6 +122,10 @@ class IntegerRelation {
   /// current space; this will result in an assert failure.
   void setSpaceExceptLocals(const PresburgerSpace &oSpace);
 
+  /// Set the ith identifier of the IntegerRelation's PresburgerSpace. The index
+  /// is relative to the kind of the variable.
+  void setId(VarKind kind, unsigned i, Identifier id);
+
   /// Returns a copy of the space without locals.
   PresburgerSpace getSpaceWithoutLocals() const {
     return PresburgerSpace::getRelationSpace(space.getNumDomainVars(),
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index f78e21ccd38eb9..7d2a63d17676f5 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -66,6 +66,14 @@ void IntegerRelation::setSpaceExceptLocals(const PresburgerSpace &oSpace) {
   space.insertVar(VarKind::Local, 0, newNumLocals);
 }
 
+void IntegerRelation::setId(VarKind kind, unsigned i, Identifier id) {
+  assert(space.isUsingIds() &&
+         "space must be using identifiers to set an identifier");
+  assert(kind != VarKind::Local && "local variables cannot have identifiers");
+  assert(i < space.getNumVarKind(kind) && "invalid variable index");
+  space.getId(kind, i) = id;
+}
+
 void IntegerRelation::append(const IntegerRelation &other) {
   assert(space.isEqual(other.getSpace()) && "Spaces must be equal.");
 
diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index dff092b3204bb3..0d568a525f60cf 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -451,3 +451,39 @@ TEST(IntegerRelationTest, mergeAndAlignCommonSuffixSymbols) {
   EXPECT_EQ(otherSpace.getId(VarKind::Range, 1),
             Identifier(&otherIdentifiers[3]));
 }
+
+TEST(IntegerRelationTest, setId) {
+  IntegerRelation rel = parseRelationFromSet(
+      "(x, y, z)[A, B, C, D] : (x + A - C - y + D - z >= 0)", 2);
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 4, 0);
+  space.resetIds();
+
+  // Attach identifiers.
+  int identifiers[7] = {'x', 'y', 'z', 'A', 'B', 'C', 'D'};
+  space.getId(VarKind::Domain, 0) = Identifier(&identifiers[0]);
+  space.getId(VarKind::Domain, 1) = Identifier(&identifiers[1]);
+  space.getId(VarKind::Range, 0) = Identifier(&identifiers[2]);
+  space.getId(VarKind::Symbol, 0) = Identifier(&identifiers[3]);
+  space.getId(VarKind::Symbol, 1) = Identifier(&identifiers[4]);
+  space.getId(VarKind::Symbol, 2) = Identifier(&identifiers[5]);
+  space.getId(VarKind::Symbol, 3) = Identifier(&identifiers[6]);
+  rel.setSpace(space);
+
+  rel.getSpace().dump();
+  int newIdentifiers[3] = {1, 2, 3};
+  rel.setId(VarKind::Domain, 1, Identifier(&newIdentifiers[0]));
+  rel.setId(VarKind::Range, 0, Identifier(&newIdentifiers[1]));
+  rel.setId(VarKind::Symbol, 2, Identifier(&newIdentifiers[2]));
+  rel.getSpace().dump();
+
+  space = rel.getSpace();
+  // Check that new identifiers are set correctly.
+  EXPECT_EQ(space.getId(VarKind::Domain, 1), Identifier(&newIdentifiers[0]));
+  EXPECT_EQ(space.getId(VarKind::Range, 0), Identifier(&newIdentifiers[1]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 2), Identifier(&newIdentifiers[2]));
+  // Check that old identifier are not changed.
+  EXPECT_EQ(space.getId(VarKind::Domain, 0), Identifier(&identifiers[0]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 0), Identifier(&identifiers[3]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 1), Identifier(&identifiers[4]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 3), Identifier(&identifiers[6]));
+}

>From 094fca35de2b48e5ee468ba455ce14e8354e6816 Mon Sep 17 00:00:00 2001
From: iambrj <joshibharathiramana at gmail.com>
Date: Fri, 12 Jan 2024 23:28:59 +0530
Subject: [PATCH 2/2] [MLIR][Presburger] Use Identifiers outside Presburger
 library

---
 .../Analysis/FlatLinearValueConstraints.h     |  92 +++++++++------
 .../Dialect/Affine/Analysis/AffineAnalysis.h  |   3 +-
 .../Affine/Analysis/AffineStructures.h        |   4 +-
 .../Analysis/FlatLinearValueConstraints.cpp   | 109 +++++-------------
 .../Affine/Analysis/AffineAnalysis.cpp        |  42 ++++---
 .../Affine/Analysis/AffineStructures.cpp      |  31 +++--
 mlir/lib/Dialect/Affine/Utils/Utils.cpp       |   2 +
 7 files changed, 136 insertions(+), 147 deletions(-)

diff --git a/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h b/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
index e4de5b0661571c..3cfc435310ffd2 100644
--- a/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
+++ b/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
@@ -205,6 +205,8 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
 /// where each non-local variable can have an SSA Value attached to it.
 class FlatLinearValueConstraints : public FlatLinearConstraints {
 public:
+  using Identifier = presburger::Identifier;
+
   /// Constructs a constraint system reserving memory for the specified number
   /// of constraints and variables. `valArgs` are the optional SSA values
   /// associated with each dimension/symbol. These must either be empty or match
@@ -217,11 +219,12 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
       : FlatLinearConstraints(numReservedInequalities, numReservedEqualities,
                               numReservedCols, numDims, numSymbols, numLocals) {
     assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolVars());
-    values.reserve(numReservedCols);
-    if (valArgs.empty())
-      values.resize(getNumDimAndSymbolVars(), std::nullopt);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+    // Use values in space for FlatLinearValueConstraints.
+    space.resetIds();
+    // Set the values for the non-local variables.
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      if (valArgs[i])
+        setValue(i, *valArgs[i]);
   }
 
   /// Constructs a constraint system reserving memory for the specified number
@@ -236,11 +239,12 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
       : FlatLinearConstraints(numReservedInequalities, numReservedEqualities,
                               numReservedCols, numDims, numSymbols, numLocals) {
     assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolVars());
-    values.reserve(numReservedCols);
-    if (valArgs.empty())
-      values.resize(getNumDimAndSymbolVars(), std::nullopt);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+    // Use values in space for FlatLinearValueConstraints.
+    space.resetIds();
+    // Set the values for the non-local variables.
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      if (valArgs[i])
+        setValue(i, valArgs[i]);
   }
 
   /// Constructs a constraint system with the specified number of dimensions
@@ -273,10 +277,12 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
                              ArrayRef<std::optional<Value>> valArgs = {})
       : FlatLinearConstraints(fac) {
     assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolVars());
-    if (valArgs.empty())
-      values.resize(getNumDimAndSymbolVars(), std::nullopt);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+    // Use values in space for FlatLinearValueConstraints.
+    space.resetIds();
+    // Set the values for the non-local variables.
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      if (valArgs[i])
+        setValue(i, *valArgs[i]);
   }
 
   /// Creates an affine constraint system from an IntegerSet.
@@ -302,7 +308,9 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   inline Value getValue(unsigned pos) const {
     assert(pos < getNumDimAndSymbolVars() && "Invalid position");
     assert(hasValue(pos) && "variable's Value not set");
-    return *values[pos];
+    VarKind kind = getVarKindAt(pos);
+    unsigned relativePos = pos - getVarKindOffset(kind);
+    return space.getId(kind, relativePos).getValue<Value>();
   }
 
   /// Returns the Values associated with variables in range [start, end).
@@ -317,21 +325,42 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
       values->push_back(getValue(i));
   }
 
-  inline ArrayRef<std::optional<Value>> getMaybeValues() const {
-    return {values.data(), values.size()};
+  inline SmallVector<std::optional<Value>> getMaybeValues() const {
+    SmallVector<std::optional<Value>> maybeValues;
+    maybeValues.reserve(getNumDimAndSymbolVars());
+    for (unsigned i = 0, e = getNumDimAndSymbolVars(); i < e; i++) {
+      if (hasValue(i))
+        maybeValues.push_back(getValue(i));
+      else
+        maybeValues.push_back(std::nullopt);
+    }
+    return maybeValues;
   }
 
-  inline ArrayRef<std::optional<Value>>
-  getMaybeValues(presburger::VarKind kind) const {
-    assert(kind != VarKind::Local &&
-           "Local variables do not have any value attached to them.");
-    return {values.data() + getVarKindOffset(kind), getNumVarKind(kind)};
-  }
+  inline SmallVector<std::optional<Value>>
+    getMaybeValues(presburger::VarKind kind) const {
+      SmallVector<std::optional<Value>> maybeValues;
+      maybeValues.reserve(getNumVarKind(kind));
+      for (unsigned i = 0, e = getNumVarKind(kind); i < e; i++) {
+        Identifier id = space.getId(kind, i);
+        if (id.hasValue())
+          maybeValues.push_back(space.getId(kind, i).getValue<Value>());
+        else
+          maybeValues.push_back(std::nullopt);
+      }
+      return maybeValues;
+    }
 
   /// Returns true if the pos^th variable has an associated Value.
   inline bool hasValue(unsigned pos) const {
     assert(pos < getNumDimAndSymbolVars() && "Invalid position");
-    return values[pos].has_value();
+    VarKind kind = getVarKindAt(pos);
+    unsigned relativePos = pos - getVarKindOffset(kind);
+    return space.getId(kind, relativePos).hasValue();
+  }
+
+  void resetValues() {
+    space.resetIds();
   }
 
   unsigned appendDimVar(ValueRange vals);
@@ -360,7 +389,9 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   /// Sets the Value associated with the pos^th variable.
   inline void setValue(unsigned pos, Value val) {
     assert(pos < getNumDimAndSymbolVars() && "invalid var position");
-    values[pos] = val;
+    VarKind kind = getVarKindAt(pos);
+    unsigned relativePos = pos - getVarKindOffset(kind);
+    space.getId(kind, relativePos) = presburger::Identifier(val);
   }
 
   /// Sets the Values associated with the variables in the range [start, end).
@@ -455,17 +486,6 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   // See implementation comments for more details.
   void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false,
                                bool *isResultIntegerExact = nullptr) override;
-
-  /// Returns false if the fields corresponding to various variable counts, or
-  /// equality/inequality buffer sizes aren't consistent; true otherwise. This
-  /// is meant to be used within an assert internally.
-  bool hasConsistentState() const override;
-
-  /// Values corresponding to the (column) non-local variables of this
-  /// constraint system appearing in the order the variables correspond to
-  /// columns. Variables that aren't associated with any Value are set to
-  /// std::nullopt.
-  SmallVector<std::optional<Value>, 8> values;
 };
 
 /// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h
index a27583877b603c..4134aef8174bc1 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h
@@ -15,6 +15,7 @@
 #ifndef MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H
 #define MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H
 
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/SmallVector.h"
@@ -115,7 +116,7 @@ struct MemRefAccess {
   ///
   /// Returns failure for yet unimplemented/unsupported cases (see docs of
   /// mlir::getIndexSet and mlir::getRelationFromMap for these cases).
-  LogicalResult getAccessRelation(FlatAffineRelation &accessRel) const;
+  LogicalResult getAccessRelation(presburger::IntegerRelation &accessRel) const;
 
   /// Populates 'accessMap' with composition of AffineApplyOps reachable from
   /// 'indices'.
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
index 7c500f13895af1..efd28a88e93752 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
@@ -251,9 +251,9 @@ class FlatAffineRelation : public FlatAffineValueConstraints {
 /// For AffineValueMap, the domain and symbols have Value set corresponding to
 /// the Value in `map`. Returns failure if the AffineMap could not be flattened
 /// (i.e., semi-affine is not yet handled).
-LogicalResult getRelationFromMap(AffineMap &map, FlatAffineRelation &rel);
+LogicalResult getRelationFromMap(AffineMap &map, presburger::IntegerRelation &rel);
 LogicalResult getRelationFromMap(const AffineValueMap &map,
-                                 FlatAffineRelation &rel);
+                                 presburger::IntegerRelation &rel);
 
 } // namespace affine
 } // namespace mlir
diff --git a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
index 69846a356e0cc4..8414ca4b8b8e4a 100644
--- a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
+++ b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
@@ -817,13 +817,12 @@ FlatLinearValueConstraints::FlatLinearValueConstraints(IntegerSet set,
                             set.getNumDims() + set.getNumSymbols() + 1,
                             set.getNumDims(), set.getNumSymbols(),
                             /*numLocals=*/0) {
-  // Populate values.
-  if (operands.empty()) {
-    values.resize(getNumDimAndSymbolVars(), std::nullopt);
-  } else {
-    assert(set.getNumInputs() == operands.size() && "operand count mismatch");
-    values.assign(operands.begin(), operands.end());
-  }
+  // Use values in space for FlatLinearValueConstraints.
+  space.resetIds();
+  // Set the values for the non-local variables.
+  for (unsigned i = 0, e = operands.size(); i < e; ++i)
+    setValue(i, operands[i]);
+
 
   // Flatten expressions and add them to the constraint system.
   std::vector<SmallVector<int64_t, 8>> flatExprs;
@@ -873,11 +872,6 @@ unsigned FlatLinearValueConstraints::insertVar(VarKind kind, unsigned pos,
                                                unsigned num) {
   unsigned absolutePos = IntegerPolyhedron::insertVar(kind, pos, num);
 
-  if (kind != VarKind::Local) {
-    values.insert(values.begin() + absolutePos, num, std::nullopt);
-    assert(values.size() == getNumDimAndSymbolVars());
-  }
-
   return absolutePos;
 }
 
@@ -890,11 +884,9 @@ unsigned FlatLinearValueConstraints::insertVar(VarKind kind, unsigned pos,
   unsigned absolutePos = IntegerPolyhedron::insertVar(kind, pos, num);
 
   // If a Value is provided, insert it; otherwise use std::nullopt.
-  for (unsigned i = 0; i < num; ++i)
-    values.insert(values.begin() + absolutePos + i,
-                  vals[i] ? std::optional<Value>(vals[i]) : std::nullopt);
+  for (unsigned i = 0, e = vals.size(); i < e; ++i)
+    setValue(absolutePos + i, vals[i]);
 
-  assert(values.size() == getNumDimAndSymbolVars());
   return absolutePos;
 }
 
@@ -902,10 +894,7 @@ unsigned FlatLinearValueConstraints::insertVar(VarKind kind, unsigned pos,
 /// associated with the same set of variables, appearing in the same order.
 static bool areVarsAligned(const FlatLinearValueConstraints &a,
                            const FlatLinearValueConstraints &b) {
-  return a.getNumDimVars() == b.getNumDimVars() &&
-         a.getNumSymbolVars() == b.getNumSymbolVars() &&
-         a.getNumVars() == b.getNumVars() &&
-         a.getMaybeValues().equals(b.getMaybeValues());
+  return a.getSpace().isAligned(b.getSpace());
 }
 
 /// Calls areVarsAligned to check if two constraint systems have the same set
@@ -928,12 +917,14 @@ static bool LLVM_ATTRIBUTE_UNUSED areVarsUnique(
     return true;
 
   SmallPtrSet<Value, 8> uniqueVars;
-  ArrayRef<std::optional<Value>> maybeValues =
-      cst.getMaybeValues().slice(start, end - start);
-  for (std::optional<Value> val : maybeValues) {
+  SmallVector<std::optional<Value>, 8> maybeValuesAll = cst.getMaybeValues();
+  ArrayRef<std::optional<Value>> maybeValues = {maybeValuesAll.data() + start,
+                                                maybeValuesAll.data() + end};
+
+  for (std::optional<Value> val : maybeValues)
     if (val && !uniqueVars.insert(*val).second)
       return false;
-  }
+
   return true;
 }
 
@@ -1058,20 +1049,9 @@ void FlatLinearValueConstraints::mergeSymbolVars(
          "expected same number of symbols");
 }
 
-bool FlatLinearValueConstraints::hasConsistentState() const {
-  return IntegerPolyhedron::hasConsistentState() &&
-         values.size() == getNumDimAndSymbolVars();
-}
-
 void FlatLinearValueConstraints::removeVarRange(VarKind kind, unsigned varStart,
                                                 unsigned varLimit) {
   IntegerPolyhedron::removeVarRange(kind, varStart, varLimit);
-  unsigned offset = getVarKindOffset(kind);
-
-  if (kind != VarKind::Local) {
-    values.erase(values.begin() + varStart + offset,
-                 values.begin() + varLimit + offset);
-  }
 }
 
 AffineMap
@@ -1089,14 +1069,15 @@ FlatLinearValueConstraints::computeAlignedMap(AffineMap map,
 
   dims.reserve(getNumDimVars());
   syms.reserve(getNumSymbolVars());
-  for (unsigned i = getVarKindOffset(VarKind::SetDim),
-                e = getVarKindEnd(VarKind::SetDim);
-       i < e; ++i)
-    dims.push_back(values[i] ? *values[i] : Value());
-  for (unsigned i = getVarKindOffset(VarKind::Symbol),
-                e = getVarKindEnd(VarKind::Symbol);
-       i < e; ++i)
-    syms.push_back(values[i] ? *values[i] : Value());
+  for (unsigned i = 0, e = getNumVarKind(VarKind::SetDim); i < e; ++i) {
+    Identifier id = space.getId(VarKind::SetDim, i);
+    dims.push_back(id.hasValue() ? Value(id.getValue<Value>()) : Value());
+  }
+  for (unsigned i = 0, e = getNumVarKind(VarKind::Symbol); i < e; ++i) {
+    Identifier id = space.getId(VarKind::Symbol, i);
+    syms.push_back(id.hasValue() ? Value(id.getValue<Value>()) : Value());
+  }
+
 
   AffineMap alignedMap =
       alignAffineMapWithValues(map, operands, dims, syms, newSymsPtr);
@@ -1110,8 +1091,7 @@ FlatLinearValueConstraints::computeAlignedMap(AffineMap map,
 bool FlatLinearValueConstraints::findVar(Value val, unsigned *pos,
                                          unsigned offset) const {
   unsigned i = offset;
-  for (const auto &mayBeVar :
-       ArrayRef<std::optional<Value>>(values).drop_front(offset)) {
+  for (const auto &mayBeVar : getMaybeValues()) {
     if (mayBeVar && *mayBeVar == val) {
       *pos = i;
       return true;
@@ -1122,25 +1102,12 @@ bool FlatLinearValueConstraints::findVar(Value val, unsigned *pos,
 }
 
 bool FlatLinearValueConstraints::containsVar(Value val) const {
-  return llvm::any_of(values, [&](const std::optional<Value> &mayBeVar) {
-    return mayBeVar && *mayBeVar == val;
-  });
+  unsigned pos;
+  return findVar(val, &pos, 0);
 }
 
 void FlatLinearValueConstraints::swapVar(unsigned posA, unsigned posB) {
   IntegerPolyhedron::swapVar(posA, posB);
-
-  if (getVarKindAt(posA) == VarKind::Local &&
-      getVarKindAt(posB) == VarKind::Local)
-    return;
-
-  // Treat value of a local variable as std::nullopt.
-  if (getVarKindAt(posA) == VarKind::Local)
-    values[posB] = std::nullopt;
-  else if (getVarKindAt(posB) == VarKind::Local)
-    values[posA] = std::nullopt;
-  else
-    std::swap(values[posA], values[posB]);
 }
 
 void FlatLinearValueConstraints::addBound(BoundType type, Value val,
@@ -1182,27 +1149,13 @@ void FlatLinearValueConstraints::printSpace(raw_ostream &os) const {
 
 void FlatLinearValueConstraints::clearAndCopyFrom(
     const IntegerRelation &other) {
-
-  if (auto *otherValueSet =
-          dyn_cast<const FlatLinearValueConstraints>(&other)) {
-    *this = *otherValueSet;
-  } else {
-    *static_cast<IntegerRelation *>(this) = other;
-    values.clear();
-    values.resize(getNumDimAndSymbolVars(), std::nullopt);
-  }
+  IntegerPolyhedron::clearAndCopyFrom(other);
 }
 
 void FlatLinearValueConstraints::fourierMotzkinEliminate(
     unsigned pos, bool darkShadow, bool *isResultIntegerExact) {
-  SmallVector<std::optional<Value>, 8> newVals = values;
-  if (getVarKindAt(pos) != VarKind::Local)
-    newVals.erase(newVals.begin() + pos);
-  // Note: Base implementation discards all associated Values.
   IntegerPolyhedron::fourierMotzkinEliminate(pos, darkShadow,
                                              isResultIntegerExact);
-  values = newVals;
-  assert(values.size() == getNumDimAndSymbolVars());
 }
 
 void FlatLinearValueConstraints::projectOut(Value val) {
@@ -1215,11 +1168,7 @@ void FlatLinearValueConstraints::projectOut(Value val) {
 
 LogicalResult FlatLinearValueConstraints::unionBoundingBox(
     const FlatLinearValueConstraints &otherCst) {
-  assert(otherCst.getNumDimVars() == getNumDimVars() && "dims mismatch");
-  assert(otherCst.getMaybeValues()
-             .slice(0, getNumDimVars())
-             .equals(getMaybeValues().slice(0, getNumDimVars())) &&
-         "dim values mismatch");
+  assert(otherCst.getSpace().isAligned(getSpace(), VarKind::SetDim) && "dims mismatch");
   assert(otherCst.getNumLocalVars() == 0 && "local vars not supported here");
   assert(getNumLocalVars() == 0 && "local vars not supported yet here");
 
diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index 1ba0bc8b6bfbe5..f95fd5213227c5 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -12,7 +12,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
+#include "mlir/Analysis/Presburger/PresburgerSpace.h"
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
@@ -385,7 +388,7 @@ static void
 addOrderingConstraints(const FlatAffineValueConstraints &srcDomain,
                        const FlatAffineValueConstraints &dstDomain,
                        unsigned loopDepth,
-                       FlatAffineValueConstraints *dependenceDomain) {
+                       IntegerRelation *dependenceDomain) {
   unsigned numCols = dependenceDomain->getNumCols();
   SmallVector<int64_t, 4> eq(numCols);
   unsigned numSrcDims = srcDomain.getNumDimVars();
@@ -411,7 +414,7 @@ addOrderingConstraints(const FlatAffineValueConstraints &srcDomain,
 static void computeDirectionVector(
     const FlatAffineValueConstraints &srcDomain,
     const FlatAffineValueConstraints &dstDomain, unsigned loopDepth,
-    FlatAffineValueConstraints *dependenceDomain,
+    IntegerRelation *dependenceDomain,
     SmallVector<DependenceComponent, 2> *dependenceComponents) {
   // Find the number of common loops shared by src and dst accesses.
   SmallVector<AffineForOp, 4> commonLoops;
@@ -423,7 +426,7 @@ static void computeDirectionVector(
   unsigned numIdsToEliminate = dependenceDomain->getNumVars();
   // Add new variables to 'dependenceDomain' to represent the direction
   // constraints for each shared loop.
-  dependenceDomain->insertDimVar(/*pos=*/0, /*num=*/numCommonLoops);
+  dependenceDomain->insertVar(VarKind::Domain, /*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.
@@ -457,7 +460,7 @@ static void computeDirectionVector(
   }
 }
 
-LogicalResult MemRefAccess::getAccessRelation(FlatAffineRelation &rel) const {
+LogicalResult MemRefAccess::getAccessRelation(IntegerRelation &rel) const {
   // Create set corresponding to domain of access.
   FlatAffineValueConstraints domain;
   if (failed(getOpIndexSet(opInst, &domain)))
@@ -469,25 +472,29 @@ LogicalResult MemRefAccess::getAccessRelation(FlatAffineRelation &rel) const {
   if (failed(getRelationFromMap(accessValueMap, rel)))
     return failure();
 
-  FlatAffineRelation domainRel(rel.getNumDomainDims(), /*numRangeDims=*/0,
-                               domain);
+  IntegerRelation domainRel = domain;
+  domainRel.convertVarKind(VarKind::SetDim, 0, rel.getNumDomainVars(), VarKind::Domain);
 
-  // Merge and align domain ids of `ret` and ids of `domain`. Since the domain
+  // Merge and align domain ids of `rel` and ids of `domain`. Since the domain
   // of the access map is a subset of the domain of access, the domain ids of
-  // `ret` are guranteed to be a subset of ids of `domain`.
+  // `rel` are guranteed to be a subset of ids of `domain`.
   for (unsigned i = 0, e = domain.getNumDimVars(); i < e; ++i) {
-    unsigned loc;
-    if (rel.findVar(domain.getValue(i), &loc)) {
-      rel.swapVar(i, loc);
+    const PresburgerSpace &relSpace = rel.getSpace();
+    const Identifier *findBegin = relSpace.getIds(VarKind::Domain).begin() + i;
+    const Identifier *findEnd = relSpace.getIds(VarKind::Domain).end();
+    const Identifier domainIdi = Identifier(domain.getValue(i));
+    const Identifier *itr = std::find(findBegin, findEnd, domainIdi);
+    if (itr != findEnd) {
+      rel.swapVar(i, i + std::distance(findBegin, itr));
     } else {
-      rel.insertDomainVar(i);
-      rel.setValue(i, domain.getValue(i));
+      rel.insertVar(VarKind::Domain, i);
+      rel.setId(VarKind::Domain, i, domainIdi);
     }
   }
 
   // Append domain constraints to `rel`.
-  domainRel.appendRangeVar(rel.getNumRangeDims());
-  domainRel.mergeSymbolVars(rel);
+  domainRel.appendVar(VarKind::Range, rel.getNumRangeVars());
+  domainRel.mergeAndAlignSymbols(rel);
   domainRel.mergeLocalVars(rel);
   rel.append(domainRel);
 
@@ -624,7 +631,8 @@ DependenceResult mlir::affine::checkMemrefAccessDependence(
     return DependenceResult::Failure;
 
   // Create access relation from each MemRefAccess.
-  FlatAffineRelation srcRel, dstRel;
+  PresburgerSpace space = PresburgerSpace::getRelationSpace();
+  IntegerRelation srcRel(space), dstRel(space);
   if (failed(srcAccess.getAccessRelation(srcRel)))
     return DependenceResult::Failure;
   if (failed(dstAccess.getAccessRelation(dstRel)))
@@ -668,7 +676,7 @@ DependenceResult mlir::affine::checkMemrefAccessDependence(
   LLVM_DEBUG(dstRel.dump());
 
   if (dependenceConstraints)
-    *dependenceConstraints = dstRel;
+    *dependenceConstraints = cast<FlatAffineValueConstraints>(dstRel);
   return DependenceResult::HasDependence;
 }
 
diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
index 469298d3e8f43f..e0d03ab5bf8c83 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
@@ -11,12 +11,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/LinearTransform.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/LLVM.h"
@@ -372,9 +374,8 @@ FlatAffineValueConstraints FlatAffineRelation::getRangeSet() const {
 void FlatAffineRelation::compose(const FlatAffineRelation &other) {
   assert(getNumDomainDims() == other.getNumRangeDims() &&
          "Domain of this and range of other do not match");
-  assert(std::equal(values.begin(), values.begin() + getNumDomainDims(),
-                    other.values.begin() + other.getNumDomainDims()) &&
-         "Domain of this and range of other do not match");
+  assert(space.getDomainSpace().isAligned(other.getSpace().getRangeSpace()) &&
+      "Values of domain of this and range of other do not match");
 
   FlatAffineRelation rel = other;
 
@@ -491,12 +492,15 @@ void FlatAffineRelation::removeVarRange(VarKind kind, unsigned varStart,
 }
 
 LogicalResult mlir::affine::getRelationFromMap(AffineMap &map,
-                                               FlatAffineRelation &rel) {
+                                               IntegerRelation &rel) {
   // Get flattened affine expressions.
   std::vector<SmallVector<int64_t, 8>> flatExprs;
   FlatAffineValueConstraints localVarCst;
   if (failed(getFlattenedAffineExprs(map, &flatExprs, &localVarCst)))
     return failure();
+  // Add identifiers to the local constraints. We need to do this since
+  // getFlattenedAffineExprs creates a FlatLinearConstraints with no
+  // identifiers.
 
   unsigned oldDimNum = localVarCst.getNumDimVars();
   unsigned oldCols = localVarCst.getNumCols();
@@ -506,6 +510,10 @@ LogicalResult mlir::affine::getRelationFromMap(AffineMap &map,
   // Add range as the new expressions.
   localVarCst.appendDimVar(numRangeVars);
 
+  localVarCst.resetValues();
+  for(unsigned i = 0, e = localVarCst.getNumDimAndSymbolVars(); i < e; ++i)
+    localVarCst.setValue(i, Value());
+
   // Add equalities between source and range.
   SmallVector<int64_t, 8> eq(localVarCst.getNumCols());
   for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
@@ -522,23 +530,24 @@ LogicalResult mlir::affine::getRelationFromMap(AffineMap &map,
   }
 
   // Create relation and return success.
-  rel = FlatAffineRelation(numDomainVars, numRangeVars, localVarCst);
+  rel = localVarCst;
+  rel.convertVarKind(VarKind::SetDim, 0, numDomainVars, VarKind::Domain);
   return success();
 }
 
 LogicalResult mlir::affine::getRelationFromMap(const AffineValueMap &map,
-                                               FlatAffineRelation &rel) {
+                                               IntegerRelation &rel) {
 
   AffineMap affineMap = map.getAffineMap();
   if (failed(getRelationFromMap(affineMap, rel)))
     return failure();
 
   // Set symbol values for domain dimensions and symbols.
-  for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i)
-    rel.setValue(i, map.getOperand(i));
-  for (unsigned i = rel.getNumDimVars(), e = rel.getNumDimAndSymbolVars();
-       i < e; ++i)
-    rel.setValue(i, map.getOperand(i - rel.getNumRangeDims()));
+  for (unsigned i = 0, e = rel.getNumDomainVars(); i < e; ++i)
+    rel.setId(VarKind::Domain, i, Identifier(map.getOperand(i)));
+
+  for(unsigned i = 0, e = rel.getNumSymbolVars(); i < e; ++i)
+    rel.setId(VarKind::Symbol, i, Identifier(map.getOperand(i - rel.getNumRangeVars())));
 
   return success();
 }
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 578d03c629285a..2a35b02185190c 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -13,6 +13,8 @@
 
 #include "mlir/Dialect/Affine/Utils.h"
 
+#include "mlir/Analysis/Presburger/PresburgerRelation.h"
+#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"



More information about the Mlir-commits mailing list