Use wave id intrinsic on gfx1250

[Hardcode84]

• gfx12+ supports getting wave id directly, without doing linearized_thread_id / wave_size
• Enable it using new special symbol WAVE_ID (and also add LANE_ID while we at it)
• 0 readfirstlanes in triple buffered gemm kernel with water backend! (IREE backend actually regresses as gpu.subgroup_is -> rocdl.wave_id is not properly plumbed there yet).
• Disable this path for older arch (will need to enable later after investigating the performance, also too much churn in lit tests)
• Refactor kernel.initialize_xyz functions

[tgymnich]

LGTM, please also add WAVE_ID and LANE_ID to water.
Is this for gfx1250 only? If not please update the PR title.

[Copilot]

Pull request overview

This PR enables using GPU intrinsics for wave/lane identification on gfx1250 (via new WAVE_ID/LANE_ID symbols) to avoid computing wave_id from linearized_thread_id, and refactors kernel constraint initialization into a single entry point.

Changes:

• Add WAVE_ID/LANE_ID global symbols and plumb them through codegen/emitter substitutions.
• Switch multiple passes to use hardware_constraint.wave_id / lane_id instead of recomputing from linearized_thread_id.
• Refactor constraint initialization call sites to use initialize_all_constraints(options) and add a hardware constraint init hook.

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/kernel/wave/wave.py	Adds initialize_all_constraints and target-gated enabling of use_wave_id.
wave_lang/kernel/wave/water.py	Passes chipset through convert-gpu-to-rocdl in the Water pipeline.
wave_lang/kernel/wave/utils/general_utils.py	Treats WAVE_ID/LANE_ID as thread-level symbols when stripping indexing.
wave_lang/kernel/wave/tensor_load_to_shared.py	Uses hardware_constraint.wave_id for wave-based redistribution.
wave_lang/kernel/wave/specialize.py	Uses hardware_constraint.wave_id for specialization conditions.
wave_lang/kernel/wave/schedule_reordering.py	Uses hardware_constraint.wave_id when inserting conditional barriers.
wave_lang/kernel/wave/minimize_global_loads.py	Clears WAVE_ID/LANE_ID when building minimized access patterns.
wave_lang/kernel/wave/fuse_tensor_loads.py	Imports WAVE_ID, adjusts fusion substitutions, and uses hardware_constraint.wave_id.
wave_lang/kernel/wave/decompose_scan_ops.py	Uses hardware_constraint.lane_id / wave_id during scan decomposition.
wave_lang/kernel/wave/decompose_reduce_ops.py	Uses lane_id in reduction decomposition/shuffle config derivation.
wave_lang/kernel/wave/constraints.py	Adds use_wave_id and makes wave_id/lane_id return intrinsics when enabled.
wave_lang/kernel/wave/compile.py	Switches compilation to initialize constraints via initialize_all_constraints(options).
wave_lang/kernel/wave/asm/wave_asm/test/e2e/waveasm_e2e.py	Updates e2e capture helpers to use initialize_all_constraints(options).
wave_lang/kernel/wave/asm/wave_asm/test/e2e/compare_backends.py	Updates backend comparison tests to use initialize_all_constraints(options).
wave_lang/kernel/wave/asm/scripts/compare_backends.py	Updates script path to use initialize_all_constraints(mlir_options).
wave_lang/kernel/lang/global_symbols.py	Introduces WAVE_ID and LANE_ID symbols.
wave_lang/kernel/compiler/wave_codegen/emitter.py	Emits gpu.subgroup_id/gpu.lane_id and binds WAVE_ID/LANE_ID in substitutions.
tests/kernel/wave_gemm_test.py	Updates expected resource counts and disables use_wave_id for specialize test.
tests/kernel/wave_gemm_mxfp_test.py	Adds use_water_backend parameterization and propagates it into options.
lit_tests/kernel/wave/specialize.py	Updates FileCheck expectations to use gpu.subgroup_id.
lit_tests/kernel/wave/gemm.py	Updates FileCheck expectations to use gpu.lane_id in indexing maps.

---

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