Skip to content

Commit

Permalink
Split compilation of dense kernels
Browse files Browse the repository at this point in the history
  • Loading branch information
upsj committed Jul 30, 2023
1 parent 945a4d8 commit 331a70d
Show file tree
Hide file tree
Showing 10 changed files with 628 additions and 172 deletions.
8 changes: 8 additions & 0 deletions common/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,14 @@ set(UNIFIED_SOURCES
matrix/coo_kernels.cpp
matrix/csr_kernels.cpp
matrix/dense_kernels.cpp
matrix/dense_kernels_conv.cpp
matrix/dense_kernels_conv_ell.cpp
matrix/dense_kernels_conv_sellp.cpp
matrix/dense_kernels_dot.cpp
matrix/dense_kernels_dot_conj.cpp
matrix/dense_kernels_norm1.cpp
matrix/dense_kernels_norm2.cpp
matrix/dense_kernels_norm2_sq.cpp
matrix/ell_kernels.cpp
matrix/hybrid_kernels.cpp
matrix/sellp_kernels.cpp
Expand Down
172 changes: 0 additions & 172 deletions common/unified/matrix/dense_kernels.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


#include "common/unified/base/kernel_launch.hpp"
#include "common/unified/base/kernel_launch_reduction.hpp"
#include "core/base/mixed_precision_types.hpp"
#include "core/components/prefix_sum_kernels.hpp"


namespace gko {
Expand Down Expand Up @@ -242,176 +240,6 @@ void sub_scaled_diag(std::shared_ptr<const DefaultExecutor> exec,
GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_SUB_SCALED_DIAG_KERNEL);


template <typename ValueType>
void compute_dot(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* x,
const matrix::Dense<ValueType>* y,
matrix::Dense<ValueType>* result, array<char>& tmp)
{
run_kernel_col_reduction_cached(
exec,
[] GKO_KERNEL(auto i, auto j, auto x, auto y) {
return x(i, j) * y(i, j);
},
GKO_KERNEL_REDUCE_SUM(ValueType), result->get_values(), x->get_size(),
tmp, x, y);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_DOT_KERNEL);


template <typename ValueType>
void compute_conj_dot(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* x,
const matrix::Dense<ValueType>* y,
matrix::Dense<ValueType>* result, array<char>& tmp)
{
run_kernel_col_reduction_cached(
exec,
[] GKO_KERNEL(auto i, auto j, auto x, auto y) {
return conj(x(i, j)) * y(i, j);
},
GKO_KERNEL_REDUCE_SUM(ValueType), result->get_values(), x->get_size(),
tmp, x, y);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_CONJ_DOT_KERNEL);


template <typename ValueType>
void compute_norm2(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* x,
matrix::Dense<remove_complex<ValueType>>* result,
array<char>& tmp)
{
run_kernel_col_reduction_cached(
exec,
[] GKO_KERNEL(auto i, auto j, auto x) { return squared_norm(x(i, j)); },
[] GKO_KERNEL(auto a, auto b) { return a + b; },
[] GKO_KERNEL(auto a) { return sqrt(a); }, remove_complex<ValueType>{},
result->get_values(), x->get_size(), tmp, x);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_NORM2_KERNEL);

template <typename ValueType>
void compute_norm1(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* x,
matrix::Dense<remove_complex<ValueType>>* result,
array<char>& tmp)
{
run_kernel_col_reduction_cached(
exec, [] GKO_KERNEL(auto i, auto j, auto x) { return abs(x(i, j)); },
GKO_KERNEL_REDUCE_SUM(remove_complex<ValueType>), result->get_values(),
x->get_size(), tmp, x);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_NORM1_KERNEL);


template <typename ValueType>
void compute_max_nnz_per_row(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* source,
size_type& result)
{
array<size_type> partial{exec, source->get_size()[0] + 1};
count_nonzeros_per_row(exec, source, partial.get_data());
run_kernel_reduction(
exec, [] GKO_KERNEL(auto i, auto partial) { return partial[i]; },
GKO_KERNEL_REDUCE_MAX(size_type),
partial.get_data() + source->get_size()[0], source->get_size()[0],
partial);
result = exec->copy_val_to_host(partial.get_const_data() +
source->get_size()[0]);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
GKO_DECLARE_DENSE_COMPUTE_MAX_NNZ_PER_ROW_KERNEL);


template <typename ValueType>
void compute_slice_sets(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* source,
size_type slice_size, size_type stride_factor,
size_type* slice_sets, size_type* slice_lengths)
{
const auto num_rows = source->get_size()[0];
array<size_type> row_nnz{exec, num_rows};
count_nonzeros_per_row(exec, source, row_nnz.get_data());
const auto num_slices =
static_cast<size_type>(ceildiv(num_rows, slice_size));
run_kernel_row_reduction(
exec,
[] GKO_KERNEL(auto slice, auto local_row, auto row_nnz, auto slice_size,
auto stride_factor, auto num_rows) {
const auto row = slice * slice_size + local_row;
return row < num_rows ? static_cast<size_type>(
ceildiv(row_nnz[row], stride_factor) *
stride_factor)
: size_type{};
},
GKO_KERNEL_REDUCE_MAX(size_type), slice_lengths, 1,
gko::dim<2>{num_slices, slice_size}, row_nnz, slice_size, stride_factor,
num_rows);
exec->copy(num_slices, slice_lengths, slice_sets);
components::prefix_sum_nonnegative(exec, slice_sets, num_slices + 1);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
GKO_DECLARE_DENSE_COMPUTE_SLICE_SETS_KERNEL);


template <typename ValueType, typename IndexType>
void count_nonzeros_per_row(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* mtx,
IndexType* result)
{
run_kernel_row_reduction(
exec,
[] GKO_KERNEL(auto i, auto j, auto mtx) {
return is_nonzero(mtx(i, j)) ? 1 : 0;
},
GKO_KERNEL_REDUCE_SUM(IndexType), result, 1, mtx->get_size(), mtx);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_DENSE_COUNT_NONZEROS_PER_ROW_KERNEL);
GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
GKO_DECLARE_DENSE_COUNT_NONZEROS_PER_ROW_KERNEL_SIZE_T);


template <typename ValueType>
void compute_squared_norm2(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* x,
matrix::Dense<remove_complex<ValueType>>* result,
array<char>& tmp)
{
run_kernel_col_reduction_cached(
exec,
[] GKO_KERNEL(auto i, auto j, auto x) { return squared_norm(x(i, j)); },
GKO_KERNEL_REDUCE_SUM(remove_complex<ValueType>), result->get_values(),
x->get_size(), tmp, x);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
GKO_DECLARE_DENSE_COMPUTE_SQUARED_NORM2_KERNEL);


template <typename ValueType>
void compute_sqrt(std::shared_ptr<const DefaultExecutor> exec,
matrix::Dense<ValueType>* x)
{
run_kernel(
exec,
[] GKO_KERNEL(auto row, auto col, auto x) {
x(row, col) = sqrt(x(row, col));
},
x->get_size(), x);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_SQRT_KERNEL);


template <typename ValueType, typename IndexType>
void symm_permute(std::shared_ptr<const DefaultExecutor> exec,
const array<IndexType>* permutation_indices,
Expand Down
75 changes: 75 additions & 0 deletions common/unified/matrix/dense_kernels_conv.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/*******************************<GINKGO LICENSE>******************************
Copyright (c) 2017-2023, the Ginkgo authors
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:
1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
******************************<GINKGO LICENSE>*******************************/

#include "core/matrix/dense_kernels.hpp"


#include <ginkgo/core/base/math.hpp>


#include "common/unified/base/kernel_launch_reduction.hpp"


namespace gko {
namespace kernels {
namespace GKO_DEVICE_NAMESPACE {
/**
* @brief The Dense matrix format namespace.
*
* @ingroup dense
*/
namespace dense {


template <typename ValueType, typename IndexType>
void count_nonzeros_per_row(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* mtx,
IndexType* result)
{
run_kernel_row_reduction(
exec,
[] GKO_KERNEL(auto i, auto j, auto mtx) {
return is_nonzero(mtx(i, j)) ? 1 : 0;
},
GKO_KERNEL_REDUCE_SUM(IndexType), result, 1, mtx->get_size(), mtx);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_DENSE_COUNT_NONZEROS_PER_ROW_KERNEL);
GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
GKO_DECLARE_DENSE_COUNT_NONZEROS_PER_ROW_KERNEL_SIZE_T);


} // namespace dense
} // namespace GKO_DEVICE_NAMESPACE
} // namespace kernels
} // namespace gko
76 changes: 76 additions & 0 deletions common/unified/matrix/dense_kernels_conv_ell.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
/*******************************<GINKGO LICENSE>******************************
Copyright (c) 2017-2023, the Ginkgo authors
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:
1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
******************************<GINKGO LICENSE>*******************************/

#include "core/matrix/dense_kernels.hpp"


#include <ginkgo/core/base/math.hpp>


#include "common/unified/base/kernel_launch_reduction.hpp"


namespace gko {
namespace kernels {
namespace GKO_DEVICE_NAMESPACE {
/**
* @brief The Dense matrix format namespace.
*
* @ingroup dense
*/
namespace dense {


template <typename ValueType>
void compute_max_nnz_per_row(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Dense<ValueType>* source,
size_type& result)
{
array<size_type> partial{exec, source->get_size()[0] + 1};
count_nonzeros_per_row(exec, source, partial.get_data());
run_kernel_reduction(
exec, [] GKO_KERNEL(auto i, auto partial) { return partial[i]; },
GKO_KERNEL_REDUCE_MAX(size_type),
partial.get_data() + source->get_size()[0], source->get_size()[0],
partial);
result = exec->copy_val_to_host(partial.get_const_data() +
source->get_size()[0]);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
GKO_DECLARE_DENSE_COMPUTE_MAX_NNZ_PER_ROW_KERNEL);


} // namespace dense
} // namespace GKO_DEVICE_NAMESPACE
} // namespace kernels
} // namespace gko
Loading

0 comments on commit 331a70d

Please sign in to comment.