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

[WIP] Add umbrella optimization flags #20047

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open

[WIP] Add umbrella optimization flags #20047

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e25afda
Add functionality and tests
IanWood1 Feb 13, 2025
a9147ad
Fix tests and add another
IanWood1 Feb 21, 2025
11d6258
Use optimization for global optimization phase
IanWood1 Feb 21, 2025
882ee69
Add basic iree-compile test
IanWood1 Feb 21, 2025
1f8ae49
Add init_at_opt
IanWood1 Feb 22, 2025
a9c362b
Fix template and clean up
IanWood1 Feb 27, 2025
1a96192
Fix asan
IanWood1 Feb 27, 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
13 changes: 13 additions & 0 deletions compiler/bindings/python/test/api/api_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,19 @@ def testFlagError(self):
with self.assertRaises(ValueError):
session.set_flags("--does-not-exist=1")

def testOptFlags(self):
session = Session()
flags = session.get_flags()
self.assertIn("--iree-opt-level=O0", flags)
self.assertIn("--iree-global-optimization-opt-level=O0", flags)
self.assertIn("--iree-opt-strip-assertions=false", flags)

session.set_flags("--iree-opt-level=O2")
flags = session.get_flags()
self.assertIn("--iree-opt-level=O2", flags)
self.assertIn("--iree-global-optimization-opt-level=O0", flags)
self.assertIn("--iree-opt-strip-assertions=false", flags)

class DlInvocationTest(unittest.TestCase):
def testCreate(self):
session = Session()
Expand Down
21 changes: 19 additions & 2 deletions compiler/src/iree/compiler/API/Internal/CompilerDriver.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,7 @@ struct GlobalInit {
// Our session options can optionally be bound to the global command-line
// environment. If that is not the case, then these will be nullptr, and
// they should be default initialized at the session level.
GlobalPipelineOptions *clGlobalPipelineOptions = nullptr;
PluginManagerOptions *clPluginManagerOptions = nullptr;
BindingOptions *clBindingOptions = nullptr;
InputDialectOptions *clInputOptions = nullptr;
Expand Down Expand Up @@ -278,6 +279,7 @@ void GlobalInit::registerCommandLineOptions() {
mlir::tracing::DebugConfig::registerCLOptions();

// Bind session options to the command line environment.
clGlobalPipelineOptions = &GlobalPipelineOptions::FromFlags::get();
clPluginManagerOptions = &PluginManagerOptions::FromFlags::get();
clBindingOptions = &BindingOptions::FromFlags::get();
clInputOptions = &InputDialectOptions::FromFlags::get();
Expand Down Expand Up @@ -323,6 +325,7 @@ struct Session {
if (failed(binder.parseArguments(argc, argv, callback))) {
return new Error(std::move(errorMessage));
}

return nullptr;
}

Expand Down Expand Up @@ -387,6 +390,7 @@ struct Session {
bool pluginsActivated = false;
LogicalResult pluginActivationStatus{failure()};

GlobalPipelineOptions pipelineOptions;
BindingOptions bindingOptions;
InputDialectOptions inputOptions;
PreprocessingOptions preprocessingOptions;
Expand All @@ -410,6 +414,7 @@ Session::Session(GlobalInit &globalInit)
// Bootstrap session options from the cl environment, if enabled.
if (globalInit.usesCommandLine) {
debugConfig = mlir::tracing::DebugConfig::createFromCLOptions();
pipelineOptions = *globalInit.clGlobalPipelineOptions;
pluginManagerOptions = *globalInit.clPluginManagerOptions;
bindingOptions = *globalInit.clBindingOptions;
inputOptions = *globalInit.clInputOptions;
Expand All @@ -430,6 +435,7 @@ Session::Session(GlobalInit &globalInit)

// Register each options struct with the binder so we can manipulate
// mnemonically via the API.
pipelineOptions.bindOptions(binder);
bindingOptions.bindOptions(binder);
preprocessingOptions.bindOptions(binder);
inputOptions.bindOptions(binder);
Expand Down Expand Up @@ -938,6 +944,17 @@ void Invocation::dumpCompilationPhase(IREEVMPipelinePhase phase,

bool Invocation::runPipeline(enum iree_compiler_pipeline_t pipeline) {
auto passManager = createPassManager();

auto globalBinder = OptionsBinder::global();
auto &binder =
session.globalInit.usesCommandLine ? globalBinder : session.binder;
GlobalOptimizationOptions highLevelOptimizationOptions =
session.highLevelOptimizationOptions;

// Set optimization options on a copy of the sessions options.
highLevelOptimizationOptions.applyOptimization(binder,
session.pipelineOptions);

switch (pipeline) {
case IREE_COMPILER_PIPELINE_STD: {
IREEVMPipelinePhase compileFrom;
Expand All @@ -962,7 +979,7 @@ bool Invocation::runPipeline(enum iree_compiler_pipeline_t pipeline) {

buildIREEVMTransformPassPipeline(
session.targetRegistry, session.bindingOptions, session.inputOptions,
session.preprocessingOptions, session.highLevelOptimizationOptions,
session.preprocessingOptions, highLevelOptimizationOptions,
session.schedulingOptions, session.halTargetOptions,
session.vmTargetOptions, pipelineHooks, *passManager, compileFrom,
compileTo);
Expand Down Expand Up @@ -995,7 +1012,7 @@ bool Invocation::runPipeline(enum iree_compiler_pipeline_t pipeline) {
}
buildIREEPrecompileTransformPassPipeline(
session.targetRegistry, session.bindingOptions, session.inputOptions,
session.preprocessingOptions, session.highLevelOptimizationOptions,
session.preprocessingOptions, highLevelOptimizationOptions,
session.schedulingOptions, session.halTargetOptions, pipelineHooks,
*passManager, compileFrom, compileTo);
break;
Expand Down
28 changes: 28 additions & 0 deletions compiler/src/iree/compiler/Pipelines/Options.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,16 +5,29 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/compiler/Pipelines/Options.h"
#include "llvm/Passes/OptimizationLevel.h"

IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::BindingOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::InputDialectOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(
mlir::iree_compiler::GlobalOptimizationOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::SchedulingOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::PreprocessingOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::GlobalPipelineOptions);

namespace mlir::iree_compiler {

void GlobalPipelineOptions::bindOptions(OptionsBinder &binder) {
static llvm::cl::OptionCategory category(
"IREE global pipeline options controlling the entire compilation flow.");

binder.opt<llvm::OptimizationLevel>(
"iree-opt-level", optLevel,
llvm::cl::desc("Global optimization level to apply to the entire "
"compilation flow."),
llvm::cl::cat(category));
}
Comment on lines +20 to +29
Copy link
Member

Choose a reason for hiding this comment

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

Just dropping a reminder here: we have some existing docs for "optimization options" that are due to a rework, especially after these changes land: https://iree.dev/reference/optimization-options/

We could even add a snippet that we include on each page like https://iree.dev/guides/deployment-configurations/cpu/#compile-and-run-a-program and https://iree.dev/guides/deployment-configurations/gpu-vulkan/#compile-and-run-a-program that explains what to expect out of the box and how to optimize / tune / configure from there (linking to the optimization options reference page).


void BindingOptions::bindOptions(OptionsBinder &binder) {
static llvm::cl::OptionCategory category(
"IREE translation binding support options.");
Expand Down Expand Up @@ -119,9 +132,24 @@ void PreprocessingOptions::bindOptions(OptionsBinder &binder) {
llvm::cl::cat(category));
}

void GlobalOptimizationOptions::applyOptimization(
const OptionsBinder &binder, const GlobalPipelineOptions &globalLevel) {
binder.overrideDefault("iree-global-optimization-opt-level", optLevel,
globalLevel.optLevel);

if (optLevel != llvm::OptimizationLevel::O0) {
binder.overrideDefault("iree-opt-strip-assertions", stripAssertions, true);
llvm::dbgs() << "stripAssertions " << stripAssertions << '\n';
}
};

void GlobalOptimizationOptions::bindOptions(OptionsBinder &binder) {
static llvm::cl::OptionCategory category(
"IREE options for controlling global optimizations.");
binder.opt<llvm::OptimizationLevel>(
"iree-global-optimization-opt-level", optLevel,
llvm::cl::desc("Optimization level for the this pipeline"),
llvm::cl::cat(category));
Copy link
Contributor Author

Choose a reason for hiding this comment

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

@benvanik I attempted to do similar to what you suggested but I've run into some problems.

As an example init_at_opt_level would be the function used to tell the binder to set stripAssertions to true when at O3 (and there wouldn't be a need for the manual applyOptimization):

 binder.opt<bool>(
      "iree-opt-strip-assertions",
      stripAssertions,
      init_at_opt_level(O3, true));

The problem I'm having is that iree-opt-strip-assertions should be set based upon the value iree-global-optimization-opt-level and I don't have a good way to let the binder know that.

There is a secondary problem, too. As you mentioned, the options shouldn't be directly set but rather a copy should be made. This is needed because we don't want setting optimization levels to affect the other flags for a session. This makes it hard (impossible?) to to make it so that GlobalOptimizationOptions::stripAssertions should be set based on GlobalOptimizationOptions::optLevel from inside of of bindOptions because the GlobalOptimizationOptions object we use to bind won't be the one that has the optimizations applied.

Copy link
Collaborator

Choose a reason for hiding this comment

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

This is a real good start!

There's a few super magic ways we could do things, but for v0 I'd have your applyOptimization above walk through the options registered on a binder and if there are registered init_at_opt_level entries perform the override. It's a bit of boilerplate we could remove with template magic but copy/pasting for now is fine. something like:

applyOptimization:
  binder.overrideDefault("...-opt-level", optLevel, ...);
  binder.overrideOptimizationDefaults(optLevel);
  // if there are nested binders, call applyOptimization on them

The fancier automatic way would be to have the pipeline opt level option be stashed by the binder - instead of binder.opt<OptimizationLevel>(...) could have binder.optimizationLevel("...-opt-level") to make it explicit what the primary scoped optimization level option is and then the applyOptimization default impl above can do everything itself: the top level gets called, it checks if it has a registered optimization level option, overrides it if needed, and then calls overrideOptimizationDefaults. Then we'd only need to override it if we had nested options (but even that could be improved, and for now we could ignore).

binder.opt<bool>(
"iree-opt-aggressively-propagate-transposes",
aggressiveTransposePropagation,
Expand Down
11 changes: 11 additions & 0 deletions compiler/src/iree/compiler/Pipelines/Options.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,12 @@

namespace mlir::iree_compiler {

struct GlobalPipelineOptions {
llvm::OptimizationLevel optLevel = llvm::OptimizationLevel::O0;
void bindOptions(OptionsBinder &binder);
using FromFlags = OptionsFromFlags<GlobalPipelineOptions>;
};

struct BindingOptions {
// Whether to include runtime support functions for the IREE native ABI.
bool native = true;
Expand Down Expand Up @@ -72,6 +78,7 @@ struct InputDialectOptions {
// 2. Through a Transform dialect spec file.
// 3. Through a PDL spec file.
struct PreprocessingOptions {
llvm::OptimizationLevel optLevel = llvm::OptimizationLevel::O0;
std::string preprocessingPassPipeline;
std::string preprocessingTransformSpecFilename;
std::string preprocessingPDLSpecFilename;
Expand All @@ -82,6 +89,7 @@ struct PreprocessingOptions {

// Options controlling high level optimizations.
struct GlobalOptimizationOptions {
llvm::OptimizationLevel optLevel = llvm::OptimizationLevel::O0;
// Maximum byte size increase allowed for constant expr hoisting policy to
// allow hoisting. The threshold is 1MB by default.
int64_t constExprMaxSizeIncreaseThreshold = 1024 * 1024;
Expand Down Expand Up @@ -131,6 +139,9 @@ struct GlobalOptimizationOptions {
// Converts linalg named matmul ops to linalg generic ops.
bool generalizeMatmul = false;

void applyOptimization(const OptionsBinder &binder = OptionsBinder::global(),
const GlobalPipelineOptions &globalLevel =
GlobalPipelineOptions::FromFlags::get());
void bindOptions(OptionsBinder &binder);
using FromFlags = OptionsFromFlags<GlobalOptimizationOptions>;
};
Expand Down
54 changes: 47 additions & 7 deletions compiler/src/iree/compiler/Utils/OptionUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,14 @@

#include "iree/compiler/Utils/OptionUtils.h"

#include "llvm/ADT/StringSwitch.h"
#include "llvm/Passes/OptimizationLevel.h"
#include "llvm/Support/ManagedStatic.h"

namespace mlir::iree_compiler {

void OptionsBinder::addGlobalOption(std::unique_ptr<llvm::cl::Option> option) {
static llvm::ManagedStatic<std::vector<std::unique_ptr<llvm::cl::Option>>>
globalOptions;
globalOptions->push_back(std::move(option));
}
llvm::ManagedStatic<llvm::DenseMap<llvm::StringRef, OptionsBinder::OptionInfo>>
OptionsBinder::globalOptions;

LogicalResult OptionsBinder::parseArguments(int argc, const char *const *argv,
ErrorCallback onError) {
Expand Down Expand Up @@ -76,10 +75,10 @@ LogicalResult OptionsBinder::parseArguments(int argc, const char *const *argv,
llvm::SmallVector<std::string>
OptionsBinder::printArguments(bool nonDefaultOnly) {
llvm::SmallVector<std::string> values;
for (auto &info : localOptions) {
for (auto &[flag, info] : getOptionsStorage()) {
if (!info.print)
continue;
if (nonDefaultOnly && !info.isChanged())
if (nonDefaultOnly && !info.isDefault())
continue;

std::string s;
Expand Down Expand Up @@ -191,3 +190,44 @@ void llvm::cl::parser<PowerOf2ByteSize>::printOptionDiff(
}

void llvm::cl::parser<PowerOf2ByteSize>::anchor() {}

bool llvm::cl::parser<llvm::OptimizationLevel>::parse(
Option &O, StringRef ArgName, StringRef Arg, llvm::OptimizationLevel &Val) {
auto val = StringSwitch<std::optional<OptimizationLevel>>(Arg)
.Case("O0", OptimizationLevel::O0)
.Case("O1", OptimizationLevel::O1)
.Case("O2", OptimizationLevel::O2)
.Case("O3", OptimizationLevel::O3)
// .Case("Os", OptimizationLevel::Os)
// .Case("Oz", OptimizationLevel::Oz)
.Default(std::nullopt);
if (!val) {
return O.error("'" + Arg +
"' value not a valid optimization level, use "
"O0/O1/O2/O3");
}
Val = *val;
return false;
}

void llvm::cl::parser<llvm::OptimizationLevel>::printOptionDiff(
const Option &O, llvm::OptimizationLevel V, const OptVal &Default,
size_t GlobalWidth) const {
printOptionName(O, GlobalWidth);
std::string Str;
{
llvm::raw_string_ostream SS(Str);
SS << V.getSpeedupLevel() << "/" << V.getSizeLevel();
}
outs() << "= " << Str;
outs().indent(2) << " (default: ";
if (Default.hasValue()) {
auto defaultVal = Default.getValue();
outs() << defaultVal.getSpeedupLevel() << "/" << defaultVal.getSizeLevel();
} else {
outs() << "*no default*";
}
outs() << ")\n";
}

void llvm::cl::parser<llvm::OptimizationLevel>::anchor() {}
Loading
Loading