[Mlir-commits] [mlir] adc93e0 - [mlir][SCF] Loop lb/ub are symbols during Affine Min/Max canonicalization
Matthias Springer
llvmlistbot at llvm.org
Thu Aug 18 02:51:12 PDT 2022
Author: Matthias Springer
Date: 2022-08-18T11:44:48+02:00
New Revision: adc93e0d38b39b04c91bf7cb60275ea7708cf452
URL: https://github.com/llvm/llvm-project/commit/adc93e0d38b39b04c91bf7cb60275ea7708cf452
DIFF: https://github.com/llvm/llvm-project/commit/adc93e0d38b39b04c91bf7cb60275ea7708cf452.diff
LOG: [mlir][SCF] Loop lb/ub are symbols during Affine Min/Max canonicalization
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.)
Differential Revision: https://reviews.llvm.org/D132021
Added:
Modified:
mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp
mlir/test/Dialect/SCF/for-loop-canonicalization.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index b6922d9795f4c..320eb15361765 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -238,10 +238,10 @@ class IntegerRelation {
/// 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
index 4cbd721c993c4..3c7561ef40d72 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
@@ -299,7 +299,8 @@ class FlatAffineValueConstraints : public presburger::IntegerPolyhedron {
/// 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
index 012338b148a3e..8d7d63eead323 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
@@ -271,14 +271,12 @@ void FlatAffineValueConstraints::reset(unsigned newNumDims,
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
index f832c544d14b6..c2d004089faae 100644
--- a/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp
@@ -211,24 +211,24 @@ addLoopRangeConstraints(FlatAffineValueConstraints &constraints, Value iv,
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
@@ -238,14 +238,19 @@ addLoopRangeConstraints(FlatAffineValueConstraints &constraints, Value iv,
// 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
index 0497429152550..4638a217abbb5 100644
--- a/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir
@@ -368,3 +368,26 @@ func.func @one_trip_scf_for_canonicalize_min(%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
+}
More information about the Mlir-commits
mailing list