ArrayRef.h 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380
  1. //===--- ArrayRef.h - Array Reference Wrapper -------------------*- C++ -*-===//
  2. //
  3. // The LLVM Compiler Infrastructure
  4. //
  5. // This file is distributed under the University of Illinois Open Source
  6. // License. See LICENSE.TXT for details.
  7. //
  8. //===----------------------------------------------------------------------===//
  9. // ATen: modified from llvm::ArrayRef.
  10. // removed llvm-specific functionality
  11. // removed some implicit const -> non-const conversions that rely on
  12. // complicated std::enable_if meta-programming
  13. // removed a bunch of slice variants for simplicity...
  14. #pragma once
  15. #include <c10/macros/Macros.h>
  16. #include <c10/util/Deprecated.h>
  17. #include <c10/util/Exception.h>
  18. #include <c10/util/SmallVector.h>
  19. #include <array>
  20. #include <cstddef>
  21. #include <cstdint>
  22. #include <initializer_list>
  23. #include <iterator>
  24. #include <ostream>
  25. #include <type_traits>
  26. #include <vector>
  27. namespace c10 {
  28. /// ArrayRef - Represent a constant reference to an array (0 or more elements
  29. /// consecutively in memory), i.e. a start pointer and a length. It allows
  30. /// various APIs to take consecutive elements easily and conveniently.
  31. ///
  32. /// This class does not own the underlying data, it is expected to be used in
  33. /// situations where the data resides in some other buffer, whose lifetime
  34. /// extends past that of the ArrayRef. For this reason, it is not in general
  35. /// safe to store an ArrayRef.
  36. ///
  37. /// This is intended to be trivially copyable, so it should be passed by
  38. /// value.
  39. template <typename T>
  40. class ArrayRef final {
  41. public:
  42. using iterator = const T*;
  43. using const_iterator = const T*;
  44. using size_type = size_t;
  45. using value_type = T;
  46. using reverse_iterator = std::reverse_iterator<iterator>;
  47. private:
  48. /// The start of the array, in an external buffer.
  49. const T* Data;
  50. /// The number of elements.
  51. size_type Length;
  52. void debugCheckNullptrInvariant() {
  53. TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
  54. Data != nullptr || Length == 0,
  55. "created ArrayRef with nullptr and non-zero length! std::optional relies on this being illegal");
  56. }
  57. public:
  58. /// @name Constructors
  59. /// @{
  60. /// Construct an empty ArrayRef.
  61. /* implicit */ constexpr ArrayRef() : Data(nullptr), Length(0) {}
  62. /// Construct an ArrayRef from a single element.
  63. // TODO Make this explicit
  64. constexpr ArrayRef(const T& OneElt) : Data(&OneElt), Length(1) {}
  65. /// Construct an ArrayRef from a pointer and length.
  66. C10_HOST_CONSTEXPR_EXCEPT_WIN_CUDA ArrayRef(const T* data, size_t length)
  67. : Data(data), Length(length) {
  68. debugCheckNullptrInvariant();
  69. }
  70. /// Construct an ArrayRef from a range.
  71. C10_HOST_CONSTEXPR_EXCEPT_WIN_CUDA ArrayRef(const T* begin, const T* end)
  72. : Data(begin), Length(end - begin) {
  73. debugCheckNullptrInvariant();
  74. }
  75. /// Construct an ArrayRef from a SmallVector. This is templated in order to
  76. /// avoid instantiating SmallVectorTemplateCommon<T> whenever we
  77. /// copy-construct an ArrayRef.
  78. template <typename U>
  79. /* implicit */ ArrayRef(const SmallVectorTemplateCommon<T, U>& Vec)
  80. : Data(Vec.data()), Length(Vec.size()) {
  81. debugCheckNullptrInvariant();
  82. }
  83. template <
  84. typename Container,
  85. typename = std::enable_if_t<std::is_same_v<
  86. std::remove_const_t<decltype(std::declval<Container>().data())>,
  87. T*>>>
  88. /* implicit */ ArrayRef(const Container& container)
  89. : Data(container.data()), Length(container.size()) {
  90. debugCheckNullptrInvariant();
  91. }
  92. /// Construct an ArrayRef from a std::vector.
  93. // The enable_if stuff here makes sure that this isn't used for
  94. // std::vector<bool>, because ArrayRef can't work on a std::vector<bool>
  95. // bitfield.
  96. template <typename A>
  97. /* implicit */ ArrayRef(const std::vector<T, A>& Vec)
  98. : Data(Vec.data()), Length(Vec.size()) {
  99. static_assert(
  100. !std::is_same<T, bool>::value,
  101. "ArrayRef<bool> cannot be constructed from a std::vector<bool> bitfield.");
  102. }
  103. /// Construct an ArrayRef from a std::array
  104. template <size_t N>
  105. /* implicit */ constexpr ArrayRef(const std::array<T, N>& Arr)
  106. : Data(Arr.data()), Length(N) {}
  107. /// Construct an ArrayRef from a C array.
  108. template <size_t N>
  109. // NOLINTNEXTLINE(*c-arrays*)
  110. /* implicit */ constexpr ArrayRef(const T (&Arr)[N]) : Data(Arr), Length(N) {}
  111. /// Construct an ArrayRef from a std::initializer_list.
  112. /* implicit */ constexpr ArrayRef(const std::initializer_list<T>& Vec)
  113. : Data(
  114. std::begin(Vec) == std::end(Vec) ? static_cast<T*>(nullptr)
  115. : std::begin(Vec)),
  116. Length(Vec.size()) {}
  117. /// @}
  118. /// @name Simple Operations
  119. /// @{
  120. constexpr iterator begin() const {
  121. return Data;
  122. }
  123. constexpr iterator end() const {
  124. return Data + Length;
  125. }
  126. // These are actually the same as iterator, since ArrayRef only
  127. // gives you const iterators.
  128. constexpr const_iterator cbegin() const {
  129. return Data;
  130. }
  131. constexpr const_iterator cend() const {
  132. return Data + Length;
  133. }
  134. constexpr reverse_iterator rbegin() const {
  135. return reverse_iterator(end());
  136. }
  137. constexpr reverse_iterator rend() const {
  138. return reverse_iterator(begin());
  139. }
  140. /// empty - Check if the array is empty.
  141. constexpr bool empty() const {
  142. return Length == 0;
  143. }
  144. constexpr const T* data() const {
  145. return Data;
  146. }
  147. /// size - Get the array size.
  148. constexpr size_t size() const {
  149. return Length;
  150. }
  151. /// front - Get the first element.
  152. C10_HOST_CONSTEXPR_EXCEPT_WIN_CUDA const T& front() const {
  153. TORCH_CHECK(
  154. !empty(), "ArrayRef: attempted to access front() of empty list");
  155. return Data[0];
  156. }
  157. /// back - Get the last element.
  158. C10_HOST_CONSTEXPR_EXCEPT_WIN_CUDA const T& back() const {
  159. TORCH_CHECK(!empty(), "ArrayRef: attempted to access back() of empty list");
  160. return Data[Length - 1];
  161. }
  162. /// equals - Check for element-wise equality.
  163. constexpr bool equals(ArrayRef RHS) const {
  164. return Length == RHS.Length && std::equal(begin(), end(), RHS.begin());
  165. }
  166. /// slice(n, m) - Take M elements of the array starting at element N
  167. C10_HOST_CONSTEXPR_EXCEPT_WIN_CUDA ArrayRef<T> slice(size_t N, size_t M)
  168. const {
  169. TORCH_CHECK(
  170. N + M <= size(),
  171. "ArrayRef: invalid slice, N = ",
  172. N,
  173. "; M = ",
  174. M,
  175. "; size = ",
  176. size());
  177. return ArrayRef<T>(data() + N, M);
  178. }
  179. /// slice(n) - Chop off the first N elements of the array.
  180. C10_HOST_CONSTEXPR_EXCEPT_WIN_CUDA ArrayRef<T> slice(size_t N) const {
  181. TORCH_CHECK(
  182. N <= size(), "ArrayRef: invalid slice, N = ", N, "; size = ", size());
  183. return slice(N, size() - N);
  184. }
  185. /// @}
  186. /// @name Operator Overloads
  187. /// @{
  188. constexpr const T& operator[](size_t Index) const {
  189. return Data[Index];
  190. }
  191. /// Vector compatibility
  192. C10_HOST_CONSTEXPR_EXCEPT_WIN_CUDA const T& at(size_t Index) const {
  193. TORCH_CHECK(
  194. Index < Length,
  195. "ArrayRef: invalid index Index = ",
  196. Index,
  197. "; Length = ",
  198. Length);
  199. return Data[Index];
  200. }
  201. /// Disallow accidental assignment from a temporary.
  202. ///
  203. /// The declaration here is extra complicated so that "arrayRef = {}"
  204. /// continues to select the move assignment operator.
  205. template <typename U>
  206. std::enable_if_t<std::is_same_v<U, T>, ArrayRef<T>>& operator=(
  207. // NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
  208. U&& Temporary) = delete;
  209. /// Disallow accidental assignment from a temporary.
  210. ///
  211. /// The declaration here is extra complicated so that "arrayRef = {}"
  212. /// continues to select the move assignment operator.
  213. template <typename U>
  214. std::enable_if_t<std::is_same_v<U, T>, ArrayRef<T>>& operator=(
  215. std::initializer_list<U>) = delete;
  216. /// @}
  217. /// @name Expensive Operations
  218. /// @{
  219. std::vector<T> vec() const {
  220. return std::vector<T>(Data, Data + Length);
  221. }
  222. /// @}
  223. };
  224. template <typename T>
  225. std::ostream& operator<<(std::ostream& out, ArrayRef<T> list) {
  226. int i = 0;
  227. out << "[";
  228. for (const auto& e : list) {
  229. if (i++ > 0)
  230. out << ", ";
  231. out << e;
  232. }
  233. out << "]";
  234. return out;
  235. }
  236. /// @name ArrayRef Convenience constructors
  237. /// @{
  238. /// Construct an ArrayRef from a single element.
  239. template <typename T>
  240. ArrayRef<T> makeArrayRef(const T& OneElt) {
  241. return OneElt;
  242. }
  243. /// Construct an ArrayRef from a pointer and length.
  244. template <typename T>
  245. ArrayRef<T> makeArrayRef(const T* data, size_t length) {
  246. return ArrayRef<T>(data, length);
  247. }
  248. /// Construct an ArrayRef from a range.
  249. template <typename T>
  250. ArrayRef<T> makeArrayRef(const T* begin, const T* end) {
  251. return ArrayRef<T>(begin, end);
  252. }
  253. /// Construct an ArrayRef from a SmallVector.
  254. template <typename T>
  255. ArrayRef<T> makeArrayRef(const SmallVectorImpl<T>& Vec) {
  256. return Vec;
  257. }
  258. /// Construct an ArrayRef from a SmallVector.
  259. template <typename T, unsigned N>
  260. ArrayRef<T> makeArrayRef(const SmallVector<T, N>& Vec) {
  261. return Vec;
  262. }
  263. /// Construct an ArrayRef from a std::vector.
  264. template <typename T>
  265. ArrayRef<T> makeArrayRef(const std::vector<T>& Vec) {
  266. return Vec;
  267. }
  268. /// Construct an ArrayRef from a std::array.
  269. template <typename T, std::size_t N>
  270. ArrayRef<T> makeArrayRef(const std::array<T, N>& Arr) {
  271. return Arr;
  272. }
  273. /// Construct an ArrayRef from an ArrayRef (no-op) (const)
  274. template <typename T>
  275. ArrayRef<T> makeArrayRef(const ArrayRef<T>& Vec) {
  276. return Vec;
  277. }
  278. /// Construct an ArrayRef from an ArrayRef (no-op)
  279. template <typename T>
  280. ArrayRef<T>& makeArrayRef(ArrayRef<T>& Vec) {
  281. return Vec;
  282. }
  283. /// Construct an ArrayRef from a C array.
  284. template <typename T, size_t N>
  285. // NOLINTNEXTLINE(*c-arrays*)
  286. ArrayRef<T> makeArrayRef(const T (&Arr)[N]) {
  287. return ArrayRef<T>(Arr);
  288. }
  289. // WARNING: Template instantiation will NOT be willing to do an implicit
  290. // conversions to get you to an c10::ArrayRef, which is why we need so
  291. // many overloads.
  292. template <typename T>
  293. bool operator==(c10::ArrayRef<T> a1, c10::ArrayRef<T> a2) {
  294. return a1.equals(a2);
  295. }
  296. template <typename T>
  297. bool operator!=(c10::ArrayRef<T> a1, c10::ArrayRef<T> a2) {
  298. return !a1.equals(a2);
  299. }
  300. template <typename T>
  301. bool operator==(const std::vector<T>& a1, c10::ArrayRef<T> a2) {
  302. return c10::ArrayRef<T>(a1).equals(a2);
  303. }
  304. template <typename T>
  305. bool operator!=(const std::vector<T>& a1, c10::ArrayRef<T> a2) {
  306. return !c10::ArrayRef<T>(a1).equals(a2);
  307. }
  308. template <typename T>
  309. bool operator==(c10::ArrayRef<T> a1, const std::vector<T>& a2) {
  310. return a1.equals(c10::ArrayRef<T>(a2));
  311. }
  312. template <typename T>
  313. bool operator!=(c10::ArrayRef<T> a1, const std::vector<T>& a2) {
  314. return !a1.equals(c10::ArrayRef<T>(a2));
  315. }
  316. using IntArrayRef = ArrayRef<int64_t>;
  317. // This alias is deprecated because it doesn't make ownership
  318. // semantics obvious. Use IntArrayRef instead!
  319. C10_DEFINE_DEPRECATED_USING(IntList, ArrayRef<int64_t>)
  320. } // namespace c10