iree-org / iree-org/wave

Introduce GatherToLDS MLIR operation

Open
#1,290 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
59
Forks
32
PR merge metrics
No merged PRs in 30d

Description

The FX-level `GatherToLDS` op has no MLIR counterpart. This means any pipeline path that touches it in MLIR -- roundtrip, lowering, analysis -- is blocked .

What's needed:

- Define `wave.gather_to_lds` in `WaveOps.td` with appropriate interfaces. Possibly, it should not inherit `HasWaveIndexMapping` since it lacks a standard `index` attribute.
- Implement a C++ lowering pattern to `amdgpu.gather_to_lds` in `LowerReadWriteOps.cpp`
- Verify the op interacts correctly with existing MLIR passes and propagations.
- Add emission (`water_emitter.py`) and import (`fx_emitter.py`) support for MLIR roundtrip.

Currently, the MXFP4 roundtrip test works around the absence of this op by setting `use_global_to_shared=False` and `schedule=SchedulingType.NONE`, so does `water_e2e_test.py`. The manual schedule (`gemm_mxfp4_double_buffer.py`) depends on `GatherToLDS` nodes, so `SchedulingType.MANUAL` cannot be used without this.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.