From 71efd949fa5276e2464416fcf237f8fd2c486281 Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Wed, 15 Mar 2023 22:26:58 -0400
Subject: AD: Primal-Hoisting Rework + Checkpoint Policy Framework (#2702)

---
 source/slang/slang-ir-autodiff-region.cpp | 56 +++++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 source/slang/slang-ir-autodiff-region.cpp

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

diff --git a/source/slang/slang-ir-autodiff-region.cpp b/source/slang/slang-ir-autodiff-region.cpp
new file mode 100644
index 000000000..98b64f179
--- /dev/null
+++ b/source/slang/slang-ir-autodiff-region.cpp
@@ -0,0 +1,56 @@
+// slang-ir-autodiff-region.cpp
+#include "slang-ir-autodiff-region.h"
+
+namespace Slang{
+    RefPtr<IndexedRegionMap> buildIndexedRegionMap(IRGlobalValueWithCode* func)
+    {
+        RefPtr<IndexedRegionMap> regionMap = new IndexedRegionMap;
+
+        List<IRBlock*> workList;
+        
+        regionMap->mapBlock(func->getFirstBlock(), nullptr);
+        workList.add(func->getFirstBlock());
+
+        while (workList.getCount() > 0)
+        {
+            auto currentBlock = workList.getLast();
+            workList.removeLast();
+
+            auto terminator = currentBlock->getTerminator();
+            auto currentRegion = regionMap->getRegion(currentBlock);
+
+            switch (terminator->getOp())
+            {
+                case kIROp_loop:
+                {
+                    auto loopRegion = regionMap->newRegion(as<IRLoop>(terminator), currentRegion);
+                    auto condBlock = as<IRLoop>(terminator)->getTargetBlock();
+
+                    regionMap->mapBlock(condBlock, loopRegion);
+                    workList.add(condBlock);
+                                    
+                    auto ifElse = as<IRIfElse>(condBlock->getTerminator());
+                    SLANG_RELEASE_ASSERT(ifElse);
+
+                    // TODO: this is one of the places we'll need to change if we support loops that
+                    // loop on either the true or false side. For now, we assume the loop is on the
+                    // true side only.
+                    //
+                    regionMap->mapBlock(ifElse->getFalseBlock(), currentRegion);
+                    workList.add(ifElse->getFalseBlock());
+                }
+            }
+
+            for (auto successor : currentBlock->getSuccessors())
+            {   
+                // If already mapped, skip.
+                if (regionMap->hasMapping(successor))
+                    continue;
+                regionMap->mapBlock(successor, currentRegion);
+                workList.add(successor);
+            }
+        }
+        
+        return regionMap;
+    }
+};
-- 
cgit v1.2.3