[mlir][SCF] Loop lb/ub are symbols during Affine Min/Max canonicalization

[springerm]

Summary

This fixes a bug in SCF/AffineCanonicalizationUtils.cpp. Loop lb/ub were previously considered dimensions, which caused a crash when a (non-optimizable) affine.min / affine.max expression was processed (due to multiplication of two dims). Lb/ub are now considered symbols and symbols may be multiplied. (The scope of the analysis is "within the loop body", at which point lb/ub are constants.)

Depends On D131949

[springerm]

fix bug IntegerRelation

[Groverkss]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

271–272 ↗

(On Diff #453265)

Thanks for fixing this bug! The behaviour should be to return the absolute position of the variable to match insertVar. Instead of doing this, could you change the doc of appendVar to say "Return the absolute column position (i.e., not relative to the kind of variable) of the first added variable"? Having different return behaviour for insertVar and appendVar could be confusing.

[springerm]

address comments

[nicolasvasilache]

Thanks !

mlir/include/mlir/Analysis/Presburger/IntegerRelation.h

public:
 
   /// Insert `num` variables of the specified kind at position `pos`.
   /// Positions are relative to the kind of variable. The coefficient columns
   /// corresponding to the added variables are initialized to zero. Return the
   /// absolute column position (i.e., not relative to the kind of variable)
   /// of the first added variable.
   virtual unsigned insertVar(VarKind kind, unsigned pos, unsigned num = 1);
 
-  /// Append `num` variables of the specified kind after the last variable.
-  /// of that kind. Return the position of the first appended column relative to
-  /// the kind of variable. The coefficient columns corresponding to the added
-  /// variables are initialized to zero.
+  /// Append `num` variables of the specified kind after the last variable
+  /// of that kind. The coefficient columns corresponding to the added variables
+  /// are initialized to zero. Return the absolute column position (i.e., not
+  /// relative to the kind of variable) of the first appended variable.
   unsigned appendVar(VarKind kind, unsigned num = 1);
 
   /// Adds an inequality (>= 0) from the coefficients specified in `inEq`.
   void addInequality(ArrayRef<int64_t> inEq);
   /// Adds an equality from the coefficients specified in `eq`.
   void addEquality(ArrayRef<int64_t> eq);
 
   /// Eliminate the `posB^th` local variable, replacing every instance of it

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

public:
   unsigned insertSymbolVar(unsigned pos, ValueRange vals);
   unsigned insertVar(presburger::VarKind kind, unsigned pos,
                      unsigned num = 1) override;
   unsigned insertVar(presburger::VarKind kind, unsigned pos, ValueRange vals);
 
   /// Append variables of the specified kind after the last variable of that
   /// kind. The coefficient columns corresponding to the added variables are
   /// initialized to zero. `vals` are the Values corresponding to the
-  /// variables. Return the position of the first added column.
+  /// variables. Return the absolute column position (i.e., not relative to the
+  /// kind of variable) of the first appended variable.
   ///
   /// Note: Empty Values are allowed in `vals`.
   unsigned appendDimVar(ValueRange vals);
   unsigned appendSymbolVar(ValueRange vals);
   unsigned appendDimVar(unsigned num = 1) {
     return appendVar(VarKind::SetDim, num);
   }
   unsigned appendSymbolVar(unsigned num = 1) {

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

void FlatAffineValueConstraints::reset(unsigned newNumDims,
                                        unsigned newNumLocals,
                                        ArrayRef<Value> valArgs) {
   reset(0, 0, newNumDims + newNumSymbols + newNumLocals + 1, newNumDims,
         newNumSymbols, newNumLocals, valArgs);
 }
 
 unsigned FlatAffineValueConstraints::appendDimVar(ValueRange vals) {
   unsigned pos = getNumDimVars();
-  insertVar(VarKind::SetDim, pos, vals);
-  return pos;
+  return insertVar(VarKind::SetDim, pos, vals);
 }
 
 unsigned FlatAffineValueConstraints::appendSymbolVar(ValueRange vals) {
   unsigned pos = getNumSymbolVars();
-  insertVar(VarKind::Symbol, pos, vals);
-  return pos;
+  return insertVar(VarKind::Symbol, pos, vals);
 }
 
 unsigned FlatAffineValueConstraints::insertDimVar(unsigned pos,
                                                   ValueRange vals) {
   return insertVar(VarKind::SetDim, pos, vals);
 }
 
 unsigned FlatAffineValueConstraints::insertSymbolVar(unsigned pos,

mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp

addLoopRangeConstraints(FlatAffineValueConstraints &constraints, Value iv,
   // IntegerPolyhedron does not support semi-affine expressions.
   // Therefore, only constant step values are supported.
   auto stepInt = getConstantIntValue(step);
   if (!stepInt)
     return failure();
 
   unsigned dimIv = constraints.appendDimVar(iv);
   auto lbv = lb.dyn_cast<Value>();
-  unsigned dimLb =
-      lbv ? constraints.appendDimVar(lbv) : constraints.appendDimVar(/*num=*/1);
+  unsigned symLb = lbv ? constraints.appendSymbolVar(lbv)
+                       : constraints.appendSymbolVar(/*num=*/1);
   auto ubv = ub.dyn_cast<Value>();
-  unsigned dimUb =
-      ubv ? constraints.appendDimVar(ubv) : constraints.appendDimVar(/*num=*/1);
+  unsigned symUb = ubv ? constraints.appendSymbolVar(ubv)
+                       : constraints.appendSymbolVar(/*num=*/1);
 
   // If loop lower/upper bounds are constant: Add EQ constraint.
   Optional<int64_t> lbInt = getConstantIntValue(lb);
   Optional<int64_t> ubInt = getConstantIntValue(ub);
   if (lbInt)
-    constraints.addBound(IntegerPolyhedron::EQ, dimLb, *lbInt);
+    constraints.addBound(IntegerPolyhedron::EQ, symLb, *lbInt);
   if (ubInt)
-    constraints.addBound(IntegerPolyhedron::EQ, dimUb, *ubInt);
+    constraints.addBound(IntegerPolyhedron::EQ, symUb, *ubInt);
 
   // Lower bound: iv >= lb (equiv.: iv - lb >= 0)
   SmallVector<int64_t> ineqLb(constraints.getNumCols(), 0);
   ineqLb[dimIv] = 1;
-  ineqLb[dimLb] = -1;
+  ineqLb[symLb] = -1;
   constraints.addInequality(ineqLb);
 
   // Upper bound
   AffineExpr ivUb;
   if (lbInt && ubInt && (*lbInt + *stepInt >= *ubInt)) {
     // The loop has at most one iteration.
     // iv < lb + 1
     // TODO: Try to derive this constraint by simplifying the expression in
     // the else-branch.
-    ivUb = rewriter.getAffineDimExpr(dimLb) + 1;
+    ivUb =
+        rewriter.getAffineSymbolExpr(symLb - constraints.getNumDimVars()) + 1;
   } else {
     // The loop may have more than one iteration.
     // iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
-    AffineExpr exprLb = lbInt ? rewriter.getAffineConstantExpr(*lbInt)
-                              : rewriter.getAffineDimExpr(dimLb);
-    AffineExpr exprUb = ubInt ? rewriter.getAffineConstantExpr(*ubInt)
-                              : rewriter.getAffineDimExpr(dimUb);
+    AffineExpr exprLb =
+        lbInt
+            ? rewriter.getAffineConstantExpr(*lbInt)
+            : rewriter.getAffineSymbolExpr(symLb - constraints.getNumDimVars());
+    AffineExpr exprUb =
+        ubInt
+            ? rewriter.getAffineConstantExpr(*ubInt)
+            : rewriter.getAffineSymbolExpr(symUb - constraints.getNumDimVars());
     ivUb = exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
   }
   auto map = AffineMap::get(
       /*dimCount=*/constraints.getNumDimVars(),
       /*symbolCount=*/constraints.getNumSymbolVars(), /*result=*/ivUb);
 
   return constraints.addBound(IntegerPolyhedron::UB, dimIv, map);
 }

mlir/test/Dialect/SCF/for-loop-canonicalization.mlir

func.func @one_trip_scf_for_canonicalize_min(%A : memref<i64>) {
 
   scf.for %i = %c0 to %c4 step %c4 {
     %1 = affine.min affine_map<(d0, d1)[] -> (4, d1 - d0)> (%i, %c4)
     %2 = arith.index_cast %1: index to i64
     memref.store %2, %A[]: memref<i64>
   }
   return
 }
+
+// -----
+
+// This is a regression test to ensure that the no assertions are failing.
+
+//       CHECK: #[[$map:.+]] = affine_map<(d0)[s0] -> (-(d0 * (5 ceildiv s0)) + 5, 3)>
+// CHECK-LABEL: func @regression_multiplication_with_sym
+func.func @regression_multiplication_with_sym(%A : memref<i64>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+  // CHECK: %[[dummy:.*]] = "test.dummy"
+  %ub = "test.dummy"() : () -> (index)
+  // CHECK: scf.for %[[iv:.*]] =
+  scf.for %i = %c0 to %ub step %c1 {
+    // CHECK: affine.min #[[$map]](%[[iv]])[%[[dummy]]]
+    %1 = affine.min affine_map<(d0)[s0] -> (-(d0 * (5 ceildiv s0)) + 5, 3)>(%i)[%ub]
+    %2 = arith.index_cast %1: index to i64
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}