Specialize generic/existential calls within generic functions. (#2294)

* Expose internals of dce and use it to implement call graph walk.

* Specialize calls in generic functions.

* Fix clang error.

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

8 files changed, 204 insertions, 115 deletions

diff --git a/source/slang/slang-ir-any-value-marshalling.cpp b/source/slang/slang-ir-any-value-marshalling.cpp
index b61f0b273..ba8f49cdd 100644
--- a/source/slang/slang-ir-any-value-marshalling.cpp
+++ b/source/slang/slang-ir-any-value-marshalling.cpp
@@ -636,6 +636,7 @@ namespace Slang
             // will become duplicates with new types we introduced (e.g. PtrType(AnyValueStruct)), and therefore
             // invalidates our `globalValueNumberingMap` hash map. We need to rebuild it.
             sharedContext->sharedBuilderStorage.deduplicateAndRebuildGlobalNumberingMap();
+            sharedContext->mapInterfaceRequirementKeyValue.Clear();
         }
     };
 
diff --git a/source/slang/slang-ir-dce.cpp b/source/slang/slang-ir-dce.cpp
index 2d417221b..d58e307da 100644
--- a/source/slang/slang-ir-dce.cpp
+++ b/source/slang/slang-ir-dce.cpp
@@ -6,7 +6,6 @@
 
 namespace Slang
 {
-
 struct DeadCodeEliminationContext
 {
     // This type implements a simple global DCE pass over
@@ -158,20 +157,8 @@ struct DeadCodeEliminationContext
                 // "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;
-                }
-                markInstAsLive(inst->getOperand(ii));
+                if (!isWeakReferenceOperand(inst, ii))
+                    markInstAsLive(inst->getOperand(ii));
             }
 
             // Finally, we need to consider the children
@@ -267,119 +254,146 @@ struct DeadCodeEliminationContext
     //
     bool shouldInstBeLiveIfParentIsLive(IRInst* inst)
     {
-        // The main source of confusion/complexity here is that
-        // we are using the same routine to decide:
-        //
-        // * Should some ordinary instruction in a basic block be kept around?
-        // * Should a basic block in some function be kept around?
-        // * Should a function/type/variable in a module be kept around?
-        //
-        // Still, there are a few basic patterns we can observe.
-        // First, if `inst` is an instruction that might have some effects
-        // when it is executed, then we should keep it around.
-        //
-        if(inst->mightHaveSideEffects())
-            return true;
-        //
-        // The `mightHaveSideEffects` query is conservative, and will
-        // return `true` as its default mode, so once we are past that
-        // query we know that `inst` is either something "structural"
-        // (that makes up the program) rather than executable, or it
-        // is executable but was on an allow-list of things that are
-        // safe to eliminate.
+        return Slang::shouldInstBeLiveIfParentIsLive(inst, options);
+    }
+};
 
-        // Most top-level objects (functions, types, etc.) obviously
-        // do *not* have side effects. That creates the risk that
-        // we'll just go ahead and eliminate every single function/type
-        // in a module. There needs to be a way to identify the
-        // functions we want to keep around, and for right now
-        // that is handled with the `[keepAlive]` decoration.
-        //
-        if(inst->findDecorationImpl(kIROp_KeepAliveDecoration))
-            return true;
-        //
-        // We also consider anything with an `[export(...)]` as live,
-        // when the appropriate option has been set.
-        //
-        // Note: our current approach to linking for back-end compilation
-        // leaves many linakge decorations in place that we seemingly
-        // don't need/want, so this option currently can't be enabled
-        // unconditionally.
-        //
-        if( options.keepExportsAlive )
-        {
-            if( inst->findDecoration<IRExportDecoration>() )
-            {
-                return true;
-            }
-        }
+bool shouldInstBeLiveIfParentIsLive(IRInst* inst, IRDeadCodeEliminationOptions options)
+{
+    // The main source of confusion/complexity here is that
+    // we are using the same routine to decide:
+    //
+    // * Should some ordinary instruction in a basic block be kept around?
+    // * Should a basic block in some function be kept around?
+    // * Should a function/type/variable in a module be kept around?
+    //
+    // Still, there are a few basic patterns we can observe.
+    // First, if `inst` is an instruction that might have some effects
+    // when it is executed, then we should keep it around.
+    //
+    if (inst->mightHaveSideEffects())
+        return true;
+    //
+    // The `mightHaveSideEffects` query is conservative, and will
+    // return `true` as its default mode, so once we are past that
+    // query we know that `inst` is either something "structural"
+    // (that makes up the program) rather than executable, or it
+    // is executable but was on an allow-list of things that are
+    // safe to eliminate.
 
-        if (options.keepLayoutsAlive && inst->findDecoration<IRLayoutDecoration>())
+    // Most top-level objects (functions, types, etc.) obviously
+    // do *not* have side effects. That creates the risk that
+    // we'll just go ahead and eliminate every single function/type
+    // in a module. There needs to be a way to identify the
+    // functions we want to keep around, and for right now
+    // that is handled with the `[keepAlive]` decoration.
+    //
+    if (inst->findDecorationImpl(kIROp_KeepAliveDecoration))
+        return true;
+    //
+    // We also consider anything with an `[export(...)]` as live,
+    // when the appropriate option has been set.
+    //
+    // Note: our current approach to linking for back-end compilation
+    // leaves many linakge decorations in place that we seemingly
+    // don't need/want, so this option currently can't be enabled
+    // unconditionally.
+    //
+    if (options.keepExportsAlive)
+    {
+        if (inst->findDecoration<IRExportDecoration>())
         {
             return true;
         }
+    }
 
-        // A basic block is an interesting case. Knowing that a function
-        // is live means that its entry block is live, but the liveness
-        // of any other blocks is determined by whether they are referenced
-        // by other instructions (e.g., a branch from one block to
-        // another).
+    if (options.keepLayoutsAlive && inst->findDecoration<IRLayoutDecoration>())
+    {
+        return true;
+    }
+
+    // A basic block is an interesting case. Knowing that a function
+    // is live means that its entry block is live, but the liveness
+    // of any other blocks is determined by whether they are referenced
+    // by other instructions (e.g., a branch from one block to
+    // another).
+    //
+    if (auto block = as<IRBlock>(inst))
+    {
+        // To determine whether this is the first block in its
+        // parent function (or what-have-you) we can simply
+        // check if there is a previous block before it.
         //
-        if( auto block = as<IRBlock>(inst) )
-        {
-            // To determine whether this is the first block in its
-            // parent function (or what-have-you) we can simply
-            // check if there is a previous block before it.
-            //
-            auto prevBlock = block->getPrevBlock();
-            return prevBlock == nullptr;
-        }
+        auto prevBlock = block->getPrevBlock();
+        return prevBlock == nullptr;
+    }
 
-        // There are a few special cases of "structural" instructions
-        // that we don't want to eliminate, so we'll check for those next.
+    // There are a few special cases of "structural" instructions
+    // that we don't want to eliminate, so we'll check for those next.
+    //
+    switch (inst->getOp())
+    {
+        // Function parameters obviously shouldn't get eliminated,
+        // even if nothing references them, and block parameters
+        // (phi nodes) will be considered live when their block is,
+        // just so that we don't have to deal with any complications
+        // around re-writing the relevant inter-block argument passing.
         //
-        switch( inst->getOp() )
-        {
-            // Function parameters obviously shouldn't get eliminated,
-            // even if nothing references them, and block parameters
-            // (phi nodes) will be considered live when their block is,
-            // just so that we don't have to deal with any complications
-            // around re-writing the relevant inter-block argument passing.
-            //
-            // TODO: A smarter DCE pass could deal with this case more
-            // carefully, or we could improve the interprocedural SCCP
-            // pass to deal with block parameters instead.
-            //
-        case kIROp_Param:
-            return true;
+        // TODO: A smarter DCE pass could deal with this case more
+        // carefully, or we could improve the interprocedural SCCP
+        // pass to deal with block parameters instead.
+        //
+    case kIROp_Param:
+        return true;
 
-            // IR struct types and witness tables are currently kludged
-            // so that they have child instructions that represent their
-            // entries (effectively `(key,value)` pairs), and those child
-            // instructions are never directly referenced (e.g., an access
-            // to a struct field references the *key* but not the `(key,value)`
-            // pair that is the `IRField` instruction.
-            //
-            // TODO: at some point the IR should use a different representation
-            // for struct types and witness tables that does away with
-            // this problem.
-            //
-        case kIROp_StructField:
-        case kIROp_WitnessTableEntry:
-            return true;
+        // IR struct types and witness tables are currently kludged
+        // so that they have child instructions that represent their
+        // entries (effectively `(key,value)` pairs), and those child
+        // instructions are never directly referenced (e.g., an access
+        // to a struct field references the *key* but not the `(key,value)`
+        // pair that is the `IRField` instruction.
+        //
+        // TODO: at some point the IR should use a different representation
+        // for struct types and witness tables that does away with
+        // this problem.
+        //
+    case kIROp_StructField:
+    case kIROp_WitnessTableEntry:
+        return true;
 
-        default:
-            break;
-        }
+    default:
+        break;
+    }
 
-        // If none of the explicit cases above matched, then we will consider
-        // the instruction to not be live just because its parent is. Further
-        // analysis could still lead to a change in the status of `inst`, if
-        // an instruction that uses it as an operand is marked live.
-        //
-        return false;
+    // If none of the explicit cases above matched, then we will consider
+    // the instruction to not be live just because its parent is. Further
+    // analysis could still lead to a change in the status of `inst`, if
+    // an instruction that uses it as an operand is marked live.
+    //
+    return false;
+}
+
+bool isWeakReferenceOperand(IRInst* inst, UInt operandIndex)
+{
+    // 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(operandIndex)->getOp() == kIROp_WitnessTable)
+            return true;
+        break;
+    case kIROp_SpecializationDictionaryItem:
+        // Ignore all operands of SpecializationDictionaryItem.
+        // This inst is used as a cache and shouldn't hold anything alive.
+        return true;
+    default:
+        break;
     }
-};
+    return false;
+}
 
 // The top-level function for invoking the DCE pass
 // is straighforward. We set up the context object
diff --git a/source/slang/slang-ir-dce.h b/source/slang/slang-ir-dce.h
index aac10a20a..0abf17f9f 100644
--- a/source/slang/slang-ir-dce.h
+++ b/source/slang/slang-ir-dce.h
@@ -1,6 +1,8 @@
 // slang-ir-dce.h
 #pragma once
 
+#include "slang-ir-insts.h"
+
 namespace Slang
 {
     struct IRModule;
@@ -21,4 +23,8 @@ namespace Slang
     bool eliminateDeadCode(
         IRModule*                           module,
         IRDeadCodeEliminationOptions const& options = IRDeadCodeEliminationOptions());
+
+    bool shouldInstBeLiveIfParentIsLive(IRInst* inst, IRDeadCodeEliminationOptions options);
+
+    bool isWeakReferenceOperand(IRInst* inst, UInt operandIndex);
 }
diff --git a/source/slang/slang-ir-generics-lowering-context.h b/source/slang/slang-ir-generics-lowering-context.h
index 475464c8f..f7a8ae0ec 100644
--- a/source/slang/slang-ir-generics-lowering-context.h
+++ b/source/slang/slang-ir-generics-lowering-context.h
@@ -3,6 +3,7 @@
 
 #include "slang-ir.h"
 #include "slang-ir-insts.h"
+#include "slang-ir-dce.h"
 
 #include "slang-ir-lower-generics.h"
 
@@ -49,6 +50,8 @@ namespace Slang
         void addToWorkList(
             IRInst* inst)
         {
+            if (!inst) return;
+
             for (auto ii = inst->getParent(); ii; ii = ii->getParent())
             {
                 if (as<IRGeneric>(ii))
@@ -128,4 +131,45 @@ namespace Slang
             }
         }
     }
+
+    template<typename TFunc>
+    void workOnCallGraph(SharedGenericsLoweringContext* sharedContext, const TFunc& func)
+    {
+        SharedIRBuilder* sharedBuilder = &sharedContext->sharedBuilderStorage;
+        sharedBuilder->init(sharedContext->module);
+
+        sharedContext->addToWorkList(sharedContext->module->getModuleInst());
+        IRDeadCodeEliminationOptions dceOptions;
+        dceOptions.keepExportsAlive = true;
+        dceOptions.keepLayoutsAlive = true;
+
+        while (sharedContext->workList.getCount() != 0)
+        {
+            IRInst* inst = sharedContext->workList.getLast();
+
+            sharedContext->workList.removeLast();
+
+            sharedContext->addToWorkList(inst->parent);
+            sharedContext->addToWorkList(inst->getFullType());
+
+            UInt operandCount = inst->getOperandCount();
+            for (UInt ii = 0; ii < operandCount; ++ii)
+            {
+                if (!isWeakReferenceOperand(inst, ii))
+                    sharedContext->addToWorkList(inst->getOperand(ii));
+            }
+
+            if (auto call = as<IRCall>(inst))
+            {
+                if (func(call))
+                    return;
+            }
+
+            for (auto child = inst->getLastChild(); child; child = child->getPrevInst())
+            {
+                if (shouldInstBeLiveIfParentIsLive(child, dceOptions))
+                    sharedContext->addToWorkList(child);
+            }
+        }
+    }
 }
diff --git a/source/slang/slang-ir-lower-generic-function.cpp b/source/slang/slang-ir-lower-generic-function.cpp
index b50737c23..1bd785e8e 100644
--- a/source/slang/slang-ir-lower-generic-function.cpp
+++ b/source/slang/slang-ir-lower-generic-function.cpp
@@ -315,6 +315,7 @@ namespace Slang
             }
             // Update hash keys of globalNumberingMap, since the types are modified.
             sharedContext->sharedBuilderStorage.deduplicateAndRebuildGlobalNumberingMap();
+            sharedContext->mapInterfaceRequirementKeyValue.Clear();
         }
 
         void processModule()
diff --git a/source/slang/slang-ir-lower-generic-type.cpp b/source/slang/slang-ir-lower-generic-type.cpp
index 76233d155..24a8c38d8 100644
--- a/source/slang/slang-ir-lower-generic-type.cpp
+++ b/source/slang/slang-ir-lower-generic-type.cpp
@@ -79,6 +79,7 @@ namespace Slang
                 }
             }
             sharedContext->sharedBuilderStorage.deduplicateAndRebuildGlobalNumberingMap();
+            sharedContext->mapInterfaceRequirementKeyValue.Clear();
         }
     };
 
diff --git a/source/slang/slang-ir-lower-generics.cpp b/source/slang/slang-ir-lower-generics.cpp
index b8b2ef972..f517719ad 100644
--- a/source/slang/slang-ir-lower-generics.cpp
+++ b/source/slang/slang-ir-lower-generics.cpp
@@ -109,6 +109,7 @@ namespace Slang
             return;
 
         sharedContext->sharedBuilderStorage.deduplicateAndRebuildGlobalNumberingMap();
+        sharedContext->mapInterfaceRequirementKeyValue.Clear();
 
         specializeRTTIObjectReferences(sharedContext);
 
diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index 1d62af47e..1dd18ea47 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -79,6 +79,17 @@ struct SpecializationContext
         if(inst->getOp() == kIROp_InterfaceRequirementEntry)
             return true;
 
+        // A generic parameter is never specialized.
+        switch (inst->getOp())
+        {
+        case kIROp_GlobalGenericParam:
+            return false;
+        case kIROp_Param:
+            if (inst->getParent() && inst->getParent()->getOp() == kIROp_Block &&
+                inst->getParent()->getParent() &&
+                inst->getParent()->getParent()->getOp() == kIROp_Generic)
+                return false;
+        }
         return fullySpecializedInsts.Contains(inst);
     }
 
@@ -113,6 +124,14 @@ struct SpecializationContext
     void addToWorkList(
         IRInst* inst)
     {
+#if 0
+        // Note(Yong): we should no longer ignore generic functions
+        // because they maybe called via dynamic dispatch.
+        // We still want to specialize calls inside a generic function
+        // if we can derive its type at compile time. The following
+        // skipping logic is disabled and we should consider remove it.
+        //
+        // 
         // We will ignore any code that is nested under a generic,
         // because it doesn't make sense to perform specialization
         // on such code.
@@ -122,6 +141,7 @@ struct SpecializationContext
             if(as<IRGeneric>(ii))
                 return;
         }
+#endif
 
         if (workList.Add(inst))
         {
@@ -1203,6 +1223,7 @@ struct SpecializationContext
             case kIROp_Call:
             case kIROp_ExtractExistentialType:
             case kIROp_CreateExistentialObject:
+            case kIROp_Param:
                 return false;
             default:
                 break;