Hack handling of primal insts that has a function type. (#2728)

* Update diff-bwd material test

* Minor update

* Hack handling of primal insts that has a function type.

---------

Co-authored-by: winmad <winmad.wlf@gmail.com>
Co-authored-by: Yong He <yhe@nvidia.com>

7 files changed, 83 insertions, 6 deletions

diff --git a/source/slang/slang-ir-autodiff-transcriber-base.cpp b/source/slang/slang-ir-autodiff-transcriber-base.cpp
index 2bc67e561..552ac762c 100644
--- a/source/slang/slang-ir-autodiff-transcriber-base.cpp
+++ b/source/slang/slang-ir-autodiff-transcriber-base.cpp
@@ -1116,9 +1116,12 @@ IRInst* AutoDiffTranscriberBase::transcribe(IRBuilder* builder, IRInst* origInst
                 }
                 else
                 {
-                    if (!pair.primal->findDecoration<IRAutodiffInstDecoration>()
-                        && !as<IRConstant>(pair.differential))
+                    if (!pair.primal->findDecoration<IRAutodiffInstDecoration>())
                     {
+                        if (as<IRConstant>(pair.differential))
+                            break;
+                        if (as<IRType>(pair.differential))
+                            break;
                         auto mixedType = (IRType*)(pair.primal->getDataType());
                         builder->markInstAsMixedDifferential(pair.primal, mixedType);
                     }
diff --git a/source/slang/slang-ir-autodiff-transpose.h b/source/slang/slang-ir-autodiff-transpose.h
index 653589933..bbdb01290 100644
--- a/source/slang/slang-ir-autodiff-transpose.h
+++ b/source/slang/slang-ir-autodiff-transpose.h
@@ -950,8 +950,25 @@ struct DiffTransposePass
             // Slang doesn't support function values. So if we see a func-typed inst
             // it's proabably a reference to a function.
             // 
-            if (as<IRFuncType>(child->getDataType()))
+            switch (child->getOp())
+            {
+            /*
+               TODO: need a better way to move specialize, lookupwitness, extractExistentialType/Value/Witness
+               insts to a proper location that dominates all their use sites. Create copies of these insts
+               when necessary.
+                case kIROp_Specialize:
+                case kIROp_LookupWitness:
+                case kIROp_ExtractExistentialType:
+                case kIROp_ExtractExistentialValue:
+                case kIROp_ExtractExistentialWitnessTable:
+            */
+            case kIROp_ForwardDifferentiate:
+            case kIROp_BackwardDifferentiate:
+            case kIROp_BackwardDifferentiatePrimal:
+            case kIROp_BackwardDifferentiatePropagate:
                 typeInsts.add(child);
+                break;
+            }
         }
 
         for (auto inst : typeInsts)
@@ -959,7 +976,6 @@ struct DiffTransposePass
             inst->insertAtEnd(revBlock);
         }
 
-
         // Then, go backwards through the regular instructions, and transpose them into the new
         // rev block.
         // Note the 'reverse' traversal here.
@@ -2221,6 +2237,10 @@ struct DiffTransposePass
             case kIROp_ifElse:
             case kIROp_loop:
             case kIROp_Switch:
+            case kIROp_LookupWitness:
+            case kIROp_ExtractExistentialType:
+            case kIROp_ExtractExistentialValue:
+            case kIROp_ExtractExistentialWitnessTable:
             {
                 // Ignore. transposeBlock() should take care of adding the
                 // appropriate branch instruction.
@@ -3474,6 +3494,9 @@ struct DiffTransposePass
     List<IRUse*>                                         primalUsesToHoist;
 
     Dictionary<IRStore*, IRBlock*>                       mapStoreToDefBlock;
+
+    IRCloneEnv typeInstCloneEnv = {};
+
 };
 
 
diff --git a/tests/autodiff/material/GlossyMaterialInstance.slang b/tests/autodiff/material/GlossyMaterialInstance.slang
index b1430e62c..fbbcc1a26 100644
--- a/tests/autodiff/material/GlossyMaterialInstance.slang
+++ b/tests/autodiff/material/GlossyMaterialInstance.slang
@@ -7,18 +7,26 @@ struct GlossyBSDF : IBSDF
 {
     float3 albedo;
 
-    [BackwardDifferentiable]
     float3 getAlbedo()
     {
         return albedo;
     }
 
     [ForwardDerivativeOf(getAlbedo)]
+    [TreatAsDifferentiable]
     DifferentialPair<float3> __fwd_d_getAlbedo()
     {
         return diffPair(albedo, float3(1.f));
     }
 
+    [BackwardDerivativeOf(getAlbedo)]
+    [TreatAsDifferentiable]
+    void __bwd_d_getAlbedo(float3 dOut)
+    {
+        [unroll]
+        for (int j = 0; j < 3; j++) outputBuffer[j + 3] += dOut[j];
+    }
+
     [BackwardDifferentiable]
     float3 eval(const float3 wiLocal, const float3 woLocal)
     {
diff --git a/tests/autodiff/material/IBSDF.slang b/tests/autodiff/material/IBSDF.slang
index 41b0957db..9bdeb9197 100644
--- a/tests/autodiff/material/IBSDF.slang
+++ b/tests/autodiff/material/IBSDF.slang
@@ -1,5 +1,8 @@
 //TEST_IGNORE_FILE:
 
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
 interface IBSDF
 {
     [BackwardDifferentiable]
diff --git a/tests/autodiff/material/diff-bwd-falcor-material-system.slang b/tests/autodiff/material/diff-bwd-falcor-material-system.slang
new file mode 100644
index 000000000..6d4101d12
--- /dev/null
+++ b/tests/autodiff/material/diff-bwd-falcor-material-system.slang
@@ -0,0 +1,33 @@
+//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));
+    let mi = material.setupMaterialInstance();
+    float3 f = mi.eval(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/material/diff-bwd-falcor-material-system.slang.expected.txt b/tests/autodiff/material/diff-bwd-falcor-material-system.slang.expected.txt
new file mode 100644
index 000000000..e60005477
--- /dev/null
+++ b/tests/autodiff/material/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/material/diff-falcor-material-system.slang b/tests/autodiff/material/diff-falcor-material-system.slang
index 1cb8275fa..875432c28 100644
--- a/tests/autodiff/material/diff-falcor-material-system.slang
+++ b/tests/autodiff/material/diff-falcor-material-system.slang
@@ -1,8 +1,8 @@
 //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
-RWStructuredBuffer<float> outputBuffer;
 
 import MaterialSystem;
 import DiffuseMaterial;