[Mlir-commits] [mlir] e5a4cf6 - [mlir] Fix printer when it is a DenseElementsAttr of i1

[River Riddle]

Author: Chengji Yao
Date: 2022-05-05T16:56:43-07:00
New Revision: e5a4cf67431ac6e04d739bb5ff4e5c0c1c19f471

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

LOG: [mlir] Fix printer when it is a DenseElementsAttr of i1

A large DenseElementsAttr of i1could trigger a bug in printer/parser roundtrip.

Ex. A DenseElementsAttr of i1 with 200 elements will print as Hex format of length 400 before the fix. However, when parsing the printed text, an error will be triggered. After fix, the printed length will be 50.

Reviewed By: rriddle

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

Added: 
    mlir/test/IR/attribute-roundtrip.mlir

Modified: 
    mlir/lib/IR/BuiltinAttributes.cpp
    mlir/test/Target/SPIRV/constant.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index 6988d1f8e4c30..0004fe90fe87a 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -696,8 +696,8 @@ DenseElementsAttr DenseElementsAttr::get(ShapedType type,
   size_t storageBitWidth = getDenseElementStorageWidth(bitWidth);
 
   // Compress the attribute values into a character buffer.
-  SmallVector<char, 8> data(llvm::divideCeil(storageBitWidth, CHAR_BIT) *
-                            values.size());
+  SmallVector<char, 8> data(
+      llvm::divideCeil(storageBitWidth * values.size(), CHAR_BIT));
   APInt intVal;
   for (unsigned i = 0, e = values.size(); i < e; ++i) {
     assert(eltType == values[i].getType() &&
@@ -1027,7 +1027,7 @@ int64_t DenseElementsAttr::getNumElements() const {
 template <typename APRangeT>
 static void writeAPIntsToBuffer(size_t storageWidth, std::vector<char> &data,
                                 APRangeT &&values) {
-  data.resize(llvm::divideCeil(storageWidth, CHAR_BIT) * llvm::size(values));
+  data.resize(llvm::divideCeil(storageWidth * llvm::size(values), CHAR_BIT));
   size_t offset = 0;
   for (auto it = values.begin(), e = values.end(); it != e;
        ++it, offset += storageWidth) {
@@ -1184,7 +1184,7 @@ static ShapedType mappingHelper(Fn mapping, Attr &attr, ShapedType inType,
     assert(newArrayType && "Unhandled tensor type");
 
   size_t numRawElements = attr.isSplat() ? 1 : newArrayType.getNumElements();
-  data.resize(llvm::divideCeil(storageBitWidth, CHAR_BIT) * numRawElements);
+  data.resize(llvm::divideCeil(storageBitWidth * numRawElements, CHAR_BIT));
 
   // Functor used to process a single element value of the attribute.
   auto processElt = [&](decltype(*attr.begin()) value, size_t index) {

diff  --git a/mlir/test/IR/attribute-roundtrip.mlir b/mlir/test/IR/attribute-roundtrip.mlir
new file mode 100644
index 0000000000000..6c34fcce57e25
--- /dev/null
+++ b/mlir/test/IR/attribute-roundtrip.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt -canonicalize %s | mlir-opt | FileCheck %s
+
+// CHECK-LABEL: @large_i1_tensor_roundtrip
+func @large_i1_tensor_roundtrip() -> tensor<160xi1> {
+  %cst_0 = arith.constant dense<"0xFFF00000FF000000FF000000FF000000FF000000"> : tensor<160xi1>
+  %cst_1 = arith.constant dense<"0xFF000000FF000000FF000000FF000000FF0000F0"> : tensor<160xi1>
+  // CHECK: dense<"0xFF000000FF000000FF000000FF000000FF000000">
+  %0 = arith.andi %cst_0, %cst_1 : tensor<160xi1>
+  return %0 : tensor<160xi1>
+}

diff  --git a/mlir/test/Target/SPIRV/constant.mlir b/mlir/test/Target/SPIRV/constant.mlir
index d79fbdce98999..6ce506459f67a 100644
--- a/mlir/test/Target/SPIRV/constant.mlir
+++ b/mlir/test/Target/SPIRV/constant.mlir
@@ -159,7 +159,7 @@ spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   spv.func @bool_vector_const() -> () "None" {
     // CHECK: spv.Constant dense<false> : vector<2xi1>
     %0 = spv.Constant dense<false> : vector<2xi1>
-    // CHECK: spv.Constant dense<[true, true, true]> : vector<3xi1>
+    // CHECK: spv.Constant dense<true> : vector<3xi1>
     %1 = spv.Constant dense<true> : vector<3xi1>
     // CHECK: spv.Constant dense<[false, true]> : vector<2xi1>
     %2 = spv.Constant dense<[false, true]> : vector<2xi1>