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;
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)
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");
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:
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);
Implementation of an assert that can safely be called from device code.
KOKKOS_INLINE_FUNCTION void device_assert(bool val)
#define INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(test, x, msg)
Iterator allows linear traversal of (part of) a Kokkos View in a manner that is agnostic to its rank.
A helper class that allows iteration over three Kokkos Views simultaneously, according to tensor comb...
KOKKOS_INLINE_FUNCTION void setLocation(Kokkos::Array< int, 7 > location1, Kokkos::Array< int, 7 > location2)
KOKKOS_INLINE_FUNCTION ScalarType getView1Entry()
KOKKOS_INLINE_FUNCTION int nextIncrementRank()
KOKKOS_INLINE_FUNCTION TensorViewIterator(TensorViewType tensor_view, ViewType1 view1, ViewType2 view2, RankCombinationViewType rank_combination_types)
Constructor.
KOKKOS_INLINE_FUNCTION void setLocation(const Kokkos::Array< int, 7 > location)
KOKKOS_INLINE_FUNCTION int increment()
KOKKOS_INLINE_FUNCTION void set(ScalarType value)
KOKKOS_INLINE_FUNCTION ScalarType getView2Entry()
A helper class that allows iteration over some part of a Kokkos View, while allowing the calling code...
KOKKOS_INLINE_FUNCTION int nextIncrementRank()
KOKKOS_INLINE_FUNCTION int getExtent(int dimension)
KOKKOS_INLINE_FUNCTION void set(const ScalarType &value)
KOKKOS_INLINE_FUNCTION int increment()
KOKKOS_INLINE_FUNCTION ScalarType get()
KOKKOS_INLINE_FUNCTION void setLocation(const Kokkos::Array< int, 7 > &location)
KOKKOS_INLINE_FUNCTION void reset(unsigned from_rank_number=0)