cache

Author: yifeizh2

include/gc/Analysis/GlobalAnalysis.h

@@ -143,6 +143,9 @@ class GlobalAnalysis {
   DenseMap<Operation *, OperatorLayout> layoutCache;
 };
 
+namespace utils {
+bool isPackableNamedOp(Operation *op);
+}
 } // namespace gc
 } // namespace mlir
 

include/gc/Transforms/Passes.td

@@ -97,4 +97,15 @@ def PostProcessPackUnpack : Pass<"post-process-pack-unpack"> {
       ];
 }
 
+def LowerPackUnpack : Pass<"lower-pack-unpack"> {
+  let summary = "Lower pack and unpack ops.";
+  let description = [{
+    Lower pack and unpack into transpose and shape related ops.
+  }];
+  let dependentDialects = [
+      "mlir::tensor::TensorDialect",
+      "mlir::linalg::LinalgDialect"
+      ];
+}
+
 #endif // GC_DIALECT_GC_PASSES

lib/gc/Analysis/GlobalAnalysis.cpp

@@ -342,5 +342,17 @@ GlobalAnalysis::GlobalAnalysis(Operation *root) {
     }
   });
 }
+
+namespace utils {
+bool isPackableNamedOp(Operation *op) {
+  if ((isa<linalg::LinalgOp>(op) &&
+       !mlir::linalg::isaContractionOpInterface(
+           dyn_cast<linalg::LinalgOp>(op)) &&
+       !isa<linalgx::Mm4DVnniOp>(op)) ||
+      isa<tensor::ExpandShapeOp>(op) || isa<tensor::PadOp>(op))
+    return true;
+  return false;
+}
+} // namespace utils
 } // namespace gc
 } // namespace mlir

lib/gc/Transforms/CMakeLists.txt

@@ -14,6 +14,7 @@ add_mlir_library(GCPasses
   Pipeline.cpp
   PropagateLayout.cpp
   PostProcessPackUnpack.cpp
+  LowerPackUnpack.cpp
   DeepTileContractionNamedOp.cpp
   TilingUtil.cpp
   SinkOpIntoInnerLoop.cpp

lib/gc/Transforms/LowerPackUnpack.cpp

@@ -0,0 +1,83 @@
+//===- LowerPackUnpack.cpp - Lower pack unpack into linalg ops *---- 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_LOWERPACKUNPACK
+#include "gc/Transforms/Passes.h.inc"
+
+#define DEBUG_TYPE "lower-pack-unpack"
+
+using namespace mlir;
+
+// copied from tpp
+// A wrapper pattern that calls linalg::lowerPack on tensor::PackOp. It lowers
+// a tensor.pack op to tensor.pad + tensor.expand_shape + linalg.transpose ops.
+struct LowerPackPattern : public OpRewritePattern<tensor::PackOp> {
+  using OpRewritePattern<tensor::PackOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::PackOp op,
+                                PatternRewriter &rewriter) const override {
+    FailureOr<linalg::LowerPackResult> res = linalg::lowerPack(rewriter, op);
+    if (failed(res)) {
+      return rewriter.notifyMatchFailure(
+          op, "cannot lower to pad + expand + transpose");
+    }
+    return success();
+  }
+};
+
+// A wrapper pattern that calls linalg::lowerUnPack on tensor::UnPackOp. It
+// lowers a tensor.unpack op to tensor.empty + linalg.transpose +
+// tensor.collapse_shape + tensor.extract_slice ops.
+struct LowerUnPackPattern : public OpRewritePattern<tensor::UnPackOp> {
+  using OpRewritePattern<tensor::UnPackOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::UnPackOp op,
+                                PatternRewriter &rewriter) const override {
+    if (failed(linalg::lowerUnPack(rewriter, op))) {
+      return rewriter.notifyMatchFailure(
+          op, "cannot lower to empty + transpose + reshape + extract_slice");
+    }
+    return success();
+  }
+};
+
+class LowerPackUnpack : public impl::LowerPackUnpackBase<LowerPackUnpack> {
+public:
+  using impl::LowerPackUnpackBase<LowerPackUnpack>::LowerPackUnpackBase;
+  void runOnOperation() final;
+};
+
+void LowerPackUnpack::runOnOperation() {
+  auto *ctx = &getContext();
+  RewritePatternSet patterns(ctx);
+  patterns.add<LowerPackPattern, LowerUnPackPattern>(ctx);
+  (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+}
+
+} // namespace gc
+} // namespace mlir

lib/gc/Transforms/Pipeline.cpp

@@ -58,6 +58,8 @@ void populateTensorPasses(mlir::OpPassManager &pm) {
   pm.addNestedPass<func::FuncOp>(createLinalgGeneralizeNamedOpsPass());
   pm.addPass(createLoopInvariantCodeMotionPass());
   pm.addPass(createControlFlowSinkPass());
+  // TODO(yifei): remove lower pack here
+  pm.addPass(createLowerPackUnpack());
   populateCleanUpPasses(pm);
 }
 

lib/gc/Transforms/PostProcessPackUnpack.cpp

@@ -126,39 +126,6 @@ struct EliminateDummyUnpack : public OpRewritePattern<tensor::UnPackOp> {
   }
 };
 
-// copied from tpp
-// A wrapper pattern that calls linalg::lowerPack on tensor::PackOp. It lowers
-// a tensor.pack op to tensor.pad + tensor.expand_shape + linalg.transpose ops.
-struct LowerPackPattern : public OpRewritePattern<tensor::PackOp> {
-  using OpRewritePattern<tensor::PackOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(tensor::PackOp op,
-                                PatternRewriter &rewriter) const override {
-    FailureOr<linalg::LowerPackResult> res = linalg::lowerPack(rewriter, op);
-    if (failed(res)) {
-      return rewriter.notifyMatchFailure(
-          op, "cannot lower to pad + expand + transpose");
-    }
-    return success();
-  }
-};
-
-// A wrapper pattern that calls linalg::lowerUnPack on tensor::UnPackOp. It
-// lowers a tensor.unpack op to tensor.empty + linalg.transpose +
-// tensor.collapse_shape + tensor.extract_slice ops.
-struct LowerUnPackPattern : public OpRewritePattern<tensor::UnPackOp> {
-  using OpRewritePattern<tensor::UnPackOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(tensor::UnPackOp op,
-                                PatternRewriter &rewriter) const override {
-    if (failed(linalg::lowerUnPack(rewriter, op))) {
-      return rewriter.notifyMatchFailure(
-          op, "cannot lower to empty + transpose + reshape + extract_slice");
-    }
-    return success();
-  }
-};
-
 static void populateEliminateDummyPackUnpack(RewritePatternSet &patterns) {
   MLIRContext *ctx = patterns.getContext();
   patterns.add<EliminateDummyPack, EliminateDummyUnpack>(ctx);
@@ -205,7 +172,6 @@ void PostProcessPackUnpack::runOnOperation() {
   tppPopulateSimplifyPacking(patterns);
   // gc new packing related simplification
   populateEliminateDummyPackUnpack(patterns);
-  patterns.add<LowerPackPattern, LowerUnPackPattern>(ctx);
   (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }