sum (and potentially other) block reduction throws random type errors
- Dominant language
- Python
- Stars
- 59
- Forks
- 32
- PR merge metrics
- No merged PRs in 30d
Description
```
permute_constraints = [
tkw.WorkgroupConstraint(M, BLOCK_M, 0),
tkw.WorkgroupConstraint(N, BLOCK_N, 1),
tkw.WaveConstraint(M, sympy.floor(BLOCK_M / 2)),
tkw.WaveConstraint(N, sympy.floor(BLOCK_N / 2)),
tkw.HardwareConstraint(
threads_per_wave=64, vector_shapes={M: BLOCK_M, N: BLOCK_N}
),
]
@tkw.wave(permute_constraints)
def sum(
a: tkl.Memory[M, N, ADDRESS_SPACE_A, tkl.f16],
c: tkl.Memory[M, ADDRESS_SPACE_C, tkl.f16],
):
res = tkw.read(a)
init = tkw.read(c)
res = tkw.sum(res, init, dim=N, block=True)
tkw.write(res, c)
subs = {
ADDRESS_SPACE_A: GLOBAL_ADDRESS_SPACE,
ADDRESS_SPACE_C: GLOBAL_ADDRESS_SPACE,
BLOCK_M: 64,
BLOCK_N: 64,
M: 128,
N: 128,
}
options = WaveCompileOptions(
subs=subs,
compile_to_mlir=True, # Avoid IREE compilation
location_capture_config=LocationCaptureConfig(level=LocationCaptureLevel.NONE),
enforce_locations=False,
)
options = set_default_run_config(options)
compiled_kernel = wave_compile(options, sum)
```
results in
```
mlir_converter_sum
region_0 [root]:
graph():
%a : [num_users=1] = placeholder[target=a]
%c : [num_users=2] = placeholder[target=c]
%read_M:0_N:0 : [num_users=2] = [read](args = (%a, 1, None, (), None, MemoryAccessFlags.NONE, None, None, None), kwargs = {})
%read_1_M:0_N:0 : [num_users=2] = [read](args = (%c, 1, None, (), None, MemoryAccessFlags.NONE, None, None, None), kwargs = {})
%extract : [num_users=2] = [extract](args = (%read_M:0_N:0, [0]), kwargs = {})
%shuffle : [num_users=1] = [shuffle](args = (%extract, 1, 64, ShuffleMode.XOR), kwargs = {})
%add : [num_users=2] = [add](args = (%extract, %shuffle), kwargs = {})
%shuffle_1 : [num_users=1] = [shuffle](args = (%add, 2, 64, ShuffleMode.XOR), kwargs = {})
%add_1 : [num_users=2] = [add](args = (%add, %shuffle_1), kwargs = {})
%shuffle_2 : [num_users=1] = [shuffle](args = (%add_1, 4, 64, ShuffleMode.XOR), kwargs = {})
%add_2 : [num_users=2] = [add](args = (%add_1, %shuffle_2), kwargs = {})
%shuffle_3 : [num_users=1] = [shuffle](args = (%add_2, 8, 64, ShuffleMode.XOR), kwargs = {})
%add_3 : [num_users=2] = [add](args = (%add_2, %shuffle_3), kwargs = {})
%shuffle_4 : [num_users=1] = [shuffle](args = (%add_3, 16, 64, ShuffleMode.XOR), kwargs = {})
%add_4 : [num_users=2] = [add](args = (%add_3, %shuffle_4), kwargs = {})
%shuffle_5 : [num_users=1] = [shuffle](args = (%add_4, 32, 64, ShuffleMode.XOR), kwargs = {})
%add_5 : [num_users=1] = [add](args = (%add_4, %shuffle_5), kwargs = {})
%add_6 : [num_users=0] = [add](args = (%read_1_M:0_N:0, %add_5), kwargs = {})
%sum_1_M:0_N:0 : [num_users=1] = [sum](args = ([%read_M:0_N:0], %read_1_M:0_N:0, N, True), kwargs = {})
%write_M:0_N:0 : [num_users=0] = [write](args = (%sum_1_M:0_N:0, %c, 1, None, (), None, MemoryAccessFlags.NONE, None, None), kwargs = {})
return None
Custom format:
placeholder(_name=a, _type=Memory[M, N].of(f16)) type(Memory[M, N].of(f16))
placeholder(_name=c, _type=Memory[M].of(f16)) type(Memory[M].of(f16))
read(memory=a, elements_per_thread=1, mapping_dynamic_vals=(), flags=MemoryAccessFlags.NONE, index={M: $WG0*BLOCK_M + Mod($T0, 64) + floor($T0/64)*floor(BLOCK_M/2) : 1 : 1, N: $T1*floor(BLOCK_N/2) + $WG1*BLOCK_N + Mod($T0, 64) : 1 : 64}) type(Register[M, N].of(f16))
read(memory=c, elements_per_thread=1, mapping_dynamic_vals=(), flags=MemoryAccessFlags.NONE, index={M: $WG0*BLOCK_M + Mod($T0, 64) + floor($T0/64)*floor(BLOCK_M/2) : 1 : 1}) type(Register[M].of(f16))
extract(register_=read_M:0_N:0, offset=[0]) type(Register[M].of(f16))
shuffle(arg=extract, offset=1, width=64, mode=ShuffleMode.XOR) type(Register[M].of(f16))
add(lhs=extract, rhs=shuffle) type(Register[M].of(f16))
shuffle(arg=add, offset=2, width=64, mode=ShuffleMode.XOR) type(Register[M].of(f16))
add(lhs=add, rhs=shuffle_1) type(Register[M].of(f16))
shuffle(arg=add_1, offset=4, width=64, mode=ShuffleMode.XOR) type(Register[M].of(f16))
add(lhs=add_1, rhs=shuffle_2) type(Register[M].of(f16))
shuffle(arg=add_2, offset=8, width=64, mode=ShuffleMode.XOR) type(Register[M].of(f16))
add(lhs=add_2, rhs=shuffle_3) type(Register[M].of(f16))
shuffle(arg=add_3, offset=16, width=64, mode=ShuffleMode.XOR) type(Register[M].of(f16))
add(lhs=add_3, rhs=shuffle_4) type(Register[M].of(f16))
shuffle(arg=add_4, offset=32, width=64, mode=ShuffleMode.XOR) type(Register[M].of(f16))
add(lhs=add_4, rhs=shuffle_5) type(Register[M].of(f16))
add(lhs=read_1_M:0_N:0, rhs=add_5) type(Register[M].of(f16))
sum(arg=[read_M:0_N:0], init=read_1_M:0_N:0, dim=N, block=True, index={M: $WG0*BLOCK_M + Mod($T0, 64) + floor($T0/64)*floor(BLOCK_M/2) : 1 : 1}) type(Register[M].of(f16))
write(register_=sum_1_M:0_N:0, memory=c, elements_per_thread=1, mapping_dynamic_vals=(), flags=MemoryAccessFlags.NONE, index={M: $WG0*BLOCK_M + Mod($T0, 64) + floor($T0/64)*floor(BLOCK_M/2) : 1 : 1}) type(Memory[M].of(f16))
output(return_vals=(None,)) type(None)
Traceback (most recent call last):
File "/home/azinenko/git/wave/lit_tests/kernel/wave/mlir_converter.py", line 64, in
@run_test
^^^^^^^^
File "/home/azinenko/git/wave/wave_lang/kernel/wave/utils/general_utils.py", line 66, in run_test
func()
File "/home/azinenko/git/wave/lit_tests/kernel/wave/mlir_converter.py", line 105, in mlir_converter_sum
compiled_kernel = wave_compile(options, sum)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/azinenko/git/wave/wave_lang/kernel/wave/compile.py", line 1038, in wave_compile
) = _trace_launchable_and_get_kernel_signature(kernel, options, schedule)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/azinenko/git/wave/wave_lang/kernel/wave/compile.py", line 885, in _trace_launchable_and_get_kernel_signature
try_apply_pass(
File "/home/azinenko/git/wave/wave_lang/kernel/wave/utils/print_utils.py", line 784, in try_apply_pass
p()
File "/home/azinenko/git/wave/wave_lang/kernel/wave/decompose_reduce_ops.py", line 462, in decompose_reduce_ops
num_reduction_waves = int(
^^^^
File "/home/azinenko/git/wave/.venv/lib/python3.12/site-packages/sympy/core/expr.py", line 342, in __int__
raise TypeError("Cannot convert symbols to int")
TypeError: Cannot convert symbols to int
```
making the guts of the compiler explode into user's face, this is horrendous as user experience.
Contributor guide
Assessment
This issue has not been assessed yet.