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

Co-authored-by: Han-Chung Wang <[email protected]>
Signed-off-by: Hyunsung Lee <[email protected]>

Author: ita9naiwa

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

@@ -190,29 +190,18 @@ def Linalg_PackOp : Linalg_RelayoutOp<"pack", [
     // Method to get the `RankedTensorType` of the result based on the inner
     // tiles, position of the inner tiles (innerDimsPos)  and interchange vector
     // of outer loops (outerDimsPerm).
-    /// This method uses inferPackedShape to ensure consistency with other shape
-    /// inference methods regarding which dimensions are dynamic.
     static RankedTensorType inferPackedTensorType(RankedTensorType sourceType,
         ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
         ArrayRef<int64_t> outerDimsPerm = {});
 
     // Method to get the `MemRefType` of the result based on the inner
     // tiles, position of the inner tiles (innerDimsPos)  and interchange vector
     // of outer loops (outerDimsPerm).
-    /// This method uses inferPackedShape to ensure consistency with other shape
-    /// inference methods regarding which dimensions are dynamic.
     static MemRefType inferPackedMemRefType(MemRefType sourceType,
         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).
-
-    /// Helper for PackOp::{getResultShape, inferPackedTensorType, inferPackedMemRefType}.
-    /// Returns the shape of the packed type. Having a shared helper helps
-    /// implement these three methods in a way that ensures
-    /// that they agree on which dimensions are dynamic.
+    // Returns the shape of the packed type. It is a shared helper helps type inference methods in a way that ensures  that they agree on which dimensions are dynamic.
     static SmallVector<int64_t> inferPackedShape(ArrayRef<int64_t> inputShape,
         ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
         ArrayRef<int64_t> outerDimsPerm = {});

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

@@ -5006,7 +5006,8 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
 
   // Insert tensor.cast ops if static shape inference is available..
   SmallVector<int64_t> srcShape, destShape;
-  if (inferStaticShape(packOp, srcShape, destShape)) {
+  if (inferStaticShape(packOp, srcShape, destShape) &&
+      packOp.hasPureTensorSemantics()) {
     Location loc = packOp.getLoc();
     Value source = packOp.getSource();
     if (srcShape != packOp.getSourceType().getShape()) {
@@ -5030,20 +5031,11 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
     // Insert a cast if needed
     if (needUpdateDestType) {
       rewriter.setInsertionPointAfter(packOp);
-      Operation *castOp;
-      bool hasTensorSemantics = packOp.hasPureTensorSemantics();
-      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);
-    } else {
-      // TODO: support memref.cast if static shape inference is available.
-      return failure();
+      auto castOp =
+          rewriter.create<tensor::CastOp>(loc, originalResultType, packOp);
+      rewriter.replaceAllUsesExcept(packOp, castOp, castOp);
     }
+
     return success();
   }
 
@@ -5424,7 +5416,7 @@ struct FoldTensorCastUnPackOp : public OpRewritePattern<UnPackOp> {
     if (!tensor::hasFoldableTensorCastOperand(op))
       return failure();
 
-    // TODO: Support Memref PackOp. Temporarily return failure.
+    // TODO: Support Memref UnPackOp. Temporarily return failure.
     if (!op.hasPureTensorSemantics())
       return failure();
 

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

@@ -268,7 +268,7 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
     highs[pos] = affine::makeComposedFoldedAffineApply(
         rewriter, loc, map, {outerSize, origSize, innerSize});
   }
-  // TODO: Need memref.pad operation to support memref operands
+
   RankedTensorType collapsed = tensor::CollapseShapeOp::inferCollapsedType(
       RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape),
       packingMetadata.reassociations);
@@ -1030,9 +1030,6 @@ LogicalResult ExtractSliceOfPadTensorSwapPattern::matchAndRewrite(
 static Value getPackOpSourceOrPaddedSource(OpBuilder &builder,
                                            linalg::PackOp packOp) {
   Value input = packOp.getSource();
-  // TODO: Support Memref PackOp. Temporarily return just Op Source.
-  if (!packOp.hasPureTensorSemantics())
-    return input;
 
   if (!packOp.getPaddingValue()) {
     return input;

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

@@ -1889,9 +1889,8 @@ static LogicalResult
 vectorizeUnPackOpPrecondition(linalg::UnPackOp unpackOp,
                               ArrayRef<int64_t> inputVectorSizes) {
   // TODO: Support Memref PackOp. Temporarily return failure.
-  if (!unpackOp.hasPureTensorSemantics()) {
+  if (!unpackOp.hasPureTensorSemantics())
     return failure();
-  }
 
   if (llvm::any_of(unpackOp.getInnerTiles(), [](OpFoldResult res) {
         return !getConstantIntValue(res).has_value();

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

@@ -315,12 +315,11 @@ SmallVector<Range> SliceFromCollapseHelper::getExtractSliceParams(
     // have proven that these are not sliced. In this case we just take
     // the full extent of each dimension in the reassociation list.
     if (linearizedDimensions[it.index()]) {
-      llvm::append_range(
-          offsetsSizesAndStrides,
-          llvm::map_range(it.value(), [&](int64_t idx) -> Range {
-            return {zeroAttr, collapseShapeInputShape[idx], oneAttr};
-          }));
-
+      llvm::append_range(offsetsSizesAndStrides,
+                         llvm::map_range(it.value(), [&](int64_t idx) -> Range {
+                           return {zeroAttr, collapseShapeInputShape[idx],
+                                   oneAttr};
+                         }));
       continue;
     }
 

mlir/test/Dialect/Linalg/canonicalize.mlir

@@ -1894,9 +1894,7 @@ func.func @fold_cast_unpack_dynamic_tile_size(
     return %unpack : tensor<7x?xi32>
 }
 
-//===----------------------------------------------------------------------===//
-// linalg.unpack + linalg.pack
-//===----------------------------------------------------------------------===//
+// -----
 
 // CHECK-LABEL: func.func @fold_pack_unpack_tensor
 // CHECK-SAME:  (%[[ARG0:.*]]: tensor<2x3xf32>) -> tensor<2x3xf32>