// 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. InstWorkList workList; InstHashSet workListSet; WorkListPass(IRModule* inModule) : module(inModule), workList(inModule), workListSet(inModule), toBeDeleted(inModule), toBeDeletedSet(inModule) {} 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. InstWorkList toBeDeleted; InstHashSet 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 { BindingQueryLoweringContext(IRModule* inModule) : WorkListPass(inModule) {} // 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(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 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(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()) { if (auto layout = as(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(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(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()) { 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(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(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 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(module); context.sink = sink; context.processModule(); } }