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 @@ -1576,9 +1576,11 @@ auto prevLbMap = lbMap; auto prevUbMap = ubMap; + composeAffineMapAndOperands(&lbMap, &lbOperands); canonicalizeMapAndOperands(&lbMap, &lbOperands); lbMap = removeDuplicateExprs(lbMap); + composeAffineMapAndOperands(&ubMap, &ubOperands); canonicalizeMapAndOperands(&ubMap, &ubOperands); ubMap = removeDuplicateExprs(ubMap); 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 @@ -694,3 +694,23 @@ } return } + +// ----- + +// Ensure bounding maps of affine.for are composed. + +// CHECK-DAG: #[[$MAP0]] = affine_map<()[s0] -> (s0 - 2)> +// CHECK-DAG: #[[$MAP1]] = affine_map<()[s0] -> (s0 + 2)> + +// CHECK-LABEL: func @compose_affine_for_bounds +// CHECK-SAME: %[[N:.*]]: index) +// CHECK: affine.for %{{.*}} = #[[$MAP0]]()[%[[N]]] to #[[$MAP1]]()[%[[N]]] { + +func @compose_affine_for_bounds(%N: index) { + %u = affine.apply affine_map<(d0) -> (d0 + 2)>(%N) + %l = affine.apply affine_map<(d0) -> (d0 - 2)>(%N) + affine.for %i = %l to %u { + "foo"() : () -> () + } + return +}