[Mlir-commits] [mlir] 3be3c94 - [MLIR] FlatAffineConstraints: Ensure dimensionalities match when calling mergeLocalIds

[llvmlistbot at llvm.org]

Author: Groverkss
Date: 2021-10-30T17:39:27+05:30
New Revision: 3be3c944a5bacfd208b56853941b0fa4dec3ddcc

URL: https://github.com/llvm/llvm-project/commit/3be3c944a5bacfd208b56853941b0fa4dec3ddcc
DIFF: https://github.com/llvm/llvm-project/commit/3be3c944a5bacfd208b56853941b0fa4dec3ddcc.diff

LOG: [MLIR] FlatAffineConstraints: Ensure dimensionalities match when calling mergeLocalIds

This patch reorders mergeLocalIds usage to merge locals only after number of
dimensions and symbols are same. This does not change any functionality
because it does not matter in what order identifiers are merged, since
the reason to do it is to ensure that two FACs are aligned.

The order ensured in this patch simplifies a subsequent patch to improve
mergeLocalIds which requires dimensions and symbols to be aligned.

Reviewed By: bondhugula

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

Added: 
    

Modified: 
    mlir/include/mlir/Analysis/AffineStructures.h
    mlir/lib/Analysis/AffineAnalysis.cpp
    mlir/lib/Analysis/AffineStructures.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Analysis/AffineStructures.h b/mlir/include/mlir/Analysis/AffineStructures.h
index 59424fb9619b2..0c351c7a2959c 100644
--- a/mlir/include/mlir/Analysis/AffineStructures.h
+++ b/mlir/include/mlir/Analysis/AffineStructures.h
@@ -432,7 +432,8 @@ class FlatAffineConstraints {
 
   /// Merge local ids of `this` and `other`. This is done by appending local ids
   /// of `other` to `this` and inserting local ids of `this` to `other` at start
-  /// of its local ids.
+  /// of its local ids. Number of dimension and symbol ids should match in
+  /// `this` and `other`.
   void mergeLocalIds(FlatAffineConstraints &other);
 
   /// Removes all equalities and inequalities.

diff  --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp
index 8ba6fe216dd14..252b2b5fe0f93 100644
--- a/mlir/lib/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Analysis/AffineAnalysis.cpp
@@ -455,8 +455,8 @@ LogicalResult MemRefAccess::getAccessRelation(FlatAffineRelation &rel) const {
 
   // Append domain constraints to `ret`.
   domainRel.appendRangeId(rel.getNumRangeDims());
-  domainRel.mergeLocalIds(rel);
   domainRel.mergeSymbolIds(rel);
+  domainRel.mergeLocalIds(rel);
   rel.append(domainRel);
 
   return success();

diff  --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp
index 085396b840b9a..abb80104f0382 100644
--- a/mlir/lib/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Analysis/AffineStructures.cpp
@@ -512,9 +512,6 @@ static void mergeAndAlignIds(unsigned offset, FlatAffineValueConstraints *a,
                      b->getMaybeValues().begin() + b->getNumDimAndSymbolIds(),
                      [](Optional<Value> id) { return id.hasValue(); }));
 
-  // Bring A and B to common local space
-  a->mergeLocalIds(*b);
-
   SmallVector<Value, 4> aDimValues;
   a->getValues(offset, a->getNumDimIds(), &aDimValues);
 
@@ -543,6 +540,8 @@ static void mergeAndAlignIds(unsigned offset, FlatAffineValueConstraints *a,
 
   // Merge and align symbols of A and B
   a->mergeSymbolIds(*b);
+  // Merge and align local ids of A and B
+  a->mergeLocalIds(*b);
 
   assert(areIdsAligned(*a, *b) && "IDs expected to be aligned");
 }
@@ -1875,8 +1874,13 @@ void FlatAffineConstraints::removeRedundantConstraints() {
 
 /// Merge local ids of `this` and `other`. This is done by appending local ids
 /// of `other` to `this` and inserting local ids of `this` to `other` at start
-/// of its local ids.
+/// of its local ids. Number of dimension and symbol ids should match in
+/// `this` and `other`.
 void FlatAffineConstraints::mergeLocalIds(FlatAffineConstraints &other) {
+  assert(getNumDimIds() == other.getNumDimIds() &&
+         "Number of dimension ids should match");
+  assert(getNumSymbolIds() == other.getNumSymbolIds() &&
+         "Number of symbol ids should match");
   unsigned initLocals = getNumLocalIds();
   insertLocalId(getNumLocalIds(), other.getNumLocalIds());
   other.insertLocalId(0, initLocals);
@@ -3665,15 +3669,32 @@ void FlatAffineRelation::compose(const FlatAffineRelation &other) {
          "Domain of this and range of other do not match");
 
   FlatAffineRelation rel = other;
+
+  // Convert `rel` from
+  //    [otherDomain] -> [otherRange]
+  // to
+  //    [otherDomain] -> [otherRange thisRange]
+  // and `this` from
+  //    [thisDomain] -> [thisRange]
+  // to
+  //    [otherDomain thisDomain] -> [thisRange].
+  unsigned removeDims = rel.getNumRangeDims();
+  insertDomainId(0, rel.getNumDomainDims());
+  rel.appendRangeId(getNumRangeDims());
+
+  // Merge symbol and local identifiers.
   mergeSymbolIds(rel);
   mergeLocalIds(rel);
 
-  // Convert domain of `this` and range of `rel` to local identifiers.
-  convertDimToLocal(0, getNumDomainDims());
-  rel.convertDimToLocal(rel.getNumDomainDims(), rel.getNumDimIds());
-  // Add dimensions such that both relations become `domainRel -> rangeThis`.
-  appendDomainId(rel.getNumDomainDims());
-  rel.appendRangeId(getNumRangeDims());
+  // Convert `rel` from [otherDomain] -> [otherRange thisRange] to
+  // [otherDomain] -> [thisRange] by converting first otherRange range ids
+  // to local ids.
+  rel.convertDimToLocal(rel.getNumDomainDims(),
+                        rel.getNumDomainDims() + removeDims);
+  // Convert `this` from [otherDomain thisDomain] -> [thisRange] to
+  // [otherDomain] -> [thisRange] by converting last thisDomain domain ids
+  // to local ids.
+  convertDimToLocal(getNumDomainDims() - removeDims, getNumDomainDims());
 
   auto thisMaybeValues = getMaybeDimValues();
   auto relMaybeValues = rel.getMaybeDimValues();