Fix matrix packing/unpacking logic. (#2393)

Include indirectly referenced witness table in dynamic dispatch.

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

3 files changed, 61 insertions, 7 deletions

diff --git a/source/slang/slang-ir-any-value-marshalling.cpp b/source/slang/slang-ir-any-value-marshalling.cpp
index ea1a6cf32..f1a1982c2 100644
--- a/source/slang/slang-ir-any-value-marshalling.cpp
+++ b/source/slang/slang-ir-any-value-marshalling.cpp
@@ -160,18 +160,19 @@ namespace Slang
                 auto elementType = matrixType->getElementType();
                 auto colCount = getIntVal(matrixType->getColumnCount());
                 auto rowCount = getIntVal(matrixType->getRowCount());
+                auto rowVecType = builder->getVectorType(elementType, matrixType->getRowCount());
                 for (IRIntegerValue i = 0; i < colCount; i++)
                 {
                     auto col = builder->emitElementAddress(
-                        elementType,
+                        builder->getPtrType(rowVecType),
                         concreteTypedVar,
                         builder->getIntValue(builder->getIntType(), i));
                     for (IRIntegerValue j = 0; j < rowCount; j++)
                     {
-                        auto element = builder->emitElementExtract(
-                            elementType,
+                        auto element = builder->emitElementAddress(
+                            builder->getPtrType(elementType),
                             col,
-                            builder->getIntValue(builder->getIntType(), i));
+                            builder->getIntValue(builder->getIntType(), j));
                         emitMarshallingCode(builder, context, element);
                     }
                 }
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.
diff --git a/source/slang/slang-ir-specialize-dynamic-associatedtype-lookup.cpp b/source/slang/slang-ir-specialize-dynamic-associatedtype-lookup.cpp
index 441d32d24..89b965739 100644
--- a/source/slang/slang-ir-specialize-dynamic-associatedtype-lookup.cpp
+++ b/source/slang/slang-ir-specialize-dynamic-associatedtype-lookup.cpp
@@ -197,7 +197,7 @@ struct AssociatedTypeLookupSpecializationContext
             if (inst->getOp() == kIROp_WitnessTable)
             {
                 auto seqId = inst->findDecoration<IRSequentialIDDecoration>();
-                SLANG_ASSERT(seqId);
+                SLANG_RELEASE_ASSERT(seqId);
                 // Insert code to pack sequential ID into an uint2 at all use sites.
                 IRUse* nextUse = nullptr;
                 for (auto use = inst->firstUse; use; use = nextUse)