[MLIR][XeGPU] Refactor XeGPU layout propagation: passing lane_layout/lane_data with inst_data (#203156)
**Motivation**
Enhance setup* rules in layout propagation to pass lane_layout, and
lane_data information during inst_data propagation, so that the
propagation can have lane level information when choosing an optimal
inst_data. This branch makes that relationship explicit and uniform
across all setup rules.
**Invariant**
All setup rules now produce layouts that satisfy:
Nd ops + dpas/dpas_mx: inst_data = k * (lane_layout * lane_data), k ≥ 1
Scatter/matrix ops + non-anchor ops: inst_data = lane_layout * lane_data
**Key changes in XeGPULayoutImpl**
- New per-op anchor setup rules: setupStoreNdAnchorLayout,
setupPrefetchNdAnchorLayout,
setupLoadNdAnchorLayout (replacing a generic block-IO path) — Nd ops
have rigid lane info
and fit inst_data to the lane factorization.
- New complete*LayoutFromInstData helpers: fill in lane info on
user-provided anchors that
specify only inst_data, by re-running the op's Lane-kind setup with
inst_data as the
destination shape. Covers scatter load/store, block store/load, DPAS,
and DPAS_MX.
- BlockIOInstructionInterface (new in uArchBase.h): shared abstraction
over the
load/store/prefetch 2D-block instructions; the three Xe2 instructions
now implement it.
- setupMultiReductionResultLayout reorganized so the InstData and Lane
branches share the
same lane-layout logic (computeReductionLaneLayoutAndData).
- createScaleLayout (dpas_mx): caps the scale lane_layout by inst_data
so the scale
operand's load_nd satisfies the multiple-of invariant.
- getValidLayouts → getSgLayoutCandidates: renamed and generalized from
2D to N-D
factorization.
- inferShapeCastSourceLayout preserves lane info through 1D↔ND
collapse-style casts.
- chunkSize → contiguousChunkSize parameter rename for the scatter
inst_data decision.
- ResolveLayoutConflicts: retargets an existing convert_layout instead
of chaining a second
one; UnrollConvertLayoutOp short-circuits when input == target.
- Sub-byte (4-bit) element support added to the block-load uArch table.
- Tests updated: propagate-layout-inst-data.mlir, propagate-layout.mlir,
resolve-layout-conflicts.mlir.
Assisted-by-claude
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>