[mlir] Refactor ElementsAttr's value access API

There are several aspects of the API that either aren't easy to use, or are
deceptively easy to do the wrong thing. The main change of this commit
is to remove all of the `getValue<T>`/`getFlatValue<T>` from ElementsAttr
and instead provide operator[] methods on the ranges returned by
`getValues<T>`. This provides a much more convenient API for the value
ranges. It also removes the easy-to-be-inefficient nature of
getValue/getFlatValue, which under the hood would construct a new range for
the type `T`. Constructing a range is not necessarily cheap in all cases, and
could lead to very poor performance if used within a loop; i.e. if you were to
naively write something like:

```
DenseElementsAttr attr = ...;
for (int i = 0; i < size; ++i) {
  // We are internally rebuilding the APFloat value range on each iteration!!
  APFloat it = attr.getFlatValue<APFloat>(i);
}
```

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

Author: River707

mlir/include/mlir/IR/BuiltinAttributeInterfaces.h

@@ -227,6 +227,33 @@ class ElementsAttrIterator
   ElementsAttrIndexer indexer;
   ptrdiff_t index;
 };
+
+/// This class provides iterator utilities for an ElementsAttr range.
+template <typename IteratorT>
+class ElementsAttrRange : public llvm::iterator_range<IteratorT> {
+public:
+  using reference = typename IteratorT::reference;
+
+  ElementsAttrRange(Type shapeType,
+                    const llvm::iterator_range<IteratorT> &range)
+      : llvm::iterator_range<IteratorT>(range), shapeType(shapeType) {}
+  ElementsAttrRange(Type shapeType, IteratorT beginIt, IteratorT endIt)
+      : ElementsAttrRange(shapeType, llvm::make_range(beginIt, endIt)) {}
+
+  /// Return the value at the given index.
+  reference operator[](ArrayRef<uint64_t> index) const;
+  reference operator[](uint64_t index) const {
+    return *std::next(this->begin(), index);
+  }
+
+  /// Return the size of this range.
+  size_t size() const { return llvm::size(*this); }
+
+private:
+  /// The shaped type of the parent ElementsAttr.
+  Type shapeType;
+};
+
 } // namespace detail
 
 //===----------------------------------------------------------------------===//
@@ -256,6 +283,16 @@ verifyAffineMapAsLayout(AffineMap m, ArrayRef<int64_t> shape,
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
+namespace detail {
+/// Return the value at the given index.
+template <typename IteratorT>
+auto ElementsAttrRange<IteratorT>::operator[](ArrayRef<uint64_t> index) const
+    -> reference {
+  // Skip to the element corresponding to the flattened index.
+  return (*this)[ElementsAttr::getFlattenedIndex(shapeType, index)];
+}
+} // namespace detail
+
 /// Return the elements of this attribute as a value of type 'T'.
 template <typename T>
 auto ElementsAttr::value_begin() const -> DefaultValueCheckT<T, iterator<T>> {

mlir/include/mlir/IR/BuiltinAttributeInterfaces.td

@@ -158,27 +158,6 @@ def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
   ];
 
   string ElementsAttrInterfaceAccessors = [{
-    /// Return the attribute value at the given index. The index is expected to
-    /// refer to a valid element.
-    Attribute getValue(ArrayRef<uint64_t> index) const {
-      return getValue<Attribute>(index);
-    }
-
-    /// Return the value of type 'T' at the given index, where 'T' corresponds
-    /// to an Attribute type.
-    template <typename T>
-    std::enable_if_t<!std::is_same<T, ::mlir::Attribute>::value &&
-                     std::is_base_of<T, ::mlir::Attribute>::value>
-    getValue(ArrayRef<uint64_t> index) const {
-      return getValue(index).template dyn_cast_or_null<T>();
-    }
-
-    /// Return the value of type 'T' at the given index.
-    template <typename T>
-    T getValue(ArrayRef<uint64_t> index) const {
-      return getFlatValue<T>(getFlattenedIndex(index));
-    }
-
     /// Return the number of elements held by this attribute.
     int64_t size() const { return getNumElements(); }
 
@@ -281,6 +260,14 @@ def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
     // 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>
@@ -292,19 +279,16 @@ def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
     template <typename T>
     auto getValues() const {
       auto beginIt = $_attr.template value_begin<T>();
-      return llvm::make_range(beginIt, std::next(beginIt, size()));
-    }
-    /// Return the value at the given flattened index.
-    template <typename T> T getFlatValue(uint64_t index) const {
-      return *std::next($_attr.template value_begin<T>(), index);
+      return detail::ElementsAttrRange<decltype(beginIt)>(
+        Attribute($_attr).getType(), beginIt, std::next(beginIt, size()));
     }
   }] # ElementsAttrInterfaceAccessors;
 
   let extraClassDeclaration = [{
     template <typename T>
     using iterator = detail::ElementsAttrIterator<T>;
     template <typename T>
-    using iterator_range = llvm::iterator_range<iterator<T>>;
+    using iterator_range = detail::ElementsAttrRange<iterator<T>>;
 
     //===------------------------------------------------------------------===//
     // Accessors
@@ -329,8 +313,12 @@ def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
     uint64_t getFlattenedIndex(ArrayRef<uint64_t> index) const {
       return getFlattenedIndex(*this, index);
     }
-    static uint64_t getFlattenedIndex(Attribute elementsAttr,
+    static uint64_t getFlattenedIndex(Type type,
                                       ArrayRef<uint64_t> index);
+    static uint64_t getFlattenedIndex(Attribute elementsAttr,
+                                      ArrayRef<uint64_t> index) {
+      return getFlattenedIndex(elementsAttr.getType(), index);
+    }
 
     /// Returns the number of elements held by this attribute.
     int64_t getNumElements() const { return getNumElements(*this); }
@@ -350,13 +338,6 @@ def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
                                   !std::is_base_of<Attribute, T>::value,
                                   ResultT>;
 
-    /// Return the element of this attribute at the given index as a value of
-    /// type 'T'.
-    template <typename T>
-    T getFlatValue(uint64_t index) const {
-      return *std::next(value_begin<T>(), index);
-    }
-
     /// Return the splat value for this attribute. This asserts that the
     /// attribute corresponds to a splat.
     template <typename T>
@@ -368,7 +349,7 @@ def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
     /// Return the elements of this attribute as a value of type 'T'.
     template <typename T>
     DefaultValueCheckT<T, iterator_range<T>> getValues() const {
-      return iterator_range<T>(value_begin<T>(), value_end<T>());
+      return {Attribute::getType(), value_begin<T>(), value_end<T>()};
     }
     template <typename T>
     DefaultValueCheckT<T, iterator<T>> value_begin() const;
@@ -384,12 +365,12 @@ def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
       llvm::mapped_iterator<iterator<Attribute>, T (*)(Attribute)>;
     template <typename T>
     using DerivedAttrValueIteratorRange =
-      llvm::iterator_range<DerivedAttrValueIterator<T>>;
+      detail::ElementsAttrRange<DerivedAttrValueIterator<T>>;
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     DerivedAttrValueIteratorRange<T> getValues() const {
       auto castFn = [](Attribute attr) { return attr.template cast<T>(); };
-      return llvm::map_range(getValues<Attribute>(),
-                             static_cast<T (*)(Attribute)>(castFn));
+      return {Attribute::getType(), llvm::map_range(getValues<Attribute>(),
+                             static_cast<T (*)(Attribute)>(castFn))};
     }
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     DerivedAttrValueIterator<T> value_begin() const {
@@ -407,8 +388,10 @@ def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
     /// return the iterable range. Otherwise, return llvm::None.
     template <typename T>
     DefaultValueCheckT<T, Optional<iterator_range<T>>> tryGetValues() const {
-      if (Optional<iterator<T>> beginIt = try_value_begin<T>())
-        return iterator_range<T>(*beginIt, value_end<T>());
+      if (Optional<iterator<T>> beginIt = try_value_begin<T>()) {
+        return iterator_range<T>(Attribute::getType(), *beginIt,
+                                 value_end<T>());
+      }
       return llvm::None;
     }
     template <typename T>
@@ -418,10 +401,15 @@ def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
     /// return the iterable range. Otherwise, return llvm::None.
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     Optional<DerivedAttrValueIteratorRange<T>> tryGetValues() const {
+      auto values = tryGetValues<Attribute>();
+      if (!values)
+        return llvm::None;
+
       auto castFn = [](Attribute attr) { return attr.template cast<T>(); };
-      if (auto values = tryGetValues<Attribute>())
-        return llvm::map_range(*values, static_cast<T (*)(Attribute)>(castFn));
-      return llvm::None;
+      return DerivedAttrValueIteratorRange<T>(
+        Attribute::getType(),
+        llvm::map_range(*values, static_cast<T (*)(Attribute)>(castFn))
+      );
     }
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     Optional<DerivedAttrValueIterator<T>> try_value_begin() const {