Allow dynamic dispatch to handle nested interface-typed fields. (#2336)

8 files changed, 163 insertions, 14 deletions

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index d7ddb773a..9b7fcfa38 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -365,10 +365,6 @@ Result linkAndOptimizeIR(
     if (sink->getErrorCount() != 0)
         return SLANG_FAIL;
 
-    lowerTuples(irModule, sink);
-    if (sink->getErrorCount() != 0)
-        return SLANG_FAIL;
-
     // TODO(DG): There are multiple DCE steps here, which need to be changed
     //   so that they don't just throw out any non-entry point code
     // Debugging code for IR transformations...
diff --git a/source/slang/slang-ir-any-value-marshalling.cpp b/source/slang/slang-ir-any-value-marshalling.cpp
index ba8f49cdd..5c91766f7 100644
--- a/source/slang/slang-ir-any-value-marshalling.cpp
+++ b/source/slang/slang-ir-any-value-marshalling.cpp
@@ -215,6 +215,20 @@ namespace Slang
                 }
                 break;
             }
+            case kIROp_AnyValueType:
+            {
+                auto anyValType = cast<IRAnyValueType>(dataType);
+                auto info = ensureAnyValueType(anyValType);
+                for (auto field : info->fieldKeys)
+                {
+                    auto fieldAddr = builder->emitFieldAddress(
+                        builder->getPtrType(builder->getUIntType()),
+                        concreteTypedVar,
+                        field);
+                    emitMarshallingCode(builder, context, fieldAddr);
+                }
+                break;
+            }
             default:
                 if (as<IRTextureTypeBase>(dataType) || as<IRSamplerStateTypeBase>(dataType))
                 {
@@ -719,10 +733,34 @@ namespace Slang
             }
             return offset;
         }
+        case kIROp_AnyValueType:
+        {
+            auto anyValueType = cast<IRAnyValueType>(type);
+            return alignUp(offset, 4) + (SlangInt)getIntVal(anyValueType->getSize());
+        }
+        case kIROp_TupleType:
+        {
+            auto tupleType = cast<IRTupleType>(type);
+            for (UInt i = 0; i < tupleType->getOperandCount(); i++)
+            {
+                auto elementType = tupleType->getOperand(i);
+                offset = _getAnyValueSizeRaw((IRType*)elementType, offset);
+                if (offset < 0) return offset;
+            }
+            return offset;
+        }
+        case kIROp_WitnessTableType:
+        case kIROp_WitnessTableIDType:
+        case kIROp_RTTIHandleType:
+        {
+            return alignUp(offset, 4) + kRTTIHandleSize;
+        }
         case kIROp_InterfaceType:
         {
-            // TODO: implement anyValue packing for interface types.
-            return -1;
+            auto interfaceType = cast<IRInterfaceType>(type);
+            auto size = SharedGenericsLoweringContext::getInterfaceAnyValueSize(interfaceType, interfaceType->sourceLoc);
+            size += kRTTIHeaderSize;
+            return alignUp(offset, 4) + alignUp((SlangInt)size, 4);
         }
         default:
             if (as<IRTextureTypeBase>(type) || as<IRSamplerStateTypeBase>(type))
diff --git a/source/slang/slang-ir-generics-lowering-context.cpp b/source/slang/slang-ir-generics-lowering-context.cpp
index 6377ef00e..fbe298d3f 100644
--- a/source/slang/slang-ir-generics-lowering-context.cpp
+++ b/source/slang/slang-ir-generics-lowering-context.cpp
@@ -373,9 +373,9 @@ namespace Slang
         return witnessTables;
     }
 
-    IRIntegerValue SharedGenericsLoweringContext::getInterfaceAnyValueSize(IRInst* type, SourceLoc usageLocation)
+    IRIntegerValue SharedGenericsLoweringContext::getInterfaceAnyValueSize(IRInst* type, SourceLoc usageLoc)
     {
-        SLANG_UNUSED(usageLocation);
+        SLANG_UNUSED(usageLoc);
 
         if (auto decor = type->findDecoration<IRAnyValueSizeDecoration>())
         {
diff --git a/source/slang/slang-ir-generics-lowering-context.h b/source/slang/slang-ir-generics-lowering-context.h
index f7a8ae0ec..e373e2bec 100644
--- a/source/slang/slang-ir-generics-lowering-context.h
+++ b/source/slang/slang-ir-generics-lowering-context.h
@@ -13,6 +13,8 @@ namespace Slang
 
     constexpr IRIntegerValue kInvalidAnyValueSize = 0xFFFFFFFF;
     constexpr IRIntegerValue kDefaultAnyValueSize = 16;
+    constexpr SlangInt kRTTIHeaderSize = 16;
+    constexpr SlangInt kRTTIHandleSize = 8;
 
     struct SharedGenericsLoweringContext
     {
@@ -73,7 +75,7 @@ namespace Slang
         // Emits an IRRTTIObject containing type information for a given type.
         IRInst* maybeEmitRTTIObject(IRInst* typeInst);
 
-        IRIntegerValue getInterfaceAnyValueSize(IRInst* type, SourceLoc usageLocation);
+        static IRIntegerValue getInterfaceAnyValueSize(IRInst* type, SourceLoc usageLoc);
         IRType* lowerAssociatedType(IRBuilder* builder, IRInst* type);
 
         IRType* lowerType(IRBuilder* builder, IRInst* paramType, const Dictionary<IRInst*, IRInst*>& typeMapping, IRType* concreteType);
diff --git a/source/slang/slang-ir-layout.cpp b/source/slang/slang-ir-layout.cpp
index c89d15808..cbf0a6e61 100644
--- a/source/slang/slang-ir-layout.cpp
+++ b/source/slang/slang-ir-layout.cpp
@@ -3,6 +3,8 @@
 
 #include "slang-ir-insts.h"
 
+#include "slang-ir-generics-lowering-context.h"
+
 // This file implements facilities for computing and caching layout
 // information on IR types.
 //
@@ -194,7 +196,52 @@ static Result _calcNaturalSizeAndAlignment(
                 outSizeAndAlignment);
         }
         break;
-
+    case kIROp_AnyValueType:
+        {
+            auto anyValType = cast<IRAnyValueType>(type);
+            outSizeAndAlignment->size = getIntVal(anyValType->getSize());
+            outSizeAndAlignment->alignment = 4;
+            return SLANG_OK;
+        }
+        break;
+    case kIROp_TupleType:
+        {
+            auto tupleType = cast<IRTupleType>(type);
+            IRSizeAndAlignment resultLayout;
+            for (UInt i = 0; i < tupleType->getOperandCount(); i++)
+            {
+                auto elementType = tupleType->getOperand(i);
+                IRSizeAndAlignment fieldTypeLayout;
+                SLANG_RETURN_ON_FAIL(getNaturalSizeAndAlignment(target, (IRType*)elementType, &fieldTypeLayout));
+                resultLayout.size = align(resultLayout.size, fieldTypeLayout.alignment);
+                resultLayout.alignment = std::max(resultLayout.alignment, fieldTypeLayout.alignment);
+            }
+            *outSizeAndAlignment = resultLayout;
+            return SLANG_OK;
+        }
+        break;
+    case kIROp_WitnessTableType:
+    case kIROp_WitnessTableIDType:
+    case kIROp_RTTIHandleType:
+        {
+            outSizeAndAlignment->size = kRTTIHandleSize;
+            outSizeAndAlignment->alignment = 4;
+            return SLANG_OK;
+        }
+        break;
+    case kIROp_InterfaceType:
+        {
+            auto interfaceType = cast<IRInterfaceType>(type);
+            auto size = SharedGenericsLoweringContext::getInterfaceAnyValueSize(interfaceType, interfaceType->sourceLoc);
+            size += kRTTIHeaderSize;
+            size = align(size, 4);
+            IRSizeAndAlignment resultLayout;
+            resultLayout.size = size;
+            resultLayout.alignment = 4;
+            *outSizeAndAlignment = resultLayout;
+            return SLANG_OK;
+        }
+        break;
     case kIROp_MatrixType:
         {
             auto matType = cast<IRMatrixType>(type);
diff --git a/source/slang/slang-ir-lower-generics.cpp b/source/slang/slang-ir-lower-generics.cpp
index f517719ad..3b9a17738 100644
--- a/source/slang/slang-ir-lower-generics.cpp
+++ b/source/slang/slang-ir-lower-generics.cpp
@@ -5,6 +5,7 @@
 #include "slang-ir-augment-make-existential.h"
 #include "slang-ir-generics-lowering-context.h"
 #include "slang-ir-lower-existential.h"
+#include "slang-ir-lower-tuple-types.h"
 #include "slang-ir-lower-generic-function.h"
 #include "slang-ir-lower-generic-call.h"
 #include "slang-ir-lower-generic-type.h"
@@ -194,15 +195,19 @@ namespace Slang
         if (sink->getErrorCount() != 0)
             return;
 
-        generateAnyValueMarshallingFunctions(&sharedContext);
-        if (sink->getErrorCount() != 0)
-            return;
-
         // This optional step replaces all uses of witness tables and RTTI objects with
         // sequential IDs. Without this step, we will emit code that uses function pointers and
         // real RTTI objects and witness tables.
         specializeRTTIObjects(&sharedContext, sink);
 
+        lowerTuples(module, sink);
+        if (sink->getErrorCount() != 0)
+            return;
+
+        generateAnyValueMarshallingFunctions(&sharedContext);
+        if (sink->getErrorCount() != 0)
+            return;
+
         // We might have generated new temporary variables during lowering.
         // An SSA pass can clean up unnecessary load/stores.
         constructSSA(module);
diff --git a/tests/bugs/nested-existential-dyndispatch.slang b/tests/bugs/nested-existential-dyndispatch.slang
new file mode 100644
index 000000000..b6689b54d
--- /dev/null
+++ b/tests/bugs/nested-existential-dyndispatch.slang
@@ -0,0 +1,56 @@
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+interface IFoo
+{
+    float sample();
+}
+
+//TEST_INPUT: type_conformance Bar:IBar = 0
+[anyValueSize(64)]
+interface IBar
+{
+    float sample();
+}
+
+//TEST_INPUT: type_conformance Foo:IFoo = 0
+struct Foo : IFoo
+{
+    float3 albedo;
+    float sample()
+    {
+        return albedo.x;
+    }
+}
+
+struct Bar : IBar
+{
+    IFoo foo;
+    float sample()
+    {
+        return foo.sample();
+    }
+}
+
+struct AnyVal
+{
+    uint data[16];
+}
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    int index = dispatchThreadID.x;
+
+    Foo f;
+    f.albedo = float3(1.0);
+    Bar bar;
+    bar.foo = f;
+
+    AnyVal data = reinterpret<AnyVal, Bar>(bar);
+    IBar dynBar = createDynamicObject<IBar, AnyVal>(0, data);
+
+    outputBuffer[index] = dynBar.sample();
+}
diff --git a/tests/bugs/nested-existential-dyndispatch.slang.expected.txt b/tests/bugs/nested-existential-dyndispatch.slang.expected.txt
new file mode 100644
index 000000000..24e3a50ad
--- /dev/null
+++ b/tests/bugs/nested-existential-dyndispatch.slang.expected.txt
@@ -0,0 +1,5 @@
+type: float
+1.0
+1.0
+1.0
+1.0
\ No newline at end of file