Add support for `[PrimalSubstitute]` and `[PrimalSubstituteOf]`. (#2691)

* Add support for `[PrimalSubstitute]` and `[PrimalSubstituteOf]`.

* Fix

* Fix.

* Cleanup.

---------

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

1 files changed, 24 insertions, 12 deletions

diff --git a/source/slang/slang-ir-dce.cpp b/source/slang/slang-ir-dce.cpp
index e5c9b1fdb..1fe88e780 100644
--- a/source/slang/slang-ir-dce.cpp
+++ b/source/slang/slang-ir-dce.cpp
@@ -379,25 +379,37 @@ bool shouldInstBeLiveIfParentIsLive(IRInst* inst, IRDeadCodeEliminationOptions o
     //
     if (options.keepExportsAlive)
     {
-        if (inst->findDecoration<IRExportDecoration>())
+        bool isImported = false;
+        bool shouldKeptAliveIfImported = false;
+        IRInst* innerInst = inst;
+        if (auto genInst = as<IRGeneric>(inst))
         {
-            return true;
+            innerInst = findInnerMostGenericReturnVal(genInst);
         }
-        if (inst->findDecoration<IRImportDecoration>())
+        for (auto decor : inst->getDecorations())
         {
-            if (inst->findDecoration<IRForwardDerivativeDecoration>())
-                return true;
-            if (inst->findDecoration<IRUserDefinedBackwardDerivativeDecoration>())
+            switch (decor->getOp())
+            {
+            case kIROp_ExportDecoration:
                 return true;
-            if (auto genInst = as<IRGeneric>(inst))
+            case kIROp_ImportDecoration:
+                isImported = true;
+                break;
+            }
+        }
+        for (auto decor : innerInst->getDecorations())
+        {
+            switch (decor->getOp())
             {
-                auto inner = findInnerMostGenericReturnVal(genInst);
-                if (inner->findDecoration<IRForwardDerivativeDecoration>())
-                    return true;
-                if (inner->findDecoration<IRUserDefinedBackwardDerivativeDecoration>())
-                    return true;
+            case kIROp_ForwardDerivativeDecoration:
+            case kIROp_UserDefinedBackwardDerivativeDecoration:
+            case kIROp_PrimalSubstituteDecoration:
+                shouldKeptAliveIfImported = true;
+                break;
             }
         }
+        if (isImported && shouldKeptAliveIfImported)
+            return true;
     }
 
     if (options.keepLayoutsAlive && inst->findDecoration<IRLayoutDecoration>())