From a6deb5ed82cb8fc6b4f4c5c5fee264e09f97ff89 Mon Sep 17 00:00:00 2001 From: Yong He Date: Mon, 29 Sep 2025 17:45:08 -0700 Subject: Rewriting the lower-buffer-element-type pass to avoid unnecessary packing/unpacking. (#8526) Part of the effort to improve the performance of generated SPIRV code. The existing lower-buffer-element-type pass works by loading the entire buffer element content from memory, and translate it to logical type stored in a local variable at the earliest reference of a buffer handle. This means that is can generate inefficient code that reads more than necessary. Consider this example: ``` struct BigStruct { bool values[1024]; } ConstantBuffer cb; void test(BigStruct v) { if (v.values[0]) { printf("ok"); } } [numthreads(1,1,1)] void computeMain() { test(cb); } ``` In IR, the `computeMain` function before lower-buffer-element-type pass is something like following: ``` func test: %v = param : BigStruct %barr = fieldExtract(%v, "values") %element = elementExtract(%barr, 0) ... // uses %element func computeMain: %v = load(cb) call %test %v ``` The existing lower-buffer-element-type pass will rewrite the bool array in `BigStruct` into `int` array so it is legal in SPIRV. However, it does so by inserting the translation on the first `load` of the constant buffer: ``` struct BigStruct_std430 { int values[1024]; } var cb : ConstantBuffer; func computeMain: %tmpVar : var call %unpackStorage(%tmpVar, cb) %v : BigStruct = load %tmpVar call %test %v ``` This means that the entire array will be loaded and translated to int, before calling `test`, which only uses one element. It turns out that the downstream compiler isn't always able to optimize out this inefficient translation/copy. This PR completely rewrites the way buffer-element-type lowering is handled to avoid producing this inefficient code. It works in two parts: first we turn on the `transformParamsToConstRef` pass for SPIRV target as well, so we will translate the `test` function to take the `v` parameter as `constref`. The second part is a redesigned buffer-element-type pass that defers the storage-type to logical-type translation until a value is actually used by a `load` instruction. In this example, after `transformParamsToConstRef`, the IR is: ``` func test: %v = param : ConstRef %barr = fieldAddr(%v, "values") %elementPtr = elementAddr(%barr, 0) %element = load(%elementPtr) ... // uses %element func computeMain: call %test %cb ``` The new `buffer-element-type-lowering` pass will take this IR, and insert translation at latest possible time across the entire call graph, and translate the IR into: ``` func test: %v = param : ConstRef %barr = fieldAddr(%v, "values") %elementPtr : ptr = elementAddr(%barr, 0) %element_int = load(%elementPtr) %element = cast(%element_int) : %bool ... // uses %element func computeMain: call %test %cb ``` In this new IR, there is no longer a load and conversion of the entire array. See new comment in `slang-ir-lower-buffer-element-type.cpp` for more details of how the pass works. This PR also address many other issues surfaced by turning on `transformParamsToConstRef` pass on SPIRV backend. --------- Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com> --- source/slang/slang-ir-spirv-legalize.cpp | 46 +++++++++++++++++--------------- 1 file changed, 25 insertions(+), 21 deletions(-) (limited to 'source/slang/slang-ir-spirv-legalize.cpp') diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp index 2c4bd11cc..f795a6559 100644 --- a/source/slang/slang-ir-spirv-legalize.cpp +++ b/source/slang/slang-ir-spirv-legalize.cpp @@ -893,18 +893,22 @@ struct SPIRVLegalizationContext : public SourceEmitterBase IRBuilder builder(inst); builder.setInsertBefore(inst); auto funcType = as(funcValue->getFullType()); + bool argsChanged = false; for (UInt i = 0; i < inst->getArgCount(); i++) { auto arg = inst->getArg(i); auto paramType = funcType->getParamType(i); - if (as(paramType)) + if (auto ptrType = as(paramType)) { - // If the parameter has an explicit pointer type, - // then we know the user is using the variable pointer - // capability to pass a true pointer. - // In this case we should not rewrite the call. - newArgs.add(arg); - continue; + if (ptrType->getAddressSpace() == AddressSpace::UserPointer) + { + // If the parameter has an explicit pointer type, + // then we know the user is using the variable pointer + // capability to pass a true pointer. + // In this case we should not rewrite the call. + newArgs.add(arg); + continue; + } } auto ptrType = as(arg->getDataType()); if (!as(arg->getDataType())) @@ -953,13 +957,26 @@ struct SPIRVLegalizationContext : public SourceEmitterBase // If we reach here, we need to allocate a temp var. auto tempVar = builder.emitVar(ptrType->getValueType()); + builder.addDecoration(tempVar, kIROp_DisableCopyEliminationDecoration); auto load = builder.emitLoad(arg); builder.emitStore(tempVar, load); newArgs.add(tempVar); + argsChanged = true; + + // We may need to write the value back to the original pointer argument + // after the call. + // + // If callee doesn't modify the memory location, no need to write back. + if (funcType && funcType->getParamCount() > i && + as(funcType->getParamType(i))) + continue; + // If the buffer location is immutable, don't write back. + if (isPointerToImmutableLocation(root)) + continue; writeBacks.add(WriteBackPair{arg, tempVar}); } SLANG_ASSERT((UInt)newArgs.getCount() == inst->getArgCount()); - if (writeBacks.getCount()) + if (argsChanged) { auto newCall = builder.emitCallInst(inst->getFullType(), inst->getCallee(), newArgs); for (auto wb : writeBacks) @@ -2297,19 +2314,6 @@ struct SPIRVLegalizationContext : public SourceEmitterBase // so we need to update the function types to match that. updateFunctionTypes(); - // Lower all loads/stores from buffer pointers to use correct storage types. - // We didn't do the lowering for buffer pointers because we don't know which pointer - // types are actual storage buffer pointers until we propagated the address space of - // pointers in this pass. In the future we should consider separate out IRAddress as - // the type for IRVar, and use IRPtrType to dedicate pointers in user code, so we can - // safely lower the pointer load stores early together with other buffer types. - BufferElementTypeLoweringOptions bufferElementTypeLoweringOptions; - bufferElementTypeLoweringOptions.lowerBufferPointer = true; - lowerBufferElementTypeToStorageType( - m_sharedContext->m_targetProgram, - m_module, - bufferElementTypeLoweringOptions); - // Look for structs that are both used as fields and marked with Block // decorations, and move the Block decoration to a wrapper struct. legalizeStructBlocks(); -- cgit v1.2.3