Test collapse pack and unpack expand

Author: Jerry Wu

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

@@ -552,6 +552,192 @@ class BubbleUpPackThroughPadOp final : public OpRewritePattern<tensor::PackOp> {
   ControlPropagationFn controlFn;
 };
 
+static LogicalResult
+bubbleUpPackOpThroughCollapseShape(tensor::CollapseShapeOp collapseOp,
+                                   tensor::PackOp packOp,
+                                   PatternRewriter &rewriter) {
+  SmallVector<int64_t> innerTileSizes = packOp.getStaticTiles();
+  ArrayRef<int64_t> innerDimsPos = packOp.getInnerDimsPos();
+  ArrayRef<int64_t> outerDimsPerm = packOp.getOuterDimsPerm();
+
+  if (llvm::any_of(innerTileSizes,
+                   [](int64_t size) { return ShapedType::isDynamic(size); })) {
+    return failure();
+  }
+
+  ArrayRef<int64_t> srcShape = collapseOp.getSrcType().getShape();
+  SmallVector<ReassociationIndices> reassocIndices =
+      collapseOp.getReassociationIndices();
+  SmallVector<int64_t> baseDimsPos;
+  for (auto pos : innerDimsPos) {
+    baseDimsPos.push_back(reassocIndices[pos].back());
+  }
+  // Check if the base dims before reassociation are divisible by the inner tile
+  // sizes.
+  for (auto [basePos, tileSize] :
+       llvm::zip_equal(baseDimsPos, innerTileSizes)) {
+    int64_t dim = srcShape[basePos];
+    if (ShapedType::isDynamic(dim) || (dim % tileSize) != 0) {
+      return failure();
+    }
+  }
+  // Expand the outer dims perm with associated src dims.
+  SmallVector<int64_t> newOuterDimsPerm;
+  for (auto outerPos : outerDimsPerm) {
+    newOuterDimsPerm.insert(newOuterDimsPerm.end(),
+                            reassocIndices[outerPos].begin(),
+                            reassocIndices[outerPos].end());
+  }
+
+  auto emptyOp = tensor::PackOp::createDestinationTensor(
+      rewriter, packOp.getLoc(), collapseOp.getSrc(), packOp.getMixedTiles(), baseDimsPos,
+      newOuterDimsPerm);
+  auto newPackOp = rewriter.create<tensor::PackOp>(
+      packOp.getLoc(), collapseOp.getSrc(), emptyOp, baseDimsPos, packOp.getMixedTiles(),
+      packOp.getPaddingValue(), newOuterDimsPerm);
+
+  SmallVector<ReassociationIndices> newReassocIndices;
+  int64_t currPos = 0;
+  for (auto outerPos : outerDimsPerm) {
+    int64_t start = currPos;
+    int64_t end = start + reassocIndices[outerPos].size();
+    newReassocIndices.push_back(llvm::to_vector(llvm::seq(start, end)));
+    currPos = end;
+  }
+  for (auto unused : innerTileSizes) {
+    (void)unused;
+    newReassocIndices.push_back({currPos});
+    currPos += 1;
+  }
+
+  auto newCollapseOp = rewriter.create<tensor::CollapseShapeOp>(
+      collapseOp.getLoc(), packOp.getType(), newPackOp, newReassocIndices);
+  rewriter.replaceOp(packOp, newCollapseOp);
+
+  return success();
+}
+
+class BubbleUpPackOpThroughReshapeOp final
+    : public OpRewritePattern<tensor::PackOp> {
+public:
+  BubbleUpPackOpThroughReshapeOp(MLIRContext *context, ControlPropagationFn fun)
+      : OpRewritePattern<tensor::PackOp>(context), controlFn(std::move(fun)) {}
+
+  LogicalResult matchAndRewrite(tensor::PackOp packOp,
+                                PatternRewriter &rewriter) const override {
+    if (packOp.getPaddingValue())
+      return failure();
+
+    Operation *srcOp = packOp.getSource().getDefiningOp();
+    if (!srcOp || !(srcOp->getNumResults() == 1) ||
+        !srcOp->getResult(0).hasOneUse())
+      return failure();
+
+    if (auto collapseOp = dyn_cast<tensor::CollapseShapeOp>(srcOp)) {
+      return bubbleUpPackOpThroughCollapseShape(collapseOp, packOp, rewriter);
+    }
+    return failure();
+  }
+
+private:
+  ControlPropagationFn controlFn;
+};
+
+static LogicalResult
+pushDownUnPackOpThroughExpandShape(tensor::UnPackOp unPackOp,
+                                   tensor::ExpandShapeOp expandOp,
+                                   PatternRewriter &rewriter) {
+
+  SmallVector<int64_t> innerTileSizes = unPackOp.getStaticTiles();
+  ArrayRef<int64_t> innerDimsPos = unPackOp.getInnerDimsPos();
+  ArrayRef<int64_t> outerDimsPerm = unPackOp.getOuterDimsPerm();
+
+  if (llvm::any_of(innerTileSizes,
+                   [](int64_t size) { return ShapedType::isDynamic(size); })) {
+    return failure();
+  }
+
+  ArrayRef<int64_t> dstShape = expandOp.getType().getShape();
+  SmallVector<ReassociationIndices> reassocIndices =
+      expandOp.getReassociationIndices();
+  SmallVector<int64_t> baseDimsPos;
+  for (auto pos : innerDimsPos) {
+    baseDimsPos.push_back(reassocIndices[pos].back());
+  }
+  // Check if the base dims after reassociation are divisible by the inner tile
+  // sizes.
+  for (auto [basePos, tileSize] :
+       llvm::zip_equal(baseDimsPos, innerTileSizes)) {
+    int64_t dim = dstShape[basePos];
+    if (ShapedType::isDynamic(dim) || dstShape[basePos] % tileSize != 0) {
+      return failure();
+    }
+  }
+  // Expand the outer dims perm with associated src dims.
+  SmallVector<int64_t> newOuterDimsPerm;
+  for (auto outerPos : outerDimsPerm) {
+    newOuterDimsPerm.insert(newOuterDimsPerm.end(),
+                            reassocIndices[outerPos].begin(),
+                            reassocIndices[outerPos].end());
+  }
+
+  SmallVector<ReassociationIndices> newReassocIndices;
+  int64_t currPos = 0;
+  for (auto outerPos : outerDimsPerm) {
+    int64_t start = currPos;
+    int64_t end = start + reassocIndices[outerPos].size();
+    newReassocIndices.push_back(llvm::to_vector(llvm::seq(start, end)));
+    currPos = end;
+  }
+  for (auto unused : innerTileSizes) {
+    (void)unused;
+    newReassocIndices.push_back({currPos});
+    currPos += 1;
+  }
+
+  RankedTensorType newExpandType = tensor::PackOp::inferPackedType(
+      expandOp.getType(), innerTileSizes, baseDimsPos, newOuterDimsPerm);
+  auto newExpandOp = rewriter.create<tensor::ExpandShapeOp>(
+      expandOp.getLoc(), newExpandType, unPackOp.getSource(),
+      newReassocIndices);
+
+  auto emptyOp = tensor::UnPackOp::createDestinationTensor(
+      rewriter, unPackOp.getLoc(), newExpandOp, unPackOp.getMixedTiles(), baseDimsPos,
+      newOuterDimsPerm);
+  auto newUnPackOp = rewriter.create<tensor::UnPackOp>(
+      unPackOp.getLoc(), newExpandOp.getResult(), emptyOp, baseDimsPos,
+      unPackOp.getMixedTiles(), newOuterDimsPerm);
+  rewriter.replaceOp(expandOp, newUnPackOp);
+
+  return success();
+}
+
+class PushDownUnPackOpThroughReshapeOp final
+    : public OpRewritePattern<tensor::UnPackOp> {
+public:
+  PushDownUnPackOpThroughReshapeOp(MLIRContext *context,
+                                   ControlPropagationFn fun)
+      : OpRewritePattern<tensor::UnPackOp>(context), controlFn(std::move(fun)) {
+  }
+
+  LogicalResult matchAndRewrite(tensor::UnPackOp unPackOp,
+                                PatternRewriter &rewriter) const override {
+    Value result = unPackOp.getResult();
+    if (!result.hasOneUse()) {
+      return failure();
+    }
+    Operation *userOp = *result.user_begin();
+
+    if (auto expandOp = dyn_cast<tensor::ExpandShapeOp>(userOp)) {
+      return pushDownUnPackOpThroughExpandShape(unPackOp, expandOp, rewriter);
+    }
+    return failure();
+  }
+
+private:
+  ControlPropagationFn controlFn;
+};
+
 // TODO: Relax this restriction. We should unpack a generic op also
 // in the presence of multiple unpack ops as producers.
 /// Return the unpacked operand, if present, for the current generic op.
@@ -774,6 +960,7 @@ void mlir::linalg::populateDataLayoutPropagationPatterns(
     const ControlPropagationFn &controlPackUnPackPropagation) {
   patterns
       .insert<BubbleUpPackOpThroughGenericOpPattern, BubbleUpPackThroughPadOp,
-              PushDownUnPackOpThroughGenericOp, PushDownUnPackThroughPadOp>(
+              BubbleUpPackOpThroughReshapeOp, PushDownUnPackOpThroughGenericOp,
+              PushDownUnPackThroughPadOp, PushDownUnPackOpThroughReshapeOp>(
           patterns.getContext(), controlPackUnPackPropagation);
 }