Use recommended/max team size functions in Cuda ParallelFor and Reduce constructors

[tcclevenger]

Fixes #6814. Previous computation of m_team_size in cuda ParallelFor() and ParallelReduce() was intended to match policy.team_size_recommended(), but
missing extra scratch space allocation. Easiest to use team_size_recommended() instead. This matches exactly what is done for HIP backend.

Also, we were manually implementing policy.team_max_size() to verify m_team_size was not too large in ParallelFor(). Use team_max_size() instead.

[romintomasetti]

The changes looks good overall ! (I'm always happier with code reuse 👍)

However, I wonder about one thing that this PR will change.

I think

kokkos/core/src/Cuda/Kokkos_Cuda_Parallel_Team.hpp

Lines 144 to 147 in 9fff1e0

	space().impl_internal_space_instance(), attr, f,
	(size_t)impl_vector_length(),
	(size_t)team_scratch_size(0) + 2 * sizeof(double),
	(size_t)thread_scratch_size(0) + sizeof(double));

will set the team and thread scratch sizes regardless of whether it is a parallel for or a parallel reduce.

But I guess that for a parallel for, there is no need for the tiny additional scratch size.

[romintomasetti]

A possible fix would be to modify the recommended team size function along these lines:

template <class FunctorType, class ParallelTagType>
  int team_size_recommended(const FunctorType& f, const ParallelTagType&) const {
    using closure_type =
        Impl::ParallelFor<FunctorType, TeamPolicy<Properties...>>;
    cudaFuncAttributes attr =
        CudaParallelLaunch<closure_type, typename traits::launch_bounds>::
            get_cuda_func_attributes(space().cuda_device());
    const int block_size =
        Kokkos::Impl::cuda_get_opt_block_size<FunctorType,
                                              typename traits::launch_bounds>(
            space().impl_internal_space_instance(), attr, f,
            (size_t)impl_vector_length(),
            (size_t)team_scratch_size(0) + is it reduce ? 2 * sizeof(double) : 0,
            (size_t)thread_scratch_size(0) + is it reduce ? sizeof(double) : 0);
    return block_size / impl_vector_length();
  }

[masterleinad]

Even a TeamPolicy parallel_for can use nested parallel_reduce, parallel_scan, team_reduce, or team_scan. The scratch allocation here is independent of the allocation needed for the outer level parallel_reduce.

[crtrott]

What Daniel said.

[romintomasetti]

One major concern...

[crtrott]

What Daniel said.

[ndellingwood]

Does this merit a changelog entry for 4.4 #6914 ?

[masterleinad]

Does this merit a changelog entry for 4.4 #6914 ?

done