[Auto-diff] Overhaul auto-diff type tracking + Overhaul dynamic dispatch for differentiable functions (#5866)

* Overhauled the auto-diff system for dynamic dispatch

* More fixes

* remove intermediate dumps

* Update slang-ast-type.h

* More fixes + add a workaround for existential no-diff

* Update reverse-control-flow-3.slang

* remove dumps

* remove more dumps

* Delete working-reverse-control-flow-3.hlsl

* Cleanup comments + unused variables

* More comment cleanup

* Add support for lowering `DiffPairType(TypePack)` & `MakePair(MakeValuePack, MakeValuePack)`

* Fix array of issues in Falcor tests.

* Update slang-ir-autodiff-pairs.cpp

* More fixes for Falcor image tests

* Small fixups.

---------

Co-authored-by: Yong He <yonghe@outlook.com>

5 files changed, 163 insertions, 0 deletions

diff --git a/tests/autodiff/dynamic-dispatch-material.slang b/tests/autodiff/dynamic-dispatch-material.slang
new file mode 100644
index 000000000..1185a92e7
--- /dev/null
+++ b/tests/autodiff/dynamic-dispatch-material.slang
@@ -0,0 +1,142 @@
+// Test calling differentiable function through dynamic dispatch.
+
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+//TEST_INPUT: set g_materials = new StructuredBuffer<MaterialDataBlob>[new MaterialDataBlob{new MaterialHeader{[0, 0, 0, 0]}, new MaterialPayload{[1.0, 1.2, 0.3, 0.5]}}];
+RWStructuredBuffer<MaterialDataBlob> g_materials;
+
+public struct ShadingInput
+{
+    public float scale;
+}
+
+struct MaterialHeader
+{
+    uint4 header;
+};
+struct MaterialPayload
+{
+    float4 data;
+};
+struct MaterialDataBlob
+{
+    MaterialHeader header;   // 16B
+    MaterialPayload payload; // 16B
+};
+
+interface IMaterial : IDifferentiable
+{
+    associatedtype MaterialInstance : IMaterialInstance;
+
+    [Differentiable]
+    MaterialInstance setupMaterialInstance( ShadingInput input );
+}
+
+interface IMaterialInstance : IDifferentiable
+{
+    [Differentiable]
+    float eval( float x );
+}
+
+
+[BackwardDerivative(getMaterial_bwd)]
+IMaterial getMaterial(int id)
+{
+    return createDynamicObject<IMaterial, MaterialDataBlob>(id, g_materials[id]);
+}
+
+void getMaterial_bwd(int id, IDifferentiable d)
+{
+    // Something random
+    outputBuffer[id] = 2.f;
+}
+
+struct Material1: IMaterial
+{
+    typedef MaterialInstance1 MaterialInstance;
+
+    MaterialHeader header;
+    float a;
+    float b;
+    float c;
+
+    [Differentiable]
+    MaterialInstance1 setupMaterialInstance( ShadingInput input )
+    {
+        MaterialInstance1 instance;
+        instance.a = a * input.scale;
+        instance.b = b * input.scale;
+        instance.c = c * input.scale;
+        return instance;
+    }
+
+}
+struct MaterialInstance1: IMaterialInstance
+{
+    float a;
+    float b;
+    float c;
+
+    [Differentiable]
+    float eval( float x )
+    {
+        return a * x * x + b * x + c;
+    }
+}
+
+struct Material2: IMaterial
+{
+    typedef MaterialInstance2 MaterialInstance;
+
+    MaterialHeader header;
+    float a;
+    float b;
+
+    [Differentiable]
+    MaterialInstance2 setupMaterialInstance( ShadingInput input )
+    {
+        MaterialInstance2 instance;
+        instance.a = a * input.scale * input.scale;
+        instance.b = b * input.scale * input.scale;
+        return instance;
+    }
+
+}
+public struct MaterialInstance2: IMaterialInstance
+{
+    float a;
+    float b;
+
+    [Differentiable]
+    public float eval( float x )
+    {
+        return a * x + b;
+    }
+}
+
+[Differentiable]
+public float shade(int material, ShadingInput input, float x)
+{
+    IMaterial m = getMaterial(material);
+    IMaterialInstance mi = m.setupMaterialInstance(input);
+    return mi.eval(x);
+}
+
+//TEST_INPUT: type_conformance Material1:IMaterial = 0
+//TEST_INPUT: type_conformance Material2:IMaterial = 1
+
+[shader("compute")]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    outputBuffer[0] = shade(0, {0.5}, 0.6);
+
+    // TODO: VERIFY
+    DifferentialPair<float> dpx = diffPair(3.0);
+    bwd_diff(shade)(0, {0.5}, dpx, 1.0);
+
+    outputBuffer[3] = dpx.d;
+}
\ No newline at end of file
diff --git a/tests/autodiff/dynamic-dispatch-material.slang.expected.txt b/tests/autodiff/dynamic-dispatch-material.slang.expected.txt
new file mode 100644
index 000000000..49e4a0e6b
--- /dev/null
+++ b/tests/autodiff/dynamic-dispatch-material.slang.expected.txt
@@ -0,0 +1,6 @@
+type: float
+2.000000
+0.000000
+0.000000
+0.000000
+0.000000
diff --git a/tests/autodiff/material2/DiffuseMaterial.slang b/tests/autodiff/material2/DiffuseMaterial.slang
index d4ae579ff..1422ee30c 100644
--- a/tests/autodiff/material2/DiffuseMaterial.slang
+++ b/tests/autodiff/material2/DiffuseMaterial.slang
@@ -21,6 +21,11 @@ public struct DiffuseMaterial : IMaterial
         float3 albedo = getAlbedo(baseColor);
         DiffuseMaterialInstance mi = { baseColor };
         for (uint i = 0; i < 3; i++) miData.data[i] = albedo[i];
+
+        // unused
+        miData.data[3] = 0.f;
+        miData.data[4] = 0.f;
+        
         return mi;
     }
 };
diff --git a/tests/autodiff/material2/GlossyMaterial.slang b/tests/autodiff/material2/GlossyMaterial.slang
index e937443d4..1070c6e63 100644
--- a/tests/autodiff/material2/GlossyMaterial.slang
+++ b/tests/autodiff/material2/GlossyMaterial.slang
@@ -35,6 +35,11 @@ public struct GlossyMaterial : IMaterial
         float3 albedo = getAlbedo(baseColor);
         GlossyMaterialInstance mi = { baseColor };
         for (uint i = 0; i < 3; i++) miData.data[i] = albedo[i];
+
+        // unused
+        miData.data[3] = 0.f;
+        miData.data[4] = 0.f;
+        
         return mi;
     }
 };
diff --git a/tests/autodiff/material2/MxLayeredMaterial.slang b/tests/autodiff/material2/MxLayeredMaterial.slang
index 8c5cd6d3c..fabfde80c 100644
--- a/tests/autodiff/material2/MxLayeredMaterial.slang
+++ b/tests/autodiff/material2/MxLayeredMaterial.slang
@@ -57,6 +57,11 @@ public struct MxLayeredMaterial : IMaterial
         float3 albedo = getAlbedo(baseColor);
         UsedMaterialInstance mi;
         for (uint i = 0; i < 3; i++) miData.data[i] = albedo[i];
+
+        // unused
+        miData.data[3] = 0.f;
+        miData.data[4] = 0.f;
+        
         return mi;
     }
 }