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, 80 insertions, 30 deletions

diff --git a/source/slang/slang-ir-glsl-legalize.cpp b/source/slang/slang-ir-glsl-legalize.cpp
index a2f56cf7d..a79ca2379 100644
--- a/source/slang/slang-ir-glsl-legalize.cpp
+++ b/source/slang/slang-ir-glsl-legalize.cpp
@@ -1079,6 +1079,12 @@ IRInst* getOrCreateBuiltinParamForHullShader(
             if (sysAttr->getName().caseInsensitiveEquals(builtinSemantic))
             {
                 outputControlPointIdParam = param;
+                if (as<IRPtrTypeBase>(outputControlPointIdParam->getDataType()))
+                {
+                    IRBuilder builder(param);
+                    setInsertAfterOrdinaryInst(&builder, param);
+                    outputControlPointIdParam = builder.emitLoad(param);
+                }
                 break;
             }
         }
@@ -2348,11 +2354,11 @@ ScalarizedVal getSubscriptVal(
             auto inputAdapter = val.impl.as<ScalarizedTypeAdapterValImpl>();
             RefPtr<ScalarizedTypeAdapterValImpl> resultAdapter = new ScalarizedTypeAdapterValImpl();
 
-            resultAdapter->pretendType = inputAdapter->pretendType;
-            resultAdapter->actualType = inputAdapter->actualType;
+            resultAdapter->pretendType = elementType;
+            resultAdapter->actualType = getElementType(*builder, inputAdapter->actualType);
 
             resultAdapter->val =
-                getSubscriptVal(builder, inputAdapter->actualType, inputAdapter->val, indexVal);
+                getSubscriptVal(builder, resultAdapter->actualType, inputAdapter->val, indexVal);
             return ScalarizedVal::typeAdapter(resultAdapter);
         }
 
@@ -3127,7 +3133,7 @@ void tryReplaceUsesOfStageInput(
                 {
                     auto user = use->getUser();
                     IRBuilder builder(user);
-                    builder.setInsertBefore(user);
+                    setInsertBeforeOrdinaryInst(&builder, user);
                     builder.replaceOperand(use, val.irValue);
                 });
         }
@@ -3155,7 +3161,7 @@ void tryReplaceUsesOfStageInput(
                         return;
                     }
                     IRBuilder builder(user);
-                    builder.setInsertBefore(user);
+                    setInsertBeforeOrdinaryInst(&builder, user);
                     if (needMaterialize)
                     {
                         auto materializedVal = materializeValue(&builder, val);
@@ -3176,22 +3182,50 @@ void tryReplaceUsesOfStageInput(
                 {
                     auto user = use->getUser();
                     IRBuilder builder(user);
-                    builder.setInsertBefore(user);
+                    setInsertBeforeOrdinaryInst(&builder, user);
                     auto typeAdapter = as<ScalarizedTypeAdapterValImpl>(val.impl);
-                    auto materializedInner = materializeValue(&builder, typeAdapter->val);
-                    auto adapted = adaptType(
-                        &builder,
-                        materializedInner,
-                        typeAdapter->pretendType,
-                        typeAdapter->actualType);
-                    if (user->getOp() == kIROp_Load)
-                    {
-                        user->replaceUsesWith(adapted.irValue);
-                        user->removeAndDeallocate();
-                    }
-                    else
+                    switch (user->getOp())
                     {
-                        use->set(adapted.irValue);
+                    case kIROp_Load:
+                        {
+                            auto materialized = materializeValue(&builder, val);
+                            user->replaceUsesWith(materialized);
+                            user->removeAndDeallocate();
+                        }
+                        break;
+                    case kIROp_GetElementPtr:
+                        {
+                            auto targetType = typeAdapter->pretendType;
+                            auto elementType = getElementType(builder, targetType);
+                            SLANG_ASSERT(elementType);
+                            auto subscriptVal = getSubscriptVal(
+                                &builder,
+                                (IRType*)elementType,
+                                val,
+                                user->getOperand(1));
+                            tryReplaceUsesOfStageInput(context, subscriptVal, user);
+                        }
+                        break;
+                    case kIROp_FieldAddress:
+                        {
+                            auto targetType = as<IRStructType>(typeAdapter->pretendType);
+                            SLANG_ASSERT(targetType);
+                            auto subscriptVal = extractField(
+                                &builder,
+                                val,
+                                kMaxUInt,
+                                (IRStructKey*)user->getOperand(1));
+                            tryReplaceUsesOfStageInput(context, subscriptVal, user);
+                        }
+                        break;
+                    default:
+                        {
+                            auto materialized = materializeValue(&builder, val);
+                            auto tmpVar = builder.emitVar(materialized->getDataType());
+                            builder.emitStore(tmpVar, materialized);
+                            use->set(tmpVar);
+                        }
+                        break;
                     }
                 });
         }
@@ -3205,13 +3239,13 @@ void tryReplaceUsesOfStageInput(
                     auto arrayIndexImpl = as<ScalarizedArrayIndexValImpl>(val.impl);
                     auto user = use->getUser();
                     IRBuilder builder(user);
-                    builder.setInsertBefore(user);
+                    setInsertBeforeOrdinaryInst(&builder, user);
                     auto subscriptVal = getSubscriptVal(
                         &builder,
                         arrayIndexImpl->elementType,
                         arrayIndexImpl->arrayVal,
                         arrayIndexImpl->index);
-                    builder.setInsertBefore(user);
+                    setInsertBeforeOrdinaryInst(&builder, user);
                     auto materializedInner = materializeValue(&builder, subscriptVal);
                     if (user->getOp() == kIROp_Load)
                     {
@@ -3220,7 +3254,9 @@ void tryReplaceUsesOfStageInput(
                     }
                     else
                     {
-                        use->set(materializedInner);
+                        auto tmpVar = builder.emitVar(materializedInner->getDataType());
+                        builder.emitStore(tmpVar, materializedInner);
+                        use->set(tmpVar);
                     }
                 });
             break;
@@ -3233,6 +3269,9 @@ void tryReplaceUsesOfStageInput(
                 [&](IRUse* use)
                 {
                     auto user = use->getUser();
+                    IRBuilder builder(user);
+                    setInsertBeforeOrdinaryInst(&builder, user);
+
                     switch (user->getOp())
                     {
                     case kIROp_FieldExtract:
@@ -3270,10 +3309,20 @@ void tryReplaceUsesOfStageInput(
                             }
                         }
                         break;
+                    case kIROp_GetElementPtr:
+                        {
+                            auto arrayType = as<IRArrayTypeBase>(tupleVal->type);
+                            SLANG_ASSERT(arrayType);
+                            auto subscriptVal = getSubscriptVal(
+                                &builder,
+                                (IRType*)arrayType->getElementType(),
+                                val,
+                                user->getOperand(1));
+                            tryReplaceUsesOfStageInput(context, subscriptVal, user);
+                        }
+                        break;
                     case kIROp_Load:
                         {
-                            IRBuilder builder(user);
-                            builder.setInsertBefore(user);
                             auto materializedVal = materializeTupleValue(&builder, val);
                             user->replaceUsesWith(materializedVal);
                             user->removeAndDeallocate();
@@ -3449,7 +3498,7 @@ void legalizeEntryPointParameterForGLSL(
 
                     // Okay, we have a declaration, and we want to modify it!
 
-                    builder->setInsertBefore(ii);
+                    setInsertBeforeOrdinaryInst(builder, ii);
 
                     assign(builder, globalOutputVal, ScalarizedVal::value(ii->getOperand(2)));
                 }
@@ -3768,12 +3817,13 @@ void legalizeEntryPointParameterForGLSL(
                                 blockToMaterialized.tryGetValue(callingBlock, materialized);
                             if (!found)
                             {
-                                replaceBuilder.setInsertBefore(callingBlock->getFirstInst());
+                                replaceBuilder.setInsertBefore(
+                                    callingBlock->getFirstOrdinaryInst());
                                 materialized = materializeValue(&replaceBuilder, globalValue);
                                 blockToMaterialized.set(callingBlock, materialized);
                             }
 
-                            replaceBuilder.setInsertBefore(user);
+                            setInsertBeforeOrdinaryInst(builder, user);
                             auto field =
                                 replaceBuilder.emitFieldExtract(globalVarType, materialized, key);
                             replaceBuilder.replaceOperand(operandUse, field);
@@ -3888,7 +3938,7 @@ void assignRayPayloadHitObjectAttributeLocations(IRModule* module)
                 {
                     rayPayloadCounter++;
                 }
-                builder.setInsertBefore(inst);
+                setInsertBeforeOrdinaryInst(&builder, inst);
                 location = builder.getIntValue(builder.getIntType(), rayPayloadCounter);
                 decor->setOperand(0, location);
                 rayPayloadCounter++;
@@ -3902,7 +3952,7 @@ void assignRayPayloadHitObjectAttributeLocations(IRModule* module)
                 {
                     callablePayloadCounter++;
                 }
-                builder.setInsertBefore(inst);
+                setInsertBeforeOrdinaryInst(&builder, inst);
                 location = builder.getIntValue(builder.getIntType(), callablePayloadCounter);
                 decor->setOperand(0, location);
                 callablePayloadCounter++;
@@ -3915,7 +3965,7 @@ void assignRayPayloadHitObjectAttributeLocations(IRModule* module)
                 {
                     hitObjectAttributeCounter++;
                 }
-                builder.setInsertBefore(inst);
+                setInsertBeforeOrdinaryInst(&builder, inst);
                 location = builder.getIntValue(builder.getIntType(), hitObjectAttributeCounter);
                 decor->setOperand(0, location);
                 hitObjectAttributeCounter++;