[Mlir-commits] [mlir] [MLIR][XeGPU] Fix layout recovery for dead loop-carried values (PR #205884)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Nishant Patel (nbpatel)

<details>
<summary>Changes</summary>

A single backward scan resolves a loop-carried value only when it has a downstream consumer; with no consumer, its layout is known solely from block-arg uses seen later in the scan and gets missed. Repeat the scan to stamp the layout.

---
Full diff: https://github.com/llvm/llvm-project/pull/205884.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp (+29-6) 
- (modified) mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir (+45) 


``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 2a13997aa181f..674646eba8a6d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -319,12 +319,35 @@ bool xegpu::recoverTemporaryLayouts(Operation *rootOp) {
     });
   };
   removeTemporaryLayoutAttrs(rootOp);
-  rootOp->walk([&](func::FuncOp func) {
-    processFunc(func.getBody(), func.getSymName());
-  });
-  rootOp->walk([&](gpu::GPUFuncOp func) {
-    processFunc(func.getBody(), func.getName());
-  });
+
+  // Count layout attributes attached under `rootOp`. Recovery only adds
+  // attributes, so this count is used to detect a fixed point.
+  auto countTemporaryLayouts = [&]() {
+    unsigned count = 0;
+    rootOp->walk([&](Operation *nestOp) {
+      for (const NamedAttribute &namedAttr : nestOp->getDiscardableAttrs())
+        if (isa<xegpu::DistributeLayoutAttr>(namedAttr.getValue()))
+          ++count;
+    });
+    return count;
+  };
+
+  // A single backward sweep cannot resolve loop-carried values whose layout is
+  // only known from block-argument uses that are themselves annotated later in
+  // the sweep. Repeat until the layout count stops changing; recovery only adds
+  // attributes, so this terminates.
+  unsigned prevCount = 0;
+  unsigned curCount = countTemporaryLayouts();
+  do {
+    prevCount = curCount;
+    rootOp->walk([&](func::FuncOp func) {
+      processFunc(func.getBody(), func.getSymName());
+    });
+    rootOp->walk([&](gpu::GPUFuncOp func) {
+      processFunc(func.getBody(), func.getName());
+    });
+    curCount = countTemporaryLayouts();
+  } while (curCount != prevCount);
 
   return true;
 }
diff --git a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
index e2a4897fac519..9533037273d49 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
@@ -146,3 +146,48 @@ gpu.func @if_basic(
   gpu.return
 }
 }
+
+// -----
+// Test scf.for with a dead loop-carried result: %bias is used inside the loop
+// but the loop result %2#1 is unused. The dead chain (arith.negf -> scf.yield
+// -> %2#1) only gets its layout from the loop block-arg uses, so recovery must
+// iterate to a fixed point to fill it in.
+
+gpu.module @test_dead_carry {
+// CHECK-LABEL: gpu.func @for_dead_carry
+gpu.func @for_dead_carry(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<8x16xf32>) {
+  %c0 = arith.constant 0 : index
+  %c128 = arith.constant 128 : index
+  %c16 = arith.constant 16 : index
+  %0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
+  %1 = xegpu.create_nd_tdesc %arg1 : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16>
+  %cst = arith.constant dense<0.000000e+00> : vector<8x16xf32>
+  %cst_0 = arith.constant dense<1.000000e+00> : vector<8x16xf32>
+  // CHECK: scf.for
+  %2:2 = scf.for %arg3 = %c0 to %c128 step %c16
+      iter_args(%acc = %cst, %bias = %cst_0) -> (vector<8x16xf32>, vector<8x16xf32>) {
+    %4 = xegpu.load_nd %0[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+        : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+    %5 = xegpu.load_nd %1[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+        : !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
+    %6 = xegpu.dpas %4, %5, %acc
+        {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+         layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
+         layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+        : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+    %7 = arith.addf %6, %bias : vector<8x16xf32>
+    // CHECK: arith.negf
+    // CHECK-SAME: layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+    // CHECK-SAME: layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+    %8 = arith.negf %bias : vector<8x16xf32>
+    // CHECK: scf.yield
+    // CHECK-SAME: layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+    // CHECK-SAME: layout_operand_1 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+    scf.yield %7, %8 : vector<8x16xf32>, vector<8x16xf32>
+  }
+  %3 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
+  xegpu.store_nd %2#0, %3[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+      : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
+  gpu.return
+}
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/205884