diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp --- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp @@ -564,11 +564,19 @@ // Calculate upperBound for normalized loop. SmallVector ubOperands; AffineBound lb = op.getLowerBound(); + SmallVector origLbOperands; + origLbOperands.append(lb.getOperands().begin(), lb.getOperands().end()); AffineBound ub = op.getUpperBound(); + SmallVector origUbOperands; + origUbOperands.append(ub.getOperands().begin(), ub.getOperands().end()); ubOperands.reserve(ub.getNumOperands() + lb.getNumOperands()); AffineMap origLbMap = lb.getMap(); AffineMap origUbMap = ub.getMap(); + // Make sure to canonicalize the original affine map beforehand. + canonicalizeMapAndOperands(&origLbMap, &origLbOperands); + canonicalizeMapAndOperands(&origUbMap, &origUbOperands); + // Add dimension operands from upper/lower bound. for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j) ubOperands.push_back(ub.getOperand(j)); @@ -612,7 +620,7 @@ SmallVector lbOperands(lb.getOperands().begin(), lb.getOperands().begin() + - lb.getMap().getNumDims()); + origLbMap.getNumDims()); // Normalize the loop. op.setUpperBound(ubOperands, newUbMap); @@ -626,9 +634,9 @@ 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)); + lbOperands.push_back(origLbOperands[origLbMap.getNumDims() + j]); - AffineExpr origIVExpr = opBuilder.getAffineDimExpr(lb.getMap().getNumDims()); + AffineExpr origIVExpr = opBuilder.getAffineDimExpr(origLbMap.getNumDims()); AffineExpr newIVExpr = origIVExpr * origLoopStep + origLbMap.getResult(0); AffineMap ivMap = AffineMap::get(origLbMap.getNumDims() + 1, origLbMap.getNumSymbols(), newIVExpr); diff --git a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir --- a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir +++ b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir @@ -213,3 +213,33 @@ } 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:.*]] = 0 to 1 { + affine.for %i = %c0 to %c1 { + // CHECK-NEXT: %[[IV:.*]] = affine.apply #map(%[[ARG0]]) + } + } + return +} + +// ----- + +// CHECK-LABEL: func @ensure_canonicalize_constant_bound +func.func @ensure_canonicalize_constant_bound() { + %c0 = arith.constant 0 : index + %c2 = arith.constant 2 : index + // CHECK: affine.for %[[ARG0:.*]] = 0 to 2 { + affine.for %i = %c0 to %c2 { + // CHECK-NEXT: %[[IV:.*]] = affine.apply #map(%[[ARG0]]) + "test.foo"(%i) : (index) -> () + } + return +}