Refactor: Move StableHLO matchers to dialect utils and reorganize namespaces (#730)

- Move StablehloMatchers from Transforms to Dialect/StablehloExt/Utils
- Simplify matcher header to focus on softmax pattern matching
- Move pass implementations from mlir to mtrt namespace
- Rename LinalgExtElementwiseFusionPass to
MTRTLinalgElementwiseFusionPass
- Reorganize test infrastructure under test/lib/Dialect/StablehloExt
- Update build system and pass registrations accordingly

GitOrigin-RevId: 7ecb025b071c65fcc6ccdc9543cd6f9f3ec389eb

Author: christopherbate

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/StablehloExt/Utils/StablehloMatchers.h

@@ -0,0 +1,60 @@
+//===- StablehloMatchers.h  -----------------------------------------------===//
+//
+// SPDX-FileCopyrightText: Copyright 2024 NVIDIA CORPORATION & AFFILIATES.
+// All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// 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
+//
+// http://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.
+//
+//===----------------------------------------------------------------------===//
+///
+/// This file defines different matchers for StableHLO.
+///
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_DIALECT_STABLEHLOEXT_UTILS_STABLEHLOMATCHERS
+#define MLIR_TENSORRT_DIALECT_STABLEHLOEXT_UTILS_STABLEHLOMATCHERS
+
+#include "stablehlo/dialect/StablehloOps.h"
+
+namespace mlir::stablehlo {
+
+namespace detail {
+/// Matcher for stablehlo
+/// reduce(max)->subtract->exponential->reduce(add)->divide The Softmax Matcher
+/// is rooted at stablehlo::DivOp operation.
+template <typename BcastInDimOp, typename ReduceOp, typename SubOp,
+          typename ExpnOp, typename DivideOp, typename MaxOp, typename AddOp>
+struct HLOToSoftmaxMatcher {
+  HLOToSoftmaxMatcher(Value &softmaxInputDeduced, int64_t &reductionDim)
+      : softmaxInputOperand(softmaxInputDeduced), softmaxAxis(reductionDim) {}
+  Value &softmaxInputOperand;
+  int64_t &softmaxAxis;
+  bool match(Operation *op);
+};
+
+} // namespace detail
+
+/// Match a pattern rooted at stablhehlo.divide to a softmax
+/// operation.
+inline auto m_StableHLOSoftmaxMatcher(Value &softmaxInputOperand,
+                                      int64_t &softmax_axis) {
+  return detail::HLOToSoftmaxMatcher<
+      mlir::stablehlo::BroadcastInDimOp, mlir::stablehlo::ReduceOp,
+      mlir::stablehlo::SubtractOp, mlir::stablehlo::ExpOp,
+      mlir::stablehlo::DivOp, mlir::stablehlo::MaxOp, mlir::stablehlo::AddOp>(
+      softmaxInputOperand, softmax_axis);
+}
+
+} // namespace mlir::stablehlo
+
+#endif // MLIR_TENSORRT_DIALECT_STABLEHLOEXT_UTILS_STABLEHLOMATCHERS

mlir-tensorrt/compiler/include/mlir-tensorrt/Transforms/Passes.h

@@ -32,11 +32,12 @@ namespace mlir {
 namespace tensor {
 class TensorDialect;
 }
+} // namespace mlir
 
+namespace mtrt {
 #define GEN_PASS_DECL
 #define GEN_PASS_REGISTRATION
 #include "mlir-tensorrt/Transforms/Passes.h.inc"
-
-} // namespace mlir
+} // namespace mtrt
 
 #endif // MLIR_TENSORRT_TRANSFORMS_PASSES_H

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

@@ -22,20 +22,6 @@
 
 include "mlir/Pass/PassBase.td"
 
-#ifdef MLIR_TENSORRT_ENABLE_HLO
-//===----------------------------------------------------------------------===//
-// TestStablehloMatchersPass to test different matchers for stablehlo
-//===----------------------------------------------------------------------===//
-def TestStablehloMatchersPass : Pass<"test-stablehlo-matchers"> {
-  let summary = "tests pattern matching utilities for StableHLO.";
-  let description =[{
-   StableHLOMatchers.h defines matchers to raise to different patterns for
-   dot-product attention, softmax, etc. This pass tests a matcher by looking
-   for the `__matched__` attribute added by a matcher.
-   }];
-}
-#endif // MLIR_TENSORRT_ENABLE_HLO
-
 //===----------------------------------------------------------------------===//
 // DropNestedModulesPass
 //===----------------------------------------------------------------------===//
@@ -44,6 +30,25 @@ def DropNestedModulesPass : Pass<"drop-nested-modules", "::mlir::ModuleOp"> {
     " nested within the top-level Module";
 }
 
+//===----------------------------------------------------------------------===//
+// LinalgElementwiseFusionPass
+//===----------------------------------------------------------------------===//
+
+def LinalgElementwiseFusionPass : Pass<"mtrt-linalg-elementwise-fusion"> {
+  let summary =
+    "Performs elementwise fusion on linalg operations specific to MLIR-TensorRT";
+
+  let description = [{
+    Runs linalg elementwise fusion with a control function that is specific
+    to MLIR-TensorRT.
+  }];
+
+  let dependentDialects = [
+    "::mlir::tensor::TensorDialect",
+    "::mlir::linalg::LinalgDialect"
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // MemRefCastEliminationPass
 //===----------------------------------------------------------------------===//

mlir-tensorrt/compiler/include/mlir-tensorrt/Transforms/StablehloMatchers/StablehloMatchers.h

@@ -67,7 +67,7 @@ void mtrt::compiler::registerAllPasses() {
   mlir::registerLowerAffinePass();
   mlir::registerMLIRTensorRTCommonConversionPasses();
   mlir::registerMLIRTensorRTConversionPasses();
-  mlir::registerMLIRTensorRTGenericTransformsPasses();
+  mtrt::registerMLIRTensorRTGenericTransformsPasses();
   mlir::registerTransformsPasses();
   mlir::tensorrt::registerTensorRTPasses();
   mtrt::compiler::registerTensorRTToExecutablePasses();

mlir-tensorrt/compiler/lib/Compiler/InitAllPasses.cpp

@@ -26,6 +26,7 @@ add_mlir_tensorrt_library(MLIRTensorRTCompilerStableHloToExecutable
       MLIRTensorRTHostBackend
       MLIRTensorRTHostToEmitC
       MLIRTensorRTHostToLLVM
+      MLIRTensorRTLinalgElementwiseFusion
       MLIRTensorRTMemRefToCUDA
       MLIRTensorRTPlanToExecutor
       MLIRTensorRTPlanTransforms

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

@@ -97,12 +97,14 @@ class ProcessHostClustersPass
   ProcessHostClustersPass() {
     dynamicPM = OpPassManager("func.func");
     dynamicPM.addPass(mlir::createStablehloToLinalgPass());
+    dynamicPM.addPass(mlir::createLinalgGeneralizeNamedOpsPass());
+    dynamicPM.addPass(mtrt::createLinalgElementwiseFusionPass());
     dynamicPM.addPass(mlir::createLinalgDetensorizePass(
         mlir::LinalgDetensorizePassOptions{/*aggressiveMode=*/true}));
     dynamicPM.addPass(mlir::createConvertToLoops());
     dynamicPM.addPass(mlir::createCSEPass());
     dynamicPM.addPass(mlir::createCanonicalizerPass());
-    dynamicPM.addPass(mlir::createSCFDetensorizeLoopsPass());
+    dynamicPM.addPass(createSCFDetensorizeLoopsPass());
   }
 
   void runOnOperation() override {

mlir-tensorrt/compiler/lib/Compiler/StablehloToExecutable/Passes.cpp

@@ -93,8 +93,8 @@ void mtrt::compiler::buildStablehloPreProcessingPipeline(
   // `convert-stablehlo-to-scf`:
   if (opts.legalizeControlFlowToSCF) {
     pm.addNestedPass<func::FuncOp>(mlir::createConvertStablehloToScfPass());
-    pm.addNestedPass<func::FuncOp>(mlir::createUnrollForLoopsPass(
-        mlir::UnrollForLoopsPassOptions{opts.unrollThreshold}));
+    pm.addNestedPass<func::FuncOp>(createUnrollForLoopsPass(
+        mtrt::UnrollForLoopsPassOptions{opts.unrollThreshold}));
   }
 
   // `stablehlo-ext-constant-folding`:

mlir-tensorrt/compiler/lib/Compiler/StablehloToExecutable/StableHloInputPipelines.cpp

@@ -13,7 +13,6 @@ add_mlir_tensorrt_library(MLIRTensorRTStablehloToTensorRT
     MLIRTensorRTConvertToTensorRTCommon
     MLIRTensorRTTensorRTDialect
     MLIRTensorRTStableHloExtUtils
-    MLIRTensorRTStablehloMatchers
     MLIRTensorRTTensorRTUtils
     MLIRTensorRTUtilsShapeInfo
     StablehloOps

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

@@ -31,7 +31,7 @@
 #include "mlir-tensorrt/Conversion/Patterns.h"
 #include "mlir-tensorrt/Conversion/TensorRTCommon/ConvertToTensorRTCommon.h"
 #include "mlir-tensorrt/Dialect/StablehloExt/Utils/GatherScatterUtils.h"
-#include "mlir-tensorrt/Transforms/StablehloMatchers/StablehloMatchers.h"
+#include "mlir-tensorrt/Dialect/StablehloExt/Utils/StablehloMatchers.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/FuncConversions.h"
@@ -3713,7 +3713,7 @@ struct ConvertHLOSoftmax : public OpRewritePattern<stablehlo::DivOp> {
     mlir::Value softmaxInputOperand;
     int64_t deducedAxis = -1;
     if (!matchPattern(op.getOperation(),
-                      mlir::matchers::m_StableHLOSoftmaxMatcher(
+                      mlir::stablehlo::m_StableHLOSoftmaxMatcher(
                           softmaxInputOperand, deducedAxis)))
       return failure();