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

[llvmlistbot at llvm.org]

Author: vabridgers
Date: 2026-08-24T09:11:50-05:00
New Revision: 7e0fc20933f48be89b0f82046a1cfb1ab53ee37a

URL: https://github.com/llvm/llvm-project/commit/7e0fc20933f48be89b0f82046a1cfb1ab53ee37a
DIFF: https://github.com/llvm/llvm-project/commit/7e0fc20933f48be89b0f82046a1cfb1ab53ee37a.diff

LOG: [mlir][linalg] Fix splat fold crash on non-TypedAttr element types (#218012)

`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.

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
    mlir/test/Dialect/Linalg/canonicalize.mlir

Removed: 
    


################################################################################
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>,