From 854953fd512bf94042ea5881b96978daf684bb79 Mon Sep 17 00:00:00 2001
From: Chao Chen <116223022+chencha3@users.noreply.github.com>
Date: Tue, 9 Jul 2024 11:26:12 -0500
Subject: [PATCH] Update XeTileBlocking  (#802)

Generalize the logic for StoreTileOp
---
 lib/Dialect/XeTile/Transforms/Blocking.cpp | 13 +++----------
 1 file changed, 3 insertions(+), 10 deletions(-)

diff --git a/lib/Dialect/XeTile/Transforms/Blocking.cpp b/lib/Dialect/XeTile/Transforms/Blocking.cpp
index 28ecf4506..77537edf2 100644
--- a/lib/Dialect/XeTile/Transforms/Blocking.cpp
+++ b/lib/Dialect/XeTile/Transforms/Blocking.cpp
@@ -932,17 +932,10 @@ struct StoreTileOpPattern
     auto tileTy = llvm::dyn_cast<xetile::TileType>(adaptor.getTile().getType());
     auto innerBlocks = tileTy.getInnerBlocks();
     auto value = adaptor.getValue();
-    // its inputs has not been updated yet.
-    if (innerBlocks && value.getDefiningOp<xetile::TileUnpackOp>()) {
-      value = addPackOp(value, innerBlocks.asArrayRef(), rewriter);
-      rewriter.replaceOpWithNewOp<xetile::StoreTileOp>(op, value,
-                                                       adaptor.getTile());
-      return mlir::success();
-    }
+    auto valTy = mlir::dyn_cast<mlir::VectorType>(value.getType());
 
-    // TODO: Blocking is not applied on shapecast yet, so it needs special
-    // attention.
-    if (innerBlocks && value.getDefiningOp<mlir::vector::ShapeCastOp>()) {
+    // its inputs has not been updated yet.
+    if (innerBlocks && valTy.getRank() == 2) {
       value = addPackOp(value, innerBlocks.asArrayRef(), rewriter);
       rewriter.replaceOpWithNewOp<xetile::StoreTileOp>(op, value,
                                                        adaptor.getTile());