Don't elide splat attributes during printing

A splat attribute have a single element during printing so we should
treat it as such when we decide if we elide it or not based on the flag
intended to elide large attributes.

Reviewed By: rriddle, mehdi_amini

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

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 73d272c..9d3e261 100644 (file)
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -158,7 +158,8 @@ OpPrintingFlags &OpPrintingFlags::useLocalScope() {
 /// Return if the given ElementsAttr should be elided.
 bool OpPrintingFlags::shouldElideElementsAttr(ElementsAttr attr) const {
   return elementsAttrElementLimit.hasValue() &&
-         *elementsAttrElementLimit < int64_t(attr.getNumElements());
+         *elementsAttrElementLimit < int64_t(attr.getNumElements()) &&
+         !attr.isa<SplatElementsAttr>();
 }
 
 /// Return the size limit for printing large ElementsAttr.

diff --git a/mlir/test/IR/pretty-attributes.mlir b/mlir/test/IR/pretty-attributes.mlir
index e4e5b34..d4ac8e7 100644 (file)
--- a/mlir/test/IR/pretty-attributes.mlir
+++ b/mlir/test/IR/pretty-attributes.mlir
@@ -16,3 +16,6 @@
 
 // CHECK: opaque<"", "0xDEADBEEF"> : tensor<100xf32>
 "test.opaque_attr"() {foo.opaque_attr = opaque<"", "0xEBFE"> : tensor<100xf32> } : () -> ()
+
+// CHECK: dense<1> : tensor<3xi32>
+"test.dense_splat"() {foo.dense_attr = dense<1> : tensor<3xi32>} : () -> ()