[Mlir-commits] [mlir] 3c22857 - Revert "[mlir][SCF] Further simplify affine maps during `for-loop-canonicalization`"
Alexander Belyaev
llvmlistbot at llvm.org
Thu Nov 25 01:55:05 PST 2021
Author: Alexander Belyaev
Date: 2021-11-25T10:54:52+01:00
New Revision: 3c228573bcb8833019ab443b6ef738514688c20a
URL: https://github.com/llvm/llvm-project/commit/3c228573bcb8833019ab443b6ef738514688c20a
DIFF: https://github.com/llvm/llvm-project/commit/3c228573bcb8833019ab443b6ef738514688c20a.diff
LOG: Revert "[mlir][SCF] Further simplify affine maps during `for-loop-canonicalization`"
This reverts commit ee1bf186723abb933b2c337e589c5958167f3cbe.
It breaks IREE lowering. Reverting the commit for now while we
investigate what's going on.
Added:
Modified:
mlir/include/mlir/Analysis/AffineStructures.h
mlir/lib/Analysis/AffineStructures.cpp
mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
mlir/test/Dialect/SCF/for-loop-canonicalization.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Analysis/AffineStructures.h b/mlir/include/mlir/Analysis/AffineStructures.h
index fa58e3e3412fa..53b3d043606e4 100644
--- a/mlir/include/mlir/Analysis/AffineStructures.h
+++ b/mlir/include/mlir/Analysis/AffineStructures.h
@@ -385,6 +385,7 @@ class FlatAffineConstraints {
/// Returns the constant bound for the pos^th identifier if there is one;
/// None otherwise.
+ // TODO: Support EQ bounds.
Optional<int64_t> getConstantBound(BoundType type, unsigned pos) const;
/// Gets the lower and upper bound of the `offset` + `pos`th identifier
diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp
index 21dd469b812f9..c0465a2532b46 100644
--- a/mlir/lib/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Analysis/AffineStructures.cpp
@@ -2836,22 +2836,11 @@ FlatAffineConstraints::computeConstantLowerOrUpperBound(unsigned pos) {
Optional<int64_t> FlatAffineConstraints::getConstantBound(BoundType type,
unsigned pos) const {
+ assert(type != BoundType::EQ && "EQ not implemented");
FlatAffineConstraints tmpCst(*this);
if (type == BoundType::LB)
- return FlatAffineConstraints(*this)
- .computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
- if (type == BoundType::UB)
- return FlatAffineConstraints(*this)
- .computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
-
- assert(type == BoundType::EQ && "expected EQ");
- Optional<int64_t> lb =
- FlatAffineConstraints(*this)
- .computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
- Optional<int64_t> ub =
- FlatAffineConstraints(*this)
- .computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
- return (lb && ub && *lb == *ub) ? Optional<int64_t>(*ub) : None;
+ return tmpCst.computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
+ return tmpCst.computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
}
// A simple (naive and conservative) check for hyper-rectangularity.
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index a71d9fe045f43..c54f634dcac04 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -305,16 +305,6 @@ canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap map,
AffineMap newMap = alignedBoundMap;
SmallVector<Value> newOperands;
unpackOptionalValues(constraints.getMaybeDimAndSymbolValues(), newOperands);
- // If dims/symbols have known constant values, use those in order to simplify
- // the affine map further.
- for (int64_t i = 0; i < constraints.getNumDimAndSymbolIds(); ++i) {
- // Skip unused operands and operands that are already constants.
- if (!newOperands[i] || getConstantIntValue(newOperands[i]))
- continue;
- if (auto bound = constraints.getConstantBound(FlatAffineConstraints::EQ, i))
- newOperands[i] =
- rewriter.create<arith::ConstantIndexOp>(op->getLoc(), *bound);
- }
mlir::canonicalizeMapAndOperands(&newMap, &newOperands);
rewriter.setInsertionPoint(op);
rewriter.replaceOpWithNewOp<AffineApplyOp>(op, newMap, newOperands);
@@ -467,30 +457,19 @@ mlir::scf::canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op,
if (ubInt)
constraints.addBound(FlatAffineConstraints::EQ, dimUb, *ubInt);
- // Lower bound: iv >= lb (equiv.: iv - lb >= 0)
+ // iv >= lb (equiv.: iv - lb >= 0)
SmallVector<int64_t> ineqLb(constraints.getNumCols(), 0);
ineqLb[dimIv] = 1;
ineqLb[dimLb] = -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;
- } 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);
- ivUb =
- exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
- }
+ // 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 ivUb =
+ exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
auto map = AffineMap::get(
/*dimCount=*/constraints.getNumDimIds(),
/*symbolCount=*/constraints.getNumSymbolIds(), /*result=*/ivUb);
diff --git a/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir b/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir
index add420b3b508f..a07e5688ead36 100644
--- a/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-canonicalization.mlir
@@ -348,22 +348,3 @@ func @tensor_dim_of_loop_result_no_canonicalize(%t : tensor<?x?xf32>,
%dim = tensor.dim %1, %c0 : tensor<?x?xf32>
return %dim : index
}
-
-// -----
-
-// CHECK-LABEL: func @one_trip_scf_for_canonicalize_min
-// CHECK: %[[C4:.*]] = arith.constant 4 : i64
-// CHECK: scf.for
-// CHECK: memref.store %[[C4]], %{{.*}}[] : memref<i64>
-func @one_trip_scf_for_canonicalize_min(%A : memref<i64>) {
- %c0 = arith.constant 0 : index
- %c2 = arith.constant 2 : index
- %c4 = arith.constant 4 : index
-
- 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
-}
More information about the Mlir-commits
mailing list