Merge pull request #281 from coredac/fix/nested-loop-live-in-bug

Skip direct-dominating live-in optimization for loop headers

Author: guosran

lib/NeuraDialect/Transforms/CanonicalizeLiveInPass.cpp

@@ -436,6 +436,25 @@ identifyDirectDominatingLiveIns(Region &region, DominanceInfo &dom_info,
         continue;
       }
 
+      // If the using block is a loop header (has a back-edge), we must NOT treat
+      // any live-in as a direct dominating live-in. This is because:
+      // 1. Live-ins from outer scopes have rate mismatch and need PHI_START
+      // 2. Live-ins from inner blocks are loop-carried dependencies that need PHI
+      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/e2e/bicg/bicg_int_kernel.mlir

@@ -521,6 +521,8 @@ module attributes {torch.debug_module_name = "SimpleResNetBlock"} {
 // KERNEL-NEXT:   }
 // KERNEL-NEXT: }
 
+
+
 // STREAM:      module attributes {torch.debug_module_name = "SimpleResNetBlock"} {
 // STREAM-NEXT:   memref.global "private" constant @__constant_64xf32 : memref<64xf32> = dense<0.000000e+00> {alignment = 64 : i64}
 // STREAM-NEXT:   memref.global "private" constant @__constant_64x3x3x64xf32_0 : memref<64x3x3x64xf32> = dense<-0.0151730878> {alignment = 64 : i64}