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>

13 files changed, 264 insertions, 16 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);
 }
diff --git a/tests/autodiff/loop-init.slang b/tests/autodiff/loop-init.slang
new file mode 100644
index 000000000..26d837e75
--- /dev/null
+++ b/tests/autodiff/loop-init.slang
@@ -0,0 +1,64 @@
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+//TEST:SIMPLE(filecheck=CHECK): -target hlsl -profile cs_5_0 -entry computeMain -line-directive-mode none
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+struct A : IDifferentiable
+{
+    float data[5];
+};
+
+// Check that the intermediate context of B.eval does not have any arrays.
+// This will fail if the induction variable is not properly detected, or
+// if the various loop restructuring passes accidentally introduce additional 
+// loop state.
+// 
+
+// CHECK: struct s_bwd_B_eval_Intermediates_0
+// CHECK-NOT: int {{[A-Za-z0-9_]+}}[{{.*}}]
+// CHECK: }
+
+__generic<let TBsdfCount : int>
+struct B
+{
+    [Differentiable]
+    float3 eval(const A miData, const float3 wi, const float3 wo)
+    {
+        float3 albedo;
+        for (uint i = 0; i < 3; i++) albedo[i] = miData.data[i];
+
+        float3 result = float3(1.f);
+        [ForceUnroll] for (uint i = 0; i < TBsdfCount; i++) result *= albedo;
+        return result;
+    }
+};
+
+[Differentiable]
+float3 outerEval(const A miData, const float3 wi, const float3 wo)
+{
+    B<3> b;
+    return b.eval(miData, wi, wo);
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    float3 wi = float3(2.0, 3.0, 0);
+    float3 wo = float3(1.0, 1.0, 0);
+    float data[5] = { 1, 2, 3, 4, 5 };
+    A dataStruct = { data };
+
+    float3 val = outerEval(dataStruct, wi, wo);
+    outputBuffer[0] = val.x;
+
+    DifferentialPair<float3> dpwi = diffPair(wi);
+    DifferentialPair<float3> dpwo = diffPair(wo);
+    DifferentialPair<A> dpdata = diffPair(dataStruct);
+    float3 dOut = float3(1.0, 0.0, 0.0);
+    __bwd_diff(outerEval)(dpdata, dpwi, dpwo, dOut);
+
+    // Write output
+    outputBuffer[0] = dpdata.d.data[0];
+}
\ No newline at end of file
diff --git a/tests/autodiff/loop-init.slang.expected.txt b/tests/autodiff/loop-init.slang.expected.txt
new file mode 100644
index 000000000..857cebc03
--- /dev/null
+++ b/tests/autodiff/loop-init.slang.expected.txt
@@ -0,0 +1,5 @@
+type: float
+3.000000
+0.000000
+0.000000
+0.000000
diff --git a/tests/autodiff/material2/IBSDF.slang b/tests/autodiff/material2/IBSDF.slang
index 9bdeb9197..57cff2883 100644
--- a/tests/autodiff/material2/IBSDF.slang
+++ b/tests/autodiff/material2/IBSDF.slang
@@ -1,6 +1,6 @@
 //TEST_IGNORE_FILE:
 
-//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
 interface IBSDF
diff --git a/tests/autodiff/material2/MxLayeredMaterial.slang b/tests/autodiff/material2/MxLayeredMaterial.slang
new file mode 100644
index 000000000..5e2717b60
--- /dev/null
+++ b/tests/autodiff/material2/MxLayeredMaterial.slang
@@ -0,0 +1,62 @@
+//TEST_IGNORE_FILE:
+
+__exported import IMaterial;
+__exported import IMaterialInstance;
+import MxLayeredMaterialInstance;
+import MxWeights;
+
+struct LayeredData_mixedLobes : IMxLayeredMaterialData
+{
+    static const int bsdfCount = 3;
+    static const int layerCount = 1;
+
+    int getBsdfCount() { return bsdfCount; }
+    int getLayerCount() { return layerCount; }
+}
+#define TMxLayeredMaterialData LayeredData_mixedLobes
+
+struct Layered_mixedLobes_WeightsCalculator : IMxLayeredWeightCalculator
+{
+    void calculateWeights<let TBsdfCount : int, let TLayerCount : int, TLayeredMaterialData : IMxLayeredMaterialData>(
+        const TLayeredMaterialData data, out MxWeights<TBsdfCount> weights)
+    {
+        for (uint i = 0; i < TBsdfCount; i++) weights.weights[i] = float3(0.f);
+    }
+}
+
+struct MxLayeredMaterial : IMaterial
+{
+    float3 baseColor;
+
+    typealias UsedMaterialInstance = MxLayeredMaterialInstance<3, 1, TMxLayeredMaterialData, Layered_mixedLobes_WeightsCalculator>;
+    typedef UsedMaterialInstance MaterialInstance;    
+
+    float3 getAlbedo(float3 albedo)
+    {
+        return albedo;
+    }
+
+    [ForwardDerivativeOf(getAlbedo)]
+    [TreatAsDifferentiable]
+    DifferentialPair<float3> __fwd_d_getAlbedo(DifferentialPair<float3> dpAlbedo)
+    {
+        return diffPair(dpAlbedo.p, float3(1.f));
+    }
+
+    [BackwardDerivativeOf(getAlbedo)]
+    [TreatAsDifferentiable]
+    void __bwd_d_getAlbedo(inout DifferentialPair<float3> dpAlbedo, float3 dOut)
+    {
+        [unroll]
+        for (int j = 0; j < 3; j++) outputBuffer[j + 6] += dOut[j];
+    }
+    
+    [Differentiable]
+    UsedMaterialInstance setupMaterialInstance(out MaterialInstanceData miData)
+    {
+        float3 albedo = getAlbedo(baseColor);
+        UsedMaterialInstance mi;
+        for (uint i = 0; i < 3; i++) miData.data[i] = albedo[i];
+        return mi;
+    }
+}
diff --git a/tests/autodiff/material2/MxLayeredMaterialInstance.slang b/tests/autodiff/material2/MxLayeredMaterialInstance.slang
new file mode 100644
index 000000000..53a319d6f
--- /dev/null
+++ b/tests/autodiff/material2/MxLayeredMaterialInstance.slang
@@ -0,0 +1,35 @@
+//TEST_IGNORE_FILE:
+
+__exported import IMaterialInstance;
+__exported import IBSDF;
+
+import MxWeights;
+
+struct MxLayeredMaterialInstance<
+    let TBsdfCount : int,
+    let TLayerCount : int,
+    TLayeredMaterialData : IMxLayeredMaterialData,
+    TWeightsCalc : IMxLayeredWeightCalculator> : IMaterialInstance
+{
+    TLayeredMaterialData data;
+    TWeightsCalc calculator;
+    typealias Weights = MxWeights<TBsdfCount>;
+
+    Weights calculateWeights()
+    {
+        Weights result;
+        calculator.calculateWeights<TBsdfCount, TLayerCount, TLayeredMaterialData>(data, result);
+        return result;
+    }
+
+    [Differentiable]
+    float3 eval(const MaterialInstanceData miData, const float3 wi, const float3 wo)
+    {
+        float3 albedo;
+        for (uint i = 0; i < 3; i++) albedo[i] = miData.data[i];
+
+        float3 result = float3(1.f);
+        [ForceUnroll] for (uint i = 0; i < TBsdfCount; i++) result *= albedo;
+        return result;
+    }
+}
diff --git a/tests/autodiff/material2/MxWeights.slang b/tests/autodiff/material2/MxWeights.slang
new file mode 100644
index 000000000..1d1a9f543
--- /dev/null
+++ b/tests/autodiff/material2/MxWeights.slang
@@ -0,0 +1,18 @@
+//TEST_IGNORE_FILE:
+
+struct MxWeights<let TBsdfCount : int>
+{
+    float3 weights[TBsdfCount];
+}
+
+interface IMxLayeredMaterialData
+{
+    int getBsdfCount();
+    int getLayerCount();
+}
+
+interface IMxLayeredWeightCalculator
+{
+    void calculateWeights<let TBsdfCount : int, let TLayerCount: int, TLayeredMaterialData : IMxLayeredMaterialData>(
+        const TLayeredMaterialData data, out MxWeights<TBsdfCount> weights);
+}
diff --git a/tests/autodiff/material2/diff-bwd-falcor-material-system.slang b/tests/autodiff/material2/diff-bwd-falcor-material-system.slang
index 2b873af77..bf57b8238 100644
--- a/tests/autodiff/material2/diff-bwd-falcor-material-system.slang
+++ b/tests/autodiff/material2/diff-bwd-falcor-material-system.slang
@@ -2,16 +2,19 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 
 // outputBuffer is defined in IBSDF.slang
-//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0 0 0 0], stride=4):out,name=outputBuffer
 
 import MaterialSystem;
 import DiffuseMaterial;
 import DiffuseMaterialInstance;
 import GlossyMaterial;
 import GlossyMaterialInstance;
+import MxLayeredMaterial;
+import MxLayeredMaterialInstance;
 
 //TEST_INPUT: type_conformance DiffuseMaterial:IMaterial = 0
 //TEST_INPUT: type_conformance GlossyMaterial:IMaterial = 1
+//TEST_INPUT: type_conformance MxLayeredMaterial:IMaterial = 2
 
 [BackwardDifferentiable]
 float3 evalBSDF(int type)
@@ -31,4 +34,5 @@ void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
 {
     __bwd_diff(evalBSDF)(0, float3(1.f));
     __bwd_diff(evalBSDF)(1, float3(1.f));
+    __bwd_diff(evalBSDF)(2, float3(1.f));
 }
diff --git a/tests/autodiff/material2/diff-bwd-falcor-material-system.slang.expected.txt b/tests/autodiff/material2/diff-bwd-falcor-material-system.slang.expected.txt
index e60005477..38ced0153 100644
--- a/tests/autodiff/material2/diff-bwd-falcor-material-system.slang.expected.txt
+++ b/tests/autodiff/material2/diff-bwd-falcor-material-system.slang.expected.txt
@@ -4,4 +4,7 @@ type: float
 0.0
 1.8
 1.2
-0.4
\ No newline at end of file
+0.4
+2.43
+1.08
+0.12
\ No newline at end of file
diff --git a/tests/autodiff/material2/diff-falcor-material-system.slang b/tests/autodiff/material2/diff-falcor-material-system.slang
index 7b0f39ed1..3c7155ec2 100644
--- a/tests/autodiff/material2/diff-falcor-material-system.slang
+++ b/tests/autodiff/material2/diff-falcor-material-system.slang
@@ -2,16 +2,19 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 
 // outputBuffer is defined in IBSDF.slang
-//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0 0 0 0], stride=4):out,name=outputBuffer
 
 import MaterialSystem;
 import DiffuseMaterial;
 import DiffuseMaterialInstance;
 import GlossyMaterial;
 import GlossyMaterialInstance;
+import MxLayeredMaterial;
+import MxLayeredMaterialInstance;
 
 //TEST_INPUT: type_conformance DiffuseMaterial:IMaterial = 0
 //TEST_INPUT: type_conformance GlossyMaterial:IMaterial = 1
+//TEST_INPUT: type_conformance MxLayeredMaterial:IMaterial = 2
 
 [BackwardDifferentiable]
 float3 evalBSDF(int type)
@@ -44,4 +47,12 @@ void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
         outputBuffer[4] = dp_f.d[1];
         outputBuffer[5] = dp_f.d[2];
     }
+
+    {
+        DifferentialPair<float3> dp_f = __fwd_diff(evalBSDF)(2);
+
+        outputBuffer[6] = dp_f.d[0];
+        outputBuffer[7] = dp_f.d[1];
+        outputBuffer[8] = dp_f.d[2];
+    }
 }
diff --git a/tests/autodiff/material2/diff-falcor-material-system.slang.expected.txt b/tests/autodiff/material2/diff-falcor-material-system.slang.expected.txt
index e60005477..38ced0153 100644
--- a/tests/autodiff/material2/diff-falcor-material-system.slang.expected.txt
+++ b/tests/autodiff/material2/diff-falcor-material-system.slang.expected.txt
@@ -4,4 +4,7 @@ type: float
 0.0
 1.8
 1.2
-0.4
\ No newline at end of file
+0.4
+2.43
+1.08
+0.12
\ No newline at end of file