Migrate output layer

include/cudanet.cuh

@@ -19,7 +19,7 @@
 #include "input.hpp"
 #include "layer.hpp"
 #include "max_pooling.hpp"
-#include "output.cuh"
+#include "output.hpp"
 
 // Models
 #include "model.hpp"

include/layers/output.hpp

@@ -46,6 +46,12 @@ class Output : public SequentialLayer {
   private:
     int    inputSize;
     float* h_output;
+
+    float* forwardCPU(const float* input);
+
+#ifdef USE_CUDA
+    float* forwardCUDA(const float* input);
+#endif
 };
 
 }  // namespace CUDANet::Layers

include/model/model.hpp

@@ -8,7 +8,7 @@
 #include "input.hpp"
 #include "layer.hpp"
 #include "module.hpp"
-#include "output.cuh"
+#include "output.hpp"
 
 namespace CUDANet {
 

src/backends/cuda/layers/output.cu

@@ -0,0 +1,14 @@
+#include "output.hpp"
+
+#include "cuda_helper.cuh"
+
+using namespace CUDANet::Layers;
+
+float* Output::forwardCUDA(const float* input) {
+    CUDA_CHECK(cudaMemcpy(
+        h_output, input, sizeof(float) * inputSize, cudaMemcpyDeviceToHost
+    ));
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return h_output;
+}

src/layers/output.cu

@@ -1,6 +1,5 @@
-#include "output.cuh"
-
-#include "cuda_helper.cuh"
+#include "output.hpp"
+#include <stdexcept>
 
 using namespace CUDANet::Layers;
 
@@ -13,13 +12,16 @@ Output::~Output() {
     free(h_output);
 }
 
-float* Output::forward(const float* input) {
-    CUDA_CHECK(cudaMemcpy(
-        h_output, input, sizeof(float) * inputSize, cudaMemcpyDeviceToHost
-    ));
-    CUDA_CHECK(cudaDeviceSynchronize());
+float* Output::forwardCPU(const float* input) {
+    throw std::logic_error("Not implemented");
+}
 
-    return h_output;
+float* Output::forward(const float* input) {
+#ifdef USE_CUDA
+    return forwardCUDA(input);
+#else
+    return forwardCPU(input);
+#endif
 }
 
 int Output::getOutputSize() {

test/cuda/layers/test_output.cu

@@ -1,7 +1,7 @@
 #include <cuda_runtime_api.h>
 #include <gtest/gtest.h>
 
-#include "output.cuh"
+#include "output.hpp"
 
 TEST(OutputLayerTest, OutputForward) {
     cudaError_t cudaStatus;