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 @@ -1578,9 +1578,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 @@ -884,3 +884,23 @@ %1 = affine.max affine_map<()[s0, s1] -> (s0 + 4, 7 + s1)> ()[%0, %i2] return %1: index } + +// ----- + +// 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 +}