[mlir][xegpu] Fix crash in XeGPUPropagateLayout when module has llvm.func

[joker-eph]

updateFunctionOpInterface() called funcOp.setType(FunctionType::get(...)) on every FunctionOpInterface operation, including llvm.func. However, llvm.func stores its type as LLVMFunctionType, not the standard FunctionType. Calling setType(FunctionType{}) on it corrupts the function_type attribute, and the next call to getFunctionType() (which tries to cast the stored attribute) aborts.

Fix by skipping functions whose type is not a standard FunctionType. XeGPU layout propagation only applies to functions using MLIR's FunctionType; other function types (like LLVMFunctionType) are not expected to carry XeGPU layouts.

Fixes #177846
Fixes #177777
Fixes #181970

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-gpu

Author: Mehdi Amini (joker-eph)

Changes

updateFunctionOpInterface() called funcOp.setType(FunctionType::get(...)) on every FunctionOpInterface operation, including llvm.func. However, llvm.func stores its type as LLVMFunctionType, not the standard FunctionType. Calling setType(FunctionType{}) on it corrupts the function_type attribute, and the next call to getFunctionType() (which tries to cast<LLVMFunctionType> the stored attribute) aborts.

Fix by skipping functions whose type is not a standard FunctionType. XeGPU layout propagation only applies to functions using MLIR's FunctionType; other function types (like LLVMFunctionType) are not expected to carry XeGPU layouts.

Fixes #177846

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp (+6)
• (added) mlir/test/Dialect/XeGPU/xegpu-propagate-layout-invalid.mlir (+13)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 10cd65b080405..313c60a9028db 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1500,6 +1500,12 @@ updateControlFlowOps(mlir::OpBuilder &builder,
 static LogicalResult updateFunctionOpInterface(mlir::OpBuilder &builder,
                                                mlir::FunctionOpInterface funcOp,
                                                GetLayoutFnTy getLayoutOfValue) {
+  // Only process functions whose type is a standard MLIR FunctionType.
+  // Functions using a different type representation (e.g. llvm.func with
+  // LLVMFunctionType) are not targets for XeGPU layout propagation, and
+  // calling setType(FunctionType{}) on them would corrupt their type.
+  if (!isa<FunctionType>(funcOp.getFunctionType()))
+    return success();
   SmallVector<Type> newArgTypes;
   // Update the function arguments.
   for (BlockArgument arg : funcOp.getArguments()) {
diff --git a/mlir/test/Dialect/XeGPU/xegpu-propagate-layout-invalid.mlir b/mlir/test/Dialect/XeGPU/xegpu-propagate-layout-invalid.mlir
new file mode 100644
index 0000000000000..ded8a627f9f6a
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/xegpu-propagate-layout-invalid.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt --xegpu-propagate-layout %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
+// llvm.func declaration. updateFunctionOpInterface called setType(FunctionType)
+// on an llvm.func (whose type is LLVMFunctionType), corrupting its
+// function_type attribute; the subsequent getFunctionType() then
+// triggered cast<LLVMFunctionType> on a FunctionType and aborted.
+
+// CHECK-LABEL: llvm.func @some_function()
+module {
+  llvm.func @some_function()
+}

[charithaintc]

LGTM. Thanks!