49 #ifndef Intrepid2_TensorViewIterator_h 50 #define Intrepid2_TensorViewIterator_h 55 #include "Kokkos_Vector.hpp" 72 template<
class TensorViewType,
class ViewType1,
class ViewType2 ,
typename ScalarType>
76 enum RankCombinationType
82 using RankCombinationViewType = Kokkos::View<RankCombinationType*, typename TensorViewType::device_type>;
89 RankCombinationViewType rank_combination_types_;
110 KOKKOS_INLINE_FUNCTION
112 RankCombinationViewType rank_combination_types)
114 tensor_view_iterator_(tensor_view),
115 view1_iterator_(view1),
116 view2_iterator_(view2),
117 rank_combination_types_(rank_combination_types)
135 unsigned max_component_rank = (view1.rank() > view2.rank()) ? view1.rank() : view2.rank();
136 unsigned max_rank = (tensor_view.rank() > max_component_rank) ? tensor_view.rank() : max_component_rank;
138 INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(rank_combination_types.extent(0) != max_rank, std::invalid_argument,
"need to provide RankCombinationType for the largest-rank View");
140 unsigned expected_rank = 0;
141 bool contracting =
false;
142 for (
unsigned d=0; d<rank_combination_types.extent(0); d++)
144 if (rank_combination_types[d] == TENSOR_CONTRACTION)
147 INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(view1.extent_int(d) != view2.extent_int(d), std::invalid_argument,
"Contractions can only occur along ranks of equal length");
155 if (rank_combination_types[d] == TENSOR_PRODUCT)
157 INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(tensor_view.extent_int(d) != view1.extent_int(d) * view2.extent_int(d), std::invalid_argument,
"For TENSOR_PRODUCT rank combination, the tensor View must have length in that dimension equal to the product of the two component views in that dimension");
161 INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(view1.extent_int(d) != view2.extent_int(d), std::invalid_argument,
"For DIMENSION_MATCH rank combination, all three views must have length equal to each other in that rank");
162 INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(tensor_view.extent_int(d) != view1.extent_int(d), std::invalid_argument,
"For DIMENSION_MATCH rank combination, all three views must have length equal to each other in that rank");
167 INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(contracting, std::invalid_argument,
"encountered a non-contraction rank combination after a contraction; contractions can only go at the end");
171 INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(expected_rank != tensor_view.rank(), std::invalid_argument,
"Tensor view does not match expected rank");
176 KOKKOS_INLINE_FUNCTION
181 if (view2_next_increment_rank > view1_next_increment_rank)
return view2_next_increment_rank;
182 else return view1_next_increment_rank;
187 KOKKOS_INLINE_FUNCTION
202 if (view2_next_increment_rank > view1_next_increment_rank)
205 device_assert(rank_combination_types_[view2_next_increment_rank]==TENSOR_PRODUCT);
206 view1_iterator_.
reset(view2_next_increment_rank);
209 return view2_next_increment_rank;
213 int view1_rank_change = view1_iterator_.
increment();
214 if (view1_rank_change >= 0)
216 switch (rank_combination_types_[view1_rank_change])
218 case DIMENSION_MATCH:
226 case TENSOR_CONTRACTION:
231 return view1_rank_change;
237 KOKKOS_INLINE_FUNCTION
248 KOKKOS_INLINE_FUNCTION
249 void setLocation(Kokkos::Array<int,7> location1, Kokkos::Array<int,7> location2)
253 Kokkos::Array<int,7> tensor_location = location1;
254 for (
unsigned d=0; d<rank_combination_types_.extent(0); d++)
256 switch (rank_combination_types_[d])
260 tensor_location[d] = location2[d] * view1_iterator_.
getExtent(d) + location1[d];
262 case DIMENSION_MATCH:
265 case TENSOR_CONTRACTION:
266 tensor_location[d] = 0;
270 #ifdef HAVE_INTREPID2_DEBUG 272 for (
unsigned d=0; d<rank_combination_types_.extent(0); d++)
274 switch (rank_combination_types_[d])
279 case DIMENSION_MATCH:
280 case TENSOR_CONTRACTION:
281 device_assert(location1[d] == location2[d]);
285 device_assert(location1[d] < view1_iterator_.
getExtent(d));
286 device_assert(location2[d] < view2_iterator_.
getExtent(d));
287 device_assert(tensor_location[d] < tensor_view_iterator_.
getExtent(d));
290 tensor_view_iterator_.
setLocation(tensor_location);
295 KOKKOS_INLINE_FUNCTION
298 return view1_iterator_.
get();
303 KOKKOS_INLINE_FUNCTION
306 return view2_iterator_.
get();
311 KOKKOS_INLINE_FUNCTION
312 void set(ScalarType value)
314 tensor_view_iterator_.
set(value);
KOKKOS_INLINE_FUNCTION ScalarType getView1Entry()
KOKKOS_INLINE_FUNCTION void reset(unsigned from_rank_number=0)
KOKKOS_INLINE_FUNCTION int nextIncrementRank()
KOKKOS_INLINE_FUNCTION int getExtent(int dimension)
KOKKOS_INLINE_FUNCTION TensorViewIterator(TensorViewType tensor_view, ViewType1 view1, ViewType2 view2, RankCombinationViewType rank_combination_types)
Constructor.
KOKKOS_INLINE_FUNCTION ScalarType get()
KOKKOS_INLINE_FUNCTION void setLocation(Kokkos::Array< int, 7 > location1, Kokkos::Array< int, 7 > location2)
Implementation of an assert that can safely be called from device code.
Iterator allows linear traversal of (part of) a Kokkos View in a manner that is agnostic to its rank...
KOKKOS_INLINE_FUNCTION void set(const ScalarType &value)
KOKKOS_INLINE_FUNCTION int increment()
KOKKOS_INLINE_FUNCTION void setLocation(const Kokkos::Array< int, 7 > location)
KOKKOS_INLINE_FUNCTION ScalarType getView2Entry()
KOKKOS_INLINE_FUNCTION int increment()
KOKKOS_INLINE_FUNCTION void setLocation(const Kokkos::Array< int, 7 > &location)
KOKKOS_INLINE_FUNCTION int nextIncrementRank()
A helper class that allows iteration over three Kokkos Views simultaneously, according to tensor comb...