Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Release MLIRContext and Module before executing opt command for memory reduction #3000

Draft
wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

Release MLIRContext and Module before executing opt command for memory reduction #3000

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e3d173d
Release modle before opt command.
imaihal Nov 5, 2024
dfef65f
Merge branch 'main' into release_mod_before_opt
imaihal Nov 7, 2024
2d43a63
Add comments.
imaihal Nov 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 29 additions & 16 deletions src/Compiler/CompilerUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -452,6 +453,14 @@ static int genLLVMBitcode(const mlir::OwningOpRef<ModuleOp> &module,
llvm::WriteBitcodeToFile(*llvmModule, moduleBitcodeStream);
moduleBitcodeStream.flush();

// Free memory before using LLVM `opt` command
// TODO: It might be unexpected that the module is released in this function.
// it is better to restructure functions so that the module is
// automatically released.
llvmModule.reset();
module.release();
context.~MLIRContext();
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you know what are freed/released with these statements? It looks not too much memory saved and not sure it is worth to make a change.


// Use the LLVM's 'opt' command to optimize the bitcode.
std::string optPath = getToolPath("opt");
Command optBitcode(/*exePath=*/optPath);
Expand Down Expand Up @@ -598,10 +607,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,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing const and freeing the object internally seems difficult to follow, at least by looking at the function name. Same for other signature changes in other functions.

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(
Expand All @@ -611,11 +622,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(
Expand All @@ -627,10 +639,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(
Expand Down Expand Up @@ -790,8 +803,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)) {
Expand All @@ -806,7 +819,7 @@ static int emitOutputFiles(std::string outputNameNoExt,
case EmitLib: {
std::string sharedLibNameWithExt;
int rc = compileModuleToSharedLibrary(
module, outputNameNoExt, sharedLibNameWithExt);
module, outputNameNoExt, sharedLibNameWithExt, context);
Copy link
Collaborator

@tungld tungld Nov 8, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@imaihal just be noticed that module and context are potentially used later by outputCode (See Line 813/826). So please make sure outputCode still works since module has been released inside compileModuleToSharedLibrary.

if (rc != CompilerSuccess)
return rc;
if (keepFiles(KeepFilesOfType::MLIR)) {
Expand All @@ -819,7 +832,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)) {
Expand Down
Loading