From cfc5c46d5e2998e0b166fc6d741179e60d9f99d3 Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Mon, 4 Mar 2024 22:16:03 +0100
Subject: [PATCH] Initialize conv2d layer

---
 CMakeLists.txt            |  1 +
 include/layers/conv.cuh   | 31 --------------
 include/layers/conv2d.cuh | 60 +++++++++++++++++++++++++++
 include/layers/dense.cuh  |  2 +-
 src/layers/conv2d.cu      | 86 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 148 insertions(+), 32 deletions(-)
 delete mode 100644 include/layers/conv.cuh
 create mode 100644 include/layers/conv2d.cuh
 create mode 100644 src/layers/conv2d.cu

diff --git a/CMakeLists.txt b/CMakeLists.txt
index ed802bb..90f947d 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -13,6 +13,7 @@ set(LIBRARY_SOURCES
     src/kernels/activations.cu
     src/kernels/padding.cu
     src/layers/dense.cu
+    src/layers/conv2d.cu
 )
 
 set(CMAKE_CUDA_ARCHITECTURES 75)
diff --git a/include/layers/conv.cuh b/include/layers/conv.cuh
deleted file mode 100644
index 81bd1ad..0000000
--- a/include/layers/conv.cuh
+++ /dev/null
@@ -1,31 +0,0 @@
-#ifndef CONV_LAYER_H
-#define CONV_LAYER_H
-
-#include <cublas_v2.h>
-
-namespace Layers {
-
-class Conv {
-  public:
-    Conv(
-        int            inputSize,
-        int            outputSize,
-        int            kernelSize,
-        cublasHandle_t cublasHandle
-    );
-    ~Conv();
-
-    void forward(const float* input, float* output);
-
-  private:
-    int            inputSize;
-    int            outputSize;
-    int            kernelSize;
-    cublasHandle_t cublasHandle;
-    float*         d_weights;
-    float*         d_biases;
-};
-
-}  // namespace Layers
-
-#endif  // CONV_LAYER_H
diff --git a/include/layers/conv2d.cuh b/include/layers/conv2d.cuh
new file mode 100644
index 0000000..bf58b41
--- /dev/null
+++ b/include/layers/conv2d.cuh
@@ -0,0 +1,60 @@
+#ifndef CONV_LAYER_H
+#define CONV_LAYER_H
+
+#include <cublas_v2.h>
+
+#include <string>
+#include <vector>
+
+#include "activations.cuh"
+
+namespace Layers {
+
+class Conv2d {
+  public:
+    Conv2d(
+        int            inputSize,
+        int            inputChannels,
+        int            kernelSize,
+        int            stride,
+        std::string    padding,
+        int            numFilters,
+        Activation     activation,
+        cublasHandle_t cublasHandle
+    );
+    ~Conv2d();
+
+    void forward(const float* d_input, float* d_output);
+
+  private:
+    // Inputs
+    int inputSize;
+    int inputChannels;
+
+    // Kernel
+    int kernelSize;
+    int stride;
+    int paddingSize;
+    int numFilters;
+
+    // Outputs
+    int outputSize;
+
+    // Kernels
+    std::vector<float> kernels;
+
+    // Cuda
+    cublasHandle_t cublasHandle;
+    float*         d_kernels;
+    float*         d_padded;
+
+    // Kernels
+    Activation activation;
+
+    void initializeKernels();
+    void toCuda();
+};
+
+}  // namespace Layers
+
+#endif  // CONV_LAYER_H
diff --git a/include/layers/dense.cuh b/include/layers/dense.cuh
index 12be07d..6a2417e 100644
--- a/include/layers/dense.cuh
+++ b/include/layers/dense.cuh
@@ -21,7 +21,7 @@ class Dense : public ILayer {
     );
     ~Dense();
 
-    void forward(const float* input, float* output);
+    void forward(const float* d_input, float* d_output);
     void setWeights(const std::vector<std::vector<float>>& weights);
     void setBiases(const std::vector<float>& biases);
 
diff --git a/src/layers/conv2d.cu b/src/layers/conv2d.cu
new file mode 100644
index 0000000..6951395
--- /dev/null
+++ b/src/layers/conv2d.cu
@@ -0,0 +1,86 @@
+#include <cublas_v2.h>
+
+#include <string>
+
+#include "activations.cuh"
+#include "conv2d.cuh"
+#include "cuda_helper.cuh"
+#include "padding.cuh"
+
+Layers::Conv2d::Conv2d(
+    int            inputSize,
+    int            inputChannels,
+    int            kernelSize,
+    int            stride,
+    std::string    padding,
+    int            numFilters,
+    Activation     activation,
+    cublasHandle_t cublasHandle
+)
+    : inputSize(inputSize),
+      inputChannels(inputChannels),
+      kernelSize(kernelSize),
+      stride(stride),
+      numFilters(numFilters),
+      cublasHandle(cublasHandle),
+      activation(activation) {
+    // Allocate memory for kernels
+
+    if (padding == "SAME") {
+        outputSize  = inputSize;
+        paddingSize = ((stride - 1) * inputSize - stride + kernelSize) / 2;
+    } else if (padding == "VALID") {
+        paddingSize = 0;
+        outputSize  = (inputSize - kernelSize) / stride + 1;
+    }
+
+    kernels.resize(kernelSize * kernelSize);
+    initializeKernels();
+
+    d_kernels = nullptr;
+
+    CUDA_CHECK(
+        cudaMalloc((void**)&d_kernels, sizeof(float) * kernelSize * kernelSize)
+    );
+    toCuda();
+
+    d_padded = nullptr;
+
+    if (paddingSize > 0) {
+        CUDA_CHECK(
+            cudaMalloc((void**)&d_padded,
+                       sizeof(float) * (inputSize + 2 * paddingSize) *
+                           (inputSize + 2 * paddingSize) * inputChannels)
+        );
+    }
+}
+
+Layers::Conv2d::~Conv2d() {
+    cudaFree(d_kernels);
+    cudaFree(d_padded);
+}
+
+void Layers::Conv2d::initializeKernels() {
+    std::fill(kernels.begin(), kernels.end(), 0.0f);
+}
+
+void Layers::Conv2d::toCuda() {
+    CUDA_CHECK(cudaMemcpy(
+        d_kernels, kernels.data(), sizeof(float) * kernelSize * kernelSize,
+        cudaMemcpyHostToDevice
+    ));
+}
+
+void Layers::Conv2d::forward(const float* d_input, float* d_output) {
+
+    // Padd input
+    int THREADS_PER_BLOCK = 256;
+    int BLOCKS            = (outputSize * outputSize * inputChannels) / THREADS_PER_BLOCK + 1;
+
+    pad_matrix_kernel<<<BLOCKS, THREADS_PER_BLOCK>>>(
+        d_input, d_padded, inputSize, inputSize, inputChannels, paddingSize
+    );
+
+    // TODO: Implement 2D convolution
+
+}
\ No newline at end of file