Refactor several IR passes (#761)

* Refactor several IR passes

This change takes some IR passes that lived together in `ir.cpp` and moves them into their own files to improve clarity.
In most cases these were passes introduced early in the life of the IR, so that it didn't seem like a big deal to have them all in one file, but now that `ir.cpp` has grown unwieldly this seems like an important cleanup to make.

To give a quick rundown of the passes involved:

* The IR "linking" step has been pulled out to `ir-link.{h,cpp}`. This code for this pass is pretty much identical to what was in `ir.cpp`, and no attempt has been made to clean up or refactor it in the current change.

* The GLSL legalization step has been pulled out to `ir-glsl-legalize.{h,cpp}`. This used to be invoked directly from the linking step, but has been made a new top-level pass invoked from `emit.cpp`. Just like with the linking, the code in the new file is just a copy-paste of what was in `ir.cpp`, and no attempt at cleanup has been made. Also note that it might be a good idea to move this pass later in the overall sequence, but this PR doesn't attempt to do that as it could change results.

* The generic specialization step has been pulled out to `ir-specialize.{h,cpp}`. The file name does not explicitly reference *generic* specialization because I anticipate this pass having to perform other kinds of specialization as well. The code in this case amounts to a heavy cleanup/refactoring pass and  thus deserves careful scrutiny. The reason for the cleanup is that the generic specialization step used to be part of the "linking" step long ago, and continued to share infrastructure with it long after that stopped making sense. The newly cleaned up pass has much simpler logic that should be easy enough to follow from the comments.

* In order to reduce code dulication, the IR "cloning" part of the `ir-specialize-resources.{h,cpp}` pass was pulled into its own files (`ir-clone.{h,cpp}`) that both the generic specialization step and the resource-based specialization step now share.

The remaining changes then pertain to deleting a bunch of code out of `ir.cpp` and adding the new files to the build.

The only test that needed updating was `vkray/raygen`, where some subtle ordering change in the refactored generic specialization logic has lead to the relative order of the specialized `TraceRay` and `saturate` functions beind reversed.

* fixup: typo in assert

* fixup: typos in comments

1 files changed, 685 insertions, 0 deletions

diff --git a/source/slang/ir-specialize.cpp b/source/slang/ir-specialize.cpp
new file mode 100644
index 000000000..4a20c6195
--- /dev/null
+++ b/source/slang/ir-specialize.cpp
@@ -0,0 +1,685 @@
+// ir-specialize.cpp
+#include "ir-specialize.h"
+
+#include "ir.h"
+#include "ir-clone.h"
+#include "ir-insts.h"
+
+namespace Slang
+{
+
+// This file implements the primary specialization pass, that takes
+// generic/polymorphic Slang code and specializes/monomorphises it.
+//
+// At present this primarily means generating specialized copies
+// of generic functions/types based on the concrete types used
+// at specialization sites, and also specializing instances
+// of witness-table lookup to directly refer to the concrete
+// values for witnesses when witness tables are known.
+//
+// Eventually, this pass will also need to perform specialization
+// of functions to argument values for parameters that must
+// be compile-time constants, and simplification of code using
+// existential (interface) types for function parameters/results.
+
+struct SpecializationContext
+{
+    // We know that we can only perform specialization when all
+    // of the arguments to a generic are also fully specialized.
+    // The "is fully specialized" condition is something we
+    // need to solve for over the program, because the fully-
+    // specialized-ness of an instruction depends on the
+    // fully-specialized-ness of its operands.
+    //
+    // We will build an explicit hash set to encode those
+    // instructions that are fully specialized.
+    //
+    HashSet<IRInst*> fullySpecializedInsts;
+
+    // An instruction is then fully specialized if and only
+    // if it is in our set.
+    //
+    bool isInstFullySpecialized(
+        IRInst*                     inst)
+    {
+        // A small wrinkle is that a null instruction pointer
+        // sometimes appears a a type, and so should be treated
+        // as fully specialized too.
+        //
+        // TODO: It would be nice to remove this wrinkle.
+        //
+        if(!inst) return true;
+
+        return fullySpecializedInsts.Contains(inst);
+    }
+
+    // When an instruction isn't fully specialized, but its operands *are*
+    // then it is a candidate for specialization itself, so we will have
+    // a query to check for the "all operands fully specialized" case.
+    //
+    bool areAllOperandsFullySpecialized(
+        IRInst*                     inst)
+    {
+        if(!isInstFullySpecialized(inst->getFullType()))
+            return false;
+
+        UInt operandCount = inst->getOperandCount();
+        for(UInt ii = 0; ii < operandCount; ++ii)
+        {
+            IRInst* operand = inst->getOperand(ii);
+            if(!isInstFullySpecialized(operand))
+                return false;
+        }
+
+        return true;
+    }
+
+    // We will also maintain a work list of instructions that are
+    // not fully specialized, and that we want to consider for
+    // specialization.
+    //
+    List<IRInst*> workList;
+
+    // When we consider adding an instruction to our work list
+    // we will try to be careful and only add things that aren't
+    // already fully specialized.
+    //
+    void maybeAddToWorkList(IRInst* inst)
+    {
+        if(isInstFullySpecialized(inst))
+            return;
+
+        workList.Add(inst);
+    }
+
+    // When we go to populate the work list by recursively
+    // traversing some code, we will be careful to *not*
+    // add generics or their children to the work list,
+    // and will instead consider a generic to be "fully
+    // specialized" already (because uses of that generic
+    // as an *operand* should be seen as fully specialized
+    // references).
+    //
+    void populateWorkListRec(
+        IRInst*                     inst)
+    {
+        if(auto genericInst = as<IRGeneric>(inst))
+        {
+            fullySpecializedInsts.Add(genericInst);
+        }
+        else
+        {
+            maybeAddToWorkList(inst);
+
+            for(auto child : inst->getChildren())
+            {
+                populateWorkListRec(child);
+            }
+        }
+    }
+
+    // Of course, somewhere along the way we expect
+    // to run into uses of `specialize(...)` instructions
+    // to bind a generic to arguments that we want to
+    // specialize into concrete code.
+    //
+    // We also know that if we encouter `specialize(g, a, b, c)`
+    // and then later `specialize(g, a, b, c)` again, we
+    // only want to generate the specialized code for `g<a,b,c>`
+    // *once*, and re-use it for both versions.
+    //
+    // We will cache existing specializations of generic function/types
+    // using the simple key type defined as part of the IR cloning infrastructure.
+    //
+    typedef IRSimpleSpecializationKey Key;
+    Dictionary<Key, IRInst*> genericSpecializations;
+
+    // We will also use some shared IR building state across
+    // all of our specialization/cloning steps.
+    //
+    SharedIRBuilder sharedBuilderStorage;
+
+    // Now let's look at the task of finding or generation a
+    // specialization of some generic `g`, given a specialization
+    // instruction like `specialize(g, a, b, c)`.
+    //
+    // The `specializeGeneric` function will return a value
+    // suitable for use as a replacement for the `specialize(...)`
+    // instruction.
+    //
+    IRInst* specializeGeneric(
+        IRGeneric*      genericVal,
+        IRSpecialize*   specializeInst)
+    {
+        // First, we want to see if an existing specialization
+        // has already been made. To do that we will construct a key
+        // for lookup in the generic specialization context.
+        //
+        // Our key will consist of the identity of the generic
+        // being specialized, and each of the argument values
+        // being pased to it. In our hypothetical example of
+        // `specialize(g, a, b, c)` the key will then be
+        // the array `[g, a, b, c]`.
+        //
+        Key key;
+        key.vals.Add(specializeInst->getBase());
+        UInt argCount = specializeInst->getArgCount();
+        for( UInt ii = 0; ii < argCount; ++ii )
+        {
+            key.vals.Add(specializeInst->getArg(ii));
+        }
+
+        {
+            // We use our generated key to look for an
+            // existing specialization that has been registered.
+            // If one is found, our work is done.
+            //
+            IRInst* specializedVal = nullptr;
+            if(genericSpecializations.TryGetValue(key, specializedVal))
+                return specializedVal;
+        }
+
+        // If no existing specialization is found, we need
+        // to create the specialization instead.
+        //
+        // Effectively this amounts to "calling" the generic
+        // on its concrete argument values and computing the
+        // result it returns.
+        //
+        // For now, all of our generics consist of a single
+        // basic block, so we can "call" them just by
+        // cloning the instructions in their single block
+        // into the global scope, using an environment for
+        // cloning that maps the generic parameters to
+        // the concrete arguments that were provided
+        // by the `specialize(...)` instruction.
+        //
+        IRCloneEnv      env;
+
+        // We will walk through the parameters of the generic and
+        // register the corresponding argument of the `specialize`
+        // instruction to be used as the "cloned" value for each
+        // parameter.
+        //
+        // Suppose we are looking at `specialize(g, a, b, c)` and `g` has
+        // three generic parameters: `T`, `U`, and `V`. Then we will
+        // be initializing our environment to map `T -> a`, `U -> b`,
+        // and `V -> c`.
+        //
+        UInt argCounter = 0;
+        for( auto param : genericVal->getParams() )
+        {
+            UInt argIndex = argCounter++;
+            SLANG_ASSERT(argIndex < specializeInst->getArgCount());
+
+            IRInst* arg = specializeInst->getArg(argIndex);
+
+            env.mapOldValToNew.Add(param, arg);
+        }
+
+        // We will set up an IR builder for insertion
+        // into the global scope, at the same location
+        // as the original generic.
+        //
+        IRBuilder builderStorage;
+        IRBuilder* builder = &builderStorage;
+        builder->sharedBuilder = &sharedBuilderStorage;
+        builder->setInsertBefore(genericVal);
+
+        // Now we will run through the body of the generic and
+        // clone each of its instructions into the global scope,
+        // until we reach a `return` instruction.
+        //
+        for( auto bb : genericVal->getBlocks() )
+        {
+            // We expect a generic to only ever contain a single block.
+            //
+            SLANG_ASSERT(bb == genericVal->getFirstBlock());
+
+            // We will iterate over the non-parameter ("ordinary")
+            // instructions only, because parameters were dealt
+            // with explictly at an earlier point.
+            //
+            for( auto ii : bb->getOrdinaryInsts() )
+            {
+                // The last block of the generic is expected to end with
+                // a `return` instruction for the specialized value that
+                // comes out of the abstraction.
+                //
+                // We thus use that cloned value as the result of the
+                // specialization step.
+                //
+                if( auto returnValInst = as<IRReturnVal>(ii) )
+                {
+                    auto specializedVal = findCloneForOperand(&env, returnValInst->getVal());
+
+                    // The value that was returned from evaluating
+                    // the generic is the specialized value, and we
+                    // need to remember it in our dictionary of
+                    // specializations so that we don't instantiate
+                    // this generic again for the same arguments.
+                    //
+                    genericSpecializations.Add(key, specializedVal);
+
+                    return specializedVal;
+                }
+
+                // For any instruction other than a `return`, we will
+                // simply clone it completely into the global scope.
+                //
+                IRInst* clonedInst = cloneInst(&env, builder, ii);
+
+                // Any new instructions we create during cloning were
+                // not present when we initially built our work list,
+                // so we need to make sure to consider them now.
+                //
+                // This is important for the cases where one generic
+                // invokes another, because there will be `specialize`
+                // operations nested inside the first generic that refer
+                // to the second.
+                //
+                populateWorkListRec(clonedInst);
+            }
+        }
+
+        // If we reach this point, something went wrong, because we
+        // never encountered a `return` inside the body of the generic.
+        //
+        SLANG_UNEXPECTED("no return from generic");
+        UNREACHABLE_RETURN(nullptr);
+    }
+
+    // The logic for generating a specialization of an IR generic
+    // relies on the ability to "evaluate" the code in the body of
+    // the generic, but that obviously doesn't work if we don't
+    // actually have the full definition for the body.
+    //
+    // This can arise in particular for builtin operations/types.
+    //
+    // Before calling `specializeGeneric()` we need to make sure
+    // that the generic is actually amenable to specialization,
+    // by looking at whether it is a definition or a declaration.
+    //
+    bool canSpecializeGeneric(
+        IRGeneric*  generic)
+    {
+        // It is possible to have multiple "layers" of generics
+        // (e.g., when a generic function is nested in a generic
+        // type). Therefore we need to drill down through all
+        // of the layers present to see if at the leaf we have
+        // something that looks like a definition.
+        //
+        IRGeneric* g = generic;
+        for(;;)
+        {
+            // Given the generic `g`, we will find the value
+            // it appears to return in its body.
+            //
+            auto val = findGenericReturnVal(g);
+            if(!val)
+                return false;
+
+            // If `g` returns an inner generic, then we need
+            // to drill down further.
+            //
+            if (auto nestedGeneric = as<IRGeneric>(val))
+            {
+                g = nestedGeneric;
+                continue;
+            }
+
+            // Once we've found the leaf value that will be produced
+            // after all specialization is complete, we can check
+            // whether it looks like a definition or not.
+            //
+            return isDefinition(val);
+        }
+    }
+
+    // Now that we know when we can specialize a generic, and how
+    // to do it, we can write a subroutine that takes a
+    // `specialize(g, a, b, c, ...)` instruction and performs
+    // specialization if it is possible.
+    //
+    void maybeSpecializeGeneric(
+        IRSpecialize* specInst)
+    {
+        // The invariant that the arguments are fully specialized
+        // should mean that `a, b, c, ...` are in a form that
+        // we can work with, but it does *not* guarantee
+        // that the `g` operand is something we can work with.
+        //
+        // We can only perform specialization in the case where
+        // the base `g` is a known `generic` instruction.
+        //
+        auto baseVal = specInst->getBase();
+        auto genericVal = as<IRGeneric>(baseVal);
+        if(!genericVal)
+            return;
+
+        // We can also only specialize a generic if it
+        // represents a definition rather than a declaration.
+        //
+        if(!canSpecializeGeneric(genericVal))
+            return;
+
+        // Once we know that specialization is possible,
+        // the actual work is fairly simple.
+        //
+        // First, we find or generate a specialized
+        // version of the result of the generic (a specialized
+        // type, function, or whatever).
+        //
+        auto specializedVal = specializeGeneric(genericVal, specInst);
+
+        // Then we simply replace any uses of the `specialize(...)`
+        // instruction with the specialized value and delete
+        // the `specialize(...)` instruction from existence.
+        //
+        specInst->replaceUsesWith(specializedVal);
+        specInst->removeAndDeallocate();
+    }
+
+    // The basic rule we are following is that once all the operands
+    // to an instruction are fully specialized, we are safe
+    // to specialize the instruction itself, but the work
+    // required to specialize an instruction depends on the
+    // form of the instruction.
+    //
+    void fullySpecializeInst(
+        IRInst*                     inst)
+    {
+        // A precondition of the `fullySpecializeInst` operation
+        // is that the operands to `inst` have all been fully
+        // specialized.
+        //
+        SLANG_ASSERT(areAllOperandsFullySpecialized(inst));
+
+        switch(inst->op)
+        {
+        default:
+            // The default case is that there is nothing to
+            // be done to specialize an instruction; once all
+            // of its operands are specialized it is safe
+            // to consider the instruction itself as fully
+            // specialized.
+            //
+            break;
+
+        case kIROp_Specialize:
+            // The logic for specializing a `specialize(...)`
+            // instruction has already been elaborated above.
+            //
+            maybeSpecializeGeneric(cast<IRSpecialize>(inst));
+            break;
+
+        case kIROp_lookup_interface_method:
+            // The remaining case we need to consider here
+            // is when we have a `lookup_witness_method` instruction
+            // that is being applied to a concrete witness table,
+            // because we can specialize it to just be a direct
+            // reference to the actual witness value from the table.
+            //
+            maybeSpecializeWitnessLookup(cast<IRLookupWitnessMethod>(inst));
+            break;
+        }
+    }
+
+    void maybeSpecializeWitnessLookup(
+        IRLookupWitnessMethod* lookupInst)
+    {
+        // Note: While we currently have named the instruction
+        // `lookup_witness_method`, the `method` part is a misnomer
+        // and the same instruction can look up *any* interface
+        // requirement based on the witness table that provides
+        // a conformance, and the "key" that indicates the interface
+        // requirement.
+
+        // We can only specialize in the case where the lookup
+        // is being done on a concrete witness table, and not
+        // the result of a `specialize` instruction or other
+        // operation that will yield such a table.
+        //
+        auto witnessTable = as<IRWitnessTable>(lookupInst->getWitnessTable());
+        if(!witnessTable)
+            return;
+
+        // Because we have a concrete witness table, we can
+        // use it to look up the IR value that satisfies
+        // the given interface requirement.
+        //
+        auto requirementKey = lookupInst->getRequirementKey();
+        auto satisfyingVal = findWitnessVal(witnessTable, requirementKey);
+
+        // We expect to always find a satisfying value, but
+        // we will go ahead and code defensively so that
+        // we leave "correct" but unspecialized code if
+        // we cannot find a concrete value to use.
+        //
+        if(!satisfyingVal)
+            return;
+
+        // At this point, we know that `satisfyingVal` is what
+        // would result from executing this `lookup_witness_method`
+        // instruciton dynamically, so we can go ahead and
+        // replace the original instruction with that value.
+        //
+        lookupInst->replaceUsesWith(satisfyingVal);
+        lookupInst->removeAndDeallocate();
+    }
+
+    // The above subroutine needed a way to look up
+    // the satisfying value for a given requirement
+    // key in a concrete witness table, so let's
+    // define that now.
+    //
+    IRInst* findWitnessVal(
+        IRWitnessTable* witnessTable,
+        IRInst*         requirementKey)
+    {
+        // A witness table is basically just a container
+        // for key-value pairs, and so the best we can
+        // do for now is a naive linear search.
+        //
+        for( auto entry : witnessTable->getEntries() )
+        {
+            if (requirementKey == entry->getRequirementKey())
+            {
+                return entry->getSatisfyingVal();
+            }
+        }
+        return nullptr;
+    }
+
+    // All of the machinery has been defined above, so
+    // we can now walk through the flow of the overall
+    // specialization pass.
+    //
+    void processModule(IRModule* module)
+    {
+        // We start by initializing our shared IR building state,
+        // since we will re-use that state for any code we
+        // generate along the way.
+        //
+        SharedIRBuilder* sharedBuilder = &sharedBuilderStorage;
+        sharedBuilder->module = module;
+        sharedBuilder->session = module->session;
+
+        // The unspecialized IR we receive as input will have
+        // `IRBindGlobalGenericParam` instructions that associate
+        // each global-scope generic parameter (a type, witness
+        // table, or what-have-you) with the value that it should
+        // be bound to for the purposes of this code-generation
+        // pass.
+        //
+        // Before doing any other specialization work, we will
+        // iterate over these instructions (which may only
+        // appear at the global scope) and use them to drive
+        // replacement of the given generic type parameter with
+        // the desired concrete value.
+        //
+        auto moduleInst = module->getModuleInst();
+        for(auto inst : moduleInst->getChildren())
+        {
+            // We only want to consider the `bind_global_generic_param`
+            // instructions, and ignore everything else.
+            //
+            auto bindInst = as<IRBindGlobalGenericParam>(inst);
+            if(!bindInst)
+                continue;
+
+            // HACK: Our current front-end emit logic can end up emitting multiple
+            // `bind_global_generic_param` instructions for the same parameter. This is
+            // a buggy behavior, but a real fix would require refactoring the way
+            // global generic arguments are specified today.
+            //
+            // For now we will do a sanity check to detect parameters that
+            // have already been specialized.
+            //
+            if( !as<IRGlobalGenericParam>(bindInst->getOperand(0)) )
+            {
+                // The "parameter" operand is no longer a parameter, so it
+                // seems things must have been specialized already.
+                //
+                continue;
+            }
+
+            // The actual logic for applying the substitution is
+            // almost trivial: we will replace any uses of the
+            // global generic parameter with its desired value.
+            //
+            auto param = bindInst->getParam();
+            auto val = bindInst->getVal();
+            param->replaceUsesWith(val);
+        }
+        {
+            // Now that we've replaced any uses of global generic
+            // parameters, we will do a second pass to remove
+            // the parameters and any `bind_global_generic_param`
+            // instructions, since both should be dead/unused.
+            //
+            IRInst* next = nullptr;
+            for(auto inst = moduleInst->getFirstChild(); inst; inst = next)
+            {
+                next = inst->getNextInst();
+
+                switch(inst->op)
+                {
+                default:
+                    break;
+
+                case kIROp_GlobalGenericParam:
+                case kIROp_BindGlobalGenericParam:
+                    // A `bind_global_generic_param` instruction should
+                    // have no uses in the first place, and all the global
+                    // generic parameters should have had their uses replaced.
+                    //
+                    SLANG_ASSERT(!inst->firstUse);
+                    inst->removeAndDeallocate();
+                    break;
+                }
+            }
+        }
+
+        // Now that we've eliminated all cases of global generic parameters,
+        // we should now have the properties that:
+        //
+        // 1. Execution starts in non-generic code, with no unbound
+        //    generic parameters in scope.
+        //
+        // 2. Any case where non-generic code makes use of a generic
+        //    type/function, there will be a `specialize` instruction
+        //    that specifies both the generic and the (concrete) type
+        //    arguments that should be provided to it.
+        //
+        // Our primary goal is then to find `specialize` instructions that
+        // can be replaced with references to, e.g., a suitably
+        // specialized function, and to resolve any `lookup_interface_method`
+        // instructions to the concrete value fetched from a witness
+        // table.
+        //
+        // We need to be careful of a few things:
+        //
+        // * It would not in general make sense to consider specialize-able
+        //   instructions under an `IRGeneric`, since that could mean "specializing"
+        //   code to parameter values that are still unknown.
+        //
+        // * We *also* need to be careful not to specialize something when one
+        //   or more of its inputs is also a `specialize` or `lookup_interface_method`
+        //   instruction, because then we'd be propagating through non-concrete
+        //   values.
+        //
+        // The approach we use here is to build a work list of instructions
+        // that are candidates for specialization, and then process them one
+        // at a time to see if we can make some forward progress.
+        //
+        // We will start by recursively walking all the instructions to add
+        // the appropriate ones to  our work list:
+        //
+        populateWorkListRec(moduleInst);
+
+        // We want to treat our work list like a queue rather than
+        // a stack, because we populated it in program order,
+        // and fully-specialized-ness will tend to flow top-down.
+        //
+        // To accomplish this we will ping-pong between the
+        // real work list and a copy so that we can iterate over
+        // one list while adding to the other.
+        //
+        List<IRInst*> workListCopy;
+        while(workList.Count() != 0)
+        {
+            workListCopy.Clear();
+            workListCopy.SwapWith(workList);
+
+            for( auto inst : workListCopy )
+            {
+                // We need to check whether it is possible to specialize
+                // the instruction yet. It might not be specializable
+                // because its operands haven't been specialized.
+                //
+                if(!areAllOperandsFullySpecialized(inst))
+                {
+                    // If we can't fully specialize this instruction
+                    // yet, then we need to toss it back onto the
+                    // work list to be considered in the next round.
+                    //
+                    // TODO: We need to carefully vet that this can
+                    // never lead to infinite looping
+                    //
+                    workList.Add(inst);
+                }
+                else
+                {
+                    // If all of the operands to `inst` are
+                    // fully specialized, then we can go
+                    // ahead and do whatever is required
+                    // to "fully specialize" `inst` itself.
+                    //
+                    fullySpecializeInst(inst);
+
+                    // At this point, we want to start
+                    // considering `inst` as fully specialized,
+                    // so let's add it to our set.
+                    //
+                    fullySpecializedInsts.Add(inst);
+                }
+            }
+        }
+
+        // Once the work list has gone dry, we should have the invariant
+        // that there are no `specialize` instructions inside of non-generic
+        // functions that in turn reference a generic type/function, *except*
+        // in the case where that generic is for a builtin type/function, in
+        // which case we wouldn't want to specialize it anyway.
+    }
+};
+
+void specializeGenerics(
+    IRModule*   module)
+{
+    SpecializationContext context;
+    context.processModule(module);
+}
+
+} // namespace Slang