diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp index 6cda7100fe073..fa82bcb816a2a 100644 --- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp +++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp @@ -361,7 +361,7 @@ static void generateUnrolledLoop( // Update yielded values. for (unsigned i = 0, e = lastYielded.size(); i < e; i++) - lastYielded[i] = operandMap.lookup(yieldedValues[i]); + lastYielded[i] = operandMap.lookupOrDefault(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 68a11fb6a72c6..0368505a1b70d 100644 --- a/mlir/test/Dialect/SCF/loop-unroll.mlir +++ b/mlir/test/Dialect/SCF/loop-unroll.mlir @@ -489,3 +489,32 @@ func.func @static_loop_unroll_with_integer_iv() -> (f32, f32) { // UNROLL-BY-3-NEXT: scf.yield %[[EADD]], %[[EMUL]] : f32, f32 // UNROLL-BY-3-NEXT: } // UNROLL-BY-3-NEXT: return %[[EFOR]]#0, %[[EFOR]]#1 : f32, f32 + +// ----- + +// Test loop unrolling when the yielded value is defined above the loop. +func.func @loop_unroll_static_yield_value_defined_above(%init: i32) { + %c42 = arith.constant 42 : i32 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c4 = arith.constant 4 : index + %103:2 = scf.for %i = %c0 to %c4 step %c1 + iter_args(%iter1 = %c42, %iter2 = %init) -> (i32, i32) { + %0 = arith.andi %iter2, %iter1 : i32 + scf.yield %0, %init : i32, i32 + } + return +} +// UNROLL-OUTER-BY-2-LABEL: @loop_unroll_static_yield_value_defined_above( +// UNROLL-OUTER-BY-2-SAME: %[[INIT:.*]]: i32) { +// UNROLL-OUTER-BY-2-DAG: %[[C42:.*]] = arith.constant 42 : i32 +// UNROLL-OUTER-BY-2-DAG: %[[C0:.*]] = arith.constant 0 : index +// UNROLL-OUTER-BY-2-DAG: %[[C4:.*]] = arith.constant 4 : index +// UNROLL-OUTER-BY-2-DAG: %[[C2:.*]] = arith.constant 2 : index +// UNROLL-OUTER-BY-2: scf.for %{{.*}} = %[[C0]] to %[[C4]] step %[[C2]] +// UNROLL-OUTER-BY-2-SAME: iter_args(%[[ITER1:.*]] = %[[C42]], +// UNROLL-OUTER-BY-2-SAME: %[[ITER2:.*]] = %[[INIT]]) +// UNROLL-OUTER-BY-2: %[[SUM:.*]] = arith.andi %[[ITER2]], %[[ITER1]] +// UNROLL-OUTER-BY-2: %[[SUM1:.*]] = arith.andi %[[INIT]], %[[SUM]] +// UNROLL-OUTER-BY-2: scf.yield %[[SUM1]], %[[INIT]] : i32, i32 +