[llvm-branch-commits] [mlir] 75eca67 - [mlir][Parser] Fix crash in DenseElementsAttr parser when no elements are parsed

[River Riddle via llvm-branch-commits]

Author: River Riddle
Date: 2020-12-10T12:48:37-08:00
New Revision: 75eca67c1c4b53a07a70cd3c8036713aec537769

URL: https://github.com/llvm/llvm-project/commit/75eca67c1c4b53a07a70cd3c8036713aec537769
DIFF: https://github.com/llvm/llvm-project/commit/75eca67c1c4b53a07a70cd3c8036713aec537769.diff

LOG: [mlir][Parser] Fix crash in DenseElementsAttr parser when no elements are parsed

This fixes a crash when no elements are parsed, but the type expects at least one.

Fixes PR#47763

Differential Revision: https://reviews.llvm.org/D92982

Added: 
    

Modified: 
    mlir/lib/Parser/AttributeParser.cpp
    mlir/test/IR/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Parser/AttributeParser.cpp b/mlir/lib/Parser/AttributeParser.cpp
index 2e385cfa30e7..e78237e8e5a0 100644
--- a/mlir/lib/Parser/AttributeParser.cpp
+++ b/mlir/lib/Parser/AttributeParser.cpp
@@ -531,6 +531,13 @@ DenseElementsAttr TensorLiteralParser::getAttr(llvm::SMLoc loc,
     return nullptr;
   }
 
+  // Handle the case where no elements were parsed.
+  if (!hexStorage.hasValue() && storage.empty() && type.getNumElements()) {
+    p.emitError(loc) << "parsed zero elements, but type (" << type
+                     << ") expected at least 1";
+    return nullptr;
+  }
+
   // Handle complex types in the specific element type cases below.
   bool isComplex = false;
   if (ComplexType complexTy = eltType.dyn_cast<ComplexType>()) {

diff  --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir
index 4930341a8b64..6b28b33e7c78 100644
--- a/mlir/test/IR/invalid.mlir
+++ b/mlir/test/IR/invalid.mlir
@@ -687,6 +687,11 @@ func @elementsattr_toolarge1() -> () {
 
 // -----
 
+// expected-error at +1 {{parsed zero elements, but type ('tensor<i64>') expected at least 1}}
+#attr = dense<> : tensor<i64>
+
+// -----
+
 func @elementsattr_toolarge2() -> () {
 ^bb0:
   "foo"(){bar = dense<[-777]> : tensor<1xi8>} : () -> () // expected-error {{integer constant out of range}}