add DecomposeAggregatedOps

Author: xurui1995

include/gc/Transforms/Passes.td

@@ -92,6 +92,12 @@ def VerifyTargetDescription : Pass<"verify-target-description", "ModuleOp"> {
            /*default=*/"\"CPU\"",
            "The device to verify. Supported device: CPU, ">,
   ];
+
+def DecomposeAggregatedOps : Pass<"decompose-aggregated-ops", "func::FuncOp"> {
+  let summary = "Decompose aggregated operations.";
+  let description = [{
+    Decompose operations that implement the `AggregatedOpInterface`.
+  }];
 }
 
 #endif // GC_DIALECT_GC_PASSES

lib/gc/Transforms/CMakeLists.txt

@@ -16,7 +16,8 @@ gc_add_mlir_library(GcPasses
   IterativeTilingAndFusion.cpp
   TilingUsingInterfaceX.cpp
   VerifyTargetDescription.cpp
-
+  DecomposeAggregatedOps.cpp
+  
   DEPENDS
     GraphCompilerPassIncGen
 

lib/gc/Transforms/DecomposeAggregatedOps.cpp

@@ -0,0 +1,49 @@
+//===- DecomposeAggregatedOps.cpp --------------------------------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "gc/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+namespace mlir {
+namespace gc {
+#define GEN_PASS_DEF_DECOMPOSEAGGREGATEDOPS
+#include "gc/Transforms/Passes.h.inc"
+} // namespace gc
+} // namespace mlir
+
+namespace {
+
+struct DecomposeAggregateOpsImpl : public OpRewritePattern<linalg::SoftmaxOp> {
+  using OpRewritePattern<linalg::SoftmaxOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(linalg::SoftmaxOp softmaxOp,
+                                PatternRewriter &rewriter) const override {
+    auto decomposableOp =
+        cast<linalg::AggregatedOpInterface>(softmaxOp.getOperation());
+    FailureOr<SmallVector<Value>> maybeNewResult =
+        decomposableOp.decomposeOperation(rewriter);
+    if (failed(maybeNewResult))
+      return failure();
+    rewriter.replaceOp(softmaxOp, *maybeNewResult);
+    return success();
+  }
+};
+
+struct DecomposeAggregatedOps
+    : public gc::impl::DecomposeAggregatedOpsBase<DecomposeAggregatedOps> {
+  void runOnOperation() override {
+    RewritePatternSet patterns(getOperation().getContext());
+    patterns.add<DecomposeAggregateOpsImpl>(patterns.getContext());
+    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+  }
+};
+
+} // namespace

lib/gc/Transforms/Pipeline.cpp

@@ -54,6 +54,8 @@ void populateTensorPasses(mlir::OpPassManager &pm) {
   // REMOVE this pass after the above passes are added. Currently we add this
   // pass to make the pipeline work properly
   pm.addNestedPass<func::FuncOp>(createLinalgGeneralizeNamedOpsPass());
+  // copied from tpp project
+  pm.addNestedPass<func::FuncOp>(createDecomposeAggregatedOps());
 }
 
 // scf + arith + math + vector + tensor + linalg.brgemm