1 #pragma once 2 3 // A fixed-size array type usable from both host and 4 // device code. 5 6 #include <c10/macros/Macros.h> 7 #include <c10/util/irange.h> 8 9 namespace at::detail { 10 11 template <typename T, int size_> 12 struct Array { 13 // NOLINTNEXTLINE(*c-array*) 14 T data[size_]; 15 16 C10_HOST_DEVICE T operator[](int i) const { 17 return data[i]; 18 } 19 C10_HOST_DEVICE T& operator[](int i) { 20 return data[i]; 21 } 22 #if defined(USE_ROCM) 23 C10_HOST_DEVICE Array() = default; 24 C10_HOST_DEVICE Array(const Array&) = default; 25 C10_HOST_DEVICE Array& operator=(const Array&) = default; 26 #else 27 Array() = default; 28 Array(const Array&) = default; 29 Array& operator=(const Array&) = default; 30 #endif sizeArray31 static constexpr int size() { 32 return size_; 33 } 34 // Fill the array with x. ArrayArray35 C10_HOST_DEVICE Array(T x) { 36 for (int i = 0; i < size_; i++) { 37 data[i] = x; 38 } 39 } 40 }; 41 42 } // namespace at::detail 43