enable GPU solves

Author: fritzgoebel

common/cuda_hip/distributed/preconditioner/bddc_kernels.cpp

@@ -6,9 +6,14 @@
 
 #include <algorithm>
 
+#include "common/cuda_hip/base/config.hpp"
+#include "common/cuda_hip/base/runtime.hpp"
+#include "common/cuda_hip/base/types.hpp"
+#include "common/cuda_hip/components/thread_ids.hpp"
 #include "core/base/allocator.hpp"
 #include "core/base/device_matrix_data_kernels.hpp"
 #include "core/base/iterator_factory.hpp"
+#include "core/base/utils.hpp"
 #include "core/components/prefix_sum_kernels.hpp"
 #include "reference/distributed/partition_helpers.hpp"
 
@@ -19,6 +24,51 @@ namespace GKO_DEVICE_NAMESPACE {
 namespace bddc {
 
 
+constexpr int default_block_size = 512;
+
+
+namespace kernel {
+
+
+template <typename ValueType, typename IndexType>
+__global__ __launch_bounds__(default_block_size) void generate_constraints(
+    const IndexType* interface_offsets, size_type n_inactive_idxs,
+    size_type n_edges_faces, IndexType* row_idxs, IndexType* col_idxs,
+    ValueType* values)
+{
+    const auto interface_idx = thread::get_thread_id_flat();
+    if (interface_idx < n_edges_faces) {
+        const auto start = interface_offsets[interface_idx];
+        const auto stop = interface_offsets[interface_idx + 1];
+        const ValueType val =
+            one<ValueType>() / static_cast<ValueType>(stop - start);
+        for (size_type idx = start; idx < stop; idx++) {
+            row_idxs[idx] = interface_idx;
+            col_idxs[idx] = n_inactive_idxs + idx;
+            values[idx] = val;
+        }
+    }
+}
+
+
+template <typename ValueType>
+__global__ __launch_bounds__(default_block_size) void fill_coarse_data(
+    size_type n_edges_faces, size_type n_corners, size_type lambda_stride,
+    size_type phi_stride, ValueType* lambda, ValueType* phi)
+{
+    const auto i = thread::get_thread_id_flat();
+    if (i < n_edges_faces) {
+        lambda[i * lambda_stride + i] = one<ValueType>();
+    }
+    if (i < n_corners) {
+        phi[i * phi_stride + n_edges_faces + i] = one<ValueType>();
+    }
+}
+
+
+}  // namespace kernel
+
+
 template <typename ValueType, typename IndexType>
 void classify_dofs(
     std::shared_ptr<const DefaultExecutor> exec,
@@ -38,11 +88,27 @@ GKO_INSTANTIATE_FOR_EACH_NON_COMPLEX_VALUE_AND_INDEX_TYPE_BASE(
 
 
 template <typename ValueType, typename IndexType>
-void generate_constraints(
-    std::shared_ptr<const DefaultExecutor> exec,
-    const matrix::Dense<ValueType>* labels, size_type n_inner_idxs,
-    size_type n_edges_faces, const array<IndexType>& interface_sizes,
-    device_matrix_data<ValueType, IndexType>& constraints) GKO_NOT_IMPLEMENTED;
+void generate_constraints(std::shared_ptr<const DefaultExecutor> exec,
+                          const matrix::Dense<ValueType>* labels,
+                          size_type n_inactive_idxs, size_type n_edges_faces,
+                          const array<IndexType>& interface_sizes,
+                          device_matrix_data<ValueType, IndexType>& constraints)
+{
+    array<IndexType> interface_offsets{exec, n_edges_faces + 1};
+    exec->copy(n_edges_faces, interface_sizes.get_const_data(),
+               interface_offsets.get_data());
+    components::prefix_sum_nonnegative(exec, interface_offsets.get_data(),
+                                       n_edges_faces + 1);
+
+    const auto grid_dim = ceildiv(n_edges_faces, default_block_size);
+    if (grid_dim > 0) {
+        kernel::generate_constraints<<<grid_dim, default_block_size>>>(
+            as_device_type(interface_offsets.get_const_data()), n_inactive_idxs,
+            n_edges_faces, as_device_type(constraints.get_row_idxs()),
+            as_device_type(constraints.get_col_idxs()),
+            as_device_type(constraints.get_values()));
+    }
+}
 
 GKO_INSTANTIATE_FOR_EACH_NON_COMPLEX_VALUE_AND_INDEX_TYPE_BASE(
     GKO_DECLARE_GENERATE_CONSTRAINTS);
@@ -51,14 +117,27 @@ GKO_INSTANTIATE_FOR_EACH_NON_COMPLEX_VALUE_AND_INDEX_TYPE_BASE(
 template <typename ValueType>
 void fill_coarse_data(std::shared_ptr<const DefaultExecutor> exec,
                       matrix::Dense<ValueType>* phi_P,
-                      matrix::Dense<ValueType>* lambda_rhs) GKO_NOT_IMPLEMENTED;
+                      matrix::Dense<ValueType>* lambda_rhs)
+{
+    const auto n_edges_faces = lambda_rhs->get_size()[0];
+    const auto n_corners = phi_P->get_size()[1];
+    const auto grid_dim =
+        ceildiv(std::max(n_edges_faces, n_corners), default_block_size);
+    if (grid_dim > 0) {
+        kernel::fill_coarse_data<<<grid_dim, default_block_size>>>(
+            n_edges_faces, n_corners, lambda_rhs->get_stride(),
+            phi_P->get_stride(), as_device_type(lambda_rhs->get_values()),
+            as_device_type(phi_P->get_values()));
+    }
+}
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE_BASE(GKO_DECLARE_FILL_COARSE_DATA);
 
 
 template <typename ValueType, typename IndexType>
 void build_coarse_contribution(
     std::shared_ptr<const DefaultExecutor> exec,
+    const array<experimental::distributed::preconditioner::dof_type>& dof_types,
     const array<remove_complex<ValueType>>& local_labels,
     const array<IndexType>& local_tags,
     const array<remove_complex<ValueType>>& global_labels,

core/distributed/assembly.cpp

@@ -81,17 +81,17 @@ device_matrix_data<ValueType, GlobalIndexType> assemble_rows_from_neighbors(
         recv_values.set_executor(exec->get_master());
     }
     auto row_req = comm.i_all_to_all_v(
-        use_host_buffer ? exec : exec->get_master(),
+        use_host_buffer ? exec->get_master() : exec,
         send_row_idxs.get_const_data(), send_sizes.get_data(),
         send_offsets.get_data(), recv_row_idxs.get_data(),
         recv_sizes.get_data(), recv_offsets.get_data());
     auto col_req = comm.i_all_to_all_v(
-        use_host_buffer ? exec : exec->get_master(),
+        use_host_buffer ? exec->get_master() : exec,
         send_col_idxs.get_const_data(), send_sizes.get_data(),
         send_offsets.get_data(), recv_col_idxs.get_data(),
         recv_sizes.get_data(), recv_offsets.get_data());
     auto val_req =
-        comm.i_all_to_all_v(use_host_buffer ? exec : exec->get_master(),
+        comm.i_all_to_all_v(use_host_buffer ? exec->get_master() : exec,
                             send_values.get_const_data(), send_sizes.get_data(),
                             send_offsets.get_data(), recv_values.get_data(),
                             recv_sizes.get_data(), recv_offsets.get_data());

core/distributed/dd_matrix.cpp

@@ -121,7 +121,6 @@ void DdMatrix<ValueType, LocalIndexType, GlobalIndexType>::read_distributed(
     GKO_ASSERT_EQ(comm.size(), partition->get_num_parts());
     auto exec = this->get_executor();
     auto local_part = comm.rank();
-    auto use_host_buffer = mpi::requires_host_buffer(exec, comm);
     auto tmp_partition = make_temporary_clone(exec, partition);
 
     // set up LinOp sizes
@@ -163,9 +162,9 @@ void DdMatrix<ValueType, LocalIndexType, GlobalIndexType>::read_distributed(
 
     // Gather local sizes from all ranks and build the partition in the enriched
     // space.
-    array<GlobalIndexType> range_bounds{
-        use_host_buffer ? exec->get_master() : exec, num_parts + 1};
-    comm.all_gather(exec, &local_num_rows, 1, range_bounds.get_data(), 1);
+    array<GlobalIndexType> range_bounds{exec->get_master(), num_parts + 1};
+    comm.all_gather(exec->get_master(), &local_num_rows, 1,
+                    range_bounds.get_data(), 1);
     range_bounds.set_executor(exec);
     exec->run(dd_matrix::make_prefix_sum_nonnegative(range_bounds.get_data(),
                                                      num_parts + 1));