TeamScan for CUDA, Pthreads, OpenMPTarget, HIP

[jrmadsen]

• implemented team-level parallel scan in CUDA

[jrmadsen]

• Addresses Team-level prefix scan not supported in hierarchical parallelism #3197
• Addresses parallel_scan on the device with a TeamThreadRange #45

[crtrott]

impl scatch
Threaded:

parallel_scan(..., N,...) {
  my_sum;
  for(int i = N/team_size*team_rank; ... ) {

     f(i,my_sum,false);
  }
  offset = team.team_scan(my_sum);
  my_sum = 0;
  for(int i = N/team_size*team_rank; ... ) {
     f(i,my_sum+offset,false);
  }
}

Cuda:

parallel_scan(..., N,...) {
  my_sum;
  offset = 0;
  for(int chunks = N/team_size) {
     f(chunks*team_size+team_rank,my_sum,false);
    local_offset = team.team_scan(my_sum);
     f(chunks*team_size+team_rank,my_sum+local_offset+offset,true);
     if(team_rank()==team_size-1)
     offset += ...
team.team_broadcast(offset);
  }
}

[jrmadsen]

@crtrott I ended up with a slightly different implementation that you recommended. This appears to work as long as the team-size is a power of 2, otherwise the team_scan throws an error.

[dalg24]

Retest this please

[crtrott]

We met today and discussed correctness issues. Jonathan is working on this.

[jrmadsen]

@dalg24 @masterleinad Do either of y'all know why the Jenkins build failed? I searched for all instances of the word "error" and "fail" in the Jenkins log and basically the only "failure" was "script exited with error code 2" but the build and tests appear to be fine

[masterleinad]

/var/jenkins/workspace/Kokkos/core/unit_test/TestTeamScan.hpp:66:16: error: unused variable 'teamSize' [clang-diagnostic-unused-variable]
          auto teamSize   = team.team_size();
               ^
/var/jenkins/workspace/Kokkos/core/unit_test/TestTeamScan.hpp:87:5: note: in instantiation of member function 'Test::TestTeamScan<Kokkos::Serial, short>::operator()' requested here
    (*this)(M, N, a_d, a_r);
    ^

[masterleinad]

/var/jenkins/workspace/Kokkos/install/include/Cuda/Kokkos_Cuda_Parallel.hpp(680): error: calling a constexpr __host__ function("operator()") from a __device__ function("exec_team") is not allowed. The experimental flag '--expt-relaxed-constexpr' can be used to allow this.
          detected during:
            instantiation of "std::enable_if<std::is_same<TagType, void>::value, void>::type Kokkos::Impl::ParallelFor<FunctorType, Kokkos::TeamPolicy<Properties...>, Kokkos::Cuda>::exec_team<TagType>(const Kokkos::Impl::ParallelFor<FunctorType, Kokkos::TeamPolicy<Properties...>, Kokkos::Cuda>::Member &) const [with FunctorType=lambda [](const Kokkos::Impl::CudaTeamMember &)->void, Properties=<Kokkos::Cuda>, TagType=void]" 
(731): here
            instantiation of "void Kokkos::Impl::ParallelFor<FunctorType, Kokkos::TeamPolicy<Properties...>, Kokkos::Cuda>::operator()() const [with FunctorType=lambda [](const Kokkos::Impl::CudaTeamMember &)->void, Properties=<Kokkos::Cuda>]" 
/var/jenkins/workspace/Kokkos/install/include/Cuda/Kokkos_Cuda_KernelLaunch.hpp(121): here
            instantiation of "void Kokkos::Impl::cuda_parallel_launch_local_memory(DriverType) [with DriverType=Kokkos::Impl::ParallelFor<lambda [](const Kokkos::Impl::CudaTeamMember &)->void, Kokkos::TeamPolicy<Kokkos::Cuda>, Kokkos::Cuda>]" 
/var/jenkins/workspace/Kokkos/install/include/Cuda/Kokkos_Cuda_KernelLaunch.hpp(319): here
            instantiation of "std::decay_t<decltype((<expression>))> Kokkos::Impl::CudaParallelLaunchKernelFunc<DriverType, Kokkos::LaunchBounds<0U, 0U>, Kokkos::Impl::Experimental::CudaLaunchMechanism::LocalMemory>::get_kernel_func() [with DriverType=Kokkos::Impl::ParallelFor<lambda [](const Kokkos::Impl::CudaTeamMember &)->void, Kokkos::TeamPolicy<Kokkos::Cuda>, Kokkos::Cuda>]" 
/var/jenkins/workspace/Kokkos/install/include/Cuda/Kokkos_Cuda_KernelLaunch.hpp(646): here
            instantiation of "cudaFuncAttributes Kokkos::Impl::CudaParallelLaunchImpl<DriverType, Kokkos::LaunchBounds<MaxThreadsPerBlock, MinBlocksPerSM>, LaunchMechanism>::get_cuda_func_attributes() [with DriverType=Kokkos::Impl::ParallelFor<lambda [](const Kokkos::Impl::CudaTeamMember &)->void, Kokkos::TeamPolicy<Kokkos::Cuda>, Kokkos::Cuda>, MaxThreadsPerBlock=0U, MinBlocksPerSM=0U, LaunchMechanism=Kokkos::Impl::Experimental::CudaLaunchMechanism::LocalMemory]" 
(764): here
            instantiation of "Kokkos::Impl::ParallelFor<FunctorType, Kokkos::TeamPolicy<Properties...>, Kokkos::Cuda>::ParallelFor(const FunctorType &, const Kokkos::Impl::ParallelFor<FunctorType, Kokkos::TeamPolicy<Properties...>, Kokkos::Cuda>::Policy &) [with FunctorType=lambda [](const Kokkos::Impl::CudaTeamMember &)->void, Properties=<Kokkos::Cuda>]" 
/var/jenkins/workspace/Kokkos/install/include/Kokkos_Parallel.hpp(168): here
            instantiation of "void Kokkos::parallel_for(const ExecPolicy &, const FunctorType &, const std::__cxx11::string &, std::enable_if<Kokkos::is_execution_policy<ExecPolicy>::value, void>::type *) [with ExecPolicy=Kokkos::TeamPolicy<Kokkos::Cuda>, FunctorType=lambda [](const Kokkos::Impl::CudaTeamMember &)->void]" 
/var/jenkins/workspace/Kokkos/core/unit_test/TestTeamScan.hpp(80): here
            instantiation of "void Test::TestTeamScan<Device, DataType>::operator()(int32_t, int32_t, Test::TestTeamScan<Device, DataType>::view_type, Test::TestTeamScan<Device, DataType>::view_type) const [with Device=Kokkos::Cuda, DataType=int16_t]" 
/var/jenkins/workspace/Kokkos/core/unit_test/TestTeamScan.hpp(87): here

[masterleinad]

/var/jenkins/workspace/Kokkos/core/unit_test/TestTeamScan.hpp:125:57: error: typedef 'using TEST_POLICY = class Kokkos::TeamPolicy<Kokkos::OpenMP>' locally defined but not used [-Werror=unused-local-typedefs]
   using TEST_POLICY = Kokkos::TeamPolicy<TEST_EXECSPACE>;
                                                         ^

[jrmadsen]

@crtrott @dalg24 So everything is passing here except for some SYCL stuff, which I didn't touch but it looks like the new tests instantiate something that is incomplete in the SYCL backend. How do we proceed here?

[crtrott]

Minus the "exclude the test in the exclude list for SYCL" this looks good!

[Rombur]

The HIP backend looks good

[dalg24]

Please clarify the tolerance with floating point numbers

[crtrott]

Retest this please.

[crtrott]

So intereseting question: should we just merge? Windows likely failing because of the MSVC stuff, Jenkins its the AMD node, and travis is one timeout ...

[crtrott]

Also: Jonathan you want to rewrite history or should I squash commit?

[Char-Aznable]

Hi @jrmadsen @crtrott , I have code using something like parallel_scan(TeamThreadRange(team, n), [&](const int i, double& udpate, const bool isFinal) { ... }); with n = 200 and it aborts at the line

kokkos/core/src/Cuda/Kokkos_Cuda_ReduceScan.hpp

Lines 698 to 700 in e483144

	if (BlockSizeMask & blockDim.y) {
	Kokkos::abort("Cuda::cuda_intra_block_scan requires power-of-two blockDim");
	}

and

(cuda-gdb) p blockDim.y
$4 = 224
(cuda-gdb) p blockDim.y - 1
$5 = 223

any idea what went wrong here? Does the code in this PR require the item counts n to be the same as number of threads in the team because it seems to be calling team_member.team_scan() judging from the stack trace?

[Char-Aznable]

Judging from the unit test cases, I guess this implementation only works if the work item counts to TeamThreadRange is power of 2?

[jrmadsen]

Judging from the unit test cases, I guess this implementation only works if the work item counts to TeamThreadRange is power of 2?

Yes

[Char-Aznable]

Would it be difficult to support non-power-of-2 work counts? I can help implementing it with the CUDA backend at least. I have a lot of small loops of a few hundreds work counts and it would be a big hit in performance if forced to use power-of-2 loop, especially on the host because I need to convert all the host side to be power of 2 to be portable

[masterleinad]

We will discuss this but in general, contributions are very welcome.

[masterleinad]

@Char-Aznable We decided that we are interested in making non-power-2 team sizes work, see #4146. Any help in implementing that is very welcome!

[Char-Aznable]

Great! I'll take a look at the code and see what I can do