Introduce canonical Wave region formats and annotated pass adapters (iree-org#1120)

## Summary

This is a step towards a canonical form for Wave FX graphs. This change
only tackles region structure for now, not all graph structure more
broadly.

The canonical region form chosen here is an isolated region
representation: nested regions do not directly reference outer values,
and all captures are represented explicitly as lifted `Placeholder`s at
the start of the region signature. This gives us one stable structural
form to target for verification, printing, and roundtrips.

We are introducing this gradually. Not all existing passes operate on
that canonical region form yet, and this change does not require
rewriting them all at once. Instead, passes can now declare which region
form they expect, and the pass boundary will temporarily adapt the graph
into that form before the pass runs and return to canonical form
afterwards in the normal pipeline.

This change introduces four region forms:
- `ISOLATED`: the canonical form, with explicit lifted placeholder
captures
- `LEGACY_PLACEHOLDERS`: the older placeholder-based capture form
expected by several existing passes
- `DIRECT_OUTER_REF`: a legacy form where nested regions directly
reference outer values
- `SCHEDULE_SIGNATURE_PLACEHOLDERS`: a hybrid form that keeps
placeholders only for schedule-signature sources

## What this PR does
- introduce canonical region-capture handling and verification for
nested Wave regions
- let passes declare their required region form directly via annotation
through a decorator
- adapt the pipeline to canonicalize at pass boundaries while still
supporting legacy region forms where needed
- update lit tests to request raw post-pass output explicitly at the
call site


## Additional Change
`minimize_shared_allocs` now skips dead shared allocs instead of
assuming every alloc still has a live first/last use. This became
necessary once the new canonical-region migration flow exposed
intermediate graphs where an alloc may survive temporarily after its
uses are gone, making an old pass assumption explicit.

---------

Signed-off-by: Martin Lücke <martin.luecke@amd.com>

Author: martin-luecke

docs/wave/canonical_ir_format.md

@@ -0,0 +1,173 @@
+# Canonical Region IR Format
+
+This document describes the canonical region structure used by the Wave FX pipeline.
+
+It is intentionally limited in scope: it only covers nested region interfaces
+and capture structure. It does not attempt to define a full canonical form for
+all FX graph details such as node ordering, `vector_shapes`, or write
+dependencies.
+
+## Goal
+
+Wave is moving toward a single canonical region form so that:
+
+- nested regions have one stable structural representation in Python
+- FX <-> MLIR roundtrips preserve region structure in a form that later Python
+  passes can continue to process
+- for migration purposes, passes may declare which temporary non-canonical region view they need
+  without forcing the whole pipeline to stay in that view
+
+The canonical form is represented by `RegionFormat.ISOLATED` in
+`wave_lang/kernel/wave/region_canonicalization.py`.
+
+## Terms
+
+- Outer source: a node defined outside a nested region but used by that region
+- Local capture: the region-local representative of an outer source
+- Capture signature: the ordered list stored on the parent `NestedRegionOp` in
+  `implicit_captures`
+- Direct outer reference: a region node operand that points straight to a node
+  in the outer graph
+
+## Canonical Form: `ISOLATED`
+
+`ISOLATED` is the canonical/default region form.
+
+Structural invariants:
+
+- nested region bodies do not directly reference outer graph nodes
+- every captured outer value used inside the region is represented by a region-local
+  `Placeholder`
+- these capture placeholders form the leading non-`IterArg` input prefix of the
+  subgraph
+- the parent `NestedRegionOp.implicit_captures` list is the authoritative
+  ordered capture signature: it defines *which* outer values are captured and in
+  *what order*
+- each local capture placeholder carries a `meta["lifted"]` link to its outer
+  source. This per-placeholder metadata is derived from `implicit_captures` and
+  must stay consistent with it (the verifier checks this)
+
+In other words, the region interface is explicit and isolated from above.
+
+Schematic shape:
+
+```text
+root graph:
+  %outer_a
+  %outer_b
+  %region = iterate(..., implicit_captures=[%outer_a, %outer_b], ...)
+
+region subgraph:
+  %iter_arg0 = placeholder(iter arg)
+  %outer_a = placeholder(lifted from outer)
+  %outer_b = placeholder(lifted from outer)
+  ...
+```
+
+## Temporary Non-Canonical Forms
+
+Not all existing passes operate on `ISOLATED` yet. To migrate incrementally,
+passes may request one of several temporary region views. The pass boundary
+adapts into that form before the pass runs and, in the normal pipeline,
+canonicalizes back to `ISOLATED` afterwards.
+
+### `LEGACY_PLACEHOLDERS`
+
+This is the older placeholder-based capture form still expected by some
+pre-existing passes.
+
+Structural properties:
+
+- captured outer values are represented by region-local placeholders
+- the mapping from a local placeholder back to its outer source may still be
+  recovered with ad-hoc conventions that pre-existing passes relied on: name
+  matching and positional fallback within the capture prefix (codified in
+  `_try_resolve_legacy_capture_source` in `region_canonicalization.py`)
+- unlike `ISOLATED`, this form does not require `implicit_captures` plus
+  `meta["lifted"]` to be the sole authoritative description of the capture
+  interface
+
+This is a weaker contract than `ISOLATED`. A pass marked
+`LEGACY_PLACEHOLDERS` may still reason about captures through placeholder
+layout or legacy lookup behavior instead of relying only on the canonical
+capture interface. An already-canonical region may also satisfy this weaker
+contract, so the adapter can be a no-op on some graphs.
+
+This mode exists to support passes that still expect legacy placeholder
+structure while the pipeline as a whole moves toward explicit canonical
+captures.
+
+### `DIRECT_OUTER_REF`
+
+This is a legacy form where region bodies directly reference outer graph nodes.
+
+Structural properties:
+
+- operands inside the region may point directly to outer nodes
+- capture placeholders are removed or bypassed where possible
+- the parent capture signature may still track those outer values, but the body
+  itself is not isolated from above
+
+This form is convenient for passes that want to inspect or mutate the original
+outer values directly, especially around captured memory operands.
+
+### `SCHEDULE_SIGNATURE_PLACEHOLDERS`
+
+This is a hybrid legacy form used by scheduling-related passes.
+
+Structural properties (schedule-signature sources are the outer values that
+define the region boundary from the scheduler's point of view, namely
+outer-graph `Placeholder`s, i.e. kernel arguments, and `NewRegister`s):
+
+- placeholders are kept only for those schedule-signature sources
+- non-signature captures are rewritten back to direct outer references
+- the region mixes explicit placeholders for signature-defining values with
+  direct outer references for everything else
+
+In practice, the schedule-signature sources are the outer values that define
+the region boundary from the scheduler's point of view, namely values such as
+root placeholders and `NewRegister`s.
+
+## Why There Are Multiple Forms
+
+The long-term goal is for passes to converge on `ISOLATED`.
+
+The intermediate forms exist because rewriting every pass at once would be too
+large and too risky. Instead:
+
+1. the pipeline keeps one canonical form
+2. each pass declares the temporary form it currently expects
+3. pass-boundary adapters convert into that form just for the duration of the
+   pass
+4. the normal pipeline returns to canonical form afterwards
+
+This makes the migration incremental while still establishing one structural
+source of truth.
+
+## Pass Contract
+
+Passes declare their required region form with `@requires_region_format(...)`.
+
+The important contract is:
+
+- if a pass does not declare a region form, it is assumed to operate on the
+  canonical `ISOLATED` form
+- in the normal pipeline, pass outputs are canonicalized back to `ISOLATED`
+- white-box tests that want to inspect a temporary legacy form must request
+  `canonicalize_output=False` explicitly at the call site
+
+This keeps the default pipeline principled while still allowing tests to inspect
+the raw intermediate structure when needed.
+
+## Non-Goals
+
+This document does not define:
+
+- a canonical ordering for all FX nodes
+- a canonical form for `GetResult` materialization beyond what is required by
+  region structure
+- a canonical form for `vector_shapes`, write dependencies, or other downstream
+  analysis state
+
+Those are separate concerns that may later build on top of this region-level
+structural baseline.

lit_tests/kernel/wave/barrier_strategies.py

@@ -179,7 +179,7 @@ def test_gemm():
         set_post_expansion_indices(trace, constraints)
         tweak_index(graph)
         hoist_loop_invariant_ops(trace, constraints)
-        add_shared_memory_barriers(trace)
+        add_shared_memory_barriers(trace, canonicalize_output=False)
         print_trace(trace, False)
 
     # CHECK-LABEL: test_gemm
@@ -231,8 +231,14 @@ def test_split_barriers():
         tweak_index(graph)
         hoist_loop_invariant_ops(trace, constraints)
         schedule_graph(trace, constraints, True, enable_scheduling)
-        schedule_reordering(trace, constraints, enable_scheduling, use_global_to_shared)
-        add_shared_memory_barriers(trace, target="gfx1201")
+        schedule_reordering(
+            trace,
+            constraints,
+            enable_scheduling,
+            use_global_to_shared,
+            canonicalize_output=False,
+        )
+        add_shared_memory_barriers(trace, target="gfx1201", canonicalize_output=False)
         print_trace(trace, False)
 
     # CHECK-LABEL: test_split_barriers
@@ -363,7 +369,7 @@ def test_existing_barrier_not_duplicated():
 
         # Now run barrier placement - it should detect the existing barrier
         # and NOT insert duplicates
-        add_shared_memory_barriers(trace)
+        add_shared_memory_barriers(trace, canonicalize_output=False)
 
         # Count barriers after
         barriers_after = count_barriers(graph)
@@ -570,7 +576,7 @@ def test_memory_counter_wait_barrier_prevents_redundant_barrier():
 
         # Now run barrier placement - it should detect the existing MemoryCounterWaitBarrier
         # and NOT insert an additional SharedMemoryBarrier (since synchronization is already provided)
-        add_shared_memory_barriers(trace)
+        add_shared_memory_barriers(trace, canonicalize_output=False)
 
         # Count barriers after
         mcw_barriers_after = count_memory_counter_wait_barriers(graph)

lit_tests/kernel/wave/barriers.py

@@ -114,7 +114,7 @@ def test_read_write_equal_sizes():
         expand_graph(trace, constraints)
         set_post_expansion_indices(trace, constraints)
         tweak_index(graph)
-        add_shared_memory_barriers(trace)
+        add_shared_memory_barriers(trace, canonicalize_output=False)
         print_trace(trace, False)
     # CHECK: %a
     # CHECK-NEXT: %c
@@ -205,7 +205,7 @@ def test_gemm():
         set_post_expansion_indices(trace, constraints)
         tweak_index(graph)
         hoist_loop_invariant_ops(trace, constraints)
-        add_shared_memory_barriers(trace)
+        add_shared_memory_barriers(trace, canonicalize_output=False)
         print_trace(trace, False)
     # Root graph:
     # CHECK: %a
@@ -330,8 +330,14 @@ def test_split_barriers():
         tweak_index(graph)
         hoist_loop_invariant_ops(trace, constraints)
         schedule_graph(trace, constraints, True, enable_scheduling)
-        schedule_reordering(trace, constraints, enable_scheduling, use_global_to_shared)
-        add_shared_memory_barriers(trace, target="gfx1201")
+        schedule_reordering(
+            trace,
+            constraints,
+            enable_scheduling,
+            use_global_to_shared,
+            canonicalize_output=False,
+        )
+        add_shared_memory_barriers(trace, target="gfx1201", canonicalize_output=False)
         print_trace(trace, False)
 
     # Note: In pipelined loops, signal/wait pairs may have operations between them

lit_tests/kernel/wave/codegen.py

@@ -900,6 +900,7 @@ def test(
             ADDRESS_SPACE: GLOBAL_ADDRESS_SPACE,
         },
         canonicalize=True,
+        compile_to_mlir=True,
     )
     options = set_default_compile_config(options)
     test = wave_compile(options, test)