[mlir] Add I1 support to DenseArrayAttr

[Mogball]

Summary

This patch adds a DenseI1ArrayAttr to support arrays of i1. Importantly,
the implementation is as a simple ArrayRef<bool> instead of using bit
compression, which was problematic in DenseElementsAttr.

[rriddle]

mlir/include/mlir/IR/BuiltinAttributes.h

808

For the resource version, we use DenseBoolResourceElementsAttr. We should try to keep these aligned.

mlir/test/IR/attribute.mlir

545

Do we intend to allow true or false values here?

[rriddle]

mlir/test/IR/attribute.mlir

545

If the API is really about bool, those would be more natural than raw integers; e.g. given that this is what DenseElementsAttr/ArrayAttr/etc. are all doing.

[Mogball]

mlir/test/IR/attribute.mlir

545

Would it make more sense for the syntax to be [:bool true, false, true]?

[rriddle]

mlir/test/IR/attribute.mlir

545

Nah, for bool as a type we always use i1. It's only values where we use the more ergonomic "true"/"false".

[mehdi_amini]

LG on the principle, with River's comments addressed.

[Mogball]

Change DenseI1ArrayAttr -> DenseBoolArrayAttr and change the values in the
assembly format to be true and false.

[rriddle]

mlir/lib/IR/BuiltinAttributes.cpp

819–821

839–843

Can we upgrade parseInteger to support bool?

[rriddle]

Also, you likely want to leave the bazel group to avoid getting hit on bazel failures (assuming you don't care about those anymore ;) )

[Mogball]

Add bool support to parseInteger

[Mogball]

I think someone has to remove me from the group :P I wasn't able to remove myself!

Also PTAL

[lattner]

+1, thank you for not doing bit compression!

mlir/include/mlir/IR/BuiltinAttributes.h

public:
 
   /// Parse the short form `42, 100, -1` without any type prefix or braces.
   static Attribute parseWithoutBraces(AsmParser &parser, Type odsType);
 
   /// Support for isa<>/cast<>.
   static bool classof(Attribute attr);
 };
 template <>
+void DenseArrayAttr<bool>::printWithoutBraces(raw_ostream &os) const;
+template <>
 void DenseArrayAttr<int8_t>::printWithoutBraces(raw_ostream &os) const;
 
+extern template class DenseArrayAttr<bool>;
 extern template class DenseArrayAttr<int8_t>;
 extern template class DenseArrayAttr<int16_t>;
 extern template class DenseArrayAttr<int32_t>;
 extern template class DenseArrayAttr<int64_t>;
 extern template class DenseArrayAttr<float>;
 extern template class DenseArrayAttr<double>;
 } // namespace detail
 
 // Public name for all the supported DenseArrayAttr
+using DenseBoolArrayAttr = detail::DenseArrayAttr<bool>;
 using DenseI8ArrayAttr = detail::DenseArrayAttr<int8_t>;
 using DenseI16ArrayAttr = detail::DenseArrayAttr<int16_t>;
 using DenseI32ArrayAttr = detail::DenseArrayAttr<int32_t>;
 using DenseI64ArrayAttr = detail::DenseArrayAttr<int64_t>;
 using DenseF32ArrayAttr = detail::DenseArrayAttr<float>;
 using DenseF64ArrayAttr = detail::DenseArrayAttr<double>;
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/BuiltinAttributes.td

let description = [{
     [1, 2, 3]
     ```
   }];
   let parameters = (ins AttributeSelfTypeParameter<"", "ShapedType">:$type,
                         "DenseArrayBaseAttr::EltType":$eltType,
                         ArrayRefParameter<"char">:$elements);
   let extraClassDeclaration = [{
     // All possible supported element type.
-    enum class EltType { I8, I16, I32, I64, F32, F64 };
+    enum class EltType { I1, I8, I16, I32, I64, F32, F64 };
 
     /// Allow implicit conversion to ElementsAttr.
     operator ElementsAttr() const {
       return *this ? cast<ElementsAttr>() : nullptr;
     }
 
     /// ElementsAttr implementation.
     using ContiguousIterableTypesT =
-        std::tuple<int8_t, int16_t, int32_t, int64_t, float, double>;
+        std::tuple<bool, int8_t, int16_t, int32_t, int64_t, float, double>;
+    const bool *value_begin_impl(OverloadToken<bool>) const;
     const int8_t *value_begin_impl(OverloadToken<int8_t>) const;
     const int16_t *value_begin_impl(OverloadToken<int16_t>) const;
     const int32_t *value_begin_impl(OverloadToken<int32_t>) const;
     const int64_t *value_begin_impl(OverloadToken<int64_t>) const;
     const float *value_begin_impl(OverloadToken<float>) const;
     const double *value_begin_impl(OverloadToken<double>) const;
 
     /// Returns the shaped type, containing the number of elements in the array

mlir/include/mlir/IR/OpBase.td

 }
 
 class DenseArrayAttrBase<string denseAttrName, string cppType, string summaryName> :
     ElementsAttrBase<CPred<"$_self.isa<::mlir::" # denseAttrName # ">()">,
                      summaryName # " dense array attribute"> {
   let storageType = "::mlir::" # denseAttrName;
   let returnType = "::llvm::ArrayRef<" # cppType # ">";
 }
+def DenseBoolArrayAttr : DenseArrayAttrBase<"DenseBoolArrayAttr", "bool", "i1">;
 def DenseI8ArrayAttr : DenseArrayAttrBase<"DenseI8ArrayAttr", "int8_t", "i8">;
 def DenseI16ArrayAttr : DenseArrayAttrBase<"DenseI16ArrayAttr", "int16_t", "i16">;
 def DenseI32ArrayAttr : DenseArrayAttrBase<"DenseI32ArrayAttr", "int32_t", "i32">;
 def DenseI64ArrayAttr : DenseArrayAttrBase<"DenseI64ArrayAttr", "int64_t", "i64">;
 def DenseF32ArrayAttr : DenseArrayAttrBase<"DenseF32ArrayAttr", "float", "f32">;
 def DenseF64ArrayAttr : DenseArrayAttrBase<"DenseF64ArrayAttr", "double", "f64">;
 
 def IndexElementsAttr

mlir/lib/AsmParser/AttributeParser.cpp

if (!type)
     return {};
   CustomAsmParser parser(*this);
   Attribute result;
   // Check for empty list.
   bool isEmptyList = getToken().is(Token::r_square);
 
   if (auto intType = type.dyn_cast<IntegerType>()) {
     switch (type.getIntOrFloatBitWidth()) {
+    case 1:
+      if (isEmptyList)
+        result = DenseBoolArrayAttr::get(parser.getContext(), {});
+      else
+        result = DenseBoolArrayAttr::parseWithoutBraces(parser, Type{});
+      break;
     case 8:
       if (isEmptyList)
         result = DenseI8ArrayAttr::get(parser.getContext(), {});
       else
         result = DenseI8ArrayAttr::parseWithoutBraces(parser, Type{});
       break;
     case 16:
       if (isEmptyList)
case 32:
       break;
     case 64:
       if (isEmptyList)
         result = DenseI64ArrayAttr::get(parser.getContext(), {});
       else
         result = DenseI64ArrayAttr::parseWithoutBraces(parser, Type{});
       break;
     default:
-      emitError(typeLoc, "expected i8, i16, i32, or i64 but got: ") << type;
+      emitError(typeLoc, "expected i1, i8, i16, i32, or i64 but got: ") << type;
       return {};
     }
   } else if (auto floatType = type.dyn_cast<FloatType>()) {
     switch (type.getIntOrFloatBitWidth()) {
     case 32:
       if (isEmptyList)
         result = DenseF32ArrayAttr::get(parser.getContext(), {});
       else

mlir/lib/AsmParser/Parser.cpp

ParseResult Parser::parseToken(Token::Kind expectedToken,
                                const Twine &message) {
   if (consumeIf(expectedToken))
     return success();
   return emitWrongTokenError(message);
 }
 
 /// Parse an optional integer value from the stream.
 OptionalParseResult Parser::parseOptionalInteger(APInt &result) {
+  // Parse `false` and `true` keywords as 0 and 1 respectively.
+  if (consumeIf(Token::kw_false)) {
+    result = false;
+    return success();
+  } else if (consumeIf(Token::kw_true)) {
+    result = true;
+    return success();
+  }
+
   Token curToken = getToken();
   if (curToken.isNot(Token::integer, Token::minus))
     return llvm::None;
 
   bool negative = consumeIf(Token::minus);
   Token curTok = getToken();
   if (parseToken(Token::integer, "expected integer value"))
     return failure();

mlir/lib/IR/AsmPrinter.cpp

if (printerFlags.shouldElideElementsAttr(sparseEltAttr.getIndices()) ||
         printDenseIntOrFPElementsAttr(indices, /*allowHex=*/false);
         os << ", ";
         printDenseElementsAttr(sparseEltAttr.getValues(), /*allowHex=*/true);
       }
       os << '>';
     }
   } else if (auto denseArrayAttr = attr.dyn_cast<DenseArrayBaseAttr>()) {
     typeElision = AttrTypeElision::Must;
-    switch (denseArrayAttr.getElementType()) {
-    case DenseArrayBaseAttr::EltType::I8:
-      os << "[:i8";
-      break;
-    case DenseArrayBaseAttr::EltType::I16:
-      os << "[:i16";
-      break;
-    case DenseArrayBaseAttr::EltType::I32:
-      os << "[:i32";
-      break;
-    case DenseArrayBaseAttr::EltType::I64:
-      os << "[:i64";
-      break;
-    case DenseArrayBaseAttr::EltType::F32:
-      os << "[:f32";
-      break;
-    case DenseArrayBaseAttr::EltType::F64:
-      os << "[:f64";
-      break;
-    }
+    os << "[:" << denseArrayAttr.getType().getElementType();
     if (denseArrayAttr.size())
       os << " ";
     denseArrayAttr.printWithoutBraces(os);
     os << "]";
   } else if (auto resourceAttr = attr.dyn_cast<DenseResourceElementsAttr>()) {
     os << "dense_resource<";
     printResourceHandle(resourceAttr.getRawHandle());
     os << ">";

mlir/lib/IR/BuiltinAttributes.cpp

 };
 
 DenseArrayBaseAttr::EltType DenseArrayBaseAttr::getElementType() const {
   return getImpl()->eltType;
 }
 
 ShapedType DenseArrayBaseAttr::getType() const { return getImpl()->type; }
 
+const bool *DenseArrayBaseAttr::value_begin_impl(OverloadToken<bool>) const {
+  return cast<DenseBoolArrayAttr>().asArrayRef().begin();
+}
 const int8_t *
 DenseArrayBaseAttr::value_begin_impl(OverloadToken<int8_t>) const {
   return cast<DenseI8ArrayAttr>().asArrayRef().begin();
 }
 const int16_t *
 DenseArrayBaseAttr::value_begin_impl(OverloadToken<int16_t>) const {
   return cast<DenseI16ArrayAttr>().asArrayRef().begin();
 }

 }
 
 void DenseArrayBaseAttr::print(AsmPrinter &printer) const {
   print(printer.getStream());
 }
 
 void DenseArrayBaseAttr::printWithoutBraces(raw_ostream &os) const {
   switch (getElementType()) {
+  case DenseArrayBaseAttr::EltType::I1:
+    this->cast<DenseBoolArrayAttr>().printWithoutBraces(os);
+    return;
   case DenseArrayBaseAttr::EltType::I8:
     this->cast<DenseI8ArrayAttr>().printWithoutBraces(os);
     return;
   case DenseArrayBaseAttr::EltType::I16:
     this->cast<DenseI16ArrayAttr>().printWithoutBraces(os);
     return;
   case DenseArrayBaseAttr::EltType::I32:
     this->cast<DenseI32ArrayAttr>().printWithoutBraces(os);

 
 template <typename T>
 void DenseArrayAttr<T>::print(AsmPrinter &printer) const {
   print(printer.getStream());
 }
 
 template <typename T>
 void DenseArrayAttr<T>::printWithoutBraces(raw_ostream &os) const {
-  ArrayRef<T> values{*this};
-  llvm::interleaveComma(values, os);
+  llvm::interleaveComma(asArrayRef(), os);
+}
+
+/// Specialization for bool to print `true` or `false`.
+template <>
+void DenseArrayAttr<bool>::printWithoutBraces(raw_ostream &os) const {
+  llvm::interleaveComma(asArrayRef(), os,
+                        [&](bool v) { os << (v ? "true" : "false"); });
 }
 
 /// Specialization for int8_t for forcing printing as number instead of chars.
 template <>
 void DenseArrayAttr<int8_t>::printWithoutBraces(raw_ostream &os) const {
-  ArrayRef<int8_t> values{*this};
-  llvm::interleaveComma(values, os, [&](int64_t v) { os << v; });
+  llvm::interleaveComma(asArrayRef(), os, [&](int64_t v) { os << v; });
 }
 
 template <typename T>
 void DenseArrayAttr<T>::print(raw_ostream &os) const {
   os << "[";
   printWithoutBraces(os);
   os << "]";
 }
 
 /// Parse a single element: generic template for int types, specialized for
-/// floating points below.
+/// floating point and boolean values below.
 template <typename T>
 static ParseResult parseDenseArrayAttrElt(AsmParser &parser, T &value) {
   return parser.parseInteger(value);
 }
 
 template <>
 ParseResult parseDenseArrayAttrElt<float>(AsmParser &parser, float &value) {
   double doubleVal;
   if (parser.parseFloat(doubleVal))
     return failure();
   value = doubleVal;
   return success();
 }
 
 template <>
 ParseResult parseDenseArrayAttrElt<double>(AsmParser &parser, double &value) {
   return parser.parseFloat(value);
 }
 
 /// Parse a DenseArrayAttr without the braces: `1, 2, 3`
 template <typename T>
return ArrayRef<T>(reinterpret_cast<const T *>(raw.data()),
                      raw.size() / sizeof(T));
 }
 
 namespace {
 /// Mapping from C++ element type to MLIR DenseArrayAttr internals.
 template <typename T>
 struct denseArrayAttrEltTypeBuilder;
 template <>
+struct denseArrayAttrEltTypeBuilder<bool> {
+  constexpr static auto eltType = DenseArrayBaseAttr::EltType::I1;
+  static ShapedType getShapedType(MLIRContext *context,
+                                  ArrayRef<int64_t> shape) {
+    return RankedTensorType::get(shape, IntegerType::get(context, 1));
+  }
+};
+template <>
 struct denseArrayAttrEltTypeBuilder<int8_t> {
   constexpr static auto eltType = DenseArrayBaseAttr::EltType::I8;
   static ShapedType getShapedType(MLIRContext *context,
                                   ArrayRef<int64_t> shape) {
     return RankedTensorType::get(shape, IntegerType::get(context, 8));
   }
 };
 template <>
bool DenseArrayAttr<T>::classof(Attribute attr) {
   return attr.isa<DenseArrayBaseAttr>() &&
          attr.cast<DenseArrayBaseAttr>().getElementType() ==
              denseArrayAttrEltTypeBuilder<T>::eltType;
 }
 
 namespace mlir {
 namespace detail {
 // Explicit instantiation for all the supported DenseArrayAttr.
+template class DenseArrayAttr<bool>;
 template class DenseArrayAttr<int8_t>;
 template class DenseArrayAttr<int16_t>;
 template class DenseArrayAttr<int32_t>;
 template class DenseArrayAttr<int64_t>;
 template class DenseArrayAttr<float>;
 template class DenseArrayAttr<double>;
 } // namespace detail
 } // namespace mlir

mlir/test/IR/attribute.mlir

 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Test DenseArrayAttr
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: func @dense_array_attr
-func.func @dense_array_attr() attributes{
+func.func @dense_array_attr() attributes {
 // CHECK-SAME: emptyf32attr = [:f32],
                emptyf32attr = [:f32],
 // CHECK-SAME: emptyf64attr = [:f64],
                emptyf64attr = [:f64],
 // CHECK-SAME: emptyi16attr = [:i16],
                emptyi16attr = [:i16],
+// CHECK-SAME: emptyi1attr = [:i1],
+               emptyi1attr = [:i1],
 // CHECK-SAME: emptyi32attr = [:i32],
                emptyi32attr = [:i32],
 // CHECK-SAME: emptyi64attr = [:i64],
                emptyi64attr = [:i64],
 // CHECK-SAME: emptyi8attr = [:i8],
                emptyi8attr = [:i8],
 // CHECK-SAME: f32attr = [:f32 1.024000e+03, 4.530000e+02, -6.435000e+03],
                f32attr = [:f32 1024., 453., -6435.],
 // CHECK-SAME: f64attr = [:f64 -1.420000e+02],
                f64attr = [:f64 -142.],
 // CHECK-SAME: i16attr = [:i16 3, 5, -4, 10],
                i16attr = [:i16 3, 5, -4, 10],
+// CHECK-SAME: i1attr = [:i1 true, false, true],
+               i1attr = [:i1 true, false, true],
 // CHECK-SAME: i32attr = [:i32 1024, 453, -6435],
                i32attr = [:i32 1024, 453, -6435],
 // CHECK-SAME: i64attr = [:i64 -142],
                i64attr = [:i64 -142],
 // CHECK-SAME: i8attr = [:i8 1, -2, 3]
                i8attr = [:i8 1, -2, 3]
  } {
 // CHECK:  test.dense_array_attr
   test.dense_array_attr
+// CHECK-SAME: i1attr = [true, false, true]
+               i1attr = [true, false, true]
 // CHECK-SAME: i8attr = [1, -2, 3]
                i8attr = [1, -2, 3]
 // CHECK-SAME: i16attr = [3, 5, -4, 10]
                i16attr = [3, 5, -4, 10]
 // CHECK-SAME: i32attr = [1024, 453, -6435]
                i32attr = [1024, 453, -6435]
 // CHECK-SAME: i64attr = [-142]
                i64attr = [-142]

mlir/test/IR/elements-attr-interface.mlir

 
 // Check that we don't crash on empty element attributes.
 // expected-error@below {{Test iterating `int64_t`: }}
 // expected-error@below {{Test iterating `uint64_t`: }}
 // expected-error@below {{Test iterating `APInt`: }}
 // expected-error@below {{Test iterating `IntegerAttr`: }}
 arith.constant dense<> : tensor<0xi64>
 
+// expected-error@below {{Test iterating `bool`: true, false, true, false, true, false}}
+arith.constant [:i1 true, false, true, false, true, false]
 // expected-error@below {{Test iterating `int8_t`: 10, 11, -12, 13, 14}}
 arith.constant [:i8 10, 11, -12, 13, 14]
 // expected-error@below {{Test iterating `int16_t`: 10, 11, -12, 13, 14}}
 arith.constant [:i16 10, 11, -12, 13, 14]
 // expected-error@below {{Test iterating `int32_t`: 10, 11, -12, 13, 14}}
 arith.constant [:i32 10, 11, -12, 13, 14]
 // expected-error@below {{Test iterating `int64_t`: 10, 11, -12, 13, 14}}
 arith.constant [:i64 10, 11, -12, 13, 14]

mlir/test/lib/Dialect/Test/TestOps.td

 def StringElementsAttrOp : TEST_Op<"string_elements_attr"> {
   let arguments = (ins
       StringElementsAttr:$scalar_string_attr
   );
 }
 
 def DenseArrayAttrOp : TEST_Op<"dense_array_attr"> {
   let arguments = (ins
+    DenseBoolArrayAttr:$i1attr,
     DenseI8ArrayAttr:$i8attr,
     DenseI16ArrayAttr:$i16attr,
     DenseI32ArrayAttr:$i32attr,
     DenseI64ArrayAttr:$i64attr,
     DenseF32ArrayAttr:$f32attr,
     DenseF64ArrayAttr:$f64attr,
     DenseI32ArrayAttr:$emptyattr
   );
   let assemblyFormat = [{
-   `i8attr` `=` $i8attr `i16attr` `=` $i16attr `i32attr` `=` $i32attr
-   `i64attr` `=` $i64attr  `f32attr` `=` $f32attr `f64attr` `=` $f64attr
-   `emptyattr` `=` $emptyattr
-   attr-dict
+   `i1attr` `=` $i1attr `i8attr` `=` $i8attr `i16attr` `=` $i16attr
+   `i32attr` `=` $i32attr `i64attr` `=` $i64attr  `f32attr` `=` $f32attr
+   `f64attr` `=` $f64attr `emptyattr` `=` $emptyattr attr-dict
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Test Enum Attributes
 //===----------------------------------------------------------------------===//
 
 // Define the C++ enum.

mlir/test/lib/IR/TestBuiltinAttributeInterfaces.cpp

void runOnOperation() override {
     getOperation().walk([&](Operation *op) {
       for (NamedAttribute attr : op->getAttrs()) {
         auto elementsAttr = attr.getValue().dyn_cast<ElementsAttr>();
         if (!elementsAttr)
           continue;
         if (auto concreteAttr =
                 attr.getValue().dyn_cast<DenseArrayBaseAttr>()) {
           switch (concreteAttr.getElementType()) {
+          case DenseArrayBaseAttr::EltType::I1:
+            testElementsAttrIteration<bool>(op, elementsAttr, "bool");
+            break;
           case DenseArrayBaseAttr::EltType::I8:
             testElementsAttrIteration<int8_t>(op, elementsAttr, "int8_t");
             break;
           case DenseArrayBaseAttr::EltType::I16:
             testElementsAttrIteration<int16_t>(op, elementsAttr, "int16_t");
             break;
           case DenseArrayBaseAttr::EltType::I32:
             testElementsAttrIteration<int32_t>(op, elementsAttr, "int32_t");