Merge pull request shader-slang#417 from csyonghe/leakfix

Fix IR memory leaks.

Author: csyonghe

slang.h

@@ -1237,6 +1237,7 @@ namespace slang
 #include "source/slang/dxc-support.cpp"
 #include "source/slang/emit.cpp"
 #include "source/slang/ir.cpp"
+#include "source/slang/memory_pool.cpp"
 #include "source/slang/ir-legalize-types.cpp"
 #include "source/slang/ir-ssa.cpp"
 #include "source/slang/legalize-types.cpp"

source/core/smart-pointer.h

@@ -35,6 +35,11 @@ namespace Slang
             referenceCount++;
         }
 
+        void decreaseReference()
+        {
+            --referenceCount;
+        }
+
         void releaseReference()
         {
             SLANG_ASSERT(referenceCount != 0);
@@ -192,6 +197,15 @@ namespace Slang
             return pointer;
         }
 
+        T* detach()
+        {
+            if (pointer)
+                dynamic_cast<RefObject*>(pointer)->decreaseReference();
+            auto rs = pointer;
+            pointer = nullptr;
+            return rs;
+        }
+
     private:
         T* pointer;
 

source/slang/compiler.h

@@ -179,7 +179,7 @@ namespace Slang
         // This will only be valid/non-null after semantic
         // checking and IR generation are complete, so it
         // is not safe to use this field without testing for NULL.
-        IRModule* irModule;
+        RefPtr<IRModule> irModule;
     };
 
     // A request to generate output in some target format
@@ -225,7 +225,7 @@ namespace Slang
         RefPtr<ModuleDecl>  moduleDecl;
 
         // The IR for the module
-        IRModule* irModule = nullptr;
+        RefPtr<IRModule> irModule = nullptr;
     };
 
     class Session;
@@ -313,6 +313,8 @@ namespace Slang
         // Map from the logical name of a module to its definition
         Dictionary<Name*, RefPtr<LoadedModule>> mapNameToLoadedModules;
 
+        // The resulting specialized IR module for each entry point request
+        List<RefPtr<IRModule>> compiledModules;
 
         CompileRequest(Session* session);
 
@@ -450,7 +452,6 @@ namespace Slang
 
         Dictionary<int, RefPtr<Type>> builtinTypes;
         Dictionary<String, Decl*> magicDecls;
-        List<RefPtr<Type>> canonicalTypes;
 
         void initializeTypes();
 
@@ -505,6 +506,7 @@ namespace Slang
             RefPtr<Scope> const&    scope,
             String const&           path,
             String const&           source);
+        ~Session();
     };
 
 }

source/slang/emit.cpp

@@ -8084,17 +8084,17 @@ String emitEntryPoint(
 
     EmitVisitor visitor(&context);
 
+    // We are going to create a fresh IR module that we will use to
+    // clone any code needed by the user's entry point.
+    IRSpecializationState* irSpecializationState = createIRSpecializationState(
+        entryPoint,
+        programLayout,
+        target,
+        targetRequest);
     {
         TypeLegalizationContext typeLegalizationContext;
         typeLegalizationContext.session = entryPoint->compileRequest->mSession;
 
-        // We are going to create a fresh IR module that we will use to
-        // clone any code needed by the user's entry point.
-        IRSpecializationState* irSpecializationState = createIRSpecializationState(
-            entryPoint,
-            programLayout,
-            target,
-            targetRequest);
         IRModule* irModule = getIRModule(irSpecializationState);
 
         typeLegalizationContext.irModule = irModule;
@@ -8158,16 +8158,19 @@ String emitEntryPoint(
         // TODO: do we want to emit directly from IR, or translate the
         // IR back into AST for emission?
         visitor.emitIRModule(&context, irModule);
+        
+        // retain the specialized ir module, because the current
+        // GlobalGenericParamSubstitution implementation may reference ir objects 
+        targetRequest->compileRequest->compiledModules.Add(irModule);
     }
+    destroyIRSpecializationState(irSpecializationState);
 
     String code = sharedContext.sb.ProduceString();
     sharedContext.sb.Clear();
 
     // Now that we've emitted the code for all the declaratiosn in the file,
     // it is time to stich together the final output.
 
-
-
     // There may be global-scope modifiers that we should emit now
     visitor.emitGLSLPreprocessorDirectives(translationUnitSyntax);
     String prefix = sharedContext.sb.ProduceString();
@@ -8189,7 +8192,7 @@ String emitEntryPoint(
     finalResultBuilder << code;
 
     String finalResult = finalResultBuilder.ProduceString();
-
+    
     return finalResult;
 }
 

source/slang/ir-insts.h

@@ -33,6 +33,11 @@ struct IRLayoutDecoration : IRDecoration
     enum { kDecorationOp = kIRDecorationOp_Layout };
 
     RefPtr<Layout>  layout;
+    virtual void dispose() override
+    {
+        IRDecoration::dispose();
+        layout = nullptr;
+    }
 };
 
 enum IRLoopControl
@@ -52,6 +57,11 @@ struct IRTargetSpecificDecoration : IRDecoration
 {
     // TODO: have a more structured representation of target specifiers
     String targetName;
+    virtual void dispose()override
+    {
+        IRDecoration::dispose();
+        targetName = String();
+    }
 };
 
 struct IRTargetDecoration : IRTargetSpecificDecoration
@@ -64,6 +74,11 @@ struct IRTargetIntrinsicDecoration : IRTargetSpecificDecoration
     enum { kDecorationOp = kIRDecorationOp_TargetIntrinsic };
 
     String definition;
+    virtual void dispose()override
+    {
+        IRTargetSpecificDecoration::dispose();
+        definition = String();
+    }
 };
 
 //
@@ -73,6 +88,11 @@ struct IRTargetIntrinsicDecoration : IRTargetSpecificDecoration
 struct IRDeclRef : IRValue
 {
     DeclRef<Decl> declRef;
+    virtual void dispose() override
+    {
+        IRValue::dispose();
+        declRef = decltype(declRef)();
+    }
 };
 
 // An instruction that specializes another IR value
@@ -333,6 +353,13 @@ struct IRWitnessTable : IRGlobalValue
     RefPtr<GenericDecl> genericDecl;
     DeclRef<Decl> subTypeDeclRef, supTypeDeclRef;
     IRValueList<IRWitnessTableEntry> entries;
+    virtual void dispose() override
+    {
+        IRGlobalValue::dispose();
+        genericDecl = decltype(genericDecl)();
+        subTypeDeclRef = decltype(subTypeDeclRef)();
+        supTypeDeclRef = decltype(supTypeDeclRef)();
+    }
 };
 
 // An instruction that yields an undefined value.
@@ -610,22 +637,26 @@ struct IRBuilder
         UInt            caseArgCount,
         IRValue* const* caseArgs);
 
-
-    IRDecoration* addDecorationImpl(
-        IRValue*        value,
-        UInt            decorationSize,
-        IRDecorationOp  op);
-
     template<typename T>
     T* addDecoration(IRValue* value, IRDecorationOp op)
     {
-        return (T*) addDecorationImpl(value, sizeof(T), op);
+        assert(getModule());
+        auto decorationSize = sizeof(T);
+        auto decoration = (T*)getModule()->memoryPool.allocZero(decorationSize);
+        new(decoration)T();
+       
+        decoration->op = op;
+
+        decoration->next = value->firstDecoration;
+        value->firstDecoration = decoration;
+        getModule()->irObjectsToFree.Add(decoration);
+        return decoration;
     }
 
     template<typename T>
     T* addDecoration(IRValue* value)
     {
-        return (T*) addDecorationImpl(value, sizeof(T), IRDecorationOp(T::kDecorationOp));
+        return addDecoration<T>(value, IRDecorationOp(T::kDecorationOp));
     }
 
     IRHighLevelDeclDecoration* addHighLevelDeclDecoration(IRValue* value, Decl* decl);