llvm-project
ebf518b3 - [MLIR][XeGPU] Refactor XeGPU layout propagation: passing lane_layout/lane_data with inst_data (#203156)

Commit
26 days ago
[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>
Author
Parents
Loading