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

[Han-Chung Wang]

================
@@ -190,7 +190,7 @@ 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).
-    static RankedTensorType inferPackedType(RankedTensorType sourceType,
+    static RankedTensorType inferPackedType(ShapedType sourceType,
----------------
hanhanW wrote:

This looks more like a `inferPackedTensorType` to me now. Because it always returns a `RankedTensorType`. In the LLVM codebase, it is used in op verification for shapes and the data-layout propagation pass. IMO, we can leave it as what it is and update the verifier.

We can add a new `inferShape` to the pack op and use it in the verifier.

https://github.com/llvm/llvm-project/blob/4c4e4e4299b16d8dd85811e9dd8697b17c95577f/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp#L4471-L4480

On the other hand, please revisit the verifier. I'd expect some changes there. E.g., we do not support unranked tensor/memref in Linalg ops, at least for these pack/unpack ops. Without your change, it emits an error like `custom op 'linalg.pack' invalid kind of type specified`. With your change, it crashes. This is because we switch to `AnyShaped`, and there is no longer "hasRank" trait in the op definition.

To repro: `mlir-opt repro.mlir`

```mlir
func.func @pack_tensor(%source: tensor<*xf32>, %dest: tensor<*xf32>) -> tensor<*xf32>{
  %0 = linalg.pack %source outer_dims_perm = [] inner_dims_pos = [] inner_tiles = []
      into %dest : tensor<*xf32> -> tensor<*xf32>                                                                                                                                                                                                              return %0 : tensor<*xf32>
}
```

IMO, we should check that the types are ranked, and the shapes are compatible.

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