AD: Various fixes around dynamic dispatch (#2820)

* Add a test for the new diff material system

* Various fixes for AD

- inout primal context params converted to out params,
- added attributed types to list of stored types
- used differentiated primal func type instead of type of differentiated func to avoid tangling with user-code differential types.

---------

Co-authored-by: Lifan Wu <lifanw@nvidia.com>

12 files changed, 244 insertions, 0 deletions

diff --git a/tests/autodiff/material2/DiffuseMaterial.slang b/tests/autodiff/material2/DiffuseMaterial.slang
new file mode 100644
index 000000000..721445249
--- /dev/null
+++ b/tests/autodiff/material2/DiffuseMaterial.slang
@@ -0,0 +1,26 @@
+//TEST_IGNORE_FILE:
+
+__exported import IMaterial;
+__exported import DiffuseMaterialInstance;
+
+struct DiffuseMaterial : IMaterial
+{
+    float3 baseColor;
+
+    typedef DiffuseMaterialInstance MaterialInstance;
+
+    [TreatAsDifferentiable]
+    float3 getAlbedo(float3 albedo)
+    {
+        return albedo;
+    }
+
+    [BackwardDifferentiable]
+    DiffuseMaterialInstance setupMaterialInstance(out MaterialInstanceData miData)
+    {
+        float3 albedo = getAlbedo(baseColor);
+        DiffuseMaterialInstance mi = { baseColor };
+        for (uint i = 0; i < 3; i++) miData.data[i] = albedo[i];
+        return mi;
+    }
+};
diff --git a/tests/autodiff/material2/DiffuseMaterialInstance.slang b/tests/autodiff/material2/DiffuseMaterialInstance.slang
new file mode 100644
index 000000000..e33e3f568
--- /dev/null
+++ b/tests/autodiff/material2/DiffuseMaterialInstance.slang
@@ -0,0 +1,17 @@
+//TEST_IGNORE_FILE:
+
+__exported import IMaterialInstance;
+__exported import IBSDF;
+
+struct DiffuseMaterialInstance : IMaterialInstance
+{
+    float3 albedo;
+
+    [BackwardDifferentiable]
+    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];
+        return albedo;
+    }
+};
diff --git a/tests/autodiff/material2/GlossyMaterial.slang b/tests/autodiff/material2/GlossyMaterial.slang
new file mode 100644
index 000000000..969370896
--- /dev/null
+++ b/tests/autodiff/material2/GlossyMaterial.slang
@@ -0,0 +1,40 @@
+//TEST_IGNORE_FILE:
+
+__exported import IMaterial;
+__exported import GlossyMaterialInstance;
+
+struct GlossyMaterial : IMaterial
+{
+    float3 baseColor;
+
+    typedef GlossyMaterialInstance 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(DifferentialPair<float3> dpAlbedo, float3 dOut)
+    {
+        [unroll]
+        for (int j = 0; j < 3; j++) outputBuffer[j + 3] += dOut[j];
+    }
+
+    [BackwardDifferentiable]
+    GlossyMaterialInstance setupMaterialInstance(out MaterialInstanceData miData)
+    {
+        float3 albedo = getAlbedo(baseColor);
+        GlossyMaterialInstance mi = { baseColor };
+        for (uint i = 0; i < 3; i++) miData.data[i] = albedo[i];
+        return mi;
+    }
+};
diff --git a/tests/autodiff/material2/GlossyMaterialInstance.slang b/tests/autodiff/material2/GlossyMaterialInstance.slang
new file mode 100644
index 000000000..3cef458ff
--- /dev/null
+++ b/tests/autodiff/material2/GlossyMaterialInstance.slang
@@ -0,0 +1,17 @@
+//TEST_IGNORE_FILE:
+
+__exported import IMaterialInstance;
+__exported import IBSDF;
+
+struct GlossyMaterialInstance : IMaterialInstance
+{
+    float3 albedo;
+
+    [BackwardDifferentiable]
+    float3 eval(const MaterialInstanceData miData, const float3 wi, const float3 wo)
+    {
+        float3 albedo = float3(0.0f, 0.0f, 0.0f);
+        for (uint i = 0; i < 3; i++) albedo[i] = miData.data[i];
+        return albedo * albedo;
+    }
+};
diff --git a/tests/autodiff/material2/IBSDF.slang b/tests/autodiff/material2/IBSDF.slang
new file mode 100644
index 000000000..9bdeb9197
--- /dev/null
+++ b/tests/autodiff/material2/IBSDF.slang
@@ -0,0 +1,10 @@
+//TEST_IGNORE_FILE:
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+interface IBSDF
+{
+    [BackwardDifferentiable]
+    float3 eval(const float3 wi, const float3 wo);
+}
diff --git a/tests/autodiff/material2/IMaterial.slang b/tests/autodiff/material2/IMaterial.slang
new file mode 100644
index 000000000..66780c9ff
--- /dev/null
+++ b/tests/autodiff/material2/IMaterial.slang
@@ -0,0 +1,11 @@
+//TEST_IGNORE_FILE:
+
+__exported import IMaterialInstance;
+
+interface IMaterial
+{
+    associatedtype MaterialInstance : IMaterialInstance;
+
+    [BackwardDifferentiable]
+    MaterialInstance setupMaterialInstance(out MaterialInstanceData miData);
+}
diff --git a/tests/autodiff/material2/IMaterialInstance.slang b/tests/autodiff/material2/IMaterialInstance.slang
new file mode 100644
index 000000000..ddf5ca0f6
--- /dev/null
+++ b/tests/autodiff/material2/IMaterialInstance.slang
@@ -0,0 +1,20 @@
+//TEST_IGNORE_FILE:
+
+struct MaterialInstanceData : IDifferentiable
+{
+    float data[5];
+}
+
+/*
+struct MaterialInstanceSetupResult : IDifferentiable
+{
+    IMaterialInstance mi;
+    MaterialInstanceData data;
+}
+*/
+
+interface IMaterialInstance
+{
+    [BackwardDifferentiable]
+    float3 eval(const MaterialInstanceData miData, const float3 wi, const float3 wo);
+}
diff --git a/tests/autodiff/material2/MaterialSystem.slang b/tests/autodiff/material2/MaterialSystem.slang
new file mode 100644
index 000000000..5e733fad9
--- /dev/null
+++ b/tests/autodiff/material2/MaterialSystem.slang
@@ -0,0 +1,8 @@
+//TEST_IGNORE_FILE:
+
+__exported import IMaterial;
+
+IMaterial createMaterialClassConformance(int type, float3 value)
+{
+    return createDynamicObject<IMaterial, float3>(type, value);
+}
diff --git a/tests/autodiff/material2/diff-bwd-falcor-material-system.slang b/tests/autodiff/material2/diff-bwd-falcor-material-system.slang
new file mode 100644
index 000000000..2b873af77
--- /dev/null
+++ b/tests/autodiff/material2/diff-bwd-falcor-material-system.slang
@@ -0,0 +1,34 @@
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+//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
+
+import MaterialSystem;
+import DiffuseMaterial;
+import DiffuseMaterialInstance;
+import GlossyMaterial;
+import GlossyMaterialInstance;
+
+//TEST_INPUT: type_conformance DiffuseMaterial:IMaterial = 0
+//TEST_INPUT: type_conformance GlossyMaterial:IMaterial = 1
+
+[BackwardDifferentiable]
+float3 evalBSDF(int type)
+{
+    float3 wi = normalize(float3(0.5, 0.2, 0.8));
+    float3 wo = normalize(float3(-0.1, -0.3, 0.9));
+
+    IMaterial material = createMaterialClassConformance(type, float3(0.9f, 0.6f, 0.2f));
+    MaterialInstanceData miData;
+    let mi = material.setupMaterialInstance(miData);
+    float3 f = mi.eval(miData, wi, wo);
+    return f;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    __bwd_diff(evalBSDF)(0, float3(1.f));
+    __bwd_diff(evalBSDF)(1, 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
new file mode 100644
index 000000000..e60005477
--- /dev/null
+++ b/tests/autodiff/material2/diff-bwd-falcor-material-system.slang.expected.txt
@@ -0,0 +1,7 @@
+type: float
+0.0
+0.0
+0.0
+1.8
+1.2
+0.4
\ 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
new file mode 100644
index 000000000..7b0f39ed1
--- /dev/null
+++ b/tests/autodiff/material2/diff-falcor-material-system.slang
@@ -0,0 +1,47 @@
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+//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
+
+import MaterialSystem;
+import DiffuseMaterial;
+import DiffuseMaterialInstance;
+import GlossyMaterial;
+import GlossyMaterialInstance;
+
+//TEST_INPUT: type_conformance DiffuseMaterial:IMaterial = 0
+//TEST_INPUT: type_conformance GlossyMaterial:IMaterial = 1
+
+[BackwardDifferentiable]
+float3 evalBSDF(int type)
+{
+    float3 wi = normalize(float3(0.5, 0.2, 0.8));
+    float3 wo = normalize(float3(-0.1, -0.3, 0.9));
+
+    IMaterial material = createMaterialClassConformance(type, float3(0.9f, 0.6f, 0.2f));
+    MaterialInstanceData miData;
+    let mi = material.setupMaterialInstance(miData);
+    float3 f = mi.eval(miData, wi, wo);
+    return f;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    {
+        DifferentialPair<float3> dp_f = __fwd_diff(evalBSDF)(0);
+
+        outputBuffer[0] = dp_f.d[0];
+        outputBuffer[1] = dp_f.d[1];
+        outputBuffer[2] = dp_f.d[2];
+    }
+
+    {
+        DifferentialPair<float3> dp_f = __fwd_diff(evalBSDF)(1);
+
+        outputBuffer[3] = dp_f.d[0];
+        outputBuffer[4] = dp_f.d[1];
+        outputBuffer[5] = 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
new file mode 100644
index 000000000..e60005477
--- /dev/null
+++ b/tests/autodiff/material2/diff-falcor-material-system.slang.expected.txt
@@ -0,0 +1,7 @@
+type: float
+0.0
+0.0
+0.0
+1.8
+1.2
+0.4
\ No newline at end of file