From ee052a9e9bd7c0d233816556ebe8f0078bd9ec4d Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Tue, 13 Aug 2024 14:01:36 -0400
Subject: FIx issue with specializing witness tables (#4839)

---
 source/slang/slang-ir-specialize.cpp               |  15 ++-
 .../unit-test-generic-interface-conformance.cpp    | 102 +++++++++++++++++++++
 2 files changed, 116 insertions(+), 1 deletion(-)
 create mode 100644 tools/slang-unit-test/unit-test-generic-interface-conformance.cpp

diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index 5c9e1ad24..c86906b2d 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -131,6 +131,15 @@ struct SpecializationContext
         return false;
     }
 
+    // Check if an inst is a dynamic dispatch witness table. 
+    // These insts may not have any uses yet, and do not have side effects,
+    // but should be specialized if necessary.
+    //
+    bool isWitnessTableType(IRInst* inst)
+    {
+        return inst->findDecoration<IRDynamicDispatchWitnessDecoration>();
+    }
+
     // When an instruction isn't fully specialized, but its operands *are*
     // then it is a candidate for specialization itself, so we will have
     // a query to check for the "all operands fully specialized" case.
@@ -826,8 +835,12 @@ struct SpecializationContext
                     // specialization opportunities (generic specialization,
                     // existential specialization, simplifications, etc.)
                     //
-                    if (inst->hasUses() || inst->mightHaveSideEffects())
+                    if (inst->hasUses() ||
+                        inst->mightHaveSideEffects() ||
+                        isWitnessTableType(inst))
+                    {
                         hasSpecialization |= maybeSpecializeInst(inst);
+                    }
 
                     // Finally, we need to make our logic recurse through
                     // the whole IR module, so we want to add the children
diff --git a/tools/slang-unit-test/unit-test-generic-interface-conformance.cpp b/tools/slang-unit-test/unit-test-generic-interface-conformance.cpp
new file mode 100644
index 000000000..740d13bef
--- /dev/null
+++ b/tools/slang-unit-test/unit-test-generic-interface-conformance.cpp
@@ -0,0 +1,102 @@
+// unit-test-translation-unit-import.cpp
+
+#include "slang.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "tools/unit-test/slang-unit-test.h"
+#include "slang-com-ptr.h"
+#include "../../source/core/slang-io.h"
+#include "../../source/core/slang-process.h"
+
+using namespace Slang;
+
+// Test that the IModule::findAndCheckEntryPoint API supports discovering 
+// entrypoints without a [shader] attribute.
+
+SLANG_UNIT_TEST(genericInterfaceConformance)
+{
+    // Source for a module that contains an undecorated entrypoint.
+    const char* userSourceBody = R"(
+        public interface ITestInterface<Real : IFloat> {
+            Real sample();
+        }
+
+        struct TestInterfaceImpl<Real : IFloat> : ITestInterface<Real> {
+            Real sample() {
+                return x;
+            }
+            Real x;
+        }
+
+        //TEST_INPUT: set data = new StructuredBuffer<ITestInterface<float> >[new TestInterfaceImpl<float>{1.0}];
+        StructuredBuffer<ITestInterface<float>> data;
+
+        //TEST_INPUT: set outputBuffer = out ubuffer(data=[0 0 0 0], stride=4);
+        RWStructuredBuffer<int> outputBuffer;
+
+        //TEST_INPUT: type_conformance TestInterfaceImpl<float>:ITestInterface<float> = 3
+
+        [numthreads(1, 1, 1)]
+        void computeMain()
+        {
+            let obj = data[0];
+            // CHECK: 1
+            outputBuffer[0] = int(obj.sample());
+        }
+        )";
+
+    ComPtr<slang::IGlobalSession> globalSession;
+    SLANG_CHECK(slang_createGlobalSession(SLANG_API_VERSION, globalSession.writeRef()) == SLANG_OK);
+    slang::TargetDesc targetDesc = {};
+    targetDesc.format = SLANG_HLSL;
+
+    slang::SessionDesc sessionDesc = {};
+    sessionDesc.targetCount = 1;
+    sessionDesc.targets = &targetDesc;
+    sessionDesc.allowGLSLSyntax = true;
+    
+    ComPtr<slang::ISession> session;
+    SLANG_CHECK(globalSession->createSession(sessionDesc, session.writeRef()) == SLANG_OK);
+
+    ComPtr<slang::IBlob> diagnosticBlob;
+    auto module = session->loadModuleFromSourceString("m", "m.slang", userSourceBody, diagnosticBlob.writeRef());
+    SLANG_CHECK(module != nullptr);
+
+    ComPtr<slang::IEntryPoint> entryPoint;
+    module->findAndCheckEntryPoint("computeMain", SLANG_STAGE_COMPUTE, entryPoint.writeRef(), diagnosticBlob.writeRef());
+    SLANG_CHECK(entryPoint != nullptr);
+
+    ComPtr<slang::IComponentType> compositeProgram;
+    slang::IComponentType* components[] = { module, entryPoint.get() };
+    session->createCompositeComponentType(components, 2, compositeProgram.writeRef(), diagnosticBlob.writeRef());
+    SLANG_CHECK(compositeProgram != nullptr);
+
+    ComPtr<slang::ITypeConformance> typeConformance;
+    auto result =  session->createTypeConformanceComponentType(
+        compositeProgram->getLayout()->findTypeByName("TestInterfaceImpl<float>"), 
+        compositeProgram->getLayout()->findTypeByName("ITestInterface<float>"),
+        typeConformance.writeRef(),
+        3,
+        diagnosticBlob.writeRef());
+    SLANG_CHECK(result == SLANG_OK);
+    SLANG_CHECK(typeConformance != nullptr);
+
+    ComPtr<slang::IComponentType> compositeProgram2;
+    slang::IComponentType* components2[] = { compositeProgram.get(), typeConformance.get() };
+    session->createCompositeComponentType(
+        components2, 2, compositeProgram2.writeRef(), diagnosticBlob.writeRef());
+
+    ComPtr<slang::IComponentType> linkedProgram;
+    compositeProgram2->link(linkedProgram.writeRef(), diagnosticBlob.writeRef());
+    SLANG_CHECK(linkedProgram != nullptr);
+
+    ComPtr<slang::IBlob> code;
+    linkedProgram->getEntryPointCode(0, 0, code.writeRef(), diagnosticBlob.writeRef());
+    SLANG_CHECK(code != nullptr);
+
+    auto codeSrc = UnownedStringSlice((const char*)code->getBufferPointer());
+    SLANG_CHECK(codeSrc.indexOf(toSlice("computeMain")) != -1);
+}
+
-- 
cgit v1.2.3