Separate primal computations from unzipped function into an explicit function. (shader-slang#2569)

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

Author: csyonghe

build/visual-studio/slang/slang.vcxproj

@@ -523,6 +523,7 @@ IF EXIST ..\..\..\external\slang-glslang\bin\windows-aarch64\release\slang-glsla
     <ClCompile Include="..\..\..\source\slang\slang-ir-autodiff-fwd.cpp" />
     <ClCompile Include="..\..\..\source\slang\slang-ir-autodiff-pairs.cpp" />
     <ClCompile Include="..\..\..\source\slang\slang-ir-autodiff-rev.cpp" />
+    <ClCompile Include="..\..\..\source\slang\slang-ir-autodiff-unzip.cpp" />
     <ClCompile Include="..\..\..\source\slang\slang-ir-autodiff.cpp" />
     <ClCompile Include="..\..\..\source\slang\slang-ir-bind-existentials.cpp" />
     <ClCompile Include="..\..\..\source\slang\slang-ir-byte-address-legalize.cpp" />

build/visual-studio/slang/slang.vcxproj.filters

@@ -671,6 +671,9 @@
     <ClCompile Include="..\..\..\source\slang\slang-ir-autodiff-rev.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\..\..\source\slang\slang-ir-autodiff-unzip.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
     <ClCompile Include="..\..\..\source\slang\slang-ir-autodiff.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>

source/slang/slang-ir-autodiff-propagate.h

@@ -10,12 +10,12 @@
 namespace Slang
 {
 
-bool isDifferentialInst(IRInst* inst)
+inline bool isDifferentialInst(IRInst* inst)
 {
     return inst->findDecoration<IRDifferentialInstDecoration>();
 }
 
-bool isMixedDifferentialInst(IRInst* inst)
+inline bool isMixedDifferentialInst(IRInst* inst)
 {
     return inst->findDecoration<IRMixedDifferentialInstDecoration>();
 }
@@ -104,4 +104,4 @@ struct DiffPropagationPass : InstPassBase
     }
 };
 
-}
+}

source/slang/slang-ir-autodiff-rev.cpp

@@ -491,6 +491,98 @@ struct BackwardDiffTranscriber
         builder.emitBranch(firstBlock);
     }
 
+    void cleanUpUnusedPrimalIntermediate(IRInst* func, IRInst* primalFunc, IRInst* intermediateType)
+    {
+        IRStructType* structType = as<IRStructType>(intermediateType);
+        if (!structType)
+        {
+            auto genType = as<IRGeneric>(intermediateType);
+            structType = as<IRStructType>(findGenericReturnVal(genType));
+            SLANG_RELEASE_ASSERT(structType);
+        }
+
+        // Collect fields that are never fetched by reverse func.
+        OrderedHashSet<IRStructKey*> fieldsToCleanup;
+        for (auto children : structType->getChildren())
+        {
+            if (auto field = as<IRStructField>(children))
+            {
+                auto structKey = field->getKey();
+                bool usedByRevFunc = false;
+                for (auto use = structKey->firstUse; use; use = use->nextUse)
+                {
+                    if (isChildInstOf(use->getUser(), func))
+                    {
+                        usedByRevFunc = true;
+                        break;
+                    }
+                }
+                if (!usedByRevFunc)
+                {
+                    List<IRInst*> users;
+                    for (auto use = structKey->firstUse; use; use = use->nextUse)
+                    {
+                        users.add(use->getUser());
+                    }
+                    for (auto user : users)
+                    {
+                        if (!isChildInstOf(user, primalFunc))
+                            continue;
+                        if (auto addr = as<IRFieldAddress>(user))
+                        {
+                            if (addr->hasMoreThanOneUse())
+                                continue;
+                            if (addr->firstUse)
+                            {
+                                if (addr->firstUse->getUser()->getOp() == kIROp_Store)
+                                {
+                                    addr->firstUse->getUser()->removeAndDeallocate();
+                                }
+                                addr->removeAndDeallocate();
+                            }
+                        }
+                    }
+
+                    bool hasNonTrivialUse = false;
+                    for (auto use = structKey->firstUse; use; use = use->nextUse)
+                    {
+                        switch (use->getUser()->getOp())
+                        {
+                        case kIROp_PrimalValueStructKeyDecoration:
+                        case kIROp_StructField:
+                            continue;
+                        default:
+                            hasNonTrivialUse = true;
+                            break;
+                        }
+                    }
+                    if (!hasNonTrivialUse)
+                    {
+                        fieldsToCleanup.Add(structKey);
+                    }
+                }
+            }
+        }
+
+        // Actually remove fields from struct.
+        for (auto children : structType->getChildren())
+        {
+            if (auto field = as<IRStructField>(children))
+            {
+                if (fieldsToCleanup.Contains(field->getKey()))
+                {
+                    auto key = field->getKey();
+                    List<IRInst*> keyUsers;
+                    for (auto use = key->firstUse; use; use = use->nextUse)
+                        keyUsers.add(use->getUser());
+                    for (auto keyUser : keyUsers)
+                        keyUser->removeAndDeallocate();
+                    key->removeAndDeallocate();
+                }
+            }
+        }
+    }
+
     // Transcribe a function definition.
     InstPair transcribeFunc(IRBuilder* builder, IRFunc* primalFunc, IRFunc* diffFunc)
     {
@@ -520,12 +612,9 @@ struct BackwardDiffTranscriber
         // second block of the unzipped function.
         // 
         IRFunc* unzippedFwdDiffFunc = diffUnzipPass->unzipDiffInsts(fwdDiffFunc);
-        
+
         // Clone the primal blocks from unzippedFwdDiffFunc
         // to the reverse-mode function.
-        // TODO: This is the spot where we can make a decision to split
-        // the primal and differential into two different funcitons
-        // instead of two blocks in the same function.
         // 
         // Special care needs to be taken for the first block since it holds the parameters
 
@@ -547,6 +636,11 @@ struct BackwardDiffTranscriber
                 block->insertAtEnd(diffFunc);
         }
 
+        // Extracts the primal computations into its own func, and replace the primal insts
+        // with the intermediate results computed from the extracted func.
+        IRInst* intermediateType = nullptr;
+        auto extractedPrimalFunc = diffUnzipPass->extractPrimalFunc(diffFunc, unzippedFwdDiffFunc, intermediateType);
+
         // Transpose the first block (parameter block)
         transcribeParameterBlock(builder, diffFunc);
 
@@ -563,6 +657,9 @@ struct BackwardDiffTranscriber
         unzippedFwdDiffFunc->removeAndDeallocate();
         fwdDiffFunc->removeAndDeallocate();
 
+        eliminateDeadCode(diffFunc);
+        cleanUpUnusedPrimalIntermediate(diffFunc, extractedPrimalFunc, intermediateType);
+
         return InstPair(primalFunc, diffFunc);
     }