Inconsistent results from TeamThreadRange reduction

[Char-Aznable]

Why do the following 2 slightly different usages of TeamThreadRange reduction give different results:

//some array to contribute to the sum
Kokkos::DualView<T[nTeams][teamSize]> v("v")
//some DualView to store the results
Kokkos::DualView<T[nTeams]> sum("sum")
...
parallel_for(policy, KOKKOS_LAMBDA(const auto& team) {
  int s = 0;
  parallel_reduce(TeamThreadRange(team,teamSize),[&](const auto& i, auto& update) {
     update += v.d_view(team.league_rank(),i);
  },s);
  team.team_barrier();
  Kokkos::single(Kokkos::PerTeam(team),[&](){
      sum.d_view(team.league_rank()) = s;
  });
});
...

vs.

//some array to contribute to the sum
Kokkos::DualView<T[nTeams][teamSize]> v("v")
//some DualView to store the results
Kokkos::DualView<T[nTeams]> sum("sum")
...
parallel_for(policy, KOKKOS_LAMBDA(const auto& team) {
   parallel_reduce(TeamThreadRange(team,teamSize),[&](const auto& i, auto& update) {
     update += v.d_view(team.league_rank(),i);
  },sum.d_view(team.league_rank()));
  team.team_barrier();
});
...

The first case gives the correct answer but the second case gives gibberish. What makes the difference here?

Here's a complete example to try it out:

#include <Kokkos_Core.hpp>
#include <Kokkos_DualView.hpp>
#include <Kokkos_Random.hpp>

template < class DualView >
void syncToHost(DualView&& data) {
  using MemorySpace = typename std::decay_t<DualView>::memory_space;
  using HostMirrorSpace = typename std::decay_t<DualView>::host_mirror_space;
  std::forward<DualView>(data).template modify<MemorySpace>();
  std::forward<DualView>(data).template sync<HostMirrorSpace>();
}

template < class DualView >
void syncToDevice(DualView&& data) {
  using MemorySpace = typename std::decay_t<DualView>::memory_space;
  using HostMirrorSpace = typename std::decay_t<DualView>::host_mirror_space;
  std::forward<DualView>(data).template modify<HostMirrorSpace>();
  std::forward<DualView>(data).template sync<MemorySpace>();
}

using ExecutionSpace = Kokkos::DefaultExecutionSpace;
using namespace std;

static constexpr std::size_t nTeams = 8;
static constexpr std::size_t teamSize = 32;

int main(int argc, char* argv[]) {
  Kokkos::initialize(argc,argv); {
  Kokkos::DualView<std::size_t[nTeams][teamSize],ExecutionSpace> v("v");
  Kokkos::DualView<std::size_t[nTeams],ExecutionSpace> sum("sum");

  for(std::size_t i = 0; i < nTeams; ++i) {
    for(std::size_t j = 0; j < teamSize; ++j) {
      v.h_view(i,j) = 1;
    }
  }
  syncToDevice(v);

  auto policy = Kokkos::TeamPolicy<ExecutionSpace>(nTeams,teamSize);
  Kokkos::parallel_for(policy,KOKKOS_LAMBDA (const auto& team) {
    const auto nT = teamSize;
    const auto iTeam = team.league_rank();
    //std::size_t s = 0;
    Kokkos::parallel_reduce(Kokkos::TeamThreadRange(team,nT),
      [&](const auto& iT, auto& update) {
      update += v.d_view(iTeam,iT);
    },sum.d_view(iTeam));
    team.team_barrier();
    //sum.d_view(iTeam) = s;
  });
  Kokkos::fence();
  syncToHost(sum);

  for(std::size_t i = 0; i < nTeams; ++i) {
    printf("%10lu%10lu\n",i,sum.h_view(i));
  }

  } Kokkos::finalize();
}

[mhoemmen]

This isn't relevant to your question, but just curious: What's with the rvalue references && on DualView? If you made it an lvalue reference &, you could skip the std::forward thing.

[Char-Aznable]

@mhoemmen It's a good point. I did it because I didn't know enough what the modify and sync function do (what if, e.g., they are ref qualified) and I was trying to keep it generic. Then I simply copy the snippet all over the place.

[mhoemmen]

You definitely don't need the && there. Plain old & suffices.

[jeffmiles63]

You cannot update sum.d_view(team.league_rank() directly within a team parallel_for because all of the threads in the team will be live and will access sum.d_view at some point. The Kokkos::Single is necessary to limit access to sum.d_view to only once per team. See pages 132- 136 in the KokkosTutorial_ORNL18.pdf for more info.

[Char-Aznable]

@jeffmiles63 Thanks for the reply. I think Kokkos::Single isn't strictly necessary here because s is the same for all threads after reduction and I believe there's at least one thread will succeed in writing the correct answer to sum.d_view(team.league_rank()). I tried it myself that:

//some array to contribute to the sum
Kokkos::DualView<T[nTeams][teamSize]> v("v")
//some DualView to store the results
Kokkos::DualView<T[nTeams]> sum("sum")
...
parallel_for(policy, KOKKOS_LAMBDA(const auto& team) {
  int s = 0;
  parallel_reduce(TeamThreadRange(team,teamSize),[&](const auto& i, auto& update) {
     update += v.d_view(team.league_rank(),i);
  },s);
  team.team_barrier();
  sum.d_view(team.league_rank()) = s;
});
...

will give the same answer too.

I think something inside parallel_reduce expects the 3rd argument to be a thread-local variable but not something in global memory.

[crtrott]

This is weird. I would have thought that this should all give the same answer as long as T is word size or smaller.

[crtrott]

Got an even simpler reproducer .. . Only fails with OpenMP. Serial, Pthreads and CUDA are fine.

[crtrott]

#include<Kokkos_Core.hpp>

int main(int argc, char* argv[]) {
  Kokkos::initialize(argc,argv);
  {
     int N = (argc>1) ? atoi(argv[1]) : 10;
     int M = (argc>2) ? atoi(argv[2]) : 64;
     int R = (argc>3) ? atoi(argv[3]) : 8;

     Kokkos::View<double*> results1("r1",N);
     Kokkos::View<double*> results2("r2",N);
     Kokkos::View<double**> data("d",N,M);
     Kokkos::deep_copy(data,1);
     Kokkos::parallel_for(Kokkos::TeamPolicy<>(N,R), KOKKOS_LAMBDA (const Kokkos::TeamPolicy<>::member_type& team) {
       const int i = team.league_rank();
       Kokkos::parallel_reduce(Kokkos::TeamThreadRange(team,M), [&] (const int j, double& update) {
         update += data(i,j) + 1000*i + j;
       },results1(i));
       team.team_barrier();
       double s;
       Kokkos::parallel_reduce(Kokkos::TeamThreadRange(team,M), [&] (const int j, double& update) {
         update += data(i,j) + 1000*i + j;
       },s);
       results2(i) = s;
     });

     Kokkos::parallel_for(N, KOKKOS_LAMBDA (const int i) {
       printf("%i %lf %lf %lf\n",i,results1(i),results2(i),1.0*(M*(1000*i+1)+M*(M-1)/2));
     });
  }
  Kokkos::finalize();
}

[crtrott]

Root cause is that the OpenMP backend didn't produce a local variable for intermediate results but simply used the argument the user provided. I.e. the different threads would stomp over each other.

[crtrott]

Pull request issued #2079