Shared Memory Allocation Error at parallel_reduce

[mndevec]

Team scratch memory allocations fail with parallel_reduce. Below is a simple example, where memory allocation works for parallel_for but fails for parallel_reduce. Also, for some reason I get a compile error if I uncomment the below line
Kokkos::parallel_reduce( team_policy_t(1, 8, 32), myfunc1(), reduction);
in the code.

#include <stdio.h>
#include <Kokkos_Core.hpp>
typedef typename Kokkos::Cuda MyMemorySpace;
typedef typename Kokkos::Cuda MyExecSpace;
typedef Kokkos::TeamPolicy<MyExecSpace> team_policy_t ;
typedef typename team_policy_t::member_type team_member_t ;
#define VECTORSIZE 32
#define TEAMSIZE 8
struct myfunc1{
  myfunc1(){}
  KOKKOS_INLINE_FUNCTION
  void operator()(const team_member_t & teamMember) const {
    const int team_size = teamMember.team_size();
    const int team_rank = teamMember.team_rank();
    int thread_id = (teamMember.league_rank()  * team_size + team_rank);
    int *ptr = (int *) ( (teamMember.team_shmem().get_shmem(16384) )) ;
    if (ptr == NULL && thread_id == 0){
      Kokkos::single(Kokkos::PerThread(teamMember),[=] () {
        printf("ptr is Null in parallel for\n");
      });
    }
  }
  KOKKOS_INLINE_FUNCTION
  void operator()(const team_member_t & teamMember, int  &reduction) const {
    const int team_size = teamMember.team_size();
    const int team_rank = teamMember.team_rank();
    int thread_id = (teamMember.league_rank()  * team_size + team_rank);
    int *ptr = (int *) ( (teamMember.team_shmem().get_shmem(16384) )) ;
    reduction++;
    if (ptr == NULL && thread_id == 0){
      Kokkos::single(Kokkos::PerThread(teamMember),[=] () {
        printf("ptr is Null in parallel reduce\n");
      });
    }
  }
  size_t team_shmem_size (int team_size) const {
    return 16384;
  }
};

struct myfunc2{
  myfunc2(){}
  KOKKOS_INLINE_FUNCTION
  void operator()(const team_member_t & teamMember, int  &reduction) const {
    const int team_size = teamMember.team_size();
    const int team_rank = teamMember.team_rank();
    int thread_id = (teamMember.league_rank()  * team_size + team_rank);
    int *ptr = (int *) ( (teamMember.team_shmem().get_shmem(16384) )) ;
    reduction++;
    if (ptr == NULL && thread_id == 0){
      Kokkos::single(Kokkos::PerThread(teamMember),[=] () {
        printf("ptr is Null in parallel reduce\n");
      });
    }
  }
  size_t team_shmem_size (int team_size) const {
    return 16384;
  }
};
int  main (int  argc, char ** argv){
  Kokkos::initialize(argc, argv);
  MyExecSpace::print_configuration(std::cout);
  Kokkos::parallel_for( team_policy_t(1, TEAMSIZE, VECTORSIZE), myfunc1());
  MyExecSpace::fence();
  int reduction = 0;
  //Kokkos::parallel_reduce( team_policy_t(1, 8, 32), myfunc1(), reduction);
  Kokkos::parallel_reduce( team_policy_t(1, TEAMSIZE, VECTORSIZE), myfunc2(), reduction);
  MyExecSpace::fence();
  Kokkos::finalize();
  return 0;
}

[crtrott]

Both Bugs are confirmed.

[crtrott]

Ok I know what is going on with the shared memory size not being set and I know roughly what is going on with the first one not compiling. Both of those will get addressed as part of my work on adding a better generic reduction interface, and some internal code cleanup. For now you can work around the shared memory issue by setting the scratch size via the policy instead of the functor:

Kokkos::parallel_reduce( team_policy_t(1, TEAMSIZE, VECTORSIZE).set_scratch_size(1,Kokkos::PerTeam(16384)), myfunc2(), reduction);

[crtrott]

The shared memory thingy is fixed. The "can't compile thing" would be much harder, and for now my position is to retreat to: "If you want more than one operator in a class, you must use the Tag mechanism to distinguish them, even if they are technically distinguishable through the fact that one is a reduction operator the other one is not."

The reason for that is that the internal Functor inspection would be getting more complicated, and its already complicated enough.

[mndevec]

Thanks Christian!

[crtrott]

Reopening until pushed to master