`apply_expr` typing is broken for comparisons
- 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
Assessment
This issue has not been assessed yet.