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>

10 files changed, 210 insertions, 5 deletions

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