From 254d4da00e82cf065a0a04743e2028e7ede57e6e Mon Sep 17 00:00:00 2001
From: Anders Leino <aleino@nvidia.com>
Date: Thu, 5 Dec 2024 10:51:54 +0200
Subject: [PATCH] Remove globals if they're left unused after inlining

This closes #5607.
---
 .../slang/slang-ir-legalize-global-values.cpp | 34 +++++++++++++------
 .../slang/slang-ir-legalize-global-values.h   |  3 +-
 source/slang/slang-ir-spirv-legalize.cpp      |  2 +-
 source/slang/slang-ir-wgsl-legalize.cpp       |  2 +-
 4 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/source/slang/slang-ir-legalize-global-values.cpp b/source/slang/slang-ir-legalize-global-values.cpp
index 32103305a6..45f0cdd89c 100644
--- a/source/slang/slang-ir-legalize-global-values.cpp
+++ b/source/slang/slang-ir-legalize-global-values.cpp
@@ -6,9 +6,9 @@
 namespace Slang
 {
 
-void GlobalInstInliningContextGeneric::inlineGlobalValues(IRModule * module)
+void GlobalInstInliningContextGeneric::inlineGlobalValuesAndRemoveIfUnused(IRModule* module)
 {
-    List<IRUse*> globalInstUsesToInline;
+    Dictionary<IRInst*, List<IRUse*>> globalInstUsesToInline;
 
     for (auto globalInst : module->getGlobalInsts())
     {
@@ -17,20 +17,32 @@ void GlobalInstInliningContextGeneric::inlineGlobalValues(IRModule * module)
             for (auto use = globalInst->firstUse; use; use = use->nextUse)
             {
                 if (getParentFunc(use->getUser()) != nullptr)
-                    globalInstUsesToInline.add(use);
+                {
+                    globalInstUsesToInline.getOrAddValue(globalInst, List<IRUse*>()).add(use);
+                }
             }
         }
     }
 
-    for (auto use : globalInstUsesToInline)
+    for (auto globalInstUses : globalInstUsesToInline)
     {
-        auto user = use->getUser();
-        IRBuilder builder(user);
-        builder.setInsertBefore(getOutsideASM(user));
-        IRCloneEnv cloneEnv;
-        auto val = maybeInlineGlobalValue(builder, use->getUser(), use->get(), cloneEnv);
-        if (val != use->get())
-            builder.replaceOperand(use, val);
+        auto inst = globalInstUses.first;
+        auto uses = globalInstUses.second;
+        for (auto use : uses)
+        {
+            auto user = use->getUser();
+            IRBuilder builder(user);
+            builder.setInsertBefore(getOutsideASM(user));
+            IRCloneEnv cloneEnv;
+            auto val = maybeInlineGlobalValue(builder, use->getUser(), use->get(), cloneEnv);
+            if (val != use->get())
+                builder.replaceOperand(use, val);
+        }
+
+        // Since certain globals that appear in the IR are considered illegal for all targets,
+        // e.g. calls to functions, we delete globals which no longer have uses after inlining.
+        if (!inst->hasUses())
+            inst->removeAndDeallocate();
     }
 }
 
diff --git a/source/slang/slang-ir-legalize-global-values.h b/source/slang/slang-ir-legalize-global-values.h
index aca08115fd..8dbafc0c0f 100644
--- a/source/slang/slang-ir-legalize-global-values.h
+++ b/source/slang/slang-ir-legalize-global-values.h
@@ -20,7 +20,8 @@ struct GlobalInstInliningContextGeneric
     virtual IRInst* getOutsideASM(IRInst* beforeInst) =0;
 
     // Inline global values that can't represented by the target to their use sites.
-    void inlineGlobalValues(IRModule * module);
+    // If this leaves any global unused, then remove it.
+    void inlineGlobalValuesAndRemoveIfUnused(IRModule * module);
 
     // Opcodes that can exist in global scope, as long as the operands are.
     bool isLegalGlobalInst(IRInst* inst);
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index 7bbfc2064f..c8c8fb1211 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -2016,7 +2016,7 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
             }
         }
 
-        GlobalInstInliningContext().inlineGlobalValues(m_module);
+        GlobalInstInliningContext().inlineGlobalValuesAndRemoveIfUnused(m_module);
 
         // Some legalization processing may change the function parameter types,
         // so we need to update the function types to match that.
diff --git a/source/slang/slang-ir-wgsl-legalize.cpp b/source/slang/slang-ir-wgsl-legalize.cpp
index 3a300a9792..ff15ba9e9d 100644
--- a/source/slang/slang-ir-wgsl-legalize.cpp
+++ b/source/slang/slang-ir-wgsl-legalize.cpp
@@ -1688,7 +1688,7 @@ void legalizeIRForWGSL(TargetProgram * targetProgram, IRModule* module, Diagnost
     // Go through every instruction in the module and legalize them as needed.
     context.processInst(module->getModuleInst());
 
-    GlobalInstInliningContext().inlineGlobalValues(module);
+    GlobalInstInliningContext().inlineGlobalValuesAndRemoveIfUnused(module);
     simplifyIRForWGSLLegalization(targetProgram, sink, module);
 }