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

2 files changed, 28 insertions, 0 deletions

diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index 52ab736..93a24de 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 d2c1153..58d4940 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@+1 {{parsed 1 elements, but type ('complex<i64>') expected 2 elements}}
+#attr = dense<0> : tensor<2xcomplex<i64>>
+
+// -----
+
+// expected-error@+1 {{parsed 2 elements, but type ('tensor<2xcomplex<i64>>') expected 4 elements}}
+#attr = dense<[0, 1]> : tensor<2xcomplex<i64>>
+
+// -----
+
+// expected-error@+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}}
 }