iree-org / iree-org/wave

Model `start` and `condition` in MLIR `IterateOp`

Open
#870 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
59
Forks
32
PR merge metrics
No merged PRs in 30d

Description

The FX-side `Iterate` class has `start` and `condition` fields that allow specifying a custom loop start value and termination condition. These have no counterpart in the MLIR `IterateOp` definition, which currently only models the iteration axis (`iterator`), carried values (`iter_args`), and captured values (`captures`).

The MLIR op description says "bounds that are not yet specified" -- the iteration range is resolved later when the symbolic dimension is instantiated. However, some FX graphs may set explicit `start` or `condition` values, and these are silently lost in the FX-to-MLIR-to-FX roundtrip.

We should extend `IterateOp` with optional attributes for `start` (an expression for the loop lower bound) and `condition` (an expression for the termination predicate), and update both `water_emitter.py` and `fx_emitter.py` to serialize/deserialize them.

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.