Implement backend factory

include/backend.hpp

@@ -1,14 +1,55 @@
 #pragma once
 
 #include <cstddef>
+#include <stdexcept>
 
 #include "shape.hpp"
+#ifdef USE_CUDA
+#include "backend/cuda/cuda.cuh"
+#endif
 
 namespace CUDANet {
 
 // Forward declaration
 class Tensor;
 
+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) {
+        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);
+            cuda.initialize();
+
+            return cuda;
+
+            #else
+            throw std::runtime_error("Library was compiled without CUDA support.");
+            #endif
+
+            break;
+        
+        default:
+            break;
+        }
+
+        return nullptr;
+    }
+};
+
 class Backend {
   public:
     // Memory management

include/backend/cuda/cuda.cuh

@@ -27,7 +27,14 @@ do { \
 namespace CUDANet::Backends {
 
 class CUDA : public Backend {
+  private:
+    int device_id;
   public:
+    CUDA(const BackendConfig& config);
+
+    static bool is_cuda_available();
+    void initialize();
+
     // Memory management
     void* allocate(size_t bytes) override;
     void  deallocate(void* ptr) override;

include/cudanet.hpp

@@ -41,15 +41,15 @@
 #include "layers/concat.hpp"
 
 // ============================================================================
-// Utilities
+// Dataset Labels
 // ============================================================================
 
-#include "utils/imagenet.hpp"
+#include "datasets/imagenet.hpp"
 
 // ============================================================================
 // Backend-Specific Includes (conditionally compiled)
 // ============================================================================
 
 #ifdef USE_CUDA
-#include "backend/cuda/cuda_backend.cuh"
+#include "backend/cuda/all.cuh"
 #endif

include/shape.hpp

@@ -9,6 +9,7 @@
 #endif
 
 #include <format>
+#include <stdexcept>
 #include <vector>
 
 namespace CUDANet {

src/backends/cuda/cuda.cu

@@ -0,0 +1,52 @@
+#include <cuda_runtime.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <format>
+
+#include "backend/cuda/cuda.cuh"
+
+using namespace CUDANet::Backends;
+
+
+CUDA::CUDA(const BackendConfig& config) {
+    device_id = config.device_id < 0 ? 0 : config.device_id;
+    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);
+}
+
+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,38 +0,0 @@
-#include <cuda_runtime.h>
-
-#include <cstdio>
-#include <cstdlib>
-
-#include "backend/cuda/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::Backends;
-
-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));
-}