Make backward differentiation work with generics. (#2586)

* Make backward differentiation work with generics.

* Fix.

* Another fix.

* More fix.

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

1 files changed, 121 insertions, 42 deletions

diff --git a/source/slang/slang-ir-autodiff-rev.cpp b/source/slang/slang-ir-autodiff-rev.cpp
index 817534065..af408a5b3 100644
--- a/source/slang/slang-ir-autodiff-rev.cpp
+++ b/source/slang/slang-ir-autodiff-rev.cpp
@@ -20,27 +20,29 @@ namespace Slang
         {
             bool noDiff = false;
             auto origType = funcType->getParamType(i);
-            if (auto attrType = as<IRAttributedType>(origType))
+            auto primalType = (IRType*)findOrTranscribePrimalInst(builder, origType);
+
+            if (auto attrType = as<IRAttributedType>(primalType))
             {
                 if (attrType->findAttr<IRNoDiffAttr>())
                 {
                     noDiff = true;
-                    origType = attrType->getBaseType();
+                    primalType = attrType->getBaseType();
                 }
             }
             if (noDiff)
             {
-                newParameterTypes.add(origType);
+                newParameterTypes.add(primalType);
             }
             else
             {
-                if (auto diffPairType = tryGetDiffPairType(builder, origType))
+                if (auto diffPairType = tryGetDiffPairType(builder, primalType))
                 {
                     auto inoutDiffPairType = builder->getPtrType(kIROp_InOutType, diffPairType);
                     newParameterTypes.add(inoutDiffPairType);
                 }
                 else
-                    newParameterTypes.add(origType);
+                    newParameterTypes.add(primalType);
             }
         }
 
@@ -55,35 +57,47 @@ namespace Slang
 
         return builder->getFuncType(newParameterTypes, diffReturnType);
     }
+
+    static IRInst* getOriginalFuncRef(IRBuilder& builder, IRInst* func, IRInst* useSite)
+    {
+        if (!func) return nullptr;
+        auto userGeneric = findOuterGeneric(useSite);
+        if (!userGeneric) return func;
+        auto funcGen = findOuterGeneric(func);
+        SLANG_RELEASE_ASSERT(funcGen);
+        return maybeSpecializeWithGeneric(builder, funcGen, userGeneric);
+    }
     
     IRFuncType* BackwardDiffPrimalTranscriber::differentiateFunctionType(IRBuilder* builder, IRInst* func, IRFuncType* funcType)
     {
-        auto intermediateType = builder->getBackwardDiffIntermediateContextType(func);
+        auto funcRef = getOriginalFuncRef(*builder, func, builder->getInsertLoc().getParent());
+        auto intermediateType = builder->getBackwardDiffIntermediateContextType(funcRef);
         auto outType = builder->getOutType(intermediateType);
         List<IRType*> paramTypes;
         for (UInt i = 0; i < funcType->getParamCount(); i++)
         {
-            paramTypes.add(funcType->getParamType(i));
+            auto origType = funcType->getParamType(i);
+            auto primalType = (IRType*)findOrTranscribePrimalInst(builder, origType);
+            paramTypes.add(primalType);
         }
         paramTypes.add(outType);
         IRFuncType* primalFuncType = builder->getFuncType(
-            paramTypes, funcType->getResultType());
+            paramTypes, (IRType*)findOrTranscribePrimalInst(builder, funcType->getResultType()));
         return primalFuncType;
     }
 
     InstPair BackwardDiffPrimalTranscriber::transcribeFunc(IRBuilder* builder, IRFunc* primalFunc, IRFunc* diffFunc)
     {
-        SLANG_UNUSED(builder);
-        SLANG_UNUSED(diffFunc);
-        auto intermediateTypeDecor = primalFunc->findDecoration<IRBackwardDerivativeIntermediateTypeDecoration>();
-        SLANG_RELEASE_ASSERT(intermediateTypeDecor);
-        auto primalDecor = primalFunc->findDecoration<IRBackwardDerivativePrimalDecoration>();
-        return InstPair(primalFunc, primalDecor->getBackwardDerivativePrimalFunc());
+        // Don't need to do anything other than add a decoration in the original func to point to the primal func.
+        // The body of the primal func will be generated by propagateTranscriber together with propagate func.
+        addTranscribedFuncDecoration(*builder, primalFunc, diffFunc);
+        return InstPair(primalFunc, primalFunc);
     }
 
     IRFuncType* BackwardDiffPropagateTranscriber::differentiateFunctionType(IRBuilder* builder, IRInst* func, IRFuncType* funcType)
     {
-        auto intermediateType = builder->getBackwardDiffIntermediateContextType(func);
+        auto funcRef = getOriginalFuncRef(*builder, func, builder->getInsertLoc().getParent());
+        auto intermediateType = builder->getBackwardDiffIntermediateContextType(funcRef);
         return differentiateFunctionTypeImpl(builder, funcType, intermediateType);
     }
 
@@ -96,6 +110,7 @@ namespace Slang
     InstPair BackwardDiffPropagateTranscriber::transcribeFunc(IRBuilder* builder, IRFunc* primalFunc, IRFunc* diffFunc)
     {
         IRGlobalValueWithCode* diffPrimalFunc = nullptr;
+        addTranscribedFuncDecoration(*builder, primalFunc, diffFunc);
         transcribeFuncImpl(builder, primalFunc, diffFunc, diffPrimalFunc);
         return InstPair(primalFunc, diffFunc);
     }
@@ -211,8 +226,7 @@ namespace Slang
         if (!isMarkedForBackwardDifferentiation(origFunc))
             return InstPair(nullptr, nullptr);
 
-        IRBuilder builder(inBuilder->getSharedBuilder());
-        builder.setInsertBefore(origFunc);
+        IRBuilder builder = *inBuilder;
 
         IRFunc* primalFunc = origFunc;
 
@@ -221,6 +235,8 @@ namespace Slang
         auto diffFunc = builder.createFunc();
 
         SLANG_ASSERT(as<IRFuncType>(origFunc->getFullType()));
+        builder.setInsertBefore(diffFunc);
+
         IRType* diffFuncType = this->differentiateFunctionType(
             &builder,
             origFunc,
@@ -235,18 +251,6 @@ namespace Slang
             builder.addNameHintDecoration(diffFunc, newNameSb.getUnownedSlice());
         }
 
-        if (auto outerGen = findOuterGeneric(diffFunc))
-        {
-            builder.setInsertBefore(origFunc);
-            auto specialized =
-                specializeWithGeneric(builder, outerGen, as<IRGeneric>(findOuterGeneric(origFunc)));
-            addExistingDiffFuncDecor(&builder, origFunc, specialized);
-        }
-        else
-        {
-            addExistingDiffFuncDecor(&builder, origFunc, diffFunc);
-        }
-
         // Mark the generated derivative function itself as differentiable.
         builder.addBackwardDifferentiableDecoration(diffFunc);
 
@@ -259,6 +263,22 @@ namespace Slang
         return InstPair(primalFunc, diffFunc);
     }
 
+    void BackwardDiffTranscriberBase::addTranscribedFuncDecoration(IRBuilder& builder, IRFunc* origFunc, IRFunc* transcribedFunc)
+    {
+        IRBuilder subBuilder = builder;
+        if (auto outerGen = findOuterGeneric(transcribedFunc))
+        {
+            subBuilder.setInsertBefore(origFunc);
+            auto specialized =
+                specializeWithGeneric(subBuilder, outerGen, as<IRGeneric>(findOuterGeneric(origFunc)));
+            addExistingDiffFuncDecor(&subBuilder, origFunc, specialized);
+        }
+        else
+        {
+            addExistingDiffFuncDecor(&subBuilder, origFunc, transcribedFunc);
+        }
+    }
+
     InstPair BackwardDiffTranscriberBase::transcribeFuncHeader(IRBuilder* inBuilder, IRFunc* origFunc)
     {
         auto result = transcribeFuncHeaderImpl(inBuilder, origFunc);
@@ -288,7 +308,7 @@ namespace Slang
         List<IRType*> primalTypes, propagateTypes;
         for (UInt i = 0; i < funcType->getParamCount(); i++)
         {
-            auto paramType = funcType->getParamType(i);
+            auto paramType = (IRType*)findOrTranscribePrimalInst(&builder, funcType->getParamType(i));
             auto param = builder.emitParam(paramType);
             if (i != funcType->getParamCount() - 1)
             {
@@ -368,10 +388,8 @@ namespace Slang
         {
             IRParam* nextParam = param->getNextParam();
 
-            // Copy inst into the new parameter block.
-            auto clonedParam = cloneInst(&cloneEnv, &builder, param);
-            param->replaceUsesWith(clonedParam);
-            param->removeAndDeallocate();
+            // Move inst into the new parameter block.
+            param->insertAtEnd(paramBlock);
 
             param = nextParam;
         }
@@ -383,6 +401,62 @@ namespace Slang
         builder.emitBranch(firstBlock);
     }
 
+    // Create a copy of originalFunc's forward derivative in the same generic context (if any) of
+    // `diffPropagateFunc`.
+    IRFunc* BackwardDiffTranscriberBase::generateNewForwardDerivativeForFunc(
+        IRBuilder* builder, IRFunc* originalFunc, IRFunc* diffPropagateFunc)
+    {
+        auto primalOuterParent = findOuterGeneric(originalFunc);
+        if (!primalOuterParent)
+            primalOuterParent = originalFunc;
+
+        // Make a clone of original func so we won't modify the original.
+        IRCloneEnv originalCloneEnv;
+        primalOuterParent = cloneInst(&originalCloneEnv, builder, primalOuterParent);
+        auto primalFunc = as<IRFunc>(getGenericReturnVal(primalOuterParent));
+
+        // Strip any existing derivative decorations off the clone.
+        stripDerivativeDecorations(primalFunc);
+        eliminateDeadCode(primalOuterParent);
+
+        // Forward transcribe the clone of the original func.
+        ForwardDiffTranscriber fwdTranscriber(autoDiffSharedContext, builder->getSharedBuilder(), sink);
+        fwdTranscriber.pairBuilder = pairBuilder;
+        IRFunc* fwdDiffFunc = as<IRFunc>(getGenericReturnVal(fwdTranscriber.transcribe(builder, primalOuterParent)));
+        SLANG_ASSERT(fwdDiffFunc);
+        fwdTranscriber.transcribeFunc(builder, primalFunc, fwdDiffFunc);
+
+        // Remove the clone of original func.
+        primalOuterParent->removeAndDeallocate();
+
+        // Migrate the new forward derivative function into the generic parent of `diffPropagateFunc`.
+        if (auto fwdParentGeneric = as<IRGeneric>(findOuterGeneric(fwdDiffFunc)))
+        {
+            // Clone forward derivative func from its own generic into current generic parent.
+            GenericChildrenMigrationContext migrationContext;
+            auto diffOuterGeneric = as<IRGeneric>(findOuterGeneric(diffPropagateFunc));
+            SLANG_RELEASE_ASSERT(diffOuterGeneric);
+
+            migrationContext.init(fwdParentGeneric, diffOuterGeneric);
+            auto inst = fwdParentGeneric->getFirstBlock()->getFirstOrdinaryInst();
+            builder->setInsertBefore(diffPropagateFunc);
+            while (inst)
+            {
+                auto next = inst->getNextInst();
+                auto cloned = migrationContext.cloneInst(builder, inst);
+                if (inst == fwdDiffFunc)
+                {
+                    fwdDiffFunc = as<IRFunc>(cloned);
+                    break;
+                }
+                inst = next;
+            }
+            fwdParentGeneric->removeAndDeallocate();
+        }
+
+        return fwdDiffFunc;
+    }
+
     // Transcribe a function definition.
     void BackwardDiffTranscriberBase::transcribeFuncImpl(IRBuilder* builder, IRFunc* primalFunc, IRFunc* diffPropagateFunc, IRGlobalValueWithCode*& diffPrimalFunc)
     {
@@ -393,12 +467,16 @@ namespace Slang
 
         // Generate a temporary forward derivative function as an intermediate step.
         IRBuilder tempBuilder = *builder;
-        tempBuilder.setInsertBefore(diffPropagateFunc);
-        ForwardDiffTranscriber* fwdTranscriber = static_cast<ForwardDiffTranscriber*>(autoDiffSharedContext->transcriberSet.forwardTranscriber);
-        IRFunc* fwdDiffFunc = as<IRFunc>(fwdTranscriber->transcribeFuncHeaderImpl(&tempBuilder, primalFunc));
-        SLANG_ASSERT(fwdDiffFunc);
+        if (auto outerGeneric = findOuterGeneric(diffPropagateFunc))
+        {
+            tempBuilder.setInsertBefore(outerGeneric);
+        }
+        else
+        {
+            tempBuilder.setInsertBefore(diffPropagateFunc);
+        }
 
-        fwdTranscriber->transcribeFunc(&tempBuilder, primalFunc, fwdDiffFunc);
+        auto fwdDiffFunc = generateNewForwardDerivativeForFunc(&tempBuilder, primalFunc, diffPropagateFunc);
         
         // Split first block into a paramter block.
         this->makeParameterBlock(&tempBuilder, as<IRFunc>(fwdDiffFunc));
@@ -466,11 +544,12 @@ namespace Slang
         // we have just created.
         auto primalOuterGeneric = findOuterGeneric(primalFunc);
         IRInst* specializedFunc = nullptr;
-        auto intermediateTypeGeneric = hoistValueFromGeneric(*builder, intermediateType, specializedFunc);
+        auto intermediateTypeGeneric = hoistValueFromGeneric(*builder, intermediateType, specializedFunc, true);
+        builder->setInsertBefore(primalFunc);
         auto specializedIntermeidateType = maybeSpecializeWithGeneric(*builder, intermediateTypeGeneric, primalOuterGeneric);
         builder->addBackwardDerivativeIntermediateTypeDecoration(primalFunc, specializedIntermeidateType);
 
-        auto primalFuncGeneric = hoistValueFromGeneric(*builder, extractedPrimalFunc, specializedFunc);
+        auto primalFuncGeneric = hoistValueFromGeneric(*builder, extractedPrimalFunc, specializedFunc, true);
         builder->setInsertBefore(primalFunc);
 
         if (auto existingDecor = primalFunc->findDecoration<IRBackwardDerivativePrimalDecoration>())
@@ -568,7 +647,7 @@ namespace Slang
             return diffParam;
         }
 
-        return cloneInst(&cloneEnv, builder, origParam);
+        return maybeCloneForPrimalInst(builder, origParam);
     }
 
     InstPair BackwardDiffTranscriberBase::copyBinaryArith(IRBuilder* builder, IRInst* origArith)