avoid double propagation

Author: yifeizh2

lib/gc/Analysis/GlobalAnalysis.cpp

@@ -22,27 +22,15 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &ss,
   SmallVector<int64_t> innerAxis = layoutCache.getInnerAxis();
   SmallVector<OpFoldResult> tileSizes = layoutCache.getTileSizes();
   ss << "[";
-  for (size_t i = 0; i < outerAxis.size(); ++i) {
-    if (i != 0) {
-      ss << ", ";
-    }
-    ss << outerAxis[i];
-  }
-  for (size_t i = 0; i < innerAxis.size(); ++i) {
-    ss << (i == 0 ? "; " : ", ");
-    ss << innerAxis[i];
+  llvm::interleaveComma(outerAxis, ss);
+  if (!innerAxis.empty()) {
+    ss << "; ";
+    llvm::interleaveComma(innerAxis, ss);
   }
   ss << "]";
   if (!tileSizes.empty()) {
     ss << "; {";
-    for (size_t i = 0; i < tileSizes.size(); ++i) {
-      if (i != 0) {
-        ss << ", ";
-      }
-      if (getConstantIntValue(tileSizes[i]).has_value()) {
-        ss << *getConstantIntValue(tileSizes[i]);
-      }
-    }
+    llvm::interleaveComma(tileSizes, ss);
     ss << "}";
   }
   return ss;
@@ -58,11 +46,11 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &ss,
                               const OperatorLayout &opLayout) {
   for (auto &&[idx, layoutCache] :
        llvm::enumerate(opLayout.getSupportedInputLayouts())) {
-    ss << "input " << idx << "'s layoutCache: " << layoutCache << "\n";
+    ss << "input " << idx << "'s layout: " << layoutCache << "\n";
   }
   for (auto &&[idx, layoutCache] :
        llvm::enumerate(opLayout.getSupportedOutputLayouts())) {
-    ss << "output " << idx << "'s layoutCache: " << layoutCache << "\n";
+    ss << "output " << idx << "'s layout: " << layoutCache << "\n";
   }
   return ss;
 }
@@ -156,15 +144,15 @@ inferTargetLayout(TensorLayout layoutBase,
 }
 
 static size_t getTargetInputIdx(ArrayRef<TensorLayout> curInputLayouts) {
-  for (auto i = 0; i < curInputLayouts.size(); ++i) {
+  for (size_t i = 0; i < curInputLayouts.size(); ++i) {
     if (!curInputLayouts[i].isPlainLayout()) {
       return i;
     }
   }
   return 0;
 }
 
-static bool supportedContractionOpList(linalg::LinalgOp &linalgOp) {
+static bool supportedContractionNamedOpList(linalg::LinalgOp &linalgOp) {
   if (isa<linalg::MatmulOp, linalg::MatmulTransposeAOp,
           linalg::MatmulTransposeBOp, linalg::BatchMatmulOp,
           linalg::BatchMatmulTransposeAOp, linalg::BatchMatmulTransposeBOp>(
@@ -211,7 +199,7 @@ GlobalAnalysis::GlobalAnalysis(Operation *root) {
       // ------ Get Current Op's Suggested Layout & Do Propagation ------
       IRRewriter rewriter(linalgOp);
       // TODO: extend to packed/vnni matmul ops
-      if (supportedContractionOpList(linalgOp)) {
+      if (supportedContractionNamedOpList(linalgOp)) {
         // get input and output rank
         auto ARank = cast<ShapedType>(linalgOp.getDpsInputs()[0].getType())
                          .getShape()
@@ -253,7 +241,8 @@ GlobalAnalysis::GlobalAnalysis(Operation *root) {
                                       rewriter.getIndexAttr(iin)});
         OperatorLayout suggestedLayout({ALayout, BLayout}, {CLayout});
         layoutCache[linalgOp] = suggestedLayout;
-      } else if (!mlir::linalg::isaContractionOpInterface(linalgOp)) {
+      } else if (!mlir::linalg::isaContractionOpInterface(linalgOp) &&
+                 !supportedContractionNamedOpList(linalgOp)) {
         SmallVector<TensorLayout> inputLayouts, outputLayouts;
         size_t targetIdx = getTargetInputIdx(curInputLayouts);
         // TODO(yifei): wisely choose the input format basis
@@ -345,11 +334,12 @@ 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))
+  if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) {
+    if (!supportedContractionNamedOpList(linalgOp)) {
+      return true;
+    }
+  } else if (isa<tensor::ExpandShapeOp, tensor::CollapseShapeOp, tensor::PadOp>(
+                 op))
     return true;
   return false;
 }

lib/gc/Transforms/PropagateLayout.cpp

@@ -6,7 +6,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include <iostream>
 #include <numeric>
 
 #include "gc/Transforms/Transforms.h"
@@ -209,8 +208,34 @@ FailureOr<linalg::PackResult> packNamedOp(RewriterBase &rewriter,
 }
 
 // check whether the op is already packed or not
-static bool checkPacked(Operation *op, const OperatorLayout &layout) {
+static bool checkPacked(Operation *op, const OperatorLayout &opLayout) {
   // check whether rank match
+  if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) {
+    assert(linalgOp.getDpsInits().size() ==
+               opLayout.getSupportedOutputLayouts().size() &&
+           linalgOp.getDpsInputs().size() ==
+               opLayout.getSupportedInputLayouts().size());
+    for (auto [index, layout] :
+         llvm::enumerate(opLayout.getSupportedInputLayouts())) {
+      // if dimension mismatch, then the op itself is already packed
+      if (layout.getOuterAxis().size() !=
+          cast<RankedTensorType>(linalgOp.getDpsInputs()[index].getType())
+              .getShape()
+              .size())
+        return true;
+    }
+    for (auto [index, layout] :
+         llvm::enumerate(opLayout.getSupportedOutputLayouts())) {
+      // if dimension mismatch, then the op itself is already packed
+      if (layout.getOuterAxis().size() !=
+          cast<RankedTensorType>(linalgOp.getDpsInits()[index].getType())
+              .getShape()
+              .size())
+        return true;
+    }
+  } else {
+    assert(op->getNumOperands() == 1 && op->getNumResults() == 1);
+  }
   return false;
 }
 
@@ -225,11 +250,28 @@ class PropagateLayoutOnNamedOps
   void runOnOperation() final;
 };
 
+LogicalResult graphAlreadyPacked(MLIRContext *ctx, mlir::Operation *graph) {
+  IRRewriter rewriter(ctx);
+  auto walk = graph->walk([&](Operation *op) {
+    if (mlir::gc::utils::isPackableNamedOp(op) && op->hasAttr("packed")) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "Graph already packed. Stop layout propagation.\n");
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+  if (walk.wasInterrupted()) {
+    return failure();
+  }
+  return success();
+}
+
 LogicalResult namedOpLayoutPropagation(MLIRContext *ctx, mlir::Operation *graph,
                                        ControlPackNamedOpsFn controlFn) {
   IRRewriter rewriter(ctx);
   auto walk = graph->walk([&](Operation *op) {
     if (mlir::gc::utils::isPackableNamedOp(op)) {
+      LLVM_DEBUG(llvm::dbgs() << "Op " << op->getName() << " visited.\n");
       FailureOr<OperatorLayout> opLayout = controlFn(op);
       if (failed(opLayout)) {
         LLVM_DEBUG(llvm::dbgs() << "Op " << op->getName()
@@ -258,6 +300,9 @@ LogicalResult namedOpLayoutPropagation(MLIRContext *ctx, mlir::Operation *graph,
             packNamedOp(rewriter, linalgOp, *opLayout);
         if (failed(packedOp)) {
           return WalkResult::skip();
+        } else {
+          packedOp->packedLinalgOp->setAttr("packed",
+                                            rewriter.getBoolAttr(true));
         }
       } else if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(op)) {
         // Location loc = expandShapeOp->getLoc();
@@ -291,8 +336,6 @@ LogicalResult namedOpLayoutPropagation(MLIRContext *ctx, mlir::Operation *graph,
     }
     return WalkResult::advance();
   });
-  if (walk.wasSkipped())
-    return failure();
   return success();
 }
 
@@ -576,6 +619,9 @@ struct UpliftPackOverBroadcast : public OpRewritePattern<tensor::PackOp> {
 void PropagateLayoutOnNamedOps::runOnOperation() {
   MLIRContext *ctx = &getContext();
   mlir::Operation *graph = getOperation();
+  // stage0: check if the graph has been packed
+  if (failed(graphAlreadyPacked(ctx, graph)))
+    return;
   // stage1: pack matmul
   RewritePatternSet packMatmulPatterns(&getContext());
   mlir::linalg::ControlBlockPackMatmulFn packMatmulControlFn =