[water] Lower wave constraints to MLIR

[tgymnich]

Stacked PRs:

• ->[water] Lower wave constraints to MLIR #422
• [water] handle device constraint #426

---

[water] Lower wave constraints to MLIR

• emit the appropriate mlir dialect attributes for wave DSL constraints

Signed-off-by: Tim Gymnich tim@gymni.ch

[Copilot]

Pull Request Overview

This PR adds support for passing constraints through the MLIR converter pipeline. The constraints are now serialized alongside the trace and options, then deserialized in the water emitter where they are converted into MLIR attributes that are attached to the function operation.

• Constraint objects (HardwareConstraint, WorkgroupConstraint) are now passed through the emit_wave_dialect function
• Constraints are serialized/deserialized via dill along with traces and options
• Constraints are converted to MLIR attributes and attached to the function operation

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
wave_lang/kernel/wave/mlir_converter/water_emitter.py	Added constraint parameter throughout, validation logic, and MLIR attribute conversion for constraints
wave_lang/kernel/wave/mlir_converter/mlir_converter.py	Updated emit_wave_dialect signature to accept constraints and serialize them
lit_tests/kernel/wave/mlir_converter.py	Updated test calls to pass constraints from kernel objects

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull Request Overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull Request Overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull Request Overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Comments suppressed due to low confidence (1)

wave_lang/kernel/wave/mlir_converter/mlir_converter.py:37

• The docstring for emit_wave_dialect should be updated to document the new constraints parameter. The docstring currently only mentions "trace and options" but the function now also accepts constraints.

Suggested update:

"""Emit Wave MLIR by sending the pickled trace, constraints, and options to the emitter.

Args:
    trace: The captured trace to convert
    constraints: List of Wave constraints to apply
    options: Compile options (only the `subs` field is used during emission)
    test_diagnostic_emission: Whether to test diagnostic emission

Returns:
    A tuple of (MLIR module string, list of diagnostics)
"""

    """Emit Wave MLIR by sending the pickled trace and options to the emitter.

    The `subs` field of options is the only option used during emission."""

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.