[wave2water] block sum reduction fails to convert
- Dominant language
- Python
- Stars
- 59
- Forks
- 32
- PR merge metrics
- No merged PRs in 30d
Description
Complains about structural error with placeholders in the exporter code. Perhaps related to missing ops as reductions are expanded by python passes.
Low-priority for now.
```python
@run_test
def mlir_converter_sum():
"""Test MLIR converter with sum kernel."""
permute_constraints = [
tkw.WorkgroupConstraint(M, 128, 0),
tkw.WorkgroupConstraint(N, 128, 1),
tkw.WaveConstraint(M, 64),
tkw.WaveConstraint(N, 64),
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: 128,
BLOCK_N: 128,
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,
print_ir_after="all",
)
options = set_default_run_config(options)
compiled_kernel = wave_compile(options, sum)
# Get the trace from the compiled kernel
trace = compiled_kernel.compiled_graph
constraints = sum.constraints
# Use the mlir_converter to emit wave MLIR dialect
mlir_output, diagnostics, _ = emit_wave_dialect(trace, constraints, options)
if diagnostics:
print(format_diagnostics(diagnostics, use_color=False), file=sys.stderr)
assert (
len(diagnostics) == 0
), "dialect emission should create valid IR, therefore diagnostics should be empty"
# Print to stdout for FileCheck.
print(mlir_output)
sys.exit(0)
```
Contributor guide
Assessment
This issue has not been assessed yet.