Skip direct-dominating live-in optimization for loop headers

lib/NeuraDialect/Transforms/CanonicalizeLiveInPass.cpp

@@ -436,6 +436,28 @@ identifyDirectDominatingLiveIns(Region &region, DominanceInfo &dom_info,
         continue;
       }
 
+      // If the using block is a loop header (has a back-edge from itself or
+      // other blocks), we must NOT treat any live-in as a direct dominating
+      // live-in. In the dataflow model, loop headers operate at the inner
+      // loop rate, so live-in values from outer blocks must be promoted to
+      // block arguments to get proper rate-matched PHI_START operations
+      // during the ctrl-to-data-flow transformation.
+      // See: https://github.com/coredac/dataflow/issues/270
+      bool using_block_is_loop_header = false;
+      for (Block *pred : block.getPredecessors()) {
+        if (dom_info.dominates(&block, pred)) {
+          using_block_is_loop_header = true;
+          break;
+        }
+      }
+
+      if (using_block_is_loop_header) {
+        // Skips direct dominating live-in optimization for loop headers.
+        // The value must be promoted to a block argument so that the
+        // transform-ctrl-to-data-flow pass creates an inner-rate PHI_START.
+        continue;
+      }
+
       // Pattern 1: Single-Source-Single-Sink with one conditional branch.
       if (isSingleSourceSingleSinkPattern(defining_block, &block, dom_info,
                                           post_dom_info)) {

test/neura/fusion/test.mlir

@@ -29,9 +29,9 @@
 
 // CHECK-FUSED: func.func @_Z6kernelPA1024_iPiS1_S1_S1_
 // CHECK-FUSED: accelerator = "neura"
-// CHECK-FUSED-DAG: %91 = neura.load_indexed %89[%90 : !neura.data<i64, i1>] !neura.data<!llvm.ptr, i1> : !neura.data<i32, i1>
-// CHECK-FUSED-DAG: %82 = "neura.mul_add"(%79, %80, %81) : (!neura.data<i32, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CHECK-FUSED-DAG: %95 = "neura.mul_add"(%92, %93, %94) : (!neura.data<i32, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CHECK-FUSED-DAG: %102 = neura.load_indexed %100[%101 : !neura.data<i64, i1>] !neura.data<!llvm.ptr, i1> : !neura.data<i32, i1>
+// CHECK-FUSED-DAG: %93 = "neura.mul_add"(%90, %91, %92) : (!neura.data<i32, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CHECK-FUSED-DAG: %106 = "neura.mul_add"(%103, %104, %105) : (!neura.data<i32, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
 
 // CHECK-MAPPING: mapping_info = {compiled_ii = 12 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 9 : i32, res_mii = 5 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}