[DispatchCreation] Fold unit dims with encoding

[sommerlukas]

Extend the FoldUnitExtentDims pass to handle tensors with encoding attached.

Encodings can opt into folding of unit extent dimensions by adding a new interface that is used to query whether a tensor with encoding can be collapsed. If any of the tensor operands or results has an encoding that doesn't implement the interface, no folding happens to avoid issues with 0-entries in the indexing maps and passes that don't handle that case well.

Signed-off-by: Lukas Sommer [email protected]

[sommerlukas]

PR for the necessary changes to upstream MLIR: https://github.com/llvm/llvm-project/pull/170104

[sommerlukas]

After discussion with @hanhanW and @MaheshRavishankar, I took a different approach to exposing more control upstream, see the new PR https://github.com/llvm/llvm-project/pull/171796.

[hanhanW]

I haven't reviewed the details yet, because I need to review the upstream one first.

High-level comment: it is better to split it into two PRs:

• One is for LLVM integrate, which always returns false if encoding is present. I.e., it matches the current behavior on the main branch.
• The other is for adding the interface method and relevant tests, which is an improvement to IREE.

I'm treating this PR more like a demo.

@sommerlukas can you prepare a NFC PR for integrate before you land the upstream change? https://github.com/llvm/llvm-project/pull/171796

Thanks!