1 files changed, 68 insertions, 26 deletions

diff --git a/source/slang/slang-ir-util.cpp b/source/slang/slang-ir-util.cpp
index bf5b25d9c..39c1c5bb1 100644
--- a/source/slang/slang-ir-util.cpp
+++ b/source/slang/slang-ir-util.cpp
@@ -2078,9 +2078,36 @@ Int getSpecializationConstantId(IRGlobalParam* param)
     return offset->getOffset();
 }
 
+IRBlock* getLoopHeaderForConditionBlock(IRBlock* block)
+{
+    // Go through uses and check if any of them are a loop condition block.
+    for (auto use = block->firstUse; use; use = use->nextUse)
+    {
+        if (auto loop = as<IRLoop>(use->getUser()))
+        {
+            if (loop->getTargetBlock() == block)
+                return cast<IRBlock>(loop->getParent());
+        }
+    }
+    return nullptr;
+}
+
 void legalizeDefUse(IRGlobalValueWithCode* func)
 {
     auto dom = computeDominatorTree(func);
+
+    // Make a map of loop condition blocks to their loop header.
+    // We need this because we'll be treating loop condition blocks as
+    // special cases (they are the special blocks since they "dominate" themselves,
+    // in the dominator tree sense)
+    //
+    Dictionary<IRBlock*, IRBlock*> loopHeaderBlockMap;
+    for (auto block : func->getBlocks())
+    {
+        if (auto header = getLoopHeaderForConditionBlock(block))
+            loopHeaderBlockMap.add(block, header);
+    }
+
     for (auto block : func->getBlocks())
     {
         for (auto inst : block->getModifiableChildren())
@@ -2099,16 +2126,22 @@ void legalizeDefUse(IRGlobalValueWithCode* func)
             }
             SLANG_ASSERT(commonDominator);
 
-            if (commonDominator == block)
+            // If commonDominator is 'block' and if the inst is not a Var in
+            // a loop condition block, we can skip the legalization.
+            //
+            if (commonDominator == block &&
+                !(as<IRVar>(inst) && loopHeaderBlockMap.containsKey(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.
+            // Normally, if the common dominator is not `block`, 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.
+            // Another exception is when a var in the loop condition block is accessed both inside
+            // and outside the loop. It is technically visible, but effects on the 'var' are not
+            // visible outside the loop, so we'll need to hoist it out of the loop.
+            //
+            // Note that after normalization, loops are in the form of:
             // ```
             // loop { if (condition) block; else break; }
             // ```
@@ -2117,38 +2150,47 @@ void legalizeDefUse(IRGlobalValueWithCode* func)
             // 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)
+            //
+            bool shouldInitializeVar = false;
+            if (loopHeaderBlockMap.containsKey(commonDominator))
             {
-                if (auto loopUser = as<IRLoop>(use->getUser()))
+                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 (loopUser->getTargetBlock() == commonDominator)
+                    if (dom->dominates(
+                            as<IRLoop>(loopHeaderBlockMap[commonDominator]->getTerminator())
+                                ->getBreakBlock(),
+                            _use->getUser()->getParent()))
                     {
-                        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());
+                        shouldMoveToHeader = true;
                         break;
                     }
                 }
+                if (shouldMoveToHeader)
+                {
+                    commonDominator = loopHeaderBlockMap[commonDominator];
+                    shouldInitializeVar = true;
+                }
             }
+
             // 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());
+                if (shouldInitializeVar)
+                {
+                    IRBuilder builder(func);
+                    builder.setInsertAfter(var);
+                    builder.emitStore(
+                        var,
+                        builder.emitDefaultConstruct(
+                            as<IRPtrTypeBase>(var->getDataType())->getValueType()));
+                }
             }
             else
             {