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

Option to set the number of threads for parallel compilation #3048

Merged
merged 15 commits into from
Jan 30, 2025
Merged

Option to set the number of threads for parallel compilation #3048

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
ff4bbf1
Add -j option for specifying the number of threads for compilation.
imaihal Jan 21, 2025
da52a06
Update copyright.
imaihal Jan 21, 2025
2414108
Update the way to create pool.
imaihal Jan 23, 2025
1725657
Change default value
imaihal Jan 24, 2025
4c453c3
Merge branch 'main' into option_to_set_num_threads_for_compilation
imaihal Jan 24, 2025
e98f2cd
Change '-j' option as common option.
imaihal Jan 27, 2025
6e6e249
Fix duplicated definitions.
imaihal Jan 27, 2025
470ee6e
Fix duplicated definition again.
imaihal Jan 27, 2025
76119b3
Set the value of option as uint64_t.
imaihal Jan 27, 2025
b176efd
Merge branch 'main' into option_to_set_num_threads_for_compilation
imaihal Jan 28, 2025
e8bc8bf
Set number of threads in onnx-mlir-opt.
imaihal Jan 29, 2025
f5f880d
Set nullptr for initializing pool.
imaihal Jan 29, 2025
2331d91
Restore include sentence for Thread.h
imaihal Jan 29, 2025
bc5f0d3
Merge branch 'main' into option_to_set_num_threads_for_compilation
imaihal Jan 29, 2025
0dfc4e4
Merge branch 'main' into option_to_set_num_threads_for_compilation
imaihal Jan 30, 2025
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
10 changes: 9 additions & 1 deletion src/Compiler/CompilerOptions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

//===------------------------ CompilerOptions.cpp -------------------------===//
//
// Copyright 2022, 2024 The IBM Research Authors.
// Copyright 2022-2025 The IBM Research Authors.
//
// =============================================================================
//
Expand Down Expand Up @@ -92,6 +92,7 @@ OptReport optReport; // onnx-mlir only
bool useOldBufferization; // onnx-mlir only
bool enableTiming; // onnx-mlir only
bool enableBoundCheck; // onnx-mlir only
int compilationNumThreads; // onnx-mlir only
Copy link
Collaborator

Choose a reason for hiding this comment

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

Can we apply this option to onnx-mlir-opt as well so that we can check performance by calling onnx-mlir-opt --constprop-onnx -j8 for a single operation? It's ok if it needs another PR.

Copy link
Collaborator

Choose a reason for hiding this comment

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

@tungld I was thinking along the same lines, also for example to test that we get the same result for a test that uses parallelism.

@imaihal I would then add a test that tries the operation that you parallelized and try it with -j 4 for example. You may have to copy a large constant to do that... can lift it from a model.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

@tungld @AlexandreEichenberger Thanks for the comments.

I applied this option to onnx-mlir-opt, but parallelization of constprop is not included in this PR (It is included in PR #3042). So I want to add the tests in PR #3042.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Absolutely, this PR set the options, and it will be tested once there is a new opt that can use parallelism.

Copy link
Collaborator

@tungld tungld Jan 28, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

@imaihal I didn't see any change to onnx-mlir-opt.cpp (e.g. MLIRContext in onnx-mlir-opt.cpp). I am curious on how this option is effective inonnx-mlir-opt command line?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

@tungld New thread pool is created in MlirOptMain.cpp of MLIR. This is due to handle-split-input-file option (https://github.com/llvm/llvm-project/blob/main/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp#L609-L613). So, it seems to be difficult to set our own threadpool in onnx-mlir-opt.cpp.

In lit tests, how about using--mlir-disable-threading for single thread and using mlir default(use all cpus) for multi-thread? Or, using pass local option might be possible to specify our own thread pool.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

@tungld Sorry, it is possible just by modifying onnx-mlir-opt.cpp. I updated.

bool split_input_file; // onnx-mlir-opt only
bool verify_diagnostics; // onnx-mlir-opt only
bool verify_passes; // onnx-mlir-opt only
Expand Down Expand Up @@ -665,6 +666,13 @@ static llvm::cl::opt<bool, true> enable_bound_check("enable-bound-check",
llvm::cl::location(enableBoundCheck), llvm::cl::init(false),
llvm::cl::cat(OnnxMlirOptions));

static llvm::cl::opt<int, true> compilation_num_threads("j",
AlexandreEichenberger marked this conversation as resolved.
Show resolved Hide resolved
llvm::cl::desc(
"Enable parallel compilation with <int> concurrent processes. "
"<int> threads are created for parallel compilation.\n"),
llvm::cl::location(compilationNumThreads), llvm::cl::init(1),
llvm::cl::cat(OnnxMlirOptions));

#if defined(_DEBUG)
// Option only available in debug mode: set using command options.
static llvm::cl::opt<bool, true> test_compiler_opt("test-compiler-opt",
Expand Down
3 changes: 2 additions & 1 deletion src/Compiler/CompilerOptions.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

//===------------------------ CompilerOptions.hpp -------------------------===//
//
// Copyright 2022-2024 The IBM Research Authors.
// Copyright 2022-2025 The IBM Research Authors.
//
// =============================================================================
//
Expand Down Expand Up @@ -139,6 +139,7 @@ extern bool useOldBufferization; // onnx-mlir only
extern bool enableTiming; // onnx-mlir only
extern bool enableBoundCheck; // onnx-mlir only
extern bool debugTestCompilerOpt; // onnx-mlir only
extern int compilationNumThreads; // onnx-mlir only

extern bool split_input_file; // onnx-mlir-opt only
extern bool verify_diagnostics; // onnx-mlir-opt only
Expand Down
12 changes: 10 additions & 2 deletions src/onnx-mlir.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

//===------------------ onnx-mlir.cpp - Compiler Driver ------------------===//
//
// Copyright 2019-2022 The IBM Research Authors.
// Copyright 2019-2025 The IBM Research Authors.
//
// =============================================================================
// Main function for onnx-mlir.
Expand All @@ -15,6 +15,7 @@
#include <regex>

#include "mlir/IR/AsmState.h"
#include "mlir/IR/Threading.h"
#include "mlir/Support/Timing.h"
#include "src/Compiler/CompilerOptions.hpp"
#include "src/Compiler/CompilerUtils.hpp"
Expand Down Expand Up @@ -68,7 +69,14 @@ int main(int argc, char *argv[]) {
}

// Create context after MLIRContextCLOptions are registered and parsed.
mlir::MLIRContext context;
mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
std::unique_ptr<llvm::ThreadPoolInterface> threadPoolPtr;
imaihal marked this conversation as resolved.
Show resolved Hide resolved
if (compilationNumThreads > 1) {
threadPoolPtr = std::make_unique<llvm::DefaultThreadPool>(
llvm::hardware_concurrency(compilationNumThreads));
context.setThreadPool(*threadPoolPtr);
}

if (!context.isMultithreadingEnabled()) {
assert(context.getNumThreads() == 1 && "1 thread if no multithreading");
LLVM_DEBUG(llvm::dbgs() << "multithreading is disabled\n");
Expand Down
Loading