49 #ifndef Intrepid2_TensorViewIterator_h
50 #define Intrepid2_TensorViewIterator_h
55 #include <Kokkos_Core.hpp>
73 template<
class TensorViewType,
class ViewType1,
class ViewType2 ,
typename ScalarType>
77 enum RankCombinationType :
int
83 using RankCombinationViewType = Kokkos::View<RankCombinationType*, typename TensorViewType::device_type>;
90 RankCombinationViewType rank_combination_types_;
111 KOKKOS_INLINE_FUNCTION
113 RankCombinationViewType rank_combination_types)
115 tensor_view_iterator_(tensor_view),
116 view1_iterator_(view1),
117 view2_iterator_(view2),
118 rank_combination_types_(rank_combination_types)
136 unsigned max_component_rank = (view1.rank() > view2.rank()) ? view1.rank() : view2.rank();
137 unsigned max_rank = (tensor_view.rank() > max_component_rank) ? tensor_view.rank() : max_component_rank;
141 unsigned expected_rank = 0;
142 bool contracting =
false;
143 for (
unsigned d=0; d<rank_combination_types.extent(0); d++)
145 if (rank_combination_types[d] == TENSOR_CONTRACTION)
156 if (rank_combination_types[d] == TENSOR_PRODUCT)
158 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");
162 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");
163 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");
177 KOKKOS_INLINE_FUNCTION
182 if (view2_next_increment_rank > view1_next_increment_rank)
return view2_next_increment_rank;
183 else return view1_next_increment_rank;
188 KOKKOS_INLINE_FUNCTION
203 if (view2_next_increment_rank > view1_next_increment_rank)
206 device_assert(rank_combination_types_[view2_next_increment_rank]==TENSOR_PRODUCT);
207 view1_iterator_.
reset(view2_next_increment_rank);
210 return view2_next_increment_rank;
214 int view1_rank_change = view1_iterator_.
increment();
215 if (view1_rank_change >= 0)
217 switch (rank_combination_types_[view1_rank_change])
219 case DIMENSION_MATCH:
227 case TENSOR_CONTRACTION:
232 return view1_rank_change;
238 KOKKOS_INLINE_FUNCTION
249 KOKKOS_INLINE_FUNCTION
250 void setLocation(Kokkos::Array<int,7> location1, Kokkos::Array<int,7> location2)
254 Kokkos::Array<int,7> tensor_location = location1;
255 for (
unsigned d=0; d<rank_combination_types_.extent(0); d++)
257 switch (rank_combination_types_[d])
261 tensor_location[d] = location2[d] * view1_iterator_.
getExtent(d) + location1[d];
263 case DIMENSION_MATCH:
266 case TENSOR_CONTRACTION:
267 tensor_location[d] = 0;
271 #ifdef HAVE_INTREPID2_DEBUG
273 for (
unsigned d=0; d<rank_combination_types_.extent(0); d++)
275 switch (rank_combination_types_[d])
280 case DIMENSION_MATCH:
281 case TENSOR_CONTRACTION:
282 device_assert(location1[d] == location2[d]);
286 device_assert(location1[d] < view1_iterator_.
getExtent(d));
287 device_assert(location2[d] < view2_iterator_.
getExtent(d));
288 device_assert(tensor_location[d] < tensor_view_iterator_.
getExtent(d));
291 tensor_view_iterator_.
setLocation(tensor_location);
296 KOKKOS_INLINE_FUNCTION
299 return view1_iterator_.
get();
304 KOKKOS_INLINE_FUNCTION
307 return view2_iterator_.
get();
312 KOKKOS_INLINE_FUNCTION
313 void set(ScalarType value)
315 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.
#define INTREPID2_TEST_FOR_EXCEPTION_DEVICE_SAFE(test, x, msg)
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 void set(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...