fix upon review

Author: ita9naiwa

mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td

@@ -201,6 +201,13 @@ def Linalg_PackOp : Linalg_RelayoutOp<"pack", [
         ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
         ArrayRef<int64_t> outerDimsPerm = {});
 
+    // Method to get the Shape of the result based on the input shape, inner
+    // tiles, position of the inner tiles (innerDimsPos)  and interchange vector
+    // of outer loops (outerDimsPerm).
+    static SmallVector<int64_t> inferPackedShape(ArrayRef<int64_t> inputShape,
+        ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
+        ArrayRef<int64_t> outerDimsPerm = {});
+
     // Returns true if we have enough static information to catch undefined
     // behavior when the tile size does not divide perfectly the dimension of
     // the input tensor. Detecting UB requires that the input size and either

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

@@ -4487,17 +4487,13 @@ static LogicalResult commonVerifierPackAndUnPackOp(OpTy packOrUnPack) {
   // Verify result shape is greater than the minimum expected
   // by the pack operation, and that the output shape
   // represents full tiles.
-  if (hasTensorSemantics) {
-    RankedTensorType expectedPackedType = PackOp::inferPackedTensorType(
-        cast<RankedTensorType>(unpackedType), packOrUnPack.getStaticTiles(),
-        innerDimsPos, outerDimPerm);
-    if (!areAllInBound(expectedPackedType.getShape(), packedType.getShape())) {
-      return op->emitError(
-                 "the shape of output is not large enough to hold the "
-                 "packed data. Expected at least ")
-             << expectedPackedType << ", got " << packedType;
-    }
-  } else {
+  auto expectedPackedShape = PackOp::inferPackedShape(
+      unpackedType.getShape(), packOrUnPack.getStaticTiles(),
+      packOrUnPack.getInnerDimsPos(), packOrUnPack.getOuterDimsPerm());
+  if (!areAllInBound(expectedPackedShape, packedType.getShape())) {
+    return op->emitError("the shape of output is not large enough to hold the "
+                         "packed data. Expected at least ")
+           << expectedPackedShape << ", got " << packedType.getShape();
   }
   if (!llvm::all_of(
           llvm::zip(packedType.getShape().take_back(mixedTiles.size()),
@@ -4784,6 +4780,14 @@ MemRefType PackOp::inferPackedMemRefType(MemRefType sourceType,
   return MemRefType::get(resultShape, sourceType.getElementType());
 }
 
+SmallVector<int64_t> PackOp::inferPackedShape(ArrayRef<int64_t> inputShape,
+                                              ArrayRef<int64_t> innerTileSizes,
+                                              ArrayRef<int64_t> innerDimsPos,
+                                              ArrayRef<int64_t> outerDimsPerm) {
+  return getPackOpResultTypeShape(inputShape, innerTileSizes, innerDimsPos,
+                                  outerDimsPerm);
+}
+
 Value PackOp::createDestinationTensor(OpBuilder &b, Location loc, Value source,
                                       ArrayRef<OpFoldResult> innerTileSizes,
                                       ArrayRef<int64_t> innerDimsPos,
@@ -5030,7 +5034,6 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
     // Insert a cast if needed
     if (needUpdateDestType) {
       rewriter.setInsertionPointAfter(packOp);
-      /// 1
       if (hasTensorSemantics) {
         auto castOp =
             rewriter.create<tensor::CastOp>(loc, originalResultType, packOp);
@@ -5040,16 +5043,6 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
             rewriter.create<memref::CastOp>(loc, originalResultType, packOp);
         rewriter.replaceAllUsesExcept(packOp, castOp, castOp);
       }
-      /// 2
-      Operation *castOp;
-      if (hasTensorSemantics) {
-        castOp =
-            rewriter.create<tensor::CastOp>(loc, originalResultType, packOp);
-      } else {
-        castOp =
-            rewriter.create<memref::CastOp>(loc, originalResultType, packOp);
-      }
-      rewriter.replaceAllUsesExcept(packOp, castOp->getResult(0), castOp);
     }
     return success();
   }

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

@@ -63,13 +63,12 @@ getPackingInfoFromOperand(OpOperand *opOperand, linalg::GenericOp genericOp,
                           OpTy packOrUnPackOp) {
   static_assert(llvm::is_one_of<OpTy, linalg::PackOp, linalg::UnPackOp>::value,
                 "applies to only pack or unpack operations");
-  if (PackOp packOp = dyn_cast<PackOp>(packOrUnPackOp)) {
-    if (!packOp.hasPureTensorSemantics())
+  if constexpr (std::is_same_v<OpTy, linalg::PackOp>) {
+    if (!packOrUnPackOp.hasPureTensorSemantics())
       return failure();
   }
-
-  if (UnPackOp unpackOp = dyn_cast<UnPackOp>(packOrUnPackOp)) {
-    if (!unpackOp.hasPureTensorSemantics())
+  if constexpr (std::is_same_v<OpTy, linalg::UnPackOp>) {
+    if (!packOrUnPackOp.hasPureTensorSemantics())
       return failure();
   }