[llvm] ADT/Matrix: two-dimensional Container with View (PR #98893)

via llvm-commits llvm-commits at lists.llvm.org
Mon Jul 15 05:56:07 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-adt

Author: Ramkumar Ramachandra (artagnon)

<details>
<summary>Changes</summary>

Leverage the excellent SmallVector infrastructure to write a two-dimensional container, along with a View that abstracts out indexing-arithmetic, eliminating memory operations on the underlying storage. The immediate applicability of Matrix is to replace uses of the vector-of-vectors idiom, with one caveat: an upper bound on the number of columns should be known ahead of time.

-- 8< --
Based on #<!-- -->98874.

---

Patch is 33.00 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/98893.diff


4 Files Affected:

- (modified) llvm/include/llvm/ADT/ArrayRef.h (+56-44) 
- (added) llvm/include/llvm/ADT/Matrix.h (+339) 
- (modified) llvm/unittests/ADT/CMakeLists.txt (+1) 
- (added) llvm/unittests/ADT/MatrixTest.cpp (+325) 


``````````diff
diff --git a/llvm/include/llvm/ADT/ArrayRef.h b/llvm/include/llvm/ADT/ArrayRef.h
index 1c6799f1c56ed..6fc7acff7ba8c 100644
--- a/llvm/include/llvm/ADT/ArrayRef.h
+++ b/llvm/include/llvm/ADT/ArrayRef.h
@@ -25,6 +25,7 @@
 
 namespace llvm {
   template<typename T> class [[nodiscard]] MutableArrayRef;
+  template <typename T> struct [[nodiscard]] MatrixRowView;
 
   /// ArrayRef - Represent a constant reference to an array (0 or more elements
   /// consecutively in memory), i.e. a start pointer and a length.  It allows
@@ -59,19 +60,21 @@ namespace llvm {
     /// The number of elements.
     size_type Length = 0;
 
+    friend MatrixRowView<T>;
+
   public:
     /// @name Constructors
     /// @{
 
     /// Construct an empty ArrayRef.
-    /*implicit*/ ArrayRef() = default;
+    /*implicit*/ constexpr ArrayRef() = default;
 
     /// Construct an empty ArrayRef from std::nullopt.
-    /*implicit*/ ArrayRef(std::nullopt_t) {}
+    /*implicit*/ constexpr ArrayRef(std::nullopt_t) {}
 
     /// Construct an ArrayRef from a single element.
-    /*implicit*/ ArrayRef(const T &OneElt)
-      : Data(&OneElt), Length(1) {}
+    /*implicit*/ constexpr ArrayRef(const T &OneElt)
+        : Data(&OneElt), Length(1) {}
 
     /// Construct an ArrayRef from a pointer and length.
     constexpr /*implicit*/ ArrayRef(const T *data, size_t length)
@@ -123,9 +126,10 @@ namespace llvm {
     /// Construct an ArrayRef<const T*> from ArrayRef<T*>. This uses SFINAE to
     /// ensure that only ArrayRefs of pointers can be converted.
     template <typename U>
-    ArrayRef(const ArrayRef<U *> &A,
-             std::enable_if_t<std::is_convertible<U *const *, T const *>::value>
-                 * = nullptr)
+    constexpr ArrayRef(
+        const ArrayRef<U *> &A,
+        std::enable_if_t<std::is_convertible<U *const *, T const *>::value> * =
+            nullptr)
         : Data(A.data()), Length(A.size()) {}
 
     /// Construct an ArrayRef<const T*> from a SmallVector<T*>. This is
@@ -150,28 +154,32 @@ namespace llvm {
     /// @name Simple Operations
     /// @{
 
-    iterator begin() const { return Data; }
-    iterator end() const { return Data + Length; }
+    constexpr iterator begin() const { return Data; }
+    constexpr iterator end() const { return Data + Length; }
 
-    reverse_iterator rbegin() const { return reverse_iterator(end()); }
-    reverse_iterator rend() const { return reverse_iterator(begin()); }
+    constexpr reverse_iterator rbegin() const {
+      return reverse_iterator(end());
+    }
+    constexpr reverse_iterator rend() const {
+      return reverse_iterator(begin());
+    }
 
     /// empty - Check if the array is empty.
-    bool empty() const { return Length == 0; }
+    constexpr bool empty() const { return Length == 0; }
 
-    const T *data() const { return Data; }
+    constexpr const T *data() const { return Data; }
 
     /// size - Get the array size.
-    size_t size() const { return Length; }
+    constexpr size_t size() const { return Length; }
 
     /// front - Get the first element.
-    const T &front() const {
+    constexpr const T &front() const {
       assert(!empty());
       return Data[0];
     }
 
     /// back - Get the last element.
-    const T &back() const {
+    constexpr const T &back() const {
       assert(!empty());
       return Data[Length-1];
     }
@@ -184,7 +192,7 @@ namespace llvm {
     }
 
     /// equals - Check for element-wise equality.
-    bool equals(ArrayRef RHS) const {
+    constexpr bool equals(ArrayRef RHS) const {
       if (Length != RHS.Length)
         return false;
       return std::equal(begin(), end(), RHS.begin());
@@ -192,22 +200,22 @@ namespace llvm {
 
     /// slice(n, m) - Chop off the first N elements of the array, and keep M
     /// elements in the array.
-    ArrayRef<T> slice(size_t N, size_t M) const {
+    constexpr ArrayRef<T> slice(size_t N, size_t M) const {
       assert(N+M <= size() && "Invalid specifier");
       return ArrayRef<T>(data()+N, M);
     }
 
     /// slice(n) - Chop off the first N elements of the array.
-    ArrayRef<T> slice(size_t N) const { return slice(N, size() - N); }
+    constexpr ArrayRef<T> slice(size_t N) const { return slice(N, size() - N); }
 
     /// Drop the first \p N elements of the array.
-    ArrayRef<T> drop_front(size_t N = 1) const {
+    constexpr ArrayRef<T> drop_front(size_t N = 1) const {
       assert(size() >= N && "Dropping more elements than exist");
       return slice(N, size() - N);
     }
 
     /// Drop the last \p N elements of the array.
-    ArrayRef<T> drop_back(size_t N = 1) const {
+    constexpr ArrayRef<T> drop_back(size_t N = 1) const {
       assert(size() >= N && "Dropping more elements than exist");
       return slice(0, size() - N);
     }
@@ -225,14 +233,14 @@ namespace llvm {
     }
 
     /// Return a copy of *this with only the first \p N elements.
-    ArrayRef<T> take_front(size_t N = 1) const {
+    constexpr ArrayRef<T> take_front(size_t N = 1) const {
       if (N >= size())
         return *this;
       return drop_back(size() - N);
     }
 
     /// Return a copy of *this with only the last \p N elements.
-    ArrayRef<T> take_back(size_t N = 1) const {
+    constexpr ArrayRef<T> take_back(size_t N = 1) const {
       if (N >= size())
         return *this;
       return drop_front(size() - N);
@@ -253,7 +261,7 @@ namespace llvm {
     /// @}
     /// @name Operator Overloads
     /// @{
-    const T &operator[](size_t Index) const {
+    constexpr const T &operator[](size_t Index) const {
       assert(Index < Length && "Invalid index!");
       return Data[Index];
     }
@@ -319,20 +327,20 @@ namespace llvm {
     using difference_type = ptrdiff_t;
 
     /// Construct an empty MutableArrayRef.
-    /*implicit*/ MutableArrayRef() = default;
+    /*implicit*/ constexpr MutableArrayRef() = default;
 
     /// Construct an empty MutableArrayRef from std::nullopt.
-    /*implicit*/ MutableArrayRef(std::nullopt_t) : ArrayRef<T>() {}
+    /*implicit*/ constexpr MutableArrayRef(std::nullopt_t) : ArrayRef<T>() {}
 
     /// Construct a MutableArrayRef from a single element.
-    /*implicit*/ MutableArrayRef(T &OneElt) : ArrayRef<T>(OneElt) {}
+    /*implicit*/ constexpr MutableArrayRef(T &OneElt) : ArrayRef<T>(OneElt) {}
 
     /// Construct a MutableArrayRef from a pointer and length.
-    /*implicit*/ MutableArrayRef(T *data, size_t length)
-      : ArrayRef<T>(data, length) {}
+    /*implicit*/ constexpr MutableArrayRef(T *data, size_t length)
+        : ArrayRef<T>(data, length) {}
 
     /// Construct a MutableArrayRef from a range.
-    MutableArrayRef(T *begin, T *end) : ArrayRef<T>(begin, end) {}
+    constexpr MutableArrayRef(T *begin, T *end) : ArrayRef<T>(begin, end) {}
 
     /// Construct a MutableArrayRef from a SmallVector.
     /*implicit*/ MutableArrayRef(SmallVectorImpl<T> &Vec)
@@ -351,45 +359,49 @@ namespace llvm {
     template <size_t N>
     /*implicit*/ constexpr MutableArrayRef(T (&Arr)[N]) : ArrayRef<T>(Arr) {}
 
-    T *data() const { return const_cast<T*>(ArrayRef<T>::data()); }
+    constexpr T *data() const { return const_cast<T *>(ArrayRef<T>::data()); }
 
-    iterator begin() const { return data(); }
-    iterator end() const { return data() + this->size(); }
+    constexpr iterator begin() const { return data(); }
+    constexpr iterator end() const { return data() + this->size(); }
 
-    reverse_iterator rbegin() const { return reverse_iterator(end()); }
-    reverse_iterator rend() const { return reverse_iterator(begin()); }
+    constexpr reverse_iterator rbegin() const {
+      return reverse_iterator(end());
+    }
+    constexpr reverse_iterator rend() const {
+      return reverse_iterator(begin());
+    }
 
     /// front - Get the first element.
-    T &front() const {
+    constexpr T &front() const {
       assert(!this->empty());
       return data()[0];
     }
 
     /// back - Get the last element.
-    T &back() const {
+    constexpr T &back() const {
       assert(!this->empty());
       return data()[this->size()-1];
     }
 
     /// slice(n, m) - Chop off the first N elements of the array, and keep M
     /// elements in the array.
-    MutableArrayRef<T> slice(size_t N, size_t M) const {
+    constexpr MutableArrayRef<T> slice(size_t N, size_t M) const {
       assert(N + M <= this->size() && "Invalid specifier");
       return MutableArrayRef<T>(this->data() + N, M);
     }
 
     /// slice(n) - Chop off the first N elements of the array.
-    MutableArrayRef<T> slice(size_t N) const {
+    constexpr MutableArrayRef<T> slice(size_t N) const {
       return slice(N, this->size() - N);
     }
 
     /// Drop the first \p N elements of the array.
-    MutableArrayRef<T> drop_front(size_t N = 1) const {
+    constexpr MutableArrayRef<T> drop_front(size_t N = 1) const {
       assert(this->size() >= N && "Dropping more elements than exist");
       return slice(N, this->size() - N);
     }
 
-    MutableArrayRef<T> drop_back(size_t N = 1) const {
+    constexpr MutableArrayRef<T> drop_back(size_t N = 1) const {
       assert(this->size() >= N && "Dropping more elements than exist");
       return slice(0, this->size() - N);
     }
@@ -409,14 +421,14 @@ namespace llvm {
     }
 
     /// Return a copy of *this with only the first \p N elements.
-    MutableArrayRef<T> take_front(size_t N = 1) const {
+    constexpr MutableArrayRef<T> take_front(size_t N = 1) const {
       if (N >= this->size())
         return *this;
       return drop_back(this->size() - N);
     }
 
     /// Return a copy of *this with only the last \p N elements.
-    MutableArrayRef<T> take_back(size_t N = 1) const {
+    constexpr MutableArrayRef<T> take_back(size_t N = 1) const {
       if (N >= this->size())
         return *this;
       return drop_front(this->size() - N);
@@ -439,7 +451,7 @@ namespace llvm {
     /// @}
     /// @name Operator Overloads
     /// @{
-    T &operator[](size_t Index) const {
+    constexpr T &operator[](size_t Index) const {
       assert(Index < this->size() && "Invalid index!");
       return data()[Index];
     }
diff --git a/llvm/include/llvm/ADT/Matrix.h b/llvm/include/llvm/ADT/Matrix.h
new file mode 100644
index 0000000000000..eeb9702772738
--- /dev/null
+++ b/llvm/include/llvm/ADT/Matrix.h
@@ -0,0 +1,339 @@
+//===- Matrix.h - Two-dimensional Container with View -----------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ADT_MATRIX_H
+#define LLVM_ADT_MATRIX_H
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace llvm {
+template <typename T, size_t M, size_t NStorageInline> class MatrixView;
+
+/// Due to the SmallVector infrastructure using SmallVectorAlignmentAndOffset
+/// that depends on the exact data layout, no derived classes can have extra
+/// members.
+template <typename T, size_t N>
+struct MatrixStorageBase : public SmallVectorImpl<T>, SmallVectorStorage<T, N> {
+  MatrixStorageBase() : SmallVectorImpl<T>(N) {}
+  MatrixStorageBase(size_t Size) : SmallVectorImpl<T>(N) { resize(Size); }
+  ~MatrixStorageBase() { destroy_range(this->begin(), this->end()); }
+  MatrixStorageBase(const MatrixStorageBase &RHS) : SmallVectorImpl<T>(N) {
+    if (!RHS.empty())
+      SmallVectorImpl<T>::operator=(RHS);
+  }
+  MatrixStorageBase(MatrixStorageBase &&RHS) : SmallVectorImpl<T>(N) {
+    if (!RHS.empty())
+      SmallVectorImpl<T>::operator=(::std::move(RHS));
+  }
+  using SmallVectorImpl<T>::size;
+  using SmallVectorImpl<T>::resize;
+  using SmallVectorImpl<T>::append;
+  using SmallVectorImpl<T>::erase;
+  using SmallVectorImpl<T>::destroy_range;
+
+  T *begin() const { return const_cast<T *>(SmallVectorImpl<T>::begin()); }
+  T *end() const { return const_cast<T *>(SmallVectorImpl<T>::end()); }
+};
+
+/// A two-dimensional container storage, whose upper bound on the number of
+/// columns should be known ahead of time. Not menat to be used directly: the
+/// primary usage API is MatrixView.
+template <typename T,
+          size_t N = CalculateSmallVectorDefaultInlinedElements<T>::value>
+class MatrixStorage {
+public:
+  MatrixStorage() = delete;
+  MatrixStorage(size_t NRows, size_t NCols)
+      : Base(NRows * NCols), NCols(NCols) {}
+  MatrixStorage(size_t NCols) : Base(), NCols(NCols) {}
+
+  size_t size() const { return Base.size(); }
+  bool empty() const { return !size(); }
+  size_t getNumRows() const { return size() / NCols; }
+  size_t getNumCols() const { return NCols; }
+  void setNumCols(size_t NCols) {
+    assert(empty() && "Column-resizing a non-empty MatrixStorage");
+    this->NCols = NCols;
+  }
+  void resize(size_t NRows) { Base.resize(NCols * NRows); }
+
+protected:
+  template <typename U, size_t M, size_t NStorageInline>
+  friend class MatrixView;
+
+  T *begin() const { return Base.begin(); }
+  T *rowFromIdx(size_t RowIdx, size_t Offset = 0) const {
+    return begin() + RowIdx * NCols + Offset;
+  }
+  std::pair<size_t, size_t> idxFromRow(T *Ptr) const {
+    assert(Ptr >= begin() && "Internal error");
+    size_t Offset = (Ptr - begin()) % NCols;
+    return {(Ptr - begin()) / NCols, Offset};
+  }
+
+  // If Arg.size() < NCols, the number of columns won't be changed, and the
+  // difference is default-constructed.
+  void addRow(const SmallVectorImpl<T> &Arg) {
+    assert(Arg.size() <= NCols &&
+           "MatrixStorage has insufficient number of columns");
+    size_t Diff = NCols - Arg.size();
+    Base.append(Arg.begin(), Arg.end());
+    Base.append(Diff, T());
+  }
+
+private:
+  MatrixStorageBase<T, N> Base;
+  size_t NCols;
+};
+
+/// MutableArrayRef with a copy-assign, and extra APIs.
+template <typename T>
+struct [[nodiscard]] MatrixRowView : public MutableArrayRef<T> {
+  using pointer = typename MutableArrayRef<T>::pointer;
+  using iterator = typename MutableArrayRef<T>::iterator;
+  using const_iterator = typename MutableArrayRef<T>::const_iterator;
+
+  constexpr MatrixRowView() = delete;
+  constexpr MatrixRowView(pointer Data, size_t Length)
+      : MutableArrayRef<T>(Data, Length) {}
+  constexpr MatrixRowView(iterator Begin, iterator End)
+      : MutableArrayRef<T>(Begin, End) {}
+  constexpr MatrixRowView(const_iterator Begin, const_iterator End)
+      : MutableArrayRef<T>(Begin, End) {}
+  constexpr MatrixRowView(MutableArrayRef<T> Other)
+      : MutableArrayRef<T>(Other.data(), Other.size()) {}
+  MatrixRowView(const SmallVectorImpl<T> &Vec) : MutableArrayRef<T>(Vec) {}
+
+  using MutableArrayRef<T>::size;
+  using MutableArrayRef<T>::data;
+
+  constexpr T &back() const { return MutableArrayRef<T>::back(); }
+  constexpr T &front() const { return MutableArrayRef<T>::front(); }
+  constexpr MatrixRowView<T> drop_back(size_t N = 1) const { // NOLINT
+    return MutableArrayRef<T>::drop_back(N);
+  }
+  constexpr MatrixRowView<T> drop_front(size_t N = 1) const { // NOLINT
+    return MutableArrayRef<T>::drop_front(N);
+  }
+  // This slice is different from the MutableArrayRef slice, and specifies a
+  // Begin and End index, instead of a Begin and Length.
+  constexpr MatrixRowView<T> slice(size_t Begin, size_t End) {
+    return MutableArrayRef<T>::slice(Begin, End - Begin);
+  }
+  constexpr void pop_back(size_t N = 1) { // NOLINT
+    this->Length -= N;
+  }
+  constexpr void pop_front(size_t N = 1) { // NOLINT
+    this->Data += N;
+    this->Length -= N;
+  }
+
+  MatrixRowView &operator=(const SmallVectorImpl<T> &Vec) {
+    copy_assign(Vec.begin(), Vec.end());
+    return *this;
+  }
+  MatrixRowView &operator=(std::initializer_list<T> IL) {
+    copy_assign(IL.begin(), IL.end());
+    return *this;
+  }
+
+  void swap(MatrixRowView<T> &Other) {
+    std::swap(this->Data, Other.Data);
+    std::swap(this->Length, Other.Length);
+  }
+
+protected:
+  void copy_assign(iterator Begin, iterator End) { // NOLINT
+    std::uninitialized_copy(Begin, End, data());
+    this->Length = End - Begin;
+  }
+  void copy_assign(const_iterator Begin, const_iterator End) { // NOLINT
+    std::uninitialized_copy(Begin, End, data());
+    this->Length = End - Begin;
+  }
+};
+
+/// The primary usage API of MatrixStorage. Abstracts out indexing-arithmetic,
+/// eliminating memory operations on the underlying data. Supports
+/// variable-length columns.
+template <typename T,
+          size_t N = CalculateSmallVectorDefaultInlinedElements<T>::value,
+          size_t NStorageInline =
+              CalculateSmallVectorDefaultInlinedElements<T>::value>
+class [[nodiscard]] MatrixView {
+public:
+  using row_type = MatrixRowView<T>;
+  using container_type = SmallVector<row_type, N>;
+  using iterator = typename container_type::iterator;
+  using const_iterator = typename container_type::const_iterator;
+
+  constexpr MatrixView(MatrixStorage<T, NStorageInline> &Mat, size_t RowSpan,
+                       size_t ColSpan)
+      : Mat(Mat) {
+    RowView.reserve(RowSpan);
+    for (size_t RowIdx = 0; RowIdx < RowSpan; ++RowIdx) {
+      auto RangeBegin = Mat.begin() + RowIdx * ColSpan;
+      RowView.emplace_back(RangeBegin, RangeBegin + ColSpan);
+    }
+  }
+
+  // Constructor with a full View of the underlying MatrixStorage, if
+  // MatrixStorage has a non-zero number of Columns. Otherwise, creates an empty
+  // view.
+  constexpr MatrixView(MatrixStorage<T, NStorageInline> &Mat)
+      : MatrixView(Mat, Mat.getNumRows(), Mat.getNumCols()) {}
+
+  // Obvious copy-construator is deleted, since the underlying storage could
+  // have changed.
+  constexpr MatrixView(const MatrixView &) = delete;
+
+  // Copy-assignment operator should not be used when the underlying storage
+  // changes.
+  constexpr MatrixView &operator=(const MatrixView &Other) {
+    assert(Mat.begin() == Other.Mat.begin() &&
+           "Underlying storage has changed: use custom copy-constructor");
+    RowView = Other.RowView;
+    return *this;
+  }
+
+  // The actual copy-constructor: to be used when the underlying storage is
+  // copy-constructed.
+  MatrixView(const MatrixView &OldView,
+             MatrixStorage<T, NStorageInline> &NewMat)
+      : Mat(NewMat) {
+    assert(OldView.Mat.size() == Mat.size() &&
+           "Custom copy-constructor called on non-copied storage");
+
+    // The underlying storage will change. Construct a new RowView by performing
+    // pointer-arithmetic on the underlying storage of OldView, using pointers
+    // from OldVie.
+    for (const auto &R : OldView.RowView) {
+      auto [StorageIdx, StartOffset] = OldView.Mat.idxFromRow(R.data());
+      RowView.emplace_back(Mat.rowFromIdx(StorageIdx, StartOffset), R.size());
+    }
+  }
+
+  void addRow(const SmallVectorImpl<T> &Row) {
+    // The underlying storage may be resized, performing reallocations. The
+    // pointers in RowView will no longer be valid, so save and restore the
+    // data. Construct RestoreData by performing pointer-arithmetic on the
+    // underlying storgge.
+    SmallVector<std::tuple<size_t, size_t, size_t>> RestoreData;
+    RestoreData.reserve(RowView.size());
+    for (const auto &R : RowView) {
+      auto [StorageIdx, StartOffset] = Mat.idxFromRow(R.data());
+      RestoreData.emplace_back(StorageIdx, StartOffset, R.size());
+    }
+
+    Mat.addRow(Row);
+
+    // Restore the RowView by performing pointer-arithmetic on the
+    // possibly-reallocated storage, using information from RestoreData.
+    RowView.clear();
+    for (const auto &[StorageIdx, StartOffset, Len] : RestoreData)
+      RowView.emplace_back(Mat.rowFromIdx(StorageIdx, StartOffset), Len);
+
+    // Finally, add the new row to the VRowView.
+    RowView.emplace_back(Mat.rowFromIdx(Mat.getNumRows() - 1), Row.size());
+  }
+
+  // To support addRow(View[Idx]).
+  void addRow(const row_type &Row) { addRow(SmallVector<T>{Row}); }
+
+  void addRow(std::initializer_list<T> Row) { addRow(SmallVector<T>{Row}); }
+
+  constexpr row_type &operator[](size_t RowIdx) {
+    assert(RowIdx < RowView.size() && "Indexing out of bounds");
+    return RowView[RowIdx];
+  }
+
+  constexpr T *data() const {
+    assert(!empty() && "Non-empty view expected");
+    return RowView.front().data();
+  }
+  size_t size() const { return getRowSpan() * getMaxColSpan(); }
+...
[truncated]

``````````

</details>


https://github.com/llvm/llvm-project/pull/98893


More information about the llvm-commits mailing list