Unique token improvement

[crtrott]

Align interface with other Kokkos interfaces which default the execution space.

This also removes the usage of bitset in favor of using straight up int arrays as locks for UniqueToken in HIP and CUDA.
As a consequence atomic conflicts are vastly reduced, and also by default you are getting much better aligned indicies within a warp, improving data access patterns in common usecases.

Here is a performance example:

#include<Kokkos_Core.hpp>

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


     Kokkos::Experimental::UniqueToken<> token;
     Kokkos::View<double*> a("A",token.size());

     auto f1 = KOKKOS_LAMBDA(int) {
       int idx = token.acquire();
       for(int r=0; r<R; r++)
         a(idx) += a(idx)+1;
       Kokkos::memory_fence();
       token.release(idx);
     };
     auto f2 = KOKKOS_LAMBDA(int i) {
       int idx = i % a.extent(0);
       for(int r=0; r<R; r++)
         a(idx) += a(idx)+1;
       Kokkos::memory_fence();
     };

     Kokkos::parallel_for(N, f1);
     Kokkos::parallel_for(N, f2);
     Kokkos::fence();

     Kokkos::Timer timer;
     Kokkos::parallel_for(N, f1);
     Kokkos::fence();
     double time1 = timer.seconds();
     timer.reset();
     Kokkos::parallel_for(N, f2);
     Kokkos::fence();
     double time2 = timer.seconds();

     printf("%lf %lf\n",time1*1.e6, time2*1.e6);
  }
  Kokkos::finalize();
}

Prior to the change:
A100:

bash-4.2$ ./test.cuda 30000 1
1823.098000 16.880000
bash-4.2$ ./test.cuda 30000000 1
152508.359000 207.732000

MI100:

// MI100
bash-4.2$ ./test.host 30000 1
1585.862000 25.080000
bash-4.2$ ./test.host 30000000 1
495600.250000 357.516000

Post change:
A100

bash-4.2$ ./test.cuda 30000 1
17.940000 14.940000
bash-4.2$ ./test.cuda 30000000 1
836.698000 370.924000

MI100

bash-4.2$ ./test.host 30000 1
26.680000 24.060000
bash-4.2$ ./test.host 30000000 1
958.009000 353.497000

[masterleinad]

Is there a good reason to not also do that for SYCL?

[Rombur]

This is not needed. If HIP is not enabled, CMake will ignore the file

[crtrott]

removed

[dalg24]

Why did you change? This is not in the description of the PR.

[crtrott]

It is now ...

Everything else in Kokkos defaults the execution space. I think this should too.

[ajpowelsnl]

Hi @crtrott and @dalg24 a couple of questions about this issue: 1) will this item be in Kokkos 3.6? 2) if so, is it still blocked?

[crtrott]

@dalg24 @masterleinad finally passed!!

[masterleinad]

I slightly prefer to define and call protected constructors here to avoid allocating memory for m_locks twice like in https://github.com/kokkos/kokkos/pull/4748/files.

[crtrott]

what do you mean allocating twice?

[masterleinad]

It looks like m_scratchConcurrentBitset is unused now.

[masterleinad]

We should still be able to get rid of the now unused m_scratchConcurrentBitset.

[dalg24]

Retest this please

[dalg24]

4: [ RUN      ] openmptarget.unique_token_global
4: /var/jenkins/workspace/Kokkos/core/unit_test/TestUniqueToken.hpp:151: Failure
4: Expected equality of these values:
4:   sum
4:     Which is: 9999980
4:   int64_t(N) * R
4:     Which is: 10000000
4: [  FAILED  ] openmptarget.unique_token_global (280 ms)

[crtrott]

I don't get how either of these guys is failing (OpenMPTarget and OpenMP for NVHPC), I didn't really touch that stuff at all did I?

[dalg24]

Retest this please

[masterleinad]

We have seen that OpenMPTarget test failure occasionally.

[masterleinad]

Why should we keep m_scratchConcurrentBitset?

[crtrott]

Another PR failing UNiqueTOken in OpenMP here: https://cloud.cees.ornl.gov/jenkins-ci/blue/organizations/jenkins/Kokkos/detail/Kokkos/8028/pipeline/52

4: [ RUN      ] openmptarget.unique_token_global

4: /var/jenkins/workspace/Kokkos/core/unit_test/TestUniqueToken.hpp:151: Failure

4: Expected equality of these values:

4:   sum

4:     Which is: 9999990

4:   int64_t(N) * R

4:     Which is: 10000000

4: [  FAILED  ] openmptarget.unique_token_global (261 ms)

[dalg24]

Failure is unrelated (Clang crashing when compiling Kokkos::complex unit test, sigh)