[Mlir-commits] [mlir] 18d10fb - [mlir][affine] addLowerOrUpperBound: Make map+operand composing optional

[Matthias Springer]

Author: Matthias Springer
Date: 2021-08-03T11:31:00+09:00
New Revision: 18d10fbe87b36cd922faeeb04d18078aea071c95

URL: https://github.com/llvm/llvm-project/commit/18d10fbe87b36cd922faeeb04d18078aea071c95
DIFF: https://github.com/llvm/llvm-project/commit/18d10fbe87b36cd922faeeb04d18078aea071c95.diff

LOG: [mlir][affine] addLowerOrUpperBound: Make map+operand composing optional

There are cases in which it is not desirable to fully compose the bound map with the operands when adding lower/upper bounds to a `FlatAffineConstraints`.

E.g., this is the case when bounds should be expressed in terms of the operands only (and not the operands' dependencies). This also makes `addLowerOrUpperBound` useable together with operands that are defined through semi-affine expressions.

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Analysis/AffineStructures.h b/mlir/include/mlir/Analysis/AffineStructures.h
index 1d857d59c34ff..302e5494eb959 100644
--- a/mlir/include/mlir/Analysis/AffineStructures.h
+++ b/mlir/include/mlir/Analysis/AffineStructures.h
@@ -237,10 +237,13 @@ class FlatAffineConstraints {
   /// Adds a lower or an upper bound for the identifier at the specified
   /// position with constraints being drawn from the specified bound map and
   /// operands. If `eq` is true, add a single equality equal to the bound map's
-  /// first result expr.
+  /// first result expr. By default, the bound map is fully composed with the
+  /// operands before adding any bounds. This allows for bounds being expressed
+  /// in terms of values that are used by the operands.
   LogicalResult addLowerOrUpperBound(unsigned pos, AffineMap boundMap,
                                      ValueRange operands, bool eq,
-                                     bool lower = true);
+                                     bool lower = true,
+                                     bool composeMapAndOperands = true);
 
   /// Returns the bound for the identifier at `pos` from the inequality at
   /// `ineqPos` as a 1-d affine value map (affine map + operands). The returned
@@ -334,7 +337,9 @@ class FlatAffineConstraints {
   /// symbols or loop IVs. The identifier is added to the end of the existing
   /// dims or symbols. Additional information on the identifier is extracted
   /// from the IR and added to the constraint system.
-  void addInductionVarOrTerminalSymbol(Value id);
+  /// Note: If `allowNonTerminal`, any symbol (incl. potentially non-terminal
+  /// ones) is allowed.
+  void addInductionVarOrTerminalSymbol(Value id, bool allowNonTerminal = false);
 
   /// Composes the affine value map with this FlatAffineConstrains, adding the
   /// results of the map as dimensions at the front [0, vMap->getNumResults())

diff  --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp
index 598bf30d5cd8d..82bf861ab063c 100644
--- a/mlir/lib/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Analysis/AffineStructures.cpp
@@ -560,12 +560,13 @@ void FlatAffineConstraints::convertLoopIVSymbolsToDims() {
   }
 }
 
-void FlatAffineConstraints::addInductionVarOrTerminalSymbol(Value id) {
+void FlatAffineConstraints::addInductionVarOrTerminalSymbol(
+    Value id, bool allowNonTerminal) {
   if (containsId(id))
     return;
 
   // Caller is expected to fully compose map/operands if necessary.
-  assert((isTopLevelValue(id) || isForInductionVar(id)) &&
+  assert((allowNonTerminal || isTopLevelValue(id) || isForInductionVar(id)) &&
          "non-terminal symbol / loop IV expected");
   // Outer loop IVs could be used in forOp's bounds.
   if (auto loop = getForInductionVarOwner(id)) {
@@ -1944,10 +1945,9 @@ void FlatAffineConstraints::getSliceBounds(unsigned offset, unsigned num,
   }
 }
 
-LogicalResult
-FlatAffineConstraints::addLowerOrUpperBound(unsigned pos, AffineMap boundMap,
-                                            ValueRange boundOperands, bool eq,
-                                            bool lower) {
+LogicalResult FlatAffineConstraints::addLowerOrUpperBound(
+    unsigned pos, AffineMap boundMap, ValueRange boundOperands, bool eq,
+    bool lower, bool composeMapAndOperands) {
   assert(pos < getNumDimAndSymbolIds() && "invalid position");
   // Equality follows the logic of lower bound except that we add an equality
   // instead of an inequality.
@@ -1959,11 +1959,13 @@ FlatAffineConstraints::addLowerOrUpperBound(unsigned pos, AffineMap boundMap,
   // transitively get to terminal symbols or loop IVs.
   auto map = boundMap;
   SmallVector<Value, 4> operands(boundOperands.begin(), boundOperands.end());
-  fullyComposeAffineMapAndOperands(&map, &operands);
+  if (composeMapAndOperands)
+    fullyComposeAffineMapAndOperands(&map, &operands);
   map = simplifyAffineMap(map);
   canonicalizeMapAndOperands(&map, &operands);
   for (auto operand : operands)
-    addInductionVarOrTerminalSymbol(operand);
+    addInductionVarOrTerminalSymbol(
+        operand, /*allowNonTerminal=*/!composeMapAndOperands);
 
   FlatAffineConstraints localVarCst;
   std::vector<SmallVector<int64_t, 8>> flatExprs;