Fix issue with nested loop unrolling (#3100)

* Do not eliminate single-iter-loops that have inner loops using their break label.

* Add test

* Delete out-old.hlsl

* Update slang-ir-autodiff-cfg-norm.cpp

* Fix whitespace

2 files changed, 31 insertions, 2 deletions

diff --git a/source/slang/slang-ir-eliminate-multilevel-break.cpp b/source/slang/slang-ir-eliminate-multilevel-break.cpp
index 7db517309..8f7307f1b 100644
--- a/source/slang/slang-ir-eliminate-multilevel-break.cpp
+++ b/source/slang/slang-ir-eliminate-multilevel-break.cpp
@@ -59,6 +59,11 @@ struct EliminateMultiLevelBreakContext
         HashSet<IRBlock*> processedBlocks;
         List<MultiLevelBreakInfo> multiLevelBreaks;
 
+        bool isUnreachable(IRBlock* block)
+        {
+            return block->getPredecessors().getCount() == 0;
+        }
+
         void collectBreakableRegionBlocks(BreakableRegionInfo& info)
         {
             // Push break block to a stack so we can easily check if a block is a break block in its
@@ -92,7 +97,7 @@ struct EliminateMultiLevelBreakContext
                         collectBreakableRegionBlocks(*childRegion);
                         info.childRegions.add(childRegion);
                         block = childRegion->getBreakBlock();
-                        if (info.blockSet.add(block))
+                        if (!isUnreachable(block) && info.blockSet.add(block))
                         {
                             info.blocks.add(block);
                         }
@@ -142,7 +147,8 @@ struct EliminateMultiLevelBreakContext
                 l->forEach(
                     [&](BreakableRegionInfo* region)
                     {
-                        mapBreakBlockToRegion.add(region->getBreakBlock(), region);
+                        if(!isUnreachable(region->getBreakBlock()))
+                            mapBreakBlockToRegion.add(region->getBreakBlock(), region);
                         for (auto block : region->blocks)
                             mapBlockToRegion.add(block, region);
                     });
diff --git a/source/slang/slang-ir-simplify-cfg.cpp b/source/slang/slang-ir-simplify-cfg.cpp
index c4c5b584e..f2d0c4555 100644
--- a/source/slang/slang-ir-simplify-cfg.cpp
+++ b/source/slang/slang-ir-simplify-cfg.cpp
@@ -94,6 +94,29 @@ static bool isTrivialSingleIterationLoop(
             }
         }
     }
+    
+    // We'll also check if there's an inner loop that is breaking out into this loop's break block.
+    // If so, we cannot remove it right away since it interferes with the multi-level break elimination
+    // logic.
+    //
+    // Track the break block backwards through the dominator tree, and see if we find a loop block
+    // that is not the current loop.
+    //
+    auto currBlock = loop->getBreakBlock();
+    for (;;)
+    {
+        auto parent = context.domTree->getImmediateDominator(currBlock);
+        if (!parent)
+            break;
+        currBlock = parent;
+        if (auto _loop = as<IRLoop>(currBlock->getTerminator()))
+        {
+            if (loop != _loop)
+                return false;
+            if (loop == _loop)
+                break;
+        }
+    }
 
     return true;
 }