[mlir] Fix DenseArrayAttr's ElementsAttr implementation

[Mogball]

Summary

Previously, the dispatch to subclasses was unchecked, resulting in
tryGetValues<int32_t> on a DenseF32ArrayAttr asserting instead of
resulting failure(). This changes the implementation of
getValuesImpl to do dynamic dispatch to the subclasses, failing where
appropriate.

Depends on D132758

[rriddle]

Users are not supposed to override this. The long standing TODO has been to change ElementsAttr to check for a failable version.

[rriddle]

More specifically, what I've wanted to do (but haven't had time) is to change the value_begin_impl to return FailureOr instead of the iterator directly (and rename to try_value_begin_impl or something). That is described here: https://github.com/llvm/llvm-project/blob/3155e3070c49b31b8d189f5f6caf376fcb82fcfb/mlir/include/mlir/IR/BuiltinAttributeInterfaces.td#L57

[Mogball]

I can make the change. You just mean to change iterator value_begin_impl to FailureOr<iterator> try_value_begin_impl?

[rriddle]

In D132759#3752796, @Mogball wrote:

I can make the change. You just mean to change iterator value_begin_impl to FailureOr<iterator> try_value_begin_impl?

Essentially, yeah.

[Mogball]

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

[Mogball]

This has been superseded by https://reviews.llvm.org/D132904

mlir/include/mlir/IR/BuiltinAttributes.td

if (!$_self.empty()) {
       auto *alloc = static_cast<char *>(
           $_allocator.allocate($_self.size(), alignof(uint64_t)));
       std::uninitialized_copy($_self.begin(), $_self.end(), alloc);
       $_dst = ArrayRef<char>(alloc, $_self.size());
     }
   }];
 }
 
-def Builtin_DenseArray : Builtin_Attr<
-    "DenseArray", [ElementsAttrInterface, TypedAttrInterface]> {
+def Builtin_DenseArray : Builtin_Attr<"DenseArray", [
+    DeclareAttrInterfaceMethods<ElementsAttrInterface,
+                                ["getValuesImpl", "isSplat"]>,
+    TypedAttrInterface
+  ]> {
   let summary = "A dense array of integer or floating point elements.";
   let description = [{
     A dense array attribute is an attribute that represents a dense array of
     primitive element types. Contrary to DenseIntOrFPElementsAttr this is a
     flat unidimensional array which does not have a storage optimization for
     splat. This allows to expose the raw array through a C++ API as
     `ArrayRef<T>` for compatible types. The element type must be bool or an
     integer or float whose bitwidth is a multiple of 8. Bool elements are stored
def Builtin_DenseArray : Builtin_Attr<"DenseArray", [
 
   let builders = [
     AttrBuilderWithInferredContext<(ins "RankedTensorType":$type,
                                         "ArrayRef<char>":$rawData), [{
       return $_get(type.getContext(), type, rawData);
     }]>,
   ];
 
-  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;
-  }];
-
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DenseIntOrFPElementsAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_DenseIntOrFPElementsAttr : Builtin_Attr<

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();
-}
-const int8_t *DenseArrayAttr::value_begin_impl(OverloadToken<int8_t>) const {
-  return cast<DenseI8ArrayAttr>().asArrayRef().begin();
-}
-const int16_t *DenseArrayAttr::value_begin_impl(OverloadToken<int16_t>) const {
-  return cast<DenseI16ArrayAttr>().asArrayRef().begin();
-}
-const int32_t *DenseArrayAttr::value_begin_impl(OverloadToken<int32_t>) const {
-  return cast<DenseI32ArrayAttr>().asArrayRef().begin();
-}
-const int64_t *DenseArrayAttr::value_begin_impl(OverloadToken<int64_t>) const {
-  return cast<DenseI64ArrayAttr>().asArrayRef().begin();
-}
-const float *DenseArrayAttr::value_begin_impl(OverloadToken<float>) const {
-  return cast<DenseF32ArrayAttr>().asArrayRef().begin();
+//===----------------------------------------------------------------------===//
+// ElementsAttr Implementation
+
+/// A dense array is never a splat.
+bool DenseArrayAttr::isSplat() const { return false; }
+
+/// Try to dispatch to the given dense array subclass to contiguously iterate
+/// its elements.
+template <typename DenseArrayT>
+static FailureOr<detail::ElementsAttrIndexer>
+tryContainerBegin(DenseArrayAttr attr) {
+  if (auto container = attr.dyn_cast<DenseArrayT>())
+    return detail::ElementsAttrIndexer::contiguous(
+        /*isSplat=*/false, container.asArrayRef().begin());
+  return failure();
 }
-const double *DenseArrayAttr::value_begin_impl(OverloadToken<double>) const {
-  return cast<DenseF64ArrayAttr>().asArrayRef().begin();
+
+/// Re-implement `getValuesImpl` to dynamically dispatch to subclasses.
+FailureOr<detail::ElementsAttrIndexer>
+DenseArrayAttr::getValuesImpl(TypeID elementID) const {
+  if (elementID == TypeID::get<bool>())
+    return tryContainerBegin<DenseBoolArrayAttr>(*this);
+  if (elementID == TypeID::get<int8_t>())
+    return tryContainerBegin<DenseI8ArrayAttr>(*this);
+  if (elementID == TypeID::get<int16_t>())
+    return tryContainerBegin<DenseI16ArrayAttr>(*this);
+  if (elementID == TypeID::get<int32_t>())
+    return tryContainerBegin<DenseI32ArrayAttr>(*this);
+  if (elementID == TypeID::get<int64_t>())
+    return tryContainerBegin<DenseI64ArrayAttr>(*this);
+  if (elementID == TypeID::get<float>())
+    return tryContainerBegin<DenseF32ArrayAttr>(*this);
+  if (elementID == TypeID::get<double>())
+    return tryContainerBegin<DenseF64ArrayAttr>(*this);
+  return failure();
 }
 
+//===----------------------------------------------------------------------===//
+// DenseArrayAttrUtil
+
 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 {
   static bool checkElementType(Type eltType) {
     auto type = eltType.dyn_cast<IntegerType>();
struct DenseArrayAttrUtil<double> {
   static Type getElementType(MLIRContext *ctx) { return Float64Type::get(ctx); }
   static void printElement(raw_ostream &os, float value) { os << value; }
   static ParseResult parseElement(AsmParser &parser, double &value) {
     return parser.parseFloat(value);
   }
 };
 } // namespace
 
+//===----------------------------------------------------------------------===//
+// DenseArrayAttrImpl
+
 template <typename T>
 void DenseArrayAttrImpl<T>::print(AsmPrinter &printer) const {
   print(printer.getStream());
 }
 
 template <typename T>
 void DenseArrayAttrImpl<T>::printWithoutBraces(raw_ostream &os) const {
   llvm::interleaveComma(asArrayRef(), os, [&](T value) {