Move scheduling passes to the end of the pipeline

[adedespirlet]

This PR moves the scheduling passes to the end of the compilation pipeline, after all optimization and analysis passes. This lets the schedule reason about the final optimized graph : merged reads, simplified indices, computed bounds ,... rather than the intermediate state that existed before.

Moving these passes earlier exposed several issues that required fixes:

1. Pipelined bug in the 4 wave assymetric schedule:
   After merge_contiguous_reads, ExtractSlice nodes are created between Reads and their consumers Bitcasts. The manual schedule's set_stage only assigns scheduling_parameters to nodes it's given explicitly and ExtractSlice was not part of them. This PR now makes sure the scheduling parameters are propagated properly with propagate_scheduling_parameters_to_bridge_nodes from source Reads to eahc extract_strided slice.
   With this fix, the liveness_anaysis in Constructpipelined loop, sees the stage gap betwen extacts strided slice and bitcast and thus creates rotating registers to carry value across pipeline iterations.

2. Placing partition_gather_like_ops pass before the manual shecudling pass created problems as well. It's a similar problem as in point 1) however now its about Reshape nodes. Reshape nodes were not assigned the correct stage as their source nodes. Fix: make sure the scheduling parameters are propagated from the source nodes to the reshape nodes , this is done again with propagate_scheduling_parameters_to_bridge_nodes

3. The other other failure happened in scaled_gemm. The auto-scheduler's create_scheduling_edges skips nodes in ignore_nodes. ExtractSlice was unknown to get_custom_operation_type (returned None), so it landed in ignore_nodes. This broke the dependency chain: edges from Read → ExtractSlice were created, but ExtractSlice → Bitcast edges were not (since ExtractSlice was skipped as a source). Without that edge, Bitcast lost its ordering constraint relative to Read and breaking the stage-transition validation. The fix made get_custom_operation_type resolve ExtractSlice recursively to its source Read's operation keeping it out of ignore_nodes and preserving the full dependency chain.

4. Add numeric probing fallback in check_is_mapping_contiguous for dynamic cases
   partition_gather_like_ops pass calls is_contiguous_vec on each Read/Write to decide whether to split it into scalar operations. This contiguity check relies on check_is_mapping_contiguous() which tries to prove symbolically that consecutive elements map to consecutive memory addresses. With static shapes, sympy can simplify and proof contiguity and so partition_gather_like_ops correctly leaves the reads intact. With dynamic shapes, the index expressions are too complex that sympy cannot simplify. The contiguity check would fail, causing partition_gather_like_ops to unnecessarily split reads into scalars.
   The fix : add numeric probing fallback in _check_contiguous_with_aligned_base when symbolic simplification fails. This uses the same principle as ProbeEvaluator in merge_contiguous_reads.

5. For tile size 32x64 in 4wave schedule, moving merge_contiguous_reads before manual scheduling causes it to merge 2 Ascale reads into a single wide read. When the schedule later calls partition_by_dim to split Ascale operations along M into 2 partitions, it fails because the merged read only has 1 unique M ID instead of the expected 2. The fix adds a conditional in the schedule that checks the number of unique M IDs before partitioning: if there aren't enough to split, it skips the partition and assigns all Ascale nodes to a single group. Other fix is in : wave_lang/kernel/ops/wave_schedule_ops.py : when merge_contiguous_reads merges reads before scheduling, filter_nodes can now legitimately return empty list

6. When flatten_read_indices runs before scheduling, GatherToLDS nodes now carry expression in the src_bounds key's. The unrolling pass was updating src_index and dst_index with the unroll offset but not
   src_bounds, so unrolled copies reused stale bounds from the original iteration. This caused the
   last G2S prefetch to use the wrong OOB guard, reading garbage into shared memory. Fix in update_index_for_unroll() to substitute in dict keys, not just values. And Unroll() to handle bounds expressions too

[harsh-nod]

Verdict: REQUEST CHANGES (confidence: 85/100)

The core architectural change — moving merge_contiguous_reads before manual scheduling — is well-designed with thoughtful cascading fixes across 10 files. The scheduling parameter propagation, ExtractSlice handling, and conditional a_scale partitioning are solid work.

However, there is one blocking issue:

• P0: 11 MI35X (CDNA1) e2e test failures producing incorrect GEMM results on 128x256x256 block shapes with dynamic dimensions. Main branch CI passes cleanly, confirming these are regressions introduced by this PR. The likely root cause is the interaction between the b_scale_shuffling_factor removal and how filter_nodes + numeric probing affect vmcnt counts on CDNA1.

Minor non-blocking items include adding a clarifying comment on the get_use() relaxation (P2), and auditing is_bscale_shuffled for dead code after the fix (P3).