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>

1 files changed, 84 insertions, 4 deletions

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