[Mlir-commits] [mlir] [mlir] Implement indexed access op interfaces for memref, vector, gpu, nvgpu (PR #177014)

[Krzysztof Drewniak]

krzysz00 wrote:

Re the non-unit dimensions:

When I was designing this interface in the previous PR, the point of getAccesedShape() was partly to have something like what https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp would need, and also to represent "these are the dimensions of the memref being accessed and how many values are being accessed along those dimensions". That information is important because it tells you which dimensions you're allowed to modify (split up, merge, change their strides, etc.) - and, when such modificationsw are being performed, "accesses" of length 1 **don't matter**.

For example, the accessed shape of a scalar `memref.load` is `{}` - you con fold in any aliasing operation you'd like, so long as the resulting memref/index collection points to the same element. However, a `vector.load` of a `vector<NxT>` (for N > 1) accesses `{N}` elements, so if you were going to fold in a `memref.subview` that had a final stride of `2`, you'd violate the semantics of the operation. (Similarly, with the gather semantics you've proposed, vector.gather accesses `{kDynamic}` elements, and loading a `vector<MxNxT>` accesses `{M, N}` elements when M > 1)

Now, as to why all this shape-casting, it's clearly true that a `vector.load` of a `vector<1x1x...T>` and one of a `vector<...xT>` are equivalent up to shape casts - for example, `vector<1xT>` and `vector<T>` are both effectively scalar loads. So, the interface implementation, in order to enable stuff like strided subview folding on the last dimension - where the current code isn't being precise here - we'll want to strip out those leading dimensions from the accessed shape, which means we'll need to introduce shape casts to keep the types behaving once the replacement is performed.

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