From 74f705cac99fea3df289b1a622dd7ff765eb6b71 Mon Sep 17 00:00:00 2001
From: Tom Deakin <tom.deakin@bristol.ac.uk>
Date: Mon, 7 Dec 2020 10:41:48 +0000
Subject: [PATCH] Update OpenMP for Issue #80

---
 OMPStream.cpp | 27 +++++++++++++--------------
 OMPStream.h   |  2 +-
 main.cpp      |  4 ++--
 3 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/OMPStream.cpp b/OMPStream.cpp
index e10c8c1..cdfce20 100644
--- a/OMPStream.cpp
+++ b/OMPStream.cpp
@@ -12,24 +12,22 @@
 #endif
 
 template <class T>
-OMPStream<T>::OMPStream(const unsigned int ARRAY_SIZE, T *a, T *b, T *c, int device)
+OMPStream<T>::OMPStream(const unsigned int ARRAY_SIZE, int device)
 {
   array_size = ARRAY_SIZE;
 
-#ifdef OMP_TARGET_GPU
-  omp_set_default_device(device);
-  // Set up data region on device
-  this->a = a;
-  this->b = b;
-  this->c = c;
-  #pragma omp target enter data map(alloc: a[0:array_size], b[0:array_size], c[0:array_size])
-  {}
-#else
   // Allocate on the host
   this->a = (T*)aligned_alloc(ALIGNMENT, sizeof(T)*array_size);
   this->b = (T*)aligned_alloc(ALIGNMENT, sizeof(T)*array_size);
   this->c = (T*)aligned_alloc(ALIGNMENT, sizeof(T)*array_size);
+
+#ifdef OMP_TARGET_GPU
+  omp_set_default_device(device);
+  // Set up data region on device
+  #pragma omp target enter data map(alloc: a[0:array_size], b[0:array_size], c[0:array_size])
+  {}
 #endif
+
 }
 
 template <class T>
@@ -43,11 +41,10 @@ OMPStream<T>::~OMPStream()
   T *c = this->c;
   #pragma omp target exit data map(release: a[0:array_size], b[0:array_size], c[0:array_size])
   {}
-#else
+#endif
   free(a);
   free(b);
   free(c);
-#endif
 }
 
 template <class T>
@@ -78,13 +75,15 @@ void OMPStream<T>::init_arrays(T initA, T initB, T initC)
 template <class T>
 void OMPStream<T>::read_arrays(std::vector<T>& h_a, std::vector<T>& h_b, std::vector<T>& h_c)
 {
+
 #ifdef OMP_TARGET_GPU
   T *a = this->a;
   T *b = this->b;
   T *c = this->c;
   #pragma omp target update from(a[0:array_size], b[0:array_size], c[0:array_size])
   {}
-#else
+#endif
+
   #pragma omp parallel for
   for (int i = 0; i < array_size; i++)
   {
@@ -92,7 +91,7 @@ void OMPStream<T>::read_arrays(std::vector<T>& h_a, std::vector<T>& h_b, std::ve
     h_b[i] = b[i];
     h_c[i] = c[i];
   }
-#endif
+
 }
 
 template <class T>
diff --git a/OMPStream.h b/OMPStream.h
index c475274..d8e9dc4 100644
--- a/OMPStream.h
+++ b/OMPStream.h
@@ -29,7 +29,7 @@ class OMPStream : public Stream<T>
     T *c;
 
   public:
-    OMPStream(const unsigned int, T*, T*, T*, int);
+    OMPStream(const unsigned int, int);
     ~OMPStream();
 
     virtual void copy() override;
diff --git a/main.cpp b/main.cpp
index 6d2f679..0139714 100644
--- a/main.cpp
+++ b/main.cpp
@@ -170,7 +170,7 @@ void run()
 
 #elif defined(OMP)
   // Use the OpenMP implementation
-  stream = new OMPStream<T>(ARRAY_SIZE, a.data(), b.data(), c.data(), deviceIndex);
+  stream = new OMPStream<T>(ARRAY_SIZE, deviceIndex);
 
 #endif
 
@@ -374,7 +374,7 @@ void run_triad()
 
 #elif defined(OMP)
   // Use the OpenMP implementation
-  stream = new OMPStream<T>(ARRAY_SIZE, a.data(), b.data(), c.data(), deviceIndex);
+  stream = new OMPStream<T>(ARRAY_SIZE, deviceIndex);
 
 #endif