matrix_gather_scatter_helper.hpp

// FEAT3: Finite Element Analysis Toolbox, Version 3
// Copyright (C) 2010 by Stefan Turek & the FEAT group
// FEAT3 is released under the GNU General Public License version 3,
// see the file 'copyright.txt' in the top level directory for details.
 
#pragma once
 
#include <kernel/base_header.hpp>
#include <kernel/util/tiny_algebra.hpp>
 
#ifdef __CUDACC__
#include <cuda/std/type_traits>
#endif
 
 
 
namespace FEAT
{
 
  namespace Intern
  {
    enum MatrixGatherScatterPolicy
    {
      useLocalOps = 0,
      useLocalSortHelper = 1,
      useColPtr = 2
    };
  }
 
  namespace LAFEM
  {
    template<typename Space_, typename  DT_, typename IT_, FEAT::Intern::MatrixGatherScatterPolicy policy_ = FEAT::Intern::MatrixGatherScatterPolicy::useLocalOps>
    struct MatrixGatherScatterHelper DOXY({});
 
    template<typename Space_, typename  DT_, typename IT_>
    struct MatrixGatherScatterHelper<Space_, DT_, IT_, FEAT::Intern::MatrixGatherScatterPolicy::useLocalOps>
    {
      typedef Space_ SpaceType;
      typedef DT_ DataType;
      typedef IT_ IndexType;
 
 
      template<typename InnerType_, int numr_, int numc_ = numr_>
      CUDA_HOST_DEVICE static void scatter_matrix_csr(const Tiny::Matrix<InnerType_, numr_, numc_>& loc_mat, InnerType_* matrix_data, const IndexType* row_map, const IndexType* col_map,
                                      [[maybe_unused]] Index matrix_num_rows, [[maybe_unused]] Index matrix_num_cols, const IndexType* matrix_row_ptr,
                                      const IndexType* matrix_col_idx, DataType alpha = DataType(1), [[maybe_unused]] IndexType* dummy_ptr = nullptr)
      {
        #ifndef __CUDACC__
        static_assert(std::is_same<typename Tiny::Intern::DataTypeExtractor<InnerType_>::MyDataType, DataType>(), "Inner Datatype does not match!");
        #else
        static_assert(::cuda::std::is_same<typename Tiny::Intern::DataTypeExtractor<InnerType_>::MyDataType, DataType>(), "Inner Datatype does not match!");
        #endif
        IndexType loc_idx_map[numc_];
        for(int i = 0; i < numr_; ++i)
        {
          const Index ix = row_map[i];
          for(IndexType k = matrix_row_ptr[ix]; k < matrix_row_ptr[ix+1]; ++k)
          {
            for(int k_ptr = 0; k_ptr < numc_; ++k_ptr)
            {
              if(matrix_col_idx[k] == col_map[k_ptr])
              {
                loc_idx_map[k_ptr] = k;
                break;
              }
            }
          }
 
          //now loop over all local columns
          for(int j = 0; j < numc_; ++j)
          {
            Tiny::axpy(matrix_data[loc_idx_map[j]], loc_mat[i][j], alpha);
          }
        }
      }
 
      #if defined(__CUDACC__) || defined(DOXYGEN)
      template<typename ThreadGroup_, int numr_, int numc_=numr_>
      CUDA_DEVICE __forceinline__ static void grouped_scatter_matrix_csr(const ThreadGroup_& tg, const int scatter_size, const int scatter_offset, const DataType* loc_mat, DataType* matrix_data, const IndexType* row_map, const IndexType* col_map,
                                      [[maybe_unused]] Index matrix_num_rows, [[maybe_unused]] Index matrix_num_cols, const IndexType* matrix_row_ptr,
                                      const IndexType* matrix_col_idx, int num_data_row, int num_data_col, DataType alpha = DataType(1), [[maybe_unused]] IndexType* dummy_ptr = nullptr)
      {
        for(int idx = tg.thread_rank(); (idx < scatter_size*numr_*numc_)/* && ((idx + scatter_offset*numr_*numc_) < num_data_row*num_data_col*numr_*numc_)*/; idx += tg.num_threads())
        {
          IndexType loc_idx_map = ~IndexType(0);
          const int i = ((idx/(numr_*numc_)+scatter_offset))/num_data_row;
          const int j = ((idx/(numr_*numc_)+scatter_offset))%num_data_row;
          const Index ix = row_map[i];
          // brute force search for the correct value
          for(IndexType k = matrix_row_ptr[ix]; k < matrix_row_ptr[ix+1]; ++k)
          {
            loc_idx_map = matrix_col_idx[k] == col_map[j] ? k : loc_idx_map;
          }
 
          // ASSERT(loc_idx_map != ~IndexType(0));
 
 
          // and now add our value strided to the correct length
          matrix_data[loc_idx_map * numr_ * numc_ + idx%(numr_*numc_)] += alpha * loc_mat[idx];
        }
      }
      #endif
 
      template<typename InnerType_, int numr_, int numc_ = numr_>
      CUDA_HOST_DEVICE static void gather_matrix_csr(Tiny::Matrix<InnerType_, numr_, numc_>& loc_mat, const InnerType_* matrix_data, const IndexType* row_map, const IndexType* col_map,
                                      [[maybe_unused]] Index matrix_num_rows, [[maybe_unused]] Index matrix_num_cols, const IndexType* matrix_row_ptr,
                                      const IndexType* matrix_col_idx, DataType alpha = DataType(1), [[maybe_unused]] const IndexType* dummy_ptr = nullptr)
      {
        #ifndef __CUDACC__
        static_assert(std::is_same<typename Tiny::Intern::DataTypeExtractor<InnerType_>::MyDataType, DataType>(), "Inner Datatype does not match!");
        #else
        static_assert(::cuda::std::is_same<typename Tiny::Intern::DataTypeExtractor<InnerType_>::MyDataType, DataType>(), "Inner Datatype does not match!");
        #endif
        IndexType loc_idx_map[numc_];
        // loop over all local row entries
        for(int i(0); i < numr_; ++i)
        {
          // fetch row index
          const Index ix = row_map[i];
 
          // build column pointer for this row entry contribution
          for(IndexType k = matrix_row_ptr[ix]; k < matrix_row_ptr[ix + 1]; ++k)
          {
            for(int k_ptr = 0; k_ptr < numc_; ++k_ptr)
            {
              if(matrix_col_idx[k] == col_map[k_ptr])
              {
                loc_idx_map[k_ptr] = k;
                break;
              }
            }
          }
 
          // loop over all local column entries
          for(int j(0); j < numc_; ++j)
          {
            Tiny::axpy(loc_mat[i][j], matrix_data[loc_idx_map[j]], alpha);
          }
          // continue with next row entry
        }
      }
 
      #if defined(__CUDACC__) || defined(DOXYGEN)
      template<typename ThreadGroup_, int numr_, int numc_=numr_>
      CUDA_HOST_DEVICE static void grouped_gather_matrix_csr(const ThreadGroup_& tg, const int scatter_size, const int scatter_offset, DataType* loc_mat, const DataType* matrix_data, const IndexType* row_map, const IndexType* col_map,
                                      [[maybe_unused]] Index matrix_num_rows, [[maybe_unused]] Index matrix_num_cols, const IndexType* matrix_row_ptr,
                                      const IndexType* matrix_col_idx, int num_data_row, int num_data_col, DataType alpha = DataType(1), [[maybe_unused]] IndexType* dummy_ptr = nullptr)
      {
        for(int idx = tg.thread_rank(); (idx < scatter_size*numr_*numc_) && ((idx + scatter_offset*numr_*numc_) < num_data_row*num_data_col*numr_*numc_); idx += tg.num_threads())
        {
          IndexType loc_idx_map = ~IndexType(0);
          const int i = ((idx/(numr_*numc_)+scatter_offset))/num_data_row;
          const int j = ((idx/(numr_*numc_)+scatter_offset))%num_data_row;
          const Index ix = row_map[i];
          // brute force search for the correct value
          for(IndexType k = matrix_row_ptr[ix]; k < matrix_row_ptr[ix+1]; ++k)
          {
            loc_idx_map = matrix_col_idx[k] == col_map[j] ? k : loc_idx_map;
          }
 
          // and now add our value strided to the correct length
          loc_mat[idx] += alpha * matrix_data[loc_idx_map * numr_ * numc_ + idx%(numr_*numc_)];
        }
      }
      #endif
    }; // struct MatrixGatherScatterHelper<localOps>
 
   template<typename Space_, typename  DT_, typename IT_>
   struct MatrixGatherScatterHelper<Space_, DT_, IT_, FEAT::Intern::MatrixGatherScatterPolicy::useLocalSortHelper>
   {
      typedef DT_ DataType;
      typedef IT_ IndexType;
      typedef Space_ SpaceType;
 
 
      template<typename InnerType_, int numr_, int numc_ = numr_>
      CUDA_HOST_DEVICE static void scatter_matrix_csr(const Tiny::Matrix<InnerType_, numr_, numc_>& loc_mat, InnerType_* matrix_data, const IndexType* row_map, const IndexType* col_map,
                                      [[maybe_unused]] Index matrix_num_rows, [[maybe_unused]] Index matrix_num_cols, const IndexType* matrix_row_ptr,
                                      const IndexType* matrix_col_idx, DataType alpha, const IndexType* col_map_sorter)
      {
        #ifndef __CUDACC__
        static_assert(std::is_same<typename Tiny::Intern::DataTypeExtractor<InnerType_>::MyDataType, DataType>(), "Inner Datatype does not match!");
        #else
        static_assert(::cuda::std::is_same<typename Tiny::Intern::DataTypeExtractor<InnerType_>::MyDataType, DataType>(), "Inner Datatype does not match!");
        #endif
        IndexType loc_idx_map[numc_];
 
        // loop over all local row entries
        for(int i(0); i < numr_; ++i)
        {
          // fetch row index
          Index k = matrix_row_ptr[row_map[i]];
          for(Index k_ptr = 0; k_ptr < numc_; ++k_ptr)
          {
            const Index real_dof = col_map_sorter[k_ptr];
            //search for our column value, no boundary checks, so be damn sure the value is inside matrix_col_idx
            while(matrix_col_idx[k] < col_map[real_dof])
            {
              ++k;
            }
            loc_idx_map[real_dof] = IndexType(k++);
          }
 
          // loop over all local column entries
          for(int j(0); j < numc_; ++j)
          {
            Tiny::axpy(matrix_data[loc_idx_map[j]], loc_mat[i][j], alpha);
          }
          // continue with next row entry
        }
      }
 
      #if defined(__CUDACC__) || defined(DOXYGEN)
      template<typename ThreadGroup_, int numr_, int numc_=numr_>
      CUDA_HOST_DEVICE static __forceinline__ void grouped_scatter_matrix_csr(const ThreadGroup_& tg, const int scatter_size, const int scatter_offset, const DataType* loc_mat, DataType* matrix_data, const IndexType* row_map, const IndexType* col_map,
                                      [[maybe_unused]] Index matrix_num_rows, [[maybe_unused]] Index matrix_num_cols, const IndexType* matrix_row_ptr,
                                      const IndexType* matrix_col_idx, int num_data_row, int num_data_col, DataType alpha = DataType(1), [[maybe_unused]] IndexType* dummy_ptr = nullptr)
      {
        MatrixGatherScatterHelper<Space_, DT_, IT_, FEAT::Intern::MatrixGatherScatterPolicy::useLocalOps>::template grouped_scatter_matrix_csr<ThreadGroup_, numr_, numc_>(tg, scatter_size, scatter_offset, loc_mat,
              matrix_data, row_map, col_map, matrix_num_rows, matrix_num_cols, matrix_row_ptr, matrix_col_idx, num_data_row, num_data_col, alpha, nullptr);
      }
      #endif
 
      template<typename InnerType_, int numr_, int numc_ = numr_>
      CUDA_HOST_DEVICE static void gather_matrix_csr(Tiny::Matrix<InnerType_, numr_, numc_>& loc_mat, const InnerType_* matrix_data, const IndexType* row_map, const IndexType* col_map,
                                      [[maybe_unused]] Index matrix_num_rows, [[maybe_unused]] Index matrix_num_cols, const IndexType* matrix_row_ptr,
                                      const IndexType* matrix_col_idx, DataType alpha, const IndexType* col_map_sorter)
      {
        #ifndef __CUDACC__
        static_assert(std::is_same<typename Tiny::Intern::DataTypeExtractor<InnerType_>::MyDataType, DataType>(), "Inner Datatype does not match!");
        #else
        static_assert(::cuda::std::is_same<typename Tiny::Intern::DataTypeExtractor<InnerType_>::MyDataType, DataType>(), "Inner Datatype does not match!");
        #endif
 
        IndexType loc_idx_map[numc_];
 
        // loop over all local row entries
        for(int i(0); i < numr_; ++i)
        {
          // fetch row index
          Index k = matrix_row_ptr[row_map[i]];
          for(Index k_ptr = 0; k_ptr < numc_; ++k_ptr)
          {
            const Index real_dof = col_map_sorter[k_ptr];
            //search for our column value, no boundary checks, so be damn sure the value is inside matrix_col_idx
            while(matrix_col_idx[k] < col_map[real_dof])
            {
              ++k;
            }
            loc_idx_map[real_dof] = IndexType(k++);
          }
 
          // loop over all local column entries
          for(int j(0); j < numc_; ++j)
          {
            Tiny::axpy(loc_mat[i][j], matrix_data[loc_idx_map[j]], alpha);
          }
          // continue with next row entry
        }
      }
 
      #if defined(__CUDACC__) || defined(DOXYGEN)
      template<typename ThreadGroup_, int numr_, int numc_=numr_>
      CUDA_HOST_DEVICE static void grouped_gather_matrix_csr(const ThreadGroup_& tg, const int scatter_size, const int scatter_offset, DataType* loc_mat, const DataType* matrix_data, const IndexType* row_map, const IndexType* col_map,
                                      [[maybe_unused]] Index matrix_num_rows, [[maybe_unused]] Index matrix_num_cols, const IndexType* matrix_row_ptr,
                                      const IndexType* matrix_col_idx, int num_data_row, int num_data_col, DataType alpha = DataType(1), [[maybe_unused]] IndexType* dummy_ptr = nullptr)
      {
        MatrixGatherScatterHelper<Space_, DT_, IT_, FEAT::Intern::MatrixGatherScatterPolicy::useLocalOps>::template grouped_gather_matrix_csr<ThreadGroup_, numr_, numc_>(tg, scatter_size, scatter_offset, loc_mat,
              matrix_data, row_map, col_map, matrix_num_rows, matrix_num_cols, matrix_row_ptr, matrix_col_idx, num_data_row, num_data_col, alpha, nullptr);
      }
      #endif
   }; // struct MatrixGatherScatterHelper
  }
}