[Mlir-commits] [mlir] [MLIR][XeGPU] Preserve anchor layouts in recoverTemporaryLayout (PR #182186)

[Nishant Patel]

https://github.com/nbpatel created https://github.com/llvm/llvm-project/pull/182186

None

>From 54790e9c306bc1d1725afc91c6cb032a365b8591 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Wed, 18 Feb 2026 23:30:31 +0000
Subject: [PATCH] Fix recoverTemporaryLayout

---
 .../lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 10 +++++++++-
 .../test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir | 12 ++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index eb7fab3610218..75a9149e93c1a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -93,7 +93,15 @@ bool xegpu::recoverTemporaryLayouts(Operation *rootOp) {
       // layout attributes to.
       if (isa<BlockArgument>(operand.get()))
         continue;
-      auto layout = xegpu::getDistributeLayoutAttr(operand.get());
+      // First, try to get the layout from the consumer (operand) side. This
+      // preserves anchor layouts on ops like load/store, which may carry
+      // user-specified inst_data/lane fields that the defining op (e.g.,
+      // multi_reduction) does not have.
+      auto layout = xegpu::getDistributeLayoutAttr(operand);
+      // Fall back to the producer (Value) side if the consumer side has no
+      // layout.
+      if (!layout)
+        layout = xegpu::getDistributeLayoutAttr(operand.get());
       if (!layout) {
         op->emitWarning("Could not find layout attribute for operand ")
             << operand.getOperandNumber() << " of operation " << op->getName();
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
index e2e94c5f0300f..199bbfd8a6d37 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
@@ -969,4 +969,16 @@ gpu.module @test_distribution {
     gpu.return
   }
 
+  // CHECK-LABEL: @preserve_anchor_layout
+  // CHECK: arith.constant dense<1.000000e+00> : vector<16x128xf32>
+  // CHECK: xegpu.store_nd %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] <{layout = #xegpu.layout<inst_data = [8, 16]>}>
+  gpu.func @preserve_anchor_layout(%dst: memref<256x128xf32>) {
+    %val = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 128]>} dense<1.0> : vector<256x128xf32>
+    %tdesc = xegpu.create_nd_tdesc %dst : memref<256x128xf32>
+      -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 128], inst_data = [8, 16]>>
+    xegpu.store_nd %val, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 128], inst_data = [8, 16]>}>
+      : vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 128], inst_data = [8, 16]>>
+    gpu.return
+  }
+
 }