text: /*******************************<GINKGO LICENSE>******************************
  Copyright (c) 2017-2022, 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/coo_kernels.hpp"
  
  
  #include <array>
  
  
  #include <omp.h>
  
  
  #include <ginkgo/core/base/exception_helpers.hpp>
  #include <ginkgo/core/base/math.hpp>
  #include <ginkgo/core/matrix/csr.hpp>
  #include <ginkgo/core/matrix/dense.hpp>
  
  
  #include "core/matrix/dense_kernels.hpp"
  #include "omp/components/atomic.hpp"
  
  
  namespace gko {
  namespace kernels {
  /**
   * @brief OpenMP namespace.
   *
   * @ingroup omp
   */
  namespace omp {
  /**
   * @brief The Coordinate matrix format namespace.
   *
   * @ingroup coo
   */
  namespace coo {
  
  
  template <typename ValueType, typename IndexType>
  void spmv(std::shared_ptr<const OmpExecutor> exec,
            const matrix::Coo<ValueType, IndexType>* a,
            const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* c)
  {
      dense::fill(exec, c, zero<ValueType>());
      spmv2(exec, a, b, c);
  }
  
  GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_COO_SPMV_KERNEL);
  
  
  template <typename ValueType, typename IndexType>
  void advanced_spmv(std::shared_ptr<const OmpExecutor> exec,
                     const matrix::Dense<ValueType>* alpha,
                     const matrix::Coo<ValueType, IndexType>* a,
                     const matrix::Dense<ValueType>* b,
                     const matrix::Dense<ValueType>* beta,
                     matrix::Dense<ValueType>* c)
  {
      dense::scale(exec, beta, c);
      advanced_spmv2(exec, alpha, a, b, c);
  }
  
  GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
      GKO_DECLARE_COO_ADVANCED_SPMV_KERNEL);
  
  
  template <int block_size, typename ValueType, typename IndexType>
  void spmv2_blocked(std::shared_ptr<const OmpExecutor> exec,
                     const matrix::Coo<ValueType, IndexType>* a,
                     const matrix::Dense<ValueType>* b,
                     matrix::Dense<ValueType>* c, ValueType scale)
  {
      GKO_ASSERT(b->get_size()[1] > block_size);
      const auto coo_val = a->get_const_values();
      const auto coo_col = a->get_const_col_idxs();
      const auto coo_row = a->get_const_row_idxs();
      const auto num_rhs = b->get_size()[1];
      const auto rounded_rhs = num_rhs / block_size * block_size;
      const auto sentinel_row = a->get_size()[0] + 1;
      const auto nnz = a->get_num_stored_elements();
  
  #pragma omp parallel
      {
          const auto num_threads = omp_get_num_threads();
          const auto work_per_thread =
              static_cast<size_type>(ceildiv(nnz, num_threads));
          const auto thread_id = static_cast<size_type>(omp_get_thread_num());
          const auto begin = work_per_thread * thread_id;
          const auto end = std::min(begin + work_per_thread, nnz);
          if (begin < end) {
              const auto first = begin > 0 ? coo_row[begin - 1] : sentinel_row;
              const auto last = end < nnz ? coo_row[end] : sentinel_row;
              auto nz = begin;
              std::array<ValueType, block_size> partial_sum;
              if (first != sentinel_row) {
                  for (size_type rhs_base = 0; rhs_base < rounded_rhs;
                       rhs_base += block_size) {
                      // handle row overlap with previous thread: block partial
                      // sums
                      partial_sum.fill(zero<ValueType>());
                      for (auto local_nz = nz;
                           local_nz < end && coo_row[local_nz] == first;
                           local_nz++) {
                          const auto col = coo_col[local_nz];
  #pragma unroll
                          for (size_type i = 0; i < block_size; i++) {
                              const auto rhs = i + rhs_base;
                              partial_sum[i] +=
                                  scale * coo_val[local_nz] * b->at(col, rhs);
                          }
                      }
                      // handle row overlap with previous thread: block add to
                      // memory
  #pragma unroll
                      for (size_type i = 0; i < block_size; i++) {
                          const auto rhs = i + rhs_base;
                          atomic_add(c->at(first, rhs), partial_sum[i]);
                      }
                  }
                  // handle row overlap with previous thread: remainder partial
                  // sums
                  partial_sum.fill(zero<ValueType>());
                  for (; nz < end && coo_row[nz] == first; nz++) {
                      const auto row = first;
                      const auto col = coo_col[nz];
                      for (size_type rhs = rounded_rhs; rhs < num_rhs; rhs++) {
                          partial_sum[rhs - rounded_rhs] +=
                              scale * coo_val[nz] * b->at(col, rhs);
                      }
                  }
                  // handle row overlap with previous thread: remainder add to
                  // memory
                  for (size_type rhs = rounded_rhs; rhs < num_rhs; rhs++) {
                      atomic_add(c->at(first, rhs),
                                 partial_sum[rhs - rounded_rhs]);
                  }
              }
              // handle non-overlapping rows
              for (; nz < end && coo_row[nz] != last; nz++) {
                  const auto row = coo_row[nz];
                  const auto col = coo_col[nz];
                  for (size_type rhs_base = 0; rhs_base < rounded_rhs;
                       rhs_base += block_size) {
  #pragma unroll
                      for (size_type i = 0; i < block_size; i++) {
                          const auto rhs = i + rhs_base;
                          c->at(row, rhs) +=
                              scale * coo_val[nz] * b->at(col, rhs);
                      }
                  }
                  for (size_type rhs = rounded_rhs; rhs < num_rhs; rhs++) {
                      c->at(row, rhs) += scale * coo_val[nz] * b->at(col, rhs);
                  }
              }
              if (last != sentinel_row) {
                  for (size_type rhs_base = 0; rhs_base < rounded_rhs;
                       rhs_base += block_size) {
                      // handle row overlap with following thread: block partial
                      // sums
                      partial_sum.fill(zero<ValueType>());
                      for (auto local_nz = nz; local_nz < end; local_nz++) {
                          const auto col = coo_col[local_nz];
  #pragma unroll
                          for (size_type i = 0; i < block_size; i++) {
                              const auto rhs = i + rhs_base;
                              partial_sum[i] +=
                                  scale * coo_val[local_nz] * b->at(col, rhs);
                          }
                      }
                      // handle row overlap with following thread: block add to
                      // memory
  #pragma unroll
                      for (size_type i = 0; i < block_size; i++) {
                          const auto rhs = i + rhs_base;
                          const auto row = last;
                          atomic_add(c->at(row, rhs), partial_sum[i]);
                      }
                  }
                  // handle row overlap with following thread: block partial sums
                  partial_sum.fill(zero<ValueType>());
                  for (; nz < end; nz++) {
                      const auto col = coo_col[nz];
                      for (size_type rhs = rounded_rhs; rhs < num_rhs; rhs++) {
                          partial_sum[rhs - rounded_rhs] +=
                              scale * coo_val[nz] * b->at(col, rhs);
                      }
                  }
                  // handle row overlap with following thread: block add to memory
                  for (size_type rhs = rounded_rhs; rhs < num_rhs; rhs++) {
                      const auto row = last;
                      atomic_add(c->at(row, rhs), partial_sum[rhs - rounded_rhs]);
                  }
              }
          }
      }
  }
  
  
  template <int num_rhs, typename ValueType, typename IndexType>
  void spmv2_small_rhs(std::shared_ptr<const OmpExecutor> exec,
                       const matrix::Coo<ValueType, IndexType>* a,
                       const matrix::Dense<ValueType>* b,
                       matrix::Dense<ValueType>* c, ValueType scale)
  {
      GKO_ASSERT(b->get_size()[1] == num_rhs);
      const auto coo_val = a->get_const_values();
      const auto coo_col = a->get_const_col_idxs();
      const auto coo_row = a->get_const_row_idxs();
      const auto sentinel_row = a->get_size()[0] + 1;
      const auto nnz = a->get_num_stored_elements();
  
  #pragma omp parallel
      {
          const auto num_threads = omp_get_num_threads();
          const auto work_per_thread =
              static_cast<size_type>(ceildiv(nnz, num_threads));
          const auto thread_id = static_cast<size_type>(omp_get_thread_num());
          const auto begin = work_per_thread * thread_id;
          const auto end = std::min(begin + work_per_thread, nnz);
          if (begin < end) {
              const auto first = begin > 0 ? coo_row[begin - 1] : sentinel_row;
              const auto last = end < nnz ? coo_row[end] : sentinel_row;
              auto nz = begin;
              std::array<ValueType, num_rhs> partial_sum;
              if (first != sentinel_row) {
                  // handle row overlap with previous thread: partial sums
                  partial_sum.fill(zero<ValueType>());
                  for (; nz < end && coo_row[nz] == first; nz++) {
                      const auto col = coo_col[nz];
  #pragma unroll
                      for (size_type rhs = 0; rhs < num_rhs; rhs++) {
                          partial_sum[rhs] +=
                              scale * coo_val[nz] * b->at(col, rhs);
                      }
                  }
                  // handle row overlap with previous thread: add to memory
  #pragma unroll
                  for (size_type rhs = 0; rhs < num_rhs; rhs++) {
                      atomic_add(c->at(first, rhs), partial_sum[rhs]);
                  }
              }
              // handle non-overlapping rows
              for (; nz < end && coo_row[nz] != last; nz++) {
                  const auto row = coo_row[nz];
                  const auto col = coo_col[nz];
  #pragma unroll
                  for (size_type rhs = 0; rhs < num_rhs; rhs++) {
                      c->at(row, rhs) += scale * coo_val[nz] * b->at(col, rhs);
                  }
              }
              if (last != sentinel_row) {
                  // handle row overlap with following thread: partial sums
                  partial_sum.fill(zero<ValueType>());
                  for (; nz < end; nz++) {
                      const auto col = coo_col[nz];
  #pragma unroll
                      for (size_type rhs = 0; rhs < num_rhs; rhs++) {
                          partial_sum[rhs] +=
                              scale * coo_val[nz] * b->at(col, rhs);
                      }
                  }
                  // handle row overlap with following thread: add to memory
  #pragma unroll
                  for (size_type rhs = 0; rhs < num_rhs; rhs++) {
                      const auto row = last;
                      atomic_add(c->at(row, rhs), partial_sum[rhs]);
                  }
              }
          }
      }
  }
  
  
  template <typename ValueType, typename IndexType>
  void generic_spmv2(std::shared_ptr<const OmpExecutor> exec,
                     const matrix::Coo<ValueType, IndexType>* a,
                     const matrix::Dense<ValueType>* b,
                     matrix::Dense<ValueType>* c, ValueType scale)
  {
      const auto num_rhs = b->get_size()[1];
      if (num_rhs <= 0) {
          return;
      }
      if (num_rhs == 1) {
          spmv2_small_rhs<1>(exec, a, b, c, scale);
          return;
      }
      if (num_rhs == 2) {
          spmv2_small_rhs<2>(exec, a, b, c, scale);
          return;
      }
      if (num_rhs == 3) {
          spmv2_small_rhs<3>(exec, a, b, c, scale);
          return;
      }
      if (num_rhs == 4) {
          spmv2_small_rhs<4>(exec, a, b, c, scale);
          return;
      }
      spmv2_blocked<4>(exec, a, b, c, scale);
  }
  
  
  template <typename ValueType, typename IndexType>
  void spmv2(std::shared_ptr<const OmpExecutor> exec,
             const matrix::Coo<ValueType, IndexType>* a,
             const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* c)
  {
      generic_spmv2(exec, a, b, c, one<ValueType>());
  }
  
  GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_COO_SPMV2_KERNEL);
  
  
  template <typename ValueType, typename IndexType>
  void advanced_spmv2(std::shared_ptr<const OmpExecutor> exec,
                      const matrix::Dense<ValueType>* alpha,
                      const matrix::Coo<ValueType, IndexType>* a,
                      const matrix::Dense<ValueType>* b,
                      matrix::Dense<ValueType>* c)
  {
      generic_spmv2(exec, a, b, c, alpha->at(0, 0));
  }
  
  GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
      GKO_DECLARE_COO_ADVANCED_SPMV2_KERNEL);
  
  
  }  // namespace coo
  }  // namespace omp
  }  // namespace kernels
  }  // namespace gko
  