Add target independent optimization pass

BUILD.bazel

@@ -316,11 +316,11 @@ gentbl_cc_library(
     tbl_outs = [
         (
             ["--gen-rewriters"],
-            "stablehlo/transforms/StablehloAggressiveSimplificationPatterns.h.inc",
+            "stablehlo/transforms/optimization/StablehloAggressiveSimplificationPatterns.h.inc",
         ),
     ],
     tblgen = "@llvm-project//mlir:mlir-tblgen",
-    td_file = "stablehlo/transforms/StablehloAggressiveSimplificationPatterns.td",
+    td_file = "stablehlo/transforms/optimization/StablehloAggressiveSimplificationPatterns.td",
     deps = [
         ":stablehlo_ops_td_files",
     ],
@@ -1125,15 +1125,33 @@ gentbl_cc_library(
     deps = ["@llvm-project//mlir:PassBaseTdFiles"],
 )
 
+cc_library(
+    name = "stablehlo_pass_utils",
+    srcs = [
+        "stablehlo/transforms/PassUtils.cpp",
+    ],
+    hdrs = [
+        "stablehlo/transforms/PassUtils.h",
+    ],
+    strip_include_prefix = ".",
+    deps = [
+        ":base",
+        ":chlo_ops",
+        ":stablehlo_ops",
+        ":stablehlo_pass_inc_gen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ComplexDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
 cc_library(
     name = "stablehlo_passes",
     srcs = [
         "stablehlo/transforms/ChloLegalizeToStablehlo.cpp",
         "stablehlo/transforms/PassPipelines.cpp",
-        "stablehlo/transforms/PassUtils.cpp",
         "stablehlo/transforms/ShapeLegalizeToStablehlo.cpp",
-        "stablehlo/transforms/StablehloAggressiveFolder.cpp",
-        "stablehlo/transforms/StablehloAggressiveSimplification.cpp",
         "stablehlo/transforms/StablehloCanonicalizeDynamism.cpp",
         "stablehlo/transforms/StablehloCompatibilityExpander.cpp",
         "stablehlo/transforms/StablehloComplexMathExpander.cpp",
@@ -1163,13 +1181,14 @@ cc_library(
         ":chlo_ops",
         ":chlo_rewriters_inc_gen",
         ":linalg_passes",
-        ":stablehlo_aggressive_simplification_inc_gen",
         ":stablehlo_create_compatibility_expander_inc_gen",
         ":stablehlo_create_complex_math_expander_inc_gen",
         ":stablehlo_legalize_deprecated_ops_inc_gen",
         ":stablehlo_ops",
         ":stablehlo_ops_inc_gen",
         ":stablehlo_pass_inc_gen",
+        ":stablehlo_pass_utils",
+        ":stablehlo_passes_optimization",
         ":stablehlo_type_inference",
         ":version",
         ":vhlo_ops",
@@ -1198,6 +1217,61 @@ cc_library(
     ],
 )
 
+gentbl_cc_library(
+    name = "stablehlo_passes_optimization_inc_gen",
+    strip_include_prefix = ".",
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=Optimization",
+            ],
+            "stablehlo/transforms/optimization/Passes.h.inc",
+        ),
+        (
+            [
+                "-gen-pass-doc",
+            ],
+            "stablehlo/transforms/optimization/stablehlo_optimization_passes.md",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "stablehlo/transforms/optimization/Passes.td",
+    deps = ["@llvm-project//mlir:PassBaseTdFiles"],
+)
+
+cc_library(
+    name = "stablehlo_passes_optimization",
+    srcs = [
+        "stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp",
+        "stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp",
+        "stablehlo/transforms/optimization/StablehloTargetIndependentOptimization.cpp",
+    ],
+    hdrs = [
+        "stablehlo/transforms/optimization/Passes.h",
+    ],
+    strip_include_prefix = ".",
+    deps = [
+        ":base",
+        ":stablehlo_aggressive_simplification_inc_gen",
+        ":stablehlo_ops",
+        ":stablehlo_pass_inc_gen",
+        ":stablehlo_pass_utils",
+        ":stablehlo_passes_optimization_inc_gen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:CommonFolders",
+        "@llvm-project//mlir:DialectUtils",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Rewrite",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
 cc_library(
     name = "stablehlo_portable_api",
     srcs = [
@@ -1364,6 +1438,7 @@ cc_binary(
         ":linalg_passes",
         ":register",
         ":stablehlo_passes",
+        ":stablehlo_passes_optimization",
         ":tosa_passes",
         "//stablehlo/tests:check_ops",
         "//stablehlo/tests:test_utils",

build_tools/github_actions/ci_build_docs.sh

@@ -41,6 +41,7 @@ fi
 
 declare -A targets
 targets[":stablehlo_pass_inc_gen_filegroup"]="bazel-bin/stablehlo/transforms/stablehlo_passes.md"
+targets[":stablehlo_passes_optimization_inc_gen_filegroup"]="bazel-bin/stablehlo/transforms/optimization/stablehlo_optimization_passes.md"
 targets[":interpreter_pass_inc_gen_filegroup"]="bazel-bin/stablehlo/reference/interpreter_passes.md"
 targets[":linalg_pass_inc_gen_filegroup"]="bazel-bin/stablehlo/conversions/linalg/transforms/stablehlo_linalg_passes.md"
 targets[":tosa_pass_inc_gen_filegroup"]="bazel-bin/stablehlo/conversions/tosa/transforms/stablehlo_tosa_passes.md"

docs/_toc.yaml

@@ -59,6 +59,8 @@ toc:
     path: /stablehlo/ide
   - title: StableHLO Passes
     path: /stablehlo/generated/stablehlo_passes
+  - title: StableHLO Optimization Passes
+    path: /stablehlo/generated/stablehlo_optimization_passes
   - title: StableHLO Interpreter Passes
     path: /stablehlo/generated/interpreter_passes
   - title: StableHLO Linalg Passes

docs/generated/stablehlo_optimization_passes.md

@@ -0,0 +1,21 @@
+<!-- Autogenerated by mlir-tblgen; don't manually edit -->
+### `-stablehlo-aggressive-folder`
+
+_Folds StableHLO operations_
+
+
+
+#### Options
+```
+-fold-float : Allow for potentially lossy computations using float type.
+```
+### `-stablehlo-aggressive-simplification`
+
+_Canonicalizes StableHLO operations_
+
+
+### `-stablehlo-target-independent-optimization`
+
+_Runs canonicalizers, folders, and other target-independent optimizations._
+
+

docs/generated/stablehlo_passes.md

@@ -13,21 +13,6 @@ An experimental pass that legalizes shape-related ops to StableHLO ops.
 Bringing shape and data computations together via an optional pass will
 make it possible for the StableHLO ecosystem to potentially leverage the
 compilation pipelines that use StableHLO operations to model dynamism.
-### `-stablehlo-aggressive-folder`
-
-_Folds StableHLO operations_
-
-
-
-#### Options
-```
--fold-float : Allow for potentially lossy computations using float type.
-```
-### `-stablehlo-aggressive-simplification`
-
-_Canonicalizes StableHLO operations_
-
-
 ### `-stablehlo-canonicalize-dynamism`
 
 _Canonicalizes dynamic StableHLO ops into static ops._

stablehlo/tests/transforms/stablehlo_target_independent_optimization.mlir

@@ -0,0 +1,12 @@
+// RUN: stablehlo-opt --stablehlo-target-independent-optimization --split-input-file %s | FileCheck %s
+
+// Check that simplificaiton and folding are both applied.
+
+// CHECK-LABEL: @add_cst_on_rhs
+func.func @add_cst_on_rhs(%arg0: tensor<f32>) -> tensor<f32> {
+  %cst = stablehlo.constant dense<1.0> : tensor<f32>
+  %0 = stablehlo.add %cst, %cst : tensor<f32>
+  // CHECK: stablehlo.add %arg0, %cst : tensor<f32>
+  %1 = stablehlo.add %0, %arg0 : tensor<f32>
+  return %1 : tensor<f32>
+}

stablehlo/tools/StablehloOptMain.cpp

@@ -26,12 +26,14 @@ limitations under the License.
 #include "stablehlo/tests/CheckOps.h"
 #include "stablehlo/tests/TestUtils.h"
 #include "stablehlo/transforms/Passes.h"
+#include "stablehlo/transforms/optimization/Passes.h"
 
 int main(int argc, char **argv) {
   mlir::registerAllPasses();
   mlir::hlo::registerAllTestPasses();
   mlir::stablehlo::registerPassPipelines();
   mlir::stablehlo::registerPasses();
+  mlir::stablehlo::registerOptimizationPasses();
   mlir::stablehlo::registerStablehloLinalgTransformsPasses();
   mlir::stablehlo::registerInterpreterTransformsPasses();
   mlir::tosa::registerStablehloTOSATransformsPasses();

stablehlo/transforms/CMakeLists.txt

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+add_subdirectory(optimization)
+
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls)
 add_public_tablegen_target(PassesIncGen)
@@ -20,10 +22,6 @@ set(LLVM_TARGET_DEFINITIONS ChloDecompositionPatterns.td)
 mlir_tablegen(ChloDecompositionPatterns.h.inc --gen-rewriters)
 add_public_tablegen_target(ChloDecompositionPatternsIncGen)
 
-set(LLVM_TARGET_DEFINITIONS StablehloAggressiveSimplificationPatterns.td)
-mlir_tablegen(StablehloAggressiveSimplificationPatterns.h.inc --gen-rewriters)
-add_public_tablegen_target(StablehloAggressiveSimplificationPatternsIncGen)
-
 set(LLVM_TARGET_DEFINITIONS StablehloCompatibilityExpanderPatterns.td)
 mlir_tablegen(StablehloCompatibilityExpanderPatterns.h.inc --gen-rewriters)
 add_public_tablegen_target(StablehloCompatibilityExpanderPatternsIncGen)
@@ -46,8 +44,6 @@ add_mlir_dialect_library(StablehloPasses
   ChloLegalizeToStablehlo.cpp
   PassPipelines.cpp
   ShapeLegalizeToStablehlo.cpp
-  StablehloAggressiveFolder.cpp
-  StablehloAggressiveSimplification.cpp
   StablehloCanonicalizeDynamism.cpp
   StablehloConvertToSignless.cpp
   StablehloCompatibilityExpander.cpp
@@ -67,7 +63,6 @@ add_mlir_dialect_library(StablehloPasses
   DEPENDS
   ChloDecompositionPatternsIncGen
   PassesIncGen
-  StablehloAggressiveSimplificationPatternsIncGen
   StablehloCompatibilityExpanderPatternsIncGen
   StablehloComplexMathExpanderPatternsIncGen
   StablehloLegalizeDeprecatedOpsPatternsIncGen
@@ -93,6 +88,7 @@ add_mlir_dialect_library(StablehloPasses
   StablehloBroadcastUtils
   StablehloLinalgTransforms
   StablehloOps
+  StablehloOptimizationPasses
   StablehloTypeInference
   VhloOps
 )

stablehlo/transforms/Passes.h

@@ -72,11 +72,6 @@ void populateChloToStablehloPatterns(MLIRContext *context,
 void populateChloConstantLikePattern(MLIRContext *context,
                                      RewritePatternSet *patterns);
 
-/// Collection of folding patterns for StableHLO.
-void populateStablehloAggressiveFolderPatterns(RewritePatternSet *patterns,
-                                               MLIRContext *context,
-                                               bool foldFloat);
-
 /// Collection of rewrite patterns for lowering quantized StableHLO operations
 /// using uniform dequantize/quantize operations.
 void populateStablehloLegalizeQuantizedOpToQDQPatterns(
@@ -88,17 +83,6 @@ void populateStablehloLegalizeQuantizedOpToQDQPatterns(
 void populateStablehloLegalizeQDQToQuantizedOpPatterns(
     RewritePatternSet *patterns, MLIRContext *context);
 
-/// A subset of folding patterns for StableHLO that is necessary for shape
-/// refinement.
-void populateStablehloShapeFolderPatterns(RewritePatternSet *patterns,
-                                          MLIRContext *context,
-                                          bool foldFloat = false);
-
-/// Collection of canonicalization patterns for StableHLO.
-void populateStablehloCanonicalizationPatterns(MLIRContext *context,
-                                               RewritePatternSet *patterns,
-                                               PatternBenefit benefit = 1);
-
 /// Collection of patterns to upgrade deprecated ops to long-term supported ops.
 void populateStablehloLegalizeDeprecatedOpsPatterns(
     MLIRContext *context, RewritePatternSet *patterns);

stablehlo/transforms/Passes.td

@@ -36,27 +36,6 @@ def ShapeLegalizeToStablehloPass : Pass<"shape-legalize-to-stablehlo", "func::Fu
   let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
 }
 
-def StablehloAggressiveFolderPass
-    : Pass<"stablehlo-aggressive-folder", "func::FuncOp"> {
-  let summary = "Folds StableHLO operations";
-  let dependentDialects = [
-    "mlir::stablehlo::StablehloDialect",
-    "mlir::tensor::TensorDialect",
-  ];
-  let options = [
-    Option<"foldFloat", "fold-float", "bool", /*default=*/"true",
-           "Allow for potentially lossy computations using float type.">,
-  ];
-}
-
-def StablehloAggressiveSimplificationPass
-    : Pass<"stablehlo-aggressive-simplification", "func::FuncOp"> {
-  let summary = "Canonicalizes StableHLO operations";
-  let dependentDialects = [
-    "mlir::stablehlo::StablehloDialect",
-  ];
-}
-
 def StablehloCanonicalizeDynamismPass : Pass<"stablehlo-canonicalize-dynamism", "func::FuncOp"> {
   let summary = "Canonicalizes dynamic StableHLO ops into static ops.";
   let description = [{

stablehlo/transforms/StablehloRefineShapes.cpp

@@ -53,6 +53,7 @@ limitations under the License.
 #include "stablehlo/dialect/StablehloOps.h"
 #include "stablehlo/dialect/TypeInference.h"
 #include "stablehlo/transforms/Passes.h"
+#include "stablehlo/transforms/optimization/Passes.h"
 
 #define DEBUG_TYPE "stablehlo-refine-shapes"
 

stablehlo/transforms/optimization/CMakeLists.txt

@@ -0,0 +1,45 @@
+# Copyright 2022 The StableHLO Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name=Optimization)
+add_public_tablegen_target(OptimizationPassesIncGen)
+
+set(LLVM_TARGET_DEFINITIONS StablehloAggressiveSimplificationPatterns.td)
+mlir_tablegen(StablehloAggressiveSimplificationPatterns.h.inc --gen-rewriters)
+add_public_tablegen_target(StablehloAggressiveSimplificationPatternsIncGen)
+
+add_mlir_dialect_library(StablehloOptimizationPasses
+  PARTIAL_SOURCES_INTENDED
+  StablehloAggressiveFolder.cpp
+  StablehloAggressiveSimplification.cpp
+  StablehloTargetIndependentOptimization.cpp
+
+  DEPENDS
+  OptimizationPassesIncGen
+  StablehloAggressiveSimplificationPatternsIncGen
+
+  LINK_LIBS PUBLIC
+  ChloOps
+  MLIRArithDialect
+  MLIRDialectUtils
+  MLIRFuncDialect
+  MLIRIR
+  MLIRRewrite
+  MLIRSupport
+  MLIRTransformUtils
+  StablehloBase
+  StablehloOps
+  StablehloTypeInference
+)