From f7b525e494681da8860e48f3c0acde5bc97ede4e Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Tue, 10 Sep 2024 19:11:21 +0200
Subject: [PATCH] Migrate input layer

---
 include/cudanet.cuh                     |  2 +-
 include/layers/{input.cuh => input.hpp} | 15 ++++++++--
 include/model/model.hpp                 |  2 +-
 src/{ => backends/cuda}/layers/input.cu | 17 +++---------
 src/layers/input.cpp                    | 37 +++++++++++++++++++++++++
 src/model/model.cpp                     |  2 +-
 test/cuda/layers/test_input.cu          |  2 +-
 7 files changed, 57 insertions(+), 20 deletions(-)
 rename include/layers/{input.cuh => input.hpp} (84%)
 rename src/{ => backends/cuda}/layers/input.cu (59%)
 create mode 100644 src/layers/input.cpp

diff --git a/include/cudanet.cuh b/include/cudanet.cuh
index c5dcb78..b08f7be 100644
--- a/include/cudanet.cuh
+++ b/include/cudanet.cuh
@@ -16,7 +16,7 @@
 #include "concat.hpp"
 #include "conv2d.cuh"
 #include "dense.hpp"
-#include "input.cuh"
+#include "input.hpp"
 #include "layer.hpp"
 #include "max_pooling.hpp"
 #include "output.cuh"
diff --git a/include/layers/input.cuh b/include/layers/input.hpp
similarity index 84%
rename from include/layers/input.cuh
rename to include/layers/input.hpp
index 21fb5c4..638b9e8 100644
--- a/include/layers/input.cuh
+++ b/include/layers/input.hpp
@@ -40,16 +40,25 @@ class Input : public SequentialLayer {
      */
     int getOutputSize();
 
-        /**
+    /**
      * @brief Get input size
-     * 
+     *
      * @return int input size
      */
     int getInputSize();
 
   private:
-    int    inputSize;
+    int inputSize;
+
+    float* forwardCPU(const float* input);
+
+#ifdef USE_CUDA
     float* d_output;
+
+    float* forwardCUDA(const float* input);
+    void   initCUDA();
+    void   delCUDA();
+#endif
 };
 
 }  // namespace CUDANet::Layers
diff --git a/include/model/model.hpp b/include/model/model.hpp
index a6ea7fd..6e50e4b 100644
--- a/include/model/model.hpp
+++ b/include/model/model.hpp
@@ -5,7 +5,7 @@
 #include <unordered_map>
 #include <vector>
 
-#include "input.cuh"
+#include "input.hpp"
 #include "layer.hpp"
 #include "module.hpp"
 #include "output.cuh"
diff --git a/src/layers/input.cu b/src/backends/cuda/layers/input.cu
similarity index 59%
rename from src/layers/input.cu
rename to src/backends/cuda/layers/input.cu
index 39531ad..05b44be 100644
--- a/src/layers/input.cu
+++ b/src/backends/cuda/layers/input.cu
@@ -1,31 +1,22 @@
 #include "cuda_helper.cuh"
-#include "input.cuh"
+#include "input.hpp"
 
 using namespace CUDANet::Layers;
 
-Input::Input(int inputSize) : inputSize(inputSize) {
+void Input::initCUDA() {
     d_output = nullptr;
     CUDA_CHECK(cudaMalloc((void**)&d_output, sizeof(float) * inputSize));
 }
 
-Input::~Input() {
+void Input::delCUDA() {
     cudaFree(d_output);
 }
 
-float* Input::forward(const float* input) {
+float* Input::forwardCUDA(const float* input) {
     CUDA_CHECK(cudaMemcpy(
         d_output, input, sizeof(float) * inputSize, cudaMemcpyHostToDevice
     ));
     CUDA_CHECK(cudaDeviceSynchronize());
 
     return d_output;
-}
-
-int Input::getOutputSize() {
-    return inputSize;
-}
-
-
-int Input::getInputSize() {
-    return inputSize;
 }
\ No newline at end of file
diff --git a/src/layers/input.cpp b/src/layers/input.cpp
new file mode 100644
index 0000000..537302a
--- /dev/null
+++ b/src/layers/input.cpp
@@ -0,0 +1,37 @@
+#include <stdexcept>
+
+#include "input.hpp"
+
+using namespace CUDANet::Layers;
+
+Input::Input(int inputSize) : inputSize(inputSize) {
+#ifdef USE_CUDA
+    initCUDA();
+#endif
+}
+
+Input::~Input() {
+#ifdef USE_CUDA
+    delCUDA();
+#endif
+}
+
+float* Input::forwardCPU(const float* input) {
+    throw std::logic_error("Not implemented");
+}
+
+float* Input::forward(const float* input) {
+#ifdef USE_CUDA
+    return forwardCUDA(input);
+#else
+    return forwardCPU(input);
+#endif
+}
+
+int Input::getOutputSize() {
+    return inputSize;
+}
+
+int Input::getInputSize() {
+    return inputSize;
+}
\ No newline at end of file
diff --git a/src/model/model.cpp b/src/model/model.cpp
index 729ac81..c5ef706 100644
--- a/src/model/model.cpp
+++ b/src/model/model.cpp
@@ -7,7 +7,7 @@
 #include <unordered_map>
 #include <vector>
 
-#include "input.cuh"
+#include "input.hpp"
 #include "layer.hpp"
 #include "batch_norm.cuh"
 
diff --git a/test/cuda/layers/test_input.cu b/test/cuda/layers/test_input.cu
index e10ebb3..a793f1a 100644
--- a/test/cuda/layers/test_input.cu
+++ b/test/cuda/layers/test_input.cu
@@ -1,7 +1,7 @@
 #include <cuda_runtime_api.h>
 #include <gtest/gtest.h>
 
-#include "input.cuh"
+#include "input.hpp"
 
 TEST(InputLayerTest, InputForward) {
     std::vector<float> input = {0.573f, 0.619f, 0.732f, 0.055f, 0.243f, 0.316f};