// 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 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 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(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` // *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 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(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(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(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(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(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(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(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(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 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