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

[Han-Chung Wang]

https://github.com/hanhanW commented:

Sorry that I made a mistake in the review about the folding. After taking a look at the real IR tests, I think there are two issues.

1. The memref version should not return a memref. By definition, the op performs packing from the source buffer and store the result to the destination buffer. Like other linalg operations, it should not return any value when it has buffer semantics. In other linalg ops, the `Variadic<AnyRankedTensor>` adds the check. It was pointed out in [a comment](https://github.com/llvm/llvm-project/pull/129036/files#r1977098388) from @adam-smnk, and I think it is not resolved.

Invalid case:

```mlir
%packed = linalg.pack %unpacked inner_dims_pos = [0, 1] inner_tiles = [4, 4] into %buf_packed : memref<40x80xf32> -> memref<10x20x4x4xf32>
```

It should be like this valid case:

```mlir
linalg.pack %unpacked
  inner_dims_pos = [0, 1] inner_tiles = [4, 4]
  into %buf_packed
  : memref<40x80xf32> -> memref<10x20x4x4xf32>
```

2. The second issue is about the folding on memrefs. TLDR is that we should disable it. Because the current logic is not correct. How I'd implement the folding is that we need to look at the source memref and other uses of the source memref. It could be expensive and the logic is more complicated on memrefs.  The below is not foldable if there are other ops using `%buf_unpacked` in the middle. Things become way more complicated when control flow is involved. I think this kind of folding should be implemented in a pass manner. (Please remember to remove the test cases in `canonicalzation.mlir` if they are disabled. I also did not see the reason why you added some folding tests to tensor, while the PR is scoped to add the support for memref. Maybe we can drop the tests for tensors as well?)

```mlir
linalg.unpack %t
  inner_dims_pos = [0, 1] inner_tiles = [4, 4]
  into %buf_unpacked
  : memref<10x20x4x4xf32> -> memref<40x80xf32>
linalg.pack %unpacked
  inner_dims_pos = [0, 1] inner_tiles = [4, 4]
  into %buf_packed
  : memref<40x80xf32> -> memref<10x20x4x4xf32>
```

(The rest issue is about reverting non-related code changes, I think you mentioned that you will remove them before landing the PR.)

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