Fix def-use issue from multi-level break elimination (#6134)

4 files changed, 112 insertions, 107 deletions

diff --git a/source/slang/slang-ir-autodiff-cfg-norm.cpp b/source/slang/slang-ir-autodiff-cfg-norm.cpp
index 08a6c84e7..30e832719 100644
--- a/source/slang/slang-ir-autodiff-cfg-norm.cpp
+++ b/source/slang/slang-ir-autodiff-cfg-norm.cpp
@@ -707,113 +707,6 @@ struct CFGNormalizationPass
     }
 };
 
-static void legalizeDefUse(IRGlobalValueWithCode* func)
-{
-    auto dom = computeDominatorTree(func);
-    for (auto block : func->getBlocks())
-    {
-        for (auto inst : block->getModifiableChildren())
-        {
-            // Inspect all uses of `inst` and find the common dominator of all use sites.
-            IRBlock* commonDominator = block;
-            for (auto use = inst->firstUse; use; use = use->nextUse)
-            {
-                auto userBlock = as<IRBlock>(use->getUser()->getParent());
-                if (!userBlock)
-                    continue;
-                while (commonDominator && !dom->dominates(commonDominator, userBlock))
-                {
-                    commonDominator = dom->getImmediateDominator(commonDominator);
-                }
-            }
-            SLANG_ASSERT(commonDominator);
-
-            if (commonDominator == block)
-                continue;
-
-            // If the common dominator is not `block`, it means we have detected
-            // uses that is no longer dominated by the current definition, and need
-            // to be fixed.
-
-            // Normally, we can simply move the definition to the common dominator.
-            // An exception is when the common dominator is the target block of a
-            // loop. Note that after normalization, loops are in the form of:
-            // ```
-            // loop { if (condition) block; else break; }
-            // ```
-            // If we find ourselves needing to make the inst available right before
-            // the `if`, it means we are seeing uses of the inst outside the loop.
-            // In this case, we should insert a var/move the inst before the loop
-            // instead of before the `if`. This situation can occur in the IR if
-            // the original code is lowered from a `do-while` loop.
-            for (auto use = commonDominator->firstUse; use; use = use->nextUse)
-            {
-                if (auto loopUser = as<IRLoop>(use->getUser()))
-                {
-                    if (loopUser->getTargetBlock() == commonDominator)
-                    {
-                        bool shouldMoveToHeader = false;
-                        // Check that the break-block dominates any of the uses are past the break
-                        // block
-                        for (auto _use = inst->firstUse; _use; _use = _use->nextUse)
-                        {
-                            if (dom->dominates(
-                                    loopUser->getBreakBlock(),
-                                    _use->getUser()->getParent()))
-                            {
-                                shouldMoveToHeader = true;
-                                break;
-                            }
-                        }
-
-                        if (shouldMoveToHeader)
-                            commonDominator = as<IRBlock>(loopUser->getParent());
-                        break;
-                    }
-                }
-            }
-            // Now we can legalize uses based on the type of `inst`.
-            if (auto var = as<IRVar>(inst))
-            {
-                // If inst is an var, this is easy, we just move it to the
-                // common dominator.
-                var->insertBefore(commonDominator->getTerminator());
-            }
-            else
-            {
-                // For all other insts, we need to create a local var for it,
-                // and replace all uses with a load from the local var.
-                IRBuilder builder(func);
-                builder.setInsertBefore(commonDominator->getTerminator());
-                IRVar* tempVar = builder.emitVar(inst->getFullType());
-                auto defaultVal = builder.emitDefaultConstruct(inst->getFullType());
-                builder.emitStore(tempVar, defaultVal);
-
-                builder.setInsertAfter(inst);
-                builder.emitStore(tempVar, inst);
-
-                traverseUses(
-                    inst,
-                    [&](IRUse* use)
-                    {
-                        auto userBlock = as<IRBlock>(use->getUser()->getParent());
-                        if (!userBlock)
-                            return;
-                        // Only fix the use of the current definition of `inst` does not
-                        // dominate it.
-                        if (!dom->dominates(block, userBlock))
-                        {
-                            // Replace the use with a load of tempVar.
-                            builder.setInsertBefore(use->getUser());
-                            auto load = builder.emitLoad(tempVar);
-                            builder.replaceOperand(use, load);
-                        }
-                    });
-            }
-        }
-    }
-}
-
 void normalizeCFG(
     IRModule* module,
     IRGlobalValueWithCode* func,
diff --git a/source/slang/slang-ir-eliminate-multilevel-break.cpp b/source/slang/slang-ir-eliminate-multilevel-break.cpp
index a54d1b46d..741b7fbd0 100644
--- a/source/slang/slang-ir-eliminate-multilevel-break.cpp
+++ b/source/slang/slang-ir-eliminate-multilevel-break.cpp
@@ -5,6 +5,7 @@
 #include "slang-ir-dominators.h"
 #include "slang-ir-eliminate-phis.h"
 #include "slang-ir-insts.h"
+#include "slang-ir-util.h"
 #include "slang-ir.h"
 
 namespace Slang
@@ -475,6 +476,8 @@ struct EliminateMultiLevelBreakContext
                 }
             }
         }
+
+        legalizeDefUse(func);
     }
 };
 
diff --git a/source/slang/slang-ir-util.cpp b/source/slang/slang-ir-util.cpp
index d05e1db7d..3043e7e31 100644
--- a/source/slang/slang-ir-util.cpp
+++ b/source/slang/slang-ir-util.cpp
@@ -1986,4 +1986,111 @@ Int getSpecializationConstantId(IRGlobalParam* param)
     return offset->getOffset();
 }
 
+void legalizeDefUse(IRGlobalValueWithCode* func)
+{
+    auto dom = computeDominatorTree(func);
+    for (auto block : func->getBlocks())
+    {
+        for (auto inst : block->getModifiableChildren())
+        {
+            // Inspect all uses of `inst` and find the common dominator of all use sites.
+            IRBlock* commonDominator = block;
+            for (auto use = inst->firstUse; use; use = use->nextUse)
+            {
+                auto userBlock = as<IRBlock>(use->getUser()->getParent());
+                if (!userBlock)
+                    continue;
+                while (commonDominator && !dom->dominates(commonDominator, userBlock))
+                {
+                    commonDominator = dom->getImmediateDominator(commonDominator);
+                }
+            }
+            SLANG_ASSERT(commonDominator);
+
+            if (commonDominator == block)
+                continue;
+
+            // If the common dominator is not `block`, it means we have detected
+            // uses that is no longer dominated by the current definition, and need
+            // to be fixed.
+
+            // Normally, we can simply move the definition to the common dominator.
+            // An exception is when the common dominator is the target block of a
+            // loop. Note that after normalization, loops are in the form of:
+            // ```
+            // loop { if (condition) block; else break; }
+            // ```
+            // If we find ourselves needing to make the inst available right before
+            // the `if`, it means we are seeing uses of the inst outside the loop.
+            // In this case, we should insert a var/move the inst before the loop
+            // instead of before the `if`. This situation can occur in the IR if
+            // the original code is lowered from a `do-while` loop.
+            for (auto use = commonDominator->firstUse; use; use = use->nextUse)
+            {
+                if (auto loopUser = as<IRLoop>(use->getUser()))
+                {
+                    if (loopUser->getTargetBlock() == commonDominator)
+                    {
+                        bool shouldMoveToHeader = false;
+                        // Check that the break-block dominates any of the uses are past the break
+                        // block
+                        for (auto _use = inst->firstUse; _use; _use = _use->nextUse)
+                        {
+                            if (dom->dominates(
+                                    loopUser->getBreakBlock(),
+                                    _use->getUser()->getParent()))
+                            {
+                                shouldMoveToHeader = true;
+                                break;
+                            }
+                        }
+
+                        if (shouldMoveToHeader)
+                            commonDominator = as<IRBlock>(loopUser->getParent());
+                        break;
+                    }
+                }
+            }
+            // Now we can legalize uses based on the type of `inst`.
+            if (auto var = as<IRVar>(inst))
+            {
+                // If inst is an var, this is easy, we just move it to the
+                // common dominator.
+                var->insertBefore(commonDominator->getTerminator());
+            }
+            else
+            {
+                // For all other insts, we need to create a local var for it,
+                // and replace all uses with a load from the local var.
+                IRBuilder builder(func);
+                builder.setInsertBefore(commonDominator->getTerminator());
+                IRVar* tempVar = builder.emitVar(inst->getFullType());
+                auto defaultVal = builder.emitDefaultConstruct(inst->getFullType());
+                builder.emitStore(tempVar, defaultVal);
+
+                builder.setInsertAfter(inst);
+                builder.emitStore(tempVar, inst);
+
+                traverseUses(
+                    inst,
+                    [&](IRUse* use)
+                    {
+                        auto userBlock = as<IRBlock>(use->getUser()->getParent());
+                        if (!userBlock)
+                            return;
+                        // Only fix the use of the current definition of `inst` does not
+                        // dominate it.
+                        if (!dom->dominates(block, userBlock))
+                        {
+                            // Replace the use with a load of tempVar.
+                            builder.setInsertBefore(use->getUser());
+                            auto load = builder.emitLoad(tempVar);
+                            builder.replaceOperand(use, load);
+                        }
+                    });
+            }
+        }
+    }
+}
+
 } // namespace Slang
diff --git a/source/slang/slang-ir-util.h b/source/slang/slang-ir-util.h
index 666ac71c0..46e0105f5 100644
--- a/source/slang/slang-ir-util.h
+++ b/source/slang/slang-ir-util.h
@@ -375,6 +375,8 @@ IRType* getIRVectorBaseType(IRType* type);
 
 Int getSpecializationConstantId(IRGlobalParam* param);
 
+void legalizeDefUse(IRGlobalValueWithCode* func);
+
 } // namespace Slang
 
 #endif