2 Commits

Author SHA1 Message Date
Charitha Saumya
eb33585b7d
[mlir][xegpu] Remove unused headers added in #177492 (#178719)
Remove unused headers added in #177492
2026-01-29 10:47:26 -08:00
Charitha Saumya
d9c65f94b1
[mlir][xegpu] Add XeGPUSgToWiDistributeExperimental pass. (#177492)
Currently XeGPU lowering pipeline uses `XeGPUSubgroupDistribute` pass to
subgroup to work item distribution of ops. This pass is well established
and relies on vector distribution's `WarpOp` based distribution
mechanism. However, recent experiments with larger kernels have shown
that this pass is very expensive in terms of compile time (see below).

This prompted us to create a new pass that does not rely on `WarpOp`
based distribution. This PR adds the initial infra to move away from the
old way and align Wg To WI distribution with Wg to Sg distribution. New
pass also uses context-aware type conversion based on XeGPU layouts to
distributed vector types from SG to WI.

This PR adds the following changes:
* SG to WI distribution pass based on context-aware type conversions
using `OpConversionPatterns`
* Test pass for testing individual patterns
(`TestXeGPUSgToWiDistributeExperimental`)
* `XeGPUSgToWiDistributeExperimentalPass` which will eventually replace
`XeGPUSubgroupDistribute`

Flash attention e2e compilations stats:
```
----Wall Time----  ----Name----
    0.0032 (  0.2%)  Parser
    0.0008 (  0.0%)  CSE
    0.0000 (  0.0%)    (A) DominanceInfo
    0.0002 (  0.0%)  GpuXeVMAttachTarget
    1.1427 ( 58.7%)  'gpu.module' Pipeline
    0.0019 (  0.1%)    XeGPUWgToSgDistribute
    0.0003 (  0.0%)    CSE
    0.0000 (  0.0%)      (A) DominanceInfo
    0.0002 (  0.0%)    LowerAffinePass
    0.0001 (  0.0%)    CSE
    0.0000 (  0.0%)      (A) DominanceInfo
    0.0008 (  0.0%)    XeGPUPropagateLayout
    0.0056 (  0.3%)    XeGPUBlocking
    0.0010 (  0.1%)    Canonicalizer
    0.0004 (  0.0%)    CSE
    0.0000 (  0.0%)      (A) DominanceInfo
    0.0015 (  0.1%)    XeGPUPropagateLayout
    0.0007 (  0.0%)    XeGPUOptimizeBlockLoads
    0.0010 (  0.0%)    Canonicalizer
    0.0004 (  0.0%)    CSE
    0.0000 (  0.0%)      (A) DominanceInfo
    0.0015 (  0.1%)    XeGPUPropagateLayout
    1.1274 ( 57.9%)    XeGPUSubgroupDistribute
    0.7959 ( 40.9%)  Output
    0.0022 (  0.1%)  Rest
    1.9461 (100.0%)  Total
```
2026-01-29 09:57:01 -08:00