diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h --- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h +++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h @@ -238,10 +238,10 @@ /// 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`. diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h --- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h @@ -299,7 +299,8 @@ /// 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); diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp --- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp @@ -271,14 +271,12 @@ 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, diff --git a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp --- a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp +++ b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp @@ -211,24 +211,24 @@ unsigned dimIv = constraints.appendDimVar(iv); auto lbv = lb.dyn_cast(); - 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(); - 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 lbInt = getConstantIntValue(lb); Optional 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 ineqLb(constraints.getNumCols(), 0); ineqLb[dimIv] = 1; - ineqLb[dimLb] = -1; + ineqLb[symLb] = -1; constraints.addInequality(ineqLb); // Upper bound @@ -238,14 +238,19 @@ // 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( diff --git a/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir b/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir --- a/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir +++ b/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir @@ -368,3 +368,26 @@ } 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) { + %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 + } + return +}