From 186c3d6315f01669a46413df49d98366e33aa208 Mon Sep 17 00:00:00 2001 From: Akshay Baviskar Date: Thu, 6 Apr 2023 10:50:51 +0530 Subject: [PATCH] [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 --- mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 12 +++-- mlir/test/Dialect/SCF/loop-unroll.mlir | 80 ++++++++++++++++++++++++++++- 2 files changed, 88 insertions(+), 4 deletions(-) diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp index d826064..1513102 100644 --- 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 --- 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) { @@ -339,4 +340,81 @@ func.func @static_loop_unroll_by_5_with_cleanup(%arg0 : memref) { // CLEANUP-UNROLL-BY-5-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V1]]] : memref // CLEANUP-UNROLL-BY-5-NEXT: %[[V2:.*]] = affine.apply {{.*}} // CLEANUP-UNROLL-BY-5-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V2]]] : memref -// 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 -- 2.7.4