[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 then the cloned value and it's corresponding mapping
won't exist, resulting in a crash. Fix this.

Reviewed By: bondhugula

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

diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index d826064..1513102 100644 (file)
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -954,9 +954,15 @@ static void generateUnrolledLoop(
       annotateFn(i, clonedOp, builder);
     }
 
-    // Update yielded values.
-    for (unsigned i = 0, e = lastYielded.size(); i < e; i++)
-      lastYielded[i] = operandMap.lookup(yieldedValues[i]);
+    // 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++) {
+      Operation *defOp = yieldedValues[i].getDefiningOp();
+      if (defOp && defOp->getBlock() == loopBodyBlock)
+        lastYielded[i] = operandMap.lookup(yieldedValues[i]);
+    }
   }
 
   // Make sure we annotate the Ops in the original body. We do this last so that

diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir
index b2928f5..c83e33d 100644 (file)
--- a/mlir/test/Dialect/SCF/loop-unroll.mlir
+++ b/mlir/test/Dialect/SCF/loop-unroll.mlir
@@ -5,6 +5,7 @@
 // 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>) {
@@ -339,4 +340,81 @@ func.func @static_loop_unroll_by_5_with_cleanup(%arg0 : memref<?xf32>) {
 //   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
\ No newline at end of file
+//   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