[Mlir-commits] [mlir] [MLIR][Linalg] pack, unpack to take memref inputs (PR #129036)

[Han-Chung Wang]

================
@@ -4942,7 +5015,7 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
           rewriter.create<tensor::CastOp>(loc, newSrcType, packOp.getSource());
     }
     Value dest = packOp.getDest();
-    RankedTensorType originalResultType = packOp.getDestType();
+    ShapedType originalResultType = packOp.getDestType();
----------------
hanhanW wrote:

This canonicalization pattern is broken when pack/unpack are on memrefs. We should only update the result type when they are on tensors. The memref version does not return any SSA value.

```
    rewriter.modifyOpInPlace(packOp, [&] {
      packOp.getSourceMutable().assign(source);
      packOp.getDestMutable().assign(dest);
      packOp.getResult().setType(cast<RankedTensorType>(dest.getType()));
    });
```

Also, the above cast should be memref.cast if we are going to support it in the PR. If you don't know how to implement it, we can bail out in line 5009 and I can help add the support later.


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