[Mlir-commits] [mlir] [mlir] Canonicalize extract_slice(unpack) (PR #133777)

[Han-Chung Wang]

================
@@ -5243,6 +5243,26 @@ LogicalResult UnPackOp::canonicalize(UnPackOp unPackOp,
                              [&]() { unPackOp.setDpsInitOperand(0, newDest); });
     return success();
   }
+  /// extract_slice(unpack(x)) -> unpack(x)
+  if (unPackOp->hasOneUse()) {
+    auto extractSliceUser =
+        dyn_cast<tensor::ExtractSliceOp>(*unPackOp->getUsers().begin());
+    if (extractSliceUser &&
+        areAllConstantIntValue(extractSliceUser.getMixedOffsets(), 0) &&
+        areAllConstantIntValue(extractSliceUser.getMixedStrides(), 1) &&
+        extractSliceUser.getSourceType().getRank() ==
+            extractSliceUser.getResultType().getRank()) {
+      auto newDest = rewriter.create<tensor::ExtractSliceOp>(
+          unPackOp->getLoc(), unPackOp.getDest(),
+          extractSliceUser.getMixedOffsets(), extractSliceUser.getMixedSizes(),
+          extractSliceUser.getMixedStrides());
+      rewriter.replaceOpWithNewOp<UnPackOp>(
+          extractSliceUser, unPackOp.getSource(), newDest,
+          unPackOp.getInnerDimsPos(), unPackOp.getMixedTiles(),
+          unPackOp.getOuterDimsPerm());
----------------
hanhanW wrote:

I think it is better to set insertion point to before `unpack` op and apply in-place modification.

https://github.com/llvm/llvm-project/blob/0b31f08537746beff4d5e0df44221cbe5a9237c5/mlir/include/mlir/IR/Builders.h#L402-L406

https://github.com/llvm/llvm-project/pull/133777