Use const T* for input tensors in layer and tensor operations

CMakeLists.txt

@@ -23,8 +23,8 @@ endif()
 
 
 file(GLOB_RECURSE CPU_SOURCES
+    src/*.cpp
     src/layers/*.cpp
-    src/model/*.cpp
 )
 
 set(LIBRARY_SOURCES ${CPU_SOURCES})
@@ -32,10 +32,7 @@ set(LIBRARY_SOURCES ${CPU_SOURCES})
 if(USE_CUDA)
     file(GLOB_RECURSE CUDA_SOURCES
         src/backends/cuda/*.cu
-        src/backends/cuda/utils/*.cu
         src/backends/cuda/kernels/*.cu
-        src/backends/cuda/layers/*.cu
-        src/layers/*.cu  # To be removed
     )
     set(LIBRARY_SOURCES ${LIBRARY_SOURCES} ${CUDA_SOURCES})
 endif()
@@ -46,17 +43,17 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 add_library(${PROJECT_NAME} STATIC ${LIBRARY_SOURCES})
 
 if(USE_CUDA)
+    # Enable relocatable device code for proper template instantiation across translation units
+    set_target_properties(${PROJECT_NAME} PROPERTIES
+        CUDA_SEPARABLE_COMPILATION ON
+        CUDA_RUNTIME_LIBRARY Shared
+    )
     target_link_libraries(${PROJECT_NAME} CUDA::cudart)
 endif()
 
 # Set include directories for the library
 target_include_directories(${PROJECT_NAME} PUBLIC
     ${CMAKE_CURRENT_SOURCE_DIR}/include
-    ${CMAKE_CURRENT_SOURCE_DIR}/include/utils
-    ${CMAKE_CURRENT_SOURCE_DIR}/include/kernels
-    ${CMAKE_CURRENT_SOURCE_DIR}/include/layers
-    ${CMAKE_CURRENT_SOURCE_DIR}/include/model
-    ${CMAKE_CURRENT_SOURCE_DIR}/src
 )
 
 set_property(TARGET ${PROJECT_NAME} PROPERTY CXX_STANDARD 20)

examples/inception_v3/inception_modules.cpp

@@ -213,7 +213,7 @@ InceptionB::InceptionB(
         branch3x3->getOutputSize(), branch3x3dbl_3->getOutputSize()
     );
     concat_2 = new CUDANet::Layers::Concat(
-        concat_1->getOutputSize(), branchPool->getOutputSize()
+        concat_1->getOutputSize(), branchPool->get_output_size()
     );
 
     outputSize = concat_2->getOutputSize();
@@ -441,7 +441,7 @@ InceptionD::InceptionD(
         branch3x3_2->getOutputSize(), branch7x7x3_4->getOutputSize()
     );
     concat_2 = new CUDANet::Layers::Concat(
-        concat_1->getOutputSize(), branchPool->getOutputSize()
+        concat_1->getOutputSize(), branchPool->get_output_size()
     );
 
     outputSize = concat_2->getOutputSize();
@@ -707,7 +707,7 @@ InceptionV3::InceptionV3(
     addLayer("AveragePool", avgpool);
 
     fc = new CUDANet::Layers::Dense(
-        avgpool->getOutputSize(), 1000, CUDANet::Layers::ActivationType::NONE
+        avgpool->get_output_size(), 1000, CUDANet::Layers::ActivationType::NONE
     );
     addLayer("fc", fc);
 }

include/backend.hpp

@@ -0,0 +1,126 @@
+#pragma once
+
+#include <memory>
+#include <optional>
+
+#include "shape.hpp"
+#include "tensor.hpp"
+
+namespace CUDANet {
+
+// Forward declarations
+class Backend;
+class Tensor;
+enum class DType;
+
+enum BackendType { CUDA_BACKEND, CPU_BACKEND };
+
+struct BackendConfig {
+    int device_id = 0;
+};
+
+class BackendFactory {
+  public:
+    static std::unique_ptr<Backend> create(BackendType backend_type, const BackendConfig& config);
+};
+
+class Backend {
+  protected:
+    std::optional<DType> default_dtype;
+  public:
+
+    // Dtypes
+    virtual bool supports_dtype(DType dtype) const = 0;
+    virtual void set_default_dtype(DType dtype) = 0;
+    virtual DType get_default_dtype() const = 0;
+
+    // Memory management
+    virtual void* allocate(size_t bytes) = 0;
+    virtual void  deallocate(void* ptr)  = 0;
+
+    // Tensor ops
+    virtual void print(const CUDANet::Tensor& input) = 0;
+    virtual void zero(CUDANet::Tensor& input)        = 0;
+    virtual void fill(CUDANet::Tensor& input, int data) = 0;
+
+    virtual void
+    copy_to_device(CUDANet::Tensor& tensor, void* data, size_t size) = 0;
+
+    virtual void sum(const CUDANet::Tensor& input, CUDANet::Tensor& sum) = 0;
+    virtual void max(const CUDANet::Tensor& input, CUDANet::Tensor& max) = 0;
+
+    // Layer ops
+    virtual void relu(CUDANet::Tensor& tensor)    = 0;
+    virtual void sigmoid(CUDANet::Tensor& tensor) = 0;
+    virtual void softmax(
+        CUDANet::Tensor& tensor,
+        CUDANet::Tensor& temp_max,
+        CUDANet::Tensor& temp_sum
+    ) = 0;
+
+    virtual CUDANet::Tensor& dense(
+        const CUDANet::Tensor& weights,
+        const CUDANet::Tensor& biases,
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor& output,
+        const size_t           input_size,
+        const size_t           output_size
+    ) = 0;
+
+    virtual CUDANet::Tensor& conv2d(
+        const CUDANet::Tensor& weights,
+        const CUDANet::Tensor& biases,
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor& output,
+        const CUDANet::Shape in_shape,
+        const CUDANet::Shape padding_shape,
+        const CUDANet::Shape kernel_shape,
+        const CUDANet::Shape stride_shape,
+        const CUDANet::Shape out_shape
+    ) = 0;
+
+    virtual CUDANet::Tensor& max_pool2d(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor& output,
+        CUDANet::Shape input_shape,
+        CUDANet::Shape pool_shape,
+        CUDANet::Shape stride_shape,
+        CUDANet::Shape padding_shape,
+        CUDANet::Shape output_shape
+    ) = 0;
+
+    virtual CUDANet::Tensor& avg_pool2d(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor& output,
+        CUDANet::Shape input_shape,
+        CUDANet::Shape pool_shape,
+        CUDANet::Shape stride_shape,
+        CUDANet::Shape padding_shape,
+        CUDANet::Shape output_shape
+    ) = 0;
+
+    virtual CUDANet::Tensor& batch_norm(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor& output,
+        CUDANet::Shape input_shape,
+        CUDANet::Tensor& weights,
+        CUDANet::Tensor& biases,
+        CUDANet::Tensor& running_mean,
+        CUDANet::Tensor& running_var,
+        CUDANet::Tensor& epsilon
+    ) = 0;
+
+    virtual CUDANet::Tensor& concat(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    ) = 0;
+
+    virtual CUDANet::Tensor& add(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    ) = 0;
+};
+
+}  // namespace CUDANet

include/backend/backend.hpp

@@ -1,29 +0,0 @@
-#pragma once
-
-#include <cstddef>
-#include "backend/tensor.hpp"
-
-namespace CUDANet::Backend
-{   
-
-class IBackend
-{
-public:
-
-    // Memory management
-    virtual void* allocate(size_t bytes) = 0;
-    virtual void deallocate(void* ptr) = 0;
-
-    // Tensor ops
-    virtual void print(const CUDANet::Backend::Tensor &input) = 0;
-    virtual void clear(CUDANet::Backend::Tensor &input) = 0;
-    virtual void sum(const CUDANet::Backend::Tensor &input, CUDANet::Backend::Tensor &sum) = 0;
-    virtual void max(const CUDANet::Backend::Tensor &input, CUDANet::Backend::Tensor &max) = 0;
-
-    // Layer ops
-    virtual void relu(CUDANet::Backend::Tensor &tensor) = 0;
-    virtual void sigmoid(CUDANet::Backend::Tensor &tensor) = 0;
-    virtual void softmax(CUDANet::Backend::Tensor &tensor, CUDANet::Backend::Tensor &temp_max, CUDANet::Backend::Tensor &temp_sum) = 0;
-};
-
-} // namespace CUDANet::Backend

include/backend/cpu/cpu.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "backend.hpp"
+#include "tensor.hpp"
+
+namespace CUDANet::Backends {
+
+class CPU : public Backend {
+public:
+    // Memory management
+    void* allocate(size_t bytes) override;
+    void deallocate(void* ptr) override;
+
+    // Tensor ops
+    void print(const CUDANet::Tensor &input) override;
+    void zero(CUDANet::Tensor &input) override;
+    void sum(const CUDANet::Tensor &input, CUDANet::Tensor &sum) override;
+    void max(const CUDANet::Tensor &input, CUDANet::Tensor &max) override;
+
+    // Layer ops
+    void relu(CUDANet::Tensor &tensor) override;
+    void sigmoid(CUDANet::Tensor &tensor) override;
+    void softmax(CUDANet::Tensor &tensor, CUDANet::Tensor &temp_max, CUDANet::Tensor &temp_sum) override;
+};
+
+}

include/backend/cuda.cuh

@@ -1,29 +0,0 @@
-#pragma once
-
-#include "backend/backend.hpp"
-#include "backend/tensor.hpp"
-
-namespace CUDANet::Backend {
-
-class CUDABackend : public IBackend {
-public:
-    // Memory management
-    void* allocate(size_t bytes) override;
-    void deallocate(void* ptr) override;
-
-    // Tensor ops
-    void print(const CUDANet::Backend::Tensor &input) override;
-    void clear(CUDANet::Backend::Tensor &input) override;
-    void sum(const CUDANet::Backend::Tensor &input, CUDANet::Backend::Tensor &sum) override;
-    void max(const CUDANet::Backend::Tensor &input, CUDANet::Backend::Tensor &max) override;
-
-    // Layer ops
-    void relu(CUDANet::Backend::Tensor &tensor) override;
-    void sigmoid(CUDANet::Backend::Tensor &tensor) override;
-    void softmax(CUDANet::Backend::Tensor &tensor, CUDANet::Backend::Tensor &temp_max, CUDANet::Backend::Tensor &temp_sum) override;
-
-private:
-    static constexpr int BLOCK_SIZE = 256;
-};
-
-}  // namespace CUDANet::Backend

include/backend/cuda/all.cuh

@@ -0,0 +1,11 @@
+#pragma once
+
+// CUDA Backend Implementation
+#include "backend/cuda/cuda.cuh"
+
+// CUDA Kernels
+#include "backend/cuda/kernels/activation_functions.cuh"
+#include "backend/cuda/kernels/convolution.cuh"
+#include "backend/cuda/kernels/matmul.cuh"
+#include "backend/cuda/kernels/pool.cuh"
+

include/backend/cuda/cuda.cuh

@@ -0,0 +1,243 @@
+#pragma once
+
+#include <cstdio>
+#include <set>
+
+#include "backend.hpp"
+#include "tensor.hpp"
+
+#ifndef BLOCK_SIZE
+#define BLOCK_SIZE 128
+#endif  // BLOCK_SIZE
+
+/**
+ * @brief CUDA error checking macro
+ *
+ */
+#define CUDA_CHECK(call)                                                       \
+    do {                                                                       \
+        cudaError_t result = call;                                             \
+        if (result != cudaSuccess) {                                           \
+            fprintf(                                                           \
+                stderr, "CUDA error at %s:%d code=%d(%s) \"%s\" \n", __FILE__, \
+                __LINE__, static_cast<unsigned int>(result),                   \
+                cudaGetErrorString(result), #call                              \
+            );                                                                 \
+            exit(EXIT_FAILURE);                                                \
+        }                                                                      \
+    } while (0)
+
+namespace CUDANet::Backends {
+
+template <DType dtype>
+struct cuda_dtype_map;
+
+template <>
+struct cuda_dtype_map<DType::FLOAT32> {
+    using type = float;
+};
+
+class CUDA : public Backend {
+  public:
+    CUDA(const BackendConfig& config);
+
+    bool  supports_dtype(DType dtype) const override;
+    void  set_default_dtype(DType dtype) override;
+    DType get_default_dtype() const override;
+
+    static bool is_cuda_available();
+    void        initialize();
+
+    // Memory management
+    void* allocate(size_t bytes) override;
+    void  deallocate(void* ptr) override;
+
+    // Tensor ops dispatchers
+    void print(const CUDANet::Tensor& input) override;
+    void zero(CUDANet::Tensor& input) override;
+    void fill(CUDANet::Tensor& input, int value) override;
+    void
+    copy_to_device(CUDANet::Tensor& tensor, void* data, size_t size) override;
+    void sum(const CUDANet::Tensor& input, CUDANet::Tensor& sum) override;
+    void max(const CUDANet::Tensor& input, CUDANet::Tensor& max) override;
+
+    // Layer ops dispatchers
+    void relu(CUDANet::Tensor& tensor) override;
+    void sigmoid(CUDANet::Tensor& tensor) override;
+    void softmax(
+        CUDANet::Tensor& tensor,
+        CUDANet::Tensor& temp_max,
+        CUDANet::Tensor& temp_sum
+    ) override;
+
+    CUDANet::Tensor& dense(
+        const CUDANet::Tensor& weights,
+        const CUDANet::Tensor& biases,
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        const size_t           input_size,
+        const size_t           output_size
+    ) override;
+
+    CUDANet::Tensor& conv2d(
+        const CUDANet::Tensor& weights,
+        const CUDANet::Tensor& biases,
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        const CUDANet::Shape   in_shape,
+        const CUDANet::Shape   padding_shape,
+        const CUDANet::Shape   kernel_shape,
+        const CUDANet::Shape   stride_shape,
+        const CUDANet::Shape   out_shape
+    ) override;
+
+    CUDANet::Tensor& max_pool2d(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        CUDANet::Shape         input_shape,
+        CUDANet::Shape         pool_shape,
+        CUDANet::Shape         stride_shape,
+        CUDANet::Shape         padding_shape,
+        CUDANet::Shape         output_shape
+    ) override;
+
+    CUDANet::Tensor& avg_pool2d(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        CUDANet::Shape         input_shape,
+        CUDANet::Shape         pool_shape,
+        CUDANet::Shape         stride_shape,
+        CUDANet::Shape         padding_shape,
+        CUDANet::Shape         output_shape
+    ) override;
+
+    CUDANet::Tensor& batch_norm(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        CUDANet::Shape         input_shape,
+        CUDANet::Tensor&       weights,
+        CUDANet::Tensor&       biases,
+        CUDANet::Tensor&       running_mean,
+        CUDANet::Tensor&       running_var,
+        CUDANet::Tensor&       epsilon
+    ) override;
+
+    CUDANet::Tensor& concat(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    ) override;
+
+    CUDANet::Tensor& add(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    ) override;
+
+  private:
+    int             device_id;
+    std::set<DType> supported_dtypes;
+
+    // Tensor ops template impls
+    template <typename T>
+    void print_impl(const CUDANet::Tensor& input);
+
+    template <typename T>
+    void fill_impl(CUDANet::Tensor& input, int value);
+
+    template <typename T>
+    void copy_to_device_impl(CUDANet::Tensor& tensor, void* data, size_t size);
+
+    template <typename T>
+    void sum_impl(const CUDANet::Tensor& input, CUDANet::Tensor& sum);
+
+    template <typename T>
+    void max_impl(const CUDANet::Tensor& input, CUDANet::Tensor& max);
+
+    // Layer ops template impls
+    template <typename T>
+    void relu_impl(CUDANet::Tensor& tensor);
+
+    template <typename T>
+    void sigmoid_impl(CUDANet::Tensor& tensor);
+
+    template <typename T>
+    void softmax_impl(
+        CUDANet::Tensor& tensor,
+        CUDANet::Tensor& temp_max,
+        CUDANet::Tensor& temp_sum
+    );
+
+    template <typename T>
+    CUDANet::Tensor& dense_impl(
+        const CUDANet::Tensor& weights,
+        const CUDANet::Tensor& biases,
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        const size_t           input_size,
+        const size_t           output_size
+    );
+
+    template <typename T>
+    CUDANet::Tensor& conv2d_impl(
+        const CUDANet::Tensor& weights,
+        const CUDANet::Tensor& biases,
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        const CUDANet::Shape   in_shape,
+        const CUDANet::Shape   padding_shape,
+        const CUDANet::Shape   kernel_shape,
+        const CUDANet::Shape   stride_shape,
+        const CUDANet::Shape   out_shape
+    );
+
+    template <typename T>
+    CUDANet::Tensor& max_pool2d_impl(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        CUDANet::Shape         input_shape,
+        CUDANet::Shape         pool_shape,
+        CUDANet::Shape         stride_shape,
+        CUDANet::Shape         padding_shape,
+        CUDANet::Shape         output_shape
+    );
+
+    template <typename T>
+    CUDANet::Tensor& avg_pool2d_impl(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        CUDANet::Shape         input_shape,
+        CUDANet::Shape         pool_shape,
+        CUDANet::Shape         stride_shape,
+        CUDANet::Shape         padding_shape,
+        CUDANet::Shape         output_shape
+    );
+
+    template <typename T>
+    CUDANet::Tensor& batch_norm_impl(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor&       output,
+        CUDANet::Shape         input_shape,
+        CUDANet::Tensor&       weights,
+        CUDANet::Tensor&       biases,
+        CUDANet::Tensor&       running_mean,
+        CUDANet::Tensor&       running_var,
+        CUDANet::Tensor&       epsilon
+    );
+
+    template <typename T>
+    CUDANet::Tensor& concat_impl(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    );
+
+    template <typename T>
+    CUDANet::Tensor& add_impl(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    );
+};
+
+}  // namespace CUDANet::Backends

include/backend/cuda/kernels/activation_functions.cuh

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace CUDANet::Kernels {
+
+
+template <typename T>
+__global__ void sigmoid(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const unsigned int len
+);
+
+template <typename T>
+__global__ void relu(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const unsigned int len
+);
+
+}  // namespace CUDANet::Kernels

include/backend/cuda/kernels/convolution.cuh

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include "layer.hpp"
+
+namespace CUDANet::Kernels {
+
+template <typename T>
+__global__ void convolution(
+    const T* __restrict__ d_input,
+    const T* __restrict__ d_kernel,
+    const T* __restrict__ d_bias,
+    T* __restrict__ d_output,
+    const Shape input_shape,
+    const Shape padding_shape,
+    const Shape kernel_shape,
+    const Shape stride_shape,
+    const Shape output_shape
+);
+
+}  // namespace CUDANet::Kernels

include/backend/cuda/kernels/matmul.cuh

@@ -0,0 +1,109 @@
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace CUDANet::Kernels {
+
+template <typename T>
+__global__ void mat_vec_mul(
+    const T* __restrict__ d_matrix,
+    const T* __restrict__ d_vector,
+    T* __restrict__ d_output,
+    const unsigned int w,
+    const unsigned int h
+);
+
+template <typename T>
+__global__ void vec_vec_add(
+    const T* __restrict__ d_vector1,
+    const T* __restrict__ d_vector2,
+    T* __restrict__ d_output,
+    const unsigned int w
+);
+
+template <typename T>
+__global__ void vec_vec_sub(
+    const T* __restrict__ d_vector1,
+    const T* __restrict__ d_vector2,
+    T* __restrict__ d_output,
+    const unsigned int w
+);
+
+template <typename T>
+__global__ void vec_vec_mul(
+    const T* __restrict__ d_vector1,
+    const T* __restrict__ d_vector2,
+    T* __restrict__ d_output,
+    const unsigned int w
+);
+
+template <typename T>
+__global__ void vec_scalar_sub(
+    const T* __restrict__ d_src,
+    T* __restrict__ d_out,
+    const T* __restrict__ d_scalar,
+    const unsigned int len
+);
+
+template <typename T>
+__global__ void vec_scalar_add(
+    const T* __restrict__ d_src,
+    T* __restrict__ d_out,
+    const T* __restrict__ d_scalar,
+    const unsigned int len
+);
+
+template <typename T>
+__global__ void vec_scalar_div(
+    const T* __restrict__ d_src,
+    T* __restrict__ d_out,
+    const T* __restrict__ d_scalar,
+    const unsigned int len
+);
+
+template <typename T>
+__global__ void vec_scalar_mul(
+    const T* __restrict__ d_src,
+    T* __restrict__ d_out,
+    const T* __restrict__ d_scalar,
+    const unsigned int len
+);
+
+template <typename T>
+__global__ void vec_exp(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const unsigned int len
+);
+
+template <typename T>
+__global__ void vec_sqrt(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const unsigned int len
+);
+
+template <typename T>
+__global__ void vec_scale(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const T* __restrict__ scale,
+    const T* epsilon,
+    const unsigned int len
+);
+
+template <typename T>
+__global__ void max_reduce(
+    const T* __restrict__ d_vector,
+    T* __restrict__ d_output,
+    const unsigned int len
+);
+
+template <typename T>
+__global__ void sum_reduce(
+    const T* __restrict__ d_vector,
+    T* __restrict__ d_output,
+    const unsigned int len
+);
+
+}  // namespace CUDANet::Kernels

include/backend/cuda/kernels/pool.cuh

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include "layer.hpp"
+
+namespace CUDANet::Kernels {
+
+template <typename T>
+__global__ void max_pool(
+    const T* __restrict__ d_input,
+    T* __restrict__ d_output,
+    const Shape input_shape,
+    const Shape output_shape,
+    const Shape pool_shape,
+    const Shape stride_shape,
+    const Shape padding_shape
+);
+
+template <typename T>
+__global__ void avg_pool(
+    const T* __restrict__ d_input,
+    T* __restrict__ d_output,
+    const Shape input_shape,
+    const Shape output_shape,
+    const Shape pool_shape,
+    const Shape stride_shape,
+    const Shape padding_shape
+);
+
+}  // namespace CUDANet::Kernels

include/backend/tensor.hpp

@@ -1,46 +0,0 @@
-#pragma once
-#include <cstddef>
-#include "backend/backend.hpp"
-#include <vector>
-
-namespace CUDANet::Backend
-{
-
-enum class DType
-{
-    FLOAT32,
-    // FLOAT16,  // Not implemented yet
-    // INT32,  // Not implemented yet
-};
-
-typedef std::vector<size_t> Shape;
-
-class Tensor
-{
-public:
-
-    Tensor() = default;
-    Tensor(Shape shape, DType dtype, IBackend* backend);
-    ~Tensor();
-
-    size_t size() const;
-    size_t numel() const;
-
-    template <typename T>
-    const T* data() const;
-
-    template <typename T>
-    T* data();
-
-private:
-    Shape       shape;
-    DType       dtype;
-
-    size_t total_elms;
-    size_t total_size;
-
-    IBackend*   backend;
-    void*       d_ptr;
-};
-
-} // namespace CUDANet::Backend

include/cudanet.cuh

@@ -1,35 +0,0 @@
-#ifndef CUDANET_H
-#define CUDANET_H
-
-#ifdef USE_CUDA
-#include "activation_functions.cuh"
-#include "convolution.cuh"
-#include "matmul.cuh"
-#include "pooling.cuh"
-#endif
-
-// Layers
-#include "activation.hpp"
-#include "add.hpp"
-#include "avg_pooling.hpp"
-#include "batch_norm.hpp"
-#include "concat.hpp"
-#include "conv2d.hpp"
-#include "dense.hpp"
-#include "input.hpp"
-#include "layer.hpp"
-#include "max_pooling.hpp"
-#include "output.hpp"
-
-// Models
-#include "model.hpp"
-#include "module.hpp"
-
-// Utils
-#include "imagenet.hpp"
-#ifdef USE_CUDA
-#include "cuda_helper.cuh"
-#include "vector.cuh"
-#endif
-
-#endif  // CUDANET_H

include/cudanet.hpp

@@ -0,0 +1,55 @@
+#pragma once
+
+// ============================================================================
+// Core Data Structures & Abstractions (BACKEND-INDEPENDENT)
+// ============================================================================
+
+#include "shape.hpp"
+#include "backend.hpp"
+#include "tensor.hpp"
+#include "layer.hpp"
+
+// ============================================================================
+// Container Classes
+// ============================================================================
+
+#include "module.hpp"
+#include "model.hpp"
+
+// ============================================================================
+// Layer Implementations
+// ============================================================================
+
+// Activation
+#include "layers/activation.hpp"
+
+// Normalization
+#include "layers/batch_norm.hpp"
+
+// Linear
+#include "layers/dense.hpp"
+
+// Convolutional
+#include "layers/conv2d.hpp"
+
+// Pooling
+#include "layers/max_pool.hpp"
+#include "layers/avg_pool.hpp"
+
+// Composition (element-wise operations)
+#include "layers/add.hpp"
+#include "layers/concat.hpp"
+
+// ============================================================================
+// Dataset Labels
+// ============================================================================
+
+#include "datasets/imagenet.hpp"
+
+// ============================================================================
+// Backend-Specific Includes (conditionally compiled)
+// ============================================================================
+
+#ifdef USE_CUDA
+#include "backend/cuda/all.cuh"
+#endif

include/datasets/imagenet.hpp

@@ -1,5 +1,4 @@
-#ifndef CUDANET_IMAGENET_H
-#define CUDANET_IMAGENET_H
+#pragma once
 
 #include <map>
 #include <string>
@@ -1012,5 +1011,3 @@ const std::map <int, std::string> IMAGENET_CLASS_MAP = {
 // clang-format on
 
 }
-
-#endif // CUDANET_IMAGENET_H

include/kernels/activation_functions.cuh

@@ -1,36 +0,0 @@
-#ifndef CUDANET_ACTIVATION_FUNCTIONS_H
-#define CUDANET_ACTIVATION_FUNCTIONS_H
-
-#include <cuda_runtime.h>
-
-namespace CUDANet::Kernels {
-
-/**
- * @brief Sigmoid activation function kernel
- *
- * @param src Pointer to the source array
- * @param dst Pointer to the destination array
- * @param len Length of the arrays
- */
-__global__ void sigmoid(
-    const float* __restrict__ src,
-    float* __restrict__ dst,
-    const unsigned int len
-);
-
-/**
- * @brief Relu activation function kernel
- *
- * @param src Pointer to the source array
- * @param dst Pointer to the destination array
- * @param len Length of the arrays
- */
-__global__ void relu(
-    const float* __restrict__ src,
-    float* __restrict__ dst,
-    const unsigned int len
-);
-
-}  // namespace CUDANet::Kernels
-
-#endif  // CUDANET_ACTIVATION_FUNCTIONS_H

include/kernels/convolution.cuh

@@ -1,39 +0,0 @@
-#ifndef CUDANET_CONVOLUTION_H
-#define CUDANET_CONVOLUTION_H
-
-#include <cuda_runtime.h>
-#include "layer.hpp"
-
-namespace CUDANet::Kernels {
-
-/**
- * @brief Convolution kernel
- *
- * @param d_input Device pointer to the input matrix
- * @param d_kernel Device pointer to the convolution kernel
- * @param d_bias Device pointer to the bias
- * @param d_output Device pointer to the output matrix
- * @param inputSize Width and height of the input matrix
- * @param nChannels Number of channels in the input matrix
- * @param kernelSize Width and height of the convolution kernel
- * @param stride Convolution stride
- * @param nFilters Number of output filters
- * @param outputSize Width and height of the output matrix
- */
-__global__ void convolution(
-    const float* __restrict__ d_input,
-    const float* __restrict__ d_kernel,
-    const float* __restrict__ d_bias,
-    float* __restrict__ d_output,
-    const shape2d inputSize,
-    const int nChannels,
-    const shape2d paddingSize,
-    const shape2d kernelSize,
-    const shape2d stride,
-    const int nFilters,
-    const shape2d outputSize
-);
-
-}  // namespace CUDANet::Kernels
-
-#endif  // CUDANET_CONVOLUTION_H

include/kernels/matmul.cuh

@@ -1,195 +0,0 @@
-#ifndef CUDANET_MATMUL_H
-#define CUDANET_MATMUL_H
-
-#include <cuda_runtime.h>
-
-namespace CUDANet::Kernels {
-
-/**
- * @brief Matrix vector multiplication kernel
- *
- * @param d_matrix Device pointer to matrix
- * @param d_vector Device pointer to vector
- * @param d_output Device pointer to output vector
- * @param w Width of the matrix
- * @param h Height of the matrix
- */
-__global__ void mat_vec_mul(
-    const float* __restrict__ d_matrix,
-    const float* __restrict__ d_vector,
-    float* __restrict__ d_output,
-    const unsigned int w,
-    const unsigned int h
-);
-
-/**
- * @brief Vector vector addition kernel
- *
- * @param d_vector1 Device pointer to first vector
- * @param d_vector2 Device pointer to second vector
- * @param d_output Device pointer to output vector
- * @param w Length of the vectors
- */
-__global__ void vec_vec_add(
-    const float* __restrict__ d_vector1,
-    const float* __restrict__ d_vector2,
-    float* __restrict__ d_output,
-    const unsigned int w
-);
-
-/**
- * @brief Vector vector subtraction kernel
- * 
- * @param d_vector1 
- * @param d_vector2 
- * @param d_output 
- * @param w 
- * @return __global__ 
- */
-__global__ void vec_vec_sub(
-    const float* __restrict__ d_vector1,
-    const float* __restrict__ d_vector2,
-    float* __restrict__ d_output,
-    const unsigned int w
-);
-
-__global__ void vec_vec_mul(
-    const float* __restrict__ d_vector1,
-    const float* __restrict__ d_vector2,
-    float* __restrict__ d_output,
-    const unsigned int w
-);
-
-/**
- * @brief Sub scalar from each element of the vector
- * 
- * @param d_vector 
- * @param d_scalar 
- * @param d_output 
- * @param w 
- * @return __global__ 
- */
-__global__ void vec_scalar_sub(
-    const float* __restrict__ d_src,
-    float* __restrict__ d_out,
-    const float* __restrict__ d_scalar,
-    const unsigned int len
-);
-
-/**
- * @brief Add scalar to each element of the vector
- * 
- * @param d_src 
- * @param d_out 
- * @param d_scalar 
- * @param len 
- * @return __global__ 
- */
-__global__ void vec_scalar_add(
-    const float* __restrict__ d_src,
-    float* __restrict__ d_out,
-    const float* __restrict__ d_scalar,
-    const unsigned int len
-);
-
-/**
- * @brief Divide each element of the vector by a scalar
- *
- * @param src Pointer to the source array
- * @param dst Pointer to the destination array
- * @param len Length of the arrays
- */
-__global__ void vec_scalar_div(
-    const float* __restrict__ d_src,
-    float* __restrict__ d_out,
-    const float* __restrict__ d_scalar,
-    const unsigned int len
-);
-
-/**
- * @brief Multiply each element of the vector by a scalar
- * 
- * @param d_src 
- * @param d_out 
- * @param d_scalar 
- * @param len 
- * @return __global__ 
- */
-__global__ void vec_scalar_mul(
-    const float* __restrict__ d_src,
-    float* __restrict__ d_out,
-    const float* __restrict__ d_scalar,
-    const unsigned int len
-);
-
-/**
- * @brief Exponentiate each element of the vector
- *
- * @param src Pointer to the source array
- * @param dst Pointer to the destination array
- * @param len Length of the arrays
- */
-__global__ void vec_exp(
-    const float* __restrict__ src,
-    float* __restrict__ dst,
-    const unsigned int len
-);
-
-/**
- * @brief Compute the square root of each element of the vector
- * 
- * @param src Device pointer to source vector
- * @param dst Device pointer to destination vector
- * @param len Length of the vector
- */
-__global__ void vec_sqrt(
-    const float* __restrict__ src,
-    float* __restrict__ dst,
-    const unsigned int len
-);
-
-/**
- * @brief Scales the vector by 1/sqrt(scale + epsilon)
- * 
- * @param src Device pointer to source vector
- * @param dst Device pointer to destination vector
- * @param scale Scale
- * @param epsilon Epsilon
- * @param len Length of the vector
- */
-__global__ void vec_scale(
-    const float* __restrict__ src,
-    float* __restrict__ dst,
-    const float* __restrict__ scale,
-    const float* epsilon,
-    const unsigned int len
-);
-
-/**
- * @brief Max reduction kernel
- *
- * @param d_vector Device pointer to vector
- * @param d_output Device pointer to output vector
- */
-__global__ void max_reduce(
-    const float* __restrict__ d_vector,
-    float* __restrict__ d_output,
-    const unsigned int len
-);
-
-/**
- * @brief 
- * 
- * @param d_vector Device pointer to vector
- * @param d_output Device pointer to output vector
- * @param len Length of the vector
- */
-__global__ void sum_reduce(
-    const float* __restrict__ d_vector,
-    float* __restrict__ d_output,
-    const unsigned int len
-);
-
-}  // namespace CUDANet::Kernels
-
-#endif  // CUDANET_MATMUL_H

include/kernels/pooling.cuh

@@ -1,33 +0,0 @@
-#ifndef CUDANET_POOLING_H
-#define CUDANET_POOLING_H
-
-#include <cuda_runtime.h>
-#include "layer.hpp"
-
-namespace CUDANet::Kernels {
-
-__global__ void max_pooling(
-    const float* __restrict__ d_input,
-    float* __restrict__ d_output,
-    const shape2d inputSize,
-    const shape2d outputSize,
-    const int nChannels,
-    const shape2d poolingSize,
-    const shape2d stride,
-    const shape2d padding
-);
-
-__global__ void avg_pooling(
-    const float* __restrict__ d_input,
-    float* __restrict__ d_output,
-    const shape2d inputSize,
-    const shape2d outputSize,
-    const int nChannels,
-    const shape2d poolingSize,
-    const shape2d stride,
-    const shape2d padding
-);
-
-}  // namespace CUDANet::Kernels
-
-#endif // CUDANET_POOLING_H

include/layer.hpp

@@ -0,0 +1,44 @@
+#pragma once
+
+#include <vector>
+
+#include "shape.hpp"
+#include "tensor.hpp"
+
+#define CUDANET_SAME_PADDING(inputSize, kernelSize, stride) \
+    ((stride - 1) * inputSize - stride + kernelSize) / 2;
+
+
+namespace CUDANet {
+
+/**
+ * @brief Basic Layer
+ *
+ */
+class Layer {
+  protected:
+    CUDANet::DType dtype;
+  public:
+
+    virtual ~Layer(){};
+
+    virtual CUDANet::Tensor& forward(CUDANet::Tensor &input) = 0;
+    
+    virtual CUDANet::Shape input_shape() = 0;
+
+    virtual CUDANet::Shape output_shape() = 0;
+
+    virtual size_t input_size() = 0;
+
+    virtual size_t output_size() = 0;
+
+    virtual void set_weights(void *input) = 0;
+
+    virtual size_t get_weights_size() = 0;
+
+    virtual void set_biases(void *input) = 0;
+
+    virtual size_t get_biases_size() = 0;
+};
+
+}  // namespace CUDANet

include/layers/activation.hpp

@@ -1,7 +1,8 @@
 #pragma once
 
-#include "backend/tensor.hpp"
-#include "backend/backend.hpp"
+#include "tensor.hpp"
+#include "backend.hpp"
+#include "layer.hpp"
 
 namespace CUDANet::Layers {
 
@@ -19,40 +20,42 @@ enum ActivationType { SIGMOID, RELU, SOFTMAX, NONE };
  * @brief Utility class that performs activation
  * 
  */
-class Activation {
+class Activation : public CUDANet::Layer {
   public:
 
     Activation() = default;
 
-    /**
-     * @brief Construct a new Activation object
-     * 
-     * @param activation Type of activation
-     * @param length     Length of the input
-     */
-    Activation(CUDANet::Backend::IBackend* backend, ActivationType activation, const int length);
+    Activation(ActivationType activation, const CUDANet::Shape &shape, CUDANet::Backend* backend);
+    Activation(ActivationType activation, const CUDANet::Shape &shape, CUDANet::DType dtype, CUDANet::Backend* backend);
 
-    /**
-     * @brief Destroy the Activation object
-     * 
-     */
-    ~Activation();
+    ~Activation() = default;
 
-    /**
-     * @brief Run the activation function on the input
-     * 
-     * @param d_input Pointer to the input vector on the device
-     */
-    void activate(CUDANet::Backend::Tensor input);
+    CUDANet::Tensor& forward(CUDANet::Tensor &input) override;
+    
+    CUDANet::Shape input_shape() override;
+
+    CUDANet::Shape output_shape() override;
+
+    size_t input_size() override;
+
+    size_t output_size() override;
+
+    void set_weights(void *input) override;
+
+    size_t get_weights_size() override;
+
+    void set_biases(void *input) override;
+
+    size_t get_biases_size() override;
 
 
   private:
-    CUDANet::Backend::IBackend* backend;
-    ActivationType activationType;
-    int length;
+    CUDANet::Backend* backend;
+    ActivationType activation_type;
+    CUDANet::Shape shape;
 
-    CUDANet::Backend::Tensor softmax_sum;
-    CUDANet::Backend::Tensor tensor_max;
+    CUDANet::Tensor softmax_sum;
+    CUDANet::Tensor tensor_max;
 };
 
 }  // namespace CUDANet::Layers

include/layers/add.hpp

@@ -1,49 +1,27 @@
-#ifndef CUDANET_ADD_LAYER_H
-#define CUDANET_ADD_LAYER_H
+#pragma once
+
+#include "shape.hpp"
+#include "tensor.hpp"
 
 namespace CUDANet::Layers {
 
 class Add {
   public:
-    /**
-     * @brief Create a new Add layer
-     *
-     * @param inputSize Size of the input arrays
-     */
-    Add(int inputSize);
+    Add(CUDANet::Shape a_shape, CUDANet::Shape b_shape, CUDANet::Backend* backend);
+    Add(CUDANet::Shape a_shape, CUDANet::Shape b_shape, CUDANet::DType dtype, CUDANet::Backend* backend);
 
-    /**
-     * @brief Destroy the Add layer
-     *
-     */
     ~Add();
 
-    /**
-     * @brief Adds first input to second input
-     *
-     * @param d_inputA Device pointer to the first input
-     * @param d_inputB Device pointer to the second input
-     *
-     */
-    float* forward(const float* inputA, const float* inputB);
+    CUDANet::Tensor&
+    forward(CUDANet::Tensor& input_a, CUDANet::Tensor& input_b);
 
   private:
-    int inputSize;
+    CUDANet::Shape out_shape;
+    CUDANet::Tensor output;
 
-    float* output;
+    CUDANet::Backend *backend;
 
-    float* forwardCPU(const float* inputA, const float* inputB);
-
-#ifdef USE_CUDA
-    float* d_output;
-    int gridSize;
-
-    float* forwardCUDA(const float* d_inputA, const float* d_inputB);
-    void initCUDA();
-    void delCUDA();
-#endif
+    CUDANet::DType dtype;
 };
 
 }  // namespace CUDANet::Layers
-
-#endif  // CUDANET_ADD_LAYER_H

include/layers/avg_pool.hpp

@@ -0,0 +1,64 @@
+#pragma once
+
+#include "layer.hpp"
+
+namespace CUDANet::Layers {
+
+class AvgPool2d : public CUDANet::Layer {
+  public:
+    AvgPool2d(
+        CUDANet::Shape input_shape,
+        CUDANet::Shape pool_shape,
+        CUDANet::Shape stride_shape,
+        CUDANet::Shape padding_shape,
+        CUDANet::Backend *backend
+    );
+    AvgPool2d(
+        CUDANet::Shape input_shape,
+        CUDANet::Shape pool_shape,
+        CUDANet::Shape stride_shape,
+        CUDANet::Shape padding_shape,
+        CUDANet::DType dtype,
+        CUDANet::Backend *backend
+    );
+
+    ~AvgPool2d();
+
+    CUDANet::Tensor& forward(CUDANet::Tensor& input) override;
+
+    CUDANet::Shape input_shape() override;
+
+    CUDANet::Shape output_shape() override;
+
+    size_t input_size() override;
+
+    size_t output_size() override;
+
+    void set_weights(void* input) override;
+
+    size_t get_weights_size() override;
+
+    void set_biases(void* input) override;
+
+    size_t get_biases_size() override;
+
+  protected:
+    CUDANet::Shape in_shape;
+
+    CUDANet::Shape pool_shape;
+    CUDANet::Shape stride_shape;
+    CUDANet::Shape padding_shape;
+
+    CUDANet::Shape  out_shape;
+    CUDANet::Tensor output;
+
+    CUDANet::Backend *backend;
+};
+
+class AdaptiveAvgPool2d : public AvgPool2d {
+  public:
+    AdaptiveAvgPool2d(CUDANet::Shape input_shape, CUDANet::Shape output_shape, CUDANet::Backend *backend);
+    AdaptiveAvgPool2d(CUDANet::Shape input_shape, CUDANet::Shape output_shape, CUDANet::DType dtype, CUDANet::Backend *backend);
+};
+
+}  // namespace CUDANet::Layers

include/layers/avg_pooling.hpp

@@ -1,78 +0,0 @@
-#ifndef CUDANET_AVG_POOLING_H
-#define CUDANET_AVG_POOLING_H
-
-#include "activation.hpp"
-#include "layer.hpp"
-
-namespace CUDANet::Layers {
-
-class AvgPooling2d : public SequentialLayer, public TwoDLayer {
-  public:
-    AvgPooling2d(
-        shape2d        inputSize,
-        int            nChannels,
-        shape2d        poolingSize,
-        shape2d        stride,
-        shape2d        padding,
-        ActivationType activationType
-    );
-    ~AvgPooling2d();
-
-    float* forward(const float* input);
-
-    /**
-     * @brief Get output size
-     *
-     * @return int output size
-     */
-    int getOutputSize();
-
-    /**
-     * @brief Get input size
-     *
-     * @return int input size
-     */
-    int getInputSize();
-
-    shape2d getOutputDims();
-
-  protected:
-    shape2d inputSize;
-    int     nChannels;
-    shape2d poolingSize;
-    shape2d stride;
-    shape2d padding;
-
-    shape2d outputSize;
-
-    Activation* activation;
-
-    float* forwardCPU(const float* input);
-
-#ifdef USE_CUDA
-    float* d_output;
-    float* forwardCUDA(const float* d_input);
-
-    void initCUDA();
-    void delCUDA();
-#endif
-};
-
-class AdaptiveAvgPooling2d : public AvgPooling2d {
-  public:
-    AdaptiveAvgPooling2d(
-        shape2d        inputShape,
-        int            nChannels,
-        shape2d        outputShape,
-        ActivationType activationType
-    );
-
-  private:
-#ifdef USE_CUDA
-    void initCUDA();
-#endif
-};
-
-}  // namespace CUDANet::Layers
-
-#endif  // CUDANET_AVG_POOLING_H

include/layers/batch_norm.hpp

@@ -1,170 +1,55 @@
-#ifndef CUDANET_BATCH_NORM_H
-#define CUDANET_BATCH_NORM_H
+#pragma once
 
-#include <vector>
-
-#include "activation.hpp"
 #include "layer.hpp"
 
 namespace CUDANet::Layers {
 
-class BatchNorm2d : public WeightedLayer, public TwoDLayer {
+class BatchNorm2d : public CUDANet::Layer {
   public:
-    BatchNorm2d(
-        shape2d        inputSize,
-        int            inputChannels,
-        float          epsilon,
-        ActivationType activationType
-    );
+    BatchNorm2d(CUDANet::Shape input_shape, float epsilon, CUDANet::Backend *backend);
+    BatchNorm2d(CUDANet::Shape input_shape, float epsilon, CUDANet::DType dtype, CUDANet::Backend *backend);
 
     ~BatchNorm2d();
 
-    /**
-     * @brief Compute the forward pass of the batchnorm layer
-     *
-     * @param d_input Device pointer to the input
-     * @return float* Device pointer to the output
-     */
-    float* forward(const float* d_input);
+    CUDANet::Tensor& forward(CUDANet::Tensor& input) override;
 
-    /**
-     * @brief Set the weights of the batchnorm layer
-     *
-     * @param weights_input Pointer to the weights
-     */
-    void setWeights(const float* weights_input);
+    CUDANet::Shape input_shape() override;
 
-    /**
-     * @brief Get the weights of the batchnorm layer
-     *
-     * @return std::vector<float>
-     */
-    std::vector<float> getWeights();
+    CUDANet::Shape output_shape() override;
 
-    /**
-     * @brief Set the biases of the batchnorm layer
-     *
-     * @param biases_input Pointer to the biases
-     */
-    void setBiases(const float* biases_input);
+    size_t input_size() override;
 
-    /**
-     * @brief Get the biases of the batchnorm layer
-     *
-     * @return std::vector<float>
-     */
-    std::vector<float> getBiases();
+    size_t output_size() override;
 
-    /**
-     * @brief Set the Running Mean
-     *
-     * @param running_mean_input
-     */
-    void setRunningMean(const float* running_mean_input);
+    void set_weights(void* input) override;
 
-    /**
-     * @brief Get the Running Mean
-     *
-     */
-    std::vector<float> getRunningMean();
+    size_t get_weights_size() override;
 
-    /**
-     * @brief Set the Running Var
-     *
-     * @param running_mean_input
-     */
-    void setRunningVar(const float* running_mean_input);
+    void set_biases(void* input) override;
 
-    /**
-     * @brief Get the Running Var
-     *
-     */
-    std::vector<float> getRunningVar();
+    size_t get_biases_size() override;
 
-    /**
-     * @brief Get output size
-     *
-     * @return int output size
-     */
-    int getOutputSize();
+    void set_running_mean(void* input);
 
-    /**
-     * @brief Get input size
-     *
-     * @return int input size
-     */
-    int getInputSize();
+    size_t get_running_mean_size();
 
-    shape2d getOutputDims();
+    void set_running_var(void* input);
+
+    size_t get_running_var_size();
 
   private:
-    shape2d inputSize;
-    int     inputChannels;
-    float   epsilon;
+    CUDANet::Shape  in_shape;
+    CUDANet::Tensor epsilon;
 
-    int gridSize;
+    CUDANet::Tensor running_mean;
+    CUDANet::Tensor running_var;
 
-#ifdef USE_CUDA
+    CUDANet::Tensor weights;
+    CUDANet::Tensor biases;
 
-    float* d_output;
+    CUDANet::Tensor output;
 
-    float* d_running_mean;
-    float* d_running_var;
-
-    float* d_length;
-    float* d_epsilon;
-
-    float* d_weights;
-    float* d_biases;
-
-    void initCUDA();
-    void delCUDA();
-
-    /**
-     * @brief Copy weights and biases to the device
-     *
-     */
-    void toCuda();
-
-    float* forwardCUDA(const float* d_input);
-
-#endif
-
-    std::vector<float> weights;
-    std::vector<float> biases;
-
-    std::vector<float> running_mean;
-    std::vector<float> running_var;
-
-    Activation* activation;
-
-    float* forwardCPU(const float* input);
-
-    /**
-     * @brief Initialize weights of the batchnorm layer with zeros
-     *
-     */
-    void initializeWeights();
-
-    /**
-     * @brief Initialize biases of the batchnorm layer with zeros
-     *
-     */
-    void initializeBiases();
-
-    /**
-     * @brief Initialize mean of the batchnorm layer with zeros
-     *
-     */
-    void initializeRunningMean();
-
-    /**
-     * @brief Initialize sqrt of variance of the batchnorm layer with ones
-     *
-     */
-    void initializeRunningVar();
+    CUDANet::Backend *backend;
 };
 
 }  // namespace CUDANet::Layers
-
-#endif  // CUDANET_BATCH_NORM_H

include/layers/concat.hpp

@@ -1,5 +1,4 @@
-#ifndef CUDANET_CONCAT_LAYER_H
-#define CUDANET_CONCAT_LAYER_H
+#pragma once
 
 #include "layer.hpp"
 
@@ -11,47 +10,27 @@ namespace CUDANet::Layers {
  */
 class Concat {
   public:
-    /**
-     * @brief Create a new Concat layer
-     *
-     * @param inputASize Size of the first input
-     * @param inputBSize Size of the second input
-     */
-    Concat(const int inputASize, const int inputBSize);
 
-    /**
-     * @brief Destroy the Concat layer
-     *
-     */
+    Concat(const CUDANet::Shape a_shape, const CUDANet::Shape b_shape, CUDANet::Backend *backend);
+    Concat(const CUDANet::Shape a_shape, const CUDANet::Shape b_shape, CUDANet::DType dtype, CUDANet::Backend *backend);
+
     ~Concat();
 
-    /**
-     * @brief Concatenates the two inputs
-     *
-     * @param d_input_A Device pointer to the first input
-     * @param d_input_B Device pointer to the second input
-     *
-     * @return Device pointer to the output
-     */
-    float* forward(const float* d_input_A, const float* d_input_B);
+    CUDANet::Tensor& forward(CUDANet::Tensor& input_a, CUDANet::Tensor& input_b);
 
-    int getOutputSize();
+    CUDANet::Shape output_shape();
 
   private:
-    int inputASize;
-    int inputBSize;
+    CUDANet::Shape a_shape;
+    CUDANet::Shape b_shape;
 
-    float* forwardCPU(const float* input_A, const float* input_B);
+    CUDANet::Shape out_shape;
+    CUDANet::Tensor output;
 
-#ifdef USE_CUDA
-    float* d_output;
-    float* forwardCUDA(const float* d_input_A, const float* d_input_B);
+    CUDANet::Backend *backend;
 
-    void initCUDA();
-    void delCUDA();
-#endif
+    CUDANet::DType dtype;
 };
 
 }  // namespace CUDANet::Layers
 
-#endif  // CUDANET_CONCAT_LAYER_H

include/layers/conv2d.hpp

@@ -1,9 +1,5 @@
-#ifndef CUDANET_CONV_LAYER_H
-#define CUDANET_CONV_LAYER_H
+#pragma once
 
-#include <vector>
-
-#include "activation.hpp"
 #include "layer.hpp"
 
 namespace CUDANet::Layers {
@@ -12,149 +8,60 @@ namespace CUDANet::Layers {
  * @brief 2D convolutional layer
  *
  */
-class Conv2d : public WeightedLayer, public TwoDLayer {
+class Conv2d : public CUDANet::Layer {
   public:
-    /**
-     * @brief Construct a new Conv 2d layer
-     *
-     * @param inputSize Width and height of the input matrix
-     * @param inputChannels Number of channels in the input matrix
-     * @param kernelSize Width and height of the convolution kernel
-     * @param stride Convolution stride
-     * @param numFilters Number of output filters
-     * @param paddingSize Padding size
-     * @param activationType Activation function type ('RELU', 'SIGMOID',
-     * 'SOFTMAX' or 'NONE')
-     */
     Conv2d(
-        shape2d        inputSize,
-        int            inputChannels,
-        shape2d        kernelSize,
-        shape2d        stride,
-        int            numFilters,
-        shape2d        paddingSize,
-        ActivationType activationType
+        CUDANet::Shape    input_shape,
+        CUDANet::Shape    kernel_shape,
+        CUDANet::Shape    stride_shape,
+        CUDANet::Shape    padding_shape,
+        CUDANet::Backend* backend
+    );
+    Conv2d(
+        CUDANet::Shape    input_shape,
+        CUDANet::Shape    kernel_shape,
+        CUDANet::Shape    stride_shape,
+        CUDANet::Shape    padding_shape,
+        CUDANet::DType    dtype,
+        CUDANet::Backend* backend
     );
 
-    /**
-     * @brief Destroy the Conv 2d object
-     *
-     */
     ~Conv2d();
 
-    /**
-     * @brief Forward pass of the convolutional layer
-     *
-     * @param d_input Device pointer to the input matrix
-     * @return Device pointer to the output matrix
-     */
-    float* forward(const float* d_input);
+    CUDANet::Tensor& forward(CUDANet::Tensor& input) override;
 
-    /**
-     * @brief Set the weights of the convolutional layer
-     *
-     * @param weights_input Pointer to the weights
-     */
-    void setWeights(const float* weights_input);
+    CUDANet::Shape input_shape() override;
 
-    /**
-     * @brief Get the weights of the convolutional layer
-     *
-     * @return std::vector<float>
-     */
-    std::vector<float> getWeights();
+    CUDANet::Shape output_shape() override;
 
-    /**
-     * @brief Set the biases of the convolutional layer
-     *
-     * @param biases_input Pointer to the biases
-     */
-    void setBiases(const float* biases_input);
+    size_t input_size() override;
 
-    /**
-     * @brief Get the biases of the convolutional layer
-     *
-     * @return std::vector<float>
-     */
-    std::vector<float> getBiases();
+    size_t output_size();
 
-    /**
-     * @brief Get output size
-     *
-     * @return int output size
-     */
-    int getOutputSize();
+    void set_weights(void* input) override;
 
-    /**
-     * @brief Get input size
-     *
-     * @return int input size
-     */
-    int getInputSize();
+    size_t get_weights_size() override;
 
-    /**
-     * @brief Get the padding size of the layer
-     *
-     * @return int
-     */
-    shape2d getPaddingSize() {
-        return paddingSize;
-    }
+    void set_biases(void* input) override;
 
-    shape2d getOutputDims();
+    size_t get_biases_size() override;
+
+    CUDANet::Shape get_padding_shape();
 
   private:
-    // Inputs
-    shape2d inputSize;
-    int     inputChannels;
+    CUDANet::Backend* backend;
 
-    // Outputs
-    shape2d outputSize;
+    CUDANet::Shape in_shape;
+    CUDANet::Shape out_shape;
 
-    // Kernel
-    shape2d kernelSize;
-    shape2d stride;
-    shape2d paddingSize;
-    int     numFilters;
+    CUDANet::Shape kernel_shape;
+    CUDANet::Shape stride_shape;
+    CUDANet::Shape padding_shape;
 
-    // Kernels
-    std::vector<float> weights;
-    std::vector<float> biases;
+    CUDANet::Tensor weights;
+    CUDANet::Tensor biases;
 
-    float* forwardCPU(const float* input);
-
-// Cuda
-#ifdef USE_CUDA
-    float* d_output;
-    float* d_weights;
-    float* d_biases;
-
-    float* forwardCUDA(const float* d_input);
-    void   initCUDA();
-    void   delCUDA();
-
-    /**
-     * @brief Copy weights and biases to the device
-     *
-     */
-    void toCuda();
-#endif
-
-    Activation* activation;
-
-    /**
-     * @brief Initialize weights of the convolutional layer with zeros
-     *
-     */
-    void initializeWeights();
-
-    /**
-     * @brief Initialize biases of the convolutional layer with zeros
-     *
-     */
-    void initializeBiases();
+    CUDANet::Tensor output;
 };
 
 }  // namespace CUDANet::Layers
-
-#endif  // CUDANET_CONV_LAYER_H

include/layers/dense.hpp

@@ -1,9 +1,6 @@
-#ifndef CUDANET_DENSE_LAYER_H
-#define CUDANET_DENSE_LAYER_H
+#pragma once
 
-#include <vector>
-
-#include "activation.hpp"
+#include "backend.hpp"
 #include "layer.hpp"
 
 namespace CUDANet::Layers {
@@ -12,121 +9,43 @@ namespace CUDANet::Layers {
  * @brief Dense (fully connected) layer
  *
  */
-class Dense : public WeightedLayer {
+class Dense : public CUDANet::Layer {
   public:
-    /**
-     * @brief Construct a new Dense layer
-     *
-     * @param inputSize Size of the input vector
-     * @param outputSize Size of the output vector
-     * @param activationType Activation function type ('RELU', 'SIGMOID',
-     * 'SOFTMAX' or 'NONE')
-     */
-    Dense(int inputSize, int outputSize, Layers::ActivationType activationType);
 
-    /**
-     * @brief Destroy the Dense layer
-     *
-     */
+    Dense(CUDANet::Shape input_shape, CUDANet::Shape output_shape, CUDANet::Backend *backend);
+    Dense(CUDANet::Shape input_shape, CUDANet::Shape output_shape, CUDANet::DType dtype, CUDANet::Backend *backend);
+
     ~Dense();
 
-    /**
-     * @brief Forward pass of the dense layer
-     *
-     * @param d_input Device pointer to the input vector
-     * @return Device pointer to the output vector
-     */
-    float* forward(const float* d_input);
+    CUDANet::Tensor& forward(CUDANet::Tensor &input) override;
 
-    /**
-     * @brief Set the weights of the layer
-     *
-     * @param weights Pointer to vector of weights
-     */
-    void setWeights(const float* weights);
+    CUDANet::Shape input_shape() override;
 
-    /**
-     * @brief Get the weights of the layer
-     *
-     * @return Vector of weights
-     */
-    std::vector<float> getWeights();
+    CUDANet::Shape output_shape() override;
 
-    /**
-     * @brief Set the biases of the layer
-     *
-     * @param biases Pointer to vector of biases
-     */
-    void setBiases(const float* biases);
+    size_t input_size() override;
 
-    /**
-     * @brief Get the biases of the layer
-     *
-     * @return Vector of biases
-     */
-    std::vector<float> getBiases();
+    size_t output_size() override;
 
-    /**
-     * @brief Get output size
-     *
-     * @return int output size
-     */
-    int getOutputSize();
+    void set_weights(void *input) override;
 
-    /**
-     * @brief Get input size
-     *
-     * @return int input size
-     */
-    int getInputSize();
+    size_t get_weights_size() override;
+
+    void set_biases(void *input) override;
+
+    size_t get_biases_size() override;
 
   private:
-    int inputSize;
-    int outputSize;
+    CUDANet::Backend *backend;
 
-    std::vector<float> weights;
-    std::vector<float> biases;
+    CUDANet::Shape in_shape;
+    CUDANet::Shape out_shape;
 
-    Layers::Activation* activation;
-
-    /**
-     * @brief Initialize the weights to zeros
-     *
-     */
-    void initializeWeights();
-
-    /**
-     * @brief Initialize the biases to zeros
-     *
-     */
-    void initializeBiases();
-
-    float* forwardCPU(const float* input);
-
-#ifdef USE_CUDA
-    float* d_output;
-
-    float* d_weights;
-    float* d_biases;
-
-    // Precompute kernel launch parameters
-    int forwardGridSize;
-    int biasGridSize;
-
-    /**
-     * @brief Copy the weights and biases to the device
-     *
-     */
-    void toCuda();
-
-    void initCUDA();
-    void delCUDA();
-
-    float* forwardCUDA(const float* d_input);
-#endif
+    CUDANet::Tensor weights;
+    CUDANet::Tensor biases;
 
+    CUDANet::Tensor output;
 };
 
 }  // namespace CUDANet::Layers
 
-#endif  // CUDANET_DENSE_LAYER_H

include/layers/input.hpp

@@ -1,66 +0,0 @@
-#ifndef CUDANET_INPUT_LAYER_H
-#define CUDANET_INPUT_LAYER_H
-
-#include "layer.hpp"
-
-namespace CUDANet::Layers {
-
-/**
- * @brief Input layer, just copies the input to the device
- *
- */
-class Input : public SequentialLayer {
-  public:
-    /**
-     * @brief Create a new Input layer
-     *
-     * @param inputSize Size of the input vector
-     */
-    explicit Input(int inputSize);
-
-    /**
-     * @brief Destroy the Input layer
-     *
-     */
-    ~Input();
-
-    /**
-     * @brief Forward pass of the input layer. Just copies the input to the
-     * device
-     *
-     * @param input Host pointer to the input vector
-     * @return Device pointer to the output vector
-     */
-    float* forward(const float* input);
-
-    /**
-     * @brief Get output size
-     *
-     * @return int output size
-     */
-    int getOutputSize();
-
-    /**
-     * @brief Get input size
-     *
-     * @return int input size
-     */
-    int getInputSize();
-
-  private:
-    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
-
-#endif  // CUDANET_INPUT_LAYER_H

include/layers/layer.hpp

@@ -1,124 +0,0 @@
-
-#ifndef CUDANET_I_LAYER_H
-#define CUDANET_I_LAYER_H
-
-#include <vector>
-
-#define CUDANET_SAME_PADDING(inputSize, kernelSize, stride) \
-    ((stride - 1) * inputSize - stride + kernelSize) / 2;
-
-typedef std::pair<int, int> shape2d;
-
-namespace CUDANet::Layers {
-
-
-class TwoDLayer {
-
-  public:
-    virtual shape2d getOutputDims() = 0;
-
-};
-
-
-/**
- * @brief Basic Sequential Layer
- *
- */
-class SequentialLayer {
-  public:
-    /**
-     * @brief Destroy the Sequential Layer
-     *
-     */
-    virtual ~SequentialLayer(){};
-
-    /**
-     * @brief Forward propagation virtual function
-     *
-     * @param input Device pointer to the input
-     * @return float* Device pointer to the output
-     */
-    virtual float* forward(const float* input) = 0;
-
-    /**
-     * @brief Get output size
-     *
-     * @return int output size
-     */
-    virtual int getOutputSize() = 0;
-
-    /**
-     * @brief Get input size
-     * 
-     * @return int input size
-     */
-    virtual int getInputSize() = 0;
-};
-
-/**
- * @brief Base class for layers with weights and biases
- */
-class WeightedLayer : public SequentialLayer {
-  public:
-    /**
-     * @brief Destroy the ILayer object
-     *
-     */
-    virtual ~WeightedLayer(){};
-
-    /**
-     * @brief Virtual function for forward pass
-     *
-     * @param input (Device) Pointer to the input
-     * @return float* Device pointer to the output
-     */
-    virtual float* forward(const float* input) = 0;
-
-    /**
-     * @brief Virtual function for setting weights
-     *
-     * @param weights Pointer to the weights
-     */
-    virtual void setWeights(const float* weights) = 0;
-
-    /**
-     * @brief Virtual function for getting weights
-     *
-     */
-    virtual std::vector<float> getWeights() = 0;
-
-    /**
-     * @brief Virtual function for setting biases
-     *
-     * @param biases Pointer to the biases
-     */
-    virtual void setBiases(const float* biases) = 0;
-
-    /**
-     * @brief Virtual function for getting biases
-     *
-     */
-    virtual std::vector<float> getBiases() = 0;
-
-  private:
-    /**
-     * @brief Initialize the weights
-     */
-    virtual void initializeWeights() = 0;
-
-    /**
-     * @brief Initialize the biases
-     */
-    virtual void initializeBiases() = 0;
-
-#ifdef USE_CUDA
-    /**
-     * @brief Copy the weights and biases to the device
-     */
-    virtual void toCuda() = 0;
-#endif
-};
-
-}  // namespace CUDANet::Layers
-
-#endif  // CUDANET_I_LAYERH

include/layers/max_pool.hpp

@@ -0,0 +1,59 @@
+#pragma once
+
+#include "layer.hpp"
+
+namespace CUDANet::Layers {
+
+class MaxPool2d : public CUDANet::Layer {
+  public:
+    MaxPool2d(
+        CUDANet::Shape        input_shape,
+        CUDANet::Shape        pool_shape,
+        CUDANet::Shape        stride_shape,
+        CUDANet::Shape        padding_shape,
+        CUDANet::Backend* backend
+    );
+    MaxPool2d(
+        CUDANet::Shape        input_shape,
+        CUDANet::Shape        pool_shape,
+        CUDANet::Shape        stride_shape,
+        CUDANet::Shape        padding_shape,
+        CUDANet::DType        dtype,
+        CUDANet::Backend* backend
+    );
+    ~MaxPool2d();
+
+    CUDANet::Tensor& forward(CUDANet::Tensor &input) override;
+    
+    CUDANet::Shape input_shape() override;
+
+    CUDANet::Shape output_shape() override;
+
+    size_t input_size() override;
+
+    size_t output_size() override;
+
+    void set_weights(void *input) override;
+
+    size_t get_weights_size() override;
+
+    void set_biases(void *input) override;
+
+    size_t get_biases_size() override;
+
+
+
+  private:
+    CUDANet::Shape in_shape;
+
+    CUDANet::Shape pool_shape;
+    CUDANet::Shape stride_shape;
+    CUDANet::Shape padding_shape;
+
+    CUDANet::Shape out_shape;
+    CUDANet::Tensor output;
+
+    CUDANet::Backend *backend;
+};
+
+}  // namespace CUDANet::Layers

include/layers/max_pooling.hpp

@@ -1,63 +0,0 @@
-#ifndef CUDANET_MAX_POOLING_H
-#define CUDANET_MAX_POOLING_H
-
-#include "activation.hpp"
-#include "layer.hpp"
-
-namespace CUDANet::Layers {
-
-class MaxPooling2d : public SequentialLayer, public TwoDLayer {
-  public:
-    MaxPooling2d(
-        shape2d        inputSize,
-        int            nChannels,
-        shape2d        poolingSize,
-        shape2d        stride,
-        shape2d        padding,
-        ActivationType activationType
-    );
-    ~MaxPooling2d();
-
-    float* forward(const float* input);
-
-    /**
-     * @brief Get output size
-     *
-     * @return int output size
-     */
-    int getOutputSize();
-
-    /**
-     * @brief Get input size
-     *
-     * @return int input size
-     */
-    int getInputSize();
-
-    shape2d getOutputDims();
-
-  private:
-    shape2d inputSize;
-    int     nChannels;
-    shape2d poolingSize;
-    shape2d stride;
-    shape2d padding;
-
-    shape2d outputSize;
-
-    Activation* activation;
-
-    float* forwardCPU(const float* input);
-
-#ifdef USE_CUDA
-    float* d_output;
-    float* forwardCUDA(const float* d_input);
-
-    void initCUDA();
-    void delCUDA();
-#endif
-};
-
-}  // namespace CUDANet::Layers
-
-#endif  // CUDANET_MAX_POOLING_H

include/layers/output.hpp

@@ -1,59 +0,0 @@
-#ifndef CUDANET_OUTPUT_LAYER_H
-#define CUDANET_OUTPUT_LAYER_H
-
-#include "layer.hpp"
-
-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);
-
-    /**
-     * @brief Get output size
-     *
-     * @return int output size
-     */
-    int getOutputSize();
-
-    /**
-     * @brief Get input size
-     *
-     * @return int input size
-     */
-    int getInputSize();
-
-  private:
-    int    inputSize;
-    float* h_output;
-
-    float* forwardCPU(const float* input);
-
-#ifdef USE_CUDA
-    float* forwardCUDA(const float* input);
-#endif
-};
-
-}  // namespace CUDANet::Layers
-
-#endif  // CUDANET_OUTPUT_LAYER_H

include/model.hpp

@@ -0,0 +1,55 @@
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "layer.hpp"
+#include "module.hpp"
+
+namespace CUDANet {
+
+enum TensorType {
+    WEIGHT,
+    BIAS,
+    RUNNING_MEAN,
+    RUNNING_VAR
+};
+
+struct TensorInfo {
+    std::string name;
+    TensorType  type;
+    int         size;
+    int         offset;
+};
+
+class Model {
+  public:
+    Model(const CUDANet::Shape input_shape, const CUDANet::Shape output_shape);
+    ~Model();
+
+    virtual CUDANet::Tensor& predict(CUDANet::Tensor& input);
+
+    CUDANet::Layer* get_layer(const std::string& name);
+
+    void register_layer(const std::string& name, Layer* layer);
+
+    void register_module(Module& module);
+
+    void load_weights(const std::string& path);
+
+    bool validate();
+
+    void print_summary();
+
+  protected:
+    CUDANet::Shape in_shape;
+    CUDANet::Shape out_shape;
+
+    CUDANet::Tensor output;
+
+    std::vector<std::pair<std::string, Layer*>> layers;
+    std::unordered_map<std::string, Layer*>     layer_map;
+};
+
+}  // namespace CUDANet

include/model/model.hpp

@@ -1,61 +0,0 @@
-#ifndef CUDANET_MODEL_H
-#define CUDANET_MODEL_H
-
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-#include "input.hpp"
-#include "layer.hpp"
-#include "module.hpp"
-#include "output.hpp"
-
-namespace CUDANet {
-
-enum TensorType {
-    WEIGHT,
-    BIAS,
-    RUNNING_MEAN,
-    RUNNING_VAR
-};
-
-struct TensorInfo {
-    std::string name;
-    TensorType  type;
-    int         size;
-    int         offset;
-};
-
-class Model {
-  public:
-    Model(const shape2d inputSize, const int inputChannels, const int outputSize);
-    Model(const Model& other);
-    ~Model();
-
-    virtual float* predict(const float* input);
-
-    void addLayer(const std::string& name, Layers::SequentialLayer* layer);
-    Layers::SequentialLayer* getLayer(const std::string& name);
-
-    void loadWeights(const std::string& path);
-
-    bool validate();
-
-    void printSummary();
-
-  protected:
-    Layers::Input*  inputLayer;
-    Layers::Output* outputLayer;
-
-    shape2d inputSize;
-    int inputChannels;
-
-    int outputSize;
-
-    std::vector<std::pair<std::string, Layers::SequentialLayer*>> layers;
-    std::unordered_map<std::string, Layers::SequentialLayer*>     layerMap;
-};
-
-}  // namespace CUDANet
-
-#endif  // CUDANET_MODEL_H

include/model/module.hpp

@@ -1,32 +0,0 @@
-#ifndef CUDANET_MODULE_H
-#define CUDANET_MODULE_H
-
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-#include "layer.hpp"
-
-namespace CUDANet {
-
-class Module : public Layers::SequentialLayer {
-  public:
-    virtual float* forward(const float* d_input) = 0;
-
-    int getOutputSize();
-    int getInputSize();
-
-    void addLayer(const std::string& name, Layers::SequentialLayer* layer);
-
-    const std::vector<std::pair<std::string, Layers::SequentialLayer*>>& getLayers() const;
-
-  protected:
-    std::vector<std::pair<std::string, Layers::SequentialLayer*>> layers;
-
-    int outputSize;
-    int inputSize;
-};
-
-}  // namespace CUDANet
-
-#endif

include/module.hpp

@@ -0,0 +1,31 @@
+#pragma once
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "layer.hpp"
+
+namespace CUDANet {
+
+class Module {
+  public:
+    CUDANet::Shape input_shape();
+
+    CUDANet::Shape output_shape();
+
+    void register_layer(const std::string& name, Layer* layer);
+
+    void register_module(Module& module);
+
+    const std::vector<std::pair<std::string, Layer*>>& get_layers() const;
+
+  protected:
+    std::vector<std::pair<std::string, Layer*>> layers;
+
+    CUDANet::Shape in_shape;
+    CUDANet::Shape out_shape;
+};
+
+}  // namespace CUDANet

include/shape.hpp

@@ -0,0 +1,117 @@
+#pragma once
+
+#ifndef __host__
+#define __host__
+#endif
+
+#ifndef __device__
+#define __device__
+#endif
+
+#include <format>
+#include <stdexcept>
+#include <vector>
+
+namespace CUDANet {
+
+struct Shape {
+    static constexpr size_t MAX_DIMS = 8;
+    
+    size_t dims[MAX_DIMS];
+    size_t ndim;
+    
+    __host__ __device__ Shape() : ndim(0) {
+        for (int i = 0; i < MAX_DIMS; i++) dims[i] = 0;
+    }
+    
+    __host__ Shape(std::initializer_list<size_t> list) : ndim(list.size()) {
+        if (ndim > MAX_DIMS) {
+            throw std::runtime_error("Too many dimensions");
+        }
+        size_t i = 0;
+        for (auto val : list) {
+            dims[i++] = val;
+        }
+        for (; i < MAX_DIMS; i++) dims[i] = 0;
+    }
+    
+    __host__ Shape(const std::vector<size_t>& vec) : ndim(vec.size()) {
+        if (ndim > MAX_DIMS) {
+            throw std::runtime_error("Too many dimensions");
+        }
+        for (size_t i = 0; i < ndim; i++) {
+            dims[i] = vec[i];
+        }
+        for (size_t i = ndim; i < MAX_DIMS; i++) dims[i] = 0;
+    }
+    
+    __host__ __device__ size_t operator[](size_t idx) const {
+        return dims[idx];
+    }
+    
+    __host__ __device__ size_t& operator[](size_t idx) {
+        return dims[idx];
+    }
+    
+    __host__ __device__ size_t size() const { return ndim; }
+    
+    __host__ bool operator==(const Shape& other) const {
+        if (ndim != other.ndim) return false;
+        for (size_t i = 0; i < ndim; i++) {
+            if (dims[i] != other.dims[i]) return false;
+        }
+        return true;
+    }
+    
+    __host__ bool operator!=(const Shape& other) const {
+        return !(*this == other);
+    }
+
+    __host__ __device__ bool empty() const {
+        return ndim == 0;
+    }
+
+
+};
+
+std::string format_shape(const Shape& shape) {
+    std::string result;
+    for (size_t i = 0; i < shape.size(); ++i) {
+        if (i > 0) result += ", ";
+        result += std::to_string(shape[i]);
+    }
+    return result;
+}
+
+class InvalidShapeException : public std::runtime_error {
+  public:
+    InvalidShapeException(
+        const std::string& param_name,
+        size_t             expected,
+        size_t             actual
+    )
+        : std::runtime_error(
+              std::format(
+                  "Invalid {} shape. Expected {}, actual {}",
+                  param_name,
+                  expected,
+                  actual
+              )
+          ) {}
+
+    InvalidShapeException(
+        const std::string& message,
+        const Shape&       shape_a,
+        const Shape&       shape_b
+    )
+        : std::runtime_error(
+              std::format(
+                  "{}. Shape A: [{}], Shape B: [{}]",
+                  message,
+                  format_shape(shape_a),
+                  format_shape(shape_b)
+              )
+          ) {}
+};
+
+}  // namespace CUDANet

include/tensor.hpp

@@ -0,0 +1,64 @@
+#pragma once
+
+#include <cstddef>
+#include <vector>
+
+#include "backend.hpp"
+#include "shape.hpp"
+
+namespace CUDANet
+{
+
+enum class DType
+{
+    FLOAT32,
+    // FLOAT16,  // Not implemented yet
+    // INT32,  // Not implemented yet
+};
+
+size_t dtype_size(DType dtype);
+
+// Forward declaration
+class Backend;
+
+class Tensor
+{
+public:
+
+    Tensor() = default;
+    Tensor(Shape shape, CUDANet::Backend* backend);
+    Tensor(Shape shape, DType dtype, CUDANet::Backend* backend);
+
+    Tensor(Tensor&& other) noexcept;
+    Tensor& operator=(Tensor&& other) noexcept;
+    Tensor(const Tensor&) = delete;
+    Tensor& operator=(const Tensor&) = delete;
+
+    ~Tensor();
+
+    DType get_dtype() const;
+
+    size_t size() const;
+    size_t numel() const;
+
+    void* device_ptr() const;
+    void* device_ptr();
+
+    void zero();
+
+    void fill(int value);
+
+    void set_data(void *data);
+
+private:
+    Shape       shape;
+    DType       dtype;
+
+    size_t total_elms;
+    size_t total_size;
+
+    CUDANet::Backend*   backend;
+    void*       d_ptr;
+};
+
+} // namespace CUDANet

include/utils/cuda_helper.cuh

@@ -1,26 +0,0 @@
-#ifndef CUDANET_HELPER_H
-#define CUDANET_HELPER_H
-
-#include <cuda_runtime.h>
-#include <cstdio>
-
-#ifndef BLOCK_SIZE
-#define BLOCK_SIZE 128
-#endif // BLOCK_SIZE
-
-/**
- * @brief CUDA error checking macro
- * 
- */
-#define CUDA_CHECK(call) \
-do { \
-    cudaError_t result = call; \
-    if (result != cudaSuccess) { \
-        fprintf(stderr, "CUDA error at %s:%d code=%d(%s) \"%s\" \n", \
-                __FILE__, __LINE__, static_cast<unsigned int>(result), \
-                cudaGetErrorString(result), #call); \
-        exit(EXIT_FAILURE); \
-    } \
-} while (0)
-
-#endif // CUDANET_HELPER_H

src/backend_factory.cpp

@@ -0,0 +1,40 @@
+#include <stdexcept>
+#include <memory>
+
+#ifdef USE_CUDA
+#include "backend/cuda/cuda.cuh"
+#endif
+
+#include "backend.hpp"
+
+namespace CUDANet {    
+    
+std::unique_ptr<Backend> BackendFactory::create(BackendType backend_type, const BackendConfig& config) {
+    switch (backend_type)
+    {
+    case BackendType::CUDA_BACKEND:
+        {
+        #ifdef USE_CUDA
+
+        if (!CUDANet::Backends::CUDA::is_cuda_available()) {
+            throw std::runtime_error("No CUDA devices found");
+        }
+
+        auto cuda = std::make_unique<CUDANet::Backends::CUDA>(config);
+        return cuda;
+
+        #else
+        throw std::runtime_error("Library was compiled without CUDA support.");
+        #endif
+        }
+        break;
+    
+    default:
+        throw std::runtime_error("Invalid backend");
+        break;
+    }
+
+    return nullptr;
+}
+
+} // namespace CUDANet

src/backends/cuda/cuda.cu

@@ -0,0 +1,76 @@
+#include <cuda_runtime.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <format>
+
+#include "backend/cuda/cuda.cuh"
+#include "tensor.hpp"
+
+using namespace CUDANet::Backends;
+
+
+CUDA::CUDA(const BackendConfig& config) {
+    device_id = config.device_id < 0 ? 0 : config.device_id;
+    supported_dtypes = {DType::FLOAT32};
+    default_dtype = DType::FLOAT32;
+    initialize();
+}
+
+bool CUDA::is_cuda_available() {
+    int device_count;
+    cudaError_t result = cudaGetDeviceCount(&device_count);
+    
+    // Return false instead of crashing
+    if (result != cudaSuccess || device_count == 0) {
+        return false;
+    }
+    return true;
+}
+
+void CUDA::initialize() {
+
+    int device_count;
+    CUDA_CHECK(cudaGetDeviceCount(&device_count));
+    if (device_id >= device_count) {
+        throw std::runtime_error(std::format("Invalid device id {}, only {} devices available", device_id, device_count));
+    }
+
+    CUDA_CHECK(cudaSetDevice(device_id));
+
+    cudaDeviceProp deviceProp;
+    CUDA_CHECK(cudaGetDeviceProperties(&deviceProp, device_id));
+
+    std::printf("Using CUDA device %d: %s\n", device_id, deviceProp.name);
+}
+
+bool CUDA::supports_dtype(DType dtype) const {
+    return supported_dtypes.contains(dtype);
+}
+
+void CUDA::set_default_dtype(DType dtype) {
+    if (!supported_dtypes.contains(dtype)) {
+        throw std::runtime_error("Unsupported dtype");
+    }
+
+    default_dtype = dtype;
+}
+
+CUDANet::DType CUDA::get_default_dtype() const {
+    if (default_dtype) {
+        return default_dtype.value();
+    }
+    
+    const_cast<CUDA*>(this)->default_dtype = DType::FLOAT32;
+    return DType::FLOAT32;
+}
+
+void* CUDA::allocate(size_t bytes) {
+    void* d_ptr = nullptr;
+    CUDA_CHECK(cudaMalloc(&d_ptr, bytes));
+    return d_ptr;
+}
+
+void CUDA::deallocate(void* ptr) {
+    CUDA_CHECK(cudaFree(ptr));
+}

src/backends/cuda/cuda_backend.cu

@@ -1,39 +0,0 @@
-#include <cuda_runtime.h>
-
-#include <cstdio>
-#include <cstdlib>
-#include <cuda_helper.cuh>
-
-#include "backend/cuda.cuh"
-
-cudaDeviceProp initializeCUDA() {
-    int deviceCount;
-    CUDA_CHECK(cudaGetDeviceCount(&deviceCount));
-
-    if (deviceCount == 0) {
-        std::fprintf(stderr, "No CUDA devices found. Exiting.\n");
-        std::exit(EXIT_FAILURE);
-    }
-
-    int device = 0;
-    CUDA_CHECK(cudaSetDevice(device));
-
-    cudaDeviceProp deviceProp;
-    CUDA_CHECK(cudaGetDeviceProperties(&deviceProp, device));
-
-    std::printf("Using CUDA device %d: %s\n", device, deviceProp.name);
-
-    return deviceProp;
-}
-
-using namespace CUDANet::Backend;
-
-void* CUDABackend::allocate(size_t bytes) {
-    void* d_ptr = nullptr;
-    CUDA_CHECK(cudaMalloc(&d_ptr, bytes));
-    return d_ptr;
-}
-
-void CUDABackend::deallocate(void* ptr) {
-    CUDA_CHECK(cudaFree(ptr));
-}

src/backends/cuda/kernels/activation_functions.cu

@@ -1,12 +1,19 @@
-#include "activation_functions.cuh"
-#include "cuda_helper.cuh"
+#include "backend/cuda/kernels/activation_functions.cuh"
 
 using namespace CUDANet;
 
-__global__ void Kernels::sigmoid(
+template 
+__global__ void Kernels::sigmoid<float>(
     const float* __restrict__ src,
     float* __restrict__ dst,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::sigmoid(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const unsigned int len
 ) {
     int stride = gridDim.x * blockDim.x;
     int tid    = blockDim.x * blockIdx.x + threadIdx.x;
@@ -16,10 +23,17 @@ __global__ void Kernels::sigmoid(
     }
 }
 
-__global__ void Kernels::relu(
+template __global__ void Kernels::relu<float>(
     const float* __restrict__ src,
     float* __restrict__ dst,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::relu(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const unsigned int len
 ) {
     int stride = gridDim.x * blockDim.x;
     int tid    = blockDim.x * blockIdx.x + threadIdx.x;

src/backends/cuda/kernels/convolution.cu

@@ -1,60 +1,71 @@
 #include <iostream>
 
-#include "convolution.cuh"
+#include "backend/cuda/kernels/convolution.cuh"
 
 using namespace CUDANet;
 
-__global__ void Kernels::convolution(
+template __global__ void Kernels::convolution<float>(
     const float* __restrict__ d_input,
     const float* __restrict__ d_kernel,
     const float* __restrict__ d_bias,
     float* __restrict__ d_output,
-    const shape2d inputSize,
-    const int   nChannels,
-    const shape2d paddingSize,
-    const shape2d kernelSize,
-    const shape2d stride,
-    const int   nFilters,
-    const shape2d outputSize
+    const Shape input_shape,
+    const Shape padding_shape,
+    const Shape kernel_shape,
+    const Shape stride_shape,
+    const Shape output_shape
+);
+
+template <typename T>
+__global__ void Kernels::convolution(
+    const T* __restrict__ d_input,
+    const T* __restrict__ d_kernel,
+    const T* __restrict__ d_bias,
+    T* __restrict__ d_output,
+    const Shape input_shape,
+    const Shape padding_shape,
+    const Shape kernel_shape,
+    const Shape stride_shape,
+    const Shape output_shape
 ) {
     int j = blockDim.x * blockIdx.x + threadIdx.x;
     int i = blockDim.y * blockIdx.y + threadIdx.y;
     int f = blockDim.z * blockIdx.z + threadIdx.z;
 
-    if (i >= outputSize.first || j >= outputSize.second || f >= nFilters) {
+    if (i >= output_shape[0] || j >= output_shape[1] || f >= output_shape[2]) {
         return;
     }
 
-    float sum = 0.0f;
+    T sum = static_cast<T>(0);
 
     // Iterate over kernel and input matrix
-    for (int c = 0; c < nChannels; c++) {
-        for (int k = 0; k < kernelSize.first; k++) {
-            for (int l = 0; l < kernelSize.second; l++) {
+    for (int c = 0; c < input_shape[2]; c++) {
+        for (int k = 0; k < kernel_shape[0]; k++) {
+            for (int l = 0; l < kernel_shape[1]; l++) {
                 // if i, j is in the padding region
-                if (i * stride.first + k < paddingSize.first ||
-                    i * stride.first + k >=
-                        (inputSize.first + paddingSize.first) ||
-                    j * stride.second + l < paddingSize.second ||
-                    j * stride.second + l >=
-                        (inputSize.second + paddingSize.second)) {
+                if (i * stride_shape[0] + k < padding_shape[0] ||
+                    i * stride_shape[0] + k >=
+                        (input_shape[0] + padding_shape[0]) ||
+                    j * stride_shape[1] + l < padding_shape[1] ||
+                    j * stride_shape[1] + l >=
+                        (input_shape[1] + padding_shape[1])) {
                     continue;
                 }
 
-                int kernelIndex =
-                    f * kernelSize.first * kernelSize.second * nChannels +
-                    c * kernelSize.first * kernelSize.second +
-                    k * kernelSize.second + l;
-                int inputIndex = c * inputSize.first * inputSize.second +
-                                 (i * stride.first + k - paddingSize.first) *
-                                     inputSize.second +
-                                 (j * stride.second + l - paddingSize.second);
+                int kernel_idx =
+                    f * kernel_shape[0] * kernel_shape[1] * input_shape[2] +
+                    c * kernel_shape[0] * kernel_shape[1] +
+                    k * kernel_shape[1] + l;
+                int inputIndex = c * input_shape[0] * input_shape[1] +
+                                 (i * stride_shape[0] + k - padding_shape[0]) *
+                                     input_shape[1] +
+                                 (j * stride_shape[1] + l - padding_shape[1]);
 
-                sum += d_kernel[kernelIndex] * d_input[inputIndex];
+                sum += d_kernel[kernel_idx] * d_input[inputIndex];
             }
         }
     }
 
-    d_output[f * outputSize.first * outputSize.second + i * outputSize.second + j] =
+    d_output[f * output_shape[0] * output_shape[1] + i * output_shape[1] + j] =
         sum + d_bias[f];
 }

src/backends/cuda/kernels/matmul.cu

@@ -1,19 +1,28 @@
-#include "cuda_helper.cuh"
-#include "matmul.cuh"
+#include "backend/cuda/cuda.cuh"
+#include "backend/cuda/kernels/matmul.cuh"
 
 using namespace CUDANet;
 
-__global__ void Kernels::mat_vec_mul(
+template __global__ void Kernels::mat_vec_mul<float>(
     const float* __restrict__ d_matrix,
     const float* __restrict__ d_vector,
     float* __restrict__ d_output,
     const unsigned int w,
     const unsigned int h
+);
+
+template <typename T>
+__global__ void Kernels::mat_vec_mul(
+    const T* __restrict__ d_matrix,
+    const T* __restrict__ d_vector,
+    T* __restrict__ d_output,
+    const unsigned int w,
+    const unsigned int h
 ) {
     int tid = blockDim.x * blockIdx.x + threadIdx.x;
 
     if (tid < h) {
-        float temp = 0.0f;
+        T temp = static_cast<T>(0);
 
         for (unsigned int j = 0; j < w; j++) {
             temp += d_matrix[tid * w + j] * d_vector[j];
@@ -23,11 +32,19 @@ __global__ void Kernels::mat_vec_mul(
     }
 }
 
-__global__ void Kernels::vec_vec_add(
+template __global__ void Kernels::vec_vec_add<float>(
     const float* __restrict__ d_vector1,
     const float* __restrict__ d_vector2,
     float* __restrict__ d_output,
     const unsigned int w
+);
+
+template <typename T>
+__global__ void Kernels::vec_vec_add(
+    const T* __restrict__ d_vector1,
+    const T* __restrict__ d_vector2,
+    T* __restrict__ d_output,
+    const unsigned int w
 ) {
     int tid = blockDim.x * blockIdx.x + threadIdx.x;
     if (tid >= w) {
@@ -36,11 +53,19 @@ __global__ void Kernels::vec_vec_add(
     d_output[tid] = d_vector1[tid] + d_vector2[tid];
 }
 
-__global__ void Kernels::vec_vec_sub(
+template __global__ void Kernels::vec_vec_sub<float>(
     const float* __restrict__ d_vector1,
     const float* __restrict__ d_vector2,
     float* __restrict__ d_output,
     const unsigned int w
+);
+
+template <typename T>
+__global__ void Kernels::vec_vec_sub(
+    const T* __restrict__ d_vector1,
+    const T* __restrict__ d_vector2,
+    T* __restrict__ d_output,
+    const unsigned int w
 ) {
     int tid = blockDim.x * blockIdx.x + threadIdx.x;
     if (tid >= w) {
@@ -49,11 +74,19 @@ __global__ void Kernels::vec_vec_sub(
     d_output[tid] = d_vector1[tid] - d_vector2[tid];
 }
 
-__global__ void Kernels::vec_vec_mul(
+template __global__ void Kernels::vec_vec_mul<float>(
     const float* __restrict__ d_vector1,
     const float* __restrict__ d_vector2,
     float* __restrict__ d_output,
     const unsigned int w
+);
+
+template <typename T>
+__global__ void Kernels::vec_vec_mul(
+    const T* __restrict__ d_vector1,
+    const T* __restrict__ d_vector2,
+    T* __restrict__ d_output,
+    const unsigned int w
 ) {
     int tid = blockDim.x * blockIdx.x + threadIdx.x;
     if (tid >= w) {
@@ -62,11 +95,19 @@ __global__ void Kernels::vec_vec_mul(
     d_output[tid] = d_vector1[tid] * d_vector2[tid];
 }
 
-__global__ void Kernels::vec_scalar_sub(
+template __global__ void Kernels::vec_scalar_sub<float>(
     const float* __restrict__ d_src,
     float* __restrict__ d_out,
     const float* __restrict__ d_scalar,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::vec_scalar_sub(
+    const T* __restrict__ d_src,
+    T* __restrict__ d_out,
+    const T* __restrict__ d_scalar,
+    const unsigned int len
 ) {
     int tid = blockDim.x * blockIdx.x + threadIdx.x;
     if (tid >= len) {
@@ -75,11 +116,19 @@ __global__ void Kernels::vec_scalar_sub(
     d_out[tid] = d_src[tid] - *d_scalar;
 }
 
-__global__ void Kernels::vec_scalar_add(
+template __global__ void Kernels::vec_scalar_add<float>(
     const float* __restrict__ d_src,
     float* __restrict__ d_out,
     const float* __restrict__ d_scalar,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::vec_scalar_add(
+    const T* __restrict__ d_src,
+    T* __restrict__ d_out,
+    const T* __restrict__ d_scalar,
+    const unsigned int len
 ) {
     int tid = blockDim.x * blockIdx.x + threadIdx.x;
     if (tid >= len) {
@@ -88,11 +137,19 @@ __global__ void Kernels::vec_scalar_add(
     d_out[tid] = d_src[tid] + *d_scalar;
 }
 
-__global__ void Kernels::vec_scalar_div(
+template __global__ void Kernels::vec_scalar_div<float>(
     const float* __restrict__ d_src,
     float* __restrict__ d_out,
     const float* __restrict__ d_scalar,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::vec_scalar_div(
+    const T* __restrict__ d_src,
+    T* __restrict__ d_out,
+    const T* __restrict__ d_scalar,
+    const unsigned int len
 ) {
     int tid = blockDim.x * blockIdx.x + threadIdx.x;
     if (tid >= len) {
@@ -101,11 +158,19 @@ __global__ void Kernels::vec_scalar_div(
     d_out[tid] = d_src[tid] / *d_scalar;
 }
 
-__global__ void Kernels::vec_scalar_mul(
+template __global__ void Kernels::vec_scalar_mul<float>(
     const float* __restrict__ d_src,
     float* __restrict__ d_out,
     const float* __restrict__ d_scalar,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::vec_scalar_mul(
+    const T* __restrict__ d_src,
+    T* __restrict__ d_out,
+    const T* __restrict__ d_scalar,
+    const unsigned int len
 ) {
     int tid = blockDim.x * blockIdx.x + threadIdx.x;
     if (tid >= len) {
@@ -114,52 +179,85 @@ __global__ void Kernels::vec_scalar_mul(
     d_out[tid] = d_src[tid] * *d_scalar;
 }
 
-__global__ void Kernels::vec_exp(
+template __global__ void Kernels::vec_exp<float>(
     const float* __restrict__ src,
     float* __restrict__ dst,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::vec_exp(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const unsigned int len
 ) {
     int stride = gridDim.x * blockDim.x;
     int tid    = blockDim.x * blockIdx.x + threadIdx.x;
 
     for (int i = tid; i < len; i += stride) {
+        // TODO: separate implementation for __half
         dst[i] = expf(src[i]);
     }
 }
 
-__global__ void Kernels::vec_sqrt(
+template __global__ void Kernels::vec_sqrt<float>(
     const float* __restrict__ src,
     float* __restrict__ dst,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::vec_sqrt(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const unsigned int len
 ) {
     int stride = gridDim.x * blockDim.x;
     int tid    = blockDim.x * blockIdx.x + threadIdx.x;
 
     for (int i = tid; i < len; i += stride) {
+        // TODO: separate implementation for __half
         dst[i] = sqrtf(src[i]);
     }
 }
 
-__global__ void Kernels::vec_scale(
+template __global__ void Kernels::vec_scale<float>(
     const float* __restrict__ src,
     float* __restrict__ dst,
     const float* __restrict__ scale,
     const float* epsilon,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::vec_scale(
+    const T* __restrict__ src,
+    T* __restrict__ dst,
+    const T* __restrict__ scale,
+    const T* epsilon,
+    const unsigned int len
 ) {
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (idx < len) {
+        // TODO: separate implementation for __half
         float inv_std = rsqrtf(*scale + *epsilon);
         dst[idx] = src[idx] * inv_std;
     }
 }
 
-__global__ void Kernels::max_reduce(
+template __global__ void Kernels::max_reduce<float>(
     const float* __restrict__ d_vector,
     float* __restrict__ d_output,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::max_reduce(
+    const T* __restrict__ d_vector,
+    T* __restrict__ d_output,
+    const unsigned int len
 ) {
-    __shared__ float shared_max[BLOCK_SIZE];
+    __shared__ T shared_max[BLOCK_SIZE];
     int i       = blockIdx.x * blockDim.x + threadIdx.x;
 
     if (i < len) {
@@ -172,6 +270,7 @@ __global__ void Kernels::max_reduce(
 
     for (int s = blockDim.x / 2; s > 0; s >>= 1) {
         if (threadIdx.x < s) {
+            // TODO: separate implementation for __half
             shared_max[threadIdx.x] = fmaxf(shared_max[threadIdx.x], shared_max[threadIdx.x + s]);
         }
         __syncthreads();
@@ -182,18 +281,25 @@ __global__ void Kernels::max_reduce(
     }
 }
 
-__global__ void Kernels::sum_reduce(
+template __global__ void Kernels::sum_reduce<float>(
     const float* __restrict__ d_vector,
     float* __restrict__ d_output,
     const unsigned int len
+);
+
+template <typename T>
+__global__ void Kernels::sum_reduce(
+    const T* __restrict__ d_vector,
+    T* __restrict__ d_output,
+    const unsigned int len
 ) {
-    __shared__ float partial_sum[BLOCK_SIZE];
+    __shared__ T partial_sum[BLOCK_SIZE];
     int              i = blockIdx.x * blockDim.x + threadIdx.x;
 
     if (i < len) {
         partial_sum[threadIdx.x] = d_vector[i];
     } else {
-        partial_sum[threadIdx.x] = 0.0f;
+        partial_sum[threadIdx.x] = static_cast<T>(0);
     }
 
     __syncthreads();

src/backends/cuda/kernels/pool.cu

@@ -0,0 +1,104 @@
+#include "layer.hpp"
+#include "backend/cuda/kernels/pool.cuh"
+
+using namespace CUDANet;
+
+template __global__ void Kernels::max_pool<float>(
+    const float* __restrict__ d_input,
+    float* __restrict__ d_output,
+    const Shape input_shape,
+    const Shape output_shape,
+    const Shape pool_shape,
+    const Shape stride_shape,
+    const Shape padding_shape
+);
+
+template <typename T>
+__global__ void Kernels::max_pool(
+    const T* __restrict__ d_input,
+    T* __restrict__ d_output,
+    const Shape input_shape,
+    const Shape output_shape,
+    const Shape pool_shape,
+    const Shape stride_shape,
+    const Shape padding_shape
+) {
+    int j = blockDim.x * blockIdx.x + threadIdx.x;
+    int i = blockDim.y * blockIdx.y + threadIdx.y;
+    int c = blockDim.z * blockIdx.z + threadIdx.z;
+
+    if (i >= output_shape[0] || j >= output_shape[1] || c >= output_shape[2]) {
+        return;
+    }
+
+    T max = static_cast<T>(0);
+
+    for (int k = 0; k < pool_shape[0]; k++) {
+        for (int l = 0; l < pool_shape[1]; l++) {
+            int inputRow = i * stride_shape[0] + k - padding_shape[0];
+            int inputCol = j * stride_shape[1] + l - padding_shape[1];
+
+            if (inputRow >= 0 && inputRow < input_shape[0] && inputCol >= 0 &&
+                inputCol < input_shape[1]) {
+                int inputIndex = c * input_shape[0] * input_shape[1] +
+                                 inputRow * input_shape[1] + inputCol;
+                if (d_input[inputIndex] > max) {
+                    max = d_input[inputIndex];
+                }
+            }
+        }
+    }
+
+    d_output
+        [c * output_shape[0] * output_shape[1] + i * output_shape[1] + j] =
+            max;
+}
+
+template __global__ void Kernels::avg_pool<float>(
+    const float* __restrict__ d_input,
+    float* __restrict__ d_output,
+    const Shape input_shape,
+    const Shape output_shape,
+    const Shape pool_shape,
+    const Shape stride_shape,
+    const Shape padding_shape
+);
+
+template <typename T>
+__global__ void Kernels::avg_pool(
+    const T* __restrict__ d_input,
+    T* __restrict__ d_output,
+    const Shape input_shape,
+    const Shape output_shape,
+    const Shape pool_shape,
+    const Shape stride_shape,
+    const Shape padding_shape
+) {
+    int j = blockDim.x * blockIdx.x + threadIdx.x;
+    int i = blockDim.y * blockIdx.y + threadIdx.y;
+    int c = blockDim.z * blockIdx.z + threadIdx.z;
+
+    if (i >= output_shape[0] || j >= output_shape[1] || c >= output_shape[2]) {
+        return;
+    }
+
+    T sum = static_cast<T>(0);
+
+    for (int k = 0; k < pool_shape[0]; k++) {
+        for (int l = 0; l < pool_shape[1]; l++) {
+            int inputRow = i * stride_shape[0] + k - padding_shape[0];
+            int inputCol = j * stride_shape[1] + l - padding_shape[1];
+
+            if (inputRow >= 0 && inputRow < input_shape[0] && inputCol >= 0 &&
+                inputCol < input_shape[1]) {
+                int inputIndex = c * input_shape[0] * input_shape[1] +
+                                 inputRow * input_shape[1] + inputCol;
+                sum += d_input[inputIndex];
+            }
+        }
+    }
+
+    d_output
+        [c * output_shape[0] * output_shape[1] + i * output_shape[1] + j] =
+            sum / (pool_shape[0] * pool_shape[1]);
+}

src/backends/cuda/kernels/pooling.cu

@@ -1,85 +0,0 @@
-#include "cuda_helper.cuh"
-#include "layer.hpp"
-#include "pooling.cuh"
-
-using namespace CUDANet;
-
-__global__ void Kernels::max_pooling(
-    const float* __restrict__ d_input,
-    float* __restrict__ d_output,
-    const shape2d inputSize,
-    const shape2d outputSize,
-    const int   nChannels,
-    const shape2d poolingSize,
-    const shape2d stride,
-    const shape2d padding
-) {
-    int j = blockDim.x * blockIdx.x + threadIdx.x;
-    int i = blockDim.y * blockIdx.y + threadIdx.y;
-    int c = blockDim.z * blockIdx.z + threadIdx.z;
-
-    if (i >= outputSize.first || j >= outputSize.second || c >= nChannels) {
-        return;
-    }
-
-    float max = 0.0f;
-
-    for (int k = 0; k < poolingSize.first; k++) {
-        for (int l = 0; l < poolingSize.second; l++) {
-            int inputRow = i * stride.first + k - padding.first;
-            int inputCol = j * stride.second + l - padding.second;
-
-            if (inputRow >= 0 && inputRow < inputSize.first && inputCol >= 0 &&
-                inputCol < inputSize.second) {
-                int inputIndex = c * inputSize.first * inputSize.second +
-                                 inputRow * inputSize.second + inputCol;
-                if (d_input[inputIndex] > max) {
-                    max = d_input[inputIndex];
-                }
-            }
-        }
-    }
-
-    d_output
-        [c * outputSize.first * outputSize.second + i * outputSize.second + j] =
-            max;
-}
-
-__global__ void Kernels::avg_pooling(
-    const float* __restrict__ d_input,
-    float* __restrict__ d_output,
-    const shape2d inputSize,
-    const shape2d outputSize,
-    const int   nChannels,
-    const shape2d poolingSize,
-    const shape2d stride,
-    const shape2d padding
-) {
-    int j = blockDim.x * blockIdx.x + threadIdx.x;
-    int i = blockDim.y * blockIdx.y + threadIdx.y;
-    int c = blockDim.z * blockIdx.z + threadIdx.z;
-
-    if (i >= outputSize.first || j >= outputSize.second || c >= nChannels) {
-        return;
-    }
-
-    float sum = 0.0f;
-
-    for (int k = 0; k < poolingSize.first; k++) {
-        for (int l = 0; l < poolingSize.second; l++) {
-            int inputRow = i * stride.first + k - padding.first;
-            int inputCol = j * stride.second + l - padding.second;
-
-            if (inputRow >= 0 && inputRow < inputSize.first && inputCol >= 0 &&
-                inputCol < inputSize.second) {
-                int inputIndex = c * inputSize.first * inputSize.second +
-                                 inputRow * inputSize.second + inputCol;
-                sum += d_input[inputIndex];
-            }
-        }
-    }
-
-    d_output
-        [c * outputSize.first * outputSize.second + i * outputSize.second + j] =
-            sum / (poolingSize.first * poolingSize.second);
-}

src/backends/cuda/layer_ops.cu

@@ -1,25 +1,76 @@
-#include "backend/cuda.cuh"
-#include "utils/cuda_helper.cuh"
-#include "kernels/activation_functions.cuh"
-#include "kernels/matmul.cuh"
+#include "backend/cuda/cuda.cuh"
+#include "backend/cuda/kernels/activation_functions.cuh"
+#include "backend/cuda/kernels/convolution.cuh"
+#include "backend/cuda/kernels/matmul.cuh"
+#include "backend/cuda/kernels/pool.cuh"
 
-using namespace CUDANet::Backend;
+using namespace CUDANet::Backends;
 
-void CUDABackend::relu(Tensor &tensor) {
+void CUDA::relu(Tensor& tensor) {
+    switch (tensor.get_dtype()) {
+        case DType::FLOAT32:
+            relu_impl<float>(tensor);
+            break;
+
+        default:
+            throw std::runtime_error("Unsupported dtype");
+            break;
+    }
+}
+
+template void CUDA::relu_impl<float>(Tensor& tensor);
+
+template <typename T>
+void CUDA::relu_impl(Tensor& tensor) {
     int gridSize = (tensor.numel() + BLOCK_SIZE - 1) / BLOCK_SIZE;
-    Kernels::relu<<<gridSize, BLOCK_SIZE>>>(tensor.data<float>(), tensor.data<float>(), tensor.numel());
+    Kernels::relu<<<gridSize, BLOCK_SIZE>>>(
+        static_cast<T*>(tensor.device_ptr()), static_cast<T*>(tensor.device_ptr()), tensor.numel()
+    );
     CUDA_CHECK(cudaGetLastError());
     CUDA_CHECK(cudaDeviceSynchronize());
 }
 
-void CUDABackend::sigmoid(Tensor &tensor) {
+void CUDA::sigmoid(CUDANet::Tensor& tensor) {
+    switch (tensor.get_dtype()) {
+        case DType::FLOAT32:
+            sigmoid_impl<float>(tensor);
+            break;
+
+        default:
+            throw std::runtime_error("Unsupported dtype");
+            break;
+    }
+}
+
+template void CUDA::sigmoid_impl<float>(Tensor& tensor);
+
+template <typename T>
+void CUDA::sigmoid_impl(CUDANet::Tensor& tensor) {
     int gridSize = (tensor.numel() + BLOCK_SIZE - 1) / BLOCK_SIZE;
-    Kernels::sigmoid<<<gridSize, BLOCK_SIZE>>>(tensor.data<float>(), tensor.data<float>(), tensor.numel());
+    Kernels::sigmoid<<<gridSize, BLOCK_SIZE>>>(
+        static_cast<T*>(tensor.device_ptr()), static_cast<T*>(tensor.device_ptr()), tensor.numel()
+    );
     CUDA_CHECK(cudaGetLastError());
     CUDA_CHECK(cudaDeviceSynchronize());
 }
 
-void CUDABackend::softmax(Tensor &tensor, Tensor &temp_max, Tensor &temp_sum) {
+void CUDA::softmax(Tensor& tensor, Tensor& temp_max, Tensor& temp_sum) {
+    switch (tensor.get_dtype()) {
+        case DType::FLOAT32:
+            softmax_impl<float>(tensor, temp_max, temp_sum);
+            break;
+
+        default:
+            throw std::runtime_error("Unsupported dtype");
+            break;
+    }
+}
+
+template void
+CUDA::softmax_impl<float>(Tensor& tensor, Tensor& temp_max, Tensor& temp_sum);
+
+template <typename T>
+void CUDA::softmax_impl(Tensor& tensor, Tensor& temp_max, Tensor& temp_sum) {
     int gridSize = (tensor.numel() + BLOCK_SIZE - 1) / BLOCK_SIZE;
 
     // Find max value
@@ -27,13 +78,13 @@ void CUDABackend::softmax(Tensor &tensor, Tensor &temp_max, Tensor &temp_sum) {
 
     // Subtract max value to improve numerical stability
     Kernels::vec_scalar_sub<<<gridSize, BLOCK_SIZE>>>(
-        tensor.data<float>(), tensor.data<float>(), temp_max.data<float>(), tensor.numel()
+        static_cast<T*>(tensor.device_ptr()), static_cast<T*>(tensor.device_ptr()), static_cast<T*>(temp_max.device_ptr()), tensor.numel()
     );
     CUDA_CHECK(cudaGetLastError());
 
     // Compute exponentials
     Kernels::vec_exp<<<gridSize, BLOCK_SIZE>>>(
-        tensor.data<float>(), tensor.data<float>(), tensor.numel()
+        static_cast<T*>(tensor.device_ptr()), static_cast<T*>(tensor.device_ptr()), tensor.numel()
     );
     CUDA_CHECK(cudaGetLastError());
 
@@ -41,8 +92,426 @@ void CUDABackend::softmax(Tensor &tensor, Tensor &temp_max, Tensor &temp_sum) {
     sum(tensor, temp_sum);
 
     Kernels::vec_scalar_div<<<gridSize, BLOCK_SIZE>>>(
-        tensor.data<float>(), tensor.data<float>(), temp_sum.data<float>(), tensor.numel()
+        static_cast<T*>(tensor.device_ptr()), static_cast<T*>(tensor.device_ptr()), static_cast<T*>(temp_sum.device_ptr()), tensor.numel()
     );
     CUDA_CHECK(cudaGetLastError());
     CUDA_CHECK(cudaDeviceSynchronize());
 }
+
+CUDANet::Tensor& CUDA::dense(
+    const CUDANet::Tensor& weights,
+    const CUDANet::Tensor& biases,
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    const size_t           input_size,
+    const size_t           output_size
+) {
+    switch (input.get_dtype()) {
+        case DType::FLOAT32:
+            return dense_impl<float>(
+                weights, biases, input, output, input_size, output_size
+            );
+            break;
+
+        default:
+            throw std::runtime_error("Unsupported dtype");
+            break;
+    }
+}
+
+template CUDANet::Tensor& CUDA::dense_impl<float>(
+    const CUDANet::Tensor& weights,
+    const CUDANet::Tensor& biases,
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    const size_t           input_size,
+    const size_t           output_size
+);
+
+template <typename T>
+CUDANet::Tensor& CUDA::dense_impl(
+    const CUDANet::Tensor& weights,
+    const CUDANet::Tensor& biases,
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    const size_t           input_size,
+    const size_t           output_size
+) {
+    auto forwardGridSize =
+        (std::max(input_size, output_size) + BLOCK_SIZE - 1) / BLOCK_SIZE;
+    auto biasGridSize = (output_size + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+    Kernels::mat_vec_mul<<<forwardGridSize, BLOCK_SIZE>>>(
+        static_cast<const T*>(weights.device_ptr()), static_cast<const T*>(input.device_ptr()), static_cast<T*>(output.device_ptr()), input_size,
+        output_size
+    );
+    CUDA_CHECK(cudaGetLastError());
+
+    Kernels::vec_vec_add<<<biasGridSize, BLOCK_SIZE>>>(
+        static_cast<const T*>(biases.device_ptr()), static_cast<T*>(output.device_ptr()), static_cast<T*>(output.device_ptr()), output_size
+    );
+    CUDA_CHECK(cudaGetLastError());
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return output;
+}
+
+CUDANet::Tensor& CUDA::conv2d(
+    const CUDANet::Tensor& weights,
+    const CUDANet::Tensor& biases,
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    const CUDANet::Shape   in_shape,
+    const CUDANet::Shape   padding_shape,
+    const CUDANet::Shape   kernel_shape,
+    const CUDANet::Shape   stride_shape,
+    const CUDANet::Shape   out_shape
+) {
+    switch (input.get_dtype()) {
+        case DType::FLOAT32:
+            return conv2d_impl<float>(
+                weights, biases, input, output, in_shape, padding_shape,
+                kernel_shape, stride_shape, out_shape
+            );
+            break;
+
+        default:
+            throw std::runtime_error("Unsupported dtype");
+            break;
+    }
+}
+
+template CUDANet::Tensor& CUDA::conv2d_impl<float>(
+    const CUDANet::Tensor& weights,
+    const CUDANet::Tensor& biases,
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    const CUDANet::Shape   in_shape,
+    const CUDANet::Shape   padding_shape,
+    const CUDANet::Shape   kernel_shape,
+    const CUDANet::Shape   stride_shape,
+    const CUDANet::Shape   out_shape
+);
+
+template <typename T>
+CUDANet::Tensor& CUDA::conv2d_impl(
+    const CUDANet::Tensor& weights,
+    const CUDANet::Tensor& biases,
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    const CUDANet::Shape   in_shape,
+    const CUDANet::Shape   padding_shape,
+    const CUDANet::Shape   kernel_shape,
+    const CUDANet::Shape   stride_shape,
+    const CUDANet::Shape   out_shape
+) {
+    dim3 block(8, 8, 8);
+    dim3 grid(
+        (out_shape[0] + block.x - 1) / block.x,
+        (out_shape[1] + block.y - 1) / block.y,
+        (out_shape[2] + block.z - 1) / block.z
+    );
+
+    Kernels::convolution<<<grid, block>>>(
+        static_cast<const T*>(input.device_ptr()), static_cast<const T*>(weights.device_ptr()), static_cast<const T*>(biases.device_ptr()), static_cast<T*>(output.device_ptr()),
+        in_shape, padding_shape, kernel_shape, stride_shape, out_shape
+    );
+    CUDA_CHECK(cudaGetLastError());
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return output;
+}
+
+CUDANet::Tensor& CUDA::max_pool2d(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    CUDANet::Shape         input_shape,
+    CUDANet::Shape         pool_shape,
+    CUDANet::Shape         stride_shape,
+    CUDANet::Shape         padding_shape,
+    CUDANet::Shape         output_shape
+) {
+    switch (input.get_dtype()) {
+        case DType::FLOAT32:
+            return max_pool2d_impl<float>(
+                input, output, input_shape, pool_shape, stride_shape,
+                padding_shape, output_shape
+            );
+            break;
+
+        default:
+            throw std::runtime_error("Unsupported dtype");
+            break;
+    }
+}
+
+template CUDANet::Tensor& CUDA::max_pool2d_impl<float>(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    CUDANet::Shape         input_shape,
+    CUDANet::Shape         pool_shape,
+    CUDANet::Shape         stride_shape,
+    CUDANet::Shape         padding_shape,
+    CUDANet::Shape         output_shape
+);
+
+template <typename T>
+CUDANet::Tensor& CUDA::max_pool2d_impl(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    CUDANet::Shape         input_shape,
+    CUDANet::Shape         pool_shape,
+    CUDANet::Shape         stride_shape,
+    CUDANet::Shape         padding_shape,
+    CUDANet::Shape         output_shape
+) {
+    dim3 block(8, 8, 8);
+    dim3 grid(
+        (output_shape[0] + block.x - 1) / block.x,
+        (output_shape[1] + block.y - 1) / block.y,
+        (output_shape[2] + block.z - 1) / block.z
+    );
+
+    Kernels::max_pool<<<grid, block>>>(
+        static_cast<const T*>(input.device_ptr()), static_cast<T*>(output.device_ptr()), input_shape, output_shape,
+        pool_shape, stride_shape, padding_shape
+    );
+    CUDA_CHECK(cudaGetLastError());
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return output;
+}
+
+CUDANet::Tensor& CUDA::avg_pool2d(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    CUDANet::Shape         input_shape,
+    CUDANet::Shape         pool_shape,
+    CUDANet::Shape         stride_shape,
+    CUDANet::Shape         padding_shape,
+    CUDANet::Shape         output_shape
+) {
+    switch (input.get_dtype()) {
+        case DType::FLOAT32:
+            return avg_pool2d_impl<float>(
+                input, output, input_shape, pool_shape, stride_shape,
+                padding_shape, output_shape
+            );
+            break;
+
+        default:
+            throw std::runtime_error("Unsupported dtype");
+            break;
+    }
+}
+
+template CUDANet::Tensor& CUDA::avg_pool2d_impl<float>(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    CUDANet::Shape         input_shape,
+    CUDANet::Shape         pool_shape,
+    CUDANet::Shape         stride_shape,
+    CUDANet::Shape         padding_shape,
+    CUDANet::Shape         output_shape
+);
+
+template <typename T>
+CUDANet::Tensor& CUDA::avg_pool2d_impl(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    CUDANet::Shape         input_shape,
+    CUDANet::Shape         pool_shape,
+    CUDANet::Shape         stride_shape,
+    CUDANet::Shape         padding_shape,
+    CUDANet::Shape         output_shape
+) {
+    dim3 block(8, 8, 8);
+    dim3 grid(
+        (output_shape[0] + block.x - 1) / block.x,
+        (output_shape[1] + block.y - 1) / block.y,
+        (output_shape[2] + block.z - 1) / block.z
+    );
+
+    Kernels::avg_pool<<<grid, block>>>(
+        static_cast<const T*>(input.device_ptr()), static_cast<T*>(output.device_ptr()), input_shape, output_shape,
+        pool_shape, stride_shape, padding_shape
+    );
+    CUDA_CHECK(cudaGetLastError());
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return output;
+}
+
+CUDANet::Tensor& CUDA::batch_norm(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    CUDANet::Shape         input_shape,
+    CUDANet::Tensor&       weights,
+    CUDANet::Tensor&       biases,
+    CUDANet::Tensor&       running_mean,
+    CUDANet::Tensor&       running_var,
+    CUDANet::Tensor&       epsilon
+) {
+    switch (input.get_dtype()) {
+        case DType::FLOAT32:
+            return batch_norm_impl<float>(
+                input, output, input_shape, weights, biases, running_mean,
+                running_var, epsilon
+            );
+            break;
+
+        default:
+            throw std::runtime_error("Unsupported dtype");
+            break;
+    }
+}
+
+template CUDANet::Tensor& CUDA::batch_norm_impl<float>(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    CUDANet::Shape         input_shape,
+    CUDANet::Tensor&       weights,
+    CUDANet::Tensor&       biases,
+    CUDANet::Tensor&       running_mean,
+    CUDANet::Tensor&       running_var,
+    CUDANet::Tensor&       epsilon
+);
+
+template <typename T>
+CUDANet::Tensor& CUDA::batch_norm_impl(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor&       output,
+    CUDANet::Shape         input_shape,
+    CUDANet::Tensor&       weights,
+    CUDANet::Tensor&       biases,
+    CUDANet::Tensor&       running_mean,
+    CUDANet::Tensor&       running_var,
+    CUDANet::Tensor&       epsilon
+) {
+    auto gridSize =
+        (input_shape[0] * input_shape[1] + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+    for (int i = 0; i < input_shape[2]; i++) {
+        // Subtract mean from input
+        Kernels::vec_scalar_sub<<<gridSize, BLOCK_SIZE>>>(
+            static_cast<const T*>(input.device_ptr()) + i * input_shape[0] * input_shape[1],
+            static_cast<T*>(output.device_ptr()) + i * input_shape[0] * input_shape[1],
+            &static_cast<T*>(running_mean.device_ptr())[i], input_shape[0] * input_shape[1]
+        );
+        CUDA_CHECK(cudaGetLastError());
+
+        // Divide by sqrt(running_var + epsilon)
+        Kernels::vec_scale<<<gridSize, BLOCK_SIZE>>>(
+            static_cast<T*>(output.device_ptr()) + i * input_shape[0] * input_shape[1],
+            static_cast<T*>(output.device_ptr()) + i * input_shape[0] * input_shape[1],
+            &static_cast<T*>(running_var.device_ptr())[i], static_cast<T*>(epsilon.device_ptr()),
+            input_shape[0] * input_shape[1]
+        );
+        CUDA_CHECK(cudaGetLastError());
+
+        // Multiply by weights
+        Kernels::vec_scalar_mul<<<gridSize, BLOCK_SIZE>>>(
+            static_cast<T*>(output.device_ptr()) + i * input_shape[0] * input_shape[1],
+            static_cast<T*>(output.device_ptr()) + i * input_shape[0] * input_shape[1],
+            &static_cast<T*>(weights.device_ptr())[i], input_shape[0] * input_shape[1]
+        );
+        CUDA_CHECK(cudaGetLastError());
+
+        // Add biases
+        Kernels::vec_scalar_add<<<gridSize, BLOCK_SIZE>>>(
+            static_cast<T*>(output.device_ptr()) + i * input_shape[0] * input_shape[1],
+            static_cast<T*>(output.device_ptr()) + i * input_shape[0] * input_shape[1],
+            &static_cast<T*>(biases.device_ptr())[i], input_shape[0] * input_shape[1]
+        );
+        CUDA_CHECK(cudaGetLastError());
+    }
+    CUDA_CHECK(cudaDeviceSynchronize());
+    return output;
+}
+
+CUDANet::Tensor& CUDA::concat(
+    CUDANet::Tensor& input_a,
+    CUDANet::Tensor& input_b,
+    CUDANet::Tensor& output
+) {
+    switch (input_a.get_dtype()) {
+        case DType::FLOAT32:
+            return concat_impl<float>(
+                input_a, input_b, output
+            );
+            break;
+
+        default:
+            throw std::runtime_error("Unsupported dtype");
+            break;
+    }
+}
+
+template CUDANet::Tensor& CUDA::concat_impl<float>(
+    CUDANet::Tensor& input_a,
+    CUDANet::Tensor& input_b,
+    CUDANet::Tensor& output
+);
+
+template <typename T>
+CUDANet::Tensor& CUDA::concat_impl(
+    CUDANet::Tensor& input_a,
+    CUDANet::Tensor& input_b,
+    CUDANet::Tensor& output
+) {
+    CUDA_CHECK(cudaMemcpy(
+        static_cast<T*>(output.device_ptr()), static_cast<const T*>(input_a.device_ptr()), input_a.size(),
+        cudaMemcpyDeviceToDevice
+    ));
+
+    CUDA_CHECK(cudaMemcpy(
+        static_cast<T*>(output.device_ptr()) + input_a.numel(), static_cast<const T*>(input_b.device_ptr()), input_b.size(),
+        cudaMemcpyDeviceToDevice
+    ));
+
+    CUDA_CHECK(cudaGetLastError());
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return output;
+}
+
+CUDANet::Tensor& CUDA::add(
+    CUDANet::Tensor& input_a,
+    CUDANet::Tensor& input_b,
+    CUDANet::Tensor& output
+) {
+    switch (input_a.get_dtype()) {
+    case DType::FLOAT32:
+        return add_impl<float>(
+            input_a, input_b, output
+        );
+        break;
+
+    default:
+        throw std::runtime_error("Unsupported dtype");
+        break;
+    }
+}
+
+template CUDANet::Tensor& CUDA::add_impl<float>(
+    CUDANet::Tensor& input_a,
+    CUDANet::Tensor& input_b,
+    CUDANet::Tensor& output
+);
+
+template <typename T>
+CUDANet::Tensor& CUDA::add_impl(
+    CUDANet::Tensor& input_a,
+    CUDANet::Tensor& input_b,
+    CUDANet::Tensor& output
+) {
+    auto gridSize = (input_a.numel() + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+    Kernels::vec_vec_add<<<gridSize, BLOCK_SIZE>>>(
+        static_cast<const T*>(input_a.device_ptr()), static_cast<const T*>(input_b.device_ptr()), static_cast<T*>(output.device_ptr()), input_a.numel()
+    );
+    CUDA_CHECK(cudaGetLastError());
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return output;
+}

src/backends/cuda/layers/activation.cu

@@ -1,77 +0,0 @@
-#include <vector>
-
-#include "activation.hpp"
-#include "activation_functions.cuh"
-#include "cuda_helper.cuh"
-#include "matmul.cuh"
-#include "vector.cuh"
-
-using namespace CUDANet::Layers;
-
-void Activation::initCUDA() {
-    if (activationType == SOFTMAX) {
-        d_softmax_sum = nullptr;
-        CUDA_CHECK(cudaMalloc((void**)&d_softmax_sum, sizeof(float) * length));
-
-        d_max = nullptr;
-        CUDA_CHECK(cudaMalloc((void**)&d_max, sizeof(float) * length));
-    }
-
-    gridSize = (length + BLOCK_SIZE - 1) / BLOCK_SIZE;
-}
-
-void Activation::delCUDA() {
-    if (activationType == SOFTMAX) {
-        CUDA_CHECK(cudaFree(d_softmax_sum));
-        CUDA_CHECK(cudaFree(d_max));
-    }
-}
-
-void Activation::activateCUDA(float* d_input) {
-
-    // float sum = 0.0f;
-
-    switch (activationType) {
-        case SIGMOID:
-            Kernels::sigmoid<<<gridSize, BLOCK_SIZE>>>(
-                d_input, d_input, length
-            );
-            CUDA_CHECK(cudaGetLastError());
-            break;
-
-        case RELU:
-            Kernels::relu<<<gridSize, BLOCK_SIZE>>>(d_input, d_input, length);
-            CUDA_CHECK(cudaGetLastError());
-            break;
-        case SOFTMAX:
-
-            // Find max value
-            Utils::max(d_input, d_max, length);
-
-            // Subtract max value to improve numerical stability
-            Kernels::vec_scalar_sub<<<gridSize, BLOCK_SIZE>>>(
-                d_input, d_input, &d_max[0], length
-            );
-            CUDA_CHECK(cudaGetLastError());
-
-            // Compute exponentials
-            Kernels::vec_exp<<<gridSize, BLOCK_SIZE>>>(
-                d_input, d_input, length
-            );
-            CUDA_CHECK(cudaGetLastError());
-            
-            // Find sum
-            Utils::sum(d_input, d_softmax_sum, length);
-
-            Kernels::vec_scalar_div<<<gridSize, BLOCK_SIZE>>>(
-                d_input, d_input, &d_softmax_sum[0], length
-            );
-            CUDA_CHECK(cudaGetLastError());
-            break;
-
-        default:
-            break;    
-    }
-
-    CUDA_CHECK(cudaDeviceSynchronize());
-}

src/backends/cuda/layers/add.cu

@@ -1,28 +0,0 @@
-#include "add.hpp"
-#include "matmul.cuh"
-#include "cuda_helper.cuh"
-
-using namespace CUDANet::Layers;
-
-void Add::initCUDA() {
-    d_output = nullptr;
-    CUDA_CHECK(cudaMalloc((void**)&d_output, sizeof(float) * inputSize));
-
-    gridSize = (inputSize + BLOCK_SIZE - 1) / BLOCK_SIZE;
-}
-
-void Add::delCUDA() {
-    cudaFree(d_output);
-}
-
-float* Add::forwardCUDA(const float* d_inputA, const float* d_inputB) {
-
-    Kernels::vec_vec_add<<<gridSize, BLOCK_SIZE>>>(
-        d_inputA, d_inputB, d_output, inputSize
-    );
-    CUDA_CHECK(cudaGetLastError());
-    CUDA_CHECK(cudaDeviceSynchronize());
-
-    return d_output;
-
-}

src/backends/cuda/layers/avg_pooling.cu

@@ -1,45 +0,0 @@
-#include "avg_pooling.hpp"
-#include "cuda_helper.cuh"
-#include "pooling.cuh"
-
-using namespace CUDANet::Layers;
-
-void AvgPooling2d::initCUDA() {
-    d_output = nullptr;
-    CUDA_CHECK(cudaMalloc(
-        (void**)&d_output,
-        sizeof(float) * outputSize.first * outputSize.second * nChannels
-    ));
-}
-
-void AvgPooling2d::delCUDA() {
-    cudaFree(d_output);
-}
-
-float* AvgPooling2d::forwardCUDA(const float* d_input) {
-    dim3 block(8, 8, 8);
-    dim3 grid(
-        (outputSize.first + block.x - 1) / block.x,
-        (outputSize.second + block.y - 1) / block.y,
-        (nChannels + block.z - 1) / block.z
-    );
-
-    Kernels::avg_pooling<<<grid, block>>>(
-        d_input, d_output, inputSize, outputSize, nChannels, poolingSize,
-        stride, padding
-    );
-    CUDA_CHECK(cudaGetLastError());
-
-    activation->activate(d_output);
-    CUDA_CHECK(cudaDeviceSynchronize());
-
-    return d_output;
-}
-
-void AdaptiveAvgPooling2d::initCUDA() {
-    cudaFree(d_output);
-    cudaMalloc(
-        (void**)&d_output,
-        sizeof(float) * outputSize.first * outputSize.second * nChannels
-    );
-}

src/backends/cuda/layers/batch_norm.cu

@@ -1,120 +0,0 @@
-#include <vector>
-
-#include "activation.hpp"
-#include "batch_norm.hpp"
-#include "cuda_helper.cuh"
-#include "layer.hpp"
-#include "matmul.cuh"
-#include "vector.cuh"
-
-using namespace CUDANet::Layers;
-
-void BatchNorm2d::initCUDA() {
-    d_output = nullptr;
-    CUDA_CHECK(cudaMalloc(
-        (void **)&d_output,
-        sizeof(float) * inputSize.first * inputSize.second * inputChannels
-    ));
-
-    d_running_mean = nullptr;
-    CUDA_CHECK(
-        cudaMalloc((void **)&d_running_mean, sizeof(float) * inputChannels)
-    );
-
-    d_running_var = nullptr;
-    CUDA_CHECK(
-        cudaMalloc((void **)&d_running_var, sizeof(float) * inputChannels)
-    );
-
-    d_weights = nullptr;
-    CUDA_CHECK(cudaMalloc((void **)&d_weights, sizeof(float) * inputChannels));
-
-    d_biases = nullptr;
-    CUDA_CHECK(cudaMalloc((void **)&d_biases, sizeof(float) * inputChannels));
-
-    d_length     = nullptr;
-    float length = (float)inputSize.first * inputSize.second;
-    CUDA_CHECK(cudaMalloc((void **)&d_length, sizeof(float)));
-    CUDA_CHECK(
-        cudaMemcpy(d_length, &length, sizeof(float), cudaMemcpyHostToDevice)
-    );
-
-    d_epsilon = nullptr;
-    CUDA_CHECK(cudaMalloc((void **)&d_epsilon, sizeof(float)));
-    CUDA_CHECK(
-        cudaMemcpy(d_epsilon, &epsilon, sizeof(float), cudaMemcpyHostToDevice)
-    );
-
-    gridSize =
-        (inputSize.first * inputSize.second + BLOCK_SIZE - 1) / BLOCK_SIZE;
-}
-
-void BatchNorm2d::delCUDA() {
-    cudaFree(d_output);
-    cudaFree(d_running_mean);
-    cudaFree(d_running_var);
-    cudaFree(d_weights);
-    cudaFree(d_biases);
-    cudaFree(d_length);
-    cudaFree(d_epsilon);
-}
-
-void BatchNorm2d::toCuda() {
-    CUDA_CHECK(cudaMemcpy(
-        d_weights, weights.data(), sizeof(float) * inputChannels,
-        cudaMemcpyHostToDevice
-    ));
-    CUDA_CHECK(cudaMemcpy(
-        d_biases, biases.data(), sizeof(float) * inputChannels,
-        cudaMemcpyHostToDevice
-    ));
-    CUDA_CHECK(cudaMemcpy(
-        d_running_mean, running_mean.data(), sizeof(float) * inputChannels,
-        cudaMemcpyHostToDevice
-    ));
-    CUDA_CHECK(cudaMemcpy(
-        d_running_var, running_var.data(), sizeof(float) * inputChannels,
-        cudaMemcpyHostToDevice
-    ));
-}
-
-float *BatchNorm2d::forwardCUDA(const float *d_input) {
-    // Compute per-channel batch normalization
-    for (int i = 0; i < inputChannels; i++) {
-        // Subtract mean from input
-        Kernels::vec_scalar_sub<<<gridSize, BLOCK_SIZE>>>(
-            d_input + i * inputSize.first * inputSize.second,
-            d_output + i * inputSize.first * inputSize.second,
-            &d_running_mean[i], inputSize.first * inputSize.second
-        );
-        CUDA_CHECK(cudaGetLastError());
-
-        // Divide by sqrt(running_var + epsilon)
-        Kernels::vec_scale<<<gridSize, BLOCK_SIZE>>>(
-            d_output + i * inputSize.first * inputSize.second,
-            d_output + i * inputSize.first * inputSize.second,
-            &d_running_var[i], d_epsilon, inputSize.first * inputSize.second
-        );
-        CUDA_CHECK(cudaGetLastError());
-
-        // Multiply by weights
-        Kernels::vec_scalar_mul<<<gridSize, BLOCK_SIZE>>>(
-            d_output + i * inputSize.first * inputSize.second,
-            d_output + i * inputSize.first * inputSize.second, &d_weights[i],
-            inputSize.first * inputSize.second
-        );
-        CUDA_CHECK(cudaGetLastError());
-
-        // Add biases
-        Kernels::vec_scalar_add<<<gridSize, BLOCK_SIZE>>>(
-            d_output + i * inputSize.first * inputSize.second,
-            d_output + i * inputSize.first * inputSize.second, &d_biases[i],
-            inputSize.first * inputSize.second
-        );
-        CUDA_CHECK(cudaGetLastError());
-    }
-
-    activation->activate(d_output);
-
-    return d_output;
-}

src/backends/cuda/layers/concat.cu

@@ -1,31 +0,0 @@
-#include "concat.hpp"
-#include "cuda_helper.cuh"
-
-using namespace CUDANet::Layers;
-
-void Concat::initCUDA() {
-    d_output = nullptr;
-    CUDA_CHECK(
-        cudaMalloc((void**)&d_output, sizeof(float) * (inputASize + inputBSize))
-    );
-}
-
-void Concat::delCUDA() {
-    cudaFree(d_output);
-}
-
-float* Concat::forwardCUDA(const float* d_input_A, const float* d_input_B) {
-    CUDA_CHECK(cudaMemcpy(
-        d_output, d_input_A, sizeof(float) * inputASize,
-        cudaMemcpyDeviceToDevice
-    ));
-
-    CUDA_CHECK(cudaMemcpy(
-        d_output + inputASize, d_input_B, sizeof(float) * inputBSize,
-        cudaMemcpyDeviceToDevice
-    ));
-
-    CUDA_CHECK(cudaDeviceSynchronize());
-
-    return d_output;
-}

src/backends/cuda/layers/conv2d.cu

@@ -1,73 +0,0 @@
-#include <vector>
-
-#include "activation.hpp"
-#include "conv2d.hpp"
-#include "convolution.cuh"
-#include "cuda_helper.cuh"
-#include "layer.hpp"
-#include "matmul.cuh"
-#include "vector.cuh"
-
-using namespace CUDANet::Layers;
-
-void Conv2d::initCUDA() {
-    d_output = nullptr;
-    CUDA_CHECK(cudaMalloc(
-        (void**)&d_output,
-        sizeof(float) * outputSize.first * outputSize.second * numFilters
-    ));
-
-    d_weights = nullptr;
-    CUDA_CHECK(cudaMalloc(
-        (void**)&d_weights, sizeof(float) * kernelSize.first *
-                                kernelSize.second * inputChannels * numFilters
-    ));
-
-    d_biases = nullptr;
-    CUDA_CHECK(cudaMalloc((void**)&d_biases, sizeof(float) * numFilters));
-}
-
-void Conv2d::delCUDA() {
-    cudaFree(d_output);
-    cudaFree(d_weights);
-    cudaFree(d_biases);
-}
-
-void Conv2d::toCuda() {
-    CUDA_CHECK(cudaMemcpy(
-        d_weights, weights.data(),
-        sizeof(float) * kernelSize.first * kernelSize.second * inputChannels *
-            numFilters,
-        cudaMemcpyHostToDevice
-    ));
-
-    CUDA_CHECK(cudaMemcpy(
-        d_biases, biases.data(), sizeof(float) * numFilters,
-        cudaMemcpyHostToDevice
-    ));
-}
-
-float* Conv2d::forwardCUDA(const float* d_input) {
-    // Convolve
-    dim3 block(8, 8, 8);
-    dim3 grid(
-        (outputSize.first + block.x - 1) / block.x,
-        (outputSize.second + block.y - 1) / block.y,
-        (numFilters + block.z - 1) / block.z
-    );
-
-    CUDANet::Utils::clear(d_output, outputSize.first * outputSize.second * numFilters);
-
-    Kernels::convolution<<<grid, block>>>(
-        d_input, d_weights, d_biases, d_output, inputSize, inputChannels,
-        paddingSize, kernelSize, stride, numFilters, outputSize
-    );
-    CUDA_CHECK(cudaGetLastError());
-
-    // Apply activation
-    activation->activate(d_output);
-
-    CUDA_CHECK(cudaDeviceSynchronize());
-
-    return d_output;
-}

src/backends/cuda/layers/dense.cu

@@ -1,69 +0,0 @@
-#include <cuda_runtime.h>
-
-#include <cstdio>
-#include <cstdlib>
-#include <functional>
-#include <iostream>
-
-#include "vector.cuh"
-#include "activation.hpp"
-#include "cuda_helper.cuh"
-#include "dense.hpp"
-#include "matmul.cuh"
-
-using namespace CUDANet::Layers;
-
-void Dense::initCUDA() {
-    d_output = nullptr;
-
-    CUDA_CHECK(cudaMalloc((void**)&d_output, sizeof(float) * outputSize));
-
-    d_weights = nullptr;
-    d_biases  = nullptr;
-
-    // Allocate GPU memory for weights and biases
-    CUDA_CHECK(
-        cudaMalloc((void**)&d_weights, sizeof(float) * inputSize * outputSize)
-    );
-    CUDA_CHECK(cudaMalloc((void**)&d_biases, sizeof(float) * outputSize));
-    toCuda();
-
-    // Calculate block and grid sizes
-    forwardGridSize =
-        (std::max(inputSize, outputSize) + BLOCK_SIZE - 1) / BLOCK_SIZE;
-    biasGridSize = (outputSize + BLOCK_SIZE - 1) / BLOCK_SIZE;
-}
-
-void Dense::delCUDA() {
-    cudaFree(d_output);
-    cudaFree(d_weights);
-    cudaFree(d_biases);
-}
-
-void Dense::toCuda() {
-    CUDA_CHECK(cudaMemcpy(
-        d_weights, weights.data(), sizeof(float) * inputSize * outputSize,
-        cudaMemcpyHostToDevice
-    ));
-    CUDA_CHECK(cudaMemcpy(
-        d_biases, biases.data(), sizeof(float) * outputSize,
-        cudaMemcpyHostToDevice
-    ));
-}
-
-float* Dense::forwardCUDA(const float* d_input) {
-    Kernels::mat_vec_mul<<<forwardGridSize, BLOCK_SIZE>>>(
-        d_weights, d_input, d_output, inputSize, outputSize
-    );
-    CUDA_CHECK(cudaGetLastError());
-
-    Kernels::vec_vec_add<<<biasGridSize, BLOCK_SIZE>>>(
-        d_biases, d_output, d_output, outputSize
-    );
-    CUDA_CHECK(cudaGetLastError());
-
-    activation->activate(d_output);
-    CUDA_CHECK(cudaDeviceSynchronize());
-
-    return d_output;
-}

src/backends/cuda/layers/input.cu

@@ -1,22 +0,0 @@
-#include "cuda_helper.cuh"
-#include "input.hpp"
-
-using namespace CUDANet::Layers;
-
-void Input::initCUDA() {
-    d_output = nullptr;
-    CUDA_CHECK(cudaMalloc((void**)&d_output, sizeof(float) * inputSize));
-}
-
-void Input::delCUDA() {
-    cudaFree(d_output);
-}
-
-float* Input::forwardCUDA(const float* input) {
-    CUDA_CHECK(cudaMemcpy(
-        d_output, input, sizeof(float) * inputSize, cudaMemcpyHostToDevice
-    ));
-    CUDA_CHECK(cudaDeviceSynchronize());
-
-    return d_output;
-}

src/backends/cuda/layers/max_pooling.cu

@@ -1,38 +0,0 @@
-#include "cuda_helper.cuh"
-#include "max_pooling.hpp"
-#include "pooling.cuh"
-
-using namespace CUDANet::Layers;
-
-void MaxPooling2d::initCUDA() {
-        d_output = nullptr;
-    CUDA_CHECK(cudaMalloc(
-        (void**)&d_output,
-        sizeof(float) * outputSize.first * outputSize.second * nChannels
-    ));
-}
-
-void MaxPooling2d::delCUDA() {
-    cudaFree(d_output);
-}
-
-
-float* MaxPooling2d::forwardCUDA(const float* d_input) {
-    dim3 block(8, 8, 8);
-    dim3 grid(
-        (outputSize.first + block.x - 1) / block.x,
-        (outputSize.second + block.y - 1) / block.y,
-        (nChannels + block.z - 1) / block.z
-    );
-
-    Kernels::max_pooling<<<grid, block>>>(
-        d_input, d_output, inputSize, outputSize, nChannels, poolingSize,
-        stride, padding
-    );
-    CUDA_CHECK(cudaGetLastError());
-
-    activation->activate(d_output);
-    CUDA_CHECK(cudaDeviceSynchronize());
-
-    return d_output;
-}

src/backends/cuda/layers/output.cu

@@ -1,14 +0,0 @@
-#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/backends/cuda/tensor_ops.cu

@@ -1,18 +1,32 @@
 #include <iostream>
 
-#include "backend/backend.hpp"
-#include "backend/cuda.cuh"
-#include "utils/cuda_helper.cuh"
-#include "kernels/matmul.cuh"
+#include "backend.hpp"
+#include "backend/cuda/cuda.cuh"
+#include "backend/cuda/kernels/matmul.cuh"
 
-using namespace CUDANet::Backend;
+using namespace CUDANet::Backends;
 
-void CUDABackend::print(const CUDANet::Backend::Tensor &input) {
+void CUDA::print(const CUDANet::Tensor &input) {
+    switch (input.get_dtype()) {
+    case DType::FLOAT32:
+        print_impl<float>(input);
+        break;
+
+    default:
+        throw std::runtime_error("Unsupported dtype");
+        break;
+    }
+}
+
+template void CUDA::print_impl<float> (const CUDANet::Tensor &input);
+
+template <typename T>
+void CUDA::print_impl(const CUDANet::Tensor &input) {
     auto length = input.numel();
-    std::vector<float> h_vec(input.numel());
+    std::vector<T> h_vec(input.numel());
 
     CUDA_CHECK(cudaMemcpy(
-        h_vec.data(), input.data<float>(), sizeof(float) * length, cudaMemcpyDeviceToHost
+        h_vec.data(), static_cast<const T*>(input.device_ptr()), sizeof(T) * length, cudaMemcpyDeviceToHost
     ));
 
     for (int i = 0; i < length; ++i) {
@@ -22,41 +36,109 @@ void CUDABackend::print(const CUDANet::Backend::Tensor &input) {
     std::cout << std::endl;
 }
 
-void CUDABackend::clear(CUDANet::Backend::Tensor &input) {
-    CUDA_CHECK(cudaMemset(input.data<float>(), 0, sizeof(float) * input.numel()));
+void CUDA::zero(CUDANet::Tensor &input) {
+    fill(input, 0);
 }
 
-void CUDABackend::sum(const CUDANet::Backend::Tensor &input, CUDANet::Backend::Tensor &sum) {
+void CUDA::fill(CUDANet::Tensor &input, int value) {
+    switch (input.get_dtype()) {
+    case DType::FLOAT32:
+        fill_impl<float>(input, value);
+        break;
+
+    default:
+        throw std::runtime_error("Unsupported dtype");
+        break;
+    }
+}
+
+template void CUDA::fill_impl<float>(CUDANet::Tensor &input, int value);
+
+template <typename T>
+void CUDA::fill_impl(CUDANet::Tensor &input, int value) {
+    CUDA_CHECK(cudaMemset(static_cast<T*>(input.device_ptr()), value, sizeof(T) * input.numel()));
+}
+
+void CUDA::copy_to_device(CUDANet::Tensor &tensor, void *data, size_t size) {
+    switch (tensor.get_dtype()) {
+    case DType::FLOAT32:
+        copy_to_device_impl<float>(tensor, data, size);
+        break;
+
+    default:
+        throw std::runtime_error("Unsupported dtype");
+        break;
+    }
+}
+
+template void CUDA::copy_to_device_impl<float>(CUDANet::Tensor &tensor, void *data, size_t size);
+
+template <typename T>
+void CUDA::copy_to_device_impl(CUDANet::Tensor &tensor, void *data, size_t size) {
+    CUDA_CHECK(cudaMemcpy(static_cast<T*>(tensor.device_ptr()), data, size, cudaMemcpyHostToDevice));
+}
+
+void CUDA::sum(const CUDANet::Tensor &input, CUDANet::Tensor &sum) {
+    switch (input.get_dtype()) {
+    case DType::FLOAT32:
+        sum_impl<float>(input, sum);
+        break;
+
+    default:
+        throw std::runtime_error("Unsupported dtype");
+        break;
+    }
+}
+
+template void CUDA::sum_impl<float>(const CUDANet::Tensor &input, CUDANet::Tensor &sum);
+
+template <typename T>
+void CUDA::sum_impl(const CUDANet::Tensor &input, CUDANet::Tensor &sum) {
     auto length = input.numel();
-    const int gridSize = ( + BLOCK_SIZE - 1) / BLOCK_SIZE;
+    const int gridSize = (length + BLOCK_SIZE - 1) / BLOCK_SIZE;
 
     CUDANet::Kernels::sum_reduce<<<gridSize, BLOCK_SIZE>>>(
-        input.data<float>(), sum.data<float>(), length
+        static_cast<const T*>(input.device_ptr()), static_cast<T*>(sum.device_ptr()), length
     );
     CUDA_CHECK(cudaGetLastError());
 
     int remaining = gridSize;
     while (remaining > 1) {
         int blocks_needed = (remaining + BLOCK_SIZE - 1) / BLOCK_SIZE;
-        CUDANet::Kernels::sum_reduce<<<blocks_needed, BLOCK_SIZE>>>(sum.data<float>(), sum.data<float>(), remaining);
+        CUDANet::Kernels::sum_reduce<<<blocks_needed, BLOCK_SIZE>>>(static_cast<T*>(sum.device_ptr()), static_cast<T*>(sum.device_ptr()), remaining);
         CUDA_CHECK(cudaGetLastError());
 
         remaining = blocks_needed;
     }
 }
 
-void CUDABackend::max(const CUDANet::Backend::Tensor &input, CUDANet::Backend::Tensor &max) {
+void CUDA::max(const CUDANet::Tensor &input, CUDANet::Tensor &max) {
+    switch (input.get_dtype()) {
+    case DType::FLOAT32:
+        max_impl<float>(input, max);
+        break;
+
+    default:
+        throw std::runtime_error("Unsupported dtype");
+        break;
+    }
+}
+
+template void CUDA::max_impl<float>(const CUDANet::Tensor &input, CUDANet::Tensor &max);
+
+template <typename T>
+void CUDA::max_impl(const CUDANet::Tensor &input, CUDANet::Tensor &max) {
     auto length = input.numel();
     const int grid_size = (length + BLOCK_SIZE - 1) / BLOCK_SIZE;
 
-    Kernels::max_reduce<<<grid_size, BLOCK_SIZE>>>(input.data<float>(), max.data<float>(), length);
+    Kernels::max_reduce<<<grid_size, BLOCK_SIZE>>>(static_cast<const T*>(input.device_ptr()), static_cast<T*>(max.device_ptr()), length);
     CUDA_CHECK(cudaGetLastError());
 
     int remaining = grid_size;
 
     while (remaining > 1) {
         int blocks_needed = (remaining + BLOCK_SIZE - 1) / BLOCK_SIZE;
-        CUDANet::Kernels::max_reduce<<<blocks_needed, BLOCK_SIZE>>>(max.data<float>(), max.data<float>(), remaining);
+        CUDANet::Kernels::max_reduce<<<blocks_needed, BLOCK_SIZE>>>(static_cast<T*>(max.device_ptr()), static_cast<T*>(max.device_ptr()), remaining);
         CUDA_CHECK(cudaGetLastError());
 
         remaining = blocks_needed;

src/backends/tensor.cpp

@@ -1,52 +0,0 @@
-#include "backend/tensor.hpp"
-
-#include <stdexcept>
-
-using namespace CUDANet::Backend;
-
-Tensor::Tensor(Shape shape, DType dtype, IBackend* backend)
-    : shape(shape), dtype(dtype), backend(backend), d_ptr(nullptr) {
-    // Count total elements
-    size_t count = 1;
-    for (const auto& dim : shape) {
-        count *= dim;
-    }
-    total_elms = count;
-
-    // Compute total size (bytes)
-    size_t type_size = 0;
-    switch (dtype) {
-        case DType::FLOAT32:
-            type_size = 4;
-            break;
-        default:
-            throw std::runtime_error("Unsupported data type");
-    }
-    total_size = total_elms * type_size;
-
-    // Allocate memory on backend
-    d_ptr = backend->allocate(total_size);
-}
-
-Tensor::~Tensor() {
-    backend->deallocate(d_ptr);
-    d_ptr = nullptr;
-}
-
-size_t Tensor::numel() const {
-    return total_elms;
-}
-
-size_t Tensor::size() const {
-    return total_size;
-}
-
-template <typename T>
-const T* Tensor::data() const {
-    return static_cast<T*>(d_ptr);
-}
-
-template <typename T>
-T* Tensor::data() {
-    return static_cast<T*>(d_ptr);
-}

src/layers/activation.cpp

@@ -1,34 +1,87 @@
+#include <format>
 #include <stdexcept>
 #include <vector>
 
-#include "activation.hpp"
-#include "backend/tensor.hpp"
+#include "layers/activation.hpp"
+#include "tensor.hpp"
+
 
 using namespace CUDANet::Layers;
 
-Activation::Activation(CUDANet::Backend::IBackend* backend, ActivationType activation, const int length)
-    : backend(backend), activationType(activation), length(length) {
+Activation::Activation(
+    ActivationType        activation,
+    const CUDANet::Shape& shape,
+    CUDANet::Backend*     backend
+)
+    : Activation(activation, shape, backend->get_default_dtype(), backend) {}
 
+Activation::Activation(
+    ActivationType        activation,
+    const CUDANet::Shape& shape,
+    CUDANet::DType        dtype,
+    CUDANet::Backend*     backend
+)
+    : activation_type(activation),
+      shape(shape),
+      backend(backend) {
+    this->dtype = dtype;
     
-    if (activationType == SOFTMAX) {
-      softmax_sum = CUDANet::Backend::Tensor({static_cast<size_t>(length)}, CUDANet::Backend::DType::FLOAT32, backend);
-      tensor_max = CUDANet::Backend::Tensor({static_cast<size_t>(length)}, CUDANet::Backend::DType::FLOAT32, backend);
+    if (shape.size() != 1) {
+        throw InvalidShapeException("input", 1, shape.size());
+    }
+
+    auto length = shape[0];
+
+    if (activation_type == SOFTMAX) {
+        softmax_sum =
+            CUDANet::Tensor({static_cast<size_t>(length)}, dtype, backend);
+        tensor_max =
+            CUDANet::Tensor({static_cast<size_t>(length)}, dtype, backend);
     }
 }
 
-void Activation::activate(CUDANet::Backend::Tensor input) {
-    switch (activationType)
-    {
-    case ActivationType::SIGMOID:
-        backend->sigmoid(input);
-        break;
-    case ActivationType::RELU:
-        backend->relu(input);
-        break;
-    case ActivationType::SOFTMAX:
-        backend->softmax(input, tensor_max, softmax_sum);
-        break;
-    default:
-        break;
+CUDANet::Tensor& Activation::forward(CUDANet::Tensor& input) {
+    switch (activation_type) {
+        case ActivationType::SIGMOID:
+            backend->sigmoid(input);
+            break;
+        case ActivationType::RELU:
+            backend->relu(input);
+            break;
+        case ActivationType::SOFTMAX:
+            backend->softmax(input, tensor_max, softmax_sum);
+            break;
+        default:
+            break;
     }
+
+    return input;
+}
+
+CUDANet::Shape Activation::input_shape() {
+    return shape;
+}
+
+CUDANet::Shape Activation::output_shape() {
+    return shape;
+}
+
+size_t Activation::input_size() {
+    return shape[0];
+}
+
+size_t Activation::output_size() {
+    return shape[0];
+}
+
+void Activation::set_weights(void* input) {}
+
+size_t Activation::get_weights_size() {
+    return 0;
+}
+
+void Activation::set_biases(void* input) {}
+
+size_t Activation::get_biases_size() {
+    return 0;
 }

src/layers/add.cpp

@@ -1,44 +1,32 @@
-#include "add.hpp"
-
-#include <stddef.h>
+#include "layers/add.hpp"
 
 using namespace CUDANet::Layers;
 
 
-Add::Add(int inputSize)
-    : inputSize(inputSize) {
+Add::Add(CUDANet::Shape a_shape, CUDANet::Shape b_shape, CUDANet::Backend* backend)
+    : Add(a_shape, b_shape, backend->get_default_dtype(), backend) {}
 
-    output = new float[inputSize];
-
-#ifdef USE_CUDA
-    initCUDA();
-#endif
-    
-}
-
-
-Add::~Add() {
-#ifdef USE_CUDA
-    delCUDA();
-#endif
-}
-
-
-float* Add::forward(const float* inputA, const float* inputB) {
-
-#ifdef USE_CUDA
-    return forwardCUDA(inputA, inputB);
-#else
-    return forwardCPU(inputA, inputB);
-#endif
-
-}
-
-float* Add::forwardCPU(const float* inputA, const float* inputB) {
-    for (size_t i = 0; i < inputSize; i++)
-    {
-        output[i] = inputA[i] + inputB[i];
+Add::Add(CUDANet::Shape a_shape, CUDANet::Shape b_shape, CUDANet::DType dtype, CUDANet::Backend* backend)
+    : backend(backend), dtype(dtype) {
+    if (a_shape != b_shape) {
+        throw InvalidShapeException(
+            "Add requires matching dimensions", a_shape, b_shape
+        );
     }
 
+    out_shape = a_shape;
+    output = CUDANet::Tensor(out_shape, dtype, backend);
+}
+
+Add::~Add() {}
+
+CUDANet::Tensor&
+Add::forward(CUDANet::Tensor& input_a, CUDANet::Tensor& input_b) {
+    output.zero();
+    backend->add(
+        input_a,
+        input_b,
+        output
+    );
     return output;
 }

src/layers/avg_pooling.cpp

@@ -1,101 +1,146 @@
+#include <format>
 #include <stdexcept>
 
-#include "avg_pooling.hpp"
+#include "layers/avg_pool.hpp"
 
 using namespace CUDANet::Layers;
 
-AvgPooling2d::AvgPooling2d(
-    shape2d        inputSize,
-    int            nChannels,
-    shape2d        poolingSize,
-    shape2d        stride,
-    shape2d        padding,
-    ActivationType activationType
+AvgPool2d::AvgPool2d(
+    CUDANet::Shape    input_shape,
+    CUDANet::Shape    pool_shape,
+    CUDANet::Shape    stride_shape,
+    CUDANet::Shape    padding_shape,
+    CUDANet::Backend* backend
 )
-    : inputSize(inputSize),
-      nChannels(nChannels),
-      poolingSize(poolingSize),
-      stride(stride),
-      padding(padding) {
-    outputSize = {
-        (inputSize.first + 2 * padding.first - poolingSize.first) /
-                stride.first +
+    : AvgPool2d(input_shape, pool_shape, stride_shape, padding_shape, backend->get_default_dtype(), backend) {}
+
+AvgPool2d::AvgPool2d(
+    CUDANet::Shape    input_shape,
+    CUDANet::Shape    pool_shape,
+    CUDANet::Shape    stride_shape,
+    CUDANet::Shape    padding_shape,
+    CUDANet::DType    dtype,
+    CUDANet::Backend* backend
+)
+    : in_shape(input_shape),
+      pool_shape(pool_shape),
+      stride_shape(stride_shape),
+      padding_shape(padding_shape),
+      backend(backend) {
+    if (in_shape.size() != 3) {
+        throw InvalidShapeException("input", 3, in_shape.size());
+    }
+
+    if (pool_shape.size() != 2) {
+        throw InvalidShapeException("pool", 2, pool_shape.size());
+    }
+
+    if (stride_shape.size() != 2) {
+        throw InvalidShapeException("stride", 2, stride_shape.size());
+    }
+
+    if (padding_shape.size() != 2) {
+        throw InvalidShapeException("padding", 2, padding_shape.size());
+    }
+
+    this->dtype = dtype;
+
+    out_shape = {
+        (in_shape[0] + 2 * padding_shape[0] - pool_shape[0]) / stride_shape[0] +
             1,
-        (inputSize.second + 2 * padding.second - poolingSize.second) /
-                stride.second +
-            1
+        (in_shape[1] + 2 * padding_shape[1] - pool_shape[1]) / stride_shape[1] +
+            1,
+        in_shape[2]
     };
 
-    activation = new Activation(
-        activationType, outputSize.first * outputSize.second * nChannels
+    output = CUDANet::Tensor(
+        Shape{out_shape[0] * out_shape[1] * out_shape[2]},
+        dtype, backend
     );
-
-#ifdef USE_CUDA
-    initCUDA();
-#endif
 }
 
-AvgPooling2d::~AvgPooling2d() {
-#ifdef USE_CUDA
-    delCUDA();
-#endif
-    delete activation;
+AvgPool2d::~AvgPool2d() {}
+
+CUDANet::Tensor& AvgPool2d::forward(CUDANet::Tensor& input) {
+    output.zero();
+    backend->avg_pool2d(
+        input,
+        output,
+        in_shape,
+        pool_shape,
+        stride_shape,
+        padding_shape,
+        out_shape
+    );
+    return output;
 }
 
-float* AvgPooling2d::forwardCPU(const float* input) {
-    throw std::logic_error("Not implemented");
+CUDANet::Shape AvgPool2d::input_shape() {
+    return in_shape;
 }
 
-float* AvgPooling2d::forward(const float* input) {
-#ifdef USE_CUDA
-    return forwardCUDA(input);
-#else
-    return forwardCPU(input);
-#endif
+CUDANet::Shape AvgPool2d::output_shape() {
+    return out_shape;
 }
 
-int AvgPooling2d::getOutputSize() {
-    return outputSize.first * outputSize.second * nChannels;
+size_t AvgPool2d::input_size() {
+    return dtype_size(dtype) * in_shape[0] * in_shape[1] * in_shape[2];
 }
 
-int AvgPooling2d::getInputSize() {
-    return inputSize.first * inputSize.second * nChannels;
+size_t AvgPool2d::output_size() {
+    return dtype_size(dtype) * out_shape[0] * out_shape[1] * out_shape[2];
 }
 
-shape2d AvgPooling2d::getOutputDims() {
-    return outputSize;
+void AvgPool2d::set_weights(void* input) {}
+
+size_t AvgPool2d::get_weights_size() {
+    return 0;
 }
 
-AdaptiveAvgPooling2d::AdaptiveAvgPooling2d(
-    shape2d        inputShape,
-    int            nChannels,
-    shape2d        outputShape,
-    ActivationType activationType
+void AvgPool2d::set_biases(void* input) {}
+
+size_t AvgPool2d::get_biases_size() {
+    return 0;
+}
+
+
+AdaptiveAvgPool2d::AdaptiveAvgPool2d(
+    CUDANet::Shape        input_shape,
+    CUDANet::Shape        output_shape,
+    CUDANet::Backend *backend
 )
-    : AvgPooling2d(
-          inputShape,
-          nChannels,
-          {1, 1},
-          {1, 1},
-          {0, 0},
-          activationType
-      ) {
-    stride = {
-        inputShape.first / outputShape.first,
-        inputShape.second / outputShape.second
-    };
-    poolingSize = {
-        inputShape.first - (outputShape.first - 1) * stride.first,
-        inputShape.second - (outputShape.second - 1) * stride.second
-    };
-    padding    = {(poolingSize.first - 1) / 2, (poolingSize.second - 1) / 2};
-    outputSize = outputShape;
+    : AdaptiveAvgPool2d(input_shape, output_shape, backend->get_default_dtype(), backend) {}
 
-    activation = new Activation(
-        activationType, outputSize.first * outputSize.second * nChannels
+AdaptiveAvgPool2d::AdaptiveAvgPool2d(
+    CUDANet::Shape        input_shape,
+    CUDANet::Shape        output_shape,
+    CUDANet::DType        dtype,
+    CUDANet::Backend *backend
+)
+    : AvgPool2d(
+        input_shape,
+        // pool_shape
+        {
+            input_shape[0] - (output_shape[0] - 1) * (input_shape[0] / output_shape[0]),
+            input_shape[1] - (output_shape[1] - 1) * (input_shape[1] / output_shape[1])
+        },
+        // stride_shape
+        {
+            input_shape[0] / output_shape[0],
+            input_shape[1] / output_shape[1]
+        },
+        // padding_shape
+        {
+            (input_shape[0] - (output_shape[0] - 1) * (input_shape[0] / output_shape[0]) - 1) / 2,
+            (input_shape[1] - (output_shape[1] - 1) * (input_shape[1] / output_shape[1]) - 1) / 2
+        },
+        dtype,
+        backend
+    ) {
+    out_shape = output_shape;
+
+    output = CUDANet::Tensor(
+        Shape{out_shape[0] * out_shape[1] * out_shape[2]},
+        dtype, backend
     );
-
-#ifdef USE_CUDA
-    initCUDA();
-#endif
 }

src/layers/batch_norm.cpp

@@ -1,133 +1,111 @@
-#include "batch_norm.hpp"
-
 #include <stdexcept>
 #include <vector>
 
-#include "activation.hpp"
+#include "layers/batch_norm.hpp"
 #include "layer.hpp"
 
 using namespace CUDANet::Layers;
 
 BatchNorm2d::BatchNorm2d(
-    shape2d        inputSize,
-    int            inputChannels,
-    float          epsilon,
-    ActivationType activationType
+    CUDANet::Shape input_shape,
+    float          eps,
+    CUDANet::Backend *backend
 )
-    : inputSize(inputSize), inputChannels(inputChannels), epsilon(epsilon) {
-    activation = new Activation(
-        activationType, inputSize.first * inputSize.second * inputChannels
+    : BatchNorm2d(input_shape, eps, backend->get_default_dtype(), backend) {}
+
+BatchNorm2d::BatchNorm2d(
+    CUDANet::Shape input_shape,
+    float          eps,
+    CUDANet::DType dtype,
+    CUDANet::Backend *backend
+)
+    : in_shape(input_shape), backend(backend)  {
+
+    if (in_shape.size() != 3) {
+        throw InvalidShapeException("input", 3, in_shape.size());
+    }
+
+    this->dtype = dtype;
+
+    epsilon = CUDANet::Tensor({1}, dtype, backend);
+    epsilon.set_data(&eps);
+
+    running_mean = CUDANet::Tensor({in_shape[2]}, dtype, backend);
+    running_mean.zero();
+
+    running_var = CUDANet::Tensor({in_shape[2]}, dtype, backend);
+    running_var.fill(1);
+
+    weights = CUDANet::Tensor({in_shape[2]}, dtype, backend);
+    weights.fill(1);
+
+    biases = CUDANet::Tensor({in_shape[2]}, dtype, backend);
+    biases.zero();
+
+    output = CUDANet::Tensor(in_shape, dtype, backend);
+}
+
+BatchNorm2d::~BatchNorm2d() {}
+
+CUDANet::Tensor& BatchNorm2d::forward(CUDANet::Tensor& input) {
+    output.zero();
+    backend->batch_norm(
+        input,
+        output,
+        in_shape,
+        weights,
+        biases,
+        running_mean,
+        running_var,
+        epsilon
     );
-
-    weights.resize(inputChannels);
-    biases.resize(inputChannels);
-
-    running_mean.resize(inputChannels);
-    running_var.resize(inputChannels);
-
-    initializeWeights();
-    initializeBiases();
-    initializeRunningMean();
-    initializeRunningVar();
-
-#ifdef USE_CUDA
-    initCUDA();
-    toCuda();
-#endif
+    return output;
 }
 
-BatchNorm2d::~BatchNorm2d() {
-#ifdef USE_CUDA
-    delCUDA();
-#endif
+CUDANet::Shape BatchNorm2d::input_shape() {
+    return in_shape;
 }
 
-void BatchNorm2d::initializeWeights() {
-    std::fill(weights.begin(), weights.end(), 1.0f);
+CUDANet::Shape BatchNorm2d::output_shape() {
+    return in_shape;
 }
 
-void BatchNorm2d::initializeBiases() {
-    std::fill(biases.begin(), biases.end(), 0.0f);
+size_t BatchNorm2d::input_size() {
+    return dtype_size(dtype) * in_shape[0] * in_shape[1] * in_shape[2];
 }
 
-void BatchNorm2d::initializeRunningMean() {
-    std::fill(running_mean.begin(), running_mean.end(), 0.0f);
+size_t BatchNorm2d::output_size() {
+    return dtype_size(dtype) * in_shape[0] * in_shape[1] * in_shape[2];
 }
 
-void BatchNorm2d::initializeRunningVar() {
-    std::fill(running_var.begin(), running_var.end(), 1.0f);
+void BatchNorm2d::set_weights(void* input) {
+    weights.set_data(input);
 }
 
-void BatchNorm2d::setWeights(const float* weights_input) {
-    std::copy(weights_input, weights_input + weights.size(), weights.begin());
-#ifdef USE_CUDA
-    toCuda();
-#endif
+size_t BatchNorm2d::get_weights_size() {
+    return weights.size();
 }
 
-std::vector<float> BatchNorm2d::getWeights() {
-    return weights;
+void BatchNorm2d::set_biases(void* input) {
+    biases.set_data(input);
 }
 
-void BatchNorm2d::setBiases(const float* biases_input) {
-    std::copy(biases_input, biases_input + biases.size(), biases.begin());
-#ifdef USE_CUDA
-    toCuda();
-#endif
+size_t BatchNorm2d::get_biases_size() {
+    return biases.size();
 }
 
-std::vector<float> BatchNorm2d::getBiases() {
-    return biases;
+void BatchNorm2d::set_running_mean(void* input) {
+    running_mean.set_data(input);
 }
 
-void BatchNorm2d::setRunningMean(const float* running_mean_input) {
-    std::copy(
-        running_mean_input, running_mean_input + inputChannels,
-        running_mean.begin()
-    );
-#ifdef USE_CUDA
-    toCuda();
-#endif
+size_t BatchNorm2d::get_running_mean_size() {
+    return running_mean.size();
 }
 
-std::vector<float> BatchNorm2d::getRunningMean() {
-    return running_mean;
+void BatchNorm2d::set_running_var(void* input) {
+    running_var.set_data(input);
 }
 
-void BatchNorm2d::setRunningVar(const float* running_var_input) {
-    std::copy(
-        running_var_input, running_var_input + inputChannels,
-        running_var.begin()
-    );
-#ifdef USE_CUDA
-    toCuda();
-#endif
-}
-
-std::vector<float> BatchNorm2d::getRunningVar() {
-    return running_var;
-}
-
-int BatchNorm2d::getInputSize() {
-    return inputSize.first * inputSize.second * inputChannels;
-}
-
-int BatchNorm2d::getOutputSize() {
-    return inputSize.first * inputSize.second * inputChannels;
-}
-
-shape2d BatchNorm2d::getOutputDims() {
-    return inputSize;
-}
-
-float* BatchNorm2d::forwardCPU(const float* input) {
-    throw std::logic_error("Not implemented");
-}
-
-float* BatchNorm2d::forward(const float* input) {
-#ifdef USE_CUDA
-    return forwardCUDA(input);
-#else
-    return forwardCPU(input);
-#endif
+size_t BatchNorm2d::get_running_var_size() {
+    return running_var.size();
 }

src/layers/concat.cpp

@@ -1,34 +1,35 @@
-#include <stdexcept>
-
-#include "concat.hpp"
+#include "layers/concat.hpp"
 
 using namespace CUDANet::Layers;
 
-Concat::Concat(const int inputASize, const int inputBSize)
-    : inputASize(inputASize), inputBSize(inputBSize) {
-#ifdef USE_CUDA
-    initCUDA();
-#endif
+Concat::Concat(const CUDANet::Shape a_shape, const CUDANet::Shape b_shape, CUDANet::Backend *backend)
+    : Concat(a_shape, b_shape, backend->get_default_dtype(), backend) {}
+
+Concat::Concat(const CUDANet::Shape a_shape, const CUDANet::Shape b_shape, CUDANet::DType dtype, CUDANet::Backend *backend)
+    : a_shape(a_shape), b_shape(b_shape), backend(backend), dtype(dtype) {
+    if (a_shape[0] != b_shape[0] || a_shape[1] != b_shape[1]) {
+        throw InvalidShapeException(
+            "Concat requires matching height and width dimensions", a_shape,
+            b_shape
+        );
+    }
+
+    out_shape = {a_shape[0], a_shape[1], a_shape[2] + b_shape[2]};
+    output = CUDANet::Tensor(out_shape, dtype, backend);
 }
 
-Concat::~Concat() {
-#ifdef USE_CUDA
-    delCUDA();
-#endif
+Concat::~Concat() {}
+
+CUDANet::Tensor& Concat::forward(CUDANet::Tensor& input_a, CUDANet::Tensor& input_b) {
+    output.zero();
+    backend->concat(
+        input_a,
+        input_b,
+        output
+    );
+    return output;
 }
 
-float* Concat::forwardCPU(const float* input_A, const float* input_B) {
-    throw std::logic_error("Not implemented");
+CUDANet::Shape Concat::output_shape() {
+    return out_shape;
 }
-
-float* Concat::forward(const float* input_A, const float* input_B) {
-#ifdef USE_CUDA
-    return forwardCUDA(input_A, input_B);
-#else
-    return forwardCPU(input_A, input_B);
-#endif
-}
-
-int Concat::getOutputSize() {
-    return inputASize + inputBSize;
-};

src/layers/conv2d.cpp

@@ -1,111 +1,124 @@
+#include <format>
 #include <stdexcept>
-#include <vector>
 
-#include "activation.hpp"
-#include "conv2d.hpp"
+#include "layers/conv2d.hpp"
 #include "layer.hpp"
+#include "tensor.hpp"
 
 using namespace CUDANet::Layers;
 
 Conv2d::Conv2d(
-    shape2d        inputSize,
-    int            inputChannels,
-    shape2d        kernelSize,
-    shape2d        stride,
-    int            numFilters,
-    shape2d        paddingSize,
-    ActivationType activationType
+    CUDANet::Shape    input_shape,
+    CUDANet::Shape    kernel_shape,
+    CUDANet::Shape    stride_shape,
+    CUDANet::Shape    padding_shape,
+    CUDANet::Backend* backend
 )
-    : inputSize(inputSize),
-      inputChannels(inputChannels),
-      kernelSize(kernelSize),
-      stride(stride),
-      numFilters(numFilters),
-      paddingSize(paddingSize) {
-    outputSize = {
-        (inputSize.first - kernelSize.first + 2 * paddingSize.first) /
-                stride.first +
+    : Conv2d(input_shape, kernel_shape, stride_shape, padding_shape, backend->get_default_dtype(), backend) {}
+
+Conv2d::Conv2d(
+    CUDANet::Shape    input_shape,
+    CUDANet::Shape    kernel_shape,
+    CUDANet::Shape    stride_shape,
+    CUDANet::Shape    padding_shape,
+    CUDANet::DType    dtype,
+    CUDANet::Backend* backend
+)
+    : in_shape(input_shape),
+      kernel_shape(kernel_shape),
+      stride_shape(stride_shape),
+      padding_shape(padding_shape),
+      backend(backend) {
+    if (in_shape.size() != 3) {
+        throw InvalidShapeException("input", 3, in_shape.size());
+    }
+
+    if (kernel_shape.size() != 3) {
+        throw InvalidShapeException("kernel", 3, kernel_shape.size());
+    }
+
+    if (stride_shape.size() != 2) {
+        throw InvalidShapeException("stride", 3, stride_shape.size());
+    }
+
+    if (padding_shape.size() != 2) {
+        throw InvalidShapeException("padding", 3, padding_shape.size());
+    }
+
+    this->dtype = dtype;
+
+    out_shape = {
+        (in_shape[0] - kernel_shape[0] + 2 * padding_shape[0]) /
+                stride_shape[0] +
             1,
-        (inputSize.second - kernelSize.second + 2 * paddingSize.second) /
-                stride.second +
-            1
+        (in_shape[1] - kernel_shape[1] + 2 * padding_shape[1]) /
+                stride_shape[1] +
+            1,
+        kernel_shape[2]
     };
 
-    activation = new Activation(
-        activationType, outputSize.first * outputSize.second * numFilters
+    output = CUDANet::Tensor(
+        Shape{out_shape[0], out_shape[1], out_shape[2]},
+        dtype, backend
     );
 
-    weights.resize(
-        kernelSize.first * kernelSize.second * inputChannels * numFilters
+    weights = CUDANet::Tensor(
+        Shape{
+            kernel_shape[0], kernel_shape[1], kernel_shape[2], in_shape[2]
+        },
+        dtype, backend
+    );
+    biases = CUDANet::Tensor(
+        Shape{kernel_shape[2]}, dtype, backend
     );
-    initializeWeights();
 
-    biases.resize(numFilters);
-    initializeBiases();
-
-#ifdef USE_CUDA
-    initCUDA();
-    toCuda();
-#endif
+    weights.zero();
+    biases.zero();
 }
 
-Conv2d::~Conv2d() {
-#ifdef USE_CUDA
-    delCUDA();
-#endif
-    delete activation;
+Conv2d::~Conv2d() {}
+
+CUDANet::Tensor& Conv2d::forward(CUDANet::Tensor& input) {
+    output.zero();
+    backend->conv2d(
+        weights, biases, input, output, in_shape, padding_shape, kernel_shape,
+        stride_shape, out_shape
+    );
+    return output;
 }
 
-void Conv2d::initializeWeights() {
-    std::fill(weights.begin(), weights.end(), 0.0f);
+CUDANet::Shape Conv2d::input_shape() {
+    return in_shape;
 }
 
-void Conv2d::initializeBiases() {
-    std::fill(biases.begin(), biases.end(), 0.0f);
+CUDANet::Shape Conv2d::output_shape() {
+    return out_shape;
 }
 
-void Conv2d::setWeights(const float* weights_input) {
-    std::copy(weights_input, weights_input + weights.size(), weights.begin());
-#ifdef USE_CUDA
-    toCuda();
-#endif
+size_t Conv2d::input_size() {
+    return dtype_size(dtype) * in_shape[0] * in_shape[1] * in_shape[2];
 }
 
-std::vector<float> Conv2d::getWeights() {
-    return weights;
+size_t Conv2d::output_size() {
+    return dtype_size(dtype) * out_shape[0] * out_shape[1] * out_shape[2];
 }
 
-void Conv2d::setBiases(const float* biases_input) {
-    std::copy(biases_input, biases_input + biases.size(), biases.begin());
-#ifdef USE_CUDA
-    toCuda();
-#endif
+void Conv2d::set_weights(void* input) {
+    weights.set_data(input);
 }
 
-std::vector<float> Conv2d::getBiases() {
-    return biases;
+size_t Conv2d::get_weights_size() {
+    return weights.size();
 }
 
-float* Conv2d::forwardCPU(const float* input) {
-    throw std::logic_error("Not implemented");
+void Conv2d::set_biases(void* input) {
+    biases.set_data(input);
 }
 
-float* Conv2d::forward(const float* input) {
-#ifdef USE_CUDA
-    return forwardCUDA(input);
-#else
-    return forwardCPU(input);
-#endif
+size_t Conv2d::get_biases_size() {
+    return biases.size();
 }
 
-int Conv2d::getOutputSize() {
-    return outputSize.first * outputSize.second * numFilters;
-}
-
-int Conv2d::getInputSize() {
-    return inputSize.first * inputSize.second * inputChannels;
-}
-
-shape2d Conv2d::getOutputDims() {
-    return outputSize;
+CUDANet::Shape Conv2d::get_padding_shape() {
+    return padding_shape;
 }

src/layers/dense.cpp

@@ -1,80 +1,74 @@
-#include "dense.hpp"
-
+#include <format>
 #include <stdexcept>
 
-#include "activation.hpp"
+#include "layers/dense.hpp"
 
 using namespace CUDANet::Layers;
 
-Dense::Dense(int inputSize, int outputSize, ActivationType activationType)
-    : inputSize(inputSize), outputSize(outputSize) {
-    // Allocate memory for weights and biases
-    weights.resize(outputSize * inputSize);
-    biases.resize(outputSize);
+Dense::Dense(CUDANet::Shape in_shape, CUDANet::Shape out_shape, CUDANet::Backend* backend)
+    : Dense(in_shape, out_shape, backend->get_default_dtype(), backend) {}
 
-    initializeWeights();
-    initializeBiases();
+Dense::Dense(CUDANet::Shape in_shape, CUDANet::Shape out_shape, CUDANet::DType dtype, CUDANet::Backend* backend)
+    : backend(backend),
+      in_shape(in_shape),
+      out_shape(out_shape) {
 
-    activation = new Activation(activationType, outputSize);
+    if (in_shape.size() != 1) {
+        throw InvalidShapeException("input", 1, in_shape.size());
+    }
 
-#ifdef USE_CUDA
-    initCUDA();
-#endif
+    if (out_shape.size() != 1) {
+        throw InvalidShapeException("output", 1, out_shape.size());
+    }
+
+    this->dtype = dtype;
+
+    weights = CUDANet::Tensor(Shape{out_shape[0], in_shape[0]}, dtype, backend);
+    biases = CUDANet::Tensor(Shape{out_shape[0]}, dtype, backend);
+    output = CUDANet::Tensor(Shape{out_shape[0]}, dtype, backend);
+
+    weights.zero();
+    biases.zero();
+    output.zero();
 }
 
-Dense::~Dense() {
-    delete activation;
-#ifdef USE_CUDA
-    delCUDA();
-#endif
+Dense::~Dense() {}
+
+CUDANet::Tensor& Dense::forward(CUDANet::Tensor& input) {
+    output.zero();
+    backend->dense(weights, biases, input, output, in_shape[0], out_shape[0]);
+    return output;
 }
 
-void Dense::initializeWeights() {
-    std::fill(weights.begin(), weights.end(), 0.0f);
+CUDANet::Shape Dense::input_shape() {
+    return in_shape;
 }
 
-void Dense::initializeBiases() {
-    std::fill(biases.begin(), biases.end(), 0.0f);
+CUDANet::Shape Dense::output_shape() {
+    return out_shape;
 }
 
-float* Dense::forwardCPU(const float* input) {
-    throw std::logic_error("Not implemented");
+size_t Dense::input_size() {
+    return in_shape[0];
+};
+
+size_t Dense::output_size() {
+    return out_shape[0];
+};
+
+// TODO: Use dtype
+void Dense::set_weights(void* input) {
+    weights.set_data(input);
 }
 
-float* Dense::forward(const float* input) {
-#ifdef USE_CUDA
-    return forwardCUDA(input);
-#else
-    return forwardCPU(input);
-#endif
+size_t Dense::get_weights_size() {
+    return weights.size();
 }
 
-void Dense::setWeights(const float* weights_input) {
-    std::copy(weights_input, weights_input + weights.size(), weights.begin());
-#ifdef USE_CUDA
-    toCuda();
-#endif
+void Dense::set_biases(void* input) {
+    biases.set_data(input);
 }
 
-std::vector<float> Dense::getWeights() {
-    return weights;
-}
-
-void Dense::setBiases(const float* biases_input) {
-    std::copy(biases_input, biases_input + biases.size(), biases.begin());
-#ifdef USE_CUDA
-    toCuda();
-#endif
-}
-
-std::vector<float> Dense::getBiases() {
-    return biases;
-}
-
-int Dense::getOutputSize() {
-    return outputSize;
-}
-
-int Dense::getInputSize() {
-    return inputSize;
+size_t Dense::get_biases_size() {
+    return biases.size();
 }

src/layers/input.cpp

@@ -1,37 +0,0 @@
-#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;
-}

src/layers/max_pool.cpp

@@ -0,0 +1,98 @@
+#include <stdexcept>
+
+#include "layers/max_pool.hpp"
+
+using namespace CUDANet::Layers;
+
+MaxPool2d::MaxPool2d(
+    CUDANet::Shape    input_shape,
+    CUDANet::Shape    pool_shape,
+    CUDANet::Shape    stride_shape,
+    CUDANet::Shape    padding_shape,
+    CUDANet::Backend* backend
+)
+    : MaxPool2d(input_shape, pool_shape, stride_shape, padding_shape, backend->get_default_dtype(), backend) {}
+
+MaxPool2d::MaxPool2d(
+    CUDANet::Shape    input_shape,
+    CUDANet::Shape    pool_shape,
+    CUDANet::Shape    stride_shape,
+    CUDANet::Shape    padding_shape,
+    CUDANet::DType    dtype,
+    CUDANet::Backend* backend
+)
+    : in_shape(input_shape),
+      pool_shape(pool_shape),
+      stride_shape(stride_shape),
+      padding_shape(padding_shape),
+      backend(backend) {
+    if (in_shape.size() != 3) {
+        throw InvalidShapeException("input", 3, in_shape.size());
+    }
+
+    if (pool_shape.size() != 2) {
+        throw InvalidShapeException("pool", 2, pool_shape.size());
+    }
+
+    if (stride_shape.size() != 2) {
+        throw InvalidShapeException("stride", 2, stride_shape.size());
+    }
+
+    if (padding_shape.size() != 2) {
+        throw InvalidShapeException("padding", 2, padding_shape.size());
+    }
+
+    this->dtype = dtype;
+
+    out_shape = {
+        (in_shape[0] + 2 * padding_shape[0] - pool_shape[0]) / stride_shape[0] +
+            1,
+        (in_shape[1] + 2 * padding_shape[1] - pool_shape[1]) / stride_shape[1] +
+            1,
+        in_shape[2]
+    };
+
+    output = CUDANet::Tensor(
+        Shape{out_shape[0] * out_shape[1] * out_shape[2]},
+        dtype, backend
+    );
+}
+
+MaxPool2d::~MaxPool2d() {}
+
+CUDANet::Tensor& MaxPool2d::forward(CUDANet::Tensor& input) {
+    output.zero();
+    backend->max_pool2d(
+        input, output, in_shape, pool_shape, stride_shape, padding_shape,
+        out_shape
+    );
+    return output;
+}
+
+CUDANet::Shape MaxPool2d::input_shape() {
+    return in_shape;
+}
+
+CUDANet::Shape MaxPool2d::output_shape() {
+    return out_shape;
+}
+
+size_t MaxPool2d::input_size() {
+    return dtype_size(dtype) * in_shape[0] * in_shape[1] * in_shape[2];
+}
+
+size_t MaxPool2d::output_size() {
+    return dtype_size(dtype) * out_shape[0] * out_shape[1] * out_shape[2];
+}
+
+void MaxPool2d::set_weights(void* input) {}
+
+size_t MaxPool2d::get_weights_size() {
+    return 0;
+}
+
+void MaxPool2d::set_biases(void* input) {}
+
+size_t MaxPool2d::get_biases_size() {
+    return 0;
+}

src/layers/max_pooling.cpp

@@ -1,67 +0,0 @@
-#include "max_pooling.hpp"
-#include <stdexcept>
-
-using namespace CUDANet::Layers;
-
-MaxPooling2d::MaxPooling2d(
-    shape2d          inputSize,
-    int            nChannels,
-    shape2d          poolingSize,
-    shape2d          stride,
-    shape2d          padding,
-    ActivationType activationType
-)
-    : inputSize(inputSize),
-      nChannels(nChannels),
-      poolingSize(poolingSize),
-      stride(stride),
-      padding(padding) {
-    outputSize = {
-        (inputSize.first + 2 * padding.first - poolingSize.first) /
-                stride.first +
-            1,
-        (inputSize.second + 2 * padding.second - poolingSize.second) /
-                stride.second +
-            1
-    };
-
-    activation = new Activation(
-        activationType, outputSize.first * outputSize.second * nChannels
-    );
-
-    #ifdef USE_CUDA
-    initCUDA();
-#endif
-}
-
-MaxPooling2d::~MaxPooling2d() {
-#ifdef USE_CUDA
-    delCUDA();
-#endif
-    delete activation;
-}
-
-float* MaxPooling2d::forwardCPU(const float* input) {
-    throw std::logic_error("Not implemented");
-}
-
-float* MaxPooling2d::forward(const float* input) {
-#ifdef USE_CUDA
-    return forwardCUDA(input);
-#else
-    return forwardCPU(input);
-#endif
-}
-
-
-int MaxPooling2d::getOutputSize() {
-    return outputSize.first * outputSize.second * nChannels;
-}
-
-int MaxPooling2d::getInputSize() {
-    return inputSize.first * inputSize.second * nChannels;
-}
-
-shape2d MaxPooling2d::getOutputDims() {
-    return outputSize;
-}

src/layers/output.cu

@@ -1,34 +0,0 @@
-#include "output.hpp"
-#include <stdexcept>
-
-using namespace CUDANet::Layers;
-
-
-Output::Output(int inputSize) : inputSize(inputSize) {
-    h_output = (float*) malloc(sizeof(float) * inputSize);
-}
-
-Output::~Output() {
-    free(h_output);
-}
-
-float* Output::forwardCPU(const float* input) {
-    throw std::logic_error("Not implemented");
-}
-
-float* Output::forward(const float* input) {
-#ifdef USE_CUDA
-    return forwardCUDA(input);
-#else
-    return forwardCPU(input);
-#endif
-}
-
-int Output::getOutputSize() {
-    return inputSize;
-}
-
-
-int Output::getInputSize() {
-    return inputSize;
-}

src/model.cpp

@@ -0,0 +1,244 @@
+#include <fstream>
+#include <iostream>
+#include <iomanip>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "layer.hpp"
+#include "layers/batch_norm.hpp"
+
+#include "model.hpp"
+
+using namespace CUDANet;
+
+Model::Model(
+    const CUDANet::Shape input_shape,
+    const CUDANet::Shape output_shape
+)
+    : in_shape(input_shape),
+      out_shape(out_shape),
+      layers(std::vector<std::pair<std::string, Layer*>>()),
+      layer_map(std::unordered_map<std::string, Layer*>()) {};
+
+Model::~Model() {};
+
+CUDANet::Tensor& Model::predict(CUDANet::Tensor& input) {
+    CUDANet::Tensor* current = &input;
+    for (const auto& [name, layer_ptr] : layers) {
+        current = &(layer_ptr->forward(*current));
+    }
+    return *current;
+}
+
+void Model::register_layer(const std::string& name, Layer* layer) {
+    layers.push_back({name, layer});
+    layer_map[name] = layer;
+}
+
+void Model::register_module(Module& module) {
+    for (const auto& [name, layer_ptr] : module.get_layers()) {
+        layer_map[name] = layer_ptr;
+        layers.push_back({name, layer_ptr});
+    }
+
+    return;
+}
+
+Layer* Model::get_layer(const std::string& name) {
+    return layer_map[name];
+}
+
+void Model::load_weights(const std::string& path) {
+    std::ifstream file(path, std::ios::binary);
+
+    if (!file.is_open()) {
+        std::cerr << "Failed to open file: " << path << std::endl;
+        return;
+    }
+
+    u_short version;
+    file.read(reinterpret_cast<char*>(&version), sizeof(version));
+
+    if (version != 1) {
+        std::cerr << "Unsupported model version: " << version << std::endl;
+        return;
+    }
+
+    auto get_tensor_type = [](const std::string& type_str) {
+        if (type_str == "weight") return TensorType::WEIGHT;
+        if (type_str == "bias") return TensorType::BIAS;
+        if (type_str == "running_mean") return TensorType::RUNNING_MEAN;
+        if (type_str == "running_var") return TensorType::RUNNING_VAR;
+        throw std::runtime_error("Unknown tensor type: " + type_str);
+    };
+
+    u_int64_t header_size;
+    file.read(reinterpret_cast<char*>(&header_size), sizeof(header_size));
+
+    std::string header(header_size, '\0');
+    file.read(&header[0], header_size);
+
+    std::vector<TensorInfo> tensor_infos;
+    size_t                  pos = 0;
+
+    while (pos < header.size()) {
+        size_t next_pos = header.find('\n', pos);
+        if (next_pos == std::string::npos) break;
+
+        std::string line = header.substr(pos, next_pos - pos);
+        pos              = next_pos + 1;
+
+        size_t comma_pos = line.find(',');
+        if (comma_pos == std::string::npos) continue;
+
+        // Parse tensor name into name and type
+        std::string name_str = line.substr(0, comma_pos);
+        size_t      dot_pos  = name_str.find_last_of('.');
+        if (dot_pos == std::string::npos) continue;
+        std::string name = name_str.substr(0, dot_pos);
+
+        TensorType  type = get_tensor_type(name_str.substr(dot_pos + 1));
+
+        line = line.substr(comma_pos + 1);
+
+        comma_pos = line.find(',');
+        if (comma_pos == std::string::npos) continue;
+
+        int size   = std::stoi(line.substr(0, comma_pos));
+        int offset = std::stoi(line.substr(comma_pos + 1));
+
+        tensor_infos.push_back({name, type, size, offset});
+    }
+
+    for (const auto& tensor_info : tensor_infos) {
+        std::vector<float> values(tensor_info.size);
+
+        file.seekg(
+            sizeof(version) + sizeof(header_size) + header.size() +
+            tensor_info.offset
+        );
+        file.read(
+            reinterpret_cast<char*>(values.data()),
+            tensor_info.size * sizeof(float)
+        );
+
+        if (layer_map.find(tensor_info.name) != layer_map.end()) {
+
+            Layer* layer = layer_map[tensor_info.name];
+
+            if (tensor_info.type == TensorType::WEIGHT) {
+                if (layer->get_weights_size() != values.size()) {
+                    std::cerr << "Layer: " << tensor_info.name
+                              << " has incorrect number of weights, expected "
+                              << layer->get_weights_size() << " but got "
+                              << values.size() << ", skipping" << std::endl;
+                    continue;
+                }
+
+                layer->set_weights(values.data());
+            } else if (tensor_info.type == TensorType::BIAS) {
+                if (layer->get_biases_size() != values.size()) {
+                    std::cerr << "Layer: " << tensor_info.name
+                              << " has incorrect number of biases, expected "
+                              << layer->get_biases_size() << " but got "
+                              << values.size() << ", skipping" << std::endl;
+                    continue;
+                }
+
+                layer->set_biases(values.data());
+            }
+
+            Layers::BatchNorm2d* bn_layer = dynamic_cast<Layers::BatchNorm2d*>(layer);
+            if (bn_layer == nullptr) {
+                continue;
+            }
+
+            if (tensor_info.type == TensorType::RUNNING_MEAN) {
+                if (bn_layer->get_running_mean_size() != values.size()) {
+                    std::cerr << "Layer: " << tensor_info.name << " has incorrect number of running mean values, expected "
+                                << bn_layer->get_running_mean_size() << " but got " << values.size() << ", skipping" << std::endl;
+                    continue;
+                }
+                bn_layer->set_running_mean(values.data());
+            } else if (tensor_info.type == TensorType::RUNNING_VAR) {
+                if (bn_layer->get_running_var_size() != values.size()) {
+                    std::cerr << "Layer: " << tensor_info.name << " has incorrect number of running var values, expected "
+                                << bn_layer->get_running_var_size() << " but got " << values.size() << ", skipping" << std::endl;
+                    continue;
+                }
+                bn_layer->set_running_var(values.data());
+            }
+
+
+        } else {
+            std::cerr << "Layer: " << tensor_info.name
+                      << " does not exist, skipping" << std::endl;
+        }
+    }
+
+    file.close();
+}
+
+bool Model::validate() {
+    bool valid = true;
+    CUDANet::Shape shape = in_shape;
+
+    for (const auto& [name, layer_ptr] : layers) {
+        if (layer_ptr->input_shape() != shape) {
+            valid = false;
+            std::cerr << "Layer: " << name
+                      << " has incorrect input shape, expected " << format_shape(shape)
+                      << " but got " << format_shape(layer_ptr->input_shape())
+                      << std::endl;
+            break;
+        }
+
+        shape = layer_ptr->output_shape();
+    }
+
+    return valid;
+}
+
+void Model::print_summary() {
+    struct layer_info {
+        std::string name;
+        std::string input_shape;
+        std::string output_shape;
+    };
+
+    std::vector<layer_info> layer_infos;
+
+    int max_name_length   = 0;
+    int max_input_length  = 0;
+    int max_output_length = 0;
+
+    for (const auto& [name, layer_ptr] : layers) {
+        layer_info li = {
+            name, format_shape(layer_ptr->input_shape()),
+            format_shape(layer_ptr->output_shape())
+        };
+        layer_infos.push_back(li);
+
+        max_name_length = std::max(max_name_length, (int)li.name.size());
+        max_input_length =
+            std::max(max_input_length, (int)li.input_shape.size());
+        max_output_length =
+            std::max(max_output_length, (int)li.output_shape.size());
+    }
+
+    int row_length = max_name_length + max_input_length + max_output_length + 6;
+
+    std::cout << "Model Summary:" << std::endl              
+              << std::string(row_length, '-') << std::endl;
+
+    for (const auto& li : layer_infos) {
+        std::cout << std::left
+                  << std::setw(max_name_length) << li.name
+                  << " | " << std::right
+                  << std::setw(max_input_length) << li.input_shape
+                  << " | "
+                  << std::setw(max_output_length) << li.output_shape
+                  << std::endl;
+    }
+}

src/model/model.cpp

@@ -1,277 +0,0 @@
-#include "model.hpp"
-
-#include <fstream>
-#include <iostream>
-#include <iomanip>
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-#include "input.hpp"
-#include "layer.hpp"
-#include "batch_norm.hpp"
-
-using namespace CUDANet;
-
-Model::Model(
-    const shape2d inputSize,
-    const int     inputChannels,
-    const int     outputSize
-)
-    : inputSize(inputSize),
-      inputChannels(inputChannels),
-      outputSize(outputSize),
-      layers(std::vector<std::pair<std::string, Layers::SequentialLayer*>>()),
-      layerMap(std::unordered_map<std::string, Layers::SequentialLayer*>()) {
-    inputLayer =
-        new Layers::Input(inputSize.first * inputSize.second * inputChannels);
-    outputLayer = new Layers::Output(outputSize);
-};
-
-Model::Model(const Model& other)
-    : inputSize(other.inputSize),
-      inputChannels(other.inputChannels),
-      outputSize(other.outputSize),
-      layers(std::vector<std::pair<std::string, Layers::SequentialLayer*>>()),
-      layerMap(std::unordered_map<std::string, Layers::SequentialLayer*>()) {
-    inputLayer  = new Layers::Input(*other.inputLayer);
-    outputLayer = new Layers::Output(*other.outputLayer);
-}
-
-Model::~Model() {
-    delete inputLayer;
-    delete outputLayer;
-    for (const auto& layer : layers) {
-        delete layer.second;
-    }
-};
-
-float* Model::predict(const float* input) {
-    float* d_input = inputLayer->forward(input);
-
-    for (auto& layer : layers) {
-        d_input = layer.second->forward(d_input);
-    }
-
-    return outputLayer->forward(d_input);
-}
-
-void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
-    const Module* module = dynamic_cast<Module*>(layer);
-
-    if (module != nullptr) {
-        for (const auto& moduleLayer : module->getLayers()) {
-            layerMap[moduleLayer.first] = moduleLayer.second;
-            layers.push_back({moduleLayer.first, moduleLayer.second});
-        }
-
-        return;
-    }
-
-    layers.push_back({name, layer});
-    layerMap[name] = layer;
-}
-
-Layers::SequentialLayer* Model::getLayer(const std::string& name) {
-    return layerMap[name];
-}
-
-void Model::loadWeights(const std::string& path) {
-    std::ifstream file(path, std::ios::binary);
-
-    if (!file.is_open()) {
-        std::cerr << "Failed to open file: " << path << std::endl;
-        return;
-    }
-
-    u_short version;
-    file.read(reinterpret_cast<char*>(&version), sizeof(version));
-
-    if (version != 1) {
-        std::cerr << "Unsupported model version: " << version << std::endl;
-        return;
-    }
-
-    auto getTensorType = [](const std::string& typeStr) {
-        if (typeStr == "weight") return TensorType::WEIGHT;
-        if (typeStr == "bias") return TensorType::BIAS;
-        if (typeStr == "running_mean") return TensorType::RUNNING_MEAN;
-        if (typeStr == "running_var") return TensorType::RUNNING_VAR;
-        throw std::runtime_error("Unknown tensor type: " + typeStr);
-    };
-
-    u_int64_t headerSize;
-    file.read(reinterpret_cast<char*>(&headerSize), sizeof(headerSize));
-
-    std::string header(headerSize, '\0');
-    file.read(&header[0], headerSize);
-
-    std::vector<TensorInfo> tensorInfos;
-    size_t                  pos = 0;
-
-    while (pos < header.size()) {
-        size_t nextPos = header.find('\n', pos);
-        if (nextPos == std::string::npos) break;
-
-        std::string line = header.substr(pos, nextPos - pos);
-        pos              = nextPos + 1;
-
-        size_t commaPos = line.find(',');
-        if (commaPos == std::string::npos) continue;
-
-        // Parse tensor name into name and type
-        std::string nameStr = line.substr(0, commaPos);
-        size_t      dotPos  = nameStr.find_last_of('.');
-        if (dotPos == std::string::npos) continue;
-        std::string name = nameStr.substr(0, dotPos);
-
-        TensorType  type = getTensorType(nameStr.substr(dotPos + 1));        
-
-        line = line.substr(commaPos + 1);
-
-        commaPos = line.find(',');
-        if (commaPos == std::string::npos) continue;
-
-        int size   = std::stoi(line.substr(0, commaPos));
-        int offset = std::stoi(line.substr(commaPos + 1));
-
-        tensorInfos.push_back({name, type, size, offset});
-    }
-
-    for (const auto& tensorInfo : tensorInfos) {
-        std::vector<float> values(tensorInfo.size);
-
-        file.seekg(
-            sizeof(version) + sizeof(headerSize) + header.size() +
-            tensorInfo.offset
-        );
-        file.read(
-            reinterpret_cast<char*>(values.data()),
-            tensorInfo.size * sizeof(float)
-        );
-
-        if (layerMap.find(tensorInfo.name) != layerMap.end()) {
-            Layers::WeightedLayer* wLayer =
-                dynamic_cast<Layers::WeightedLayer*>(layerMap[tensorInfo.name]);
-
-            if (wLayer == nullptr) {
-                std::cerr << "Layer: " << tensorInfo.name
-                          << " does not have weights" << std::endl;
-                continue;
-            }
-
-            if (tensorInfo.type == TensorType::WEIGHT) {
-                if (wLayer->getWeights().size() != values.size()) {
-                    std::cerr << "Layer: " << tensorInfo.name
-                              << " has incorrect number of weights, expected "
-                              << wLayer->getWeights().size() << " but got "
-                              << values.size() << ", skipping" << std::endl;
-                    continue;
-                }
-
-                wLayer->setWeights(values.data());
-            } else if (tensorInfo.type == TensorType::BIAS) {
-                if (wLayer->getBiases().size() != values.size()) {
-                    std::cerr << "Layer: " << tensorInfo.name
-                              << " has incorrect number of biases, expected "
-                              << wLayer->getBiases().size() << " but got "
-                              << values.size() << ", skipping" << std::endl;
-                    continue;
-                }
-
-                wLayer->setBiases(values.data());
-            }
-
-            Layers::BatchNorm2d* bnLayer = dynamic_cast<Layers::BatchNorm2d*>(wLayer);
-            if (bnLayer == nullptr) {
-                continue;
-            }
-
-            if (tensorInfo.type == TensorType::RUNNING_MEAN) {
-                if (bnLayer->getRunningMean().size() != values.size()) {
-                    std::cerr << "Layer: " << tensorInfo.name << " has incorrect number of running mean values, expected "
-                                << bnLayer->getRunningMean().size() << " but got " << values.size() << ", skipping" << std::endl;
-                    continue;
-                }
-                bnLayer->setRunningMean(values.data());
-            } else if (tensorInfo.type == TensorType::RUNNING_VAR) {
-                if (bnLayer->getRunningVar().size() != values.size()) {
-                    std::cerr << "Layer: " << tensorInfo.name << " has incorrect number of running var values, expected "
-                                << bnLayer->getRunningVar().size() << " but got " << values.size() << ", skipping" << std::endl;
-                    continue;
-                }
-                bnLayer->setRunningVar(values.data());
-            }
-
-
-        } else {
-            std::cerr << "Layer: " << tensorInfo.name
-                      << " does not exist, skipping" << std::endl;
-        }
-    }
-
-    file.close();
-}
-
-bool Model::validate() {
-    bool valid = true;
-    int  size  = inputLayer->getInputSize();
-
-    for (const auto& layer : layers) {
-        if (layer.second->getInputSize() != size) {
-            valid = false;
-            std::cerr << "Layer: " << layer.first
-                      << " has incorrect input size, expected " << size
-                      << " but got " << layer.second->getInputSize()
-                      << std::endl;
-            break;
-        }
-
-        size = layer.second->getOutputSize();
-    }
-
-    return valid;
-}
-
-void Model::printSummary() {
-    struct layer_info {
-        std::string name;
-        std::string inputSize;
-        std::string outputSize;
-    };
-
-    std::vector<layer_info> layerInfos;
-
-    int maxNameLength   = 0;
-    int maxInputLength  = 0;
-    int maxOutputLength = 0;
-
-    for (const auto& layer : layers) {
-        layer_info layerInfo = {
-            layer.first, std::to_string(layer.second->getInputSize()),
-            std::to_string(layer.second->getOutputSize())
-        };
-        layerInfos.push_back(layerInfo);
-
-        maxNameLength = std::max(maxNameLength, (int)layerInfo.name.size());
-        maxInputLength =
-            std::max(maxInputLength, (int)layerInfo.inputSize.size());
-        maxOutputLength =
-            std::max(maxOutputLength, (int)layerInfo.outputSize.size());
-    }
-
-    int rowLength = maxNameLength + maxInputLength + maxOutputLength + 6;
-
-    std::cout << "Model Summary:" << std::endl              
-              << std::string(rowLength, '-') << std::endl;
-
-    for (const auto& layerInfo : layerInfos) {
-        std::cout << std::left
-                  << std::setw(maxNameLength) << layerInfo.name
-                  << " | " << std::right
-                  << std::setw(maxInputLength) << layerInfo.inputSize
-                  << " | "
-                  << std::setw(maxOutputLength) << layerInfo.outputSize
-                  << std::endl;
-    }
-}

src/model/module.cpp

@@ -1,32 +0,0 @@
-#include "module.hpp"
-
-#include <algorithm>
-
-using namespace CUDANet;
-
-void Module::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
-    const Module* module = dynamic_cast<Module*>(layer);
-
-    if (module != nullptr) {
-        for (const auto& moduleLayer : module->getLayers()) {
-            layers.push_back({moduleLayer.first, moduleLayer.second});
-        }
-
-        return;
-    }
-
-    layers.push_back({name, layer});
-}
-
-const std::vector<std::pair<std::string, Layers::SequentialLayer*>>&
-Module::getLayers() const {
-    return layers;
-}
-
-int Module::getInputSize() {
-    return inputSize;
-}
-
-int Module::getOutputSize() {
-    return outputSize;
-}

src/module.cpp

@@ -0,0 +1,28 @@
+#include <algorithm>
+
+#include "module.hpp"
+
+using namespace CUDANet;
+
+CUDANet::Shape Module::input_shape() {
+    return in_shape;
+}
+
+CUDANet::Shape Module::output_shape() {
+    return out_shape;
+}
+
+void Module::register_layer(const std::string& name, Layer* layer) {
+    layers.push_back({name, layer});
+}
+
+void Module::register_module(Module& module) {
+    for (const auto& moduleLayer : module.get_layers()) {
+        layers.push_back({moduleLayer.first, moduleLayer.second});
+    }
+}
+
+const std::vector<std::pair<std::string, Layer*>>&
+Module::get_layers() const {
+    return layers;
+}

src/tensor.cpp

@@ -0,0 +1,126 @@
+#include <stdexcept>
+
+#include "tensor.hpp"
+
+using namespace CUDANet;
+
+size_t dtype_size(DType dtype) {
+    switch (dtype)
+    {
+    case DType::FLOAT32:
+        return 4;
+        break;
+    
+    default:
+        throw std::runtime_error("Unknown DType");
+        break;
+    }
+}
+
+Tensor::Tensor(Shape shape, CUDANet::Backend* backend)
+    : Tensor(shape, backend->get_default_dtype(), backend) {}
+
+Tensor::Tensor(Shape shape, DType dtype, Backend* backend)
+    : shape(shape), dtype(dtype), backend(backend), d_ptr(nullptr) {
+    if (shape.empty()) {
+        throw std::runtime_error("Tensor shape cannot be empty");
+    }
+
+    // Check if backend supports DType
+    if (!backend->supports_dtype(dtype)) {
+        throw std::runtime_error("Unsupported DType");
+    }
+
+    // Count total elements
+    size_t count = 1;
+    for (size_t i = 0; i < shape.size(); ++i) {
+        count *= shape[i];
+    }
+    total_elms = count;
+
+    // Compute total size (bytes)
+    size_t type_size = 0;
+    switch (dtype) {
+        case DType::FLOAT32:
+            type_size = 4;
+            break;
+        default:
+            throw std::runtime_error("Unsupported data type");
+    }
+    total_size = total_elms * type_size;
+
+    // Allocate memory on backend
+    d_ptr = backend->allocate(total_size);
+}
+
+Tensor::Tensor(Tensor&& other) noexcept
+    : shape(std::move(other.shape)),
+      dtype(other.dtype),
+      total_elms(other.total_elms),
+      total_size(other.total_size),
+      backend(other.backend),
+      d_ptr(other.d_ptr) {
+    other.d_ptr   = nullptr;
+    other.backend = nullptr;
+}
+
+Tensor& Tensor::operator=(Tensor&& other) noexcept {
+    if (this != &other) {
+        // Clean up our current resources
+        if (d_ptr != nullptr && backend != nullptr) {
+            backend->deallocate(d_ptr);
+        }
+
+        // Steal other's resources
+        shape      = std::move(other.shape);
+        dtype      = other.dtype;
+        total_elms = other.total_elms;
+        total_size = other.total_size;
+        backend    = other.backend;
+        d_ptr      = other.d_ptr;
+
+        // Leave other in valid but empty state
+        other.d_ptr   = nullptr;
+        other.backend = nullptr;
+    }
+    return *this;
+}
+
+Tensor::~Tensor() {
+    if (backend && d_ptr) {
+        backend->deallocate(d_ptr);
+        d_ptr = nullptr;
+    }
+}
+
+DType Tensor::get_dtype() const {
+    return dtype;
+}
+
+size_t Tensor::numel() const {
+    return total_elms;
+}
+
+size_t Tensor::size() const {
+    return total_size;
+}
+
+void* Tensor::device_ptr() const {
+    return d_ptr;
+}
+
+void* Tensor::device_ptr() {
+    return d_ptr;
+}
+
+void Tensor::zero() {
+    backend->zero(*this);
+}
+
+void Tensor::fill(int value) {
+    backend->fill(*this, value);
+}
+
+void Tensor::set_data(void *data) {
+    backend->copy_to_device(*this, data, total_size);
+}

test/cuda/layers/test_avg_pooling.cu

@@ -54,7 +54,7 @@ class AvgPoolingLayerTest : public ::testing::Test {
 
         d_output = avgPoolingLayer->forward(d_input);
 
-        int outputSize = avgPoolingLayer->getOutputSize();
+        int outputSize = avgPoolingLayer->get_output_size();
 
         std::vector<float> output(outputSize);
         cudaStatus = cudaMemcpy(
@@ -229,7 +229,7 @@ class AdaptiveAvgPoolingLayerTest : public ::testing::Test {
 
         d_output = adaptiveAvgPoolingLayer->forward(d_input);
 
-        int outputSize = adaptiveAvgPoolingLayer->getOutputSize();
+        int outputSize = adaptiveAvgPoolingLayer->get_output_size();
 
         std::vector<float> output(outputSize);
         cudaStatus = cudaMemcpy(

test/cuda/layers/test_max_pooling.cu

@@ -52,7 +52,7 @@ class MaxPoolingLayerTest : public ::testing::Test {
 
         d_output = maxPoolingLayer->forward(d_input);
 
-        int outputSize = maxPoolingLayer->getOutputSize();
+        int outputSize = maxPoolingLayer->get_output_size();
 
         std::vector<float> output(outputSize);
         cudaStatus = cudaMemcpy(