[all-commits] [llvm/llvm-project] d33b80: [MLIR][XeGPU] Add unrolling/blocking support for 3...

[Jianhui Li via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: d33b8093ce336b92d01f55033aed7021f5244eec
      https://github.com/llvm/llvm-project/commit/d33b8093ce336b92d01f55033aed7021f5244eec
  Author: Jianhui Li <jian.hui.li at intel.com>
  Date:   2026-06-08 (Mon, 08 Jun 2026)

  Changed paths:
    M mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
    M mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
    M mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
    M mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
    M mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
    M mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
    M mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
    M mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
    M mlir/test/Dialect/XeGPU/invalid.mlir
    M mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
    A mlir/test/Integration/Dialect/XeGPU/WG/simple_3d_gemm.mlir
    A mlir/test/Integration/Dialect/XeGPU/WG/simple_3d_mxfp_gemm.mlir
    M mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp

  Log Message:
  -----------
  [MLIR][XeGPU] Add unrolling/blocking support for 3D+ batched operations (#201725)

**Summary**
Add complete transform pass and lowering support for 3D+ batched
operations, building on the operation definition extensions for
load_nd/store_nd/prefetch_nd/dpas/dpas_mx. This enables end-to-end
compilation of batched GEMM workloads (e.g., [4, 64, 32] × [4, 32, 64] →
[4, 64, 64]).

**Key changes:**
Transform passes (XeGPUUnroll.cpp):

> Implement 3D batch unrolling using memref.subview to handle batch
offsets
> UnrollCreateNdOp: For rank > 2 with memref source, create per-batch
memref.subview slices and corresponding create_nd_tdesc ops
> UnrollLoadNdOp/StoreNdOp/PrefetchNdOp: Iterate over batch dimension
then inner 2D tile offsets, reusing batch tdescs across inner tiles
> UnrollDpasOp/UnrollDpasMxOp: Add outer batch loop before M/K/N
unrolling
> Remove returnSingleType parameter, compute correct batch counts
directly
>    Use targetShape (inst_data) for tdesc window size
> 

Blocking pass (XeGPUBlocking.cpp):

> Update getUnrolledTypes() to handle batch dimensions in TensorDescType
>    Extend inst_data extraction for >2D vectors with leading batch dims
> Fix getDpasInstDataVectors() to use correct M dimension for >2D types

Layout propagation (XeGPULayoutImpl.cpp, XeGPUPropagateLayout.cpp):

>    Extend createScaleLayout() to handle rank >= 2
>    Fix store_nd inst_data propagation for rank > 2
> Update lane layout logic to handle >2D vectors with leading unit dims

Distribution (XeGPUWgToSgDistribute.cpp):
>   Extend WgToSgDpasOp and WgToSgDpasMxOp to handle >2D result shapes
Lowering (XeGPUToXeVM.cpp):
>   Keep batch offset computation in unroll pass via memref.subview 
> Use last 2 dims for H/W offsets and shape in CreateNdDescToXeVMPattern
>   Handle tileRank >= 2 in load/store/prefetch lowering
Utilities (XeGPUUtils.cpp):
> Update getDistributedVectorType for batch-aware shapes by trimming
leading layout dims

Tests:
> Add xegpu-blocking.mlir tests for 3D batch dpas and dpas_mx unrolling
> Add simple_3d_gemm.mlir integration test (4-batch GEMM)
> Add simple_3d_mxfp_gemm.mlir integration test (XFAIL)
> Update invalid.mlir with 3D shape mismatch tests

Assisted-by-Claude

---------

Co-authored-by: Claude Sonnet 4.5 <noreply at anthropic.com>



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications