From 504f331d77acbfb8d2abc5ca041814930aa02fc6 Mon Sep 17 00:00:00 2001
From: Jiahan Xie <88367305+jiahanxie353@users.noreply.github.com>
Date: Mon, 13 Jan 2025 06:26:20 -0500
Subject: [PATCH] [SCFToCalyx] Re-initialize IRmapping at the start of each
 loop iteration to map block arguments in region-based operations (#8075)

---
 lib/Conversion/SCFToCalyx/SCFToCalyx.cpp      |  6 ++-
 .../Conversion/SCFToCalyx/convert_simple.mlir | 51 +++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/lib/Conversion/SCFToCalyx/SCFToCalyx.cpp b/lib/Conversion/SCFToCalyx/SCFToCalyx.cpp
index d9779eff8aa2..46be615549d8 100644
--- a/lib/Conversion/SCFToCalyx/SCFToCalyx.cpp
+++ b/lib/Conversion/SCFToCalyx/SCFToCalyx.cpp
@@ -1894,7 +1894,6 @@ class BuildParGroups : public calyx::FuncOpPartialLoweringPattern {
     OpBuilder insideBuilder(newParOp);
     Block *currBlock = nullptr;
     auto &region = newParOp.getRegion();
-    IRMapping operandMap;
 
     // extract lower bounds, upper bounds, and steps as integer index values
     SmallVector<int64_t> lbVals, ubVals, stepVals;
@@ -1920,6 +1919,11 @@ class BuildParGroups : public calyx::FuncOpPartialLoweringPattern {
     SmallVector<int64_t> indices = lbVals;
 
     while (true) {
+      // Each iteration starts with a fresh mapping, so each new block’s
+      // argument of a region-based operation (such as `scf.for`) get re-mapped
+      // independently.
+      IRMapping operandMap;
+
       // Create a new block in the region for the current combination of indices
       currBlock = &region.emplaceBlock();
       insideBuilder.setInsertionPointToEnd(currBlock);
diff --git a/test/Conversion/SCFToCalyx/convert_simple.mlir b/test/Conversion/SCFToCalyx/convert_simple.mlir
index 9fe5c39826ee..b3531e778dfb 100644
--- a/test/Conversion/SCFToCalyx/convert_simple.mlir
+++ b/test/Conversion/SCFToCalyx/convert_simple.mlir
@@ -530,3 +530,54 @@ module {
     return %res : si32
   }
 }
+
+// Test parallel op lowering when it has region-based nested ops, such as `scf.for`
+
+// -----
+
+// CHECK:           calyx.control {
+// CHECK:             calyx.seq {
+// CHECK:               calyx.par {
+// CHECK:                 calyx.seq {
+// CHECK:                   calyx.enable @init_for_0_induction_var
+// CHECK:                   calyx.repeat 2 {
+// CHECK:                     calyx.seq {
+// CHECK:                       calyx.enable @bb0_0
+// CHECK:                       calyx.enable @bb0_1
+// CHECK:                       calyx.enable @incr_for_0_induction_var
+// CHECK:                     }
+// CHECK:                   }
+// CHECK:                 }
+// CHECK:                 calyx.seq {
+// CHECK:                   calyx.enable @init_for_1_induction_var
+// CHECK:                   calyx.repeat 2 {
+// CHECK:                     calyx.seq {
+// CHECK:                       calyx.enable @bb0_2
+// CHECK:                       calyx.enable @bb0_3
+// CHECK:                       calyx.enable @incr_for_1_induction_var
+// CHECK:                     }
+// CHECK:                   }
+// CHECK:                 }
+// CHECK:               }
+// CHECK:             }
+// CHECK:           }
+// CHECK:         } {toplevel}
+
+module {
+  func.func @main() {
+    %c2 = arith.constant 2 : index
+    %c1 = arith.constant 1 : index
+    %c0 = arith.constant 0 : index
+    %alloc = memref.alloc() : memref<6xi32>
+    %alloc_1 = memref.alloc() : memref<6xi32>
+    scf.parallel (%arg2) = (%c0) to (%c2) step (%c1) {
+      scf.for %arg3 = %c0 to %c2 step %c1 {
+        %1 = memref.load %alloc_1[%arg3] : memref<6xi32>
+        %2 = arith.shli %arg2, %c1 : index
+        memref.store %1, %alloc[%2] : memref<6xi32>
+      }
+      scf.reduce
+    }
+    return
+  }
+}