[Mlir-commits] [mlir] 48f04ff - [mlir][affine] Normalize constant valued bound loop
Kai Sasaki
llvmlistbot at llvm.org
Mon Feb 13 19:27:34 PST 2023
Author: Kai Sasaki
Date: 2023-02-14T12:23:32+09:00
New Revision: 48f04ffa9f3611a60371c968498c5d3326466f9f
URL: https://github.com/llvm/llvm-project/commit/48f04ffa9f3611a60371c968498c5d3326466f9f
DIFF: https://github.com/llvm/llvm-project/commit/48f04ffa9f3611a60371c968498c5d3326466f9f.diff
LOG: [mlir][affine] Normalize constant valued bound loop
This change aims to resolve the issue reported in https://github.com/llvm/llvm-project/issues/59994.
After calling AffineForOp#setUpperBound and setLowerBound, it makes the original affine map structure inconsistent. It's necessary to keep the original map to build the new induction vars.
Reviewed By: bondhugula
Differential Revision: https://reviews.llvm.org/D142082
Added:
Modified:
mlir/lib/Dialect/Affine/Utils/Utils.cpp
mlir/test/Dialect/Affine/affine-loop-normalize.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index d731241a62661..5830e8cef99a0 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -564,31 +564,37 @@ LogicalResult mlir::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) {
OpBuilder opBuilder(op);
int64_t origLoopStep = op.getStep();
- // Calculate upperBound for normalized loop.
- SmallVector<Value, 4> ubOperands;
AffineBound lb = op.getLowerBound();
+ AffineMap originalLbMap = lb.getMap();
+ SmallVector<Value, 4> origLbOperands;
+ llvm::append_range(origLbOperands, lb.getOperands());
+
AffineBound ub = op.getUpperBound();
+ AffineMap originalUbMap = ub.getMap();
+ SmallVector<Value, 4> origUbOperands;
+ llvm::append_range(origUbOperands, ub.getOperands());
+
+ // Calculate upperBound for normalized loop.
+ SmallVector<Value, 4> ubOperands;
ubOperands.reserve(ub.getNumOperands() + lb.getNumOperands());
- AffineMap origLbMap = lb.getMap();
- AffineMap origUbMap = ub.getMap();
// Add dimension operands from upper/lower bound.
- for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j)
+ for (unsigned j = 0, e = originalUbMap.getNumDims(); j < e; ++j)
ubOperands.push_back(ub.getOperand(j));
- for (unsigned j = 0, e = origLbMap.getNumDims(); j < e; ++j)
+ for (unsigned j = 0, e = originalLbMap.getNumDims(); j < e; ++j)
ubOperands.push_back(lb.getOperand(j));
// Add symbol operands from upper/lower bound.
- for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j)
- ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j));
- for (unsigned j = 0, e = origLbMap.getNumSymbols(); j < e; ++j)
- ubOperands.push_back(lb.getOperand(origLbMap.getNumDims() + j));
+ for (unsigned j = 0, e = originalUbMap.getNumSymbols(); j < e; ++j)
+ ubOperands.push_back(ub.getOperand(originalUbMap.getNumDims() + j));
+ for (unsigned j = 0, e = originalLbMap.getNumSymbols(); j < e; ++j)
+ ubOperands.push_back(lb.getOperand(originalLbMap.getNumDims() + j));
// Add original result expressions from lower/upper bound map.
- SmallVector<AffineExpr, 1> origLbExprs(origLbMap.getResults().begin(),
- origLbMap.getResults().end());
- SmallVector<AffineExpr, 2> origUbExprs(origUbMap.getResults().begin(),
- origUbMap.getResults().end());
+ SmallVector<AffineExpr, 1> origLbExprs(originalLbMap.getResults().begin(),
+ originalLbMap.getResults().end());
+ SmallVector<AffineExpr, 2> origUbExprs(originalUbMap.getResults().begin(),
+ originalUbMap.getResults().end());
SmallVector<AffineExpr, 4> newUbExprs;
// The original upperBound can have more than one result. For the new
@@ -607,15 +613,15 @@ LogicalResult mlir::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) {
}
// Construct newUbMap.
- AffineMap newUbMap =
- AffineMap::get(origLbMap.getNumDims() + origUbMap.getNumDims(),
- origLbMap.getNumSymbols() + origUbMap.getNumSymbols(),
- newUbExprs, opBuilder.getContext());
+ AffineMap newUbMap = AffineMap::get(
+ originalLbMap.getNumDims() + originalUbMap.getNumDims(),
+ originalLbMap.getNumSymbols() + originalUbMap.getNumSymbols(), newUbExprs,
+ opBuilder.getContext());
canonicalizeMapAndOperands(&newUbMap, &ubOperands);
SmallVector<Value, 4> lbOperands(lb.getOperands().begin(),
lb.getOperands().begin() +
- lb.getMap().getNumDims());
+ originalLbMap.getNumDims());
// Normalize the loop.
op.setUpperBound(ubOperands, newUbMap);
@@ -628,13 +634,14 @@ LogicalResult mlir::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) {
// Add an extra dim operand for loopIV.
lbOperands.push_back(op.getInductionVar());
// Add symbol operands from lower bound.
- for (unsigned j = 0, e = origLbMap.getNumSymbols(); j < e; ++j)
- lbOperands.push_back(lb.getOperand(origLbMap.getNumDims() + j));
-
- AffineExpr origIVExpr = opBuilder.getAffineDimExpr(lb.getMap().getNumDims());
- AffineExpr newIVExpr = origIVExpr * origLoopStep + origLbMap.getResult(0);
- AffineMap ivMap = AffineMap::get(origLbMap.getNumDims() + 1,
- origLbMap.getNumSymbols(), newIVExpr);
+ for (unsigned j = 0, e = originalLbMap.getNumSymbols(); j < e; ++j)
+ lbOperands.push_back(origLbOperands[originalLbMap.getNumDims() + j]);
+
+ AffineExpr origIVExpr =
+ opBuilder.getAffineDimExpr(originalLbMap.getNumDims());
+ AffineExpr newIVExpr = origIVExpr * origLoopStep + originalLbMap.getResult(0);
+ AffineMap ivMap = AffineMap::get(originalLbMap.getNumDims() + 1,
+ originalLbMap.getNumSymbols(), newIVExpr);
canonicalizeMapAndOperands(&ivMap, &lbOperands);
Operation *newIV = opBuilder.create<AffineApplyOp>(loc, ivMap, lbOperands);
op.getInductionVar().replaceAllUsesExcept(newIV->getResult(0), newIV);
diff --git a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
index 35ab7994cdf2f..12dc5c0ed5d40 100644
--- a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
@@ -213,3 +213,82 @@ func.func @tiled_matmul(%0: memref<1024x1024xf32>, %1: memref<1024x1024xf32>, %2
}
return
}
+
+// -----
+
+// CHECK-LABEL: func @constant_lower_bound
+func.func @constant_lower_bound() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.for %j = %c0 to %c1 step %c1 {
+ // CHECK: affine.for %[[ARG0:.*]] =
+ affine.for %i = %c0 to %c1 {
+ // CHECK-NEXT: %{{.*}} = affine.apply #map{{.*}}(%[[ARG0]])
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> (s0 ceildiv 4)>
+// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0) -> (d0 * 4)>
+
+// CHECK-LABEL: func @upper_bound_by_symbol
+func.func @upper_bound_by_symbol(%arg0: index, %arg1: index) {
+ // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] {
+ affine.for %i = 0 to affine_map<()[s0, s1] -> (s0)>()[%arg0, %arg1] step 4 {
+ // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])
+ // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> ()
+ "test.foo"(%i) : (index) -> ()
+ }
+ return
+}
+
+// -----
+
+// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> ((-s0 + 10) ceildiv 4)>
+// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 * 4 + s0)>
+
+// CHECK-LABEL: func @lower_bound_by_symbol
+func.func @lower_bound_by_symbol(%arg0: index, %arg1: index) {
+ // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] {
+ affine.for %i = affine_map<()[s0, s1] -> (s0)>()[%arg0, %arg1] to 10 step 4 {
+ // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])[%arg{{.*}}]
+ // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> ()
+ "test.foo"(%i) : (index) -> ()
+ }
+ return
+}
+
+// -----
+
+// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> (s0 ceildiv 4)>
+// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0) -> (d0 * 4)>
+
+// CHECK-LABEL: func @upper_bound_by_dim
+func.func @upper_bound_by_dim(%arg0: index, %arg1: index) {
+ // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] {
+ affine.for %i = 0 to affine_map<(d0, d1) -> (d0)>(%arg0, %arg1) step 4 {
+ // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])
+ // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> ()
+ "test.foo"(%i) : (index) -> ()
+ }
+ return
+}
+
+// -----
+
+// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> ((-s0 + 10) ceildiv 4)>
+// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 * 4 + s0)>
+
+// CHECK-LABEL: func @upper_bound_by_dim
+func.func @upper_bound_by_dim(%arg0: index, %arg1: index) {
+ // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] {
+ affine.for %i = affine_map<(d0, d1) -> (d0)>(%arg0, %arg1) to 10 step 4 {
+ // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])[%arg{{.*}}]
+ // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> ()
+ "test.foo"(%i) : (index) -> ()
+ }
+ return
+}
More information about the Mlir-commits
mailing list