Allow `no_diff` modifier on parameters (#2538)

1 files changed, 21 insertions, 5 deletions

diff --git a/source/slang/slang-ir-autodiff-rev.cpp b/source/slang/slang-ir-autodiff-rev.cpp
index 522c995b0..daf45e1ef 100644
--- a/source/slang/slang-ir-autodiff-rev.cpp
+++ b/source/slang/slang-ir-autodiff-rev.cpp
@@ -90,17 +90,33 @@ struct BackwardDiffTranscriber
 
         for (UIndex i = 0; i < funcType->getParamCount(); i++)
         {
+            bool noDiff = false;
             auto origType = funcType->getParamType(i);
-            if (auto diffPairType = tryGetDiffPairType(builder, origType))
+            if (auto attrType = as<IRAttributedType>(origType))
             {
-                auto inoutDiffPairType = builder->getPtrType(kIROp_InOutType, diffPairType);
-                newParameterTypes.add(inoutDiffPairType);
+                if (attrType->findAttr<IRNoDiffAttr>())
+                {
+                    noDiff = true;
+                    origType = attrType->getBaseType();
+                }
             }
-            else
+            if (noDiff)
+            {
                 newParameterTypes.add(origType);
+            }
+            else
+            {
+                if (auto diffPairType = tryGetDiffPairType(builder, origType))
+                {
+                    auto inoutDiffPairType = builder->getPtrType(kIROp_InOutType, diffPairType);
+                    newParameterTypes.add(inoutDiffPairType);
+                }
+                else
+                    newParameterTypes.add(origType);
+            }
         }
 
-        newParameterTypes.add(funcType->getResultType());
+        newParameterTypes.add(differentiateType(builder, funcType->getResultType()));
 
         diffReturnType = builder->getVoidType();