[mlir][xegpu] Fix seg-fault caused by setting a null attribute

[chencha3]

A bug fix in Wg2SgElemwiseOpPattern, check the updated attribute is not null before assigning it.

[Garra1980]

Is "else" required here?

[chencha3]

I don't think so. if the newLayout is null, we don't need to keep it in the op

[Garra1980]

Yeah, you're right, sorry for the noise

[llvmbot]

@llvm/pr-subscribers-mlir-gpu

@llvm/pr-subscribers-mlir

Author: Chao Chen (chencha3)

Changes

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp (+5-3)
• (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir (+19)

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>

[adam-smnk]

Isn't there potentially the same issue with:

if (!isa<scf::IfOp, scf::ForOp, scf::WhileOp, scf::ConditionOp>(op))
          op->setAttr(name, layout.dropSgLayoutAndData());

at the end of XeGPUWgToSgDistributePass::runOnOperation?

[chencha3]

Isn't there potentially the same issue with:

if (!isa<scf::IfOp, scf::ForOp, scf::WhileOp, scf::ConditionOp>(op))
          op->setAttr(name, layout.dropSgLayoutAndData());

at the end of XeGPUWgToSgDistributePass::runOnOperation?

Thanks for the catch up, it is fix.