Fix inlining pass. (#2506)

* Fix inlining pass.

* Add more check against corner cases.

* Revise comments.

* Fixes.

* Fix premake script.

* Fixes.

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

7 files changed, 154 insertions, 58 deletions

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 9c72f1d63..de9d23e97 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -401,8 +401,6 @@ Result linkAndOptimizeIR(
     if (sink->getErrorCount() != 0)
         return SLANG_FAIL;
 
-    eliminateMultiLevelBreak(irModule);
-
     // TODO(DG): There are multiple DCE steps here, which need to be changed
     //   so that they don't just throw out any non-entry point code
     // Debugging code for IR transformations...
@@ -804,6 +802,8 @@ Result linkAndOptimizeIR(
             LivenessUtil::addVariableRangeStarts(irModule, livenessMode);
         }
 
+        eliminateMultiLevelBreak(irModule);
+
         // As a late step, we need to take the SSA-form IR and move things *out*
         // of SSA form, by eliminating all "phi nodes" (block parameters) and
         // introducing explicit temporaries instead. Doing this at the IR level
diff --git a/source/slang/slang-ir-inline.cpp b/source/slang/slang-ir-inline.cpp
index 8f5412fa6..13221ee8c 100644
--- a/source/slang/slang-ir-inline.cpp
+++ b/source/slang/slang-ir-inline.cpp
@@ -2,7 +2,6 @@
 #include "slang-ir-inline.h"
 
 #include "slang-ir-ssa-simplification.h"
-#include "slang-ir-single-return.h"
 
 // This file provides general facilities for inlining function calls.
 
@@ -318,12 +317,7 @@ struct InliningPassBase
             SLANG_ASSERT(argCounter == (Int)call->getArgCount());
         }
 
-        // For now, our inlining pass only handles the case where
-        // the callee is a "single-return" function, which means the callee
-        // function contains only one return at the end of the body.
-        
-        convertFuncToSingleReturnForm(m_module, callSite.callee);
-        inlineSingleReturnFuncBody(callSite, &env, &builder);    
+        inlineFuncBody(callSite, &env, &builder);    
     }
 
         // When instructions are cloned, with cloneInst no sourceLoc information is copied over by default.
@@ -375,20 +369,13 @@ struct InliningPassBase
         return clonedInst;
     }
 
-        /// Inline the body of the callee for `callSite`, where the callee has a single return.
-    void inlineSingleReturnFuncBody(
+        /// Inline the body of the callee for `callSite`.
+    void inlineFuncBody(
         CallSiteInfo const& callSite, IRCloneEnv* env, IRBuilder* builder)
     {
         auto call = callSite.call;
         auto callee = callSite.callee;
 
-        // We know that the callee has a single return block, so if we encounter
-        // a `returnVal` instruction then it must be the one and only
-        // return point for the function, and its operand will be the value
-        // the callee returns.
-        //
-        IRInst* returnedValue = nullptr;
-
         // Break the basic block containing the call inst into two basic blocks.
         auto callerBlock = callSite.call->getParent();
         builder->setInsertInto(callerBlock->getParent());
@@ -399,6 +386,13 @@ struct InliningPassBase
         // second half of the basic block right after `callerBlock`.
         afterBlock->insertAfter(callerBlock);
         afterBlock->sourceLoc = callSite.call->getNextInst()->sourceLoc;
+        // Define a param in afterBlock to receive the return value from the call.
+        builder->setInsertInto(afterBlock);
+
+        IRInst* returnValParam = nullptr;
+        if (callSite.call->getDataType()->getOp() != kIROp_VoidType)
+            returnValParam = builder->emitParam(callSite.call->getDataType());
+
         // Move all insts after the call in `callerBlock` to `afterBlock`.
         {
             auto inst = callSite.call->getNextInst();
@@ -422,8 +416,10 @@ struct InliningPassBase
 
         // Insert a branch into the cloned first block at the end of `callerBlock`.
         builder->setInsertInto(callerBlock);
-        auto newBranch = builder->emitBranch(as<IRBlock>(env->mapOldValToNew[callee->getFirstBlock()].GetValue()));
+        auto mainBlock = as<IRBlock>(env->mapOldValToNew[callee->getFirstBlock()].GetValue());
+        auto newBranch = builder->emitLoop(mainBlock, afterBlock, mainBlock);
         _setSourceLoc(newBranch, call, callSite);
+
         // Clone all basic blocks over to the call site.
         bool isFirstBlock = true;
         for (auto calleeBlock : callee->getBlocks())
@@ -464,9 +460,10 @@ struct InliningPassBase
                     // of the original call.
                     //
                     {
-                        auto returnBranch = builder->emitBranch(afterBlock);
+                        auto returnedValue = findCloneForOperand(env, inst->getOperand(0));
+                        auto returnBranch = builder->emitBranch(
+                            afterBlock, returnValParam ? 1 : 0, &returnedValue);
                         _setSourceLoc(returnBranch, inst, callSite);
-                        returnedValue = findCloneForOperand(env, inst->getOperand(0));
                     }
                     break;
                 }
@@ -478,9 +475,13 @@ struct InliningPassBase
         // the return value of the inlined function, then that value
         // should be used to replace any uses of the original call.
         //
-        if( returnedValue )
+        if (returnValParam)
+        {
+            call->replaceUsesWith(returnValParam);
+        }
+        else
         {
-            call->replaceUsesWith(returnedValue);
+            call->replaceUsesWith(builder->getVoidValue());
         }
 
         // Once we've cloned the body of the callee in at the call site,
diff --git a/source/slang/slang-ir-restructure.cpp b/source/slang/slang-ir-restructure.cpp
index 7cd33406d..b0f822def 100644
--- a/source/slang/slang-ir-restructure.cpp
+++ b/source/slang/slang-ir-restructure.cpp
@@ -176,7 +176,8 @@ namespace Slang
                 //
                 if(block != registeredBlock[(int)ll->op])
                 {
-                    ctx->getSink()->diagnose(block, Diagnostics::multiLevelBreakUnsupported);
+                    if (ctx->getSink())
+                        ctx->getSink()->diagnose(block, Diagnostics::multiLevelBreakUnsupported);
                 }
 
                 // Now we need to create a structured `break` or `continue` operation
diff --git a/source/slang/slang-ir-simplify-cfg.cpp b/source/slang/slang-ir-simplify-cfg.cpp
index a1bc38b64..1e247d1d9 100644
--- a/source/slang/slang-ir-simplify-cfg.cpp
+++ b/source/slang/slang-ir-simplify-cfg.cpp
@@ -2,16 +2,109 @@
 
 #include "slang-ir-insts.h"
 #include "slang-ir.h"
+#include "slang-ir-dominators.h"
+#include "slang-ir-restructure.h"
 
 namespace Slang
 {
 
-bool processFunc(IRGlobalValueWithCode* func)
+struct CFGSimplificationContext
+{
+    RefPtr<RegionTree> regionTree;
+    RefPtr<IRDominatorTree> domTree;
+};
+
+static BreakableRegion* findBreakableRegion(Region* region)
+{
+    for (;;)
+    {
+        if (auto b = as<BreakableRegion>(region))
+            return b;
+        region = region->getParent();
+        if (!region)
+            return nullptr;
+    }
+}
+
+// Test if a loop is trivial: a trivial loop runs for a single iteration without any back edges, and
+// there is only one break out of the loop at the very end. The function generates `regionTree` if
+// it is needed and hasn't been generated yet.
+static bool isTrivialSingleIterationLoop(
+    IRGlobalValueWithCode* func,
+    IRLoop* loop,
+    CFGSimplificationContext& inoutContext)
+{
+    auto targetBlock = loop->getTargetBlock();
+    if (targetBlock->getPredecessors().getCount() != 1) return false;
+    if (*targetBlock->getPredecessors().begin() != loop->getParent()) return false;
+
+    int useCount = 0;
+    for (auto use = loop->getBreakBlock()->firstUse; use; use = use->nextUse)
+    {
+        if (use->getUser() == loop)
+            continue;
+        useCount++;
+        if (useCount > 1)
+            return false;
+    }
+
+    // The loop has passed simple test.
+    // 
+    // We need to verify this is a trivial loop by checking if there is any multi-level breaks
+    // that skips out of this loop.
+
+    if (!inoutContext.domTree)
+        inoutContext.domTree = computeDominatorTree(func);
+    if (!inoutContext.regionTree)
+        inoutContext.regionTree = generateRegionTreeForFunc(func, nullptr);
+
+    SimpleRegion* targetBlockRegion = nullptr;
+    if (!inoutContext.regionTree->mapBlockToRegion.TryGetValue(targetBlock, targetBlockRegion))
+        return false;
+    BreakableRegion* loopBreakableRegion = findBreakableRegion(targetBlockRegion);
+    LoopRegion* loopRegion = as<LoopRegion>(loopBreakableRegion);
+    if (!loopRegion)
+        return false;
+    for (auto block : func->getBlocks())
+    {
+        if (!inoutContext.domTree->dominates(loop->getTargetBlock(), block))
+            continue;
+        if (inoutContext.domTree->dominates(loop->getBreakBlock(), block))
+            continue;
+        SimpleRegion* region = nullptr;
+        if (!inoutContext.regionTree->mapBlockToRegion.TryGetValue(block, region))
+            return false;
+
+        for (auto branchTarget : block->getSuccessors())
+        {
+            SimpleRegion* targetRegion = nullptr;
+            if (!inoutContext.regionTree->mapBlockToRegion.TryGetValue(branchTarget, targetRegion))
+                return false;
+            // If multi-level break out that skips over this loop exists, then this is not a trivial loop.
+            if (targetRegion->isDescendentOf(loopRegion))
+                continue;
+            if (targetBlock != loop->getBreakBlock())
+                return false;
+            if (findBreakableRegion(region) != loopRegion)
+            {
+                // If the break is initiated from a nested region, this is not trivial.
+                return false;
+            }
+        }
+    }
+
+    return true;
+}
+
+static bool processFunc(IRGlobalValueWithCode* func)
 {
     auto firstBlock = func->getFirstBlock();
     if (!firstBlock)
         return false;
 
+    // Lazily generated region tree.
+    CFGSimplificationContext simplificationContext;
+
     SharedIRBuilder sharedBuilder(func->getModule());
     IRBuilder builder(&sharedBuilder);
 
@@ -35,10 +128,12 @@ bool processFunc(IRGlobalValueWithCode* func)
                     loop->continueBlock.set(loop->getTargetBlock());
                     continueBlock->removeAndDeallocate();
                 }
-                // If there isn't any actual back jumps into loop target, remove the header
-                // and turn it into a normal branch.
+
+                // If there isn't any actual back jumps into loop target and there is a trivial
+                // break at the end of the loop, we can remove the header and turn it into
+                // a normal branch.
                 auto targetBlock = loop->getTargetBlock();
-                if (targetBlock->getPredecessors().getCount() == 1 && *targetBlock->getPredecessors().begin() == block)
+                if (isTrivialSingleIterationLoop(func, loop, simplificationContext))
                 {
                     builder.setInsertBefore(loop);
                     List<IRInst*> args;
@@ -50,6 +145,7 @@ bool processFunc(IRGlobalValueWithCode* func)
                     loop->removeAndDeallocate();
                 }
             }
+
             // If `block` does not end with an unconditional branch, bail.
             if (block->getTerminator()->getOp() != kIROp_unconditionalBranch)
                 break;
diff --git a/source/slang/slang-ir-single-return.cpp b/source/slang/slang-ir-single-return.cpp
index a00066556..f76e35040 100644
--- a/source/slang/slang-ir-single-return.cpp
+++ b/source/slang/slang-ir-single-return.cpp
@@ -54,7 +54,8 @@ struct SingleReturnContext : public InstPassBase
         {
             loopHeaderBlock->addParam(param);
         }
-        auto loopInst = (IRLoop*)builder.emitLoop(originalStartBlock, breakBlock, originalStartBlock);
+
+        builder.emitLoop(originalStartBlock, breakBlock, originalStartBlock);
 
         // Now replace all return insts as break insts.
         processChildInstsOfType<IRReturn>(kIROp_Return, func, [&](IRReturn* returnInst)
@@ -81,16 +82,6 @@ struct SingleReturnContext : public InstPassBase
             builder.emitReturn(retValParam);
         else
             builder.emitReturn();
-
-        // Now run the multi-level-break pass.
-        eliminateMultiLevelBreakForFunc(module, func);
-
-        // Now remove the trivial loop header.
-        SLANG_RELEASE_ASSERT(loopInst->getContinueBlock() == loopInst->getTargetBlock());
-        auto targetBlock = loopInst->getTargetBlock();
-        for (auto param : params)
-            targetBlock->addParam(param);
-        loopHeaderBlock->removeAndDeallocate();
     }
 };
 
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index f9686ac5b..9d538a774 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -159,8 +159,9 @@ namespace Slang
 
             v->firstUse = this;
         }
-
+#ifdef SLANG_ENABLE_FULL_IR_VALIDATION
         debugValidate();
+#endif
     }
 
     void IRUse::set(IRInst* uv)
@@ -172,13 +173,15 @@ namespace Slang
     {
         // This `IRUse` is part of the linked list
         // of uses for  `usedValue`.
-
+#ifdef SLANG_ENABLE_FULL_IR_VALIDATION
         debugValidate();
+#endif
 
         if (usedValue)
         {
+#ifdef SLANG_ENABLE_FULL_IR_VALIDATION
             auto uv = usedValue;
-
+#endif
             *prevLink = nextUse;
             if(nextUse)
             {
@@ -190,8 +193,11 @@ namespace Slang
             nextUse     = nullptr;
             prevLink    = nullptr;
 
+#ifdef SLANG_ENABLE_FULL_IR_VALIDATION
             if(uv->firstUse)
                 uv->firstUse->debugValidate();
+#endif
+
         }
     }
 
diff --git a/source/slang/slang.natvis b/source/slang/slang.natvis
index b38bd358f..38244f9eb 100644
--- a/source/slang/slang.natvis
+++ b/source/slang/slang.natvis
@@ -72,10 +72,11 @@
     </Expand>
   </Type>
   <Type Name="Slang::IRInst">
-    <DisplayString>{{{m_op} #{_debugUID}}}</DisplayString>
-    <Expand>
+      <DisplayString>{{{m_op} {(uint32_t)(void*)this, x}}}</DisplayString>
+      <DisplayString Optional="true">{{{m_op} #{_debugUID}}}</DisplayString>
+      <Expand>
       <Item Name="[op]">m_op</Item>
-      <Item Name="[UID]">_debugUID</Item>
+      <Item Name="[UID]"  Optional="true">_debugUID</Item>
       <Item Name="[type]">typeUse.usedValue</Item>
       <CustomListItems MaxItemsPerView="3">
 		  <Variable Name="child" InitialValue="m_decorationsAndChildren.first"/>
@@ -83,20 +84,20 @@
 			  <If Condition="child == 0">
 				  <Break/>
 			  </If>
-			  <If Condition="child->m_op == SlangDebug__IROpNameHint">
-				  <Item Name="[name]">((Slang::IRStringLit*)(((Slang::IRUse*)(child + 1))->usedValue))->value.stringVal.chars,[((Slang::IRStringLit*)(((Slang::IRUse*)(child + 1))->usedValue))->value.stringVal.numChars]s8</Item>
-			  </If>
-			  <If Condition="child->m_op == SlangDebug__IROpExport">
+			  <If Condition="child->m_op == kIROp_ExportDecoration">
 				  <Item Name="[exportName]">((Slang::IRStringLit*)(((Slang::IRUse*)(child + 1))->usedValue))->value.stringVal.chars,[((Slang::IRStringLit*)(((Slang::IRUse*)(child + 1))->usedValue))->value.stringVal.numChars]s8</Item>
 			  </If>
-			  <If Condition="child->m_op == SlangDebug__IROpImport">
+			  <If Condition="child->m_op == kIROp_ImportDecoration">
 				  <Item Name="[importName]">((Slang::IRStringLit*)(((Slang::IRUse*)(child + 1))->usedValue))->value.stringVal.chars,[((Slang::IRStringLit*)(((Slang::IRUse*)(child + 1))->usedValue))->value.stringVal.numChars]s8</Item>
 			  </If>
+			  <If Condition="child->m_op == kIROp_NameHintDecoration">
+				  <Item Name="[name]">((Slang::IRStringLit*)(((Slang::IRUse*)(child + 1))->usedValue))->value.stringVal.chars,[((Slang::IRStringLit*)(((Slang::IRUse*)(child + 1))->usedValue))->value.stringVal.numChars]s8</Item>
+			  </If>
 			  <Exec>child = child->next</Exec>
 		  </Loop>
       </CustomListItems>
-      <Item Name="[value]" Condition="m_op == SlangDebug__IROpStringLit">((IRStringLit*)this)->value.stringVal.chars,[((IRStringLit*)this)->value.stringVal.numChars]s8</Item>
-      <Item Name="[value]" Condition="m_op == SlangDebug__IROpIntLit">((IRIntLit*)this)->value.intVal</Item>
+      <Item Name="[value]" Condition="m_op == kIROp_StringLit">((IRStringLit*)this)->value.stringVal.chars,[((IRStringLit*)this)->value.stringVal.numChars]s8</Item>
+      <Item Name="[value]" Condition="m_op == kIROp_IntLit">((IRIntLit*)this)->value.intVal</Item>
 	  <!--
       <Synthetic Name="[operands]">
         <DisplayString>{{count = {operandCount}}}</DisplayString>
@@ -121,14 +122,14 @@
 			<Exec>child = pOperandInst->m_decorationsAndChildren.first</Exec>
 		    <Exec>nameDecoration = 0</Exec>
 		    <Loop Condition="child != 0">
-			    <If Condition="child->m_op == SlangDebug__IROpNameHint">
+			    <If Condition="child->m_op == kIROp_NameHintDecoration">
 				    <Exec>nameDecoration = child</Exec>
 				    <Break/>
 			    </If>
-			    <If Condition="child->m_op == SlangDebug__IROpExport &amp;&amp; (nameDecoration == 0 || nameDecoration->m_op != SlangDebug__IROpNameHint)">
+			    <If Condition="child->m_op == kIROp_ExportDecoration &amp;&amp; (nameDecoration == 0 || nameDecoration->m_op != kIROp_NameHintDecoration)">
 				    <Exec>nameDecoration = child</Exec>
 			    </If>
-			    <If Condition="child->m_op == SlangDebug__IROpImport &amp;&amp; (nameDecoration == 0 || nameDecoration->m_op != SlangDebug__IROpNameHint)">
+			    <If Condition="child->m_op == kIROp_ImportDecoration &amp;&amp; (nameDecoration == 0 || nameDecoration->m_op != kIROp_NameHintDecoration)">
 				    <Exec>nameDecoration = child</Exec>
 			    </If>
 			    <Exec>child = child->next</Exec>
@@ -150,14 +151,14 @@
 				  <Exec>child = pItem->m_decorationsAndChildren.first </Exec>
 			      <Exec>nameDecoration = 0</Exec>
 				  <Loop Condition="child != 0">
-				      <If Condition="child->m_op == SlangDebug__IROpNameHint">
+				      <If Condition="child->m_op == kIROp_NameHintDecoration">
 						  <Exec>nameDecoration = child</Exec>
 						  <Break/>
 				      </If>
-				      <If Condition="child->m_op == SlangDebug__IROpExport &amp;&amp; (nameDecoration == 0 || nameDecoration->m_op != SlangDebug__IROpNameHint)">
+				      <If Condition="child->m_op == kIROp_ExportDecoration &amp;&amp; (nameDecoration == 0 || nameDecoration->m_op != kIROp_NameHintDecoration)">
 						  <Exec>nameDecoration = child</Exec>
 				      </If>
-					  <If Condition="child->m_op == SlangDebug__IROpImport &amp;&amp; (nameDecoration == 0 || nameDecoration->m_op != SlangDebug__IROpNameHint)">
+					  <If Condition="child->m_op == kIROp_ImportDecoration &amp;&amp; (nameDecoration == 0 || nameDecoration->m_op != kIROp_NameHintDecoration)">
 						  <Exec>nameDecoration = child</Exec>
 					  </If>
 				      <Exec>child = child->next</Exec>