From 86fc50c5092fbccf6072dcf7bbdfafb8915f02c8 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Wed, 8 Mar 2023 21:52:34 -0800
Subject: Add support for `[PrimalSubstitute]` and `[PrimalSubstituteOf]`.
 (#2691)

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

* Fix

* Fix.

* Cleanup.

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 source/slang/slang-ir-dce.cpp | 36 ++++++++++++++++++++++++------------
 1 file changed, 24 insertions(+), 12 deletions(-)

(limited to 'source/slang/slang-ir-dce.cpp')

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>())
-- 
cgit v1.2.3