[Mlir-commits] [mlir] [mlir][xegpu] Fix seg-fault caused by setting a null attribute (PR #146002)

[Chao Chen]

https://github.com/chencha3 updated https://github.com/llvm/llvm-project/pull/146002

>From 2e4bbeb3edccd35c1d79ff3d72f757a32814612e Mon Sep 17 00:00:00 2001
From: Chao Chen <chao.chen at intel.com>
Date: Fri, 27 Jun 2025 01:24:38 +0000
Subject: [PATCH 1/2] fix

---
 .../Transforms/XeGPUWgToSgDistribute.cpp      |  8 +++++---
 .../XeGPU/xegpu-wg-to-sg-elemwise.mlir        | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index e3563d10bc6f1..20932e05985ed 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -376,10 +376,12 @@ struct WgToSgElementwiseOp : public ConversionPattern {
       // Copy all attributes, but update "layout_result_0" to drop
       // sgLayout/sgData
       for (auto attr : op->getAttrs()) {
-        if (auto layout = dyn_cast<xegpu::LayoutAttr>(attr.getValue()))
-          state.addAttribute(attr.getName(), layout.dropSgLayoutAndData());
-        else
+        if (auto layout = dyn_cast<xegpu::LayoutAttr>(attr.getValue())) {
+          if (auto newLayout = layout.dropSgLayoutAndData())
+            state.addAttribute(attr.getName(), newLayout);
+        } else {
           state.addAttribute(attr.getName(), attr.getValue());
+        }
       }
       Operation *newOp = rewriter.create(state);
       newResults.push_back(newOp->getResult(0));
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir
index 64f01d61d6e80..09df1e4da43e2 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir
@@ -1,6 +1,25 @@
 // RUN: mlir-opt --xegpu-wg-to-sg-distribute -split-input-file %s | FileCheck %s
 
 gpu.module @test_elementwise_ops {
+
+  // CHECK-LABEL: unary_ops_sg_layout_only
+  gpu.func @unary_ops_sg_layout_only(%a: memref<24x32xf32>) {
+    %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<24x32xf32>
+      -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>>
+    %load_a = xegpu.load_nd %tdesc_a
+      : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>>
+      -> vector<24x32xf32>
+    // CHECK: math.exp {{.*}} : vector<12x8xf32>
+    %exp = math.exp %load_a
+      {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>}
+      : vector<24x32xf32>
+    // CHECK: arith.negf {{.*}} : vector<12x8xf32>
+    %negf = arith.negf %load_a
+      {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>}
+      : vector<24x32xf32>
+    gpu.return
+  }
+
   // CHECK-LABEL: unary_ops
   gpu.func @unary_ops(%a: memref<24x32xf32>) {
     %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<24x32xf32>

>From 190e8f91b0fe66263453fd440cb01dc4971684ea Mon Sep 17 00:00:00 2001
From: Chao Chen <chao.chen at intel.com>
Date: Tue, 1 Jul 2025 20:13:14 +0000
Subject: [PATCH 2/2] fix

---
 mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 20932e05985ed..be7b860dd1729 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -631,8 +631,10 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
       std::string name = xegpu::getLayoutName(result);
       if (auto layout = op->getAttrOfType<xegpu::LayoutAttr>(name)) {
         op->removeAttr(name);
-        if (!isa<scf::IfOp, scf::ForOp, scf::WhileOp, scf::ConditionOp>(op))
-          op->setAttr(name, layout.dropSgLayoutAndData());
+        if (!isa<scf::IfOp, scf::ForOp, scf::WhileOp, scf::ConditionOp>(op)) {
+          if (auto newLayout = layout.dropSgLayoutAndData())
+            op->setAttr(name, newLayout);
+        }
       }
     }
   });