[MLIR][Affine] Fix generateUnrolledLoop utility

[akshaybaviskar]

Summary

generateUnrolledLoop was assuming that the yielded value is always generated in
the Block corresponding to the loop being unrolled. Thus, it was updating the
last yielded values with it's cloned value. However, if the yielded value is not
generated in the same Block then the cloned value and it's corresponding mapping
won't exist, resulting in a crash. Fix this.

[akshaybaviskar]

Fixes https://github.com/llvm/llvm-project/issues/59234

[bondhugula]

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

1044–1048

This would be wrong. You can't assume getDefiningOp always returns non-null. This will crash if the yielded value is a BlockArgument.

[bondhugula]

mlir/test/Dialect/SCF/loop-unroll.mlir

360

Missing test cases where the yield value is a block argument - one of the iter_args itself for example which can easily happen.

360–363

Please drop the %{{.*}} = part -- it's not doing any additional testing.

[akshaybaviskar]

Fix generateUnrolledLoop utility

generateUnrolledLoop was assuming that the yielded value is always generated in
the Block corresponding to the loop being unrolled. Thus, it was updating the
last yielded values with it's cloned value. However, if the yielded value is not
generated in the same Block then the cloned value and it's corresponding mapping
won't exist, resulting in a crash. Fix this.

[akshaybaviskar]

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

1044–1048

Thanks for pointing out. I have fixed it and updated the test cases. Please review.

[bondhugula]

Mostly related to code comments and redundancy.

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

1040

If the yielded value

1040

An SSA value can never be updated. Please rephrase.

1041

Nit: it's -> its

1041–1042

Again, please rewrite. An SSA value can't be updated.

1044

Please avoid != nullptr -- it's not consistent with the style used and is also more verbose. In addition, the more idiomatic way would be: !..isa<BlockArgument>(). You don't need to obtain the defining op pointer.

1044

The check also has redundancy/repetition between the two clauses. Instead:

Operation *defOp = ...;
if (defOp && defOp->getBlock() == ...)

[bondhugula]

Please add [MLIR][Affine] to the revision title.

[akshaybaviskar]

[MLIR][Affine] Fix generateUnrolledLoop utility

generateUnrolledLoop was assuming that the yielded value is always generated in
the Block corresponding to the loop being unrolled. Thus, it was updating the
last yielded values with it's cloned value. However, if the yielded value is not
generated in the same Block or if it is a BlockArgument then the cloned value
and it's corresponding mapping won't exist, resulting in a crash. Fix this.

Differential Revision: https://reviews.llvm.org/D146931

[akshaybaviskar]

Fix generateUnrolledLoop utility

generateUnrolledLoop was assuming that the yielded value is always generated in
the Block corresponding to the loop being unrolled. Thus, it was updating the
last yielded values with it's cloned value. However, if the yielded value is not
generated in the same Block then the cloned value and it's corresponding mapping
won't exist, resulting in a crash. Fix this.

Differential Revision: https://reviews.llvm.org/D146931

[akshaybaviskar]

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

1040

Update. Please re-review.

[bondhugula]

LGTM

[bondhugula]

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

1044–1048

This will eliminate an extra check:

Operation *defOp = ...;
if (defOp && defOp->getBlock() == ...)

[akshaybaviskar]

Fix generateUnrolledLoop utility

generateUnrolledLoop was assuming that the yielded value is always generated in
the Block corresponding to the loop being unrolled. Thus, it was updating the
last yielded values with it's cloned value. However, if the yielded value is not
generated in the same Block then the cloned value and it's corresponding mapping
won't exist, resulting in a crash. Fix this.

Differential Revision: https://reviews.llvm.org/D146931

[akshaybaviskar]

Fix generateUnrolledLoop utility

generateUnrolledLoop was assuming that the yielded value is always generated in
the Block corresponding to the loop being unrolled. Thus, it was updating the
last yielded values with it's cloned value. However, if the yielded value is not
generated in the same Block then the cloned value and it's corresponding mapping
won't exist, resulting in a crash. Fix this.

[bondhugula]

git diff --check HEAD~
mlir/test/Dialect/SCF/loop-unroll.mlir:8: trailing whitespace.
+// RUN: mlir-opt %s --affine-loop-unroll --split-input-file | FileCheck %s

Please fix.

[bondhugula]

mlir/test/Dialect/SCF/loop-unroll.mlir

8

Trailing white space here.

[akshaybaviskar]

Fix generateUnrolledLoop utility

generateUnrolledLoop was assuming that the yielded value is always generated in
the Block corresponding to the loop being unrolled. Thus, it was updating the
last yielded values with it's cloned value. However, if the yielded value is not
generated in the same Block then the cloned value and it's corresponding mapping
won't exist, resulting in a crash. Fix this.

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

for (unsigned i = 1; i < unrollFactor; i++) {
     }
 
     // Clone the original body of 'forOp'.
     for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++) {
       Operation *clonedOp = builder.clone(*it, operandMap);
       annotateFn(i, clonedOp, builder);
     }
 
-    // Update yielded values.
+    // Update yielded values. If yielded value is getting updated in this block
+    // then update it with it's new value else if the yielded value is a block
+    // argument or is invariant within the block, then it remains unchanged.
     for (unsigned i = 0, e = lastYielded.size(); i < e; i++)
+      if (yieldedValues[i].getDefiningOp() != nullptr &&
+          yieldedValues[i].getDefiningOp()->getBlock() == loopBodyBlock)
         lastYielded[i] = operandMap.lookup(yieldedValues[i]);
+      else
+        lastYielded[i] = yieldedValues[i];
   }
 
   // Make sure we annotate the Ops in the original body. We do this last so that
   // any annotations are not copied into the cloned Ops above.
   for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
     annotateFn(0, &*it, builder);
 
   // Update operands of the yield statement.

mlir/test/Dialect/SCF/loop-unroll.mlir

 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2' | FileCheck %s --check-prefix UNROLL-BY-2
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=3' | FileCheck %s --check-prefix UNROLL-BY-3
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=0' | FileCheck %s --check-prefix UNROLL-OUTER-BY-2
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=1' | FileCheck %s --check-prefix UNROLL-INNER-BY-2
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 annotate=true' | FileCheck %s --check-prefix UNROLL-BY-2-ANNOTATE
 // RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=6 unroll-up-to-factor=true' | FileCheck %s --check-prefix UNROLL-UP-TO
 // RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=5 cleanup-unroll=true' | FileCheck %s --check-prefix CLEANUP-UNROLL-BY-5
+// RUN: mlir-opt %s --affine-loop-unroll --split-input-file | FileCheck %s 
 
 func.func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
                           %arg3: memref<?xf32>) {
   %0 = arith.constant 7.0 : f32
   scf.for %i0 = %arg0 to %arg1 step %arg2 {
     memref.store %0, %arg3[%i0] : memref<?xf32>
   }
   return

 //   CLEANUP-UNROLL-BY-5-DAG:  %[[C0:.*]] = arith.constant 0 : index
 //   CLEANUP-UNROLL-BY-5-DAG:  %[[C3:.*]] = arith.constant 3 : index
 //   CLEANUP-UNROLL-BY-5-NEXT: %[[V0:.*]] = affine.apply {{.*}}
 //   CLEANUP-UNROLL-BY-5-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V0]]] : memref<?xf32>
 //   CLEANUP-UNROLL-BY-5-NEXT: %[[V1:.*]] = affine.apply {{.*}}
 //   CLEANUP-UNROLL-BY-5-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>
 //   CLEANUP-UNROLL-BY-5-NEXT: %[[V2:.*]] = affine.apply {{.*}}
 //   CLEANUP-UNROLL-BY-5-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V2]]] : memref<?xf32>
-//   CLEANUP-UNROLL-BY-5-NEXT: return
+//   CLEANUP-UNROLL-BY-5-NEXT: return
+
+// -----
+
+// Test loop unrolling when the yielded value remains unchanged.
+// CHECK: [[$MAP:#map]] = affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>
+// CHECK-LABEL: func @loop_unroll_static_yield_value
+func.func @loop_unroll_static_yield_value_test1() {
+  %true_4 = arith.constant true
+  %c1 = arith.constant 1 : index
+  %103 = affine.for %arg2 = 0 to 40 iter_args(%arg3 = %true_4) -> (i1) {
+    %324 = affine.max affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>(%c1)
+    affine.yield %true_4 : i1
+  }
+  return
+}
+// CHECK:         %[[TRUE:.*]] = arith.constant true
+// CHECK-NEXT:    %[[C1:.*]] = arith.constant 1 : index
+// CHECK-NEXT:    affine.for %{{.*}} = 0 to 40 step 4 iter_args(%{{.*}} = %[[TRUE]]) -> (i1) {
+// CHECK-NEXT:      affine.max [[$MAP]](%[[C1]])
+// CHECK-NEXT:      affine.max [[$MAP]](%[[C1]])
+// CHECK-NEXT:      affine.max [[$MAP]](%[[C1]])
+// CHECK-NEXT:      affine.max [[$MAP]](%[[C1]])
+// CHECK-NEXT:      affine.yield %[[TRUE]] : i1
+// CHECK-NEXT:    }
+// CHECK-NEXT:    return
+
+// -----
+
+// Loop unrolling when the yielded value is loop iv.
+// CHECK: [[$MAP0:#map[0-9]*]] = affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>
+// CHECK: [[$MAP1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 2)>
+// CHECK: [[$MAP2:#map[0-9]*]] = affine_map<(d0) -> (d0 + 4)>
+// CHECK: [[$MAP3:#map[0-9]*]] = affine_map<(d0) -> (d0 + 6)>
+// CHECK-LABEL: func @loop_unroll_yield_loop_iv
+func.func @loop_unroll_yield_loop_iv() {
+  %c1 = arith.constant 1 : index
+  %103 = affine.for %arg2 = 0 to 40 step 2 iter_args(%arg3 = %c1) -> (index) {
+    %324 = affine.max affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>(%arg2)
+    affine.yield %arg2 : index
+  }
+  return
+}
+// CHECK:         %[[C1:.*]] = arith.constant 1 : index
+// CHECK-NEXT:    affine.for %[[LOOP_IV:.*]] = 0 to 40 step 8 iter_args(%{{.*}} = %[[C1]]) -> (index) {
+// CHECK-NEXT:    affine.max [[$MAP0]](%[[LOOP_IV]])
+// CHECK-NEXT:    %[[LOOP_IV_PLUS_2:.*]] = affine.apply [[$MAP1]](%[[LOOP_IV]])
+// CHECK-NEXT:    affine.max [[$MAP0]](%[[LOOP_IV_PLUS_2]])
+// CHECK-NEXT:    %[[LOOP_IV_PLUS_4:.*]] = affine.apply [[$MAP2]](%[[LOOP_IV]])
+// CHECK-NEXT:    affine.max [[$MAP0]](%[[LOOP_IV_PLUS_4]])
+// CHECK-NEXT:    %[[LOOP_IV_PLUS_6:.*]] = affine.apply [[$MAP3]](%[[LOOP_IV]])
+// CHECK-NEXT:    affine.max [[$MAP0]](%[[LOOP_IV_PLUS_6]])
+// CHECK-NEXT:    affine.yield %[[LOOP_IV]] : index
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return
+
+// -----
+
+// Loop unrolling when the yielded value is iter_arg.
+// CHECK: [[$MAP:#map]] = affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>
+// CHECK-LABEL: func @loop_unroll_yield_iter_arg
+func.func @loop_unroll_yield_iter_arg() {
+  %c1 = arith.constant 1 : index
+  %103 = affine.for %arg2 = 0 to 40 step 2 iter_args(%arg3 = %c1) -> (index) {
+    %324 = affine.max affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>(%arg3)
+    affine.yield %arg3 : index
+  }
+  return
+}
+// CHECK:         %[[C1:.*]] = arith.constant 1 : index
+// CHECK-NEXT:    affine.for %{{.*}} = 0 to 40 step 8 iter_args(%[[ITER_ARG:.*]] = %[[C1]]) -> (index) {
+// CHECK-NEXT:      affine.max [[$MAP]](%[[ITER_ARG]])
+// CHECK-NEXT:      affine.max [[$MAP]](%[[ITER_ARG]])
+// CHECK-NEXT:      affine.max [[$MAP]](%[[ITER_ARG]])
+// CHECK-NEXT:      affine.max [[$MAP]](%[[ITER_ARG]])
+// CHECK-NEXT:      affine.yield %[[ITER_ARG]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:    return