[StablehloExt] Refactor simplification patterns into separate files

Extract simplification patterns from ConstantFolding.cpp into a new
Simplifications.cpp file to improve code organization and maintainability.
Also create CastPatterns.cpp to separate cast-related patterns.

Changes:
- Create Simplifications.cpp with simplification patterns (transpose
  combining, min/max simplification, etc.)
- Create CastPatterns.cpp for cast-related pattern matching
- Update ConstantFolding.cpp to remove extracted patterns
- Reorganize tests: move simplification tests to simplifications.mlir,
  rename constant-folding-scatter.mlir to simplifications-scatter.mlir
- Update pass registration and CMakeLists.txt accordingly

In addition, adds a generic pattern for `stablehlo.dot_general` to `stablehlo.multiply`
along with a number of test cases. This replaces the previous pattern that
was under the "stablehlo-ext-canonicalize-dot-general" pass and was less
robust.

GitOrigin-RevId: 13535beff08e9693ab0e3c2e16317c93cef1f97e

Author: christopherbate

mlir-tensorrt/compiler/include/CMakeLists.txt

@@ -1,2 +1 @@
 add_subdirectory(mlir-tensorrt)
-

mlir-tensorrt/compiler/include/mlir-tensorrt/Conversion/CMakeLists.txt

@@ -15,3 +15,4 @@ add_mlir_tensorrt_doc(ConversionPasses
   OUTPUT_FILE docs/Passes/ConversionPasses.md
   COMMAND -gen-pass-doc ${_TABLEGEN_ARGS}
   )
+

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/CMakeLists.txt

@@ -5,3 +5,4 @@ endif()
 add_subdirectory(CUDA)
 add_subdirectory(Plan)
 add_subdirectory(TensorRTRuntime)
+

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/StablehloExt/Transforms/Passes.td

@@ -22,6 +22,30 @@
 
 include "mlir/Pass/PassBase.td"
 
+//===----------------------------------------------------------------------===//
+// StablehloExtSimplificationsPass
+//===----------------------------------------------------------------------===//
+
+def StablehloExtSimplificationsPass : Pass<"stablehlo-ext-simplifications"> {
+  let summary = "Applies Stablehlo device-agnostic simplifications along with additional patterns";
+  let description = [{
+    This pass runs the patterns defined by
+    `stablehlo-aggressive-simplification` along with additional custom
+    patterns not yet upstreamed.
+  }];
+
+  let dependentDialects = [
+    "::mlir::stablehlo::StablehloDialect",
+    "::mlir::arith::ArithDialect",
+    "::mlir::tensor::TensorDialect"
+  ];
+
+  let options = [
+    Option<"foldOpElementLimit", "fold-op-element-limit", "int64_t", "1024",
+      "The computation size limit for constant folding">
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // StablehloRaiseQDQPass
 //===----------------------------------------------------------------------===//

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/StablehloExt/Transforms/Patterns.h

@@ -33,6 +33,16 @@ namespace stablehlo_ext {
 /// `tensor.cast` producers.
 void populateStableHloAbsorbTensorCastPatterns(RewritePatternSet &patterns);
 
+/// Populate patterns that perform simplifications.
+void populateStableHloExtSimplificationsPatterns(
+    RewritePatternSet &patterns,
+    const stablehlo::StablehloAggressiveSimplificationPassOptions &options,
+    PatternBenefit benefit = 1);
+
+/// Populate patterns that simplify `stablehlo.dot_general` to
+/// `stablehlo.multiply`.
+void populateStablehloDotGeneralToMultiplyPatterns(RewritePatternSet &patterns);
+
 /// Populate patterns that perform target-independent simplifications.
 /// The `sizeLimit` is the maximum tensor volume beyond which constant folding
 /// is not attempted.

mlir-tensorrt/compiler/include/mlir-tensorrt/Transforms/CMakeLists.txt

@@ -9,3 +9,4 @@ endif()
 
 mlir_tablegen(Passes.h.inc ${_TABLEGEN_ARGS})
 mtrt_add_public_tablegen_target(MLIRTensorRTGenericTransformPassIncGen)
+

mlir-tensorrt/compiler/lib/Compiler/CMakeLists.txt

@@ -48,6 +48,9 @@ add_mlir_tensorrt_library(MLIRTensorRTInitAllPasses
   InitAllPasses.cpp
   PARTIAL_SOURCES_INTENDED
 
+  DEPENDS
+    mtrt-headers
+
   LINK_LIBS PUBLIC
     MLIRTensorRTCompilerIncludes
     ${MLIR_TENSORRT_DIALECT_LIBS}

mlir-tensorrt/compiler/lib/Conversion/CMakeLists.txt

@@ -24,3 +24,4 @@ add_subdirectory(TensorRTRuntimeToExecutor)
 add_subdirectory(TensorRTRuntimeToLLVM)
 add_subdirectory(TensorRTToEmitC)
 add_subdirectory(TensorRTToTensorRTRuntime)
+

mlir-tensorrt/compiler/lib/Dialect/CMakeLists.txt

@@ -5,3 +5,4 @@ add_subdirectory(TensorRTRuntime)
 if(MLIR_TRT_ENABLE_HLO)
   add_subdirectory(StablehloExt)
 endif()
+

mlir-tensorrt/compiler/lib/Dialect/StablehloExt/Transforms/CMakeLists.txt

@@ -1,29 +1,31 @@
 add_mlir_tensorrt_library(MLIRTensorRTStableHloExtTransforms
-    CanonicalizeConvolution.cpp
-    CanonicalizeDotGeneral.cpp
-    CanonicalizeGather.cpp
-    CanonicalizeScatter.cpp
-    CanonicalizeShapes.cpp
-    ConstantFolding.cpp
-    ExpandTuples.cpp
-    GatherToSlice.cpp
-    LowerSpecialCustomCalls.cpp
-    MaterializeDenseResourceElementsAttr.cpp
-    StablehloRaiseQDQ.cpp
-    TargetSpecificOptimizations.cpp
+  CastPatterns.cpp
+  CanonicalizeConvolution.cpp
+  CanonicalizeDotGeneral.cpp
+  CanonicalizeGather.cpp
+  CanonicalizeScatter.cpp
+  CanonicalizeShapes.cpp
+  ConstantFolding.cpp
+  ExpandTuples.cpp
+  GatherToSlice.cpp
+  LowerSpecialCustomCalls.cpp
+  MaterializeDenseResourceElementsAttr.cpp
+  Simplifications.cpp
+  StablehloRaiseQDQ.cpp
+  TargetSpecificOptimizations.cpp
 
-    LINK_LIBS PUBLIC
-      ChloOps
-      MLIRExecutorCommonUtils
-      MLIRTensorRTConstantFoldingUtils
-      MLIRTensorRTShapeUtils
-      MLIRTensorRTStableHloExtUtils
-      StablehloOps
-      StablehloPasses
-      StablehloOptimizationPasses
-    MLIR_LIBS PUBLIC
-      MLIRRewrite
-      MLIRTensorDialect
-    DEPENDS
+  LINK_LIBS PUBLIC
+    ChloOps
+    MLIRExecutorCommonUtils
+    MLIRTensorRTConstantFoldingUtils
+    MLIRTensorRTShapeUtils
+    MLIRTensorRTStableHloExtUtils
+    StablehloOps
+    StablehloPasses
+    StablehloOptimizationPasses
+  MLIR_LIBS PUBLIC
+    MLIRRewrite
+    MLIRTensorDialect
+  DEPENDS
     MLIRTensorRTStableHloExtTransformsPassIncGen
 )