iree-org / iree-org/wave

sum (and potentially other) block reduction throws random type errors

Open
#903 1 comment 0 reactions 0 assignees View on GitHub
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

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.