iree-org / iree-org/wave

`apply_expr` typing is broken for comparisons

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

Repro:

```python
@wave.wave(apply_expr_constraints)
def apply_expr_kernel(
a: Memory[M, N, ADDRESS_SPACE_A, tkl.i32],
c: Memory[M, N, ADDRESS_SPACE_C, tkl.i32],
d: Memory[M, N, ADDRESS_SPACE_C, tkl.i32], # putting tkl.bool result in an error
):
a_reg = wave.read(a)
c_reg = wave.read(c)
result1 = wave.apply_expr([a_reg, c_reg], lambda x, y: sympy.Max(x, y))
result2 = wave.apply_expr([a_reg, c_reg], lambda x, y: sympy.Min(x, y))
result3 = wave.apply_expr([a_reg, c_reg], lambda x, y: x > y)
result4 = wave.apply_expr([a_reg, c_reg], lambda x, y: x < y)
result5 = wave.apply_expr([a_reg, c_reg], lambda x, y: sympy.Eq(x, y))
result6 = wave.apply_expr([a_reg, c_reg], lambda x, y: sympy.Ne(x, y))
result7 = wave.apply_expr([a_reg, c_reg], lambda x, y: x >= y)
result8 = wave.apply_expr([a_reg, c_reg], lambda x, y: x <= y)
result = (
result1 + result2
)
comparison_result = result3 + result4 + result5 + result6 + result7 + result8
# CANNOT do result + comparison_result due to some obscure assertion explosing DSL guts
wave.write(result, c)
# explicit wave.cast(comparison_result, tkl.i32) is needed
# BUT comparison_result.type says it's already i32!
wave.write(comparison_result, d)
```

this leads to highly frustrating user experience and is a prime example of why proper typing and verification are paramount.

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.