Skip to content

Commit

Permalink
Add row and column scaling
Browse files Browse the repository at this point in the history
  • Loading branch information
@fritzgoebel
fritzgoebel committed Oct 31, 2024
1 parent 7caffaf commit 1c0e4e9
Show file tree
Hide file tree
Showing 3 changed files with 192 additions and 27 deletions.
73 changes: 46 additions & 27 deletions core/distributed/dd_matrix.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -309,21 +309,12 @@ void DdMatrix<ValueType, LocalIndexType, GlobalIndexType>::apply_impl(
const LinOp* b, LinOp* x) const
{
auto exec = this->get_executor();
auto comm = this->get_communicator();
const auto nrhs = x->get_size()[1];
if (nrhs != rhs_buffer_->get_size()[1]) {
dim<2> local_buffer_size{rhs_buffer_->get_local_vector()->get_size()[0],
nrhs};
dim<2> global_buffer_size{rhs_buffer_->get_size()[0], nrhs};
lhs_buffer_ = Vector<ValueType>::create(exec, comm, global_buffer_size,
local_buffer_size);
rhs_buffer_ = Vector<ValueType>::create(exec, comm, global_buffer_size,
local_buffer_size);
}
check_and_adjust_buffer_size(nrhs);
distributed::precision_dispatch_real_complex<ValueType>(
[this](const auto dense_b, auto dense_x) {
auto exec = this->get_executor();
this->restriction_->apply(dense_b, lhs_buffer_);
restriction_->apply(dense_b, lhs_buffer_);

auto local_b = gko::matrix::Dense<ValueType>::create(
exec, lhs_buffer_->get_local_vector()->get_size(),
Expand All @@ -342,7 +333,7 @@ void DdMatrix<ValueType, LocalIndexType, GlobalIndexType>::apply_impl(

local_mtx_->apply(local_b, local_x);

this->prolongation_->apply(rhs_buffer_, dense_x);
prolongation_->apply(rhs_buffer_, dense_x);
},
b, x);
}
Expand All @@ -353,22 +344,13 @@ void DdMatrix<ValueType, LocalIndexType, GlobalIndexType>::apply_impl(
const LinOp* alpha, const LinOp* b, const LinOp* beta, LinOp* x) const
{
auto exec = this->get_executor();
auto comm = this->get_communicator();
const auto nrhs = x->get_size()[1];
if (nrhs != rhs_buffer_->get_size()[1]) {
dim<2> local_buffer_size{rhs_buffer_->get_local_vector()->get_size()[0],
nrhs};
dim<2> global_buffer_size{rhs_buffer_->get_size()[0], nrhs};
lhs_buffer_ = Vector<ValueType>::create(exec, comm, global_buffer_size,
local_buffer_size);
rhs_buffer_ = Vector<ValueType>::create(exec, comm, global_buffer_size,
local_buffer_size);
}
check_and_adjust_buffer_size(nrhs);
distributed::precision_dispatch_real_complex<ValueType>(
[this](const auto local_alpha, const auto dense_b,
const auto local_beta, auto dense_x) {
auto exec = this->get_executor();
this->restriction_->apply(dense_b, lhs_buffer_);
restriction_->apply(dense_b, lhs_buffer_);

auto local_b = gko::matrix::Dense<ValueType>::create(
exec, lhs_buffer_->get_local_vector()->get_size(),
Expand All @@ -387,23 +369,60 @@ void DdMatrix<ValueType, LocalIndexType, GlobalIndexType>::apply_impl(

local_mtx_->apply(local_b, local_x);

this->prolongation_->apply(local_alpha, rhs_buffer_, local_beta,
dense_x);
prolongation_->apply(local_alpha, rhs_buffer_, local_beta, dense_x);
},
alpha, b, beta, x);
}


template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
void DdMatrix<ValueType, LocalIndexType, GlobalIndexType>::
check_and_adjust_buffer_size(const size_type nrhs) const
{
auto exec = this->get_executor();
auto comm = this->get_communicator();
if (nrhs != rhs_buffer_->get_size()[1]) {
dim<2> local_buffer_size{rhs_buffer_->get_local_vector()->get_size()[0],
nrhs};
dim<2> global_buffer_size{rhs_buffer_->get_size()[0], nrhs};
lhs_buffer_ = Vector<ValueType>::create(exec, comm, global_buffer_size,
local_buffer_size);
rhs_buffer_ = Vector<ValueType>::create(exec, comm, global_buffer_size,
local_buffer_size);
}
}


template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
void DdMatrix<ValueType, LocalIndexType, GlobalIndexType>::col_scale(
ptr_param<const global_vector_type> scaling_factors)
{}
{
auto exec = this->get_executor();
check_and_adjust_buffer_size(1u);
size_type n_local_cols = local_mtx_->get_size()[1];
restriction_->apply(scaling_factors, lhs_buffer_);
const auto scale_diag = gko::matrix::Diagonal<ValueType>::create_const(
exec, n_local_cols,
make_const_array_view(exec, n_local_cols,
lhs_buffer_->get_const_local_values()));
scale_diag->rapply(local_mtx_, local_mtx_);
}


template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
void DdMatrix<ValueType, LocalIndexType, GlobalIndexType>::row_scale(
ptr_param<const global_vector_type> scaling_factors)
{}
{
auto exec = this->get_executor();
check_and_adjust_buffer_size(1u);
size_type n_local_cols = local_mtx_->get_size()[1];
restriction_->apply(scaling_factors, lhs_buffer_);
const auto scale_diag = gko::matrix::Diagonal<ValueType>::create_const(
exec, n_local_cols,
make_const_array_view(exec, n_local_cols,
lhs_buffer_->get_const_local_values()));
scale_diag->apply(local_mtx_, local_mtx_);
}


template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
Expand Down
2 changes: 2 additions & 0 deletions include/ginkgo/core/distributed/dd_matrix.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -436,6 +436,8 @@ class DdMatrix
LinOp* x) const override;

private:
void check_and_adjust_buffer_size(const size_type nrhs) const;

std::vector<comm_index_type> send_offsets_;
std::vector<comm_index_type> send_sizes_;
std::vector<comm_index_type> recv_offsets_;
Expand Down
144 changes: 144 additions & 0 deletions test/mpi/dd_matrix.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,7 @@ TYPED_TEST(DdMatrix, ReadsDistributed)
non_local_prolongation[rank], 0);
}


TYPED_TEST(DdMatrix, CanApplyToSingleVector)
{
using value_type = typename TestFixture::value_type;
Expand All @@ -185,6 +186,38 @@ TYPED_TEST(DdMatrix, CanApplyToSingleVector)
GKO_ASSERT_MTX_NEAR(this->y->get_local_vector(), result[rank], 0);
}


TYPED_TEST(DdMatrix, CanApplyToMultipleVectors)
{
using value_type = typename TestFixture::value_type;
using index_type = typename TestFixture::global_index_type;
auto vec_md =
gko::matrix_data<value_type, index_type>{I<I<value_type>>{{1, 2},
{2, 4},
{3, 6},
{4, 8},
{5, 10},
{6, 12},
{7, 14},
{8, 16},
{9, 18},
{10, 20},
{11, 22},
{12, 24}}};
I<I<value_type>> result[3] = {{{-4, -8}, {-3, -6}, {-2, -4}, {-1, -2}},
{{0, 0}, {1, 2}, {-1, -2}, {0, 0}},
{{1, 2}, {2, 4}, {3, 6}, {4, 8}}};
auto rank = this->comm.rank();
this->dist_mat->read_distributed(this->dist_input[rank], this->row_part);
this->x->read_distributed(vec_md, this->row_part);
this->y->read_distributed(vec_md, this->row_part);

this->dist_mat->apply(this->x, this->y);

GKO_ASSERT_MTX_NEAR(this->y->get_local_vector(), result[rank], 0);
}


TYPED_TEST(DdMatrix, CanAdvancedApplyToSingleVector)
{
using value_type = typename TestFixture::value_type;
Expand All @@ -206,4 +239,115 @@ TYPED_TEST(DdMatrix, CanAdvancedApplyToSingleVector)
GKO_ASSERT_MTX_NEAR(this->y->get_local_vector(), result[rank], 0);
}


TYPED_TEST(DdMatrix, CanAdvancedApplyToMultipleVectors)
{
using value_type = typename TestFixture::value_type;
using index_type = typename TestFixture::global_index_type;
using dense_vec_type = typename TestFixture::dense_vec_type;
auto vec_md =
gko::matrix_data<value_type, index_type>{I<I<value_type>>{{1, 2},
{2, 4},
{3, 6},
{4, 8},
{5, 10},
{6, 12},
{7, 14},
{8, 16},
{9, 18},
{10, 20},
{11, 22},
{12, 24}}};
I<I<value_type>> result[3] = {{{-3, -6}, {-1, -2}, {1, 2}, {3, 6}},
{{5, 10}, {7, 14}, {6, 12}, {8, 16}},
{{10, 20}, {12, 24}, {14, 28}, {16, 32}}};
auto rank = this->comm.rank();
this->dist_mat->read_distributed(this->dist_input[rank], this->row_part);
this->x->read_distributed(vec_md, this->row_part);
this->y->read_distributed(vec_md, this->row_part);
auto alpha = gko::initialize<dense_vec_type>({1.0}, this->exec);
auto beta = gko::initialize<dense_vec_type>({1.0}, this->exec);

this->dist_mat->apply(alpha, this->x, beta, this->y);

GKO_ASSERT_MTX_NEAR(this->y->get_local_vector(), result[rank], 0);
}


TYPED_TEST(DdMatrix, CanColScale)
{
using value_type = typename TestFixture::value_type;
using index_type = typename TestFixture::global_index_type;
using matrix_data = typename TestFixture::local_matrix_data;
using csr = typename TestFixture::local_matrix_type;
auto local_size = gko::dim<2>{6, 6};
auto vec_md = gko::matrix_data<value_type, index_type>{I<I<value_type>>{
{1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}, {9}, {10}, {11}, {12}}};
std::array<matrix_data, 3> scaled_result{
{{local_size,
{{0, 0, 2}, {0, 1, -2}, {0, 3, -4}, {1, 0, -1}, {1, 1, 6},
{1, 2, -3}, {1, 4, -5}, {2, 1, -2}, {2, 2, 6}, {2, 5, -6},
{3, 0, -1}, {3, 3, 6}, {3, 4, -2.5}, {4, 1, -2}, {4, 3, -2},
{4, 4, 10}, {4, 5, -3}, {5, 2, -3}, {5, 4, -2.5}, {5, 5, 9}}},
{local_size,
{{0, 0, 10}, {0, 1, -3}, {0, 3, -8}, {0, 4, -2}, {1, 0, -2.5},
{1, 1, 9}, {1, 5, -9}, {2, 2, 10.5}, {2, 3, -4}, {2, 4, -4},
{3, 0, -5}, {3, 2, -3.5}, {3, 3, 16}, {3, 5, -4.5}, {4, 0, -2.5},
{4, 2, -7}, {4, 4, 6}, {5, 1, -6}, {5, 3, -4}, {5, 5, 13.5}}},
{local_size, {{0, 0, 13.5}, {0, 3, -12}, {0, 5, -4}, {1, 1, 20},
{1, 2, -11}, {1, 4, -7}, {2, 1, -10}, {2, 2, 33},
{2, 3, -12}, {2, 5, -8}, {3, 0, -9}, {3, 2, -11},
{3, 3, 24}, {4, 1, -10}, {4, 4, 10.5}, {4, 5, -4},
{5, 0, -4.5}, {5, 2, -11}, {5, 4, -3.5}, {5, 5, 16}}}}};
auto rank = this->comm.rank();
auto res_local = csr::create(this->exec);
res_local->read(scaled_result[rank]);
this->dist_mat->read_distributed(this->dist_input[rank], this->row_part);
this->x->read_distributed(vec_md, this->row_part);

this->dist_mat->col_scale(this->x);

GKO_ASSERT_MTX_NEAR(gko::as<csr>(this->dist_mat->get_local_matrix()),
res_local, 0);
}


TYPED_TEST(DdMatrix, CanRowScale)
{
using value_type = typename TestFixture::value_type;
using index_type = typename TestFixture::global_index_type;
using matrix_data = typename TestFixture::local_matrix_data;
using csr = typename TestFixture::local_matrix_type;
auto local_size = gko::dim<2>{6, 6};
auto vec_md = gko::matrix_data<value_type, index_type>{I<I<value_type>>{
{1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}, {9}, {10}, {11}, {12}}};
std::array<matrix_data, 3> scaled_result{
{{local_size,
{{0, 0, 2}, {0, 1, -1}, {0, 3, -1}, {1, 0, -2}, {1, 1, 6},
{1, 2, -2}, {1, 4, -2}, {2, 1, -3}, {2, 2, 6}, {2, 5, -3},
{3, 0, -4}, {3, 3, 6}, {3, 4, -2}, {4, 1, -5}, {4, 3, -2.5},
{4, 4, 10}, {4, 5, -2.5}, {5, 2, -6}, {5, 4, -3}, {5, 5, 9}}},
{local_size,
{{0, 0, 10}, {0, 1, -2.5}, {0, 3, -5}, {0, 4, -2.5}, {1, 0, -3},
{1, 1, 9}, {1, 5, -6}, {2, 2, 10.5}, {2, 3, -3.5}, {2, 4, -7},
{3, 0, -8}, {3, 2, -4}, {3, 3, 16}, {3, 5, -4}, {4, 0, -2},
{4, 2, -4}, {4, 4, 6}, {5, 1, -9}, {5, 3, -4.5}, {5, 5, 13.5}}},
{local_size,
{{0, 0, 13.5}, {0, 3, -9}, {0, 5, -4.5}, {1, 1, 20}, {1, 2, -10},
{1, 4, -10}, {2, 1, -11}, {2, 2, 33}, {2, 3, -11}, {2, 5, -11},
{3, 0, -12}, {3, 2, -12}, {3, 3, 24}, {4, 1, -7}, {4, 4, 10.5},
{4, 5, -3.5}, {5, 0, -4}, {5, 2, -8}, {5, 4, -4}, {5, 5, 16}}}}};
auto rank = this->comm.rank();
auto res_local = csr::create(this->exec);
res_local->read(scaled_result[rank]);
this->dist_mat->read_distributed(this->dist_input[rank], this->row_part);
this->x->read_distributed(vec_md, this->row_part);

this->dist_mat->row_scale(this->x);

GKO_ASSERT_MTX_NEAR(gko::as<csr>(this->dist_mat->get_local_matrix()),
res_local, 0);
}


#endif

0 comments on commit 1c0e4e9

Please sign in to comment.