Page MenuHomePhabricator

[MLIR]Extend vector.gather to support n-D result
ClosedPublic

Authored by pzread on Aug 15 2022, 9:45 AM.

Details

Summary

Currently vector.gather only supports reading memory into a 1-D result vector. This patch extends it to support an n-D result vector with the indices, masks, and passthroughs in n-D vectors.

As we are trying to vectorize tensor.extract with vector.gather (https://github.com/iree-org/iree/issues/9198), it will need to gather the elements into an n-D vector. Having vector.gather with n-D results allows us to avoid flatten and reshape at the vectorization stage. The backends can then decide the optimal ways to lower the vector.gather op.

Note that this is different from n-D gathering, which is about reading n-D memory with the n-D indices. The indices here are still only 1-D offsets on the base.

Diff Detail

Event Timeline

pzread created this revision.Aug 15 2022, 9:45 AM
pzread updated this revision to Diff 452752.Aug 15 2022, 11:23 AM

Add tests.

pzread updated this revision to Diff 452755.Aug 15 2022, 11:36 AM

Add basic tests.

pzread retitled this revision from [MLIR]Support vector.gather with n-D result to [MLIR]Extend vector.gather to support n-D result.Aug 15 2022, 11:54 AM
pzread edited the summary of this revision. (Show Details)
pzread added a reviewer: dcaballe.
pzread published this revision for review.Aug 15 2022, 12:25 PM

Thanks, @pzread! This is a great step towards supporting more complex gather operations!
LGTM.

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
1789

I think we should remove this reference to LLVM since the operation is evolving in a different way?

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
95

Great refactoring, thanks!

109

if isMemRefTypeSupported is a pre-condition, we should add an assert.

dcaballe accepted this revision.Aug 15 2022, 10:14 PM
This revision is now accepted and ready to land.Aug 15 2022, 10:14 PM
pzread updated this revision to Diff 453185.Aug 16 2022, 7:44 PM

Add asserts.

pzread marked 2 inline comments as done.Aug 16 2022, 7:44 PM
pzread added a comment.EditedAug 23 2022, 9:14 AM

Hi @dcaballe, could you help me commit this change? Thanks!

This revision was automatically updated to reflect the committed changes.

There was a crashing test after this change: https://buildkite.com/mlir/mlir-core/builds/24941
I think this should reproduce:

cmake -G Ninja ${src_dir}/llvm -DLLVM_TARGETS_TO_BUILD="X86;AMDGPU;NVPTX" -DLLVM_ENABLE_PROJECTS=mlir -DLLVM_BUILD_EXAMPLES=ON -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=Off -DLLVM_LINK_LLVM_DYLIB=ON

Thanks for reverting it, Mehdi. I didn't anything failing locally.