[Mlir-commits] [mlir] [mlir][linalg] Fix splat fold crash on non-TypedAttr element types (PR #218012)

[llvmlistbot at llvm.org]

https://github.com/vabridgers updated https://github.com/llvm/llvm-project/pull/218012

>From f25be110d105433b1c8df95c60dfa141771612ba Mon Sep 17 00:00:00 2001
From: Vince Bridgers <vince.a.bridgers at ericsson.com>
Date: Fri, 21 Aug 2026 21:26:25 +0200
Subject: [PATCH] [mlir][linalg] Fix splat fold crash on non-TypedAttr element
 types

`getScalarConstantAttrFromDenseSplat` returns `getSplatValue<TypedAttr>()`, which
for a derived attribute type is an unchecked `llvm::cast`. Complex element types
store their splat as an `ArrayAttr` of two values, and `ArrayAttr` does not
implement `TypedAttr`, so the cast asserts on

    %cst = arith.constant dense<(1.0,2.0)> : tensor<3xcomplex<f32>>
    %0   = linalg.broadcast ins(%cst : tensor<3xcomplex<f32>>)
            outs(%init : tensor<2x3xcomplex<f32>>) dimensions = [0]

under `mlir-opt --canonicalize`. Integer and float splats yield
`IntegerAttr`/`FloatAttr`, which are `TypedAttr`, so only non-scalar element types
are affected.

Guard with a `dyn_cast` and decline the fold when the splat value has no
`TypedAttr` representation. The helper is shared by the broadcast pattern added in
PR 195980 and the transpose patterns added in PR 195991, so both call sites are
fixed; a test is added for each, as the transpose case was not previously covered.
---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp   |  8 ++++-
 mlir/test/Dialect/Linalg/canonicalize.mlir | 34 ++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 170e1edf8a55d..af670b8d4fae7 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -98,7 +98,13 @@ getScalarConstantAttrFromDenseSplat(Value input) {
   if (!splatAttr || !splatAttr.isSplat())
     return std::nullopt;
 
-  return splatAttr.getSplatValue<TypedAttr>();
+  // Not every element type has a TypedAttr splat value: a complex splat, for
+  // one, is an ArrayAttr. Decline the fold instead of asserting in the cast.
+  auto splatValue = dyn_cast<TypedAttr>(splatAttr.getSplatValue<Attribute>());
+  if (!splatValue)
+    return std::nullopt;
+
+  return splatValue;
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index bb11ce0d4dfb8..3fe824634b18f 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1219,6 +1219,23 @@ func.func @broadcast_non_splat_constant(%init: tensor<2x3xf32>) -> tensor<2x3xf3
   return %0 : tensor<2x3xf32>
 }
 
+
+// -----
+
+// A splat of a complex element type yields an ArrayAttr, which is not a
+// TypedAttr, so the fold must decline instead of asserting in the cast.
+// CHECK-LABEL: @broadcast_splat_constant_complex
+//       CHECK:   %[[BROADCAST:.+]] = linalg.broadcast
+//       CHECK:   return %[[BROADCAST]] : tensor<2x3xcomplex<f32>>
+func.func @broadcast_splat_constant_complex(%init: tensor<2x3xcomplex<f32>>)
+    -> tensor<2x3xcomplex<f32>> {
+  %cst = arith.constant dense<(1.000000e+00,2.000000e+00)> : tensor<3xcomplex<f32>>
+  %0 = linalg.broadcast
+      ins(%cst: tensor<3xcomplex<f32>>)
+      outs(%init: tensor<2x3xcomplex<f32>>)
+      dimensions = [0]
+  return %0 : tensor<2x3xcomplex<f32>>
+}
 // -----
 
 // CHECK-LABEL: @broadcast_broadcast_fold
@@ -1350,6 +1367,23 @@ func.func @transpose_non_splat_constant(%init: tensor<3x2xf32>) -> tensor<3x2xf3
   func.return %transpose : tensor<3x2xf32>
 }
 
+
+// -----
+
+// A splat of a complex element type yields an ArrayAttr, which is not a
+// TypedAttr, so the fold must decline instead of asserting in the cast.
+// CHECK-LABEL: @transpose_splat_constant_complex
+//       CHECK:   %[[TRANSPOSE:.+]] = linalg.transpose
+//       CHECK:   return %[[TRANSPOSE]] : tensor<3x2xcomplex<f32>>
+func.func @transpose_splat_constant_complex(%init: tensor<3x2xcomplex<f32>>)
+    -> tensor<3x2xcomplex<f32>> {
+  %cst = arith.constant dense<(1.000000e+00,2.000000e+00)> : tensor<2x3xcomplex<f32>>
+  %transpose = linalg.transpose
+      ins(%cst: tensor<2x3xcomplex<f32>>)
+      outs(%init: tensor<3x2xcomplex<f32>>)
+      permutation = [1, 0]
+  func.return %transpose : tensor<3x2xcomplex<f32>>
+}
 // -----
 
 func.func @transpose_transpose_cancel(%input: tensor<5x4x3xf32>,