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>

15 files changed, 249 insertions, 4 deletions

diff --git a/source/slang/slang-ir-autodiff-fwd.cpp b/source/slang/slang-ir-autodiff-fwd.cpp
index 5cf3c1509..6025e1ccd 100644
--- a/source/slang/slang-ir-autodiff-fwd.cpp
+++ b/source/slang/slang-ir-autodiff-fwd.cpp
@@ -628,7 +628,7 @@ InstPair ForwardDiffTranscriber::transcribeCall(IRBuilder* builder, IRCall* orig
                 while (auto attrType = as<IRAttributedType>(origType))
                     origType = attrType->getBaseType();
             }
-            if (auto pairType = tryGetDiffPairType(&argBuilder, origType))
+            if (auto pairType = tryGetDiffPairType(&argBuilder, primalType))
             {
                 auto pairPtrType = as<IRPtrTypeBase>(pairType);
                 auto pairValType = as<IRDifferentialPairType>(
@@ -637,7 +637,7 @@ InstPair ForwardDiffTranscriber::transcribeCall(IRBuilder* builder, IRCall* orig
                 if (auto ptrParamType = as<IRPtrTypeBase>(diffParamType))
                 {
                     // Create temp var to pass in/out arguments.
-                    auto srcVar = argBuilder.emitVar(ptrParamType->getValueType());
+                    auto srcVar = argBuilder.emitVar(pairValType);
                     argBuilder.markInstAsMixedDifferential(srcVar, pairValType->getValueType());
 
                     auto diffArg = findOrTranscribeDiffInst(&argBuilder, origArg);
diff --git a/source/slang/slang-ir-autodiff-unzip.cpp b/source/slang/slang-ir-autodiff-unzip.cpp
index 53f0cbba2..c3ce32540 100644
--- a/source/slang/slang-ir-autodiff-unzip.cpp
+++ b/source/slang/slang-ir-autodiff-unzip.cpp
@@ -85,7 +85,7 @@ struct ExtractPrimalFuncContext
         List<IRType*> paramTypes;
         for (Index i = 0; i < ((Count) originalFuncType->getParamCount()) - 1; i++)
             paramTypes.add((IRType*)migrationContext.cloneInst(&builder, originalFuncType->getParamType(i)));
-        paramTypes.add(builder.getInOutType((IRType*)outIntermediateType));
+        paramTypes.add(builder.getOutType((IRType*)outIntermediateType));
         auto resultType = (IRType*)migrationContext.cloneInst(&builder, originalFuncType->getResultType());
         auto newFuncType = builder.getFuncType(paramTypes, resultType);
         return newFuncType;
@@ -183,7 +183,7 @@ struct ExtractPrimalFuncContext
         builder.setInsertInto(paramBlock);
         auto oldIntermediateParam = func->getLastParam();
         auto outIntermediary =
-            builder.emitParam(builder.getInOutType((IRType*)intermediateType));
+            builder.emitParam(builder.getOutType((IRType*)intermediateType));
         oldIntermediateParam->transferDecorationsTo(outIntermediary);
         primalParams.Add(outIntermediary);
         oldIntermediateParam->replaceUsesWith(outIntermediary);
diff --git a/source/slang/slang-ir-autodiff.cpp b/source/slang/slang-ir-autodiff.cpp
index 3257ee102..9a7a42619 100644
--- a/source/slang/slang-ir-autodiff.cpp
+++ b/source/slang/slang-ir-autodiff.cpp
@@ -896,6 +896,7 @@ bool canTypeBeStored(IRInst* type)
     case kIROp_DifferentialPairType:
     case kIROp_DifferentialPairUserCodeType:
     case kIROp_InterfaceType:
+    case kIROp_AssociatedType:
     case kIROp_AnyValueType:
     case kIROp_ClassType:
     case kIROp_FloatType:
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