Multi-block reverse-mode autodiff (#2576)

* Initial multi-block implementation

* Implemented multi-block reverse-mode (without loops)

* Added logic to remove block-level decorations to avoid confusing IR simplification passes

* Fixed issues with block-level decorations during IR simplification by removing them prior to simplification.

Co-authored-by: Yong He <yonghe@outlook.com>

1 files changed, 50 insertions, 1 deletions

diff --git a/source/slang/slang-ir-autodiff.cpp b/source/slang/slang-ir-autodiff.cpp
index f0ec1542e..40c24d11d 100644
--- a/source/slang/slang-ir-autodiff.cpp
+++ b/source/slang/slang-ir-autodiff.cpp
@@ -421,6 +421,32 @@ void stripAutoDiffDecorations(IRModule* module)
     stripAutoDiffDecorationsFromChildren(module->getModuleInst());
 }
 
+
+void stripBlockTypeDecorations(IRFunc* func)
+{
+    for (auto child : func->getChildren())
+    {
+        if (auto block = as<IRBlock>(child))
+        {
+            for (auto decor = block->getFirstDecoration(); decor; )
+            {
+                auto next = decor->getNextDecoration();
+                switch (decor->getOp())
+                {
+                case kIROp_DifferentialInstDecoration:
+                case kIROp_MixedDifferentialInstDecoration:
+                    decor->removeAndDeallocate();
+                    break;
+                default:
+                    break;
+                }
+                decor = next;
+            }
+        }
+    }
+}
+
+
 struct StripNoDiffTypeAttributePass : InstPassBase
 {
     StripNoDiffTypeAttributePass(IRModule* module) :
@@ -484,7 +510,7 @@ struct AutoDiffPass : public InstPassBase
         {
             bool changed = false;
             List<IRInst*> autoDiffWorkList;
-            // Collect all `ForwardDifferentiate` insts from the module.
+            // Collect all `ForwardDifferentiate`/`BackwardDifferentiate` insts from the module.
             autoDiffWorkList.clear();
             processAllInsts([&](IRInst* inst)
                 {
@@ -541,6 +567,7 @@ struct AutoDiffPass : public InstPassBase
             // Run transcription logic to generate the body of forward/backward derivatives functions.
             // While doing so, we may discover new functions to differentiate, so we keep running until
             // the worklist goes dry.
+            List<IRFunc*> autodiffCleanupList;
             while (autodiffContext->followUpFunctionsToTranscribe.getCount() != 0)
             {
                 changed = true;
@@ -549,6 +576,14 @@ struct AutoDiffPass : public InstPassBase
                 {
                     auto diffFunc = as<IRFunc>(task.resultFunc);
                     SLANG_ASSERT(diffFunc);
+
+                    // We're running in to some situations where the follow-up task
+                    // has already been completed (diffFunc has been generated, processed,
+                    // and deallocated). Skip over these for now.
+                    // 
+                    if (!diffFunc->getDataType())
+                        continue;
+
                     auto primalFunc = as<IRFunc>(task.originalFunc);
                     SLANG_ASSERT(primalFunc);
                     switch (task.type)
@@ -562,12 +597,26 @@ struct AutoDiffPass : public InstPassBase
                     default:
                         break;
                     }
+
+                    autodiffCleanupList.add(diffFunc);
                 }
             }
+
+            // Get rid of block-level decorations that are used to keep track of 
+            // different block types. These don't work well with the IR simplification
+            // passes since they don't expect decorations in blocks.
+            // 
+            for (auto diffFunc : autodiffCleanupList)
+                stripBlockTypeDecorations(diffFunc);
+
+            autodiffCleanupList.clear();
+
             if (!changed)
                 break;
             hasChanges |= changed;
         }
+
+
         return hasChanges;
     }