From 90fb104daed03f5f986199aaf8943250bb8c9639 Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Thu, 21 Mar 2024 23:07:46 +0100
Subject: [PATCH] Implement output layer

---
 include/layers/conv2d.cuh  |  1 -
 include/layers/dense.cuh   |  2 --
 include/layers/layer.cuh   |  2 --
 include/layers/output.cuh  | 39 ++++++++++++++++++++++++++++++++++++++
 include/model/model.hpp    |  2 ++
 src/layers/output.cu       | 22 +++++++++++++++++++++
 test/layers/test_input.cu  |  9 +++++----
 test/layers/test_output.cu | 24 +++++++++++++++++++++++
 8 files changed, 92 insertions(+), 9 deletions(-)
 create mode 100644 include/layers/output.cuh
 create mode 100644 src/layers/output.cu
 create mode 100644 test/layers/test_output.cu

diff --git a/include/layers/conv2d.cuh b/include/layers/conv2d.cuh
index 5cd6ecf..01ffd00 100644
--- a/include/layers/conv2d.cuh
+++ b/include/layers/conv2d.cuh
@@ -1,7 +1,6 @@
 #ifndef CUDANET_CONV_LAYER_H
 #define CUDANET_CONV_LAYER_H
 
-#include <string>
 #include <vector>
 
 #include "activation.cuh"
diff --git a/include/layers/dense.cuh b/include/layers/dense.cuh
index 3922a3f..91fb193 100644
--- a/include/layers/dense.cuh
+++ b/include/layers/dense.cuh
@@ -1,8 +1,6 @@
 #ifndef CUDANET_DENSE_LAYER_H
 #define CUDANET_DENSE_LAYER_H
 
-#include <functional>
-#include <string>
 #include <vector>
 
 #include "layer.cuh"
diff --git a/include/layers/layer.cuh b/include/layers/layer.cuh
index 3c1be19..b68beb7 100644
--- a/include/layers/layer.cuh
+++ b/include/layers/layer.cuh
@@ -2,8 +2,6 @@
 #ifndef CUDANET_I_LAYER_H
 #define CUDANET_I_LAYER_H
 
-#include <vector>
-
 namespace CUDANet::Layers {
 
 /**
diff --git a/include/layers/output.cuh b/include/layers/output.cuh
new file mode 100644
index 0000000..71ff615
--- /dev/null
+++ b/include/layers/output.cuh
@@ -0,0 +1,39 @@
+#ifndef CUDANET_OUTPUT_LAYER_H
+#define CUDANET_OUTPUT_LAYER_H
+
+#include "layer.cuh"
+
+namespace CUDANet::Layers {
+
+class Output : public SequentialLayer {
+    public:
+        /**
+         * @brief Create a new Output layer
+         * 
+         * @param inputSize Size of the input vector
+         */
+        explicit Output(int inputSize);
+
+        /**
+         * @brief Destroy the Output layer
+         * 
+         */
+        ~Output();
+
+        /**
+         * @brief Forward pass of the output layer. Just copies the input from device to host
+         * 
+         * @param input Device pointer to the input vector
+         * @return Host pointer to the output vector
+         */
+        float* forward(const float* input);
+
+    private:
+        int    inputSize;
+        float* h_output;
+};
+
+
+} // namespace CUDANet::Layers
+
+#endif  // CUDANET_OUTPUT_LAYER_H
\ No newline at end of file
diff --git a/include/model/model.hpp b/include/model/model.hpp
index 30e2a43..be7711e 100644
--- a/include/model/model.hpp
+++ b/include/model/model.hpp
@@ -4,7 +4,9 @@
 #include <string>
 #include <vector>
 #include <map>
+
 #include "layer.cuh"
+#include "input.cuh"
 
 namespace CUDANet {
 
diff --git a/src/layers/output.cu b/src/layers/output.cu
new file mode 100644
index 0000000..f032376
--- /dev/null
+++ b/src/layers/output.cu
@@ -0,0 +1,22 @@
+#include "output.cuh"
+
+#include "cuda_helper.cuh"
+
+using namespace CUDANet::Layers;
+
+
+Output::Output(int inputSize) : inputSize(inputSize) {
+    h_output = (float*) malloc(sizeof(float) * inputSize);
+}
+
+Output::~Output() {
+    free(h_output);
+}
+
+float* Output::forward(const float* input) {
+    CUDA_CHECK(cudaMemcpy(
+        h_output, input, sizeof(float) * inputSize, cudaMemcpyDeviceToHost
+    ));
+
+    return h_output;
+}
\ No newline at end of file
diff --git a/test/layers/test_input.cu b/test/layers/test_input.cu
index 1986d3c..c4d4561 100644
--- a/test/layers/test_input.cu
+++ b/test/layers/test_input.cu
@@ -1,16 +1,17 @@
+#include <cuda_runtime_api.h>
 #include <gtest/gtest.h>
 
-#include "cuda_helper.cuh"
 #include "input.cuh"
 
-TEST(InputLayerTest, Init) {
+TEST(InputLayerTest, InputForward) {
     std::vector<float> input = {0.573f, 0.619f, 0.732f, 0.055f, 0.243f, 0.316f};
     CUDANet::Layers::Input inputLayer(6);
     float*                 d_output = inputLayer.forward(input.data());
 
     std::vector<float> output(6);
-    CUDA_CHECK(cudaMemcpy(
+    cudaError_t        cudaStatus = cudaMemcpy(
         output.data(), d_output, sizeof(float) * 6, cudaMemcpyDeviceToHost
-    ));
+    );
+    EXPECT_EQ(cudaStatus, cudaSuccess);
     EXPECT_EQ(input, output);
 }
\ No newline at end of file
diff --git a/test/layers/test_output.cu b/test/layers/test_output.cu
new file mode 100644
index 0000000..3fa529d
--- /dev/null
+++ b/test/layers/test_output.cu
@@ -0,0 +1,24 @@
+#include <cuda_runtime_api.h>
+#include <gtest/gtest.h>
+
+#include "output.cuh"
+
+TEST(OutputLayerTest, OutputForward) {
+    cudaError_t cudaStatus;
+
+    std::vector<float> input = {0.573f, 0.619f, 0.732f, 0.055f, 0.243f, 0.316f};
+    float*             d_input;
+    cudaStatus = cudaMalloc((void**)&d_input, sizeof(float) * 6);
+    EXPECT_EQ(cudaStatus, cudaSuccess);
+    cudaStatus = cudaMemcpy(
+        d_input, input.data(), sizeof(float) * 6, cudaMemcpyHostToDevice
+    );
+    EXPECT_EQ(cudaStatus, cudaSuccess);
+
+    CUDANet::Layers::Output outputLayer(6);
+    float* h_output = outputLayer.forward(d_input);
+
+    for (int i = 0; i < 6; ++i) {
+        EXPECT_EQ(input[i], h_output[i]);
+    }
+}
\ No newline at end of file