Language Feature Request : Hoisting out of wave.iterate
- Dominant language
- Python
- Stars
- 59
- Forks
- 32
- PR merge metrics
- No merged PRs in 30d
Description
Hey wave-devs,
I saw your attention kernels and they look like
```
def base_attention_core(q, k, v, c):
qkv_scaling = tkl.Register[B, H, M, K1, tkl.f16](dk_sqrt * log2e)
c_reg = tkl.Register[B, H, N, M, tkl.f32](0.0)
init_sum = tkl.Register[B, H, M, tkl.f32](0.0)
init_max = tkl.Register[B, H, M, tkl.f32](-1e6)
ZEROF = tkl.Register[M, K2, tkl.f32](0.0)
MIN_INF = tkl.Register[M, K2, tkl.f32](-1e6)
# This microkernel encodes the fact that if the reduction
# dimension were tiled, then we would need to materialize a loop.
@tkw.iterate(K2, init_args=[init_max, init_sum, c_reg])
def repeat(
partial_max: tkl.Register[B, H, M, tkl.f32],
partial_sum: tkl.Register[B, H, M, tkl.f32],
acc: tkl.Register[B, H, N, M, tkl.f32],
):
imm_reg = tkl.Register[B, H, K2, M, tkl.f32](0.0)
q_reg = tkw.read(q, mapping=q_mapping)
q_reg *= qkv_scaling
k_reg = tkw.read(k, mapping=k_mapping)
inner_acc = tkw.mma(k_reg, q_reg, imm_reg, mfma_variant[0])
x_j = tkw.permute(inner_acc, target_shape=[B, H, M, K2])
k2_index = tkw.self_index(K2, tkl.i32)
mask = tkw.apply_expr(k2_index, lambda x: x < K2)
mask = tkw.broadcast(mask, target_shape=[M, K2])
if is_causal:
```
Right now q_reg is hoisted out of the iterate by the compiler, but this does not work:
```
def base_attention_core(q, k, v, c):
qkv_scaling = tkl.Register[B, H, M, K1, tkl.f16](dk_sqrt * log2e)
c_reg = tkl.Register[B, H, N, M, tkl.f32](0.0)
init_sum = tkl.Register[B, H, M, tkl.f32](0.0)
init_max = tkl.Register[B, H, M, tkl.f32](-1e6)
ZEROF = tkl.Register[M, K2, tkl.f32](0.0)
MIN_INF = tkl.Register[M, K2, tkl.f32](-1e6)
q_reg = tkw.read(q, mapping=q_mapping)
q_reg *= qkv_scaling
# This microkernel encodes the fact that if the reduction
# dimension were tiled, then we would need to materialize a loop.
@tkw.iterate(K2, init_args=[init_max, init_sum, c_reg])
def repeat(
partial_max: tkl.Register[B, H, M, tkl.f32],
partial_sum: tkl.Register[B, H, M, tkl.f32],
acc: tkl.Register[B, H, N, M, tkl.f32],
):
imm_reg = tkl.Register[B, H, K2, M, tkl.f32](0.0)
k_reg = tkw.read(k, mapping=k_mapping)
inner_acc = tkw.mma(k_reg, q_reg, imm_reg, mfma_variant[0])
x_j = tkw.permute(inner_acc, target_shape=[B, H, M, K2])
k2_index = tkw.self_index(K2, tkl.i32)
mask = tkw.apply_expr(k2_index, lambda x: x < K2)
mask = tkw.broadcast(mask, target_shape=[M, K2])
if is_causal:
```
Would be great if someone could make this work.
Contributor guide
Assessment
This issue has not been assessed yet.