diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h --- a/mlir/include/mlir/IR/AffineMap.h +++ b/mlir/include/mlir/IR/AffineMap.h @@ -215,9 +215,13 @@ MLIRContext *context; }; -/// Simplify an affine map by simplifying its underlying AffineExpr results. +/// Simplifies an affine map by simplifying its underlying AffineExpr results. AffineMap simplifyAffineMap(AffineMap map); +/// Returns a map with the same dimension and symbol count as `map`, but whose +/// results are the unique affine expressions of `map`. +AffineMap removeDuplicateExprs(AffineMap map); + /// Returns a map of codomain to domain dimensions such that the first codomain /// dimension for a particular domain dimension is selected. /// Returns an empty map if the input map is empty or if `map` is not invertible diff --git a/mlir/include/mlir/Transforms/LoopUtils.h b/mlir/include/mlir/Transforms/LoopUtils.h --- a/mlir/include/mlir/Transforms/LoopUtils.h +++ b/mlir/include/mlir/Transforms/LoopUtils.h @@ -265,8 +265,8 @@ std::vector> &depthToLoops); /// Creates an AffineForOp while ensuring that the lower and upper bounds are -/// canonicalized, i.e., unused and duplicate operands are removed, and any -/// constant operands propagated/folded in. +/// canonicalized, i.e., unused and duplicate operands are removed, any constant +/// operands propagated/folded in, and duplicate bound maps dropped. AffineForOp createCanonicalizedAffineForOp(OpBuilder b, Location loc, ValueRange lbOperands, AffineMap lbMap, diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -1383,7 +1383,10 @@ auto prevUbMap = ubMap; canonicalizeMapAndOperands(&lbMap, &lbOperands); + lbMap = removeDuplicateExprs(lbMap); + canonicalizeMapAndOperands(&ubMap, &ubOperands); + ubMap = removeDuplicateExprs(ubMap); // Any canonicalization change always leads to updated map(s). if (lbMap == prevLbMap && ubMap == prevUbMap) diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp --- a/mlir/lib/IR/AffineMap.cpp +++ b/mlir/lib/IR/AffineMap.cpp @@ -325,6 +325,15 @@ return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs); } +AffineMap mlir::removeDuplicateExprs(AffineMap map) { + auto results = map.getResults(); + SmallVector uniqueExprs(results.begin(), results.end()); + uniqueExprs.erase(std::unique(uniqueExprs.begin(), uniqueExprs.end()), + uniqueExprs.end()); + return AffineMap::get(map.getNumDims(), map.getNumSymbols(), uniqueExprs, + map.getContext()); +} + AffineMap mlir::inversePermutation(AffineMap map) { if (map.isEmpty()) return map; diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp --- a/mlir/lib/Transforms/Utils/LoopUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp @@ -2132,8 +2132,10 @@ fullyComposeAffineMapAndOperands(&lbMap, &lowerOperands); canonicalizeMapAndOperands(&lbMap, &lowerOperands); + lbMap = removeDuplicateExprs(lbMap); fullyComposeAffineMapAndOperands(&ubMap, &upperOperands); canonicalizeMapAndOperands(&ubMap, &upperOperands); + ubMap = removeDuplicateExprs(ubMap); return b.create(loc, lowerOperands, lbMap, upperOperands, ubMap, step); diff --git a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir b/mlir/test/Conversion/AffineToStandard/lower-affine.mlir --- a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir +++ b/mlir/test/Conversion/AffineToStandard/lower-affine.mlir @@ -344,7 +344,7 @@ return } -#map_7_values = affine_map<(i) -> (i, i, i, i, i, i, i)> +#map_7_values = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5, d6)> // Check that the "min" (cmpi "slt" + select) reduction sequence is emitted // correctly for a an affine map with 7 results. @@ -369,8 +369,8 @@ // CHECK-NEXT: } // CHECK-NEXT: return // CHECK-NEXT: } -func @min_reduction_tree(%v : index) { - affine.for %i = 0 to min #map_7_values(%v)[] { +func @min_reduction_tree(%v1 : index, %v2 : index, %v3 : index, %v4 : index, %v5 : index, %v6 : index, %v7 : index) { + affine.for %i = 0 to min #map_7_values(%v1, %v2, %v3, %v4, %v5, %v6, %v7)[] { call @body(%i) : (index) -> () } return @@ -619,4 +619,4 @@ // CHECK: select %[[cmp]], %[[first]], %[[second]] %0 = affine.max affine_map<(d0,d1) -> (d0 - d1, d1 - d0)>(%arg0, %arg1) return %0 : index -} \ No newline at end of file +} diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir --- a/mlir/test/Dialect/Affine/canonicalize.mlir +++ b/mlir/test/Dialect/Affine/canonicalize.mlir @@ -594,3 +594,15 @@ %1 = affine.min #map2(%0)[%arg1] return %1 : index } + +// ----- +// CHECK-DAG: #[[lb:.*]] = affine_map<()[s0] -> (s0)> +// CHECK-DAG: #[[ub:.*]] = affine_map<()[s0] -> (s0 + 2)> + +func @drop_duplicate_bounds(%N : index) { + // affine.for %i = max #lb(%arg0) to min #ub(%arg0) + affine.for %i = max affine_map<(d0) -> (d0, d0)>(%N) to min affine_map<(d0) -> (d0 + 2, d0 + 2)>(%N) { + "foo"() : () -> () + } + return +}