[mlir][affine] Fix unfolded bounding maps for affine.for

[eopXD]

Summary

Loop bounds of affine.for didn't perform foldings like affine.load, affine.store.
Bound maps shall be more composed, leaving most affine.apply become dead.

This resolves the bug listed on https://bugs.llvm.org/show_bug.cgi?id=45203

[eopXD]

This comment was removed by eopXD.

[eopXD]

revert to 1st commit of this bug fix.

[eopXD]

[mlir][affine][test] Check if affine.for has bounding maps composed

[bondhugula]

Thanks for improving this. A few comments.

mlir/test/Dialect/Affine/canonicalize.mlir

700

Nit: Terminate comment with period.
is -> are

703

No need of checking for module, but please add CHECK-LABEL's for each function.

704–709

Avoid numbered SSA values in test cases. If you don't want to capture them.

%arg1, %arg0 -> %{{.*}}

713

Please update this test case to cover lower bound composition as well.

907

New line on last line.

[eopXD]

Adjust testcase to include lower bound composition,
also adjusted testcase to avoid numbered SSA.

[bondhugula]

In D99323#2652344, @eopXD wrote:

Adjust testcase to include lower bound composition,
also adjusted testcase to avoid numbered SSA.

Please feel free to mark comments as "Done" if they are done.

[eopXD]

Yes the comments are all done.
Thank you for the comments.

[bondhugula]

Loop bounds of affine.load

Typo in commit summary.

[bondhugula]

Looks good.

[bondhugula]

@eopXD Are you able to commit this? If not, I can commit it for you.

[eopXD]

@bondhugula I am not able to commit this.
Please help me to commit, thank you.

[bondhugula]

In D99323#2670025, @eopXD wrote:

@bondhugula I am not able to commit this.
Please help me to commit, thank you.

Can you please rebase your commit on main and update this revision?

[eopXD]

Hi @bondhugula,
I am trying to rebase with $ git rebase origin/main with the current HEAD on my patching commit.
Then I encountered merge conflicts of file:

CONFLICT (content): Merge conflict in libunwind/CMakeLists.txt
CONFLICT (content): Merge conflict in libcxx/CMakeLists.txt

Is this normal or how shall I do something different to rebase to main?

[bondhugula]

In D99323#2674614, @eopXD wrote:

Hi @bondhugula,
I am trying to rebase with $ git rebase origin/main with the current HEAD on my patching commit.
Then I encountered merge conflicts of file:

CONFLICT (content): Merge conflict in libunwind/CMakeLists.txt
CONFLICT (content): Merge conflict in libcxx/CMakeLists.txt

Is this normal or how shall I do something different to rebase to main?

A rebase leading to conflicts like these in unrelated files is indication that some developer/some process really messed up the LLVM git history upstream (@mehdi_amini @rriddle may have more insights here). In this case, what you'll have to do is to just cherry-pick your commit on the upstream main tip instead of rebasing. That will avoid all these conflicts completely unrelated to your PR. You will still see some conflicts on the affine ops files you are touching because you used a pretty old base - from at least a month ago. (Everyone using a significantly old base that traverses the time point where the history got messed up would see such conflcits while rebasing I believe.)

[eopXD]

Cherry-pick commit to upstream main.

[eopXD]

[mlir][affine] Fix unfolded bounding maps for affine.for

Loop bounds of affine.load didn’t perform foldings like affine.load, affine.store.
Bound maps shall be more composed, leaving most affine.apply become dead.

This resolves the bug listed on https://bugs.llvm.org/show_bug.cgi?id=45203
Differential Revision: https://reviews.llvm.org/D99323

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
   SmallVector<Value, 4> lbOperands(forOp.getLowerBoundOperands());
   SmallVector<Value, 4> ubOperands(forOp.getUpperBoundOperands());
 
   auto lbMap = forOp.getLowerBoundMap();
   auto ubMap = forOp.getUpperBoundMap();
   auto prevLbMap = lbMap;
   auto prevUbMap = ubMap;
 
+  composeAffineMapAndOperands(&lbMap, &lbOperands);
   canonicalizeMapAndOperands(&lbMap, &lbOperands);
   lbMap = removeDuplicateExprs(lbMap);
 
+  composeAffineMapAndOperands(&ubMap, &ubOperands);
   canonicalizeMapAndOperands(&ubMap, &ubOperands);
   ubMap = removeDuplicateExprs(ubMap);
 
   // Any canonicalization change always leads to updated map(s).
   if (lbMap == prevLbMap && ubMap == prevUbMap)
     return failure();
 
   if (lbMap != prevLbMap)

mlir/test/Dialect/Affine/canonicalize.mlir

affine.for %i2 = 0 to 3 {
     %5 = index_cast %4: index to i64
     memref.store %5, %A[]: memref<i64>
   }
   return
 }
 
 // -----
 
 // CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1, s0 * s1)>
 
 // CHECK: func @deduplicate_affine_min_expressions
 // CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
 func @deduplicate_affine_min_expressions(%i0: index, %i1: index) -> index {
   // CHECK:  affine.min #[[MAP]]()[%[[I0]], %[[I1]]]
   %0 = affine.min affine_map<()[s0, s1] -> (s0 + s1, s0 * s1, s1 + s0, s0 * s1)> ()[%i0, %i1]
   return %0: index
 }
 
 // -----
 
 // CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1, s0 * s1)>
 
 // CHECK: func @deduplicate_affine_max_expressions
 // CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
 func @deduplicate_affine_max_expressions(%i0: index, %i1: index) -> index {
   // CHECK:  affine.max #[[MAP]]()[%[[I0]], %[[I1]]]
   %0 = affine.max affine_map<()[s0, s1] -> (s0 + s1, s0 * s1, s1 + s0, s0 * s1)> ()[%i0, %i1]
   return %0: index
 }
 

 
 // CHECK-LABEL: func @dont_merge_affine_max_if_not_single_sym
 func @dont_merge_affine_max_if_not_single_sym(%i0: index, %i1: index, %i2: index) -> index {
   // CHECK-COUNT-2: affine.max
   %0 = affine.max affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%i0]
   %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
+}