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<BigStruct> 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<BigStruct_std430>;
func computeMain:
   %tmpVar : var<BigStruct>
    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<BigStruct>
   %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<BigStruct_std430>
   %barr = fieldAddr(%v, "values")
   %elementPtr : ptr<int> = 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>

1 files changed, 25 insertions, 21 deletions

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<IRFuncType>(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<IRPtrType>(paramType))
+            if (auto ptrType = as<IRPtrType>(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<IRPtrTypeBase>(arg->getDataType());
             if (!as<IRPtrTypeBase>(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<IRConstRefType>(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();