Fix matrix packing/unpacking logic. (#2393)

Include indirectly referenced witness table in dynamic dispatch.

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

1 files changed, 55 insertions, 2 deletions

diff --git a/source/slang/slang-ir-specialize-dispatch.cpp b/source/slang/slang-ir-specialize-dispatch.cpp
index 0ca0933f1..b6e9d2c45 100644
--- a/source/slang/slang-ir-specialize-dispatch.cpp
+++ b/source/slang/slang-ir-specialize-dispatch.cpp
@@ -92,7 +92,7 @@ IRFunc* specializeDispatchFunction(SharedGenericsLoweringContext* sharedContext,
         auto witnessTable = witnessTables[i];
         auto seqIdDecoration = witnessTable->findDecoration<IRSequentialIDDecoration>();
         SLANG_ASSERT(seqIdDecoration);
-        
+
         if (i != witnessTables.getCount() - 1)
         {
             // Create a case block if we are not the last case.
@@ -165,6 +165,39 @@ IRFunc* specializeDispatchFunction(SharedGenericsLoweringContext* sharedContext,
     return newDispatchFunc;
 }
 
+// Returns true if the witness table is transitively referenced through a witness table with
+// linkage.
+bool _isWitnessTableTransitivelyVisible(IRInst* witness)
+{
+    if (witness->findDecoration<IRLinkageDecoration>())
+        return true;
+
+    OrderedHashSet<IRInst*> workSet;
+    List<IRInst*> workList;
+    workList.add(witness);
+    for (int i = 0; i < workList.getCount(); i++)
+    {
+        auto item = workList[i];
+        if (item->findDecoration<IRLinkageDecoration>())
+            return true;
+        for (auto use = item->firstUse; use; use = use->nextUse)
+        {
+            auto user = use->getUser();
+            if (user->getOp() == kIROp_WitnessTableEntry)
+            {
+                if (user->getParent())
+                {
+                    if (workSet.Add(user->getParent()))
+                    {
+                        workList.add(user->getParent());
+                    }
+                }
+            }
+        }
+    }
+    return false;
+}
+
 // Ensures every witness table object has been assigned a sequential ID.
 // All witness tables will have a SequentialID decoration after this function is run.
 // The sequantial ID in the decoration will be the same as the one specified in the Linkage.
@@ -173,6 +206,8 @@ IRFunc* specializeDispatchFunction(SharedGenericsLoweringContext* sharedContext,
 // can be looked up by the user via future Slang API calls.
 void ensureWitnessTableSequentialIDs(SharedGenericsLoweringContext* sharedContext)
 {
+    StringBuilder generatedMangledName;
+
     auto linkage = sharedContext->targetReq->getLinkage();
     for (auto inst : sharedContext->module->getGlobalInsts())
     {
@@ -186,8 +221,26 @@ void ensureWitnessTableSequentialIDs(SharedGenericsLoweringContext* sharedContex
             else
             {
                 // If this witness table entry does not have a linkage,
+                // we need to check if it is transitively visible via
+                // associatedtypes from an existing witness table with linkage.
+                // If so we still need to include this witness talbe, otherwise
                 // don't assign sequential ID for it.
-                continue;
+                if (_isWitnessTableTransitivelyVisible(inst))
+                {
+                    // generate a unique linkage for it.
+                    static int32_t uniqueId = 0;
+                    uniqueId++;
+                    if (auto nameHint = inst->findDecoration<IRNameHintDecoration>())
+                    {
+                        generatedMangledName << nameHint->getName();
+                    }
+                    generatedMangledName << "_generated_witness_uuid_" << uniqueId;
+                    witnessTableMangledName = generatedMangledName.getUnownedSlice();
+                }
+                else
+                {
+                    continue;
+                }
             }
 
             // If the inst already has a SequentialIDDecoration, stop now.