Fix subtle corner-case with vars getting hoisted out of the loop creating unnecessary loop state (#3165)

* Extend the unit tests for MxLayeredMaterial

* Add breaking loop test

* Fix subtle corner-case with vars getting hoisted out of the loop creating unnecessary loop state

* remove whitespace changes

* Create loop-init.slang.expected.txt

* Add filecheck tests to ensure correct loop state

* Update comment

---------

Co-authored-by: winmad <winmad.wlf@gmail.com>
Co-authored-by: Yong He <yonghe@outlook.com>

3 files changed, 54 insertions, 11 deletions

diff --git a/source/slang/slang-ir-autodiff-cfg-norm.cpp b/source/slang/slang-ir-autodiff-cfg-norm.cpp
index 8a88e69ad..a9db3aecc 100644
--- a/source/slang/slang-ir-autodiff-cfg-norm.cpp
+++ b/source/slang/slang-ir-autodiff-cfg-norm.cpp
@@ -166,7 +166,17 @@ struct CFGNormalizationPass
 
             if (as<IRVar>(child))
             {
-                child->insertBefore(region->headerBlock->getTerminator());
+                if (auto loopInst = as<IRLoop>(region->headerBlock->getTerminator()))
+                {   
+                    // In order to avoid introducing unnecessary loop state, we'll move vars
+                    // to the loop's target (first loop block) instead of the loop header.
+                    // (unless the var is already in the header or target)
+                    // 
+                    if (block != region->headerBlock && block != loopInst->getTargetBlock())
+                        child->insertBefore(loopInst->getTargetBlock()->getTerminator());
+                }
+                else
+                    child->insertBefore(region->headerBlock->getTerminator());
             }
 
             child = nextChild;
@@ -701,7 +711,19 @@ static void legalizeDefUse(IRGlobalValueWithCode* func)
                 {
                     if (loopUser->getTargetBlock() == commonDominator)
                     {
-                        commonDominator = as<IRBlock>(loopUser->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());
                         break;
                     }
                 }
@@ -751,7 +773,7 @@ void normalizeCFG(
     // Remove phis to simplify our pass. We'll add them back in later
     // with constructSSA.
     //
-    eliminatePhisInFunc(LivenessMode::Disabled, func->getModule(), func);
+    eliminatePhisInFunc(LivenessMode::Disabled, func->getModule(), func, false);
 
     CFGNormalizationContext context = {module, options.sink};
     CFGNormalizationPass cfgPass(context);
diff --git a/source/slang/slang-ir-eliminate-phis.cpp b/source/slang/slang-ir-eliminate-phis.cpp
index a17759fe6..1023e6148 100644
--- a/source/slang/slang-ir-eliminate-phis.cpp
+++ b/source/slang/slang-ir-eliminate-phis.cpp
@@ -68,11 +68,20 @@ struct PhiEliminationContext
     IRModule* m_module = nullptr;
     IRBuilder m_builder;
     LivenessMode m_livenessMode;
+    bool m_useRegisterAllocation;
 
     PhiEliminationContext(LivenessMode livenessMode, IRModule* module)
         : m_module(module)
         , m_builder(module)
         , m_livenessMode(livenessMode)
+        , m_useRegisterAllocation(true)
+    {}
+
+     PhiEliminationContext(LivenessMode livenessMode, IRModule* module, bool useRegisterAllocation)
+        : m_module(module)
+        , m_builder(module)
+        , m_livenessMode(livenessMode)
+        , m_useRegisterAllocation(useRegisterAllocation)
     {}
 
     // We start with the top-down logic of the pass, which is to process
@@ -210,8 +219,12 @@ struct PhiEliminationContext
     {
         m_func = func;
         m_dominatorTree = nullptr;
-        m_registerAllocation = allocateRegistersForFunc(func, m_dominatorTree);
-        m_mapRegToTempVar = createTempVarForInsts(func);
+
+        if (m_useRegisterAllocation)
+        {
+            m_registerAllocation = allocateRegistersForFunc(func, m_dominatorTree);
+            m_mapRegToTempVar = createTempVarForInsts(func);
+        }
     }
 
     Dictionary<RegisterInfo*, IRInst*> createTempVarForInsts(IRGlobalValueWithCode* func)
@@ -1109,15 +1122,19 @@ struct PhiEliminationContext
     }
 };
 
-void eliminatePhis(LivenessMode livenessMode, IRModule* module)
+void eliminatePhis(LivenessMode livenessMode, IRModule* module, bool useRegisterAllocation)
 {
-    PhiEliminationContext context(livenessMode, module);
+    PhiEliminationContext context(livenessMode, module, useRegisterAllocation);
     context.eliminatePhisInModule();
 }
 
-void eliminatePhisInFunc(LivenessMode livenessMode, IRModule* module, IRGlobalValueWithCode* func)
+void eliminatePhisInFunc(
+    LivenessMode livenessMode,
+    IRModule* module,
+    IRGlobalValueWithCode* func,
+    bool useRegisterAllocation)
 {
-    PhiEliminationContext context(livenessMode, module);
+    PhiEliminationContext context(livenessMode, module, useRegisterAllocation);
     context.eliminatePhisInFunc(func);
 }
 
diff --git a/source/slang/slang-ir-eliminate-phis.h b/source/slang/slang-ir-eliminate-phis.h
index ff81d5b38..9bfbe51b8 100644
--- a/source/slang/slang-ir-eliminate-phis.h
+++ b/source/slang/slang-ir-eliminate-phis.h
@@ -15,7 +15,11 @@ namespace Slang
         /// are not themselves based on an SSA representation.
         ///
         /// If livenessMode is enabled LiveRangeStarts will be inserted into the module.
-    void eliminatePhis(LivenessMode livenessMode, IRModule* module);
+    void eliminatePhis(LivenessMode livenessMode, IRModule* module, bool useRegisterAllocation = true);
 
-    void eliminatePhisInFunc(LivenessMode livenessMode, IRModule* module, IRGlobalValueWithCode* func);
+    void eliminatePhisInFunc(
+        LivenessMode livenessMode,
+        IRModule* module,
+        IRGlobalValueWithCode* func,
+        bool useRegisterAllocation = true);
 }