[mlir][ElementsAttr] Change value_begin_impl to try_value_begin_impl

[Mogball]

Summary

This patch changes value_begin_impl to a faillable
try_value_begin_impl so that specific cases can fail iteration if the
type doesn't match the internal storage.

[rriddle]

mlir/include/mlir/IR/BuiltinAttributes.td

435–437

This doesn't look right, this implementation will crash if it can't iterate. I suppose that these are more annoying to check for (DenseElementsAttr+SparseElementsAttr), but we'll need at least a TODO here. This isn't the desired end state.

mlir/include/mlir/Support/LogicalResult.h

103

Can we use is_convertible here instead of is_same?

[Mogball]

mlir/include/mlir/IR/BuiltinAttributes.td

435–437

ElementsAttr::getValues also crashes if it can't iterate. That's why there's a tryGetValues method. Should I add one here as well?

[Mogball]

mlir/include/mlir/IR/BuiltinAttributes.td

435–437

Whoops I meant try_value_begin 🙃 ⚡

[Mogball]

use std::is_convertible_v

[Mogball]

mlir/include/mlir/IR/BuiltinAttributes.td

435–437

Done in https://reviews.llvm.org/D132958

mlir/include/mlir/IR/BuiltinAttributeInterfaces.td

let description = [{
     the elements being stored within an ArrayRef. This attribute could
     potentially define the iterable types as so:
 
     ```c++
     using ContiguousIterableTypesT = std::tuple<uint64_t>;
     using NonContiguousIterableTypesT = std::tuple<APInt, Attribute>;
     ```
 
-    * Provide a `iterator value_begin_impl(OverloadToken<T>) const` overload for
-      each iterable type
+    * Provide a `FailureOr<iterator> try_value_begin_impl(OverloadToken<T>) const`
+      overload for each iterable type
 
     These overloads should return an iterator to the start of the range for the
-    respective iterable type. Consider the example i64 elements attribute
-    described in the previous section. This attribute may define the
-    value_begin_impl overloads like so:
+    respective iterable type or fail if the type cannot be iterated. Consider
+    the example i64 elements attribute described in the previous section. This
+    attribute may define the value_begin_impl overloads like so:
 
     ```c++
     /// Provide begin iterators for the various iterable types.
     /// * uint64_t
-    auto value_begin_impl(OverloadToken<uint64_t>) const {
+    FailureOr<const uint64_t *>
+    value_begin_impl(OverloadToken<uint64_t>) const {
       return getElements().begin();
     }
     /// * APInt
     auto value_begin_impl(OverloadToken<llvm::APInt>) const {
-      return llvm::map_range(getElements(), [=](uint64_t value) {
+      auto it = llvm::map_range(getElements(), [=](uint64_t value) {
         return llvm::APInt(/*numBits=*/64, value);
       }).begin();
+      return FailureOr<decltype(it)>(std::move(it));
     }
     /// * Attribute
     auto value_begin_impl(OverloadToken<mlir::Attribute>) const {
       mlir::Type elementType = getType().getElementType();
-      return llvm::map_range(getElements(), [=](uint64_t value) {
+      auto it = llvm::map_range(getElements(), [=](uint64_t value) {
         return mlir::IntegerAttr::get(elementType,
                                       llvm::APInt(/*numBits=*/64, value));
       }).begin();
+      return FailureOr<decltype(it)>(std::move(it));
     }
     ```
 
     After the above, ElementsAttr will now be able to iterate over elements
     using each of the registered iterable data types:
 
     ```c++
     ElementsAttr attr = myI64ElementsAttr;
private:
     template <typename T>
     ::mlir::FailureOr<::mlir::detail::ElementsAttrIndexer> buildValueResult(
         /*isContiguous*/std::true_type) const {
       if ($_attr.empty()) {
         return ::mlir::detail::ElementsAttrIndexer::contiguous<T>(
           /*isSplat=*/false, nullptr);
       }
 
-      auto valueIt = $_attr.value_begin_impl(OverloadToken<T>());
+      auto valueIt = $_attr.try_value_begin_impl(OverloadToken<T>());
+      if (::mlir::failed(valueIt))
+        return ::mlir::failure();
       return ::mlir::detail::ElementsAttrIndexer::contiguous(
-        $_attr.isSplat(), &*valueIt);
+        $_attr.isSplat(), &**valueIt);
     }
     /// Build an indexer for the given type `T`, which is represented via a
     /// non-contiguous range.
     template <typename T>
     ::mlir::FailureOr<::mlir::detail::ElementsAttrIndexer> buildValueResult(
         /*isContiguous*/std::false_type) const {
-      auto valueIt = $_attr.value_begin_impl(OverloadToken<T>());
+      auto valueIt = $_attr.try_value_begin_impl(OverloadToken<T>());
+      if (::mlir::failed(valueIt))
+        return ::mlir::failure();
       return ::mlir::detail::ElementsAttrIndexer::nonContiguous(
-        $_attr.isSplat(), valueIt);
+        $_attr.isSplat(), *valueIt);
     }
 
   public:
     //===------------------------------------------------------------------===//
     // Value Iteration
     //===------------------------------------------------------------------===//
 
     /// The iterator for the given element type T.
     template <typename T, typename AttrT = ConcreteAttr>
     using iterator = decltype(std::declval<AttrT>().template value_begin<T>());
     /// The iterator range over the given element T.
     template <typename T, typename AttrT = ConcreteAttr>
     using iterator_range =
         decltype(std::declval<AttrT>().template getValues<T>());
 
     /// Return an iterator to the first element of this attribute as a value of
     /// type `T`.
     template <typename T>
     auto value_begin() const {
-      return $_attr.value_begin_impl(OverloadToken<T>());
+      return *$_attr.try_value_begin_impl(OverloadToken<T>());
     }
 
     /// Return the elements of this attribute as a value of type 'T'.
     template <typename T>
     auto getValues() const {
       auto beginIt = $_attr.template value_begin<T>();
       return detail::ElementsAttrRange<decltype(beginIt)>(
         $_attr.getType(), beginIt, std::next(beginIt, size()));

mlir/include/mlir/IR/BuiltinAttributes.td

let extraClassDeclaration = [{
     /// Allow implicit conversion to ElementsAttr.
     operator ElementsAttr() const {
       return *this ? cast<ElementsAttr>() : nullptr;
     }
 
     /// ElementsAttr implementation.
     using ContiguousIterableTypesT =
         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;
+    FailureOr<const bool *>
+    try_value_begin_impl(OverloadToken<bool>) const;
+    FailureOr<const int8_t *>
+    try_value_begin_impl(OverloadToken<int8_t>) const;
+    FailureOr<const int16_t *>
+    try_value_begin_impl(OverloadToken<int16_t>) const;
+    FailureOr<const int32_t *>
+    try_value_begin_impl(OverloadToken<int32_t>) const;
+    FailureOr<const int64_t *>
+    try_value_begin_impl(OverloadToken<int64_t>) const;
+    FailureOr<const float *>
+    try_value_begin_impl(OverloadToken<float>) const;
+    FailureOr<const double *>
+    try_value_begin_impl(OverloadToken<double>) const;
   }];
 
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DenseIntOrFPElementsAttr
 //===----------------------------------------------------------------------===//
let extraClassDeclaration = [{
     using NonContiguousIterableTypesT = std::tuple<
       Attribute,
       // Integer types.
       APInt, bool, std::complex<APInt>,
       // Float types.
       APFloat, std::complex<APFloat>
     >;
 
-    /// Provide a `value_begin_impl` to enable iteration within ElementsAttr.
+    /// Provide a `try_value_begin_impl` to enable iteration within
+    /// ElementsAttr.
     template <typename T>
-    auto value_begin_impl(OverloadToken<T>) const {
-      return value_begin<T>();
+    auto try_value_begin_impl(OverloadToken<T>) const {
+      return ::mlir::success(value_begin<T>());
     }
 
     /// Convert endianess of input ArrayRef for big-endian(BE) machines. All of
     /// the elements of `inRawData` has `type`. If `inRawData` is little endian
     /// (LE), it is converted to big endian (BE). Conversely, if `inRawData` is
     /// BE, converted to LE.
     static void
     convertEndianOfArrayRefForBEmachine(ArrayRef<char> inRawData,
let extraClassDeclaration = [{
     using DenseElementsAttr::isSplat;
     using DenseElementsAttr::size;
     using DenseElementsAttr::value_begin;
 
     /// The set of data types that can be iterated by this attribute.
     using ContiguousIterableTypesT = std::tuple<StringRef>;
     using NonContiguousIterableTypesT = std::tuple<Attribute>;
 
-    /// Provide a `value_begin_impl` to enable iteration within ElementsAttr.
+    /// Provide a `try_value_begin_impl` to enable iteration within
+    /// ElementsAttr.
     template <typename T>
-    auto value_begin_impl(OverloadToken<T>) const {
-      return value_begin<T>();
+    auto try_value_begin_impl(OverloadToken<T>) const {
+      return ::mlir::success(value_begin<T>());
     }
 
   protected:
     friend DenseElementsAttr;
 
   public:
   }];
   let genAccessors = 0;
   let genStorageClass = 0;
using NonContiguousIterableTypesT = std::tuple<
       // Float types.
       APFloat, float, double,
       std::complex<APFloat>, std::complex<float>, std::complex<double>,
       // String types.
       StringRef
     >;
     using ElementsAttr::Trait<SparseElementsAttr>::getValues;
 
-    /// Provide a `value_begin_impl` to enable iteration within ElementsAttr.
+    /// Provide a `try_value_begin_impl` to enable iteration within
+    /// ElementsAttr.
     template <typename T>
-    auto value_begin_impl(OverloadToken<T>) const {
-      return value_begin<T>();
+    auto try_value_begin_impl(OverloadToken<T>) const {
+      return ::mlir::success(value_begin<T>());
     }
 
     template <typename T>
     using iterator =
         llvm::mapped_iterator<typename decltype(llvm::seq<ptrdiff_t>(0, 0))::iterator,
                               std::function<T(ptrdiff_t)>>;
 
     /// Return the values of this attribute in the form of the given type 'T'.

mlir/include/mlir/Support/LogicalResult.h

public:
   operator LogicalResult() const { return success(this->has_value()); }
 
 private:
   /// Hide the bool conversion as it easily creates confusion.
   using Optional<T>::operator bool;
   using Optional<T>::has_value;
 };
 
+/// Wrap a value on the success path in a FailureOr of the same value type.
+template <typename T,
+          typename = std::enable_if_t<!std::is_convertible_v<T, bool>>>
+inline auto success(T &&t) {
+  return FailureOr<std::decay_t<T>>(std::forward<T>(t));
+}
+
 /// This class represents success/failure for parsing-like operations that find
 /// it important to chain together failable operations with `||`.  This is an
 /// extended version of `LogicalResult` that allows for explicit conversion to
 /// bool.
 ///
 /// This class should not be used for general error handling cases - we prefer
 /// to keep the logic explicit with the `succeeded`/`failed` predicates.
 /// However, traditional monadic-style parsing logic can sometimes get

mlir/lib/IR/BuiltinAttributes.cpp

if (type.getElementType().isInteger(1)) {
     return emitError() << "expected data size (" << size << " elements, "
                        << type.getElementTypeBitWidth()
                        << " bits each) does not match: " << dataSize
                        << " bytes";
   }
   return success();
 }
 
-const bool *DenseArrayAttr::value_begin_impl(OverloadToken<bool>) const {
-  return cast<DenseBoolArrayAttr>().asArrayRef().begin();
+FailureOr<const bool *>
+DenseArrayAttr::try_value_begin_impl(OverloadToken<bool>) const {
+  if (auto attr = dyn_cast<DenseBoolArrayAttr>())
+    return attr.asArrayRef().begin();
+  return failure();
 }
-const int8_t *DenseArrayAttr::value_begin_impl(OverloadToken<int8_t>) const {
-  return cast<DenseI8ArrayAttr>().asArrayRef().begin();
+FailureOr<const int8_t *>
+DenseArrayAttr::try_value_begin_impl(OverloadToken<int8_t>) const {
+  if (auto attr = dyn_cast<DenseI8ArrayAttr>())
+    return attr.asArrayRef().begin();
+  return failure();
 }
-const int16_t *DenseArrayAttr::value_begin_impl(OverloadToken<int16_t>) const {
-  return cast<DenseI16ArrayAttr>().asArrayRef().begin();
+FailureOr<const int16_t *>
+DenseArrayAttr::try_value_begin_impl(OverloadToken<int16_t>) const {
+  if (auto attr = dyn_cast<DenseI16ArrayAttr>())
+    return attr.asArrayRef().begin();
+  return failure();
 }
-const int32_t *DenseArrayAttr::value_begin_impl(OverloadToken<int32_t>) const {
-  return cast<DenseI32ArrayAttr>().asArrayRef().begin();
+FailureOr<const int32_t *>
+DenseArrayAttr::try_value_begin_impl(OverloadToken<int32_t>) const {
+  if (auto attr = dyn_cast<DenseI32ArrayAttr>())
+    return attr.asArrayRef().begin();
+  return failure();
 }
-const int64_t *DenseArrayAttr::value_begin_impl(OverloadToken<int64_t>) const {
-  return cast<DenseI64ArrayAttr>().asArrayRef().begin();
+FailureOr<const int64_t *>
+DenseArrayAttr::try_value_begin_impl(OverloadToken<int64_t>) const {
+  if (auto attr = dyn_cast<DenseI64ArrayAttr>())
+    return attr.asArrayRef().begin();
+  return failure();
 }
-const float *DenseArrayAttr::value_begin_impl(OverloadToken<float>) const {
-  return cast<DenseF32ArrayAttr>().asArrayRef().begin();
+FailureOr<const float *>
+DenseArrayAttr::try_value_begin_impl(OverloadToken<float>) const {
+  if (auto attr = dyn_cast<DenseF32ArrayAttr>())
+    return attr.asArrayRef().begin();
+  return failure();
 }
-const double *DenseArrayAttr::value_begin_impl(OverloadToken<double>) const {
-  return cast<DenseF64ArrayAttr>().asArrayRef().begin();
+FailureOr<const double *>
+DenseArrayAttr::try_value_begin_impl(OverloadToken<double>) const {
+  if (auto attr = dyn_cast<DenseF64ArrayAttr>())
+    return attr.asArrayRef().begin();
+  return failure();
 }
 
 namespace {
 /// Instantiations of this class provide utilities for interacting with native
 /// data types in the context of DenseArrayAttr.
 template <size_t width,
           IntegerType::SignednessSemantics signedness = IntegerType::Signless>
 struct DenseArrayAttrIntUtil {

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}}
+// expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
 arith.constant array<i1: true, false, true, false, true, false>
 // expected-error@below {{Test iterating `int8_t`: 10, 11, -12, 13, 14}}
+// expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
 arith.constant array<i8: 10, 11, -12, 13, 14>
 // expected-error@below {{Test iterating `int16_t`: 10, 11, -12, 13, 14}}
+// expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
 arith.constant array<i16: 10, 11, -12, 13, 14>
 // expected-error@below {{Test iterating `int32_t`: 10, 11, -12, 13, 14}}
+// expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
 arith.constant array<i32: 10, 11, -12, 13, 14>
 // expected-error@below {{Test iterating `int64_t`: 10, 11, -12, 13, 14}}
 arith.constant array<i64: 10, 11, -12, 13, 14>
 // expected-error@below {{Test iterating `float`: 10.00, 11.00, -12.00, 13.00, 14.00}}
+// expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
 arith.constant array<f32: 10., 11., -12., 13., 14.>
 // expected-error@below {{Test iterating `double`: 10.00, 11.00, -12.00, 13.00, 14.00}}
+// expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
 arith.constant array<f64: 10., 11., -12., 13., 14.>
 
 // Check that we handle an external constant parsed from the config.
 // expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
 // expected-error@below {{Test iterating `uint64_t`: 1, 2, 3}}
 // expected-error@below {{Test iterating `APInt`: unable to iterate type}}
 // expected-error@below {{Test iterating `IntegerAttr`: unable to iterate type}}
 arith.constant #test.e1di64_elements<blob1> : tensor<3xi64>
 
 {-#
   dialect_resources: {
     test: {
       blob1: "0x08000000010000000000000002000000000000000300000000000000"
     }
   }
 #-}

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

def TestI64ElementsAttr : Test_Attr<"TestI64Elements", [
   );
   let extraClassDeclaration = [{
     /// The set of data types that can be iterated by this attribute.
     using ContiguousIterableTypesT = std::tuple<uint64_t>;
     using NonContiguousIterableTypesT = std::tuple<mlir::Attribute, llvm::APInt>;
 
     /// Provide begin iterators for the various iterable types.
     // * uint64_t
-    auto value_begin_impl(OverloadToken<uint64_t>) const {
+    mlir::FailureOr<const uint64_t *>
+    try_value_begin_impl(OverloadToken<uint64_t>) const {
       return getElements().begin();
     }
     // * Attribute
-    auto value_begin_impl(OverloadToken<mlir::Attribute>) const {
+    auto try_value_begin_impl(OverloadToken<mlir::Attribute>) const {
       mlir::Type elementType = getType().getElementType();
-      return llvm::map_range(getElements(), [=](uint64_t value) {
+      return mlir::success(llvm::map_range(getElements(), [=](uint64_t value) {
         return mlir::IntegerAttr::get(elementType,
                                       llvm::APInt(/*numBits=*/64, value));
-      }).begin();
+      }).begin());
     }
     // * APInt
-    auto value_begin_impl(OverloadToken<llvm::APInt>) const {
-      return llvm::map_range(getElements(), [=](uint64_t value) {
+    auto try_value_begin_impl(OverloadToken<llvm::APInt>) const {
+      return mlir::success(llvm::map_range(getElements(), [=](uint64_t value) {
         return llvm::APInt(/*numBits=*/64, value);
-      }).begin();
+      }).begin());
     }
   }];
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 def TestSubElementsAccessAttr : Test_Attr<"TestSubElementsAccess", [
     DeclareAttrInterfaceMethods<SubElementAttrInterface,
let extraClassDeclaration = [{
     /// Return the elements referenced by this attribute.
     llvm::ArrayRef<uint64_t> getElements() const;
 
     /// The set of data types that can be iterated by this attribute.
     using ContiguousIterableTypesT = std::tuple<uint64_t>;
 
     /// Provide begin iterators for the various iterable types.
     // * uint64_t
-    auto value_begin_impl(OverloadToken<uint64_t>) const {
+    mlir::FailureOr<const uint64_t *>
+    try_value_begin_impl(OverloadToken<uint64_t>) const {
       return getElements().begin();
     }
   }];
   let assemblyFormat = "`<` $handle `>`";
 }
 
 #endif // TEST_ATTRDEFS

mlir/test/lib/IR/TestBuiltinAttributeInterfaces.cpp

getOperation().walk([&](Operation *op) {
                 testElementsAttrIteration<int64_t>(op, attr, "int64_t");
               })
               .Case([&](DenseF32ArrayAttr attr) {
                 testElementsAttrIteration<float>(op, attr, "float");
               })
               .Case([&](DenseF64ArrayAttr attr) {
                 testElementsAttrIteration<double>(op, attr, "double");
               });
+          testElementsAttrIteration<int64_t>(op, elementsAttr, "int64_t");
           continue;
         }
         testElementsAttrIteration<int64_t>(op, elementsAttr, "int64_t");
         testElementsAttrIteration<uint64_t>(op, elementsAttr, "uint64_t");
         testElementsAttrIteration<APInt>(op, elementsAttr, "APInt");
         testElementsAttrIteration<IntegerAttr>(op, elementsAttr, "IntegerAttr");
       }
     });