iree-org / iree-org/wave

tkw.broadcast after tkw.iterate loop does not inherit correct IndexSequence / vector shape

Open
#172 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

I'm trying to tile my RMSNorm kernel using tkw.iterate, and I hit an issue when applying tkw.broadcast after the loop.
In particular, the broadcasted tensor does not inherit the correct IndexSequence which causes a shape mismatch.
This leads to a failure during a later division where the LHS is a vector<4xf16> and the RHS (the broadcast result) is incorrectly a vector<1xf16>.
```

@tkw.iterate(N, init_args=[init_red])
def repeat(
partial_red: tkl.Register[M, tkl.f16],
) -> tkl.Register[M, tkl.f16]:
lhs = tkw.read(a, elements_per_thread=ELEMS_PER_THREAD)
lhs_sq = lhs * lhs
partial_red = tkw.sum(lhs_sq, partial_red, dim=N)
return partial_red

result = repeat / length_embedding
rms = tkw.sqrt(result)
rms_broad = tkw.broadcast(rms, [M, N])

# This is where it fails: mismatch between LHS and RHS vector shape
lhs2 = tkw.read(a, elements_per_thread=ELEMS_PER_THREAD)
a_scaled = lhs2 / rms_broad
```
Interestingly, the same broadcast works fine in my non-tiled RMSNorm kernel (which doesn’t use tkw.iterate). So I suspect something about the use of tkw.iterate causes the IndexSequence state to be lost .
You can see both the tiled and non-tiled versions in this commit:
[60a868e](https://github.com/iree-org/wave/commit/60a868eb8c67cfd400ae946a65ac35be9cacb632#diff-b4e8a2e4c754ab728143f3dc72c4685768d34dae00506763c51c00b9bd63fc64)

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.