[Mlir-commits] [mlir] [mlir][xegpu] Fix crash in propagate-layout with missing layouts (PR #221644)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>

The xegpu-propagate-layout pass may create an xegpu.convert_layout operation even when the input IR does not contain any XeGPU operations. Since the pass did not declare XeGPU as a dependent dialect, the new operation could be created through UnregisteredOpModel. Its ConvertLayoutOp properties were then interpreted as Attribute properties, triggering a PropertyRef TypeID assertion.

Declare XeGPU as a dependent dialect so that xegpu.convert_layout is created using its registered operation model. Also stop layout conflict resolution when the producer has no layout, avoiding construction of a convert_layout operation with a missing required target_layout.

Add a regression test covering a vector reduction in input IR without pre-existing XeGPU operations or layout information.

Fix: https://github.com/llvm/llvm-project/issues/221142

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


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td (+1-1) 
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp (+2) 
- (modified) mlir/test/Dialect/XeGPU/xegpu-propagate-layout-invalid.mlir (+20-1) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index 36f0a131b345b..959c755edc45c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -37,7 +37,7 @@ def XeGPUPropagateLayout : Pass<"xegpu-propagate-layout"> {
       Propagate the `sg_layout` and `sg_data` fields of the layout attribute.
       Default values are selected to align with hardware.
   }];
-  let dependentDialects = ["memref::MemRefDialect", "vector::VectorDialect"];
+  let dependentDialects = ["memref::MemRefDialect", "vector::VectorDialect", "xegpu::XeGPUDialect"];
   let options = [Option<
     "printOnly", "print-analysis-only", "bool",
     /*default=*/"false",
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 2903e19335a17..f69d6742d30e7 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1637,6 +1637,8 @@ LogicalResult ResolveLayoutConflicts::run() {
 LogicalResult ResolveLayoutConflicts::assignResultLayout(OpResult &result) {
   Operation *producerOp = result.getDefiningOp();
   auto producerLayout = xegpu::getDistributeLayoutAttr(result);
+  if (!producerLayout)
+    return failure();
   // Insert a convert_layout op to assign the layout.
   builder.setInsertionPointAfterValue(result);
   auto convertOp = xegpu::ConvertLayoutOp::create(
diff --git a/mlir/test/Dialect/XeGPU/xegpu-propagate-layout-invalid.mlir b/mlir/test/Dialect/XeGPU/xegpu-propagate-layout-invalid.mlir
index ded8a627f9f6a..dad5586d4a73e 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-propagate-layout-invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-propagate-layout-invalid.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt --xegpu-propagate-layout %s | FileCheck %s
+// RUN: mlir-opt --xegpu-propagate-layout --split-input-file -verify-diagnostics %s | FileCheck %s
 
 // Regression test for https://github.com/llvm/llvm-project/issues/177846:
 // --xegpu-propagate-layout must not crash when the module contains an
@@ -11,3 +11,22 @@
 module {
   llvm.func @some_function()
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/221142:
+// --xegpu-propagate-layout must not crash on a module without any XeGPU ops.
+// assignResultLayout created a xegpu.convert_layout for the scalar result of
+// a vector.reduction even though no layout existed anywhere in the module;
+// with the XeGPU dialect never loaded (no XeGPU ops and no dependentDialects
+// entry), the op creation aborted with "Property type mismatch: TypeID does
+// not match requested type". The pass now bails out when the result carries
+// no layout and reports the missing layout via diagnostics instead.
+module {
+  func.func @reduction(%arg0: f32) -> () {
+    // expected-warning at +1 {{op has users but no layout assigned for its result}}
+    %cst_1 = arith.constant dense<1.000000e+00> : vector<16xf32>
+    %0 = vector.reduction <add>, %cst_1 : vector<16xf32> into f32
+    return
+  }
+}

``````````

</details>


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