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, 21 insertions, 18 deletions

diff --git a/source/slang/slang-ir-specialize-address-space.cpp b/source/slang/slang-ir-specialize-address-space.cpp
index 29f1ec516..c4a155eec 100644
--- a/source/slang/slang-ir-specialize-address-space.cpp
+++ b/source/slang/slang-ir-specialize-address-space.cpp
@@ -168,6 +168,7 @@ struct AddressSpaceContext : public AddressSpaceSpecializationContext
                     {
                     case kIROp_Var:
                     case kIROp_RWStructuredBufferGetElementPtr:
+                    case kIROp_Load:
                         {
                             // The address space of these insts should be assigned by the initial
                             // address space assigner.
@@ -204,16 +205,6 @@ struct AddressSpaceContext : public AddressSpaceSpecializationContext
                             }
                         }
                         break;
-                    case kIROp_Load:
-                        {
-                            if (auto addrSpace =
-                                    mapVarValueToAddrSpace.tryGetValue(inst->getOperand(0)))
-                            {
-                                mapInstToAddrSpace[inst] = *addrSpace;
-                                changed = true;
-                            }
-                        }
-                        break;
                     case kIROp_Param:
                         if (!isFirstBlock)
                         {
@@ -248,22 +239,24 @@ struct AddressSpaceContext : public AddressSpaceSpecializationContext
                             if (callee)
                             {
                                 List<AddressSpace> argAddrSpaces;
-                                bool fullySpecialized = true;
+                                bool hasSpecializableArg = false;
                                 for (UInt i = 0; i < callInst->getArgCount(); i++)
                                 {
                                     auto arg = callInst->getArg(i);
-                                    auto argAddrSpace = getAddrSpace(arg);
                                     argAddrSpaces.add(getAddrSpace(arg));
-                                    if (argAddrSpace == AddressSpace::Generic &&
-                                        as<IRPtrTypeBase>(arg->getDataType()))
+                                    if (as<IRPtrTypeBase>(arg->getDataType()))
                                     {
-                                        fullySpecialized = false;
-                                        break;
+                                        hasSpecializableArg = true;
                                     }
                                 }
-                                if (!fullySpecialized)
+                                if (!hasSpecializableArg)
+                                {
+                                    workList.add(callee);
+                                    break;
+                                }
+                                // If callee doesn't have a body, don't specialize.
+                                if (!callee->getFirstBlock())
                                     break;
-
                                 FuncSpecializationKey key(callee, argAddrSpaces);
                                 IRFunc* specializedCallee = nullptr;
                                 if (IRFunc** specializedFunc =
@@ -484,4 +477,14 @@ void propagateAddressSpaceFromInsts(List<IRInst*>&& workList)
     }
 }
 
+AddressSpace NoOpInitialAddressSpaceAssigner::getAddressSpaceFromVarType(IRInst*)
+{
+    return AddressSpace::Generic;
+}
+
+AddressSpace NoOpInitialAddressSpaceAssigner::getLeafInstAddressSpace(IRInst*)
+{
+    return AddressSpace::Generic;
+}
+
 } // namespace Slang