Fix crash: dynamic dispatch of generic interface method. (#1929)

* Fix crash: dynamic dispatch of generic interface method.

* Fix memory error.

Co-authored-by: Yong He <yhe@nvidia.com>

4 files changed, 93 insertions, 6 deletions

diff --git a/source/slang/slang-ir-lower-generic-call.cpp b/source/slang/slang-ir-lower-generic-call.cpp
index 3571b7a27..7669593c6 100644
--- a/source/slang/slang-ir-lower-generic-call.cpp
+++ b/source/slang/slang-ir-lower-generic-call.cpp
@@ -218,9 +218,6 @@ namespace Slang
             auto newCall = builder->emitCallInst(calleeRetType, newCallee, args);
             auto callInstType = callInst->getDataType();
             auto unpackInst = maybeUnpackValue(builder, callInstType, calleeRetType, newCall);
-            callInst->replaceUsesWith(unpackInst);
-            callInst->removeAndDeallocate();
-
             // Unpack other `out` arguments.
             for (auto& item : argsToUnpack)
             {
@@ -229,6 +226,8 @@ namespace Slang
                 auto unpackedVal = builder->emitUnpackAnyValue(originalValType, packedVal);
                 builder->emitStore(item.dstArg, unpackedVal);
             }
+            callInst->replaceUsesWith(unpackInst);
+            callInst->removeAndDeallocate();
         }
 
         IRInst* findInnerMostSpecializingBase(IRSpecialize* inst)
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 8a40dbde9..c1738887d 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -6702,9 +6702,18 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             {
                 HashSet<IRInst*> valuesToClone;
                 markInstsToClone(valuesToClone, parentGeneric->getFirstBlock(), returnType);
+                // For Function Types, we always clone all generic parameters regardless of whether
+                // the generic parameter appears in the function signature or not.
+                if (returnType->getOp() == kIROp_FuncType)
+                {
+                    for (auto genericParam : parentGeneric->getParams())
+                    {
+                        markInstsToClone(valuesToClone, parentGeneric->getFirstBlock(), genericParam);
+                    }
+                }
                 if (valuesToClone.Count() == 0)
                 {
-                    // If returnType is independent of generic parameters, set
+                    // If the new generic has no parameters, set
                     // the generic inst's type to just `returnType`.
                     parentGeneric->setFullType((IRType*)returnType);
                 }
@@ -6727,8 +6736,17 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
                     }
                     IRInst* clonedReturnType = nullptr;
                     cloneEnv.mapOldValToNew.TryGetValue(returnType, clonedReturnType);
-                    SLANG_ASSERT(clonedReturnType);
-                    typeBuilder.emitReturn(clonedReturnType);
+                    if (clonedReturnType)
+                    {
+                        // If the type has explicit dependency on generic parameters, use
+                        // the cloned type.
+                        typeBuilder.emitReturn(clonedReturnType);
+                    }
+                    else
+                    {
+                        // Otherwise just use the original type value directly.
+                        typeBuilder.emitReturn(returnType);
+                    }
                     parentGeneric->setFullType((IRType*)typeGeneric);
                     returnType = typeGeneric;
                 }
diff --git a/tests/compute/dynamic-dispatch-17.slang b/tests/compute/dynamic-dispatch-17.slang
new file mode 100644
index 000000000..bc2b9a6d9
--- /dev/null
+++ b/tests/compute/dynamic-dispatch-17.slang
@@ -0,0 +1,68 @@
+// Test using generic interface methods with dynamic dispatch.
+
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -vk -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12 -profile sm_6_0 -use-dxil -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx11 -profile sm_5_0 -output-using-type
+
+interface IReturnsZero
+{
+    float get();
+}
+
+[anyValueSize(16)]
+interface IInterface
+{
+    float run<Z:IReturnsZero>();
+}
+
+struct UserDefinedPackedType
+{
+    float3 val;
+    uint flags;
+};
+
+//TEST_INPUT:ubuffer(data=[0], stride=4):out,name=gOutputBuffer
+RWStructuredBuffer<float> gOutputBuffer;
+
+//TEST_INPUT: set gObj = new StructuredBuffer<UserDefinedPackedType>[new UserDefinedPackedType{[1.0, 2.0, 3.0], 3}, new UserDefinedPackedType{[2.0, 3.0, 4.0], 4}];
+RWStructuredBuffer<UserDefinedPackedType> gObj;
+
+//TEST_INPUT: type_conformance FloatVal:IInterface = 3
+//TEST_INPUT: type_conformance Float4Val:IInterface = 4
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3       dispatchThreadID : SV_DispatchThreadID)
+{
+    float result = 0.0;
+    for (int i = 0; i < 2; i++)
+    {
+        var rawObj = gObj.Load(i);
+        IInterface dynamicObj = createDynamicObject<IInterface, UserDefinedPackedType>(rawObj.flags, rawObj);
+        result += dynamicObj.run<ReturnsZero>();
+    }
+    gOutputBuffer[0] = result;
+}
+
+struct ReturnsZero : IReturnsZero
+{
+    float get() { return 0.0; }
+}
+struct FloatVal : IInterface
+{
+    float val;
+    float run<Z:IReturnsZero>()
+    {
+        Z z;
+        return val + z.get();
+    }
+};
+struct Float4Struct { float4 val; }
+struct Float4Val : IInterface
+{
+    Float4Struct val;
+    float run<Z:IReturnsZero>()
+    {
+        Z z;
+        return val.val.x + val.val.y + z.get();
+    }
+};
diff --git a/tests/compute/dynamic-dispatch-17.slang.expected.txt b/tests/compute/dynamic-dispatch-17.slang.expected.txt
new file mode 100644
index 000000000..253df0793
--- /dev/null
+++ b/tests/compute/dynamic-dispatch-17.slang.expected.txt
@@ -0,0 +1,2 @@
+type: float
+6.0