[llvm] Make Sequence reverse-iterable

[gchatelet]

Summary

This is a roll forward of D102679.
This patch simplifies the implementation of Sequence and makes it compatible with llvm::reverse.
It exposes the reverse iterators through rbegin/rend which prevents a dangling reference in std::reverse_iterator::operator++().

Note: Compared to D102679, this patch introduces a asSmallVector() member function and fixes compilation issue with GCC 5.

[gchatelet]

llvm/include/llvm/ADT/Sequence.h

17–18

@Quuxplusone @mehdi_amini Introducing the asSmallVector() method add a dependency in Sequence.
I'm having a hard time evaluating the costs/benefits balance of this move.

178

@Quuxplusone I will remove the std::move in this patch (cf. D103900)

mlir/include/mlir/IR/BuiltinAttributes.td

711

The purpose of this change was to prevent mlir from depending on implementation details of llvm.

@mehdi_amini I've added typename to disambiguate between value and type for GCC5 but I have no way of checking whether that works or not (GCC5 is not available on my machine).
Do you have any suggestions on how to test it?

[gchatelet]

• Add missing header for std::ptrdiff_t
• Remove unnecessary std::move

[• Quuxplusone]

llvm/include/llvm/ADT/Sequence.h

153

Nit: Just std::ptrdiff_t seems simpler for the reader.

159–164

I still believe this should be simply

explicit iota_range(ValueT Begin, ValueT End)
      : Begin(Begin), End(End) {}

with no templates. (For the same reason as lines 182-183.)
And explicit because all constructors should be explicit.

llvm/unittests/ADT/SequenceTest.cpp

42–48

/Dereferene/Dereference/

mlir/include/mlir/IR/BuiltinAttributes.td

711

Indeed the error message looked to me like the symptom of "forgetting to include the right header somewhere," because the compiler seemed to be complaining that it didn't know what ptrdiff_t meant. Adding typename here also LGTM.
Nit: const_iterator could be just iterator, right? There's no difference between const-iterating and regular-iterating an llvm::seq, so the difference between them is nonexistent and we might as well use the shorter simpler name.

[mehdi_amini]

mlir/include/mlir/IR/BuiltinAttributes.td

711

I usually use a docker container to test this, but feel free to land it, we can always revert it if it breaks a bot again.

[gchatelet]

• Rebase and address comments

llvm/include/llvm/ADT/Sequence.h

 //===- Sequence.h - Utility for producing sequences of values ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file
 /// This routine provides some synthesis utilities to produce sequences of
 /// values. The names are intentionally kept very short as they tend to occur
 /// in common and widely used contexts.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ADT_SEQUENCE_H
 #define LLVM_ADT_SEQUENCE_H
 
-#include "llvm/ADT/iterator.h"
-#include "llvm/ADT/iterator_range.h"
-#include <algorithm>
-#include <iterator>
-#include <utility>
+#include "llvm/ADT/SmallVector.h"
+
+#include <cstddef>  //std::ptrdiff_t
+#include <iterator> //std::random_access_iterator_tag
 
 namespace llvm {
 
 namespace detail {
 
-template <typename ValueT>
-class value_sequence_iterator
-    : public iterator_facade_base<value_sequence_iterator<ValueT>,
-                                  std::random_access_iterator_tag,
-                                  const ValueT> {
-  using BaseT = typename value_sequence_iterator::iterator_facade_base;
+template <typename T, bool IsReversed> struct iota_range_iterator {
+  using iterator_category = std::random_access_iterator_tag;
+  using value_type = T;
+  using difference_type = std::ptrdiff_t;
+  using pointer = T *;
+  using reference = T &;
+
+private:
+  struct Forward {
+    static void increment(T &V) { ++V; }
+    static void decrement(T &V) { --V; }
+    static void offset(T &V, difference_type Offset) { V += Offset; }
+    static T add(const T &V, difference_type Offset) { return V + Offset; }
+    static difference_type difference(const T &A, const T &B) { return A - B; }
+  };
+
+  struct Reverse {
+    static void increment(T &V) { --V; }
+    static void decrement(T &V) { ++V; }
+    static void offset(T &V, difference_type Offset) { V -= Offset; }
+    static T add(const T &V, difference_type Offset) { return V - Offset; }
+    static difference_type difference(const T &A, const T &B) { return B - A; }
+  };
 
-  ValueT Value;
+  using Op = std::conditional_t<!IsReversed, Forward, Reverse>;
 
 public:
-  using difference_type = typename BaseT::difference_type;
-  using reference = typename BaseT::reference;
+  // default-constructible
+  iota_range_iterator() = default;
+  // copy-constructible
+  iota_range_iterator(const iota_range_iterator &) = default;
+  // value constructor
+  explicit iota_range_iterator(T Value) : Value(Value) {}
+  // copy-assignable
+  iota_range_iterator &operator=(const iota_range_iterator &) = default;
+  // destructible
+  ~iota_range_iterator() = default;
+
+  // Can be compared for equivalence using the equality/inequality operators,
+  bool operator!=(const iota_range_iterator &RHS) const {
+    return Value != RHS.Value;
+  }
+  bool operator==(const iota_range_iterator &RHS) const {
+    return Value == RHS.Value;
+  }
 
-  value_sequence_iterator() = default;
-  value_sequence_iterator(const value_sequence_iterator &) = default;
-  value_sequence_iterator(value_sequence_iterator &&Arg)
-      : Value(std::move(Arg.Value)) {}
-  value_sequence_iterator &operator=(const value_sequence_iterator &Arg) {
-    Value = Arg.Value;
-    return *this;
+  // Comparison
+  bool operator<(const iota_range_iterator &Other) const {
+    return Op::difference(Value, Other.Value) < 0;
   }
+  bool operator<=(const iota_range_iterator &Other) const {
+    return Op::difference(Value, Other.Value) <= 0;
+  }
+  bool operator>(const iota_range_iterator &Other) const {
+    return Op::difference(Value, Other.Value) > 0;
+  }
+  bool operator>=(const iota_range_iterator &Other) const {
+    return Op::difference(Value, Other.Value) >= 0;
+  }
+
+  // Dereference
+  T operator*() const { return Value; }
+  T operator[](difference_type Offset) const { return Op::add(Value, Offset); }
 
-  template <typename U, typename Enabler = decltype(ValueT(std::declval<U>()))>
-  value_sequence_iterator(U &&Value) : Value(std::forward<U>(Value)) {}
+  // Arithmetic
+  iota_range_iterator operator+(difference_type Offset) const {
+    return {Op::add(Value, Offset)};
+  }
+  iota_range_iterator operator-(difference_type Offset) const {
+    return {Op::add(Value, -Offset)};
+  }
+
+  // Iterator difference
+  difference_type operator-(const iota_range_iterator &Other) const {
+    return Op::difference(Value, Other.Value);
+  }
 
-  value_sequence_iterator &operator+=(difference_type N) {
-    Value += N;
+  // Pre/Post Increment
+  iota_range_iterator &operator++() {
+    Op::increment(Value);
     return *this;
   }
-  value_sequence_iterator &operator-=(difference_type N) {
-    Value -= N;
+  iota_range_iterator operator++(int) {
+    iota_range_iterator Tmp = *this;
+    Op::increment(Value);
+    return Tmp;
+  }
+
+  // Pre/Post Decrement
+  iota_range_iterator &operator--() {
+    Op::decrement(Value);
     return *this;
   }
-  using BaseT::operator-;
-  difference_type operator-(const value_sequence_iterator &RHS) const {
-    return Value - RHS.Value;
+  iota_range_iterator operator--(int) {
+    iota_range_iterator Tmp = *this;
+    Op::decrement(Value);
+    return Tmp;
   }
 
-  bool operator==(const value_sequence_iterator &RHS) const {
-    return Value == RHS.Value;
+  // Compound assignment operators
+  iota_range_iterator &operator+=(difference_type Offset) {
+    Op::offset(Value, Offset);
+    return *this;
   }
-  bool operator<(const value_sequence_iterator &RHS) const {
-    return Value < RHS.Value;
+  iota_range_iterator &operator-=(difference_type Offset) {
+    Op::offset(Value, -Offset);
+    return *this;
   }
 
-  reference operator*() const { return Value; }
+private:
+  T Value;
 };
 
-} // end namespace detail
+} // namespace detail
+
+template <typename ValueT> struct iota_range {
+  static_assert(std::is_integral<ValueT>::value,
+                "ValueT must be an integral type");
+
+  using value_type = ValueT;
+  using reference = ValueT &;
+  using const_reference = const ValueT &;
+  using iterator = detail::iota_range_iterator<value_type, false>;
+  using const_iterator = iterator;
+  using reverse_iterator = detail::iota_range_iterator<value_type, true>;
+  using const_reverse_iterator = reverse_iterator;
+  using difference_type = std::ptrdiff_t;
+  using size_type = std::size_t;
+
+  value_type Begin;
+  value_type End;
+
+  explicit iota_range(ValueT Begin, ValueT End) : Begin(Begin), End(End) {}
+
+  size_t size() const { return End - Begin; }
+  bool empty() const { return Begin == End; }
+
+  auto begin() const { return const_iterator(Begin); }
+  auto end() const { return const_iterator(End); }
+
+  auto rbegin() const { return const_reverse_iterator(End - 1); }
+  auto rend() const { return const_reverse_iterator(Begin - 1); }
+
+  auto asSmallVector() const { return SmallVector<ValueT>(begin(), end()); }
+
+private:
+  static_assert(std::is_same<ValueT, std::remove_cv_t<ValueT>>::value,
+                "ValueT must not be const nor volatile");
+};
 
-template <typename ValueT>
-iterator_range<detail::value_sequence_iterator<ValueT>> seq(ValueT Begin,
-                                                            ValueT End) {
-  return make_range(detail::value_sequence_iterator<ValueT>(Begin),
-                    detail::value_sequence_iterator<ValueT>(End));
+template <typename ValueT> auto seq(ValueT Begin, ValueT End) {
+  return iota_range<ValueT>(Begin, End);
 }
 
 } // end namespace llvm
 
 #endif // LLVM_ADT_SEQUENCE_H

llvm/unittests/ADT/SequenceTest.cpp

 #include "gtest/gtest.h"
 
 #include <list>
 
 using namespace llvm;
 
 namespace {
 
-TEST(SequenceTest, Basic) {
-  int x = 0;
-  for (int i : seq(0, 10)) {
-    EXPECT_EQ(x, i);
-    x++;
-  }
-  EXPECT_EQ(10, x);
-
-  auto my_seq = seq(0, 4);
-  EXPECT_EQ(4, my_seq.end() - my_seq.begin());
-  for (int i : {0, 1, 2, 3})
-    EXPECT_EQ(i, (int)my_seq.begin()[i]);
-
-  EXPECT_TRUE(my_seq.begin() < my_seq.end());
-
-  auto adjusted_begin = my_seq.begin() + 2;
-  auto adjusted_end = my_seq.end() - 2;
-  EXPECT_TRUE(adjusted_begin == adjusted_end);
-  EXPECT_EQ(2, *adjusted_begin);
-  EXPECT_EQ(2, *adjusted_end);
+TEST(SequenceTest, Forward) {
+  int X = 0;
+  for (int I : seq(0, 10)) {
+    EXPECT_EQ(X, I);
+    ++X;
+  }
+  EXPECT_EQ(10, X);
+}
+
+TEST(SequenceTest, Backward) {
+  int X = 9;
+  for (int I : reverse(seq(0, 10))) {
+    EXPECT_EQ(X, I);
+    --X;
+  }
+  EXPECT_EQ(-1, X);
+}
+
+TEST(SequenceTest, Distance) {
+  const auto Forward = seq(0, 10);
+  EXPECT_EQ(std::distance(Forward.begin(), Forward.end()), 10);
+  EXPECT_EQ(std::distance(Forward.rbegin(), Forward.rend()), 10);
+}
+
+TEST(SequenceTest, Dereference) {
+  const auto Forward = seq(0, 10).begin();
+  EXPECT_EQ(Forward[0], 0);
+  EXPECT_EQ(Forward[2], 2);
+  const auto Backward = seq(0, 10).rbegin();
+  EXPECT_EQ(Backward[0], 9);
+  EXPECT_EQ(Backward[2], 7);
 }
 
 } // anonymous namespace

mlir/include/mlir/IR/BuiltinAttributes.td

AttrBuilderWithInferredContext<(ins "ShapedType":$type,
       assert(type.hasStaticShape() && "type must have static shape");
       return $_get(type.getContext(), type,
                    indices.cast<DenseIntElementsAttr>(), values);
     }]>,
   ];
   let extraClassDeclaration = [{
     template <typename T>
     using iterator =
-        llvm::mapped_iterator<llvm::detail::value_sequence_iterator<ptrdiff_t>,
+        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'.
     /// 'T'  may be any of Attribute, APInt, APFloat, c++ integer/float types,
     /// etc.
     template <typename T> llvm::iterator_range<iterator<T>> getValues() const;
 
     /// Return the value of the element at the given index. The 'index' is

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

void PatternLowering::generateSwitch(SwitchNode *switchNode,
                                      Block *currentBlock, Qualifier *question,
                                      Value val, Block *defaultDest) {
   // If the switch question is not an exact answer, i.e. for the `at_least`
   // cases, we generate a special block sequence.
   Predicates::Kind kind = question->getKind();
   if (kind == Predicates::OperandCountAtLeastQuestion ||
       kind == Predicates::ResultCountAtLeastQuestion) {
     // Order the children such that the cases are in reverse numerical order.
-    SmallVector<unsigned> sortedChildren(
-        llvm::seq<unsigned>(0, switchNode->getChildren().size()));
+    SmallVector<unsigned> sortedChildren =
+        llvm::seq<unsigned>(0, switchNode->getChildren().size())
+            .asSmallVector();
     llvm::sort(sortedChildren, [&](unsigned lhs, unsigned rhs) {
       return cast<UnsignedAnswer>(switchNode->getChild(lhs).first)->getValue() >
              cast<UnsignedAnswer>(switchNode->getChild(rhs).first)->getValue();
     });
 
     // Build the destination for each child using the next highest child as a
     // a failure destination. This essentially creates the following control
     // flow:

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

static VectorizationResult vectorizeLinalgIndex(OpBuilder &b, Operation *op,
                                                 LinalgOp linalgOp) {
   IndexOp indexOp = dyn_cast<linalg::IndexOp>(op);
   if (!indexOp)
     return VectorizationResult{VectorizationStatus::Failure, nullptr};
   auto loc = indexOp.getLoc();
   // Compute the static loop sizes of the index op.
   auto targetShape = linalgOp.computeStaticLoopSizes();
   // Compute a one-dimensional index vector for the index op dimension.
-  SmallVector<int64_t> constantSeq(
-      llvm::seq<int64_t>(0, targetShape[indexOp.dim()]));
+  SmallVector<int64_t> constantSeq =
+      llvm::seq<int64_t>(0, targetShape[indexOp.dim()]).asSmallVector();
   ConstantOp constantOp =
       b.create<ConstantOp>(loc, b.getIndexVectorAttr(constantSeq));
   // Return the one-dimensional index vector if it lives in the trailing
   // dimension of the iteration space since the vectorization algorithm in this
   // case can handle the broadcast.
   if (indexOp.dim() == targetShape.size() - 1)
     return VectorizationResult{VectorizationStatus::NewOp, constantOp};
   // Otherwise permute the targetShape to move the index dimension last,
   // broadcast the one-dimensional index vector to the permuted shape, and
   // finally transpose the broadcasted index vector to undo the permutation.
   std::swap(targetShape[indexOp.dim()], targetShape.back());
   auto broadCastOp = b.create<vector::BroadcastOp>(
       loc, VectorType::get(targetShape, b.getIndexType()), constantOp);
-  SmallVector<int64_t> transposition(
-      llvm::seq<int64_t>(0, linalgOp.getNumLoops()));
+  SmallVector<int64_t> transposition =
+      llvm::seq<int64_t>(0, linalgOp.getNumLoops()).asSmallVector();
   std::swap(transposition.back(), transposition[indexOp.dim()]);
   auto transposeOp =
       b.create<vector::TransposeOp>(loc, broadCastOp, transposition);
   return VectorizationResult{VectorizationStatus::NewOp, transposeOp};
 }
 
 /// Generic vectorization for a single operation `op`, given already vectorized
 /// operands carried by `bvm`. Vectorization occurs as follows: