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>

8 files changed, 141 insertions, 5 deletions

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