[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

1045–1047

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

1045–1047

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–1043

If the yielded value

1040–1043

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.

1045–1047

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.

1045–1047

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–1043

Update. Please re-review.

[bondhugula]

LGTM

[bondhugula]

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

1045

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 the yielded value is defined outside the
+    // `loopBodyBlock` or if it is a BlockArgument then it won't be cloned, thus
+    // the `lastYielded` value remains unchanged. Else, update the `lastYielded`
+    // value with the clone corresponding to the yielded value.
     for (unsigned i = 0, e = lastYielded.size(); i < e; i++)
+      if (!yieldedValues[i].isa<BlockArgument>() &&
+          yieldedValues[i].getDefiningOp()->getBlock() == loopBodyBlock)
         lastYielded[i] = operandMap.lookup(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