TeamPolicy with reducers with valuetypes without += broken on CUDA

[crtrott]

Reported here:
trilinos/Trilinos#6000

I root caused the issue to our TeamPolicy::team_size_max function not taking reducers in addition to the functor and the parallel reduce tag. Hence during the "figure out how many registers this kernel will use" step, it tries to instantiate the kernel as if it were a sum reduction.
Technically that makes that wrong for any reducer other than sum. But for every reducer which has a native value type (like double) it happens to work, though the register count determination might be slightly off - if so it would result in a runtime dispatch error.

Here is a reproducer:

#include<Kokkos_Core.hpp>
#include<cmath>

int main(int argc, char* argv[]) {
  Kokkos::initialize(argc,argv);
  {
    using ExecSpace = Kokkos::DefaultExecutionSpace;
    using ReducerType = Kokkos::MinMax<double, ExecSpace>;
    using ReducerValueType = typename ReducerType::value_type;
    using DynamicScheduleType = Kokkos::Schedule<Kokkos::Dynamic>;
    using TeamPolicyType = Kokkos::TeamPolicy<ExecSpace, DynamicScheduleType>;
    using TeamHandleType = typename TeamPolicyType::member_type;

    static constexpr int num_teams = 1000;
    static constexpr int num_threads = 1000;
    ReducerValueType val;
    ReducerType reducer(val);
    Kokkos::parallel_reduce(
      TeamPolicyType(num_teams, Kokkos::AUTO),
      KOKKOS_LAMBDA(const TeamHandleType& team, ReducerValueType& teamVal) {
    }, reducer);
  }
  Kokkos::finalize();
}

To fix this: the (non-deprecated) max and recommended team size functions need an option to take in a reducer too. There are two options:

• add an overload team_size_max(FunctorType,ParallelReduceTag,Reducer)
• make ParallelReduceTag templated on ReducerType

I am thinking option 1 is better.

What needs to happen:

• add a test which exposes this problem: i.e. test ALL reducers with TeamPolicy too
• add that overload
• use it inside the ParallelReduce class to determine max and recommended team size where appropriate

[crtrott]

I think there is a band aid solution where we just hand-code something like 128 instead of calling team_size_max.

[crtrott]

OK I put in a PR with a proper fix for CUDA (but the new function overloads were not implemented for the other backends yet - but they can simply drop the reducer and call the overloads without reducers).

[srajama1]

@crtrott @DavidPoliakoff This needs to be patched to Trilinos.

[masterleinad]

@srajama1 See #2499 that pull request will be merged eventually.

[srajama1]

@masterleinad : Thanks !