Improved SCCP, inlining and resource specialization passes, legalize `ImageSubscript` for GLSL (#2146)

1 files changed, 211 insertions, 73 deletions

diff --git a/source/slang/slang-ir-inline.cpp b/source/slang/slang-ir-inline.cpp
index a8347438b..ed4b34875 100644
--- a/source/slang/slang-ir-inline.cpp
+++ b/source/slang/slang-ir-inline.cpp
@@ -1,6 +1,8 @@
 // slang-ir-inline.cpp
 #include "slang-ir-inline.h"
 
+#include "slang-ir-ssa-simplification.h"
+
 // This file provides general facilities for inlining function calls.
 
 //
@@ -33,17 +35,18 @@ struct InliningPassBase
     }
 
         /// Consider all the call sites in the module for inliing
-    void considerAllCallSites()
+    bool considerAllCallSites()
     {
-        considerAllCallSitesRec(m_module->getModuleInst());
+        return considerAllCallSitesRec(m_module->getModuleInst());
     }
 
         /// Consider all call sites at or under `inst` for inlining
-    void considerAllCallSitesRec(IRInst* inst)
+    bool considerAllCallSitesRec(IRInst* inst)
     {
+        bool changed = false;
         if( auto call = as<IRCall>(inst) )
         {
-            considerCallSite(call);
+            changed = considerCallSite(call);
         }
 
         // Note: we defensively iterate through the child instructions
@@ -54,8 +57,9 @@ struct InliningPassBase
         for( auto child = inst->getFirstChild(); child; child = next )
         {
             next = child->getNextInst();
-            considerAllCallSitesRec(child);
+            changed |= considerAllCallSitesRec(child);
         }
+        return changed;
     }
 
     // In order to inline a call site, we need certain information
@@ -93,7 +97,7 @@ struct InliningPassBase
     // basic proces of considering a call site for inlining.
 
         /// Consider the given `call` site, and possibly inline it.
-    void considerCallSite(IRCall* call)
+    bool considerCallSite(IRCall* call)
     {
         // We start by checking if inlining would even be possible,
         // since doing so collects information about the call site
@@ -104,7 +108,7 @@ struct InliningPassBase
         //
         CallSiteInfo callSite;
         if(!canInline(call, callSite))
-            return;
+            return false;
 
         // If we've decided that we *can* inline the given call
         // site, we next need to check if we *should*. The rules
@@ -112,13 +116,14 @@ struct InliningPassBase
         // so `shouldInline` is a virtual method.
         //
         if(!shouldInline(callSite))
-            return;
+            return false;
 
         // Finally, if we both *can* and *should* inline the
         // given call site, we hand off the a worker routine
         // that does the meat of the work.
         //
         inlineCallSite(callSite);
+        return true;
     }
 
     // Every subclas of `InliningPassBase` should provide its own
@@ -313,12 +318,11 @@ struct InliningPassBase
         }
 
         // For now, our inlining pass only handles the case where
-        // the callee is a "trivial" function, which can support
-        // a very simple approach to inlining.
-        //
-        if( isTrivialFunc(callee) )
+        // the callee is a "single-return" function, which means the callee
+        // function contains only one return at the end of the body.
+        if (isSingleReturnFunc(callee))
         {
-            inlineTrivialFuncBody(callSite, &env, &builder);
+            inlineSingleReturnFuncBody(callSite, &env, &builder);    
         }
         else
         {
@@ -329,26 +333,34 @@ struct InliningPassBase
         }
     }
 
-        /// Check if `func` represents a trivial single-block callee that can be inlined simply
-    bool isTrivialFunc(IRFunc* func)
+        /// Check if `func` represents a simple callee that has only a single `return`.
+    bool isSingleReturnFunc(IRFunc* func)
     {
-        // The function must have a single bocy block to be trivial.
-        //
         auto firstBlock = func->getFirstBlock();
-        if( firstBlock->getNextBlock() )
-            return false;
 
         // If the body block is decorated (for some reason), then the function is non-trivial.
         //
         if( firstBlock->getFirstDecoration() )
             return false;
 
-        // If the body block terminates in something other than a `return` then the function is non-trivial.
-        //
-        auto terminator = firstBlock->getTerminator();
-        if( !as<IRReturn>(terminator) )
-            return false;
-
+        // If the body has more than one returns, we cannot inline it now.
+        bool returnFound = false;
+        for (auto block : func->getBlocks())
+        {
+            for (auto inst : block->getChildren())
+            {
+                if (inst->getOp() == kIROp_ReturnVal || inst->getOp() == kIROp_ReturnVoid)
+                {
+                    // If the return is not at the end of the block, we cannot handle it.
+                    if (inst != block->getTerminator())
+                        return false;
+                    // If there is already a return found, this function cannot be simple.
+                    if (returnFound)
+                        return false;
+                    returnFound = true;
+                }
+            }
+        }
         return true;
     }
 
@@ -374,13 +386,8 @@ struct InliningPassBase
         // serialization.
         // 
         // For now this punts on this, and just assumes [__unsafeForceInlineEarly] is not in user code.
-    static IRInst* _cloneInstWithSourceLoc(CallSiteInfo const& callSite,
-        IRCloneEnv*     env,
-        IRBuilder*      builder,
-        IRInst*         inst)
+    static void _setSourceLoc(IRInst* clonedInst, IRInst* srcInst, CallSiteInfo const& callSite)
     {
-        IRInst* clonedInst = cloneInst(env, builder, inst);
-
         SourceLoc sourceLoc;
 
         if (callSite.call->sourceLoc.isValid())
@@ -388,67 +395,129 @@ struct InliningPassBase
             // Default to using the source loc at the call site
             sourceLoc = callSite.call->sourceLoc;
         }
-        else if (inst->sourceLoc.isValid())
+        else if (srcInst->sourceLoc.isValid())
         {
             // If we don't have that copy the inst being cloned sourceLoc
-            sourceLoc = inst->sourceLoc;
+            sourceLoc = srcInst->sourceLoc;
         }
 
         clonedInst->sourceLoc = sourceLoc;
+    }
+    static IRInst* _cloneInstWithSourceLoc(CallSiteInfo const& callSite,
+        IRCloneEnv*     env,
+        IRBuilder*      builder,
+        IRInst*         inst)
+    {
+        IRInst* clonedInst = cloneInst(env, builder, inst);
+        _setSourceLoc(clonedInst, inst, callSite);
         return clonedInst;
     }
 
-        /// Inline the body of the callee for `callSite`, where the callee is trivial as tested by `isTrivialFunc`
-    void inlineTrivialFuncBody(CallSiteInfo const& callSite, IRCloneEnv* env, IRBuilder* builder)
+        /// Inline the body of the callee for `callSite`, where the callee has a single return.
+    void inlineSingleReturnFuncBody(
+        CallSiteInfo const& callSite, IRCloneEnv* env, IRBuilder* builder)
     {
         auto call = callSite.call;
         auto callee = callSite.callee;
-        auto firstBlock = callee->getFirstBlock();
 
-        // We know that the callee has a single block, so if we encounter
+        // 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 block, and its operand will be the value
-        // the calee returns.
+        // return point for the function, and its operand will be the value
+        // the callee returns.
         //
         IRInst* returnedValue = nullptr;
 
-        // We will loop over the instructions of the one and only block,
-        // and clone each of them appropriately.
-        //
-        for( auto inst : firstBlock->getChildren() )
+        // Break the basic block containing the call inst into two basic blocks.
+        auto callerBlock = callSite.call->getParent();
+        builder->setInsertInto(callerBlock->getParent());
+        auto afterBlock = builder->createBlock();
+        
+        // Many operations (e.g. `cloneInst`) has define-before-use assumptions on the IR.
+        // It is important to make sure we keep the ordering of blocks by inserting the
+        // second half of the basic block right after `callerBlock`.
+        afterBlock->insertAfter(callerBlock);
+        afterBlock->sourceLoc = callSite.call->getNextInst()->sourceLoc;
+        // Move all insts after the call in `callerBlock` to `afterBlock`.
         {
-            switch( inst->getOp() )
+            auto inst = callSite.call->getNextInst();
+            while (inst)
             {
-            default:
-                // The default value is to clone the instruction using
-                // the existing cloning infrastructure and the `env`
-                // we have already set up.
-                //
-                // SourceLoc information is copied if there is appropriate data available.
-                _cloneInstWithSourceLoc(callSite, env, builder, inst);
-                break;
-
-            case kIROp_Param:
-                // Parameters can be completely ignored in the single-block
-                // case, because they have all been replaced via `env`.
-                break;
-
-            case kIROp_ReturnVoid:
-                // A return with no operand can be ignored, since a return
-                // from the inlined call should just continue after the
-                // call site.
-                //
-                break;
-
-            case kIROp_ReturnVal:
-                // A return with a value is similar to `returnVoid` except
-                // that we need to note the (clone of the) value being
-                // returned, so that we can use it to replace the value
-                // of the original call.
-                //
-                returnedValue = findCloneForOperand(env, inst->getOperand(0));
-                break;
+                auto next = inst->getNextInst();
+                inst->removeFromParent();
+                inst->insertAtEnd(afterBlock);
+                inst = next;
+            }
+        }
+
+        List<IRBlock*> clonedBlocks;
+        for (auto calleeBlock : callee->getBlocks())
+        {
+            auto clonedBlock = builder->createBlock();
+            clonedBlock->insertBefore(afterBlock);
+            _setSourceLoc(clonedBlock, calleeBlock, callSite);
+            env->mapOldValToNew[calleeBlock] = clonedBlock;
+        }
+
+        // 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()));
+        _setSourceLoc(newBranch, call, callSite);
+        // Clone all basic blocks over to the call site.
+        bool isFirstBlock = true;
+        for (auto calleeBlock : callee->getBlocks())
+        {
+            auto clonedBlock = env->mapOldValToNew[calleeBlock].GetValue();
+            builder->setInsertInto(clonedBlock);
+            // We will loop over the instructions of the each block,
+            // and clone each of them appropriately.
+            //
+            for (auto inst : calleeBlock->getChildren())
+            {
+                if (inst->getOp() == kIROp_Param)
+                {
+                    // Parameters in the first block can be completely ignored
+                    // because they have all been replaced via `env`.
+                    if (isFirstBlock)
+                    {
+                        continue;
+                    }
+                }
+
+                switch (inst->getOp())
+                {
+                default:
+                    // The default value is to clone the instruction using
+                    // the existing cloning infrastructure and the `env`
+                    // we have already set up.
+                    //
+                    // SourceLoc information is copied if there is appropriate data available.
+                    _cloneInstWithSourceLoc(callSite, env, builder, inst);
+                    break;
+
+                case kIROp_ReturnVoid:
+                    // A return with no operand is replaced with a branch into `afterBlock`
+                    // to return the control flow to the location after the original `call`.
+                    {
+                        auto returnBranch = builder->emitBranch(afterBlock);
+                        _setSourceLoc(returnBranch, inst, callSite);
+                    }
+                    break;
+
+                case kIROp_ReturnVal:
+                    // A return with a value is similar to `returnVoid` except
+                    // that we need to note the (clone of the) value being
+                    // returned, so that we can use it to replace the value
+                    // of the original call.
+                    //
+                    {
+                        auto returnBranch = builder->emitBranch(afterBlock);
+                        _setSourceLoc(returnBranch, inst, callSite);
+                        returnedValue = findCloneForOperand(env, inst->getOperand(0));
+                    }
+                    break;
+                }
             }
+            isFirstBlock = false;
         }
 
         // If there was a `returnVal` instruction that established
@@ -492,4 +561,73 @@ void performMandatoryEarlyInlining(IRModule* module)
     pass.considerAllCallSites();
 }
 
+
+    // Defined in slang-ir-specialize-resource.cpp
+bool isResourceType(IRType* type);
+bool isIllegalGLSLParameterType(IRType* type);
+
+    /// An inlining pass that inlines calls functions that returns resources.
+    /// This is needed for glsl targets.
+struct GLSLResourceReturnFunctionInliningPass : InliningPassBase
+{
+    typedef InliningPassBase Super;
+
+    GLSLResourceReturnFunctionInliningPass(IRModule* module)
+        : Super(module)
+    {}
+
+    bool shouldInline(CallSiteInfo const& info)
+    {
+        if (isResourceType(info.callee->getResultType()))
+        {
+            return true;
+        }
+        for (auto param : info.callee->getParams())
+        {
+            if (isIllegalGLSLParameterType(param->getDataType()))
+                return true;
+            auto outType = as<IROutTypeBase>(param->getDataType());
+            if (!outType)
+                continue;
+            auto outValueType = outType->getValueType();
+            if (isResourceType(outValueType))
+                return true;
+        }
+        return false;
+    }
+};
+
+void performGLSLResourceReturnFunctionInlining(IRModule* module)
+{
+    GLSLResourceReturnFunctionInliningPass pass(module);
+    bool changed = true;
+
+    while (changed)
+    {
+        changed = pass.considerAllCallSites();
+        simplifyIR(module);
+    }
+}
+
+struct CustomInliningPass : InliningPassBase
+{
+    typedef InliningPassBase Super;
+
+    CustomInliningPass(IRModule* module)
+        : Super(module)
+    {}
+
+    bool shouldInline(CallSiteInfo const&)
+    {
+        return true;
+    }
+};
+
+bool inlineCall(IRCall* call)
+{
+    CustomInliningPass pass(call->getModule());
+    return pass.considerCallSite(call);
+}
+
+
 } // namespace Slang