Add pipeline options for Dispatch Creation (#20217)

This change allows specifying optimization level defaults for flags used
in dispatch creation. So far, I have only moved
`iree-dispatch-creation-enable-aggressive-fusion` and didn't change the
name as it would be a breaking change for a lot of users. Also, this
cleans up no longer needed sdxl benchmark flags.

---------

Signed-off-by: Ian Wood <[email protected]>

Author: IanWood1

compiler/src/iree/compiler/API/Internal/CompilerDriver.cpp

@@ -241,6 +241,7 @@ struct GlobalInit {
   InputDialectOptions *clInputOptions = nullptr;
   PreprocessingOptions *clPreprocessingOptions = nullptr;
   GlobalOptimizationOptions *clGlobalOptimizationOptions = nullptr;
+  DispatchCreationOptions *clDispatchCreationOptions = nullptr;
   SchedulingOptions *clSchedulingOptions = nullptr;
   IREE::HAL::TargetOptions *clHalTargetOptions = nullptr;
   IREE::VM::TargetOptions *clVmTargetOptions = nullptr;
@@ -286,6 +287,7 @@ void GlobalInit::registerCommandLineOptions() {
   clInputOptions = &InputDialectOptions::FromFlags::get();
   clPreprocessingOptions = &PreprocessingOptions::FromFlags::get();
   clGlobalOptimizationOptions = &GlobalOptimizationOptions::FromFlags::get();
+  clDispatchCreationOptions = &DispatchCreationOptions::FromFlags::get();
   clSchedulingOptions = &SchedulingOptions::FromFlags::get();
   clHalTargetOptions = &IREE::HAL::TargetOptions::FromFlags::get();
   clVmTargetOptions = &IREE::VM::TargetOptions::FromFlags::get();
@@ -396,6 +398,7 @@ struct Session {
   InputDialectOptions inputOptions;
   PreprocessingOptions preprocessingOptions;
   GlobalOptimizationOptions highLevelOptimizationOptions;
+  DispatchCreationOptions dispatchCreationOptions;
   SchedulingOptions schedulingOptions;
   IREE::HAL::TargetOptions halTargetOptions;
   IREE::VM::TargetOptions vmTargetOptions;
@@ -423,6 +426,7 @@ Session::Session(GlobalInit &globalInit)
     inputOptions = *globalInit.clInputOptions;
     preprocessingOptions = *globalInit.clPreprocessingOptions;
     highLevelOptimizationOptions = *globalInit.clGlobalOptimizationOptions;
+    dispatchCreationOptions = *globalInit.clDispatchCreationOptions;
     schedulingOptions = *globalInit.clSchedulingOptions;
     halTargetOptions = *globalInit.clHalTargetOptions;
     vmTargetOptions = *globalInit.clVmTargetOptions;
@@ -443,6 +447,7 @@ Session::Session(GlobalInit &globalInit)
   preprocessingOptions.bindOptions(binder);
   inputOptions.bindOptions(binder);
   highLevelOptimizationOptions.bindOptions(binder);
+  dispatchCreationOptions.bindOptions(binder);
   schedulingOptions.bindOptions(binder);
   halTargetOptions.bindOptions(binder);
   vmTargetOptions.bindOptions(binder);
@@ -982,9 +987,9 @@ bool Invocation::runPipeline(enum iree_compiler_pipeline_t pipeline) {
     buildIREEVMTransformPassPipeline(
         session.targetRegistry, session.bindingOptions, session.inputOptions,
         session.preprocessingOptions, session.highLevelOptimizationOptions,
-        session.schedulingOptions, session.halTargetOptions,
-        session.vmTargetOptions, pipelineHooks, *passManager, compileFrom,
-        compileTo);
+        session.dispatchCreationOptions, session.schedulingOptions,
+        session.halTargetOptions, session.vmTargetOptions, pipelineHooks,
+        *passManager, compileFrom, compileTo);
     break;
   }
   case IREE_COMPILER_PIPELINE_HAL_EXECUTABLE: {
@@ -1015,8 +1020,9 @@ bool Invocation::runPipeline(enum iree_compiler_pipeline_t pipeline) {
     buildIREEPrecompileTransformPassPipeline(
         session.targetRegistry, session.bindingOptions, session.inputOptions,
         session.preprocessingOptions, session.highLevelOptimizationOptions,
-        session.schedulingOptions, session.halTargetOptions, pipelineHooks,
-        *passManager, compileFrom, compileTo);
+        session.dispatchCreationOptions, session.schedulingOptions,
+        session.halTargetOptions, pipelineHooks, *passManager, compileFrom,
+        compileTo);
     break;
   }
   default:

compiler/src/iree/compiler/ConstEval/JitGlobals.cpp

@@ -72,6 +72,7 @@ struct CompileOptions {
   InputDialectOptions inputOptions;
   PreprocessingOptions preprocessingOptions;
   GlobalOptimizationOptions globalOptimizationOptions;
+  DispatchCreationOptions dispatchCreationOptions;
   SchedulingOptions schedulingOptions;
   IREE::HAL::TargetOptions executableOptions;
   IREE::VM::TargetOptions targetOptions;
@@ -647,6 +648,7 @@ class JitGlobalsPass final : public impl::JitGlobalsPassBase<JitGlobalsPass> {
         *targetRegistry.value, compileOptions->bindingOptions,
         compileOptions->inputOptions, compileOptions->preprocessingOptions,
         compileOptions->globalOptimizationOptions,
+        compileOptions->dispatchCreationOptions,
         compileOptions->schedulingOptions, compileOptions->executableOptions,
         compileOptions->targetOptions, compileOptions->hooks, compilePipeline);
   }

compiler/src/iree/compiler/DispatchCreation/Passes.cpp

@@ -70,12 +70,6 @@ static llvm::cl::opt<bool>
                          llvm::cl::desc("Fuse multi-use ops."),
                          llvm::cl::init(false));
 
-static llvm::cl::opt<bool> clEnableAggressiveFusion(
-    "iree-dispatch-creation-enable-aggressive-fusion",
-    llvm::cl::desc("Aggressive fusion opportunities that are behind a flag "
-                   "since all backends dont support it yet"),
-    llvm::cl::init(false));
-
 static llvm::cl::opt<bool> clEnableDataTiling(
     "iree-dispatch-creation-experimental-data-tiling",
     llvm::cl::desc("Enable data-tiling at flow level, i.e., it sets encodings "
@@ -215,16 +209,18 @@ void addDispatchRegionCreationPreprocessingPasses(OpPassManager &passManager) {
 // Note that we should not hoist out small constants before the dispatch regions
 // are converted to workgroups. E.g., the `cseConstant` option needs to be false
 // in greedy pattern rewriting drivers.
-static void addDispatchRegionCreationPasses(OpPassManager &passManager) {
+static void
+addDispatchRegionCreationPasses(OpPassManager &passManager,
+                                const DispatchCreationOptions &options) {
   FunctionLikeNest(passManager)
       // Create dispatches for scalar operations as roots.
       .addPass(DispatchCreation::createFormScalarDispatchesPass)
       // Create `flow.dispatch.region` centered around a root and fuse with
       // producers and consumers.
-      .addPass([] {
+      .addPass([&] {
         return DispatchCreation::createFormDispatchRegionsPass(
             FormDispatchRegionsPassOptions{
-                clEnableAggressiveFusion,
+                options.enableAggressiveFusion,
                 clEnableFusePaddingIntoLinalgConsumerOps,
                 clEnableFusePaddingIntoLinalgProducerOps});
       })
@@ -233,10 +229,10 @@ static void addDispatchRegionCreationPasses(OpPassManager &passManager) {
       // afterwards that would need the full dispatch content but don't want to
       // handle explicit captures as materialized as dispatch workgroup operands
       // and block arguments.
-      .addPass([] {
+      .addPass([&] {
         return DispatchCreation::createCloneProducersIntoDispatchRegionsPass(
             CloneProducersIntoDispatchRegionsPassOptions{
-                clEnableAggressiveFusion});
+                options.enableAggressiveFusion});
       })
       // Collapse dimensions of linalg Ops.
       .addPass(DispatchCreation::createCollapseDimensionsPass);
@@ -306,7 +302,7 @@ void buildDispatchCreationPassPipeline(
       .addPass(mlir::createCSEPass);
 
   addDispatchRegionCreationPreprocessingPasses(passManager);
-  addDispatchRegionCreationPasses(passManager);
+  addDispatchRegionCreationPasses(passManager, transformOptions.options);
 
   FunctionLikeNest(passManager)
       .addPass(DispatchCreation::createConvertDispatchRegionsToWorkgroupsPass)

compiler/src/iree/compiler/DispatchCreation/Passes.h

@@ -23,7 +23,9 @@ namespace mlir::iree_compiler::DispatchCreation {
 
 /// This is a placeholder for future. We should pass all the options through the
 /// struct.
-struct TransformOptions : public PassPipelineOptions<TransformOptions> {};
+struct TransformOptions : public PassPipelineOptions<TransformOptions> {
+  DispatchCreationOptions options;
+};
 
 void buildDispatchCreationPassPipeline(
     OpPassManager &passManager, const TransformOptions &transformOptions);

compiler/src/iree/compiler/Pipelines/Options.cpp

@@ -14,6 +14,7 @@ IREE_DEFINE_COMPILER_OPTION_FLAGS(
 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);
+IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::DispatchCreationOptions);
 
 namespace mlir::iree_compiler {
 
@@ -293,4 +294,19 @@ void SchedulingOptions::bindOptions(OptionsBinder &binder) {
                           llvm::cl::cat(category));
 }
 
+void DispatchCreationOptions::bindOptions(OptionsBinder &binder) {
+  static llvm::cl::OptionCategory category(
+      "IREE options for controlling dispatch region creation.");
+  auto init_at_opt = binder.optimizationLevel(
+      "iree-dispatch-creation-opt-level", optLevel,
+      llvm::cl::desc("Optimization level for the this pipeline"),
+      llvm::cl::cat(category));
+  binder.opt<bool>(
+      "iree-dispatch-creation-enable-aggressive-fusion", enableAggressiveFusion,
+      {init_at_opt(llvm::OptimizationLevel::O0, false),
+       init_at_opt(llvm::OptimizationLevel::O2, true)},
+      llvm::cl::desc("Aggressive fusion opportunities that are behind a flag "
+                     "since all backends dont support it yet"));
+}
+
 } // namespace mlir::iree_compiler

compiler/src/iree/compiler/Pipelines/Options.h

@@ -214,6 +214,15 @@ struct SchedulingOptions {
   using FromFlags = OptionsFromFlags<SchedulingOptions>;
 };
 
+struct DispatchCreationOptions {
+  llvm::OptimizationLevel optLevel;
+
+  bool enableAggressiveFusion = false;
+
+  void bindOptions(OptionsBinder &binder);
+  using FromFlags = OptionsFromFlags<DispatchCreationOptions>;
+};
+
 } // namespace mlir::iree_compiler
 
 #endif // IREE_COMPILER_PIPELINES_OPTIONS_H_

compiler/src/iree/compiler/Pipelines/Pipelines.cpp

@@ -77,6 +77,7 @@ void buildIREEPrecompileTransformPassPipeline(
     BindingOptions bindingOptions, InputDialectOptions inputOptions,
     PreprocessingOptions preprocessingOptions,
     GlobalOptimizationOptions globalOptimizationOptions,
+    DispatchCreationOptions dispatchCreationOptions,
     SchedulingOptions schedulingOptions,
     IREE::HAL::TargetOptions halTargetOptions, IREEVMPipelineHooks &hooks,
     OpPassManager &passManager, IREEVMPipelinePhase compileFrom,
@@ -247,15 +248,16 @@ void buildIREEVMTransformPassPipeline(
     BindingOptions bindingOptions, InputDialectOptions inputOptions,
     PreprocessingOptions preprocessingOptions,
     GlobalOptimizationOptions globalOptimizationOptions,
+    DispatchCreationOptions dispatchCreationOptions,
     SchedulingOptions schedulingOptions,
     IREE::HAL::TargetOptions halTargetOptions,
     IREE::VM::TargetOptions vmTargetOptions, IREEVMPipelineHooks &hooks,
     OpPassManager &passManager, IREEVMPipelinePhase compileFrom,
     IREEVMPipelinePhase compileTo) {
   buildIREEPrecompileTransformPassPipeline(
       targetRegistry, bindingOptions, inputOptions, preprocessingOptions,
-      globalOptimizationOptions, schedulingOptions, halTargetOptions, hooks,
-      passManager, compileFrom, compileTo);
+      globalOptimizationOptions, dispatchCreationOptions, schedulingOptions,
+      halTargetOptions, hooks, passManager, compileFrom, compileTo);
 
   if (compileTo <= IREEVMPipelinePhase::GlobalOptimization)
     return; // early-exit
@@ -274,13 +276,14 @@ void buildIREEVMTransformPassPipeline(
     // No flow/stream processing (implies no tensors).
     break;
   default:
-    DispatchCreation::TransformOptions dispatchCreationOptions;
+    DispatchCreation::TransformOptions dispatchTransformOptions;
+    dispatchTransformOptions.options = dispatchCreationOptions;
     if (compileFrom < IREEVMPipelinePhase::DispatchCreation) { // late-entry
       IREE_TRACE_ADD_BEGIN_FRAME_PASS(passManager, "DispatchCreation");
       if (hooks.beforePhase)
         hooks.beforePhase(IREEVMPipelinePhase::DispatchCreation, passManager);
       DispatchCreation::buildDispatchCreationPassPipeline(
-          passManager, dispatchCreationOptions);
+          passManager, dispatchTransformOptions);
       if (hooks.afterPhase)
         hooks.afterPhase(IREEVMPipelinePhase::DispatchCreation, passManager);
       IREE_TRACE_ADD_END_FRAME_PASS(passManager, "DispatchCreation");
@@ -385,6 +388,7 @@ void buildDefaultIREEVMTransformPassPipeline(OpPassManager &passManager) {
       IREE::HAL::TargetRegistry::getGlobal(), BindingOptions::FromFlags::get(),
       InputDialectOptions::FromFlags::get(),
       PreprocessingOptions::FromFlags::get(), highLevelOptimizations,
+      DispatchCreationOptions::FromFlags::get(),
       SchedulingOptions::FromFlags::get(),
       IREE::HAL::TargetOptions::FromFlags::get(),
       IREE::VM::TargetOptions::FromFlags::get(), defaultHooks, passManager);

compiler/src/iree/compiler/Pipelines/Pipelines.h

@@ -104,6 +104,7 @@ void buildIREEPrecompileTransformPassPipeline(
     BindingOptions bindingOptions, InputDialectOptions inputOptions,
     PreprocessingOptions preprocessingOptions,
     GlobalOptimizationOptions highLevelOptimizationOptions,
+    DispatchCreationOptions dispatchCreationOptions,
     SchedulingOptions schedulingOptions,
     IREE::HAL::TargetOptions halTargetOptions, IREEVMPipelineHooks &hooks,
     OpPassManager &passManager,
@@ -120,6 +121,7 @@ void buildIREEVMTransformPassPipeline(
     BindingOptions bindingOptions, InputDialectOptions inputOptions,
     PreprocessingOptions preprocessingOptions,
     GlobalOptimizationOptions highLevelOptimizationOptions,
+    DispatchCreationOptions dispatchCreationOptions,
     SchedulingOptions schedulingOptions,
     IREE::HAL::TargetOptions halTargetOptions,
     IREE::VM::TargetOptions vmTargetOptions, IREEVMPipelineHooks &hooks,

experimental/regression_suite/shark-test-suite-models/sd3/test_clip.py

@@ -104,14 +104,11 @@ def SD3_CLIP_COMMON_RUN_FLAGS(
 ROCM_COMPILE_FLAGS = [
     "--iree-hal-target-backends=rocm",
     f"--iree-hip-target={rocm_chip}",
+    "--iree-opt-level=O3",
     "--iree-input-type=torch",
     "--iree-opt-const-eval=false",
-    "--iree-global-opt-propagate-transposes=true",
-    "--iree-opt-outer-dim-concat=true",
     "--iree-hip-waves-per-eu=2",
     "--iree-llvmgpu-enable-prefetch",
-    "--iree-dispatch-creation-enable-aggressive-fusion",
-    "--iree-opt-aggressively-propagate-transposes=true",
     "--iree-codegen-llvmgpu-use-vector-distribution=true",
     "--iree-execution-model=async-external",
     "--iree-preprocessing-pass-pipeline=builtin.module(iree-preprocessing-transpose-convolution-pipeline,iree-preprocessing-pad-to-intrinsics{pad-target-type=conv})",

experimental/regression_suite/shark-test-suite-models/sd3/test_mmdit.py

@@ -83,11 +83,8 @@ def SD3_MMDIT_COMMON_RUN_FLAGS(
 ROCM_COMPILE_FLAGS = [
     "--iree-hal-target-backends=rocm",
     f"--iree-hip-target={rocm_chip}",
+    "--iree-opt-level=O3",
     "--iree-opt-const-eval=false",
-    "--iree-global-opt-propagate-transposes=true",
-    "--iree-dispatch-creation-enable-aggressive-fusion=true",
-    "--iree-opt-aggressively-propagate-transposes=true",
-    "--iree-opt-outer-dim-concat=true",
     "--iree-vm-target-truncate-unsupported-floats",
     "--iree-llvmgpu-enable-prefetch=true",
     "--iree-opt-data-tiling=false",