From ceb87b25b387411dbb7978caf04ecd9e948493e3 Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Mon, 26 Feb 2024 18:22:16 -0500
Subject: AD: Handle case where struct is created with `no_diff`-wrapped
 operands. (#3629)

* Handle case where struct is created with `no_diff`-wrapped operands.

* Add test to reproduce crash upon initializing a struct with `no_diff`-wrapped operand type

* Add expected result for test
---
 source/slang/slang-ir-autodiff-fwd.cpp             | 19 +++++++++++---
 tests/autodiff/make-struct-non-diff-operand.slang  | 30 ++++++++++++++++++++++
 ...make-struct-non-diff-operand.slang.expected.txt |  5 ++++
 3 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 tests/autodiff/make-struct-non-diff-operand.slang
 create mode 100644 tests/autodiff/make-struct-non-diff-operand.slang.expected.txt

diff --git a/source/slang/slang-ir-autodiff-fwd.cpp b/source/slang/slang-ir-autodiff-fwd.cpp
index 2f61dfa72..334d4e678 100644
--- a/source/slang/slang-ir-autodiff-fwd.cpp
+++ b/source/slang/slang-ir-autodiff-fwd.cpp
@@ -498,9 +498,22 @@ InstPair ForwardDiffTranscriber::transcribeMakeStruct(IRBuilder* builder, IRInst
             {
                 auto operandDataType = origMakeStruct->getOperand(ii)->getDataType();
                 auto diffOperandType = differentiateType(builder, operandDataType);
-                SLANG_RELEASE_ASSERT(diffOperandType);
-                operandDataType = (IRType*)findOrTranscribePrimalInst(builder, operandDataType);
-                diffOperands.add(getDifferentialZeroOfType(builder, operandDataType));
+
+                if (diffOperandType)
+                {
+                    operandDataType = (IRType*)findOrTranscribePrimalInst(builder, operandDataType);
+                    diffOperands.add(getDifferentialZeroOfType(builder, operandDataType));
+                }
+                else
+                {
+                    // This case is only hit if the field is of a differentiable type but the operand is of
+                    // a non-differentiable type. This can happen if the operand is wrapped in no_diff.
+                    // In this case, we use the derivative of the field type to synthesize the 0.
+                    // 
+                    auto diffFieldOperandType = differentiateType(builder, field->getFieldType());
+                    SLANG_RELEASE_ASSERT(diffFieldOperandType);
+                    diffOperands.add(getDifferentialZeroOfType(builder, (IRType*)diffFieldOperandType));
+                }
             }
             ii++;
         }
diff --git a/tests/autodiff/make-struct-non-diff-operand.slang b/tests/autodiff/make-struct-non-diff-operand.slang
new file mode 100644
index 000000000..cdfb6b6bc
--- /dev/null
+++ b/tests/autodiff/make-struct-non-diff-operand.slang
@@ -0,0 +1,30 @@
+//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], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+struct AggType : IDifferentiable
+{
+    float field1;
+    float field2;
+}
+
+[BackwardDifferentiable]
+AggType makeTypeWithNonDiffOps(no_diff float f1, float f2)
+{
+    AggType v = { f1, f2 };
+    return v;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    AggType v = { 0, 2.0 };
+    AggType.Differential dv = { 1.3, 4.2 };
+
+    var f2pair = diffPair(2.0, 0.0); 
+
+    __bwd_diff(makeTypeWithNonDiffOps)(1.0, f2pair, dv);
+    outputBuffer[0] = f2pair.d;
+}
\ No newline at end of file
diff --git a/tests/autodiff/make-struct-non-diff-operand.slang.expected.txt b/tests/autodiff/make-struct-non-diff-operand.slang.expected.txt
new file mode 100644
index 000000000..0f62dd9c7
--- /dev/null
+++ b/tests/autodiff/make-struct-non-diff-operand.slang.expected.txt
@@ -0,0 +1,5 @@
+type: float
+4.200000
+0.000000
+0.000000
+0.000000
-- 
cgit v1.2.3