Tpetra parallel linear algebra  Version of the Day
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Pages
Tpetra_for_each_MultiVector.hpp
Go to the documentation of this file.
1 /*
2 // @HEADER
3 // ***********************************************************************
4 //
5 // Tpetra: Templated Linear Algebra Services Package
6 // Copyright (2008) Sandia Corporation
7 //
8 // Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
9 // the U.S. Government retains certain rights in this software.
10 //
11 // Redistribution and use in source and binary forms, with or without
12 // modification, are permitted provided that the following conditions are
13 // met:
14 //
15 // 1. Redistributions of source code must retain the above copyright
16 // notice, this list of conditions and the following disclaimer.
17 //
18 // 2. Redistributions in binary form must reproduce the above copyright
19 // notice, this list of conditions and the following disclaimer in the
20 // documentation and/or other materials provided with the distribution.
21 //
22 // 3. Neither the name of the Corporation nor the names of the
23 // contributors may be used to endorse or promote products derived from
24 // this software without specific prior written permission.
25 //
26 // THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
27 // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
28 // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
29 // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
30 // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
31 // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
32 // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
33 // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
34 // LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
35 // NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
36 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
37 //
38 // Questions? Contact Michael A. Heroux (maherou@sandia.gov)
39 //
40 // ************************************************************************
41 // @HEADER
42 */
43 
44 #ifndef TPETRA_FOR_EACH_MULTIVECTOR_HPP
45 #define TPETRA_FOR_EACH_MULTIVECTOR_HPP
46 
47 #include "Tpetra_for_each.hpp"
50 #include "Tpetra_Map.hpp"
51 #include "Teuchos_Comm.hpp"
52 #include "Teuchos_TypeNameTraits.hpp"
53 #include "Kokkos_Core.hpp"
54 #include <iostream>
55 #include <sstream>
56 
72 
73 namespace Tpetra {
74  namespace Details {
75 
76  // for_each uses this for the loop body of a parallel_for or
77  // parallel_reduce over the rows of a Tpetra::MultiVector with
78  // constant stride and multiple columns.
79  template<class ViewType,
80  class UserFunctionType,
81  class LocalIndexType>
82  class MultiVectorForEachLoopBody {
83  private:
84  static_assert (static_cast<int> (ViewType::Rank) == 2,
85  "ViewType must be a rank-2 Kokkos::View.");
86  public:
87  MultiVectorForEachLoopBody (const ViewType& X_lcl,
88  UserFunctionType f) :
89  X_lcl_ (X_lcl), f_ (f)
90  {}
91  KOKKOS_INLINE_FUNCTION void
92  operator () (const LocalIndexType i) const {
93  const LocalIndexType numCols =
94  static_cast<LocalIndexType> (X_lcl_.extent (1));
95  for (LocalIndexType j = 0; j < numCols; ++j) {
96  f_ (X_lcl_(i,j));
97  }
98  }
99  private:
100  ViewType X_lcl_;
101  UserFunctionType f_;
102  };
103 
104  // for_each uses this for the loop body of a parallel_for or
105  // parallel_reduce over the rows of a Tpetra::Vector (or each
106  // column of a Tpetra::MultiVector, if the Tpetra::MultiVector has
107  // nonconstant stride or only a single column).
108  template<class ViewType,
109  class UserFunctionType,
110  class LocalIndexType>
111  class VectorForEachLoopBody {
112  private:
113  static_assert (static_cast<int> (ViewType::Rank) == 1,
114  "ViewType must be a rank-1 Kokkos::View.");
115  public:
116  VectorForEachLoopBody (const ViewType& X_lcl,
117  UserFunctionType f) :
118  X_lcl_ (X_lcl), f_ (f)
119  {}
120  KOKKOS_INLINE_FUNCTION void
121  operator () (const LocalIndexType i) const {
122  f_ (X_lcl_(i));
123  }
124  private:
125  ViewType X_lcl_;
126  UserFunctionType f_;
127  };
128 
131  template<class ExecutionSpace,
132  class SC, class LO, class GO, class NT,
133  class UserFunctionType>
134  struct ForEach<ExecutionSpace,
135  ::Tpetra::MultiVector<SC, LO, GO, NT>,
136  UserFunctionType>
137  {
138  private:
139  // Given a Kokkos execution space on which the user wants to run
140  // the for_each, and a memory space in which the MultiVector's
141  // data live, determine the memory space that for_each should
142  // use in its withLocalAccess call.
143  template<class MemorySpace>
144  using for_each_memory_space =
145  typename std::conditional<
146  Kokkos::SpaceAccessibility<
147  ExecutionSpace,
148  typename MemorySpace::memory_space>::accessible,
149  typename MemorySpace::memory_space,
150  typename ExecutionSpace::memory_space>::type;
151 
152  public:
153  static void
154  for_each (const char kernelLabel[],
155  ExecutionSpace execSpace,
157  UserFunctionType f)
158  {
159  using Teuchos::TypeNameTraits;
160  using std::endl;
162  using preferred_memory_space =
163  typename MV::device_type::memory_space;
164  using memory_space = for_each_memory_space<preferred_memory_space>;
165  using range_type = Kokkos::RangePolicy<ExecutionSpace, LO>;
166 
167  const int myRank = X.getMap ()->getComm ()->getRank ();
168  const bool verbose = ::Tpetra::Details::Behavior::verbose ();
169  if (verbose) {
170  std::ostringstream os;
171  os << "Proc " << myRank << ": Tpetra::for_each:" << endl
172  << " kernelLabel: " << kernelLabel << endl
173  << " ExecutionSpace: "
174  << TypeNameTraits<ExecutionSpace>::name () << endl
175  << " memory_space: "
176  << TypeNameTraits<memory_space>::name () << endl;
177  std::cerr << os.str ();
178  }
179 
180  // Generic lambdas need C++14, so for now, we use
181  // with_local_access_function_argument_type to get a named
182  // typedef for our withLocalAccess lambda argument(s).
183 
184  range_type range (execSpace, 0, X.getLocalLength ());
185  memory_space memSpace;
186  if (X.getNumVectors () == size_t (1) || ! X.isConstantStride ()) {
187  const size_t numVecs = X.getNumVectors ();
188  for (size_t j = 0; j < numVecs; ++j) {
189  auto X_j = X.getVectorNonConst (j);
190  // Help GCC 4.9.3 deduce the type of *X_j.
191  // See discussion here:
192  // https://github.com/trilinos/Trilinos/pull/5115
193  Tpetra::Vector<SC, LO, GO, NT>& X_j_ref = *X_j;
194  using read_write_view_type =
196  decltype (readWrite (X_j_ref).on (memSpace). at(execSpace))>;
198  ([=] (const read_write_view_type& X_j_lcl) {
199  using functor_type = VectorForEachLoopBody<
200  read_write_view_type, UserFunctionType, LO>;
201  Kokkos::parallel_for (kernelLabel, range,
202  functor_type (X_j_lcl, f));
203  }, readWrite (X_j_ref).on (memSpace). at(execSpace));
204  }
205  }
206  else {
207  // Generic lambdas need C++14, so we need a typedef here.
208  using read_write_view_type =
210  decltype (readWrite (X).on (memSpace). at(execSpace))>;
212  ([=] (const read_write_view_type& X_lcl) {
213  using functor_type = MultiVectorForEachLoopBody<
214  read_write_view_type, UserFunctionType, LO>;
215  Kokkos::parallel_for (kernelLabel, range,
216  functor_type (X_lcl, f));
217  }, readWrite (X).on (memSpace). at(execSpace));
218  }
219  }
220  };
221 
228  template<class ExecutionSpace,
229  class SC, class LO, class GO, class NT,
230  class UserFunctionType>
231  struct ForEach<ExecutionSpace,
232  ::Tpetra::Vector<SC, LO, GO, NT>,
233  UserFunctionType>
234  {
235  static void
236  for_each (const char kernelLabel[],
237  ExecutionSpace execSpace,
239  UserFunctionType f)
240  {
243  impl_type::for_each (kernelLabel, execSpace, X, f);
244  }
245  };
246 
247  } // namespace Details
248 
249 } // namespace Tpetra
250 
251 #endif // TPETRA_FOR_EACH_MULTIVECTOR_HPP
252 
Include this file to make Tpetra::MultiVector and Tpetra::Vector work with Tpetra::withLocalAccess.
Declaration and definition of Tpetra::for_each; declaration of helper classes for users to specialize...
size_t getNumVectors() const
Number of columns in the multivector.
size_t getLocalLength() const
Local number of rows on the calling process.
bool isConstantStride() const
Whether this multivector has constant stride between columns.
One or more distributed dense vectors.
Teuchos::RCP< Vector< Scalar, LocalOrdinal, GlobalOrdinal, Node > > getVectorNonConst(const size_t j)
Return a Vector which is a nonconst view of column j.
Details::LocalAccess< GlobalObjectType, Details::read_write > readWrite(GlobalObjectType &)
Declare that you want to access the given global object&#39;s local data in read-and-write mode...
static bool verbose()
Whether Tpetra is in verbose mode.
void withLocalAccess(typename Details::ArgsToFunction< LocalAccessTypes...>::type userFunction, LocalAccessTypes...localAccesses)
Get access to a Tpetra global object&#39;s local data.
A distributed dense vector.
Specialize this class to implement Tpetra::for_each for specific GlobalDataStructure types...
virtual Teuchos::RCP< const map_type > getMap() const
The Map describing the parallel distribution of this object.
typename Details::GetNonowningLocalObject< LocalAccessType >::nonowning_local_object_type with_local_access_function_argument_type
Type of the local object, that is an argument to the function the user gives to withLocalAccess.
void for_each(const char kernelLabel[], ExecutionSpace execSpace, GlobalDataStructure &X, UserFunctionType f)
Apply a function entrywise to each local entry of a Tpetra global data structure, analogously to std:...
Declaration of Tpetra::Details::Behavior, a class that describes Tpetra&#39;s behavior.