From 638e5fb000d4e242a91e8b653da4a72daec0efda Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Wed, 28 Aug 2024 09:23:08 -0700
Subject: Make tuple types work in autodiff. (#4923)

---
 source/slang/slang-ir-autodiff.cpp | 34 +++++++++++++++++++---------------
 1 file changed, 19 insertions(+), 15 deletions(-)

(limited to 'source/slang/slang-ir-autodiff.cpp')

diff --git a/source/slang/slang-ir-autodiff.cpp b/source/slang/slang-ir-autodiff.cpp
index bf83d8d7f..8ca7dbe76 100644
--- a/source/slang/slang-ir-autodiff.cpp
+++ b/source/slang/slang-ir-autodiff.cpp
@@ -448,26 +448,26 @@ void DifferentiableTypeConformanceContext::setFunc(IRGlobalValueWithCode* func)
                 IRBuilder subBuilder(item->getConcreteType());
                 if (as<IRTypePack>(concreteType) || as<IRTupleType>(concreteType))
                 {
-                    // For tuple types, register the differential type for each element, but don't register for the
+                    // For tuple types with concrete element types,
+                    // register the differential type for each element, but don't register for the
                     // tuple/typepack itself.
-                    auto witnessPack = as<IRMakeWitnessPack>(witness);
-                    SLANG_ASSERT(witnessPack);
-
-                    for (UInt i = 0; i < concreteType->getOperandCount(); i++)
+                    if (auto witnessPack = as<IRMakeWitnessPack>(witness))
                     {
-                        auto element = concreteType->getOperand(i);
-                        auto elementWitness = witnessPack->getOperand(i);
-                        differentiableWitnessDictionary.addIfNotExists(
-                            (IRType*)element,
-                            _lookupWitness(&subBuilder, elementWitness, sharedContext->differentialAssocTypeStructKey));
+
+                        for (UInt i = 0; i < concreteType->getOperandCount(); i++)
+                        {
+                            auto element = concreteType->getOperand(i);
+                            auto elementWitness = witnessPack->getOperand(i);
+                            differentiableWitnessDictionary.addIfNotExists(
+                                (IRType*)element,
+                                _lookupWitness(&subBuilder, elementWitness, sharedContext->differentialAssocTypeStructKey));
+                        }
+                        return;
                     }
-                    return;
-                }
-                else
-                {
-                    differentiableWitnessDictionary.add((IRType*)item->getConcreteType(), item->getWitness());
                 }
 
+                differentiableWitnessDictionary.add((IRType*)item->getConcreteType(), item->getWitness());
+
                 if (!as<IRInterfaceType>(item->getConcreteType()))
                 {
                     differentiableWitnessDictionary.addIfNotExists(
@@ -2241,12 +2241,16 @@ void releaseNullDifferentialType(AutoDiffSharedContext* context)
     {
         if (auto keepAliveDecoration = nullStruct->findDecoration<IRKeepAliveDecoration>())
             keepAliveDecoration->removeAndDeallocate();
+        if (auto exportDecoration = nullStruct->findDecoration<IRHLSLExportDecoration>())
+            exportDecoration->removeAndDeallocate();
     }
 
     if (auto nullWitness = context->nullDifferentialWitness)
     {
         if (auto keepAliveDecoration = nullWitness->findDecoration<IRKeepAliveDecoration>())
             keepAliveDecoration->removeAndDeallocate();
+        if (auto exportDecoration = nullWitness->findDecoration<IRHLSLExportDecoration>())
+            exportDecoration->removeAndDeallocate();
     }
 }
 
-- 
cgit v1.2.3