Cleanup global memory allocations on shutdown. (#4731)

Co-authored-by: Jay Kwak <82421531+jkwak-work@users.noreply.github.com>

15 files changed, 50 insertions, 12 deletions

diff --git a/include/slang.h b/include/slang.h
index 6fd6da7b2..849532fc8 100644
--- a/include/slang.h
+++ b/include/slang.h
@@ -5351,6 +5351,14 @@ Returns nullptr if there isn't an embedded stdlib.
 */
 SLANG_API ISlangBlob* slang_getEmbeddedStdLib();
 
+
+/* Cleanup all global allocations used by Slang, to prevent memory leak detectors from
+ reporting them as leaks. This function should only be called after all Slang objects
+ have been released. No other Slang functions such as `createGlobalSession`
+ should be called after this function.
+ */
+SLANG_EXTERN_C SLANG_API void slang_shutdown();
+
 namespace slang
 {
     inline SlangResult createGlobalSession(
@@ -5358,6 +5366,7 @@ namespace slang
     {
         return slang_createGlobalSession(SLANG_API_VERSION, outGlobalSession);
     }
+    inline void shutdown() { slang_shutdown(); }
 }
 
 /** @see slang::ICompileRequest::getProgram
diff --git a/source/compiler-core/slang-spirv-core-grammar.h b/source/compiler-core/slang-spirv-core-grammar.h
index 7fe8f6cd7..958aaaef2 100644
--- a/source/compiler-core/slang-spirv-core-grammar.h
+++ b/source/compiler-core/slang-spirv-core-grammar.h
@@ -16,7 +16,8 @@ namespace Slang
     struct SPIRVCoreGrammarInfo : public RefObject
     {
         static RefPtr<SPIRVCoreGrammarInfo> loadFromJSON(SourceView& source, DiagnosticSink& sink);
-        static RefPtr<SPIRVCoreGrammarInfo> getEmbeddedVersion();
+        static RefPtr<SPIRVCoreGrammarInfo>& getEmbeddedVersion();
+        static inline void freeEmbeddedGrammerInfo() { getEmbeddedVersion() = nullptr; }
 
         template<typename K, typename T>
         struct Lookup
diff --git a/source/core/slang-performance-profiler.cpp b/source/core/slang-performance-profiler.cpp
index faf06bb92..b480e1c8e 100644
--- a/source/core/slang-performance-profiler.cpp
+++ b/source/core/slang-performance-profiler.cpp
@@ -42,6 +42,10 @@ namespace Slang
         {
             data.clear();
         }
+        virtual void dispose() override
+        {
+            data = decltype(data)();
+        }
     };
 
     PerformanceProfiler* Slang::PerformanceProfiler::getProfiler()
diff --git a/source/core/slang-performance-profiler.h b/source/core/slang-performance-profiler.h
index 372a7b273..71b34d262 100644
--- a/source/core/slang-performance-profiler.h
+++ b/source/core/slang-performance-profiler.h
@@ -29,6 +29,7 @@ public:
     virtual void exitFunction(FuncProfileContext context) = 0;
     virtual void getResult(StringBuilder& out) = 0;
     virtual void clear() = 0;
+    virtual void dispose() = 0;
 public:
     static PerformanceProfiler* getProfiler();
 };
diff --git a/source/slang/slang-api.cpp b/source/slang/slang-api.cpp
index c4f668fb2..412cec0d4 100644
--- a/source/slang/slang-api.cpp
+++ b/source/slang/slang-api.cpp
@@ -1,9 +1,10 @@
 // slang-api.cpp
 
 #include "slang-compiler.h"
-
 #include "slang-repro.h"
-
+#include "slang-capability.h"
+#include "../core/slang-rtti-info.h"
+#include "../core/slang-performance-profiler.h"
 #include "../core/slang-shared-library.h"
 #include "../slang-record-replay/record/slang-global-session.h"
 #include "../slang-record-replay/util/record-utility.h"
@@ -139,6 +140,14 @@ SLANG_API SlangResult slang_createGlobalSession(
     return SLANG_OK;
 }
 
+SLANG_API void slang_shutdown()
+{
+    Slang::PerformanceProfiler::getProfiler()->dispose();
+    Slang::SPIRVCoreGrammarInfo::freeEmbeddedGrammerInfo();
+    Slang::RttiInfo::deallocateAll();
+    Slang::freeCapabilityDefs();
+}
+
 SLANG_API SlangResult slang_createGlobalSessionWithoutStdLib(
     SlangInt                apiVersion,
     slang::IGlobalSession** outGlobalSession)
diff --git a/source/slang/slang-capability.h b/source/slang/slang-capability.h
index 299e956e4..79b1f7068 100644
--- a/source/slang/slang-capability.h
+++ b/source/slang/slang-capability.h
@@ -371,6 +371,8 @@ const CapabilityAtomSet& getAtomSetOfStages();
 
 bool hasTargetAtom(const CapabilityAtomSet& setIn, CapabilityAtom& targetAtom);
 
+void freeCapabilityDefs();
+
 //#define UNIT_TEST_CAPABILITIES
 #ifdef UNIT_TEST_CAPABILITIES
 void TEST_CapabilitySet();
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index c6e460d02..6f5e79458 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -140,7 +140,7 @@ void Session::init()
 {
     SLANG_ASSERT(BaseTypeInfo::check());
 
-    
+
     _initCodeGenTransitionMap();
 
     ::memset(m_downstreamCompilerLocators, 0, sizeof(m_downstreamCompilerLocators));
diff --git a/source/slangc/main.cpp b/source/slangc/main.cpp
index e69a6655d..2907625ed 100644
--- a/source/slangc/main.cpp
+++ b/source/slangc/main.cpp
@@ -104,7 +104,7 @@ SLANG_TEST_TOOL_API SlangResult innerMain(StdWriters* stdWriters, slang::IGlobal
     SlangResult res = _compile(compileRequest, argc, argv);
     // Now that we are done, clean up after ourselves
     spDestroyCompileRequest(compileRequest);
-
+    
     return res;
 }
 
@@ -112,6 +112,7 @@ int MAIN(int argc, char** argv)
 {
     auto stdWriters = StdWriters::initDefaultSingleton();
     SlangResult res = innerMain(stdWriters, nullptr, argc, argv);
+    slang::shutdown();
     return (int)TestToolUtil::getReturnCode(res);
 }
 
diff --git a/tools/render-test/render-test-main.cpp b/tools/render-test/render-test-main.cpp
index 0cd94c046..6f67dfd1e 100644
--- a/tools/render-test/render-test-main.cpp
+++ b/tools/render-test/render-test-main.cpp
@@ -1545,6 +1545,7 @@ int main(int argc, char**  argv)
     SlangResult res = innerMain(stdWriters, session, argc, argv);
     spDestroySession(session);
 
+    slang::shutdown();
 	return (int)TestToolUtil::getReturnCode(res);
 }
 
diff --git a/tools/slang-capability-generator/capability-generator-main.cpp b/tools/slang-capability-generator/capability-generator-main.cpp
index b08496e35..e4a4885c6 100644
--- a/tools/slang-capability-generator/capability-generator-main.cpp
+++ b/tools/slang-capability-generator/capability-generator-main.cpp
@@ -770,11 +770,11 @@ SlangResult generateDefinitions(DiagnosticSink* sink, List<RefPtr<CapabilityDef>
         }
     }
     outputUIntSetGenerator("generatorOf_kAnyTargetUIntSetBuffer", anyTargetUIntSetHash, anyTargetAtomSet);
-    anyTargetUIntSetHash << "const static CapabilityAtomSet kAnyTargetUIntSetBuffer = generatorOf_kAnyTargetUIntSetBuffer();\n";
+    anyTargetUIntSetHash << "static CapabilityAtomSet kAnyTargetUIntSetBuffer = generatorOf_kAnyTargetUIntSetBuffer();\n";
     sbCpp << anyTargetUIntSetHash;
 
     outputUIntSetGenerator("generatorOf_kAnyStageUIntSetBuffer", anyStageUIntSetHash, anyStageAtomSet);
-    anyStageUIntSetHash << "const static CapabilityAtomSet kAnyStageUIntSetBuffer = generatorOf_kAnyStageUIntSetBuffer();\n";
+    anyStageUIntSetHash << "static CapabilityAtomSet kAnyStageUIntSetBuffer = generatorOf_kAnyStageUIntSetBuffer();\n";
     sbCpp << anyStageUIntSetHash;
 
     sbHeader << "\nenum {\n";
@@ -899,6 +899,14 @@ SlangResult generateDefinitions(DiagnosticSink* sink, List<RefPtr<CapabilityDef>
     }
     
     sbCpp << "};\n";
+
+    sbCpp
+        << "void freeCapabilityDefs()\n"
+        << "{\n"
+        << "    for (auto& cap : kCapabilityArray) { cap = CapabilityAtomSet(); }\n"
+        << "    kAnyTargetUIntSetBuffer = CapabilityAtomSet();\n"
+        << "    kAnyStageUIntSetBuffer = CapabilityAtomSet();\n"
+        << "}\n";
     return SLANG_OK;
 }
 
diff --git a/tools/slang-reflection-test/slang-reflection-test-main.cpp b/tools/slang-reflection-test/slang-reflection-test-main.cpp
index be48d8e92..8866661b8 100644
--- a/tools/slang-reflection-test/slang-reflection-test-main.cpp
+++ b/tools/slang-reflection-test/slang-reflection-test-main.cpp
@@ -1446,6 +1446,6 @@ int main(
     
     SlangResult res = innerMain(stdWriters, session, argc, argv);
     spDestroySession(session);
-
+    slang::shutdown();
     return SLANG_FAILED(res) ? 1 : 0;
 }
diff --git a/tools/slang-spirv-embed-generator/spirv-embed-generator-main.cpp b/tools/slang-spirv-embed-generator/spirv-embed-generator-main.cpp
index 20431e872..fd29ed714 100644
--- a/tools/slang-spirv-embed-generator/spirv-embed-generator-main.cpp
+++ b/tools/slang-spirv-embed-generator/spirv-embed-generator-main.cpp
@@ -418,7 +418,7 @@ void writeInfo(
     //
     // Now write out the function which holds onto the static embedded info table
     //
-    line("RefPtr<SPIRVCoreGrammarInfo> SPIRVCoreGrammarInfo::getEmbeddedVersion()");
+    line("RefPtr<SPIRVCoreGrammarInfo>& SPIRVCoreGrammarInfo::getEmbeddedVersion()");
     line("{");
     line("    static RefPtr<SPIRVCoreGrammarInfo> embedded = [](){");
     line("        RefPtr<SPIRVCoreGrammarInfo> info = new SPIRVCoreGrammarInfo();");
diff --git a/tools/slang-test/slang-test-main.cpp b/tools/slang-test/slang-test-main.cpp
index f9e07bdc9..cd3ca3d54 100644
--- a/tools/slang-test/slang-test-main.cpp
+++ b/tools/slang-test/slang-test-main.cpp
@@ -4618,7 +4618,7 @@ int main(int argc, char** argv)
 {
     const SlangResult res = innerMain(argc, argv);
 
-    Slang::RttiInfo::deallocateAll();
+    slang::shutdown();
 
 #ifdef _MSC_VER
     _CrtDumpMemoryLeaks();
diff --git a/tools/slangd/main.cpp b/tools/slangd/main.cpp
index 7727ec135..42d6e0c1b 100644
--- a/tools/slangd/main.cpp
+++ b/tools/slangd/main.cpp
@@ -23,5 +23,7 @@ int main(int argc, const char* const* argv)
     }
     Slang::LanguageServerStartupOptions options;
     options.parse(argc, argv);
-    return Slang::runLanguageServer(options);
+    auto result =  Slang::runLanguageServer(options);
+    slang::shutdown();
+    return result;
 }
diff --git a/tools/test-server/test-server-main.cpp b/tools/test-server/test-server-main.cpp
index 750fb7faf..ba77d9c84 100644
--- a/tools/test-server/test-server-main.cpp
+++ b/tools/test-server/test-server-main.cpp
@@ -563,7 +563,7 @@ SlangResult _execute(int argc, const char* const* argv)
     TestServer server;
     SLANG_RETURN_ON_FAIL(server.init(argc, argv));
     SLANG_RETURN_ON_FAIL(server.execute());
-
+    slang::shutdown();
     return SLANG_OK;
 }