Improve SSA promotion for arrays and structs (#521)

* Improve SSA promotion for arrays and structs

Fixes #518

The existing SSA pass would only handle `load(v)` and `store(v,...)`
where `v` is the variable instruction, and would bail out if `v` was
used as an operand in any other fashion.

The new pass adds support for `load(ac)` where `ac` is an "access chain"
with a gramar like:

    ac :: v
        | getElementPtr(ac, ...)
	| getFieldAddress(ac, ...)

What this means in practical terms is that we can promote a local
variable of array or structure type to an SSA temporary even if there
are loads of individual elements/fields, as along as any *assignment* to
the variable assigns the whole thing.

I've added a test case to confirm that this change fixes passing of
arrays as function parameters for Vulkan.

* Fixup: disable test on Vulkan because render-test isn't ready

This is a fix for Vulkan, but I don't think our testing setup is ready
for it.

* Fixup: error in unreachable return case, caught by clang

* Fixups based on testing

These are fixes found when testing the original changes against the user code that originated the bug report.

* `emit.cpp`: Make sure to handle array-of-texture types when deciding whether to declare a temporary as a local variable in GLSL output

* `ir-legalize-types.cpp`: Make a not of a source of validation failures that we need to clean up sooner or later (just not in scope for this bug fix change).

* `ir-ssa.cpp`:
  * When checking if something is an access chain with a promotable var at the end, make sure the recursive case recurses into the "access chain" logic instead of the leaf case
  * Add some assertions to guard the assumption that any access chain we apply has been scheduled for removal
  * Correctly emit an element *extract* instead of getting an element *address* when promoting an element access into an array being promoted
  * Eliminate a wrapper routine that was setting up an `IRBuilder` and use the one from the block being processed in the SSA pass (since it was set up for stuff just like this)

* `ir-validate.cpp`
  * Add a hack to avoid validation failures when running IR validation on the stdlib code. This case triggers for an initializer (`__init`) declaration inside an interface, since the logical "return type" is the interface type itself, which has no representation at the IR level and thus yields a null result type in a `FuncType` instruction.

6 files changed, 246 insertions, 10 deletions

diff --git a/source/slang/emit.cpp b/source/slang/emit.cpp
index 71d323c9e..be6112abe 100644
--- a/source/slang/emit.cpp
+++ b/source/slang/emit.cpp
@@ -2054,6 +2054,10 @@ struct EmitVisitor
         {
             type = ptrType->getValueType();
         }
+        while (auto ptrType = as<IRArrayTypeBase>(type))
+        {
+            type = ptrType->getElementType();
+        }
 
         if(as<IRUniformParameterGroupType>(type))
         {
diff --git a/source/slang/ir-legalize-types.cpp b/source/slang/ir-legalize-types.cpp
index a7dbf6672..25bd498b8 100644
--- a/source/slang/ir-legalize-types.cpp
+++ b/source/slang/ir-legalize-types.cpp
@@ -1329,6 +1329,11 @@ void legalizeTypes(
     // the case for `IRTypeLegalizationContext::replacedInstructions`,
     // but we haven't yet folded all the legalization logic into
     // the IR legalization pass (since it used to apply to the AST too).
+    //
+    // TODO: This code has issues that can lead to IR validation
+    // failure, because we might remove a `struct X` that has been
+    // legalized away, but leave around a `ParameterBlock<X>` instruction
+    // that is no longer valid.
     for (auto& oldInst : typeLegalizationContext->instsToRemove)
     {
         oldInst->removeAndDeallocate();
diff --git a/source/slang/ir-ssa.cpp b/source/slang/ir-ssa.cpp
index 1d049c685..e97298969 100644
--- a/source/slang/ir-ssa.cpp
+++ b/source/slang/ir-ssa.cpp
@@ -84,6 +84,9 @@ struct ConstructSSAContext
     // IR building state to use during the operation
     SharedIRBuilder sharedBuilder;
 
+    // Instructions to remove during cleanup
+    List<IRInst*> instsToRemove;
+
     IRBuilder builder;
     IRBuilder* getBuilder() { return &builder; }
 
@@ -98,20 +101,78 @@ struct ConstructSSAContext
     }
 };
 
+/// Do all uses of this instruction lead to a `load`?
+///
+/// Checks if all uses of `inst` are either loads,
+/// or get-element-address/get-field-address operations
+/// that also lead to loads.
+bool allUsesLeadToLoads(IRInst* inst)
+{
+    for (auto u = inst->firstUse; u; u = u->nextUse)
+    {
+        auto user = u->getUser();
+        switch (user->op)
+        {
+        default:
+            return false;
+
+        case kIROp_Load:
+            break;
+
+        case kIROp_getElementPtr:
+        case kIROp_FieldAddress:
+            {
+                // Sanity check: the address being used should
+                // be the base-address operand, and not the field
+                // key or index (this should never be a problem).
+                if (u != &user->getOperands()[0])
+                    return false;
+
+                if (!allUsesLeadToLoads(user))
+                    return false;
+            }
+            break;
+        }
+    }
+
+    // If all of the uses passed our checking, then
+    // we are good to go.
+    return true;
+
+}
+
 // Is the given variable one that we can promote to SSA form?
 bool isPromotableVar(
     ConstructSSAContext*    /*context*/,
     IRVar*                  var)
 {
-    // If the variable is only used directly as the pointer
-    // operand of load and store instructions, then it should
-    // be promote-able.
+    // We want to identify variables such that we can always
+    // determine what they will contain at a point in the
+    // program by directly inspecting their uses.
+    //
+    // The simplest possible answer would be instructions
+    // that are only ever used as the operand of "full"
+    // load and store instructions (loads and stores that
+    // write the entire variable). This is enough to
+    // promote simple scalar variables to SSA temporaries,
+    // but falls apart for aggregates and arrays.
+    //
+    // A slightly more powerful option (which is what we
+    // implement for now) is to promote variables when
+    // all of the stores are "full," and all other uses
+    // are in the form of a "chain" of `getElmeentAddress`
+    // or `getFieldAddress` operations that terminates
+    // with a load.
+    //
+    // An even more powerful option (which we do not yet
+    // implement) would be to handle cases where there are
+    // "chains" that end with stores, and to treat these
+    // as partial assignments (where we can still form
+    // an SSA value by creating a new temporary with just
+    // one element/field different). This kind of approach
+    // would be best if it is combined with scalarization,
+    // so that we don't need to construct aggregate temps.
     //
-    // For now, we won't deal with cases where the variable
-    // is an aggregate an we sometimes pull out individual
-    // fields or elements. This is an important extension,
-    // but we probably also need to think about scalarizing
-    // any aggregates when we promote them to registers.
 
     for (auto u = var->firstUse; u; u = u->nextUse)
     {
@@ -148,6 +209,20 @@ bool isPromotableVar(
                 assert(u == &storeInst->ptr);
             }
             break;
+
+        case kIROp_getElementPtr:
+        case kIROp_FieldAddress:
+            {
+                // Sanity check: the address being used should
+                // be the base-address operand, and not the field
+                // key or index (this should never be a problem).
+                if (u != &user->getOperands()[0])
+                    return false;
+
+                if (!allUsesLeadToLoads(user))
+                    return false;
+            }
+            break;
         }
     }
 
@@ -177,6 +252,7 @@ void identifyPromotableVars(
     }
 }
 
+/// If `value` is a promotable variable, then cast and return it.
 IRVar* asPromotableVar(
     ConstructSSAContext*    context,
     IRInst*                value)
@@ -191,6 +267,78 @@ IRVar* asPromotableVar(
     return var;
 }
 
+/// If `value` is a promotable variable or an access chain
+/// based on one, then cast and return the variable.
+IRVar* asPromotableVarAccessChain(
+    ConstructSSAContext*    context,
+    IRInst*                 value)
+{
+    switch (value->op)
+    {
+    case kIROp_Var:
+        return asPromotableVar(context, value);
+
+    case kIROp_FieldAddress:
+    case kIROp_getElementPtr:
+        return asPromotableVarAccessChain(context, value->getOperand(0));
+
+    default:
+        return nullptr;
+    }
+}
+
+/// After looking up the SSA value of avariable in some context,
+/// apply whatever "access chain" was applied at the original use site.
+///
+/// E.g., if the original operation was *((&a)->b) or *((&a) + i) and we've
+/// resolved that the value of the variable `a` should be `v`, then
+/// construct v.b or v[i].
+///
+IRInst* applyAccessChain(
+    ConstructSSAContext*    context,
+    IRBuilder*              builder,
+    IRInst*                 accessChain,
+    IRInst*                 leafVarValue)
+{
+    switch (accessChain->op)
+    {
+    default:
+        SLANG_UNEXPECTED("unexpected op along access chain");
+        UNREACHABLE_RETURN(leafVarValue);
+
+    case kIROp_Var:
+        return leafVarValue;
+
+    case kIROp_FieldAddress:
+        {
+            SLANG_ASSERT(context->instsToRemove.Contains(accessChain));
+
+            auto baseChain = accessChain->getOperand(0);
+            auto fieldKey = accessChain->getOperand(1);
+            auto type = cast<IRPtrTypeBase>(accessChain->getDataType())->getValueType();
+            auto baseValue = applyAccessChain(context, builder, baseChain, leafVarValue);
+            return builder->emitFieldExtract(
+                type,
+                baseValue,
+                fieldKey);
+        }
+
+    case kIROp_getElementPtr:
+        {
+            SLANG_ASSERT(context->instsToRemove.Contains(accessChain));
+
+            auto baseChain = accessChain->getOperand(0);
+            auto index = accessChain->getOperand(1);
+            auto type = cast<IRPtrTypeBase>(accessChain->getDataType())->getValueType();
+            auto baseValue = applyAccessChain(context, builder, baseChain, leafVarValue);
+            return builder->emitElementExtract(
+                type,
+                baseValue,
+                index);
+        }
+    }
+}
+
 // Try to read the value of an SSA variable
 // in the context of the given block. If
 // the variable is defined in the block, then
@@ -600,13 +748,15 @@ void processBlock(
                 IRLoad* loadInst = (IRLoad*)ii;
                 auto ptrArg = loadInst->ptr.get();
 
-                if (auto var = asPromotableVar(context, ptrArg))
+                if (auto var = asPromotableVarAccessChain(context, ptrArg))
                 {
                     // We are loading from a promotable variable.
                     // Look up the value in the context of this
                     // block.
                     auto val = readVar(context, blockInfo, var);
 
+                    val = applyAccessChain(context, &blockInfo->builder, ptrArg, val);
+
                     // We can just replace all uses of this
                     // load instruction with the given value.
                     loadInst->replaceUsesWith(val);
@@ -615,8 +765,21 @@ void processBlock(
                     // since it is no longer needed.
                     loadInst->removeAndDeallocate();
                 }
-        }
+            }
             break;
+
+        case kIROp_getElementPtr:
+        case kIROp_FieldAddress:
+            {
+                auto  ptrArg = ii->getOperand(0);
+                if (auto var = asPromotableVarAccessChain(context, ptrArg))
+                {
+                    context->instsToRemove.Add(ii);
+                }
+            }
+            break;
+
+
         }
     }
 
@@ -864,6 +1027,20 @@ void constructSSA(ConstructSSAContext* context)
         oldTerminator->removeAndDeallocate();
     }
 
+    // Remove all the instructions we marked for deletion along
+    // the way.
+    //
+    // Currently these are "access chain" instructions for
+    // loads from (parts of) variables that got promoted.
+    for (auto inst : context->instsToRemove)
+    {
+        // TODO: do we need to be careful here in case one
+        // of thes operations still has uses, as part of
+        // another to-be-remvoed instruction?
+
+        inst->removeAndDeallocate();
+    }
+
     // Now we should be able to go through and remove
     // of of the variables
     for (auto var : context->promotableVars)
diff --git a/source/slang/ir-validate.cpp b/source/slang/ir-validate.cpp
index 1e36322f4..7c04c0873 100644
--- a/source/slang/ir-validate.cpp
+++ b/source/slang/ir-validate.cpp
@@ -75,6 +75,16 @@ namespace Slang
         auto instParent = inst->getParent();
 
         auto operandValue = operandUse->get();
+
+        if( !operandValue )
+        {
+            // A null operand should almost always be an error, but
+            // we currently have a few cases where this arises.
+            //
+            // TODO: plug the leaks.
+            return;
+        }
+
         auto operandParent = operandValue->getParent();
 
         if (auto instParentBlock = as<IRBlock>(instParent))
diff --git a/tests/bugs/gh-518.slang b/tests/bugs/gh-518.slang
new file mode 100644
index 000000000..96a101db2
--- /dev/null
+++ b/tests/bugs/gh-518.slang
@@ -0,0 +1,36 @@
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12
+//TEST_DISABLED(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute
+
+// Note: can't actually test this on Vulkan right now because
+// the support in render-test isn't good enough.
+
+// Confirm that we can handle arrays of resources
+// being passed to a function.
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):dxbinding(0),glbinding(0),out
+RWStructuredBuffer<int> gBuffer;
+
+//TEST_INPUT: Texture2D(size=4, content=one):dxbinding(0),glbinding(0)
+//TEST_INPUT: Texture2D(size=4, content=one):dxbinding(1),glbinding(1)
+Texture2D gTextures[2];
+
+float broken(Texture2D t[2])
+{
+    return t[0].Load(int3(0)).x + t[1].Load(int3(0)).x;
+}
+
+int test(int val)
+{
+    return val*int(broken(gTextures));
+}
+
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint tid = dispatchThreadID.x;
+    int val = int(tid);
+    val = test(val);
+    gBuffer[tid] = val;
+}
\ No newline at end of file
diff --git a/tests/bugs/gh-518.slang.expected.txt b/tests/bugs/gh-518.slang.expected.txt
new file mode 100644
index 000000000..e1e8ccec4
--- /dev/null
+++ b/tests/bugs/gh-518.slang.expected.txt
@@ -0,0 +1,4 @@
+0
+2
+4
+6