From c84f58b97cceab4628361887665bef1bf974e30e Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Wed, 15 May 2024 18:52:31 +0200
Subject: [PATCH] Implement modules

---
 include/model/model.hpp  |  1 +
 include/model/module.hpp | 42 ++++++++++++++++++++++++++++++++++
 src/model/model.cpp      | 12 ++++++++++
 src/model/module.cpp     | 49 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 104 insertions(+)
 create mode 100644 include/model/module.hpp
 create mode 100644 src/model/module.cpp

diff --git a/include/model/model.hpp b/include/model/model.hpp
index cdb1158..6f234e3 100644
--- a/include/model/model.hpp
+++ b/include/model/model.hpp
@@ -7,6 +7,7 @@
 
 #include "input.cuh"
 #include "layer.cuh"
+#include "module.hpp"
 #include "output.cuh"
 
 namespace CUDANet {
diff --git a/include/model/module.hpp b/include/model/module.hpp
new file mode 100644
index 0000000..9c1337f
--- /dev/null
+++ b/include/model/module.hpp
@@ -0,0 +1,42 @@
+#ifndef CUDANET_MODULE_H
+#define CUDANET_MODULE_H
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "layer.cuh"
+
+namespace CUDANet {
+
+class Module : public Layers::SequentialLayer {
+  public:
+    Module(const int inputSize, const int inputChannels, const int outputSize, const int outputChannels);
+    ~Module();
+
+    virtual float* forward(const float* d_input) = 0;
+
+    int getOutputSize();
+    int getInputSize();
+
+    void addLayer(const std::string& name, Layers::SequentialLayer* layer);
+    Layers::SequentialLayer* getLayer(const std::string& name);
+
+    const std::unordered_map<std::string, Layers::SequentialLayer*>& getLayers() const;
+
+  private:
+    int inputSize;
+    int inputChannels;
+
+    int outputSize;
+    int outputChannels;
+
+    float *d_output;
+
+    std::vector<std::pair<std::string, Layers::SequentialLayer*>> layers;
+    std::unordered_map<std::string, Layers::SequentialLayer*>     layerMap;
+};
+
+}  // namespace CUDANet
+
+#endif
\ No newline at end of file
diff --git a/src/model/model.cpp b/src/model/model.cpp
index e992f14..6571662 100644
--- a/src/model/model.cpp
+++ b/src/model/model.cpp
@@ -50,6 +50,18 @@ float* Model::predict(const float* input) {
 }
 
 void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
+
+    Module* module = dynamic_cast<Module*>(layer);
+
+    if (module != nullptr) {
+        layers.push_back({ name, module });
+        for (const auto& moduleLayer : module->getLayers()) {
+            layerMap[moduleLayer.first] = moduleLayer.second;
+        }
+
+        return;
+    }
+
     layers.push_back({ name, layer });
     layerMap[name] = layer;
 }
diff --git a/src/model/module.cpp b/src/model/module.cpp
new file mode 100644
index 0000000..30ad356
--- /dev/null
+++ b/src/model/module.cpp
@@ -0,0 +1,49 @@
+#include "module.hpp"
+
+#include "cuda_helper.cuh"
+
+using namespace CUDANet;
+
+Module::Module(
+    const int inputSize,
+    const int inputChannels,
+    const int outputSize,
+    const int outputChannels
+)
+    : inputSize(inputSize),
+      inputChannels(inputChannels),
+      outputSize(outputSize),
+      outputChannels(outputChannels),
+      layers(std::vector<std::pair<std::string, Layers::SequentialLayer*>>()),
+      layerMap(std::unordered_map<std::string, Layers::SequentialLayer*>()) {
+    d_output = nullptr;
+    CUDA_CHECK(cudaMalloc(
+        (void**)&d_output,
+        sizeof(float) * outputSize * outputSize * outputChannels
+    ));
+}
+
+Module::~Module() {
+    cudaFree(d_output);
+}
+
+void Module::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
+    layers.push_back({ name, layer });
+    layerMap[name] = layer;
+}
+
+Layers::SequentialLayer* Module::getLayer(const std::string& name) {
+    return layerMap[name];
+}
+
+const std::unordered_map<std::string, Layers::SequentialLayer*>& Module::getLayers() const {
+    return layerMap;
+}
+
+int Module::getInputSize() {
+    return inputSize * inputSize * inputChannels;
+}
+
+int Module::getOutputSize() {
+    return outputSize * outputSize * outputChannels;
+}
\ No newline at end of file