Adjusted the Kokkos implementation to fix view initialisation, and store local copies of views for lambda scoping
This commit is contained in:
Matthew Martineau
parent
1a60f130eb
commit
894829cb05
@ -7,7 +7,7 @@
|
||||
|
||||
#include "KOKKOSStream.hpp"
|
||||
|
||||
using Kokkos::parallel_for;
|
||||
using namespace Kokkos;
|
||||
|
||||
template <class T>
|
||||
KOKKOSStream<T>::KOKKOSStream(
|
||||
@ -16,21 +16,21 @@ KOKKOSStream<T>::KOKKOSStream(
|
||||
{
|
||||
Kokkos::initialize();
|
||||
|
||||
new(d_a) Kokkos::View<double*, DEVICE>("d_a", ARRAY_SIZE);
|
||||
new(d_b) Kokkos::View<double*, DEVICE>("d_b", ARRAY_SIZE);
|
||||
new(d_c) Kokkos::View<double*, DEVICE>("d_c", ARRAY_SIZE);
|
||||
new(hm_a) Kokkos::View<double*>::HostMirror();
|
||||
new(hm_b) Kokkos::View<double*>::HostMirror();
|
||||
new(hm_c) Kokkos::View<double*>::HostMirror();
|
||||
hm_a = Kokkos::create_mirror_view(d_a);
|
||||
hm_b = Kokkos::create_mirror_view(d_b);
|
||||
hm_c = Kokkos::create_mirror_view(d_c);
|
||||
d_a = new View<double*, DEVICE>("d_a", ARRAY_SIZE);
|
||||
d_b = new View<double*, DEVICE>("d_b", ARRAY_SIZE);
|
||||
d_c = new View<double*, DEVICE>("d_c", ARRAY_SIZE);
|
||||
hm_a = new View<double*, DEVICE>::HostMirror();
|
||||
hm_b = new View<double*, DEVICE>::HostMirror();
|
||||
hm_c = new View<double*, DEVICE>::HostMirror();
|
||||
*hm_a = create_mirror_view(*d_a);
|
||||
*hm_b = create_mirror_view(*d_b);
|
||||
*hm_c = create_mirror_view(*d_c);
|
||||
}
|
||||
|
||||
template <class T>
|
||||
KOKKOSStream<T>::~KOKKOSStream()
|
||||
{
|
||||
Kokkos::finalize();
|
||||
finalize();
|
||||
}
|
||||
|
||||
template <class T>
|
||||
@ -39,65 +39,81 @@ void KOKKOSStream<T>::write_arrays(
|
||||
{
|
||||
for(int ii = 0; ii < array_size; ++ii)
|
||||
{
|
||||
hm_a(ii) = a[ii];
|
||||
hm_b(ii) = b[ii];
|
||||
hm_c(ii) = c[ii];
|
||||
(*hm_a)(ii) = a[ii];
|
||||
(*hm_b)(ii) = b[ii];
|
||||
(*hm_c)(ii) = c[ii];
|
||||
}
|
||||
Kokkos::deep_copy(hm_a, d_a);
|
||||
Kokkos::deep_copy(hm_b, d_b);
|
||||
Kokkos::deep_copy(hm_c, d_c);
|
||||
deep_copy(*hm_a, *d_a);
|
||||
deep_copy(*hm_b, *d_b);
|
||||
deep_copy(*hm_c, *d_c);
|
||||
}
|
||||
|
||||
template <class T>
|
||||
void KOKKOSStream<T>::read_arrays(
|
||||
std::vector<T>& a, std::vector<T>& b, std::vector<T>& c)
|
||||
{
|
||||
Kokkos::deep_copy(d_a, hm_a);
|
||||
Kokkos::deep_copy(d_a, hm_b);
|
||||
Kokkos::deep_copy(d_a, hm_c);
|
||||
deep_copy(*d_a, *hm_a);
|
||||
deep_copy(*d_b, *hm_b);
|
||||
deep_copy(*d_c, *hm_c);
|
||||
for(int ii = 0; ii < array_size; ++ii)
|
||||
{
|
||||
a[ii] = hm_a(ii);
|
||||
b[ii] = hm_b(ii);
|
||||
c[ii] = hm_c(ii);
|
||||
a[ii] = (*hm_a)(ii);
|
||||
b[ii] = (*hm_b)(ii);
|
||||
c[ii] = (*hm_c)(ii);
|
||||
}
|
||||
}
|
||||
|
||||
template <class T>
|
||||
void KOKKOSStream<T>::copy()
|
||||
{
|
||||
Kokkos::parallel_for(array_size, KOKKOS_LAMBDA (const int index)
|
||||
View<double*, DEVICE> a(*d_a);
|
||||
View<double*, DEVICE> b(*d_b);
|
||||
View<double*, DEVICE> c(*d_c);
|
||||
|
||||
parallel_for(array_size, KOKKOS_LAMBDA (const int index)
|
||||
{
|
||||
d_c[index] = d_a[index];
|
||||
c[index] = a[index];
|
||||
});
|
||||
}
|
||||
|
||||
template <class T>
|
||||
void KOKKOSStream<T>::mul()
|
||||
{
|
||||
View<double*, DEVICE> a(*d_a);
|
||||
View<double*, DEVICE> b(*d_b);
|
||||
View<double*, DEVICE> c(*d_c);
|
||||
|
||||
const T scalar = 3.0;
|
||||
parallel_for(array_size, KOKKOS_LAMBDA (const int index)
|
||||
{
|
||||
d_b[index] = scalar*d_c[index];
|
||||
b[index] = scalar*c[index];
|
||||
});
|
||||
}
|
||||
|
||||
template <class T>
|
||||
void KOKKOSStream<T>::add()
|
||||
{
|
||||
View<double*, DEVICE> a(*d_a);
|
||||
View<double*, DEVICE> b(*d_b);
|
||||
View<double*, DEVICE> c(*d_c);
|
||||
|
||||
parallel_for(array_size, KOKKOS_LAMBDA (const int index)
|
||||
{
|
||||
d_c[index] = d_a[index] + d_b[index];
|
||||
c[index] = a[index] + b[index];
|
||||
});
|
||||
}
|
||||
|
||||
template <class T>
|
||||
void KOKKOSStream<T>::triad()
|
||||
{
|
||||
View<double*, DEVICE> a(*d_a);
|
||||
View<double*, DEVICE> b(*d_b);
|
||||
View<double*, DEVICE> c(*d_c);
|
||||
|
||||
const T scalar = 3.0;
|
||||
parallel_for(array_size, KOKKOS_LAMBDA (const int index)
|
||||
{
|
||||
d_a[index] = d_b[index] + scalar*d_c[index];
|
||||
a[index] = b[index] + scalar*c[index];
|
||||
});
|
||||
}
|
||||
|
||||
@ -118,6 +134,6 @@ std::string getDeviceDriver(const int device)
|
||||
return "Kokkos";
|
||||
}
|
||||
|
||||
template class KOKKOSStream<float>;
|
||||
//template class KOKKOSStream<float>;
|
||||
template class KOKKOSStream<double>;
|
||||
|
||||
|
||||
@ -17,8 +17,11 @@
|
||||
|
||||
#define IMPLEMENTATION_STRING "KOKKOS"
|
||||
|
||||
#define DEVICE Kokkos::OpenMP
|
||||
|
||||
#ifdef KOKKOS_TARGET_CPU
|
||||
#define DEVICE Kokkos::OpenMP
|
||||
#else
|
||||
#define DEVICE Kokkos::Cuda
|
||||
#endif
|
||||
|
||||
template <class T>
|
||||
class KOKKOSStream : public Stream<T>
|
||||
@ -28,12 +31,12 @@ class KOKKOSStream : public Stream<T>
|
||||
unsigned int array_size;
|
||||
|
||||
// Device side pointers to arrays
|
||||
Kokkos::View<double*, DEVICE> d_a;
|
||||
Kokkos::View<double*, DEVICE> d_b;
|
||||
Kokkos::View<double*, DEVICE> d_c;
|
||||
Kokkos::View<double*>::HostMirror hm_a;
|
||||
Kokkos::View<double*>::HostMirror hm_b;
|
||||
Kokkos::View<double*>::HostMirror hm_c;
|
||||
Kokkos::View<double*, DEVICE>* d_a;
|
||||
Kokkos::View<double*, DEVICE>* d_b;
|
||||
Kokkos::View<double*, DEVICE>* d_c;
|
||||
Kokkos::View<double*>::HostMirror* hm_a;
|
||||
Kokkos::View<double*>::HostMirror* hm_b;
|
||||
Kokkos::View<double*>::HostMirror* hm_c;
|
||||
|
||||
public:
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user