[mlir][tensor] Bufferize tensor.reshape with non-identity layouts

[sabauma]

Bufferization of tensor.reshape generates a memref.reshape operation. memref.reshape requires the source memref to have an identity layout. The bufferization process may result in the source memref having a non-identity layout, resulting in a verification failure.

This change causes the bufferization interface for tensor.reshape to copy the source memref to a new buffer when the source has a non-identity layout.

[Lewuathe]

In the documentation, memref.reshape seems to expect an identity layout for both inputs. Do we need to take care of the case with the non-identity map?

https://mlir.llvm.org/docs/Dialects/MemRef/#memrefreshape-memrefreshapeop

The source and destination types are compatible if both have the same element type, same number of elements, address space and identity layout map.

[sabauma]

Hi @Lewuathe, are you asking whether the restrictions on memref.reshape need to be loosened in order to permit non-identity layout maps? That could be a potential solution, though I am not privy to the reasoning behind the layout map restriction.

At the moment, bufferization can create memref.reshape operations with non-identity layouts.
The proposed solution in this change is to allocate a new buffer for the memref.reshape when the bufferization process attempts to reuse a buffer with a non-identity layout.

[Lewuathe]

@sabauma Sorry, my comment was not clear enough.

According to the documentation, the memref.reshape expects both source and destination type to have identity layout map. I thought memref.reshape does not get non-identity map input technically, which means we might not need to take care of the case where memref.reshape has non-identity input in the bufferization pass too.

But as the test you have wrote indicated, we can pass the non-identity layout tensor to memref.reshape by using tensor.extract_slice. So my assumption might be wrong.

In that case, can we improve the documentation to say we can pass non-identity layout memref to memref.reshape op?

[sabauma]

@Lewuathe The change I made in the bufferization interface for tensor.reshape ensures that the lowering process always produces a memref.reshape with an identity affine map.

Without the change to the bufferization interface, the bufferization process will produce IR that fails verification. From my perspective, the issue I'm trying to address is that bufferization produces an invalid memref.reshape.

In that case, can we improve the documentation to say we can pass non-identity layout memref to memref.reshape op?

I don't think we want to do that. I suspect the lowerings depend on this identity map requirement. I think the fact that bufferization will try to create memref.reshape operations with non-identity layouts is a bug, as it is not conforming to the spec of memref.reshape.

It may be possible to enhance memref.reshape to handle non-identity layout maps, but that is probably a more invasive change.

[Lewuathe]

Thanks for the intention of the change. It has made things much clearer to me. This change just guards the case where the input does not have an identity layout.

It may be possible to enhance memref.reshape to handle non-identity layout maps, but that is probably a more invasive change.

Yes, I agree. Even if it's the bug of the upstream to generate the invalid memref type with non-identity map, it would be a big fix.

[sabauma]

Hi @Lewuathe, if there are no other comments, do you mind merging this PR. I do not have write access.

[Lewuathe]

@sabauma Okay, I'll merge this commit. Thanks for the contribution!