add postprocess pack passes

Author: yifeizh2

include/gc/Transforms/Passes.td

@@ -58,4 +58,15 @@ def PropagateLayoutOnNamedOps : Pass<"propagate-layout-on-named-ops"> {
       ];
 }
 
+def PostProcessPackUnpack : Pass<"post-process-pack-unpack"> {
+  let summary = "Fold and simplify pack and unpack ops.";
+  let description = [{
+    Fold and simplify pack and unpack ops.
+  }];
+  let dependentDialects = [
+      "mlir::tensor::TensorDialect",
+      "mlir::linalg::LinalgDialect"
+      ];
+}
+
 #endif // GC_DIALECT_GC_PASSES

lib/gc/Transforms/CMakeLists.txt

@@ -11,6 +11,7 @@ add_mlir_library(GCPasses
   OneDNNGraphToLinalg.cpp
   Pipeline.cpp
   PropagateLayout.cpp
+  PostProcessPackUnpack.cpp
   TileNamed.cpp
 
   ADDITIONAL_HEADER_DIRS

lib/gc/Transforms/Pipeline.cpp

@@ -40,6 +40,8 @@ void populateFrontendPasses(mlir::PassManager &pm) {
 void populateTensorPasses(mlir::PassManager &pm) {
   // todo: padding propagation pass
   // todo: layout propagation pass
+  pm.addPass(createPropagateLayoutOnNamedOps());
+  pm.addPass(createPostProcessPackUnpack());
   // todo: tensor constant propagation pass
   // todo: linalg.matmul lowering to (scf.loop + linalg.brgemm) pass
   // todo: fine-grain fusion pass

lib/gc/Transforms/PostProcessPackUnpack.cpp

@@ -0,0 +1,127 @@
+//===- PostProcessPackUnpack.cpp - Fold and simplify pack unpack *-- C++-*-===//
+//
+// This file is only temporarily used to extend upstream or upcoming utility in
+// TilingInterface, which finally aims for upstream.
+//
+//===----------------------------------------------------------------------===//
+
+#include <iostream>
+#include <numeric>
+
+#include "gc/Transforms/Transforms.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "gc/Dialect/Linalgx/LinalgxDialect.h"
+#include "gc/Dialect/Linalgx/LinalgxOps.h"
+#include "gc/Transforms/Passes.h"
+namespace mlir {
+namespace gc {
+#define GEN_PASS_DEF_POSTPROCESSPACKUNPACK
+#include "gc/Transforms/Passes.h.inc"
+
+#define DEBUG_TYPE "post-process-pack-unpack"
+
+using namespace mlir;
+
+// Helper pattern - lower tensor.pack operations that pack constants.
+struct LowerConstantPacking : public OpRewritePattern<tensor::PackOp> {
+  using OpRewritePattern<tensor::PackOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::PackOp packOp,
+                                PatternRewriter &rewriter) const override {
+    auto constOp = packOp.getSource().getDefiningOp<arith::ConstantOp>();
+    if (!constOp)
+      return failure();
+    // Must be a dense constant.
+    auto denseAttr = dyn_cast<DenseElementsAttr>(constOp.getValue());
+    if (!denseAttr)
+      return failure();
+
+    // Bail out if the pack is used as a writing operation i.e., the destination
+    // is not a tensor.empty.
+    if (!packOp.getDest().getDefiningOp<tensor::EmptyOp>())
+      return rewriter.notifyMatchFailure(packOp,
+                                         "expects empty tensor destination");
+    // Pack destination must have static shape.
+    if (!packOp.getDestType().hasStaticShape())
+      return rewriter.notifyMatchFailure(
+          packOp, "expects destination with static shape");
+
+    // If it is a splat constant, skip and let tensor.pack folder to handle this
+    // case.
+    if (denseAttr.isSplat())
+      return rewriter.notifyMatchFailure(
+          packOp, "skip pack - existing folder covers constant splats");
+
+    return linalg::lowerPack(rewriter, packOp);
+  }
+};
+
+static void tppPopulateConstantFoldPack(RewritePatternSet &patterns) {
+  MLIRContext *ctx = patterns.getContext();
+  patterns.add<LowerConstantPacking>(ctx);
+  // Apply canonicalization to fold trivial cases and linalg constant folders
+  // to cleanup lowered packs.
+  linalg::FillOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::PackOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::populateRewriteAsConstantPatterns(
+      patterns, [](OpOperand *) -> bool { return true; });
+  linalg::populateConstantFoldLinalgOperations(
+      patterns, [](OpOperand *) -> bool { return true; });
+}
+
+class PostProcessPackUnpack
+    : public impl::PostProcessPackUnpackBase<PostProcessPackUnpack> {
+public:
+  using impl::PostProcessPackUnpackBase<
+      PostProcessPackUnpack>::PostProcessPackUnpackBase;
+  void runOnOperation() final;
+};
+
+static void tppPopulateSimplifyPacking(RewritePatternSet &patterns) {
+  MLIRContext *ctx = patterns.getContext();
+  tensor::populateSimplifyPackAndUnpackPatterns(patterns);
+  tensor::populateFoldTensorEmptyPatterns(patterns);
+  tensor::PackOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::UnPackOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::CastOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::InsertSliceOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::EmptyOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::PadOp::getCanonicalizationPatterns(patterns, ctx);
+  tensor::ParallelInsertSliceOp::getCanonicalizationPatterns(patterns, ctx);
+  scf::ForallOp::getCanonicalizationPatterns(patterns, ctx);
+  // Propagate packs/unpacks only through expand shapes at this point.
+  // This captures the transformation scope of the replaced downstream pass.
+  linalg::populateDataLayoutPropagationPatterns(
+      patterns, [](Operation *op) { return isa<tensor::ExpandShapeOp>(op); });
+  ctx->getLoadedDialect<tensor::TensorDialect>()->getCanonicalizationPatterns(
+      patterns);
+  // patterns.add<FoldUnPackIntoInsertSlice>(ctx);
+  tensor::populateReassociativeReshapeFoldingPatterns(patterns);
+}
+
+void PostProcessPackUnpack::runOnOperation() {
+  MLIRContext *ctx = getOperation()->getContext();
+  RewritePatternSet patterns(ctx);
+
+  // constant fold packing
+  tppPopulateConstantFoldPack(patterns);
+  // simplify packing
+  tppPopulateSimplifyPacking(patterns);
+  (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+}
+
+} // namespace gc
+} // namespace mlir