iree-org / iree-org/wave

[wave2water] block sum reduction fails to convert

Open
#904 0 comments 0 reactions 1 assignee Claimed by @martin-luecke View on GitHub
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

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.