1 // Copyright (c) 2010-2021, Lawrence Livermore National Security, LLC. Produced
2 // at the Lawrence Livermore National Laboratory. All Rights reserved. See files
3 // LICENSE and NOTICE for details. LLNL-CODE-806117.
4 //
5 // This file is part of the MFEM library. For more information and source code
6 // availability visit https://mfem.org.
7 //
8 // MFEM is free software; you can redistribute it and/or modify it under the
9 // terms of the BSD-3 license. We welcome feedback and contributions, see file
10 // CONTRIBUTING.md for details.
11
12 #ifndef MFEM_DTENSOR
13 #define MFEM_DTENSOR
14
15 #include "../general/backends.hpp"
16
17 namespace mfem
18 {
19
20 /// A Class to compute the real index from the multi-indices of a tensor
21 template <int N, int Dim, typename T, typename... Args>
22 class TensorInd
23 {
24 public:
25 MFEM_HOST_DEVICE
result(const int * sizes,T first,Args...args)26 static inline int result(const int* sizes, T first, Args... args)
27 {
28 #if !(defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP))
29 MFEM_ASSERT(first<sizes[N-1],"Trying to access out of boundary.");
30 #endif
31 return first + sizes[N - 1] * TensorInd < N + 1, Dim, Args... >
32 ::result(sizes, args...);
33 }
34 };
35
36 // Terminal case
37 template <int Dim, typename T, typename... Args>
38 class TensorInd<Dim, Dim, T, Args...>
39 {
40 public:
41 MFEM_HOST_DEVICE
result(const int * sizes,T first,Args...args)42 static inline int result(const int* sizes, T first, Args... args)
43 {
44 #if !(defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP))
45 MFEM_ASSERT(first<sizes[Dim-1],"Trying to access out of boundary.");
46 #endif
47 return first;
48 }
49 };
50
51
52 /// A class to initialize the size of a Tensor
53 template <int N, int Dim, typename T, typename... Args>
54 class Init
55 {
56 public:
57 MFEM_HOST_DEVICE
result(int * sizes,T first,Args...args)58 static inline int result(int* sizes, T first, Args... args)
59 {
60 sizes[N - 1] = first;
61 return first * Init < N + 1, Dim, Args... >::result(sizes, args...);
62 }
63 };
64
65 // Terminal case
66 template <int Dim, typename T, typename... Args>
67 class Init<Dim, Dim, T, Args...>
68 {
69 public:
70 MFEM_HOST_DEVICE
result(int * sizes,T first,Args...args)71 static inline int result(int* sizes, T first, Args... args)
72 {
73 sizes[Dim - 1] = first;
74 return first;
75 }
76 };
77
78
79 /// A basic generic Tensor class, appropriate for use on the GPU
80 template<int Dim, typename Scalar = double>
81 class DeviceTensor
82 {
83 protected:
84 int capacity;
85 Scalar *data;
86 int sizes[Dim];
87
88 public:
89 /// Default constructor
90 DeviceTensor() = delete;
91
92 /// Constructor to initialize a tensor from the Scalar array data_
93 template <typename... Args> MFEM_HOST_DEVICE
DeviceTensor(Scalar * data_,Args...args)94 DeviceTensor(Scalar* data_, Args... args)
95 {
96 static_assert(sizeof...(args) == Dim, "Wrong number of arguments");
97 // Initialize sizes, and compute the number of values
98 const long int nb = Init<1, Dim, Args...>::result(sizes, args...);
99 capacity = nb;
100 data = (capacity > 0) ? data_ : NULL;
101 }
102
103 /// Copy constructor
DeviceTensor(const DeviceTensor & t)104 MFEM_HOST_DEVICE DeviceTensor(const DeviceTensor& t)
105 {
106 capacity = t.capacity;
107 for (int i = 0; i < Dim; ++i)
108 {
109 sizes[i] = t.sizes[i];
110 }
111 data = t.data;
112 }
113
114 /// Conversion to `Scalar *`.
operator Scalar*() const115 MFEM_HOST_DEVICE inline operator Scalar *() const { return data; }
116
117 /// Const accessor for the data
118 template <typename... Args> MFEM_HOST_DEVICE inline
operator ()(Args...args) const119 Scalar& operator()(Args... args) const
120 {
121 static_assert(sizeof...(args) == Dim, "Wrong number of arguments");
122 return data[ TensorInd<1, Dim, Args...>::result(sizes, args...) ];
123 }
124
125 /// Subscript operator where the tensor is viewed as a 1D array.
operator [](int i) const126 MFEM_HOST_DEVICE inline Scalar& operator[](int i) const
127 {
128 return data[i];
129 }
130 };
131
132
133 /** @brief Wrap a pointer as a DeviceTensor with automatically deduced template
134 parameters */
135 template <typename T, typename... Dims>
Reshape(T * ptr,Dims...dims)136 inline DeviceTensor<sizeof...(Dims),T> Reshape(T *ptr, Dims... dims)
137 {
138 return DeviceTensor<sizeof...(Dims),T>(ptr, dims...);
139 }
140
141
142 typedef DeviceTensor<1,int> DeviceArray;
143 typedef DeviceTensor<1,const int> ConstDeviceArray;
144
145 typedef DeviceTensor<1,double> DeviceVector;
146 typedef DeviceTensor<1,const double> ConstDeviceVector;
147
148 typedef DeviceTensor<2,double> DeviceMatrix;
149 typedef DeviceTensor<2,const double> ConstDeviceMatrix;
150
151 typedef DeviceTensor<3,double> DeviceCube;
152 typedef DeviceTensor<3,const double> ConstDeviceCube;
153
154 } // mfem namespace
155
156 #endif // MFEM_DTENSOR
157