iree icon indicating copy to clipboard operation
iree copied to clipboard

[GPU] Add vector distribution pattern for map_scatter

Open Max191 opened this issue 6 months ago • 1 comments

Adds a vector distribution pattern for iree_linalg_ext.map_scatter. The implementation is similar to that of vector.transfer_write without masking, and the main difference is in how the distributed offsets are handled by the distributed op.

This distribution will be used after the map_scatter op is vectorized, but before it is decomposed. This keeps the distribution pattern simple, because only the input vector needs to be distributed, and the index mapping to the distributed space is very simple.

Max191 avatar Jun 18 '25 14:06 Max191

Based on https://github.com/iree-org/iree/pull/21090 to allow vectorized map_scatter in the lit test.

EDIT: Rebased

Max191 avatar Jun 18 '25 14:06 Max191

I need to support masking, so going to convert to draft for now

EDIT: Actually, supporting masking seems like a separate concern. The vector distribution pattern should work, assuming the map_scatter has been masked properly, so I'll reopen the PR.

Max191 avatar Jul 16 '25 14:07 Max191