Enable Water-based index inference as optional secondary

[ftynse]

This starts the transition to use Water-based index expression inference by enabling it to run in parallel with the original pyWave version of the same, optionally under a flag. Only the pyWave analyses or both may run. When both analyses run, the results of the two are compared with exceptions raised on mismatches. In any case, the results of the pyWave analyses are taken.

After stabilization, appropraite extension and fixes, the following steps will be taken:

1. Turn the flag running the Water-based analysis to be on by default.
2. Make the Water-based analysis primary and keep the pyWave analysis as optional backup for comparison.
3. Deprecate and remove the pyWave analysis.

Each of these steps will come in a separate commit after a certain period of time.

The plumbing requires conversion from MLIR attributes to Wave constructs, mostly sympy expressions that is added to the extent required here. The logic in emit_wave_dialect is extended to collect attribute values for all operations in the IR based on a unique ID derived from Python hash of the object attached as attribute.

Some portion of Wave code interacting with sympy is shifted to wave_lang/support to avoid the unconditional loading of IREE libraries due to wave_lang.kernel including wave_lang.kernel.wave which, in one or multiple places, transitively imports IREE, which in turn clashes with Water or any other MLIR-based project.

[Copilot]

Pull request overview

This PR enables Water-based index inference to run in parallel with the existing pyWave implementation as an optional secondary analysis. The approach allows comparison of results between the two implementations, with exceptions raised on mismatches, while continuing to use pyWave results as primary. This sets the foundation for eventually transitioning to Water-based analysis as the primary method.

Key changes include:

• Addition of MLIR-to-Wave conversion utilities for transforming Water dialect attributes to sympy expressions
• Infrastructure for assigning unique IDs to trace nodes and collecting Water-inferred attributes
• New compilation option check_water_analysis to enable parallel Water analysis
• Refactoring of code organization to avoid IREE import conflicts

Reviewed changes

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

Show a summary per file

File	Description
wave_lang/support/indexing.py	New support module containing index symbol/expression types to avoid IREE dependency
wave_lang/support/detect_water.py	New utility for detecting Water availability without importing IREE
wave_lang/kernel/wave/wave.py	Modified to conditionally run Water analysis alongside pyWave analysis
wave_lang/kernel/wave/water.py	Refactored to use centralized Water detection utilities
wave_lang/kernel/wave/templates/gemm.py	Updated constraints to use explicit sympy.floor for wave sizing
wave_lang/kernel/wave/mlir_converter/water_emitter.py	Enhanced to collect and compare Water-inferred attributes with unique node IDs
wave_lang/kernel/wave/mlir_converter/mlir_to_wave.py	New conversion module for transforming MLIR attributes to Wave constructs
wave_lang/kernel/wave/mlir_converter/mlir_converter.py	Updated to return inferred attributes and handle Water analysis pipeline
wave_lang/kernel/wave/compile_options.py	Added check_water_analysis flag
wave_lang/kernel/wave/analysis/index_sequence_analysis.py	Added Water-checked index setting function
wave_lang/kernel/lang/global_symbols.py	Refactored to use centralized symbol name constants
wave_lang/kernel/_support/indexing.py	Refactored to import from wave_lang.support.indexing
tests/*	Updated test imports and added comprehensive tests for MLIR conversion

---

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