iree-org / iree-org/wave

mark_hardware_transpose_candidates breaks elements-per-thread invariant

Open
#982 2 comments 0 reactions 1 assignee Claimed by @harsh-nod View on GitHub
Dominant language
Python
Stars
59
Forks
32
PR merge metrics
No merged PRs in 30d

Description

After running `mark_elements_per_thread` on attention with the following config:

```
attention, hyperparams, _ = get_vanilla_attention_kernel(
AttentionShape(
num_query_heads=8,
num_kv_heads=2,
query_seq_len=256,
head_size_kv=64,
head_size=64,
kv_seq_len=256,
),
(MMAType.F32_16x16x16_F16, MMAType.F32_16x16x16_F16),
False,
)
```

reads from shared memory have `elements_per_thread = 4` and `index={B: $WG2*BLOCK_B : 1 : 1, N: $WG1*BLOCK_N + 4*(Mod($T0, 4)) : 1 : 1, K2: $ARGK2*BLOCK_K2 + floor((Mod($T0, 64))/4) : 8 : 1})`, the latter implying `elements_per_thread = 8`. Note that prior to that pass, the index uses `size=4` that _is_ consistent. It is unclear if downstream passes in the pipeline take one or the other value as source-of-truth.

This yet again highlights the blatant lack of pass-level testing and verifiers and over-reliance on e2e testing that happen to work in narrow cases.

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.