Merge Pull Request #7105 from brian-kelley/Trilinos/FixKKCrsGraphType

Automatically Merged using Trilinos Pull Request AutoTester
PR Title: KokkosKernels/Tpetra stack: local CRS device types
PR Author: brian-kelley

packages/kokkos-kernels/src/sparse/KokkosSparse_CrsMatrix.hpp

@@ -389,7 +389,7 @@ class CrsMatrix {
   typedef typename Device::execution_space execution_space;
   //! Type of the matrix's memory space.
   typedef typename Device::memory_space memory_space;
-  //! Type of the matrix's device type.
+  //! Canonical device type
   typedef Kokkos::Device<execution_space, memory_space> device_type;
 
   //! Type of each value in the matrix.
@@ -407,14 +407,14 @@ class CrsMatrix {
   typedef CrsMatrix<ScalarType, OrdinalType, host_mirror_space, MemoryTraits> HostMirror;
 #ifdef KOKKOS_ENABLE_DEPRECATED_CODE
   //! Type of the graph structure of the sparse matrix.
-  typedef Kokkos::StaticCrsGraph<ordinal_type, Kokkos::LayoutLeft, execution_space, size_type, memory_traits> StaticCrsGraphType;
+  typedef Kokkos::StaticCrsGraph<ordinal_type, Kokkos::LayoutLeft, device_type, size_type, memory_traits> StaticCrsGraphType;
   //! Type of the graph structure of the sparse matrix - consistent with Kokkos.
-  typedef Kokkos::StaticCrsGraph<ordinal_type, Kokkos::LayoutLeft, execution_space, size_type, memory_traits> staticcrsgraph_type;
+  typedef Kokkos::StaticCrsGraph<ordinal_type, Kokkos::LayoutLeft, device_type, size_type, memory_traits> staticcrsgraph_type;
 #else
   //! Type of the graph structure of the sparse matrix.
-  typedef Kokkos::StaticCrsGraph<ordinal_type, Kokkos::LayoutLeft, execution_space, memory_traits, size_type> StaticCrsGraphType;
+  typedef Kokkos::StaticCrsGraph<ordinal_type, Kokkos::LayoutLeft, device_type, memory_traits, size_type> StaticCrsGraphType;
   //! Type of the graph structure of the sparse matrix - consistent with Kokkos.
-  typedef Kokkos::StaticCrsGraph<ordinal_type, Kokkos::LayoutLeft, execution_space, memory_traits, size_type> staticcrsgraph_type;
+  typedef Kokkos::StaticCrsGraph<ordinal_type, Kokkos::LayoutLeft, device_type, memory_traits, size_type> staticcrsgraph_type;
 #endif
   //! Type of column indices in the sparse matrix.
   typedef typename staticcrsgraph_type::entries_type index_type;

packages/muelu/src/Graph/Containers/MueLu_LWGraph_kokkos_decl.hpp

@@ -79,14 +79,15 @@ namespace MueLu {
     using global_ordinal_type = GlobalOrdinal;
     using execution_space     = typename DeviceType::execution_space;
     using memory_space        = typename DeviceType::memory_space;
+    using device_type         = Kokkos::Device<execution_space, memory_space>;
     using range_type          = Kokkos::RangePolicy<local_ordinal_type, execution_space>;
     using node_type           = Kokkos::Compat::KokkosDeviceWrapperNode<DeviceType>;
     using size_type           = size_t;
 
     using map_type            = Xpetra::Map<LocalOrdinal, GlobalOrdinal, node_type>;
     using local_graph_type    = Kokkos::StaticCrsGraph<LocalOrdinal,
                                                        Kokkos::LayoutLeft,
-                                                       execution_space>;
+                                                       device_type>;
     using boundary_nodes_type = Kokkos::View<const bool*, memory_space>;
     using row_type            = Kokkos::View<const LocalOrdinal*, memory_space>;
 

packages/phalanx/example/FiniteElementAssembly/LinearObjectFactory.hpp

@@ -18,6 +18,8 @@ namespace phx_example {
     // Phalanx device spaces
     using exec_t = PHX::exec_space;
     using mem_t = PHX::mem_space;
+    using local_matrix_type = KokkosSparse::CrsMatrix<double,int,PHX::Device>;
+    using local_graph_type = typename local_matrix_type::StaticCrsGraphType;
     using team_t =  Kokkos::TeamPolicy<exec_t>::member_type;
 
     const int num_nodes_;       //! Total number of nodes in mesh
@@ -36,7 +38,7 @@ namespace phx_example {
     Kokkos::View<int*,PHX::Device> row_count_;
 
     //! Row offsets for building graph, sizeof(num_dofs_+1)
-    using row_offset_size_t = typename Kokkos::StaticCrsGraph<int,Kokkos::LayoutLeft,PHX::Device>::size_type;
+    using row_offset_size_t = typename local_graph_type::size_type;
     Kokkos::View<row_offset_size_t*,PHX::Device> row_offsets_;
 
     //! Single value for the size of the matrix
@@ -46,7 +48,7 @@ namespace phx_example {
     Kokkos::View<int*, PHX::Device> col_ids_;
 
     //! Crs Graph for Jacobian
-    Kokkos::StaticCrsGraph<int,Kokkos::LayoutLeft,PHX::Device> graph_;
+    local_graph_type graph_;
 
     public:
 
@@ -98,7 +100,7 @@ namespace phx_example {
       Kokkos::deep_copy(row_count_,0); // reuse this view
       Kokkos::fence();
       graph_.row_map = row_offsets_;
-      graph_.entries = Kokkos::StaticCrsGraph<int,Kokkos::LayoutLeft,PHX::Device>::entries_type("graph_entries",num_matrix_entries_);
+      graph_.entries = local_graph_type::entries_type("graph_entries",num_matrix_entries_);
       Kokkos::parallel_for(Kokkos::RangePolicy<exec_t,TagFillGraphEntries>(0,global_node_set_.capacity()), *this );
       exec_t().fence();
 
@@ -123,8 +125,8 @@ namespace phx_example {
     Kokkos::View<double*,PHX::Device> createSolutionVector(const std::string& name) const
     { return Kokkos::View<double*,PHX::Device>(name,num_dofs_); }
 
-    KokkosSparse::CrsMatrix<double,int,PHX::Device> createJacobianMatrix(const std::string& name) const
-    { return KokkosSparse::CrsMatrix<double,int,PHX::Device>(name,graph_); }
+    local_matrix_type createJacobianMatrix(const std::string& name) const
+    { return local_matrix_type(name,graph_); }
 
 
     // ****************************************************

packages/phalanx/test/Kokkos/tKokkos.cpp

@@ -752,7 +752,10 @@ namespace phalanx_test {
     const size_t nnz = 21;
 
     using Kokkos::View;
-    View<typename Kokkos::StaticCrsGraph<int,Kokkos::LayoutLeft,PHX::Device>::size_type*,PHX::Device> row_offsets("row_offsets",num_rows+1);
+    using local_matrix_type = KokkosSparse::CrsMatrix<double,int,PHX::Device>;
+    using local_graph_type = typename local_matrix_type::StaticCrsGraphType;
+    using size_type = typename local_matrix_type::size_type;
+    View<size_type*,PHX::Device> row_offsets("row_offsets",num_rows+1);
 
     auto host_row_offsets = Kokkos::create_mirror_view(row_offsets);
     host_row_offsets(0) = 0;
@@ -786,8 +789,8 @@ namespace phalanx_test {
     Kokkos::deep_copy(col_ids, host_col_ids);
 
     // CrsMatrix requires LayoutLeft!
-    Kokkos::StaticCrsGraph<int,Kokkos::LayoutLeft,PHX::Device> g(col_ids,row_offsets);
-    KokkosSparse::CrsMatrix<double,int,PHX::Device> J("Jacobian",g);
+    local_graph_type g(col_ids,row_offsets);
+    local_matrix_type J("Jacobian",g);
 
     TEST_EQUALITY(J.numRows(),10);
     TEST_EQUALITY(J.numCols(),10);

packages/stokhos/src/sacado/kokkos/pce/linalg/Kokkos_CrsMatrix_UQ_PCE.hpp

@@ -88,7 +88,7 @@ class Multiply< KokkosSparse::CrsMatrix< Sacado::UQ::PCE<MatrixStorage>,
   typedef Sacado::UQ::PCE<InputStorage> InputVectorValue;
   typedef Sacado::UQ::PCE<OutputStorage> OutputVectorValue;
 
-  typedef MatrixDevice execution_space;
+  typedef typename MatrixDevice::execution_space execution_space;
 
   typedef KokkosSparse::CrsMatrix< MatrixValue,
                                    MatrixOrdinal,
@@ -485,7 +485,7 @@ class Multiply< KokkosSparse::CrsMatrix< Sacado::UQ::PCE<MatrixStorage>,
   typedef Sacado::UQ::PCE<InputStorage> InputVectorValue;
   typedef Sacado::UQ::PCE<OutputStorage> OutputVectorValue;
 
-  typedef MatrixDevice execution_space;
+  typedef typename MatrixDevice::execution_space execution_space;
 
   typedef KokkosSparse::CrsMatrix< MatrixValue,
                                    MatrixOrdinal,
@@ -929,7 +929,7 @@ class MeanMultiply< KokkosSparse::CrsMatrix< Sacado::UQ::PCE<MatrixStorage>,
 
   template <int BlockSize>
   struct BlockKernel {
-    typedef MatrixDevice execution_space;
+    typedef typename MatrixDevice::execution_space execution_space;
     typedef typename Kokkos::FlatArrayType<matrix_values_type>::type matrix_array_type;
     typedef typename input_vector_type::array_type input_array_type;
     typedef typename output_vector_type::array_type output_array_type;
@@ -1052,7 +1052,7 @@ class MeanMultiply< KokkosSparse::CrsMatrix< Sacado::UQ::PCE<MatrixStorage>,
   };
 
   struct Kernel {
-    typedef MatrixDevice execution_space;
+    typedef typename MatrixDevice::execution_space execution_space;
     typedef typename Kokkos::FlatArrayType<matrix_values_type>::type matrix_array_type;
     typedef typename input_vector_type::array_type input_array_type;
     typedef typename output_vector_type::array_type output_array_type;
@@ -1189,7 +1189,7 @@ class MeanMultiply< KokkosSparse::CrsMatrix< Sacado::UQ::PCE<MatrixStorage>,
   typedef Kokkos::View< OutputVectorValue**,
                         OutputP... > output_vector_type;
 
-  typedef MatrixDevice execution_space;
+  typedef typename MatrixDevice::execution_space execution_space;
   typedef typename MatrixValue::ordinal_type size_type;
   typedef typename InputVectorValue::value_type input_scalar;
   typedef typename OutputVectorValue::value_type output_scalar;

packages/stokhos/src/sacado/kokkos/vector/linalg/Kokkos_CrsMatrix_MP_Vector.hpp

@@ -120,7 +120,7 @@ class MPMultiply< KokkosSparse::CrsMatrix< Sacado::MP::Vector<MatrixStorage>,
                   Update
 #ifdef KOKKOS_ENABLE_CUDA
                   , typename std::enable_if<
-                      !std::is_same<MatrixDevice,Kokkos::Cuda>::value >::type
+                      !std::is_same<typename MatrixDevice::execution_space,Kokkos::Cuda>::value >::type
 #endif
                   >
 {
@@ -132,7 +132,7 @@ class MPMultiply< KokkosSparse::CrsMatrix< Sacado::MP::Vector<MatrixStorage>,
   typedef Sacado::MP::Vector<OutputStorage> OutputVectorValue;
   typedef OutputVectorValue scalar_type;
 
-  typedef MatrixDevice execution_space;
+  typedef typename MatrixDevice::execution_space execution_space;
   typedef typename execution_space::size_type size_type;
 
   typedef KokkosSparse::CrsMatrix< MatrixValue,
@@ -214,7 +214,7 @@ class MPMultiply< KokkosSparse::CrsMatrix< Sacado::MP::Vector<MatrixStorage>,
                   Update
 #ifdef KOKKOS_ENABLE_CUDA
                   , typename std::enable_if<
-                    !std::is_same<MatrixDevice,Kokkos::Cuda>::value >::type
+                    !std::is_same<typename MatrixDevice::execution_space,Kokkos::Cuda>::value >::type
 #endif
                   >
 {
@@ -224,7 +224,7 @@ class MPMultiply< KokkosSparse::CrsMatrix< Sacado::MP::Vector<MatrixStorage>,
   typedef Sacado::MP::Vector<OutputStorage> OutputVectorValue;
   typedef OutputVectorValue scalar_type;
 
-  typedef MatrixDevice execution_space;
+  typedef typename MatrixDevice::execution_space execution_space;
   typedef typename execution_space::size_type size_type;
 
   typedef KokkosSparse::CrsMatrix< MatrixValue,
@@ -319,7 +319,7 @@ class Multiply< KokkosSparse::CrsMatrix< Sacado::MP::Vector<MatrixStorage>,
   typedef Sacado::MP::Vector<InputStorage> InputVectorValue;
   typedef Sacado::MP::Vector<OutputStorage> OutputVectorValue;
 
-  typedef MatrixDevice execution_space;
+  typedef typename MatrixDevice::execution_space execution_space;
   typedef typename execution_space::size_type size_type;
 
   typedef KokkosSparse::CrsMatrix< MatrixValue,
@@ -376,7 +376,7 @@ class Multiply< KokkosSparse::CrsMatrix< Sacado::MP::Vector<MatrixStorage>,
   typedef Sacado::MP::Vector<InputStorage> InputVectorValue;
   typedef Sacado::MP::Vector<OutputStorage> OutputVectorValue;
 
-  typedef MatrixDevice execution_space;
+  typedef typename MatrixDevice::execution_space execution_space;
   typedef typename execution_space::size_type size_type;
 
   typedef KokkosSparse::CrsMatrix< MatrixValue,

packages/stokhos/src/sacado/kokkos/vector/linalg/Kokkos_CrsMatrix_MP_Vector_Cuda.hpp

@@ -89,7 +89,7 @@ template <typename MatrixStorage,
           typename Update>
 class MPMultiply< KokkosSparse::CrsMatrix<Sacado::MP::Vector<MatrixStorage>,
                                     MatrixOrdinal,
-                                    Kokkos::Cuda,
+                                    Kokkos::Device<Kokkos::Cuda, typename Kokkos::Cuda::memory_space>,
                                     MatrixMemory,
                                     MatrixSize>,
                   Kokkos::View< const Sacado::MP::Vector<InputStorage>*,
@@ -106,13 +106,13 @@ class MPMultiply< KokkosSparse::CrsMatrix<Sacado::MP::Vector<MatrixStorage>,
   typedef Sacado::MP::Vector<InputStorage> InputVectorValue;
   typedef Sacado::MP::Vector<OutputStorage> OutputVectorValue;
 
-  typedef typename Kokkos::Cuda Device;
-  typedef Device execution_space;
+  typedef Kokkos::Device<Kokkos::Cuda, typename Kokkos::Cuda::memory_space> Device;
+  typedef typename Device::execution_space execution_space;
   typedef typename execution_space::size_type size_type;
 
   typedef KokkosSparse::CrsMatrix<MatrixValue,
                             MatrixOrdinal,
-                            execution_space,
+                            Device,
                             MatrixMemory,
                             MatrixSize> matrix_type;
   typedef typename matrix_type::values_type matrix_values_type;
@@ -336,7 +336,7 @@ template <typename MatrixStorage,
           typename Update>
 class MPMultiply< KokkosSparse::CrsMatrix<Sacado::MP::Vector<MatrixStorage>,
                                     MatrixOrdinal,
-                                    Kokkos::Cuda,
+                                    Kokkos::Device<Kokkos::Cuda, typename Kokkos::Cuda::memory_space>,
                                     MatrixMemory,
                                     MatrixSize>,
                   Kokkos::View< const Sacado::MP::Vector<InputStorage>**,
@@ -353,14 +353,14 @@ class MPMultiply< KokkosSparse::CrsMatrix<Sacado::MP::Vector<MatrixStorage>,
   typedef Sacado::MP::Vector<InputStorage> InputVectorValue;
   typedef Sacado::MP::Vector<OutputStorage> OutputVectorValue;
 
-  typedef typename Kokkos::Cuda Device;
-  typedef Device execution_space;
+  typedef Kokkos::Device<Kokkos::Cuda, typename Kokkos::Cuda::memory_space> Device;
+  typedef typename Device::execution_space execution_space;
   typedef typename execution_space::size_type size_type;
 
 
   typedef KokkosSparse::CrsMatrix<MatrixValue,
                             MatrixOrdinal,
-                            execution_space,
+                            Device,
                             MatrixMemory,
                             MatrixSize> matrix_type;
   typedef typename matrix_type::values_type matrix_values_type;

packages/stokhos/test/Performance/MPVectorKernels/TestSpMv.hpp

@@ -109,10 +109,11 @@ test_mpvector_spmv(const int ensemble_length,
   typedef typename storage_type::value_type value_type;
   typedef typename storage_type::ordinal_type ordinal_type;
   typedef typename storage_type::execution_space execution_space;
+  typedef Kokkos::Device<execution_space, typename execution_space::memory_space> device_type;
   typedef Sacado::MP::Vector<StorageType> VectorType;
   typedef Kokkos::LayoutRight Layout;
   typedef Kokkos::View< VectorType*, Layout, execution_space > vector_type;
-  typedef KokkosSparse::CrsMatrix< VectorType, ordinal_type, execution_space > matrix_type;
+  typedef KokkosSparse::CrsMatrix< VectorType, ordinal_type, device_type > matrix_type;
   typedef typename matrix_type::StaticCrsGraphType matrix_graph_type;
   typedef typename matrix_type::values_type matrix_values_type;
 
@@ -191,8 +192,9 @@ test_scalar_spmv(const int ensemble_length,
   typedef ScalarType value_type;
   typedef OrdinalType ordinal_type;
   typedef Device execution_space;
+  typedef Kokkos::Device<execution_space, typename execution_space::memory_space> device_type;
   typedef Kokkos::View< value_type*, execution_space > vector_type;
-  typedef KokkosSparse::CrsMatrix< value_type, ordinal_type, execution_space > matrix_type;
+  typedef KokkosSparse::CrsMatrix< value_type, ordinal_type, device_type > matrix_type;
   typedef typename matrix_type::StaticCrsGraphType matrix_graph_type;
   typedef typename matrix_type::values_type matrix_values_type;
 

packages/stokhos/test/UnitTest/Stokhos_KokkosCrsMatrixMPVectorUnitTest.hpp

@@ -351,7 +351,8 @@ TEUCHOS_UNIT_TEST_TEMPLATE_1_DECL(
   Kokkos_CrsMatrix_MP, ReplaceValues, MatrixScalar )
 {
   typedef typename MatrixScalar::ordinal_type Ordinal;
-  typedef typename MatrixScalar::execution_space Device;
+  typedef typename MatrixScalar::execution_space execution_space;
+  typedef Kokkos::Device<execution_space, typename execution_space::memory_space> Device;
   typedef KokkosSparse::CrsMatrix<MatrixScalar,Ordinal,Device> Matrix;
 
   // Build diagonal matrix
@@ -370,7 +371,8 @@ TEUCHOS_UNIT_TEST_TEMPLATE_1_DECL(
   Kokkos_CrsMatrix_MP, SumIntoValues, MatrixScalar )
 {
   typedef typename MatrixScalar::ordinal_type Ordinal;
-  typedef typename MatrixScalar::execution_space Device;
+  typedef typename MatrixScalar::execution_space execution_space;
+  typedef Kokkos::Device<execution_space, typename execution_space::memory_space> Device;
   typedef KokkosSparse::CrsMatrix<MatrixScalar,Ordinal,Device> Matrix;
 
   // Build diagonal matrix
@@ -389,7 +391,8 @@ TEUCHOS_UNIT_TEST_TEMPLATE_1_DECL(
   Kokkos_CrsMatrix_MP, SumIntoValuesAtomic, MatrixScalar )
 {
   typedef typename MatrixScalar::ordinal_type Ordinal;
-  typedef typename MatrixScalar::execution_space Device;
+  typedef typename MatrixScalar::execution_space execution_space;
+  typedef Kokkos::Device<execution_space, typename execution_space::memory_space> Device;
   typedef KokkosSparse::CrsMatrix<MatrixScalar,Ordinal,Device> Matrix;
 
   // Build diagonal matrix
@@ -444,9 +447,10 @@ bool test_embedded_vector(const typename VectorType::ordinal_type nGrid,
   typedef typename VectorType::value_type scalar_type;
   typedef typename VectorType::storage_type storage_type;
   typedef typename storage_type::execution_space execution_space;
+  typedef Kokkos::Device<execution_space, typename execution_space::memory_space> device_type;
   typedef Kokkos::LayoutRight Layout;
   typedef Kokkos::View< VectorType*, Layout, execution_space > block_vector_type;
-  typedef KokkosSparse::CrsMatrix< VectorType, ordinal_type, execution_space > block_matrix_type;
+  typedef KokkosSparse::CrsMatrix< VectorType, ordinal_type, device_type > block_matrix_type;
   typedef typename block_matrix_type::StaticCrsGraphType matrix_graph_type;
   typedef typename block_matrix_type::values_type matrix_values_type;
 

packages/tpetra/core/src/Tpetra_CrsGraph_decl.hpp

@@ -256,7 +256,7 @@ namespace Tpetra {
     //! The type of the part of the sparse graph on each MPI process.
     using local_graph_type = Kokkos::StaticCrsGraph<local_ordinal_type,
                                                     Kokkos::LayoutLeft,
-                                                    execution_space>;
+                                                    device_type>;
 
     //! The Map specialization used by this class.
     using map_type = ::Tpetra::Map<LocalOrdinal, GlobalOrdinal, Node>;

packages/tpetra/core/src/Tpetra_CrsMatrix_decl.hpp

@@ -492,11 +492,10 @@ namespace Tpetra {
     using local_matrix_type =
       KokkosSparse::CrsMatrix<impl_scalar_type,
                               local_ordinal_type,
-                              execution_space,
+                              device_type,
                               void,
                               typename local_graph_type::size_type>;
 
-
     //@}
     //! @name Constructors and destructor
     //@{

packages/tpetra/core/src/Tpetra_LocalCrsMatrixOperator_decl.hpp

@@ -64,27 +64,22 @@ namespace Tpetra {
       typename LocalOperator<MultiVectorScalar, Device>::scalar_type;
     using matrix_scalar_type =
       typename LocalOperator<MatrixScalar, Device>::scalar_type;
-
-  private:
     using array_layout =
       typename LocalOperator<MultiVectorScalar, Device>::array_layout;
     using device_type =
       typename LocalOperator<MultiVectorScalar, Device>::device_type;
     using local_ordinal_type =
       ::Tpetra::Details::DefaultTypes::local_ordinal_type;
     using execution_space = typename Device::execution_space;
-    using local_graph_type =
-      Kokkos::StaticCrsGraph<local_ordinal_type,
-                             array_layout,
-                             execution_space>;
   public:
     using local_matrix_type =
       KokkosSparse::CrsMatrix<matrix_scalar_type,
                               local_ordinal_type,
-                              execution_space,
-                              void,
-                              typename local_graph_type::size_type>;
+                              device_type>;
+  private:
+    using local_graph_type = typename local_matrix_type::StaticCrsGraphType;
 
+  public:
     LocalCrsMatrixOperator (const std::shared_ptr<local_matrix_type>& A);
     ~LocalCrsMatrixOperator () override = default;