iree-org / iree-org/wave

bounds not generated/exported for loop induction variable

Open
#841 1 comment 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

Repro: add the following to the mlir_converter.py lit test and observe a failure

```python
@run_test
def mlir_converter_emit_wave_dialect_loop_implicit_capture():
"""Test emit_wave_dialect with a kernel that has a loop using registers defined in the kernel (not function args) as implicit captures.

Function arguments are mapped separately; this exercises the water_emitter path for
placeholders in iterate subgraphs that capture values defined in the kernel body
(e.g. registers or results of ops), via implicit_captures (users/ftynse/implicit-capture).
"""
M = tkl.sym.M
N = tkl.sym.N
K = tkl.sym.K
BLOCK_M = tkl.sym.BLOCK_M
BLOCK_N = tkl.sym.BLOCK_N
BLOCK_K = tkl.sym.BLOCK_K
ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE

constraints_loop: list[tkw.Constraint] = [
tkw.WorkgroupConstraint(M, BLOCK_M, 0),
tkw.WorkgroupConstraint(N, BLOCK_N, 1),
tkw.TilingConstraint(K, BLOCK_K),
tkw.WaveConstraint(M, sympy.floor(BLOCK_M / 2)),
tkw.WaveConstraint(N, sympy.floor(BLOCK_N / 2)),
tkw.HardwareConstraint(
threads_per_wave=64,
mma_type=MMAType.F32_32x32x8_F16,
vector_shapes={M: BLOCK_M, N: BLOCK_N, K: BLOCK_K},
),
]

@tkw.wave(constraints_loop)
def kernel_loop_implicit_capture(
a: tkl.Memory[M, K, ADDRESS_SPACE, tkl.f16],
b: tkl.Memory[N, K, ADDRESS_SPACE, tkl.f16],
c: tkl.Memory[M, N, GLOBAL_ADDRESS_SPACE, tkl.f32],
):
# Registers defined in the kernel body (not function args)
acc = tkl.Register[M, N, tkl.f32](0.0)
bias = tkl.Register[M, N, tkl.f32](0.0) # kernel-defined, captured by loop
# Loop body implicitly captures bias (kernel-defined register); acc is iter arg
@tkw.iterate(K, init_args=[acc])
def k_loop(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
a_reg = tkw.read(a)
b_reg = tkw.read(b)
acc = tkw.mma(a_reg, b_reg, acc)
acc = acc + bias # implicit capture of kernel-defined register
return acc

tkw.write(k_loop, c)

subs_loop = {
ADDRESS_SPACE: GLOBAL_ADDRESS_SPACE,
BLOCK_M: 32,
BLOCK_N: 32,
BLOCK_K: 16,
M: 64,
N: 64,
K: 32,
}
options_loop = WaveCompileOptions(
subs=subs_loop,
compile_to_mlir=True,
location_capture_config=LocationCaptureConfig(level=LocationCaptureLevel.NONE),
enforce_locations=False,
)
options_loop = set_default_run_config(options_loop)

compiled = wave_compile(options_loop, kernel_loop_implicit_capture)
trace = compiled.compiled_graph

mlir_output, diagnostics, _ = emit_wave_dialect(
trace, kernel_loop_implicit_capture.constraints, options_loop
)

if diagnostics:
for d in diagnostics:
print(d, file=sys.stderr)
assert len(diagnostics) == 0, "emit_wave_dialect should succeed for loop with implicit captures"

print(mlir_output)
```

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.