Make conv2d work again

src/kernels/convolution.cu

@@ -1,4 +1,5 @@
 #include "convolution.cuh"
+#include <iostream>
 
 __global__ void convolution_kernel(
     const float* d_input,
@@ -19,35 +20,26 @@ __global__ void convolution_kernel(
 
     // Get output index
     int f = tid / (outputSize * outputSize);
-    int i = (tid % (outputSize * outputSize)) / outputSize;
-    int j = (tid % (outputSize * outputSize)) % outputSize;
+    int i = tid % (outputSize * outputSize) / outputSize;
+    int j = tid % outputSize;
 
     float sum = 0.0f;
 
-    // std::cout << "f: " << f << ", i: " << i << ", j: " << j << std::endl;
-
     // Iterate over kernel and input matrix
     for (int k = 0; k < kernelSize; k++) {
         for (int l = 0; l < kernelSize; l++) {
             for (int c = 0; c < nChannels; c++) {
-                int kernelIndex =
-                    k * (kernelSize * nChannels * nFilters) +
-                    l * (nChannels * nFilters) + c * (nFilters) + f;
-                int inputIndex =
-                    (i * stride + k) * (inputSize * nChannels) +
-                    (j * stride + l) * (nChannels) + c;
-
-                // std::cout << "kernelIndex: " << kernelIndex << ", kernel
-                // value: " << kernels[kernelIndex] << ", inputIndex: " <<
-                // inputIndex << ", input value: " << input[inputIndex] <<
-                // std::endl;
+                int kernelIndex = f * kernelSize * kernelSize * nChannels +
+                                  c * kernelSize * kernelSize + k * kernelSize +
+                                  l;
+                int inputIndex = c * inputSize * inputSize +
+                                 (i * stride + k) * inputSize +
+                                 (j * stride + l);
 
                 sum += d_kernel[kernelIndex] * d_input[inputIndex];
             }
         }
     }
 
-    // std::cout << "sum: " << sum << std::endl;
-
-    d_output[i * (outputSize * nFilters) + j * (nFilters) + f] = sum;
+    d_output[tid] = sum;
 }

src/layers/conv2d.cu

@@ -1,5 +1,5 @@
-#include <string>
 #include <iostream>
+#include <string>
 
 #include "activations.cuh"
 #include "conv2d.cuh"
@@ -13,7 +13,7 @@ Layers::Conv2d::Conv2d(
     int         inputChannels,
     int         kernelSize,
     int         stride,
-    std::string padding,
+    Padding padding,
     int         numFilters,
     Activation  activation
 )
@@ -25,34 +25,43 @@ Layers::Conv2d::Conv2d(
       activation(activation) {
     // Allocate memory for kernels
 
-    if (padding == "SAME") {
+    switch (padding)
+    {
+    case SAME:
         outputSize  = inputSize;
         paddingSize = ((stride - 1) * inputSize - stride + kernelSize) / 2;
-    } else if (padding == "VALID") {
+        break;
+
+    case VALID:
         paddingSize = 0;
         outputSize  = (inputSize - kernelSize) / stride + 1;
+        break;
+    
+    default:
+        break;
     }
 
     kernels.resize(kernelSize * kernelSize * inputChannels * numFilters);
-    initializeKernels();    
+    initializeKernels();
 
     d_kernels = nullptr;
-    CUDA_CHECK(
-        cudaMalloc((void**)&d_kernels, sizeof(float) * kernelSize * kernelSize * inputChannels * numFilters)
-    );
+    CUDA_CHECK(cudaMalloc(
+        (void**)&d_kernels,
+        sizeof(float) * kernelSize * kernelSize * inputChannels * numFilters
+    ));
 
     biases.resize(outputSize * outputSize * numFilters);
     initializeBiases();
 
     d_biases = nullptr;
-    CUDA_CHECK(
-        cudaMalloc((void**)&d_biases, sizeof(float) * outputSize * outputSize * numFilters)
-    );
+    CUDA_CHECK(cudaMalloc(
+        (void**)&d_biases, sizeof(float) * outputSize * outputSize * numFilters
+    ));
 
     d_padded = nullptr;
     CUDA_CHECK(cudaMalloc(
         (void**)&d_padded, sizeof(float) * (inputSize + 2 * paddingSize) *
-                                (inputSize + 2 * paddingSize) * inputChannels
+                               (inputSize + 2 * paddingSize) * inputChannels
     ));
 
     toCuda();
@@ -79,19 +88,22 @@ void Layers::Conv2d::setKernels(const std::vector<float>& kernels_input) {
 
 void Layers::Conv2d::toCuda() {
     CUDA_CHECK(cudaMemcpy(
-        d_kernels, kernels.data(), sizeof(float) * kernelSize * kernelSize * numFilters,
+        d_kernels, kernels.data(),
+        sizeof(float) * kernelSize * kernelSize * inputChannels * numFilters,
         cudaMemcpyHostToDevice
     ));
 
     CUDA_CHECK(cudaMemcpy(
-        d_biases, biases.data(), sizeof(float) * outputSize * outputSize * numFilters,
+        d_biases, biases.data(),
+        sizeof(float) * outputSize * outputSize * numFilters,
         cudaMemcpyHostToDevice
     ));
 }
 
 void Layers::Conv2d::forward(const float* d_input, float* d_output) {
     // Pad input
-    int THREADS_PER_BLOCK =  (inputSize + 2 * paddingSize) * (inputSize + 2 * paddingSize) * inputChannels;
+    int THREADS_PER_BLOCK = (inputSize + 2 * paddingSize) *
+                            (inputSize + 2 * paddingSize) * inputChannels;
 
     pad_matrix_kernel<<<1, THREADS_PER_BLOCK>>>(
         d_input, d_padded, inputSize, inputSize, inputChannels, paddingSize
@@ -100,11 +112,14 @@ void Layers::Conv2d::forward(const float* d_input, float* d_output) {
     // Convolve
     THREADS_PER_BLOCK = outputSize * outputSize * numFilters;
     convolution_kernel<<<1, THREADS_PER_BLOCK>>>(
-        d_padded, d_kernels, d_output, inputSize + (2 * paddingSize), inputChannels, kernelSize, stride, numFilters, outputSize
+        d_padded, d_kernels, d_output, inputSize + (2 * paddingSize),
+        inputChannels, kernelSize, stride, numFilters, outputSize
     );
 
     // Add bias
-    vec_vec_add_kernel<<<1, biases.size()>>>(d_biases, d_output, d_output, biases.size());
+    vec_vec_add_kernel<<<1, biases.size()>>>(
+        d_biases, d_output, d_output, biases.size()
+    );
 
     CUDA_CHECK(cudaDeviceSynchronize());
 }
@@ -119,27 +134,35 @@ outputSize x numFilters
 */
 void Layers::Conv2d::host_conv(const float* input, float* output) {
     // Iterate over output matrix
-    for (int f = 0; f < numFilters; f++) {
-        for (int i = 0; i < outputSize; i++) {
-            for (int j = 0; j < outputSize; j++) {
-                
-                float sum = 0.0f;
+    for (int tid = 0; tid < outputSize * outputSize * numFilters; tid++)
+    {  
+        // Get output index
+        int f = tid / (outputSize * outputSize);
+        int i = tid % (outputSize * outputSize) / outputSize;
+        int j = tid % outputSize;
 
-                // Iterate over kernel and input matrix
-                for (int k = 0; k < kernelSize; k++) {
-                    for (int l = 0; l < kernelSize; l++) {
-                        for (int c = 0; c < inputChannels; c++) {
-                            
-                            int kernelIndex = k * (kernelSize * inputChannels * numFilters) + l * (inputChannels * numFilters) + c * (numFilters) + f;
-                            int inputIndex  = (i * stride + k) * (inputSize * inputChannels) + (j * stride + l) * (inputChannels) + c;
+        float sum = 0.0f;
 
-                            sum += kernels[kernelIndex] * input[inputIndex];
-                        }                      
-                    }
+        // Iterate over kernel and input matrix
+        for (int k = 0; k < kernelSize; k++) {
+            for (int l = 0; l < kernelSize; l++) {
+                for (int c = 0; c < inputChannels; c++) {
+                    int kernelIndex =
+                        f * kernelSize * kernelSize * inputChannels +
+                        c * kernelSize * kernelSize + k * kernelSize +
+                        l;
+                    int inputIndex = c * inputSize * inputSize +
+                                        (i * stride + k) * inputSize +
+                                        (j * stride + l);
+
+                    sum += kernels[kernelIndex] * input[inputIndex];
                 }
-
-                output[i * (outputSize * numFilters) + j * (numFilters) + f] = sum;
             }
         }
+
+        int outputIndex =
+            f * outputSize * outputSize + i * outputSize + j;
+
+        output[outputIndex] = sum;
     }
 }