[Mlir-commits] [mlir] 1cf6786 - [mlir] Improve error handling for dense attribute parsing in complex types (#133220)

[llvmlistbot at llvm.org]

Author: Longsheng Mou
Date: 2025-04-01T16:27:43+08:00
New Revision: 1cf6786e322ddc787d793dbb48d59b4f9827fef3

URL: https://github.com/llvm/llvm-project/commit/1cf6786e322ddc787d793dbb48d59b4f9827fef3
DIFF: https://github.com/llvm/llvm-project/commit/1cf6786e322ddc787d793dbb48d59b4f9827fef3.diff

LOG: [mlir] Improve error handling for dense attribute parsing in complex types (#133220)

- For splat dense attributes, the number of parsed elements must be 2.
- For non-splat dense attributes, the number of parsed elements must be
twice the number of elements in the type.

Fixes #132859.

Added: 
    

Modified: 
    mlir/lib/AsmParser/AttributeParser.cpp
    mlir/test/IR/invalid-builtin-attributes.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index 52ab736bac03a..93a24dee29ad2 100644
--- a/mlir/lib/AsmParser/AttributeParser.cpp
+++ b/mlir/lib/AsmParser/AttributeParser.cpp
@@ -566,6 +566,19 @@ DenseElementsAttr TensorLiteralParser::getAttr(SMLoc loc, ShapedType type) {
   if (ComplexType complexTy = dyn_cast<ComplexType>(eltType)) {
     eltType = complexTy.getElementType();
     isComplex = true;
+    // Complex types have 2 elements.
+    if (shape.empty() && storage.size() != 2) {
+      p.emitError(loc) << "parsed " << storage.size() << " elements, but type ("
+                       << complexTy << ") expected 2 elements";
+      return nullptr;
+    }
+    if (!shape.empty() &&
+        storage.size() != static_cast<size_t>(type.getNumElements()) * 2) {
+      p.emitError(loc) << "parsed " << storage.size() << " elements, but type ("
+                       << type << ") expected " << type.getNumElements() * 2
+                       << " elements";
+      return nullptr;
+    }
   }
 
   // Handle integer and index types.

diff  --git a/mlir/test/IR/invalid-builtin-attributes.mlir b/mlir/test/IR/invalid-builtin-attributes.mlir
index d2c11536404ea..58d4940eaf874 100644
--- a/mlir/test/IR/invalid-builtin-attributes.mlir
+++ b/mlir/test/IR/invalid-builtin-attributes.mlir
@@ -63,6 +63,21 @@ func.func @elementsattr_toolarge1() -> () {
 
 // -----
 
+// expected-error at +1 {{parsed 1 elements, but type ('complex<i64>') expected 2 elements}}
+#attr = dense<0> : tensor<2xcomplex<i64>>
+
+// -----
+
+// expected-error at +1 {{parsed 2 elements, but type ('tensor<2xcomplex<i64>>') expected 4 elements}}
+#attr = dense<[0, 1]> : tensor<2xcomplex<i64>>
+
+// -----
+
+// expected-error at +1 {{parsed 3 elements, but type ('tensor<2xcomplex<i64>>') expected 4 elements}}
+#attr = dense<[0, (0, 1)]> : tensor<2xcomplex<i64>>
+
+// -----
+
 func.func @elementsattr_toolarge2() -> () {
   "foo"(){bar = dense<[-777]> : tensor<1xi8>} : () -> () // expected-error {{integer constant out of range}}
 }