Skip to content

Commit fd24805

Browse files
authored
Reapply [mlir][xegpu] Add XeGPU subgroup map propagation analysis for XeGPU SIMT distribution. (#131380)
Originally introduced in #130240 and reverted in #131364 Reproduced the issue locally in Linux by doing a shared lib build. Fixes including adding the missing LINK_LIBS. **Original commit message:** This PR adds the SG map propagation step of the XeGPU SIMT distribution. SG map propagation is a sparse backward dataflow analysis that propagate the sg_map backward starting from the operands of certain operations (DPAS, store etc.). This is the first step of XeGPU subgroup distribution. This analysis result is used to attach layout information to each XeGPU SIMD subgroup op. The lowering patterns in XeGPUSubgroupDistribute will consume these layout info to distribute SIMD ops into SIMT ops that work on work-item level data fragments. Summary of Lowering XeGPU SIMD -> SIMT Subgroup map propagation (This PR) Attach sg_map to each op in move all ops inside gpu.warp_execute_on_lane0 region. Distribute each op using sg_map Additional legalization steps to align more with Xe HW.
1 parent 8bceb77 commit fd24805

File tree

4 files changed

+1242
-0
lines changed

4 files changed

+1242
-0
lines changed

mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,19 @@ def XeGPUFoldAliasOps : Pass<"xegpu-fold-alias-ops"> {
2323
];
2424
}
2525

26+
def XeGPUSubgroupDistribute : Pass<"xegpu-subgroup-distribute"> {
27+
let summary = "Distribute XeGPU ops to work items";
28+
let description = [{
29+
The pass distributes subgroup level (SIMD) XeGPU ops to work items.
30+
}];
31+
let dependentDialects = [
32+
"memref::MemRefDialect", "xegpu::XeGPUDialect", "vector::VectorDialect"
33+
];
34+
let options = [
35+
Option<"printOnly", "print-analysis-only", "bool",
36+
/*default=*/"false",
37+
"Print the result of the subgroup map propagation analysis and exit.">
38+
];
39+
}
40+
2641
#endif // MLIR_DIALECT_XEGPU_TRANSFORMS_PASSES_TD

mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
add_mlir_dialect_library(MLIRXeGPUTransforms
22
XeGPUFoldAliasOps.cpp
3+
XeGPUSubgroupDistribute.cpp
34

45
ADDITIONAL_HEADER_DIRS
56
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/XeGPU
@@ -14,4 +15,5 @@ add_mlir_dialect_library(MLIRXeGPUTransforms
1415
MLIRXeGPUDialect
1516
MLIRPass
1617
MLIRTransforms
18+
MLIRGPUDialect
1719
)

0 commit comments

Comments
 (0)