diff options
Diffstat (limited to 'source/slang/slang-ir-lower-binding-query.cpp')
| -rw-r--r-- | source/slang/slang-ir-lower-binding-query.cpp | 610 |
1 files changed, 610 insertions, 0 deletions
diff --git a/source/slang/slang-ir-lower-binding-query.cpp b/source/slang/slang-ir-lower-binding-query.cpp new file mode 100644 index 000000000..d0128d365 --- /dev/null +++ b/source/slang/slang-ir-lower-binding-query.cpp @@ -0,0 +1,610 @@ +// slang-ir-lower-tuple-types.cpp + +#include "slang-ir-lower-tuple-types.h" +#include "slang-ir.h" +#include "slang-ir-insts.h" + +// The pass in this file lowers the `getRegisterIndex()` and +// `getSpaceIndex()` intrinsics, by replacing them with literal +// values derived from the binding information on shader parameters. +// +// If these operations are applied to a global shader parameter, +// then we can simply read the binding information from that parameter +// and use it directly. +// +// Otherwise, we expect that the opaque object (resource/sampler/etc.) +// being referenced was passed down into the current function from +// a caller. We thus introduce new function parameters after the +// resource in question, transforming, e.g., this: +// +// void doThings( +// float a, +// Texture2D t, +// float b ) +// { +// ... __getRegisterIndex(t) ... +// ... __getSpaceIndex(t) ... +// } +// ... +// doThings(myTexture); +// +// into this: +// +// void doThings( +// float a, +// Texture2D t, +// uint t_index, +// uint t_space, +// float b ) +// { +// ... t_index ... +// ... t_space ... +// } +// ... +// doThings(myTexture, __getRegisterIndex(myTexture), __getRegisterSpace(myTexture)); +// +// At that point we have removed the invocations of `getRegisterIndex` +// and `getRegisterSpace` in the callee function, but introduced new +// invocations in the caller function, so we need to iterate until +// we eventually either bottom out at a global shader parameter, or +// run into a context that we cannot simplify. + +namespace Slang +{ + // There are a ton of passes we've implemented now that use + // some basic work-list structures, and it seems a bit silly + // to be writing that code intermixed with the actual algorithm. + // + // For this file, we break the common work-list functionality + // out into a base type that we can re-use in specific passes. + // + struct WorkListPass + { + public: + IRModule* module; + DiagnosticSink* sink; + + protected: + + // The base type needs to abstract over how the + // concrete pass will process each instruction + // that gets placed into the work list. + + virtual void processInst(IRInst* inst) = 0; + + // Otherwise, the implementation of the work list + // itself is straightforward, and not anything + // that hasn't been seen in other files. + + List<IRInst*> workList; + HashSet<IRInst*> workListSet; + + void addToWorkList(IRInst* inst) + { + if (workListSet.contains(inst)) + return; + + workList.add(inst); + workListSet.add(inst); + } + + void processWorkList() + { + while (workList.getCount() != 0) + { + IRInst* inst = workList.getLast(); + + workList.removeLast(); + workListSet.remove(inst); + + processInst(inst); + + for (auto child = inst->getLastChild(); child; child = child->getPrevInst()) + { + addToWorkList(child); + } + } + } + + // As long as we are factoring out repeated cruft, + // it seems reasonable to *also* deal with the + // frequent need to buffer up instructions to + // be deleted when a pass is complete. + + List<IRInst*> toBeDeleted; + HashSet<IRInst*> toBeDeletedSet; + + void addToBeDeleted(IRInst* inst) + { + if (toBeDeletedSet.contains(inst)) + return; + + toBeDeleted.add(inst); + toBeDeletedSet.add(inst); + } + + void processDeletions() + { + for (auto inst : toBeDeleted) + { + inst->removeAndDeallocate(); + } + toBeDeleted.clear(); + toBeDeletedSet.clear(); + } + }; + + // The concrete pass will then be a specialization of + // the base work-list abstraction. + // + struct BindingQueryLoweringContext : public WorkListPass + { + // All of the intrinsics we will be processing use + // the same result type (`uint`), so it is helpful + // to cache a pointer to the IR type at the start + // of the pass and re-use it. + // + IRType* indexType = nullptr; + + void processModule() + { + IRBuilder builder(module); + indexType = builder.getUIntType(); + + // Processing the module consists of recursively + // processing all the instructions in one pass, + // and then potentially revisiting instructions + // that had new intrinsics added to their bodies. + // + addToWorkList(module->getModuleInst()); + processWorkList(); + } + + void processInst(IRInst* inst) + { + // For this pass, we really only care about + // our binding query instructions. + // + if (auto query = as<IRBindingQuery>(inst)) + { + processQueryInst(query); + } + } + + void processQueryInst(IRBindingQuery* inst) + { + // Processing one of the query instructions is conceptually + // simple: we find a compute a value to replace it with, + // and then simply *replace* the instruction. + // + auto replacementValue = findOrComputeReplacementValueFor(inst); + if (!replacementValue) + { + // If we cannot find or compute a replacement value, + // then we need to treat it as an error, since the + // binding query intrinsics don't admit any reasonable + // runtime implementation. + // + sink->diagnose( + inst, + Diagnostics::opaqueReferenceMustResolveToGlobal); + return; + } + + inst->replaceUsesWith(replacementValue); + inst->removeAndDeallocate(); + } + + // We want to cache the results of computing the binding + // information for an opaque-type value, in case doing + // so required adding or modifying code. + // + // For that purpose, we introduce a simple data structure + // to hold the two pieces of binding information we + // care about. + // + struct OpaqueValueInfo + { + IRInst* registerIndex = nullptr; + IRInst* registerSpace = nullptr; + }; + + IRInst* findOrComputeReplacementValueFor(IRBindingQuery* query) + { + // Finding the replacement for a given query instruction + // then amounts to computing (or caching) the binding + // information for the opaque-type value it queries, + // and then projecting out the appropriate field. + + auto opaqueValue = query->getOpaqueValue(); + auto opaqueValueInfo = findOrComputeOpaqueValueInfo(opaqueValue); + + switch (query->getOp()) + { + default: + SLANG_UNEXPECTED("unhandled binding query instruction type"); + UNREACHABLE_RETURN(query); + + case kIROp_GetRegisterIndex: + return opaqueValueInfo.registerIndex; + + case kIROp_GetRegisterSpace: + return opaqueValueInfo.registerSpace; + } + } + + // The information will be cached in a dictionary, + // keyed on the opaque-type value that the information + // was computed for. + // + Dictionary<IRInst*, OpaqueValueInfo> mapOpaqueValueToInfo; + + // Looking up the cached information (if any) is a simple + // matter of using the dictionary. + // + // (We have a distinct operation for lookup vs. the + // memo-cached lookup below, because we may want to + // query this information while computing an entry, + // and we don't want to introduce potential recursion. + // + OpaqueValueInfo* findOpaqueValueInfo(IRInst* opaqueValue) + { + return mapOpaqueValueToInfo.tryGetValue(opaqueValue); + } + + OpaqueValueInfo findOrComputeOpaqueValueInfo(IRInst* opaqueValue) + { + if (auto foundInfo = findOpaqueValueInfo(opaqueValue)) + return *foundInfo; + + // If there is no information registered in the cache, we + // compute it on-demand. + // + // Note that there is no potential for circularity, so + // long as the implementation of `computeOpaqueValueInfo` + // does not itself call `findOrComputeValueInfo`. + // + auto computedInfo = computeOpaqueValueInfo(opaqueValue); + mapOpaqueValueToInfo.add(opaqueValue, computedInfo); + return computedInfo; + } + + // We are now (finally) getting into the meat of what this + // pass needs to do. Given an instruction with an opaque + // type, we need to try to compute the register and space + // it is bound to, or conspire to have that information + // passed along. + // + OpaqueValueInfo computeOpaqueValueInfo(IRInst* opaqueValue) + { + if (auto globalParam = as<IRGlobalParam>(opaqueValue)) + { + // The simple/base case is when we have a global shader + // parameter that has layout information attached. + // + // Note that this pass needs to run late enough that + // shader parameters declared at other scopes will have + // been massaged into the appropriate form. + // + if (auto layoutDecoration = globalParam->findDecoration<IRLayoutDecoration>()) + { + if (auto layout = as<IRVarLayout>(layoutDecoration->getLayout())) + { + // We expect any shader parameter of an opaque type + // to have a relevant resource kind, but it isn't + // too hard to code defensively. We will iterate + // over the resource kinds that are present and + // take the first one that represents an opaque type. + // + for (auto offsetAttr : layout->getOffsetAttrs()) + { + switch (offsetAttr->getResourceKind()) + { + default: + break; + + case LayoutResourceKind::ShaderResource: + case LayoutResourceKind::UnorderedAccess: + case LayoutResourceKind::ConstantBuffer: + case LayoutResourceKind::SamplerState: + case LayoutResourceKind::DescriptorTableSlot: + { + IRBuilder builder(module); + + OpaqueValueInfo info; + info.registerIndex = builder.getIntValue(indexType, offsetAttr->getOffset()); + info.registerSpace = builder.getIntValue(indexType, offsetAttr->getSpace()); + return info; + } + break; + } + } + } + } + } + else if (auto param = as<IRParam>(opaqueValue)) + { + // The other very interesting case is when the opaque-type + // value is an `IRParam`, which indicates that it is either + // a function parameter or a phi node of a basic block. + // + // Either way, we always expect a parameter to appear as + // a child of a block. + // + auto block = as<IRBlock>(param->getParent()); + SLANG_ASSERT(block); + + // When rewriting call sites, we will need to know the + // index of `param` within the parameter list. + // + Index paramIndex = -1; + { + Count paramCounter = 0; + for (auto p : block->getParams()) + { + Index i = paramCounter++; + if (p == param) + { + paramIndex = i; + break; + } + } + SLANG_ASSERT(paramIndex >= 0); + } + + // In either case (function parameter or block parameter), + // we will insert additional parameters after the original + // parameter, so that the register index and space can + // be passed along explicitly. + // + IRBuilder builder(module); + + // We create new parameters to pass along the register index/space, + // and manually insert them where we want them in the parameter list. + // + auto registerIndexParam = builder.createParam(builder.getUIntType()); + auto registerSpaceParam = builder.createParam(builder.getUIntType()); + // + registerSpaceParam->insertAfter(param); + registerIndexParam->insertAfter(param); + + // We would like for the newly-introduced parameters to have + // nice human-readable names, if the original parameter did. + // + if (auto nameHintDecoration = param->findDecoration<IRNameHintDecoration>()) + { + String hint; + hint.append(nameHintDecoration->getName()); + hint.append("."); + builder.addNameHintDecoration(registerIndexParam, (hint + "index").getUnownedSlice()); + builder.addNameHintDecoration(registerSpaceParam, (hint + "space").getUnownedSlice()); + } + + // Similarly, the new parameters should get debugging-related + // source location information from the original parameter, + // if it had any. + // + registerIndexParam->sourceLoc = param->sourceLoc; + registerSpaceParam->sourceLoc = param->sourceLoc; + + // Now we need to scan for the places that the function or block + // that the parameter belongs to gets referenced. At each such + // location, we will pass along arguments to match the additional + // parameters. + // + if (!block->getPrevBlock()) + { + // If this is the first block in the parent function, + // then this is a function parameter, and we will + // iterate over call sites of the function and rewrite + // them to pass along arguments for the new parameters. + // + auto func = block->getParent(); + + for (auto use = func->firstUse; use; use = use->nextUse) + { + auto user = use->getUser(); + if (auto call = as<IRCall>(user)) + { + if (call->getCallee() == func) + { + rewriteCall(call, paramIndex); + } + } + } + } + else + { + // If this is a block parameter, we will iterate over + // the instructions that branch to the block, and rewrite + // their argument lists, similar to what we do for function calls. + // + for (auto use = block->firstUse; use; use = use->nextUse) + { + auto user = use->getUser(); + if (auto branch = as<IRUnconditionalBranch>(user)) + { + if (branch->getTargetBlock() == block) + { + rewriteBranch(branch, paramIndex); + } + } + } + } + + // The new parameters that we introduced will be used to + // replace any binding query intrinsics applied to + // this opaque value. + // + OpaqueValueInfo info; + info.registerIndex = registerIndexParam; + info.registerSpace = registerSpaceParam; + return info; + } + + // By default we find that we cannot query binding information + // for the given instruction. + OpaqueValueInfo info; + return info; + } + + // In our IR, there isn't a lot of difference between a `call` + // and an `unconditionalBranch`; indeed, this is part of what + // motivates the use of `IRParam`s for both function parameters + // and phi nodes. + // + // However, while both blocks and functions use the same `IRParam` + // representation, we (currently) do not have a common base + // between the `call` and `unconditionalBranch` instructions. + // + // Rather than have duplicate logic between the two cases, we + // simply observe that for our purposes rewriting either a + // `call` or an `unconditionalBranch` amounts to doing + // special-case work on *one* operand of the original, while + // copying over all the other operands as-is. + // + // Given this observation, we can bottleneck both calls and + // branches into a common worker routine by passing down + // the instruction to be rewritten and a pointer to the + // `IRUse` for the one "interesting" operand. + + void rewriteCall(IRCall* oldCall, Index paramIndex) + { + rewriteCallOrBranch( + oldCall, + oldCall->getArgs() + paramIndex); + } + + void rewriteBranch(IRUnconditionalBranch* oldBranch, Index paramIndex) + { + rewriteCallOrBranch( + oldBranch, + oldBranch->getArgs() + paramIndex); + } + + void rewriteCallOrBranch( + IRInst* oldCallOrBranch, + IRUse* oldOperandToRewrite) + { + // Our goal here is to generate a new version of + // `oldCallOrBranch` that copies over most of the + // operands as-is, but introduces our rewrites + // around the chosen operand. + + IRBuilder builder(module); + builder.setInsertBefore(oldCallOrBranch); + + // We capture the old operand list as a range of + // `IRUse`s, and set up a fresh list to hold the + // new operands. + // + auto oldOperandsBegin = oldCallOrBranch->getOperands(); + auto oldOperandsEnd = oldOperandsBegin + oldCallOrBranch->getOperandCount(); + // + List<IRInst*> newOperands; + + // All of the operands that precede the interesting + // one can be copied over from the old list to the + // new one as-is. + // + for (auto u = oldOperandsBegin; u < oldOperandToRewrite; ++u) + { + auto operand = u->get(); + newOperands.add(operand); + } + + // Next we look at the value of the "intersting" + // operand, knowing that we need to pass along + // not only the original value but also the + // binding information. + // + IRInst* arg = oldOperandToRewrite->get(); + IRInst* registerIndex = nullptr; + IRInst* registerSpace = nullptr; + + // As a simple optimization, if we have *already* + // computed and cached binding information for + // the argument, we can re-use that information + // here and now. + // + if (auto info = findOpaqueValueInfo(arg)) + { + registerIndex = info->registerIndex; + registerSpace = info->registerSpace; + } + else + { + // If there is no cached information for + // the argument, we choose *not* to make + // a recursive call into `findOrComputeOpaqueValueInfo`. + // + // Instead we will simply emit additional + // binding query intrinsics into the body + // of the caller (right before the call site), + // and add those instructions to our work + // list, to be eliminated later. + // + registerIndex = builder.emitIntrinsicInst( + indexType, + kIROp_GetRegisterIndex, + 1, &arg); + registerSpace = builder.emitIntrinsicInst( + indexType, + kIROp_GetRegisterSpace, + 1, &arg); + // + addToWorkList(registerIndex); + addToWorkList(registerSpace); + } + + // Whether we have found existing binding information, + // or emitted new intrinsics, we are now ready + // to append the argument and its binding information + // to the new operand list. + // + newOperands.add(arg); + newOperands.add(registerIndex); + newOperands.add(registerSpace); + + // Any operands of the original instruction that come + // after the one we rewrite can be copied over as-is. + // + // Note: we don't currently have any operands that would + // appear after the arguments of a `call` or `branch`, + // but the fact that we encode `IRAttr`s on an instruction + // as additional (trailing) operands means that this could + // conceivably happen at some point. + // + for (auto u = oldOperandToRewrite + 1; u < oldOperandsEnd; ++u) + { + auto operand = u->get(); + newOperands.add(operand); + } + + // Once we've built up the new operand list, we can emit + // a new instruction that has the same opcode and type, + // with the new operands, and then use it to replace + // the existing instruction. + // + auto newCallOrBranch = builder.emitIntrinsicInst( + oldCallOrBranch->getFullType(), + oldCallOrBranch->getOp(), + newOperands.getCount(), + newOperands.getBuffer()); + + oldCallOrBranch->transferDecorationsTo(newCallOrBranch); + oldCallOrBranch->replaceUsesWith(newCallOrBranch); + oldCallOrBranch->removeAndDeallocate(); + } + }; + + void lowerBindingQueries( + IRModule* module, + DiagnosticSink* sink) + { + BindingQueryLoweringContext context; + context.module = module; + context.sink = sink; + context.processModule(); + } +} |