Preserve specialization cache in IR for specialization pass. (#2293)

* Perserve specialization cache in IR for specialization pass.

* Fix compile error.

* Fix.

* Fix.

* Fix test case.

* Fix.

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

11 files changed, 196 insertions, 8 deletions

diff --git a/source/core/slang-dictionary.h b/source/core/slang-dictionary.h
index 470e5f6d9..b11341051 100644
--- a/source/core/slang-dictionary.h
+++ b/source/core/slang-dictionary.h
@@ -80,6 +80,9 @@ namespace Slang
 	{
 		friend class Iterator;
 		friend class ItemProxy;
+    public:
+        typedef TValue ValueType;
+        typedef TKey KeyType;
 	private:
 		inline int GetProbeOffset(int /*probeId*/) const
 		{
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index fb0f65c5f..d7ddb773a 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -37,6 +37,7 @@
 #include "slang-ir-specialize-resources.h"
 #include "slang-ir-ssa.h"
 #include "slang-ir-ssa-simplification.h"
+#include "slang-ir-strip-cached-dict.h"
 #include "slang-ir-strip-witness-tables.h"
 #include "slang-ir-synthesize-active-mask.h"
 #include "slang-ir-union.h"
@@ -707,6 +708,8 @@ Result linkAndOptimizeIR(
         break;
     }
 
+    stripCachedDictionaries(irModule);
+
     // TODO: our current dynamic dispatch pass will remove all uses of witness tables.
     // If we are going to support function-pointer based, "real" modular dynamic dispatch,
     // we will need to disable this pass.
diff --git a/source/slang/slang-ir-dce.cpp b/source/slang/slang-ir-dce.cpp
index 9ed5249fe..2d417221b 100644
--- a/source/slang/slang-ir-dce.cpp
+++ b/source/slang/slang-ir-dce.cpp
@@ -19,6 +19,10 @@ struct DeadCodeEliminationContext
     IRModule*                       module;
     IRDeadCodeEliminationOptions    options;
 
+    // If we removed an inst, there may be still "weak references" to the inst.
+    // These uses will be replaced with `undefInst`.
+    IRInst* undefInst = nullptr;
+
     // Our overall process is going to be to determine
     // which instructions in the module are "live"
     // and then eliminate anything that wasn't found to
@@ -77,6 +81,29 @@ struct DeadCodeEliminationContext
         workList.add(inst);
     }
 
+    IRInst* getUndefInst()
+    {
+        if (!undefInst)
+        {
+            for (auto inst : module->getModuleInst()->getChildren())
+            {
+                if (inst->getOp() == kIROp_undefined && inst->getDataType() && inst->getDataType()->getOp() == kIROp_VoidType)
+                {
+                    undefInst = inst;
+                    break;
+                }
+            }
+            if (!undefInst)
+            {
+                SharedIRBuilder builderStorage(module);
+                IRBuilder builder(&builderStorage);
+                builder.setInsertInto(module->getModuleInst());
+                undefInst = builder.emitUndefined(builder.getVoidType());
+            }
+        }
+        return undefInst;
+    }
+
     // Given the basic infrastructrure above, let's
     // dive into the task of actually finding all
     // the live code in a module.
@@ -90,6 +117,13 @@ struct DeadCodeEliminationContext
         //
         markInstAsLive(module->getModuleInst());
 
+        // Ensure there is a global undef inst that is always alive.
+        // This undef inst will be used to fill in weak-referencing uses
+        // whose used value is marked as dead and eliminated.
+        // We always make sure this undef inst is available to prevent
+        // infiniate oscilating loops.
+        markInstAsLive(getUndefInst());
+
         // Marking the module as live should have
         // seeded our work list, so we can now start
         // processing entries off of our work list
@@ -120,12 +154,20 @@ struct DeadCodeEliminationContext
             UInt operandCount = inst->getOperandCount();
             for( UInt ii = 0; ii < operandCount; ++ii )
             {
+                // There are some type of operands that needs to be treated as
+                // "weak" references -- they can never hold things alive, and
+                // whenever we delete the referenced value, these operands needs
+                // to be replaced with `undef`.
                 switch (inst->getOp())
                 {
                 case kIROp_BoundInterfaceType:
                     if (inst->getOperand(ii)->getOp() == kIROp_WitnessTable)
                         continue;
                     break;
+                case kIROp_SpecializationDictionaryItem:
+                    // Ignore all operands of SpecializationDictionaryItem.
+                    // This inst is used as a cache and shouldn't hold anything alive.
+                    continue;
                 default:
                     break;
                 }
@@ -192,6 +234,10 @@ struct DeadCodeEliminationContext
             // because they must have been dead too (since we always
             // mark the parent of a live instruction as live).
             //
+            if (inst->hasUses())
+            {
+                inst->replaceUsesWith(getUndefInst());
+            }
             inst->removeAndDeallocate();
             changed = true;
         }
diff --git a/source/slang/slang-ir-generics-lowering-context.cpp b/source/slang/slang-ir-generics-lowering-context.cpp
index e20b8b680..6377ef00e 100644
--- a/source/slang/slang-ir-generics-lowering-context.cpp
+++ b/source/slang/slang-ir-generics-lowering-context.cpp
@@ -38,6 +38,7 @@ namespace Slang
 
     bool isComInterfaceType(IRType* type)
     {
+        if (!type) return false;
         if (type->findDecoration<IRComInterfaceDecoration>() || 
             type->getOp() == kIROp_ComPtrType)
         {
@@ -309,9 +310,12 @@ namespace Slang
         case kIROp_lookup_interface_method:
         {
             auto lookupInterface = static_cast<IRLookupWitnessMethod*>(paramType);
-            auto interfaceType = cast<IRInterfaceType>(cast<IRWitnessTableType>(
-                lookupInterface->getWitnessTable()->getDataType())->getConformanceType());
-            if (isBuiltin(interfaceType))
+            auto witnessTableType = as<IRWitnessTableType>(
+                lookupInterface->getWitnessTable()->getDataType());
+            if (!witnessTableType)
+                return (IRType*)paramType;
+            auto interfaceType = as<IRInterfaceType>(witnessTableType->getConformanceType());
+            if (!interfaceType || isBuiltin(interfaceType))
                 return (IRType*)paramType;
             // Make sure we are looking up inside the original interface type (prior to lowering).
             // Only in the original interface type will an associated type entry have an IRAssociatedType value.
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index c10ae8639..6304e65d2 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -756,6 +756,12 @@ INST_RANGE(Attr, PendingLayoutAttr, FuncThrowTypeAttr)
     INST(LiveRangeEnd, liveRangeEnd, 0, 0)
 INST_RANGE(LiveRangeMarker, LiveRangeStart, LiveRangeEnd)
 
+/* IRSpecialization */
+INST(SpecializationDictionaryItem, SpecializationDictionaryItem, 0, 0)
+INST(GenericSpecializationDictionary, GenericSpecializationDictionary, 0, PARENT)
+INST(ExistentialFuncSpecializationDictionary, ExistentialFuncSpecializationDictionary, 0, PARENT)
+INST(ExistentialTypeSpecializationDictionary, ExistentialTypeSpecializationDictionary, 0, PARENT)
+
 #undef PARENT
 #undef USE_OTHER
 #undef INST_RANGE
diff --git a/source/slang/slang-ir-lower-generic-function.cpp b/source/slang/slang-ir-lower-generic-function.cpp
index bf20452d7..b50737c23 100644
--- a/source/slang/slang-ir-lower-generic-function.cpp
+++ b/source/slang/slang-ir-lower-generic-function.cpp
@@ -264,7 +264,8 @@ namespace Slang
             // If the requirement is a function, interfaceRequirementVal will be the lowered function type.
             // If the requirement is an associatedtype, interfaceRequirementVal will be Ptr<RTTIObject>.
             IRInst* interfaceRequirementVal = nullptr;
-            auto witnessTableType = cast<IRWitnessTableType>(lookupInst->getWitnessTable()->getDataType());
+            auto witnessTableType = as<IRWitnessTableType>(lookupInst->getWitnessTable()->getDataType());
+            if (!witnessTableType) return;
             auto interfaceType = maybeLowerInterfaceType(cast<IRInterfaceType>(witnessTableType->getConformanceType()));
             interfaceRequirementVal = sharedContext->findInterfaceRequirementVal(interfaceType, lookupInst->getRequirementKey());
             lookupInst->setFullType((IRType*)interfaceRequirementVal);
diff --git a/source/slang/slang-ir-lower-generics.cpp b/source/slang/slang-ir-lower-generics.cpp
index 703441252..b8b2ef972 100644
--- a/source/slang/slang-ir-lower-generics.cpp
+++ b/source/slang/slang-ir-lower-generics.cpp
@@ -136,6 +136,8 @@ namespace Slang
                 if (auto lookupWitnessMethod = as<IRLookupWitnessMethod>(inst))
                 {
                     auto witnessTableType = lookupWitnessMethod->getWitnessTable()->getDataType();
+                    if (!witnessTableType)
+                        return;
                     auto interfaceType = cast<IRWitnessTableType>(witnessTableType)->getConformanceType();
                     if (isComInterfaceType((IRType*)interfaceType))
                         return;
diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index 9cc2d9ad0..1d62af47e 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -623,6 +623,76 @@ struct SpecializationContext
         }
         return nullptr;
     }
+    template<typename TDict>
+    void _readSpecializationDictionaryImpl(TDict& dict, IRInst* dictInst)
+    {
+        for (auto child : dictInst->getChildren())
+        {
+            auto item = as<IRSpecializationDictionaryItem>(child);
+            if (!item) continue;
+            IRSimpleSpecializationKey key;
+            bool shouldSkip = false;
+            for (UInt i = 1; i < item->getOperandCount(); i++)
+            {
+                if (item->getOperand(i) == nullptr)
+                {
+                    shouldSkip = true;
+                    break;
+                }
+                key.vals.add(item->getOperand(i));
+            }
+            if (shouldSkip)
+                continue;
+            auto value = as<typename std::remove_pointer<typename TDict::ValueType>::type>(item->getOperand(0));
+            SLANG_ASSERT(value);
+            dict[key] = value;
+        }
+        dictInst->removeAndDeallocate();
+    }
+    void readSpecializationDictionaries()
+    {
+        auto moduleInst = module->getModuleInst();
+        for (auto child : moduleInst->getChildren())
+        {
+            switch (child->getOp())
+            {
+            case kIROp_GenericSpecializationDictionary:
+                _readSpecializationDictionaryImpl(genericSpecializations, child);
+                break;
+            case kIROp_ExistentialFuncSpecializationDictionary:
+                _readSpecializationDictionaryImpl(existentialSpecializedFuncs, child);
+                break;
+            case kIROp_ExistentialTypeSpecializationDictionary:
+                _readSpecializationDictionaryImpl(existentialSpecializedStructs, child);
+                break;
+            default:
+                continue;
+            }
+        }
+    }
+
+    template<typename TDict>
+    void _writeSpecializationDictionaryImpl(TDict& dict, IROp dictOp, IRInst* moduleInst)
+    {
+        IRBuilder builder(&sharedBuilderStorage);
+        builder.setInsertInto(moduleInst);
+        auto dictInst = builder.emitIntrinsicInst(nullptr, dictOp, 0, nullptr);
+        builder.setInsertInto(dictInst);
+        for (auto kv : dict)
+        {
+            List<IRInst*> args;
+            args.add(kv.Value);
+            args.addRange(kv.Key.vals);
+            builder.emitIntrinsicInst(nullptr, kIROp_SpecializationDictionaryItem, args.getCount(), args.getBuffer());
+        }
+    }
+    void writeSpecializationDictionaries()
+    {
+        auto moduleInst = module->getModuleInst();
+        _writeSpecializationDictionaryImpl(genericSpecializations, kIROp_GenericSpecializationDictionary, moduleInst);
+        _writeSpecializationDictionaryImpl(existentialSpecializedFuncs, kIROp_ExistentialFuncSpecializationDictionary, moduleInst);
+        _writeSpecializationDictionaryImpl(existentialSpecializedStructs, kIROp_ExistentialTypeSpecializationDictionary, moduleInst);
+    }
 
     // All of the machinery for generic specialization
     // has been defined above, so we will now walk
@@ -637,6 +707,11 @@ struct SpecializationContext
         SharedIRBuilder* sharedBuilder = &sharedBuilderStorage;
         sharedBuilder->init(module);
 
+        // Read specialization dictionary from module if it is defined.
+        // This prevents us from generating duplicated specializations
+        // when this pass is invoked iteratively.
+        readSpecializationDictionaries();
+
         // The unspecialized IR we receive as input will have
         // `IRBindGlobalGenericParam` instructions that associate
         // each global-scope generic parameter (a type, witness
@@ -733,9 +808,15 @@ struct SpecializationContext
 
         // Once the work list has gone dry, we should have the invariant
         // that there are no `specialize` instructions inside of non-generic
-        // functions that in turn reference a generic type/function, *except*
-        // in the case where that generic is for a builtin type/function, in
-        // which case we wouldn't want to specialize it anyway.
+        // functions that in turn reference a generic type/function unless the generic is for a
+        // builtin type/function, or some of the type arguments are unknown at compile time, in
+        // which case we will rely on a follow up pass the translate it into a dynamic dispatch
+        // function.
+
+        // For functions that still have `specialize` uses left, we need to preserve the
+        // its specializations in resulting IR so they can be reconstructed when this
+        // specialization pass gets invoked again.
+        writeSpecializationDictionaries();
     }
 
     void addDirtyInstsToWorkListRec(IRInst* inst)
@@ -2223,5 +2304,4 @@ IRInst* specializeGeneric(
     return specializeGenericImpl(baseGeneric, specializeInst, module, nullptr);
 }
 
-
 } // namespace Slang
diff --git a/source/slang/slang-ir-strip-cached-dict.cpp b/source/slang/slang-ir-strip-cached-dict.cpp
new file mode 100644
index 000000000..3b7c1479e
--- /dev/null
+++ b/source/slang/slang-ir-strip-cached-dict.cpp
@@ -0,0 +1,28 @@
+// slang-ir-strip-cached-dict.cpp
+#include "slang-ir-strip-cached-dict.h"
+#include "slang-ir-insts.h"
+
+namespace Slang
+{
+
+void stripCachedDictionaries(IRModule* module)
+{
+    List<IRInst*> toRemove;
+    for (auto inst : module->getGlobalInsts())
+    {
+        switch (inst->getOp())
+        {
+        case kIROp_GenericSpecializationDictionary:
+        case kIROp_ExistentialFuncSpecializationDictionary:
+        case kIROp_ExistentialTypeSpecializationDictionary:
+            toRemove.add(inst);
+            break;
+        default:
+            continue;
+        }
+    }
+    for (auto inst : toRemove)
+        inst->removeAndDeallocate();
+}
+
+}
diff --git a/source/slang/slang-ir-strip-cached-dict.h b/source/slang/slang-ir-strip-cached-dict.h
new file mode 100644
index 000000000..95e7787c3
--- /dev/null
+++ b/source/slang/slang-ir-strip-cached-dict.h
@@ -0,0 +1,11 @@
+// slang-ir-strip-cached-dict.h
+#pragma once
+
+namespace Slang
+{
+    struct IRModule;
+    struct IRCall;
+
+        /// Removes specialization dictionaries from module.
+    void stripCachedDictionaries(IRModule* module);
+}
diff --git a/source/slang/slang-ir.h b/source/slang/slang-ir.h
index 403376dca..1ed7e52c7 100644
--- a/source/slang/slang-ir.h
+++ b/source/slang/slang-ir.h
@@ -1728,6 +1728,10 @@ private:
     MemoryArena m_memoryArena;
 };
 
+struct IRSpecializationDictionaryItem : public IRInst
+{
+    IR_LEAF_ISA(SpecializationDictionaryItem)
+};
 
 struct IRDumpOptions
 {