[Mlir-commits] [mlir] d056c75 - [mlir][scf] Fix unrolling when the yielded value is defined above the loop. (#122177)

[llvmlistbot at llvm.org]

Author: Alexander Belyaev
Date: 2025-01-09T17:31:17+01:00
New Revision: d056c756aea3fc709cf7d6bc8acabe9a8c7218db

URL: https://github.com/llvm/llvm-project/commit/d056c756aea3fc709cf7d6bc8acabe9a8c7218db
DIFF: https://github.com/llvm/llvm-project/commit/d056c756aea3fc709cf7d6bc8acabe9a8c7218db.diff

LOG: [mlir][scf] Fix unrolling when the yielded value is defined above the loop. (#122177)

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/Utils/Utils.cpp
    mlir/test/Dialect/SCF/loop-unroll.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 6cda7100fe073d..fa82bcb816a2a2 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 68a11fb6a72c64..0368505a1b70dc 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
+