[MLIR] Move LinearTransform to Presburger/

This patch moves LinearTransform to Presburger/ and makes it use
IntegerPolyhedron instead of FlatAffineConstraints. Also modifies its usage in
`FlatAffineConstraints::findIntegerSample` to support the changes.

This patch is part of a series of patches for moving presburger math functionality into Presburger directory.

Reviewed By: arjunp

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

Author: Groverkss

mlir/include/mlir/Analysis/AffineStructures.h

@@ -81,6 +81,9 @@ class FlatAffineConstraints : public IntegerPolyhedron {
                               1,
                           numDims, numSymbols, numLocals) {}
 
+  explicit FlatAffineConstraints(const IntegerPolyhedron &poly)
+      : IntegerPolyhedron(poly) {}
+
   /// Return a system with no constraints, i.e., one which is satisfied by all
   /// points.
   static FlatAffineConstraints getUniverse(unsigned numDims = 0,
@@ -212,10 +215,6 @@ class FlatAffineConstraints : public IntegerPolyhedron {
   void projectOut(unsigned pos, unsigned num);
   inline void projectOut(unsigned pos) { return projectOut(pos, 1); }
 
-  /// Sets the `values.size()` identifiers starting at `po`s to the specified
-  /// values and removes them.
-  void setAndEliminate(unsigned pos, ArrayRef<int64_t> values);
-
   /// Changes the partition between dimensions and symbols. Depending on the new
   /// symbol count, either a chunk of trailing dimensional identifiers becomes
   /// symbols, or some of the leading symbols become dimensions.

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

@@ -185,6 +185,10 @@ class IntegerPolyhedron {
   /// Removes all equalities and inequalities.
   void clearConstraints();
 
+  /// Sets the `values.size()` identifiers starting at `po`s to the specified
+  /// values and removes them.
+  void setAndEliminate(unsigned pos, ArrayRef<int64_t> values);
+
   /// Gather positions of all lower and upper bounds of the identifier at `pos`,
   /// and optionally any equalities on it. In addition, the bounds are to be
   /// independent of identifiers in position range [`offset`, `offset` + `num`).

mlir/include/mlir/Analysis/LinearTransform.h renamed to mlir/include/mlir/Analysis/Presburger/LinearTransform.h

@@ -6,14 +6,14 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// Support for linear transforms and applying them to FlatAffineConstraints.
+// Support for linear transforms and applying them to an IntegerPolyhedron.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_LINEARTRANSFORM_H
 #define MLIR_ANALYSIS_LINEARTRANSFORM_H
 
-#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Analysis/Presburger/IntegerPolyhedron.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
 #include "llvm/ADT/SmallVector.h"
 
@@ -33,9 +33,9 @@ class LinearTransform {
   static std::pair<unsigned, LinearTransform>
   makeTransformToColumnEchelon(Matrix m);
 
-  // Returns a FlatAffineConstraints having a constraint vector vT for every
-  // constraint vector v in fac, where T is this transform.
-  FlatAffineConstraints applyTo(const FlatAffineConstraints &fac) const;
+  // Returns an IntegerPolyhedron having a constraint vector vT for every
+  // constraint vector v in poly, where T is this transform.
+  IntegerPolyhedron applyTo(const IntegerPolyhedron &poly) const;
 
   // The given vector is interpreted as a row vector v. Post-multiply v with
   // this transform, say T, and return vT.

mlir/lib/Analysis/AffineStructures.cpp

@@ -11,7 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineStructures.h"
-#include "mlir/Analysis/LinearTransform.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"
@@ -1090,11 +1090,11 @@ FlatAffineConstraints::findIntegerSample() const {
       LinearTransform::makeTransformToColumnEchelon(std::move(m));
   const LinearTransform &transform = result.second;
   // 1) Apply T to S to obtain S*T.
-  FlatAffineConstraints transformedSet = transform.applyTo(*this);
+  IntegerPolyhedron transformedSet = transform.applyTo(*this);
 
   // 2) Remove the unbounded dimensions and constraints involving them to
   // obtain a bounded set.
-  FlatAffineConstraints boundedSet = transformedSet;
+  FlatAffineConstraints boundedSet(transformedSet);
   unsigned numBoundedDims = result.first;
   unsigned numUnboundedDims = getNumIds() - numBoundedDims;
   removeConstraintsInvolvingSuffixDims(boundedSet, numUnboundedDims);
@@ -1111,7 +1111,7 @@ FlatAffineConstraints::findIntegerSample() const {
   // 4) Substitute the values of the bounded dimensions into S*T to obtain a
   // full-dimensional cone, which necessarily contains an integer sample.
   transformedSet.setAndEliminate(0, *boundedSample);
-  FlatAffineConstraints &cone = transformedSet;
+  IntegerPolyhedron &cone = transformedSet;
 
   // 5) Obtain an integer sample from the cone.
   //
@@ -1139,10 +1139,10 @@ FlatAffineConstraints::findIntegerSample() const {
   // negative a_i, so we accomodate this by shifting the inequality by this
   // amount for the shrunken cone.
   for (unsigned i = 0, e = cone.getNumInequalities(); i < e; ++i) {
-    for (unsigned j = 0; j < cone.numIds; ++j) {
+    for (unsigned j = 0; j < cone.getNumIds(); ++j) {
       int64_t coeff = cone.atIneq(i, j);
       if (coeff < 0)
-        cone.atIneq(i, cone.numIds) += coeff;
+        cone.atIneq(i, cone.getNumIds()) += coeff;
     }
   }
 
@@ -2303,24 +2303,6 @@ static int findEqualityToConstant(const FlatAffineConstraints &cst,
   return -1;
 }
 
-void FlatAffineConstraints::setAndEliminate(unsigned pos,
-                                            ArrayRef<int64_t> values) {
-  if (values.empty())
-    return;
-  assert(pos + values.size() <= getNumIds() &&
-         "invalid position or too many values");
-  // Setting x_j = p in sum_i a_i x_i + c is equivalent to adding p*a_j to the
-  // constant term and removing the id x_j. We do this for all the ids
-  // pos, pos + 1, ... pos + values.size() - 1.
-  for (unsigned r = 0, e = getNumInequalities(); r < e; r++)
-    for (unsigned i = 0, numVals = values.size(); i < numVals; ++i)
-      atIneq(r, getNumCols() - 1) += atIneq(r, pos + i) * values[i];
-  for (unsigned r = 0, e = getNumEqualities(); r < e; r++)
-    for (unsigned i = 0, numVals = values.size(); i < numVals; ++i)
-      atEq(r, getNumCols() - 1) += atEq(r, pos + i) * values[i];
-  removeIdRange(pos, pos + values.size());
-}
-
 LogicalResult FlatAffineConstraints::constantFoldId(unsigned pos) {
   assert(pos < getNumIds() && "invalid position");
   int rowIdx;

mlir/lib/Analysis/CMakeLists.txt

@@ -6,7 +6,6 @@ set(LLVM_OPTIONAL_SOURCES
   CallGraph.cpp
   DataFlowAnalysis.cpp
   DataLayoutAnalysis.cpp
-  LinearTransform.cpp
   Liveness.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
@@ -48,7 +47,6 @@ add_mlir_library(MLIRAnalysis
 add_mlir_library(MLIRLoopAnalysis
   AffineAnalysis.cpp
   AffineStructures.cpp
-  LinearTransform.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   PresburgerSet.cpp

mlir/lib/Analysis/Presburger/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_mlir_library(MLIRPresburger
   IntegerPolyhedron.cpp
+  LinearTransform.cpp
   Matrix.cpp
   Simplex.cpp
   Utils.cpp

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

@@ -285,6 +285,24 @@ bool IntegerPolyhedron::hasConsistentState() const {
   return true;
 }
 
+void IntegerPolyhedron::setAndEliminate(unsigned pos,
+                                        ArrayRef<int64_t> values) {
+  if (values.empty())
+    return;
+  assert(pos + values.size() <= getNumIds() &&
+         "invalid position or too many values");
+  // Setting x_j = p in sum_i a_i x_i + c is equivalent to adding p*a_j to the
+  // constant term and removing the id x_j. We do this for all the ids
+  // pos, pos + 1, ... pos + values.size() - 1.
+  for (unsigned r = 0, e = getNumInequalities(); r < e; r++)
+    for (unsigned i = 0, numVals = values.size(); i < numVals; ++i)
+      atIneq(r, getNumCols() - 1) += atIneq(r, pos + i) * values[i];
+  for (unsigned r = 0, e = getNumEqualities(); r < e; r++)
+    for (unsigned i = 0, numVals = values.size(); i < numVals; ++i)
+      atEq(r, getNumCols() - 1) += atEq(r, pos + i) * values[i];
+  removeIdRange(pos, pos + values.size());
+}
+
 void IntegerPolyhedron::printSpace(raw_ostream &os) const {
   os << "\nConstraints (" << getNumDimIds() << " dims, " << getNumSymbolIds()
      << " symbols, " << getNumLocalIds() << " locals), (" << getNumConstraints()

mlir/lib/Analysis/LinearTransform.cpp renamed to mlir/lib/Analysis/Presburger/LinearTransform.cpp

@@ -6,8 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Analysis/LinearTransform.h"
-#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Analysis/Presburger/LinearTransform.h"
+#include "mlir/Analysis/Presburger/IntegerPolyhedron.h"
 
 namespace mlir {
 
@@ -135,12 +135,12 @@ LinearTransform::preMultiplyColumn(ArrayRef<int64_t> colVec) const {
   return result;
 }
 
-FlatAffineConstraints
-LinearTransform::applyTo(const FlatAffineConstraints &fac) const {
-  FlatAffineConstraints result(fac.getNumIds());
+IntegerPolyhedron
+LinearTransform::applyTo(const IntegerPolyhedron &poly) const {
+  IntegerPolyhedron result(poly.getNumIds());
 
-  for (unsigned i = 0, e = fac.getNumEqualities(); i < e; ++i) {
-    ArrayRef<int64_t> eq = fac.getEquality(i);
+  for (unsigned i = 0, e = poly.getNumEqualities(); i < e; ++i) {
+    ArrayRef<int64_t> eq = poly.getEquality(i);
 
     int64_t c = eq.back();
 
@@ -149,8 +149,8 @@ LinearTransform::applyTo(const FlatAffineConstraints &fac) const {
     result.addEquality(newEq);
   }
 
-  for (unsigned i = 0, e = fac.getNumInequalities(); i < e; ++i) {
-    ArrayRef<int64_t> ineq = fac.getInequality(i);
+  for (unsigned i = 0, e = poly.getNumInequalities(); i < e; ++i) {
+    ArrayRef<int64_t> ineq = poly.getInequality(i);
 
     int64_t c = ineq.back();
 

mlir/unittests/Analysis/CMakeLists.txt

@@ -2,7 +2,6 @@ add_mlir_unittest(MLIRAnalysisTests
   AffineStructuresParser.cpp
   AffineStructuresParserTest.cpp
   AffineStructuresTest.cpp
-  LinearTransformTest.cpp
   PresburgerSetTest.cpp
 )
 

mlir/unittests/Analysis/Presburger/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_mlir_unittest(MLIRPresburgerTests
   IntegerPolyhedronTest.cpp
+  LinearTransformTest.cpp
   MatrixTest.cpp
   SimplexTest.cpp
   ../AffineStructuresParser.cpp