[Mlir-commits] [mlir] [MLIR][XeGPU] Unroll dpas with N in the outer loop (PR #235095)

[Md Abdullah Shahneous Bari]

https://github.com/mshahneo created https://github.com/llvm/llvm-project/pull/235095


When `UnrollDpasOp` splits a `dpas` into instruction-sized blocks, it walked the M blocks in the outer loop and the N blocks in the inner loop. So consecutive `dpas` ops used a different B block (the `dpas` src1 operand) each time.

This change swaps the two loops: N in the outer loop, M in the inner loop. Consecutive `dpas` ops then use the same B block. IGC groups consecutive `dpas` that share src1 into one dpas macro.

Each result is still stored at the same C block index, so the unpacked result does not change. Only the order of the `dpas` ops changes.

- New test case in `xegpu-blocking.mlir`: checks the new order and where each result is stored.
- `UnrollDpasMxOp` uses the same loop order. It is not changed here.

Tested with the IMEX XeGPU integration tests on Intel GPU: same pass set with and without this change.

🤖 Generated with [Claude Code](https://claude.com/claude-code)


>From 1a42fd54b984d44e9dccd48b42834c84bed3e27c Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Fri, 23 Oct 2026 16:02:22 +0000
Subject: [PATCH] [MLIR][XeGPU] Unroll dpas with N in the outer loop

When the unroll pattern splits a dpas into instruction-sized blocks, it now
walks the N blocks in the outer loop and the M blocks in the inner loop.
Consecutive dpas ops then use the same B block (the dpas src1 operand).
IGC groups consecutive dpas that share src1 into one dpas macro. Before
this change, consecutive ops used a new B block each time.

Each result still goes to the same C block, so the unpacked result does not
change. Only the order of the dpas ops changes.

A new test in xegpu-blocking.mlir checks the order and where each result
is stored.

Co-Authored-By: Claude Opus 5.5 <noreply at anthropic.com>
---
 .../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp  |  8 ++---
 mlir/test/Dialect/XeGPU/xegpu-blocking.mlir   | 36 +++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index aaee71bb861822..5991f1f15bd069 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -413,10 +413,10 @@ struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
     for (int64_t d = 0; d < batchRank; ++d)
       batchIters *= aShape[d] / batchDims[d];
 
-    SmallVector<Value> newOps;
+    SmallVector<Value> newOps(batchIters * mIters * nIters);
     for (int64_t batch = 0; batch < batchIters; ++batch) {
-      for (int64_t i = 0; i < mIters; ++i) {
-        for (int64_t j = 0; j < nIters; ++j) {
+      for (int64_t j = 0; j < nIters; ++j) {
+        for (int64_t i = 0; i < mIters; ++i) {
           Value tmpC;
           if (c)
             tmpC = cVals[batch * (mIters * nIters) + i * nIters + j];
@@ -435,7 +435,7 @@ struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
             xegpu::dropInstDataOnInherentAttrs(newDpasOp);
             tmpC = newDpasOp.getResult();
           }
-          newOps.push_back(tmpC);
+          newOps[batch * (mIters * nIters) + i * nIters + j] = tmpC;
         }
       }
     }
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index 70de8194ff1bc4..e2598ff9a58031 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -853,3 +853,39 @@ gpu.module @test_kernel {
     gpu.return
   }
 }
+
+// -----
+// The unrolled dpas ops walk N in the outer loop and M in the inner loop, so
+// consecutive dpas ops use the same B block. Each result still goes to its own
+// C block.
+#a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>
+#b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [16, 1]>
+#c = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>
+gpu.module @test_kernel {
+  // CHECK-LABEL: gpu.func @dpas_unroll_order
+  // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+  // CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+  // CHECK: %[[A0:.+]] = xegpu.load_nd %{{.+}}[%[[C0]], %[[C0]]] {{.*}} -> vector<8x16xf16>
+  // CHECK: %[[A1:.+]] = xegpu.load_nd %{{.+}}[%[[C8]], %[[C0]]] {{.*}} -> vector<8x16xf16>
+  // CHECK: %[[B0:.+]] = xegpu.load_nd %{{.+}}[%[[C0]], %[[C0]]] {{.*}} -> vector<16x16xf16>
+  // CHECK: %[[B1:.+]] = xegpu.load_nd %{{.+}}[%[[C0]], %[[C16]]] {{.*}} -> vector<16x16xf16>
+  // CHECK: %[[D00:.+]] = xegpu.dpas %[[A0]], %[[B0]]
+  // CHECK-NEXT: %[[D10:.+]] = xegpu.dpas %[[A1]], %[[B0]]
+  // CHECK-NEXT: %[[D01:.+]] = xegpu.dpas %[[A0]], %[[B1]]
+  // CHECK-NEXT: %[[D11:.+]] = xegpu.dpas %[[A1]], %[[B1]]
+  // CHECK: xegpu.store_nd %[[D00]], %{{.+}}[%[[C0]], %[[C0]]]
+  // CHECK: xegpu.store_nd %[[D01]], %{{.+}}[%[[C0]], %[[C16]]]
+  // CHECK: xegpu.store_nd %[[D10]], %{{.+}}[%[[C8]], %[[C0]]]
+  // CHECK: xegpu.store_nd %[[D11]], %{{.+}}[%[[C8]], %[[C16]]]
+  gpu.func @dpas_unroll_order(%A: memref<16x16xf16>, %B: memref<16x32xf16>, %C: memref<16x32xf32>) {
+    %a_tdesc = xegpu.create_nd_tdesc %A : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #a>
+    %b_tdesc = xegpu.create_nd_tdesc %B : memref<16x32xf16> -> !xegpu.tensor_desc<16x32xf16, #b>
+    %c_tdesc = xegpu.create_nd_tdesc %C : memref<16x32xf32> -> !xegpu.tensor_desc<16x32xf32, #c>
+    %a = xegpu.load_nd %a_tdesc[0, 0] <{layout = #a}>: !xegpu.tensor_desc<16x16xf16, #a> -> vector<16x16xf16>
+    %b = xegpu.load_nd %b_tdesc[0, 0] <{layout = #b}>: !xegpu.tensor_desc<16x32xf16, #b> -> vector<16x32xf16>
+    %c = xegpu.dpas %a, %b <{layout_a = #a, layout_b = #b, layout_cd = #c}>: vector<16x16xf16>, vector<16x32xf16> -> vector<16x32xf32>
+    xegpu.store_nd %c, %c_tdesc[0, 0] <{layout = #c}>: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #c>
+    gpu.return
+  }
+}