Handle structure initializers in IR type legalization (shader-slang#567)

Fixes shader-slang#566

The basic problem here is that the front-end translates a structure initializer-list expression into a `makeStruct` instruction (with one argument per field), but the IR type legalization logic wasn't handling the case where a `makeStruct` is used to construct a struct value that needs to get split by legalization.

The implementation is relatively straightforward, and like the other cases of instruction legalization for compound types, it follows the shape of the `LegalType`/`LegalVal` cases. The one interesting bit is that we need to be a bit careful and filter the single argument list for `makeStruct` into two in the case where we generate a "pair" type for something that has both "ordinary" and "special" (resource) fields. Luckily the `PairInfo` data that was generated by type legalization has exactly the information we need (by design).

This change does not address several issues that could be handled in follow-on changes:

* The `makeArray` instruction will face similar issues if it is applied to a type that requires legalization: we'd need to turn an array of `LegalVal`s into a bunch of distinct arrays.

* The error message when we hit the unimplemented case here isn't great. Ideally we should provide the line number of the instruction that fails in an error message when legalization fails.

This change tries to focus narrowly on the bug at hand, and leave these issues for later changes.

Author: Tim Foley

source/slang/ir-legalize-types.cpp

@@ -669,6 +669,122 @@ static LegalVal legalizeGetElementPtr(
         indexOperand);
 }
 
+static LegalVal legalizeMakeStruct(
+    IRTypeLegalizationContext*  context,
+    LegalType                   legalType,
+    LegalVal const*             legalArgs,
+    UInt                        argCount)
+{
+    auto builder = context->builder;
+
+    switch(legalType.flavor)
+    {
+    case LegalType::Flavor::simple:
+        {
+            List<IRInst*> args;
+            for(UInt aa = 0; aa < argCount; ++aa)
+            {
+                // Note: we assume that all the arguments
+                // must be simple here, because otherwise
+                // the `struct` type with them as fields
+                // would not be simple...
+                //
+                args.Add(legalArgs[aa].getSimple());
+            }
+            return LegalVal::simple(
+                builder->emitMakeStruct(
+                    legalType.getSimple(),
+                    argCount,
+                    args.Buffer()));
+        }
+
+    case LegalType::Flavor::pair:
+        {
+            // There are two sides, the ordinary and the special,
+            // and we basically just dispatch to both of them.
+            auto pairType = legalType.getPair();
+            auto pairInfo = pairType->pairInfo;
+            LegalType ordinaryType = pairType->ordinaryType;
+            LegalType specialType = pairType->specialType;
+
+            List<LegalVal> ordinaryArgs;
+            List<LegalVal> specialArgs;
+            UInt argCounter = 0;
+            for(auto ee : pairInfo->elements)
+            {
+                UInt argIndex = argCounter++;
+                LegalVal arg = legalArgs[argIndex];
+
+                if((ee.flags & Slang::PairInfo::kFlag_hasOrdinaryAndSpecial) == Slang::PairInfo::kFlag_hasOrdinaryAndSpecial)
+                {
+                    // The field is itself a pair type, so we expect
+                    // the argument value to be one too...
+                    auto argPair = arg.getPair();
+                    ordinaryArgs.Add(argPair->ordinaryVal);
+                    specialArgs.Add(argPair->specialVal);
+                }
+                else if(ee.flags & Slang::PairInfo::kFlag_hasOrdinary)
+                {
+                    ordinaryArgs.Add(arg);
+                }
+                else if(ee.flags & Slang::PairInfo::kFlag_hasSpecial)
+                {
+                    specialArgs.Add(arg);
+                }
+            }
+
+            LegalVal ordinaryVal = legalizeMakeStruct(
+                context,
+                ordinaryType,
+                ordinaryArgs.Buffer(),
+                ordinaryArgs.Count());
+
+            LegalVal specialVal = legalizeMakeStruct(
+                context,
+                specialType,
+                specialArgs.Buffer(),
+                specialArgs.Count());
+
+            return LegalVal::pair(ordinaryVal, specialVal, pairInfo);
+        }
+        break;
+
+    case LegalType::Flavor::tuple:
+        {
+            // We are constructing a tuple of values from
+            // the individual fields. We need to identify
+            // for each tuple element what field it uses,
+            // and then extract that field's value.
+
+            auto tupleType = legalType.getTuple();
+
+            RefPtr<TuplePseudoVal> resTupleInfo = new TuplePseudoVal();
+            UInt argCounter = 0;
+            for(auto typeElem : tupleType->elements)
+            {
+                auto elemKey = typeElem.key;
+                UInt argIndex = argCounter++;
+                SLANG_ASSERT(argIndex < argCount);
+
+                LegalVal argVal = legalArgs[argIndex];
+
+                TuplePseudoVal::Element resElem;
+                resElem.key = elemKey;
+                resElem.val = argVal;
+
+                resTupleInfo->elements.Add(resElem);
+            }
+            return LegalVal::tuple(resTupleInfo);
+        }
+
+    default:
+        SLANG_UNEXPECTED("unhandled");
+        UNREACHABLE_RETURN(LegalVal());
+    }
+}
+
+
+
 static LegalVal legalizeInst(
     IRTypeLegalizationContext*    context,
     IRInst*                     inst,
@@ -695,6 +811,13 @@ static LegalVal legalizeInst(
     case kIROp_Call:
         return legalizeCall(context, (IRCall*)inst);
 
+    case kIROp_makeStruct:
+        return legalizeMakeStruct(
+            context,
+            type,
+            args,
+            inst->getOperandCount());
+
     default:
         // TODO: produce a user-visible diagnostic here
         SLANG_UNEXPECTED("non-simple operand(s)!");

source/slang/legalize-types.h

@@ -290,15 +290,15 @@ struct LegalVal
         return result;
     }
 
-    IRInst* getSimple()
+    IRInst* getSimple() const
     {
         SLANG_ASSERT(flavor == Flavor::simple);
         return irValue;
     }
 
     static LegalVal tuple(RefPtr<TuplePseudoVal> tupleVal);
 
-    RefPtr<TuplePseudoVal> getTuple()
+    RefPtr<TuplePseudoVal> getTuple() const
     {
         SLANG_ASSERT(flavor == Flavor::tuple);
         return obj.As<TuplePseudoVal>();
@@ -313,7 +313,7 @@ struct LegalVal
         LegalVal const&     specialVal,
         RefPtr<PairInfo>    pairInfo);
 
-    RefPtr<PairPseudoVal> getPair()
+    RefPtr<PairPseudoVal> getPair() const
     {
         SLANG_ASSERT(flavor == Flavor::pair);
         return obj.As<PairPseudoVal>();

tests/bugs/gh-566.slang

@@ -0,0 +1,34 @@
+// legalize-struct-init.slang
+
+//TEST(compute):COMPARE_COMPUTE:
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):dxbinding(0),glbinding(0),out
+//TEST_INPUT:ubuffer(data=[4 3 2 1], stride=4):dxbinding(1),glbinding(1)
+
+
+RWStructuredBuffer<uint> outputBuffer;
+RWStructuredBuffer<uint> inputBuffer;
+
+struct Helper
+{
+    RWStructuredBuffer<uint> o;
+    RWStructuredBuffer<uint> i;
+    uint                    t;  
+};
+
+void test(Helper h)
+{
+    h.o[h.t] = h.i[h.t] * 16 + h.t;    
+}
+
+void test(uint t)
+{
+    Helper h = { outputBuffer, inputBuffer, t };
+    test(h);
+}
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint tid = dispatchThreadID.x;
+    test(tid);
+}

tests/bugs/gh-566.slang.expected.txt

@@ -0,0 +1,4 @@
+40
+31
+22
+13