Release modle before opt command.

Signed-off-by: Haruki Imai <[email protected]> Co-authored-by: Yasushi Negishi <[email protected]>
onnx · Nov 5, 2024 · e3d173d · e3d173d
1 parent d25c63e
commit e3d173d
Showing 1 changed file with 26 additions and 16 deletions.
diff --git a/src/Compiler/CompilerUtils.cpp b/src/Compiler/CompilerUtils.cpp
@@ -396,8 +396,9 @@ std::string getTargetFilename(
 
 // Write LLVM optimized bitcode.
 // Returns 0 on success, error code on failure.
-static int genLLVMBitcode(const mlir::OwningOpRef<ModuleOp> &module,
-    std::string outputNameNoExt, std::string optimizedBitcodeNameWithExt) {
+static int genLLVMBitcode(mlir::OwningOpRef<ModuleOp> &module,
+    std::string outputNameNoExt, std::string optimizedBitcodeNameWithExt,
+    mlir::MLIRContext &context) {
   std::string msg =
       "Translating MLIR Module to LLVM and Generating LLVM Optimized Bitcode";
   showCompilePhase(msg);
@@ -452,6 +453,11 @@ static int genLLVMBitcode(const mlir::OwningOpRef<ModuleOp> &module,
   llvm::WriteBitcodeToFile(*llvmModule, moduleBitcodeStream);
   moduleBitcodeStream.flush();
 
+  // Free memory before using LLVM `opt` command
+  llvmModule.reset();
+  module.release();
+  context.~MLIRContext();
+
   // Use the LLVM's 'opt' command to optimize the bitcode.
   std::string optPath = getToolPath("opt");
   Command optBitcode(/*exePath=*/optPath);
@@ -598,10 +604,12 @@ static int genJniJar(const mlir::OwningOpRef<ModuleOp> &module,
 }
 
 // Return 0 on success, error code on failure
-static int compileModuleToObject(const mlir::OwningOpRef<ModuleOp> &module,
-    std::string outputNameWithoutExt, std::string &objectNameWithExt) {
+static int compileModuleToObject(mlir::OwningOpRef<ModuleOp> &module,
+    std::string outputNameWithoutExt, std::string &objectNameWithExt,
+    mlir::MLIRContext &context) {
   std::string bitcodeNameWithExt = outputNameWithoutExt + ".bc";
-  int rc = genLLVMBitcode(module, outputNameWithoutExt, bitcodeNameWithExt);
+  int rc =
+      genLLVMBitcode(module, outputNameWithoutExt, bitcodeNameWithExt, context);
   if (rc != CompilerSuccess)
     return rc;
   llvm::FileRemover bitcodeRemover(
@@ -611,11 +619,12 @@ static int compileModuleToObject(const mlir::OwningOpRef<ModuleOp> &module,
 }
 
 // Return 0 on success, error code on failure
-static int compileModuleToSharedLibrary(
-    const mlir::OwningOpRef<ModuleOp> &module, std::string outputNameNoExt,
-    std::string &libNameWithExt) {
+static int compileModuleToSharedLibrary(mlir::OwningOpRef<ModuleOp> &module,
+    std::string outputNameNoExt, std::string &libNameWithExt,
+    mlir::MLIRContext &context) {
   std::string modelObjNameWithExt;
-  int rc = compileModuleToObject(module, outputNameNoExt, modelObjNameWithExt);
+  int rc = compileModuleToObject(
+      module, outputNameNoExt, modelObjNameWithExt, context);
   if (rc != CompilerSuccess)
     return rc;
   llvm::FileRemover modelObjRemover(
@@ -627,10 +636,11 @@ static int compileModuleToSharedLibrary(
 }
 
 // Return 0 on success, error code on failure
-static int compileModuleToJniJar(
-    const mlir::OwningOpRef<ModuleOp> &module, std::string outputNameNoExt) {
+static int compileModuleToJniJar(mlir::OwningOpRef<ModuleOp> &module,
+    std::string outputNameNoExt, mlir::MLIRContext &context) {
   std::string modelObjNameWithExt;
-  int rc = compileModuleToObject(module, outputNameNoExt, modelObjNameWithExt);
+  int rc = compileModuleToObject(
+      module, outputNameNoExt, modelObjNameWithExt, context);
   if (rc != CompilerSuccess)
     return rc;
   llvm::FileRemover modelObjRemover(
@@ -790,8 +800,8 @@ static int emitOutputFiles(std::string outputNameNoExt,
   switch (emissionTarget) {
   case EmitObj: {
     std::string modelObjNameWithExt;
-    int rc =
-        compileModuleToObject(module, outputNameNoExt, modelObjNameWithExt);
+    int rc = compileModuleToObject(
+        module, outputNameNoExt, modelObjNameWithExt, context);
     if (rc != CompilerSuccess)
       return rc;
     if (keepFiles(KeepFilesOfType::MLIR)) {
@@ -806,7 +816,7 @@ static int emitOutputFiles(std::string outputNameNoExt,
   case EmitLib: {
     std::string sharedLibNameWithExt;
     int rc = compileModuleToSharedLibrary(
-        module, outputNameNoExt, sharedLibNameWithExt);
+        module, outputNameNoExt, sharedLibNameWithExt, context);
     if (rc != CompilerSuccess)
       return rc;
     if (keepFiles(KeepFilesOfType::MLIR)) {
@@ -819,7 +829,7 @@ static int emitOutputFiles(std::string outputNameNoExt,
                    << "' has been compiled.\n";
   } break;
   case EmitJNI: {
-    int rc = compileModuleToJniJar(module, outputNameNoExt);
+    int rc = compileModuleToJniJar(module, outputNameNoExt, context);
     if (rc != CompilerSuccess)
       return rc;
     if (keepFiles(KeepFilesOfType::MLIR)) {