iree-org / iree-org/wave

Allow tkw.iterate to emit scf.for without reduction

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

Description

tkw.iterate should support emitting a plain scf.for loop without reduction semantics. This would enable tiling over a dimension, performing arbitrary computation and writing to memory without requiring iter_arg_sand scf.yield.

Currently from my observation tkw.iterate always assumes a reduction pattern and generates accumulator-style loops which limits its usage.
This limitation came up in the context of an RMSNorm kernel, where the first loop performs a reduction to compute the RMS value. A second loop then needs to iterate over the same input to apply elementwise division by that value. This second pass doesn't involve reduction, but tkw.iterate cannot express it currently.

Summary:
It would be useful if tkw.iterate could emit plain scf.for loops when no init_args are needed — similar to:

```
scf.for %i = ... {
// compute and store
}
```
Let me know if this is already supported and I just missed it. I didn’t see any examples of tkw.iterate emitting a plain loop without iter_args.

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.