[Mlir-commits] [mlir] 14f20ce - [MLIR][XeGPU] Remove layout attribute from scf ops after wg to sg (#180771)

[llvmlistbot at llvm.org]

Author: Nishant Patel
Date: 2026-02-12T07:26:18-08:00
New Revision: 14f20ce79508e7639bc1dbe890cd97be5a955015

URL: https://github.com/llvm/llvm-project/commit/14f20ce79508e7639bc1dbe890cd97be5a955015
DIFF: https://github.com/llvm/llvm-project/commit/14f20ce79508e7639bc1dbe890cd97be5a955015.diff

LOG: [MLIR][XeGPU] Remove layout attribute from scf ops after wg to sg (#180771)

Added: 
    

Modified: 
    mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
    mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp

Removed: 
    


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

diff  --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index cd6bc9ac4b8e0..566fb45f5bcb4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1847,4 +1847,18 @@ 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<RegionBranchOpInterface, RegionBranchTerminatorOpInterface>(op))
+      return;
+
+    SmallVector<StringAttr> attrsToRemove;
+    for (auto namedAttr : op->getDiscardableAttrs()) {
+      if (isa<xegpu::DistributeLayoutAttr>(namedAttr.getValue()))
+        attrsToRemove.push_back(namedAttr.getName());
+    }
+    for (auto attrName : attrsToRemove)
+      op->removeDiscardableAttr(attrName);
+  });
 }