Enable analysis based integer optimizations. (#18756)

* Add to the canonicalization fragment in the global optimization and
the flow pipeline.
* I attempted to add a CL flag for this but given that we are going to
end up using it everywhere, it became unwieldy. If we ever have a
"disable optimizations" it might be appropriate for that.

Signed-off-by: Stella Laurenzo <[email protected]>

compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp

@@ -80,6 +80,8 @@ using FunctionLikeNest =
 
 static void addCleanupPatterns(OpPassManager &passManager) {
   FunctionLikeNest(passManager)
+      // Simplify integer arithmetic.
+      .addPass(IREE::Util::createOptimizeIntArithmeticPass)
       // Standard MLIR cleanup.
       .addPass(IREE::Flow::createCanonicalizerPass)
       .addPass(mlir::createCSEPass)

compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp

@@ -332,7 +332,7 @@ class OptimizeIntArithmeticPass
 
 } // namespace
 
-std::unique_ptr<OperationPass<void>> createOptimizeIntArithmetic() {
+std::unique_ptr<OperationPass<void>> createOptimizeIntArithmeticPass() {
   return std::make_unique<OptimizeIntArithmeticPass>();
 }
 

compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.h

@@ -30,7 +30,7 @@ createFixedPointIteratorPass(OpPassManager pipeline);
 std::unique_ptr<OperationPass<mlir::ModuleOp>> createFoldGlobalsPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>> createFuseGlobalsPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>> createIPOPass();
-std::unique_ptr<OperationPass<void>> createOptimizeIntArithmetic();
+std::unique_ptr<OperationPass<void>> createOptimizeIntArithmeticPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>> createPropagateSubrangesPass();
 std::unique_ptr<OperationPass<void>> createSimplifyGlobalAccessesPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>>

compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.td

@@ -68,7 +68,7 @@ def IPO : Pass<"iree-util-ipo", "mlir::ModuleOp"> {
 def OptimizeIntArithmetic : Pass<"iree-util-optimize-int-arithmetic", ""> {
   let summary = "Optimizes integer arithmetic using a variety of dataflow analysis and patterns.";
   let constructor = [{
-    mlir::iree_compiler::IREE::Util::createOptimizeIntArithmetic()
+    mlir::iree_compiler::IREE::Util::createOptimizeIntArithmeticPass()
   }];
 }
 

compiler/src/iree/compiler/GlobalOptimization/Passes.cpp

@@ -95,6 +95,7 @@ void buildGlobalOptimizationPassPipeline(
 
   // Preprocessing passes to get the program into a canonical state.
   FunctionLikeNest(mainPassManager)
+      .addPass(IREE::Util::createOptimizeIntArithmeticPass)
       .addPass(createLinalgQuantizedConvToConvPass)
       .addPass(createLinalgQuantizedMatmulToMatmulPass)
       .addPass(IREE::Flow::createCanonicalizerPass)
@@ -194,8 +195,11 @@ void buildGlobalOptimizationPassPipeline(
   mainPassManager.addPass(IREE::Util::createApplyPatternsPass());
   mainPassManager.addPass(IREE::Util::createFoldGlobalsPass());
   mainPassManager.addPass(IREE::Util::createIPOPass());
-  mainPassManager.addPass(IREE::Flow::createCanonicalizerPass());
-  mainPassManager.addPass(createCSEPass());
+
+  FunctionLikeNest(mainPassManager)
+      .addPass(IREE::Util::createOptimizeIntArithmeticPass)
+      .addPass(IREE::Flow::createCanonicalizerPass)
+      .addPass(createCSEPass);
 
   if (transformOptions.options.constExprHoisting) {
     buildGlobalOptExprHoistingPassPipeline(mainPassManager, transformOptions);

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

@@ -57,6 +57,11 @@ struct InputDialectOptions {
   bool promoteF16ToF32 = false;
   bool promoteBF16ToF32 = false;
 
+  // Perfoms early optimizations geared towards optimizing/simplifying the
+  // types of integer arithmetic inefficiencies that frontends typically
+  // include and which are implicated in blocking downstream optimizations.
+  bool optimizeIndexArithmetic = true;
+
   void bindOptions(OptionsBinder &binder);
   using FromFlags = OptionsFromFlags<InputDialectOptions>;
 };
@@ -77,29 +82,6 @@ struct PreprocessingOptions {
 
 // Options controlling high level optimizations.
 struct GlobalOptimizationOptions {
-  // Enables aggressive propagation of transposes to the inputs of named ops,
-  // rewriting named ops as fused generics.
-  bool aggressiveTransposePropagation = false;
-
-  // Enables transposing all concatenations to the outer most dimension.
-  bool outerDimConcat = false;
-
-  // Enables data tiling.
-  bool dataTiling = true;
-
-  // Enables const-expr hoisting into globals.
-  bool constExprHoisting = true;
-
-  // Enables recursive evaluation of immutable globals using the compiler
-  // and runtime.
-  bool constEval = true;
-
-  // Optimizations to reduce numeric precision where it is safe to do so.
-  bool numericPrecisionReduction = false;
-
-  // Strips debug assertions after any useful information has been extracted.
-  bool stripAssertions = false;
-
   // Maximum byte size increase allowed for constant expr hoisting policy to
   // allow hoisting. The threshold is 1MB by default.
   int64_t constExprMaxSizeIncreaseThreshold = 1024 * 1024;
@@ -123,6 +105,29 @@ struct GlobalOptimizationOptions {
   // module.
   std::string parameterSplatExportFile = "";
 
+  // Enables aggressive propagation of transposes to the inputs of named ops,
+  // rewriting named ops as fused generics.
+  bool aggressiveTransposePropagation = false;
+
+  // Enables transposing all concatenations to the outer most dimension.
+  bool outerDimConcat = false;
+
+  // Enables data tiling.
+  bool dataTiling = true;
+
+  // Enables const-expr hoisting into globals.
+  bool constExprHoisting = true;
+
+  // Enables recursive evaluation of immutable globals using the compiler
+  // and runtime.
+  bool constEval = true;
+
+  // Optimizations to reduce numeric precision where it is safe to do so.
+  bool numericPrecisionReduction = false;
+
+  // Strips debug assertions after any useful information has been extracted.
+  bool stripAssertions = false;
+
   void bindOptions(OptionsBinder &binder);
   using FromFlags = OptionsFromFlags<GlobalOptimizationOptions>;
 };