[MLIR] Move LinearTransform to Presburger/

[Groverkss]

Summary

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.

[arjunp]

I would mention in the summary that this is part of a series of patches moving functionality into Presburger.

mlir/include/mlir/Analysis/AffineStructures.h

84

It's probably better for this to be explicit right?

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

36–38

mlir/lib/Analysis/AffineStructures.cpp

1093

For this patch you can keep this as an FAC right? (convert the result, then you don't need any changes below and don't have to move setAndEliminate?)

[Groverkss]

• Address Arjun's comments

[Groverkss]

mlir/lib/Analysis/AffineStructures.cpp

1093

LinearTransform returns a IntegerPolyhedron. I don't think you can convert a base class (IntegerPolyhedron) to a derived class (FlatAffineConstriants) here without doing a copy. Is there any specific way you are suggesting?

[arjunp]

Please update the summary as well, as mentioned earlier.

mlir/lib/Analysis/AffineStructures.cpp

1093

I guess this is fine then.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   FlatAffineConstraints(unsigned numDims = 0, unsigned numSymbols = 0,
                         unsigned numLocals = 0)
       : IntegerPolyhedron(/*numReservedInequalities=*/0,
                           /*numReservedEqualities=*/0,
                           /*numReservedCols=*/numDims + numSymbols + numLocals +
                               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,
                                            unsigned numSymbols = 0) {
     return FlatAffineConstraints(numDims, numSymbols);
   }
 
   /// Creates an affine constraint system from an IntegerSet.
public:
   /// Projects out (aka eliminates) `num` identifiers starting at position
   /// `pos`. The resulting constraint system is the shadow along the dimensions
   /// that still exist. This method may not always be integer exact.
   // TODO: deal with integer exactness when necessary - can return a value to
   // mark exactness for example.
   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.
   void setDimSymbolSeparation(unsigned newSymbolCount);
 
   /// Tries to fold the specified identifier to a constant using a trivial
   /// equality detection; if successful, the constant is substituted for the
   /// identifier everywhere in the constraint system and then removed from the

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

public:
   void removeInequalityRange(unsigned start, unsigned end);
 
   /// Swap the posA^th identifier with the posB^th identifier.
   virtual void swapId(unsigned posA, unsigned posB);
 
   /// 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`).
   void
   getLowerAndUpperBoundIndices(unsigned pos,
                                SmallVectorImpl<unsigned> *lbIndices,
                                SmallVectorImpl<unsigned> *ubIndices,
                                SmallVectorImpl<unsigned> *eqIndices = nullptr,

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

 //===- LinearTransform.h - MLIR LinearTransform Class -----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// 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"
 
 namespace mlir {
 
 class LinearTransform {
 public:
   explicit LinearTransform(Matrix &&oMatrix);
   explicit LinearTransform(const Matrix &oMatrix);
 
   // Returns a linear transform T such that MT is M in column echelon form.
   // Also returns the number of non-zero columns in MT.
   //
   // Specifically, T is such that in every column the first non-zero row is
   // strictly below that of the previous column, and all columns which have only
   // zeros are at the end.
   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.
   SmallVector<int64_t, 8> postMultiplyRow(ArrayRef<int64_t> rowVec) const;
 
   // The given vector is interpreted as a column vector v. Pre-multiply v with
   // this transform, say T, and return Tv.
   SmallVector<int64_t, 8> preMultiplyColumn(ArrayRef<int64_t> colVec) const;
 
 private:
   Matrix matrix;
 };
 
 } // namespace mlir
 #endif // MLIR_ANALYSIS_LINEARTRANSFORM_H

mlir/lib/Analysis/AffineStructures.cpp

 //===- AffineStructures.cpp - MLIR Affine Structures Class-----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Structures for affine/polyhedral analysis of affine dialect ops.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/IntegerSet.h"
FlatAffineConstraints::findIntegerSample() const {
   // In column echelon form, each row of m occupies only the first rank(m)
   // columns and has zeros on the other columns. The transform T that brings S
   // to column echelon form is unimodular as well, so this is a suitable
   // transform to use in step 1 of the algorithm.
   std::pair<unsigned, LinearTransform> result =
       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);
   boundedSet.removeIdRange(numBoundedDims, boundedSet.getNumIds());
 
   // 3) Try to obtain a sample from the bounded set.
   Optional<SmallVector<int64_t, 8>> boundedSample =
       Simplex(boundedSet).findIntegerSample();
   if (!boundedSample)
     return {};
   assert(boundedSet.containsPoint(*boundedSample) &&
          "Simplex returned an invalid sample!");
 
   // 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.
   //
   // We shrink the cone such that for any rational point in the shrunken cone,
   // rounding up each of the point's coordinates produces a point that still
   // lies in the original cone.
   //
   // Rounding up a point x adds a number e_i in [0, 1) to each coordinate x_i.
FlatAffineConstraints::findIntegerSample() const {
   // is minimized when we add 1 to the x_i with negative coefficient a_i and
   // keep the other x_i the same. In the example, we would get x = (3, 1, 1),
   // changing the value of the LHS by -3 + -7 = -10.
   //
   // In general, the value of the LHS can change by at most the sum of the
   // 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;
     }
   }
 
   // Obtain an integer sample in the cone by rounding up a rational point from
   // the shrunken cone. Shrinking the cone amounts to shifting its apex
   // "inwards" without changing its "shape"; the shrunken cone is still a
   // full-dimensional cone and is hence non-empty.
   Simplex shrunkenConeSimplex(cone);
for (unsigned r = 0, e = cst.getNumEqualities(); r < e; r++) {
     }
     if (c == f)
       // Equality is free of other identifiers.
       return r;
   }
   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;
   if ((rowIdx = findEqualityToConstant(*this, pos)) == -1)
     return failure();
 
   // atEq(rowIdx, pos) is either -1 or 1.
   assert(atEq(rowIdx, pos) * atEq(rowIdx, pos) == 1);

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AliasAnalysis.cpp
   AffineAnalysis.cpp
   AffineStructures.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
   DataLayoutAnalysis.cpp
-  LinearTransform.cpp
   Liveness.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   NumberOfExecutions.cpp
   PresburgerSet.cpp
   SliceAnalysis.cpp
   Utils.cpp
 
add_mlir_library(MLIRAnalysis
   MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
   MLIRViewLikeInterface
   )
 
 add_mlir_library(MLIRLoopAnalysis
   AffineAnalysis.cpp
   AffineStructures.cpp
-  LinearTransform.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   PresburgerSet.cpp
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
 

mlir/lib/Analysis/Presburger/CMakeLists.txt

 add_mlir_library(MLIRPresburger
   IntegerPolyhedron.cpp
+  LinearTransform.cpp
   Matrix.cpp
   Simplex.cpp
   Utils.cpp
 
   DEPENDS
   MLIRBuiltinLocationAttributesIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRSupport
   )

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

bool IntegerPolyhedron::hasConsistentState() const {
 
   // Catches errors where numDims, numSymbols, numIds aren't consistent.
   if (numDims > numIds || numSymbols > numIds || numDims + numSymbols > numIds)
     return false;
 
   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()
      << " constraints)\n";
 }
 
 void IntegerPolyhedron::print(raw_ostream &os) const {
   assert(hasConsistentState());

mlir/lib/Analysis/Presburger/LinearTransform.cpp

 //===- LinearTransform.cpp - MLIR LinearTransform Class -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Analysis/LinearTransform.h"
-#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Analysis/Presburger/LinearTransform.h"
+#include "mlir/Analysis/Presburger/IntegerPolyhedron.h"
 
 namespace mlir {
 
 LinearTransform::LinearTransform(Matrix &&oMatrix) : matrix(oMatrix) {}
 LinearTransform::LinearTransform(const Matrix &oMatrix) : matrix(oMatrix) {}
 
 // Set M(row, targetCol) to its remainder on division by M(row, sourceCol)
 // by subtracting from column targetCol an appropriate integer multiple of
LinearTransform::preMultiplyColumn(ArrayRef<int64_t> colVec) const {
 
   SmallVector<int64_t, 8> result(matrix.getNumRows(), 0);
   for (unsigned row = 0, e = matrix.getNumRows(); row < e; row++)
     for (unsigned i = 0, e = matrix.getNumColumns(); i < e; i++)
       result[row] += matrix(row, i) * colVec[i];
   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();
 
     SmallVector<int64_t, 8> newEq = postMultiplyRow(eq.drop_back());
     newEq.push_back(c);
     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();
 
     SmallVector<int64_t, 8> newIneq = postMultiplyRow(ineq.drop_back());
     newIneq.push_back(c);
     result.addInequality(newIneq);
   }
 
   return result;
 }
 
 } // namespace mlir

mlir/unittests/Analysis/CMakeLists.txt

 add_mlir_unittest(MLIRAnalysisTests
   AffineStructuresParser.cpp
   AffineStructuresParserTest.cpp
   AffineStructuresTest.cpp
-  LinearTransformTest.cpp
   PresburgerSetTest.cpp
 )
 
 target_link_libraries(MLIRAnalysisTests
   PRIVATE
   MLIRLoopAnalysis
   MLIRParser
   )
 
 add_subdirectory(Presburger)

mlir/unittests/Analysis/Presburger/CMakeLists.txt

 add_mlir_unittest(MLIRPresburgerTests
   IntegerPolyhedronTest.cpp
+  LinearTransformTest.cpp
   MatrixTest.cpp
   SimplexTest.cpp
   ../AffineStructuresParser.cpp
 )
 
 target_link_libraries(MLIRPresburgerTests
   PRIVATE MLIRPresburger
   MLIRLoopAnalysis
   MLIRParser
   )

mlir/unittests/Analysis/Presburger/LinearTransformTest.cpp

 //===- LinearTransformTest.cpp - Tests for LinearTransform ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Analysis/LinearTransform.h"
+#include "mlir/Analysis/Presburger/LinearTransform.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 namespace mlir {
 
 void testColumnEchelonForm(const Matrix &m, unsigned expectedRank) {
   unsigned lastAllowedNonZeroCol = 0;
   std::pair<unsigned, LinearTransform> result =