[Mlir-commits] [mlir] [MLIR][XeGPU] Skip control flow ops in recoverTemporaryLayout (PR #180771)

[Nishant Patel]

https://github.com/nbpatel updated https://github.com/llvm/llvm-project/pull/180771

>From 6c9b83257d80fe08ba16ce4b2ec01b4141795223 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 10 Feb 2026 16:30:15 +0000
Subject: [PATCH 1/3] Skip control flow ops in recoverTemporaryLayout

---
 mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index a4e47fca96d34..291a39bca9938 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -91,7 +91,7 @@ bool xegpu::recoverTemporaryLayouts(Operation *rootOp) {
         continue;
       // Skip block arguments since they don't have defining ops to attach
       // layout attributes to
-      if (isa<BlockArgument>(operand.get()))
+      if (isa<BlockArgument>(operand.get()) || isa<LoopLikeOpInterface>(op))
         continue;
       auto layout = xegpu::getDistributeLayoutAttr(operand.get());
       if (!layout) {

>From 61dbf047d1b6b41852cc8a5bab0e382b09e6ae7f Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 10 Feb 2026 20:16:50 +0000
Subject: [PATCH 2/3] Move check

---
 mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 291a39bca9938..21f96672cd062 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -85,13 +85,16 @@ xegpu::dropInstDataOnAttrs(ArrayRef<NamedAttribute> attrs) {
 // a layout attribute.
 bool xegpu::recoverTemporaryLayouts(Operation *rootOp) {
   auto result = rootOp->walk([&](Operation *op) {
+    // Skip scf loop ops (e.g., scf.for, scf.while)
+    if (isa<LoopLikeOpInterface>(op))
+      return WalkResult::advance();
     for (OpOperand &operand : op->getOpOperands()) {
       // Layouts are needed for vector type only.
       if (!isa<VectorType>(operand.get().getType()))
         continue;
       // Skip block arguments since they don't have defining ops to attach
-      // layout attributes to
-      if (isa<BlockArgument>(operand.get()) || isa<LoopLikeOpInterface>(op))
+      // layout attributes to.
+      if (isa<BlockArgument>(operand.get()))
         continue;
       auto layout = xegpu::getDistributeLayoutAttr(operand.get());
       if (!layout) {

>From 357e6e6fb7b59266b5caf5e509f84f57db9da5e8 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Wed, 11 Feb 2026 16:40:08 +0000
Subject: [PATCH 3/3] Clean up attr

---
 mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp      | 3 ---
 .../lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp | 7 +++++++
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 21f96672cd062..a6db4563faaa2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -85,9 +85,6 @@ xegpu::dropInstDataOnAttrs(ArrayRef<NamedAttribute> attrs) {
 // a layout attribute.
 bool xegpu::recoverTemporaryLayouts(Operation *rootOp) {
   auto result = rootOp->walk([&](Operation *op) {
-    // Skip scf loop ops (e.g., scf.for, scf.while)
-    if (isa<LoopLikeOpInterface>(op))
-      return WalkResult::advance();
     for (OpOperand &operand : op->getOpOperands()) {
       // Layouts are needed for vector type only.
       if (!isa<VectorType>(operand.get().getType()))
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index cd6bc9ac4b8e0..3cd7b08c8c0e0 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1847,4 +1847,11 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
   if (failed(
           applyPartialConversion(getOperation(), target, std::move(patterns))))
     return signalPassFailure();
+
+  // Remove layout attributes from SCF ops
+  getOperation()->walk([](Operation *op) {
+    if (isa<scf::IfOp, scf::ForOp, scf::WhileOp, scf::ConditionOp>(op)) {
+      xegpu::removeLayoutAttrs(op);
+    }
+  });
 }