From c53b635a3cb1df49dca3e45315515a747d7620ac Mon Sep 17 00:00:00 2001
From: Tom Deakin <tom.deakin@bristol.ac.uk>
Date: Wed, 3 Feb 2021 11:25:26 +0000
Subject: [PATCH] Add nstream kernel to HIP

---
 HIPStream.cpp | 17 +++++++++++++++++
 HIPStream.h   |  1 +
 2 files changed, 18 insertions(+)

diff --git a/HIPStream.cpp b/HIPStream.cpp
index d790ee5..fbc3b71 100644
--- a/HIPStream.cpp
+++ b/HIPStream.cpp
@@ -182,6 +182,23 @@ void HIPStream<T>::triad()
   check_error();
 }
 
+template <typename T>
+__global__ void nstream_kernel(T * a, const T * b, const T * c)
+{
+  const T scalar = startScalar;
+  const int i = hipBlockDim_x * hipBlockIdx_x + hipThreadIdx_x;
+  a[i] += b[i] + scalar * c[i];
+}
+
+template <class T>
+void HIPStream<T>::nstream()
+{
+  hipLaunchKernelGGL(HIP_KERNEL_NAME(nstream_kernel<T>), dim3(array_size/TBSIZE), dim3(TBSIZE), 0, 0, d_a, d_b, d_c);
+  check_error();
+  hipDeviceSynchronize();
+  check_error();
+}
+
 template <class T>
 __global__ void dot_kernel(const T * a, const T * b, T * sum, int array_size)
 {
diff --git a/HIPStream.h b/HIPStream.h
index fdab392..44a2893 100644
--- a/HIPStream.h
+++ b/HIPStream.h
@@ -41,6 +41,7 @@ class HIPStream : public Stream<T>
     virtual void add() override;
     virtual void mul() override;
     virtual void triad() override;
+    virtual void nstream() override;
     virtual T dot() override;
 
     virtual void init_arrays(T initA, T initB, T initC) override;