Add a basc inlining facility for use in the stdlib (#1271)

The main feature visible to the stdlib here is the `[__unsafeForceInlineEarly]` attribute, which can be attached to a function definition and forces calls to that function to be inlined immediately after initial IR lowering.

The pass is implemented in `slang-ir-inline.{h,cpp}` and currently only handles the completely trivial case of a function with no control flow that ends with a single `return`. The lack of support for any other cases motivates the `__unsafe` prefix on the attribute.

In order to test that the pass works, I modified the "comma operator" in the standard library to be defined directly (rather than relying on special-case handling in IR lowering), and then added a test that uses that operator to make sure it generates code as expected. The compute version of the test confirms that we generate semantically correct code for the operator, while the SPIR-V cross-compilation test confirms that our output matches GLSL where the comma operator has been inlined, rather than turned into a subroutine.

Notes for the future:

* With this change it should be possible (in principle) to redefine all the compound operators in the stdlib to instead be ordinary functions with the new attribute, removing the need for `slang-compound-intrinsics.h`.

* Once the compound intrinsics are defined in the stdlib, it should be easier/possible to start making built-in operators like `+` be ordinary functions from the standpoint of the IR

* The attribute and pass here could be extended to include an alternative inlining attribute that happens later in compilation (after linking) but otherwise works the same. This could in theory be used for functions where we don't want to inline the definition into generated IR, but still want to inline things berfore generating final HlSL/GLSL/whatever.

* The inlining pass itself could be generalized to work for less trivial functions pretty easily; for the most part it would just mean "splitting" the block with the call site and then inserting clones of the blocks in the callee. Any `return` instructions in the clone would become unconditional branches (with arguments) to the block after the call (which would get a parameter to represent the returned value).

* The "hard" part for such an inlining pass would be handling cases where the control flow that results from inlining can't be handled by our later restructuring passes. The long-term fix there is to implement something like the "relooper" algorithm to restructure control flow as required for specific targets.

1 files changed, 451 insertions, 0 deletions

diff --git a/source/slang/slang-ir-inline.cpp b/source/slang/slang-ir-inline.cpp
new file mode 100644
index 000000000..19083c09b
--- /dev/null
+++ b/source/slang/slang-ir-inline.cpp
@@ -0,0 +1,451 @@
+// slang-ir-inline.cpp
+#include "slang-ir-inline.h"
+
+// This file provides general facilities for inlining function calls.
+
+//
+// A  *call site* is an individual `call` instruction (`IRCall`), and the *callee*
+// for a given call site is whatever is being called. When the callee is a `func`
+// (`IRFunc`) or a specialization of a `generic` that yields a `func`, *and* the
+// function has a body, then inlinling is possible.
+//
+// Different inlining passes may apply different heuristics or rules to decide
+// which call sites should be inlined (if possible). The rules may be based
+// on user-supplied hints, or on optimization criteria like performance and
+// code size.
+
+#include "slang-ir.h"
+#include "slang-ir-clone.h"
+#include "slang-ir-insts.h"
+
+namespace Slang
+{
+    /// Base type for inlining passes, providing shared/common functionality
+struct InliningPassBase
+{
+        /// The module that we are optimizing/transforming
+    IRModule* m_module = nullptr;
+
+        /// Initialize an inlining pass to operate on the given `module`
+    InliningPassBase(IRModule* module)
+        : m_module(module)
+    {
+    }
+
+        /// Consider all the call sites in the module for inliing
+    void considerAllCallSites()
+    {
+        considerAllCallSitesRec(m_module->getModuleInst());
+    }
+
+        /// Consider all call sites at or under `inst` for inlining
+    void considerAllCallSitesRec(IRInst* inst)
+    {
+        if( auto call = as<IRCall>(inst) )
+        {
+            considerCallSite(call);
+        }
+
+        // Note: we defensively iterate through the child instructions
+        // so that even if `child` gets removed (because of inlining)
+        // we automatically start at the next instruction after it.
+        //
+        IRInst* next = nullptr;
+        for( auto child = inst->getFirstChild(); child; child = next )
+        {
+            next = child->getNextInst();
+            considerAllCallSitesRec(child);
+        }
+    }
+
+    // In order to inline a call site, we need certain information
+    // to be present/available. Most notable is that the callee must
+    // be known, and it must be in the form of an `IRFunc`.
+    //
+    // Since checking whether we *can* inline a call site involves
+    // finding all of this information, we will use that opportunity
+    // to package it all up in a `struct` that can be re-used when
+    // we actually get around to inlining a call site.
+
+        /// Information about a call site to be inlined
+    struct CallSiteInfo
+    {
+            /// The call instruction.
+        IRCall* call = nullptr;
+
+            /// The function being called.
+            ///
+            /// For an inlinable call, this must be non-null and a valid function *definition* (with a body) for inlining to proceed.
+        IRFunc* callee = nullptr;
+
+            /// The specialization of the function, if any.
+            ///
+            /// For an inlineable call, this must be non-null if the function is generic, but may be null otherwise.
+        IRSpecialize* specialize = nullptr;
+
+            /// The generic being specialized.
+            ///
+            /// For an inlineable call, this must be be non-null if `specialize` is non-null.
+        IRGeneric* generic = nullptr;
+    };
+
+    // With `CallSiteInfo` defined, we can now understand the
+    // basic proces of considering a call site for inlining.
+
+        /// Consider the given `call` site, and possibly inline it.
+    void considerCallSite(IRCall* call)
+    {
+        // We start by checking if inlining would even be possible,
+        // since doing so collects information about the call site
+        // that can simplify the following steps.
+        //
+        // If the call can't be inlined, there is nothing else
+        // to consider and we bail out.
+        //
+        CallSiteInfo callSite;
+        if(!canInline(call, callSite))
+            return;
+
+        // If we've decided that we *can* inline the given call
+        // site, we next need to check if we *should*. The rules
+        // for when we should inline may vary by subclass,
+        // so `shouldInline` is a virtual method.
+        //
+        if(!shouldInline(callSite))
+            return;
+
+        // 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);
+    }
+
+    // Every subclas of `InliningPassBase` should provide its own
+    // definition of `shouldInline`. We define a default implementation
+    // here for the benefit of passes that might implement their
+    // own logic for deciding what to inline, bypassing `considerCallSite`.
+
+        /// Determine whether `callSite` should be inlined.
+    virtual bool shouldInline(CallSiteInfo const& callSite)
+    {
+        SLANG_UNUSED(callSite);
+        return false;
+    }
+
+        /// Determine whether `call` can be inlined, and if so write information about it to `outCallSite`
+    bool canInline(IRCall* call, CallSiteInfo& outCallSite)
+    {
+        // We can start by writing the `call` instruction into our `CallSiteInfo`.
+        //
+        outCallSite.call = call;
+
+        // Next we consider the callee.
+        //
+        IRInst* callee = call->getCallee();
+
+        // If the callee is a `specialize` instruction, then we
+        // want to look at what is being specialized instead.
+        //
+        if( auto specialize = as<IRSpecialize>(callee) )
+        {
+            // If the `specialize` is applied to something other
+            // than a `generic` instruction, then we can't
+            // inline the call site. This can happen for a
+            // call to a generic method in an interface.
+            //
+            IRGeneric* generic = findSpecializedGeneric(specialize);
+            if(!generic)
+                return false;
+
+            // If we have a `generic` instruction, then we
+            // will look to see if we can determine what
+            // it returns. If a result is found, that
+            // will be used as the new callee for this
+            // call site.
+            //
+            // If we can't identify the value that the generic
+            // yields, then inlining isn't possible.
+            //
+            callee = findGenericReturnVal(generic);
+            if(!callee)
+                return false;
+
+            // If we decide to inline this call, then the information
+            // we've just extracted about generic specialization
+            // will be relevant, so we write it to the `CallSiteInfo` now.
+            //
+            outCallSite.specialize = specialize;
+            outCallSite.generic = generic;
+        }
+
+        // Once we've dispensed with any possible generic specialization
+        // we will check if the callee is a `func` instruction (`IRFunc`).
+        //
+        // If it is not, then inlining isn't possible.
+        //
+        auto calleeFunc = as<IRFunc>(callee);
+        if(!calleeFunc)
+            return false;
+        //
+        // If the callee *is* a function, then we can update
+        // the `CalleSiteInfo` with what we've found.
+        //
+        outCallSite.callee = calleeFunc;
+
+        // At this point the `CallSiteInfo` is complete and
+        // could be used for inlining, but we have additional
+        // checks to make.
+        //
+        // In particular, we should only go about inlining
+        // a call site if the callee function is a full definition
+        // in the IR (not just a declaration).
+        //
+        if(!isDefinition(calleeFunc))
+            return false;
+
+        return true;
+    }
+
+        /// Inline the given `callSite`, which is assumed to have been validated
+    void inlineCallSite(CallSiteInfo const& callSite)
+    {
+        // Information about the call site, including
+        // the `call` instruction and the callee `func`
+        // should already have been computed and stored
+        // in the `CallSiteInfo`.
+        //
+        IRCall* call = callSite.call;
+        IRFunc* callee = callSite.callee;
+
+        // We will use the existing IR cloning infrastructure to clone
+        // the body of the callee, but we need to establish an
+        // environment for cloning in which any parameters of
+        // the callee are replaced with the matching arguments
+        // at the call site.
+        //
+        IRCloneEnv env;
+
+        // We also need an `IRBuilder` to construct the cloned IR,
+        // and will set it up to insert before the `call` that
+        // is going to be replaced.
+        //
+        SharedIRBuilder sharedBuilder;
+        sharedBuilder.session = m_module->getSession();
+        sharedBuilder.module = m_module;
+        IRBuilder builder;
+        builder.sharedBuilder = &sharedBuilder;
+        builder.setInsertBefore(call);
+
+        // If the callee is a generic function, then we will
+        // need to include the substitution of generic parameters
+        // with their argument values in our cloning.
+        //
+        if( auto specialize = callSite.specialize )
+        {
+            auto generic = callSite.generic;
+
+            // We start by establishing a mapping from the
+            // generic parameters to the matching arguments.
+            //
+            Int argCounter = 0;
+            for( auto param : generic->getParams() )
+            {
+                SLANG_ASSERT(argCounter < (Int)specialize->getArgCount());
+                auto arg = specialize->getArg(argCounter++);
+
+                env.mapOldValToNew.Add(param, arg);
+            }
+            SLANG_ASSERT(argCounter == (Int)specialize->getArgCount());
+
+            // We also need to clone any instructions in the
+            // body of the `generic` being specialized, since
+            // these might construct types or constants that
+            // reference the generic parameters.
+            //
+            auto body = generic->getFirstBlock();
+            SLANG_ASSERT(!body->getNextBlock()); // All IR generics should have a single block.
+
+            for( auto inst : body->getChildren() )
+            {
+                if( inst == callee )
+                {
+                    // We don't want to create a clone of the callee
+                    // function at the call site, since it would
+                    // immediately become dead code when we inline
+                    // its body.
+                }
+                else if(as<IRReturn>(inst))
+                {
+                    // We also don't want to clone any `return`
+                    // instruction in the generic, since that is
+                    // how they yield their result (which we
+                    // already know is `callee`.
+                }
+                else
+                {
+                    // In the default case, we just clone the instruction
+                    // from the body of the generic into the call site.
+                    //
+                    // TODO: This assumes that deduplication will work
+                    // as intended, so in practice we might run into
+                    // problems if we create new instances of IR types
+                    // or constants that already exist.
+                    //
+                    cloneInst(&env, &builder, inst);
+                }
+            }
+        }
+
+        // Compared to dealing with generic parameters, the process
+        // for dealing with value parameters is much simpler.
+        //
+        {
+            // For each parameter of the callee function, we
+            // insert a mapping into `env` from that parameter to the
+            // matching argument at the call site.
+            //
+            Int argCounter = 0;
+            for(auto param : callee->getParams())
+            {
+                SLANG_ASSERT(argCounter < (Int)call->getArgCount());
+                auto arg = call->getArg(argCounter++);
+                env.mapOldValToNew.Add(param, arg);
+            }
+            SLANG_ASSERT(argCounter == (Int)call->getArgCount());
+        }
+
+        // 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) )
+        {
+            inlineTrivialFuncBody(callSite, &env, &builder);
+        }
+        else
+        {
+            // Running into any non-trivial function to be inlined
+            // is currently an internal compiler error.
+            //
+            SLANG_UNIMPLEMENTED_X("general case of inlining");
+        }
+    }
+
+        /// Check if `func` represents a trivial single-block callee that can be inlined simply
+    bool isTrivialFunc(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;
+
+        return true;
+    }
+
+        /// 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)
+    {
+        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
+        // 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.
+        //
+        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() )
+        {
+            switch( inst->op )
+            {
+            default:
+                // The default value is to clone the instruction using
+                // the existing cloning infrastructure and the `env`
+                // we have already set up.
+                //
+                cloneInst(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;
+            }
+        }
+
+        // If there was a `returnVal` instruction that established
+        // the return value of the inlined function, then that value
+        // should be used to replace any uses of the original call.
+        //
+        if( returnedValue )
+        {
+            call->replaceUsesWith(returnedValue);
+        }
+
+        // Once we've cloned the body of the callee in at the call site,
+        // there is no reason to keep around the original `call` instruction,
+        // so we remove it.
+        //
+        call->removeAndDeallocate();
+    }
+};
+
+    /// An inlining pass that inlines calls to `[unsafeForceInlineEarly]` functions
+struct MandatoryEarlyInliningPass : InliningPassBase
+{
+    typedef InliningPassBase Super;
+
+    MandatoryEarlyInliningPass(IRModule* module)
+        : Super(module)
+    {}
+
+    bool shouldInline(CallSiteInfo const& info)
+    {
+        if(info.callee->findDecoration<IRUnsafeForceInlineEarlyDecoration>())
+            return true;
+        return false;
+    }
+};
+
+
+void performMandatoryEarlyInlining(IRModule* module)
+{
+    MandatoryEarlyInliningPass pass(module);
+    pass.considerAllCallSites();
+}
+
+} // namespace Slang