Migrate concat layer

src/backends/cuda/layer_ops.cu

@@ -211,4 +211,24 @@ CUDANet::Tensor& CUDA::batch_norm(
         );
         CUDA_CHECK(cudaGetLastError());
     }
+}
+
+CUDANet::Tensor& CUDA::concat(
+    CUDANet::Tensor& input_a,
+    CUDANet::Tensor& input_b,
+    CUDANet::Tensor& output
+) {
+    CUDA_CHECK(cudaMemcpy(
+        output.data<float>(), input_a.data<float>(), input_a.size(),
+        cudaMemcpyDeviceToDevice
+    ));
+
+    CUDA_CHECK(cudaMemcpy(
+        output.data<float>() + input_a.numel(), input_b.data<float>(), input_b.size(),
+        cudaMemcpyDeviceToDevice
+    ));
+
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return 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/layers/concat.cpp

@@ -1,34 +1,32 @@
-#include <stdexcept>
-
 #include "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)
+    : a_shape(a_shape), b_shape(b_shape), backend(backend) {
+    if (a_shape[0] != b_shape[0] || a_shape[1] != b_shape[1]) {
+        throw InvalidShapeException(
+            "Concat requires matching batch and height dimensions", a_shape,
+            b_shape
+        );
+    }
+
+    out_shape = {a_shape[0], a_shape[1], a_shape[2] + b_shape[2]};
+    output = CUDANet::Tensor(out_shape, CUDANet::DType::FLOAT32, 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");
-}
-
-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;
-};
+CUDANet::Shape Concat::output_shape() {
+    return out_shape;
+}