feat(convolution - ops): add img2col operation and related test cases

- Added the img2col operation to the convolution operations suite.
- Implemented a set of test cases to verify the correctness of the img2col operation.
- These test cases cover various input scenarios to ensure the stability and accuracy of the img2col implementation.

Author: Mgepahmge

include/NeuZephyr/OperationKernels.cuh

@@ -1000,11 +1000,15 @@ namespace nz::krnl {
     void NgradCopy(dim3 gridDim, dim3 blockDim, float* out, float* in, size_t n,
                    const std::vector<size_t>& offset_o, const std::vector<size_t>& offset_i);
 
-    void Expand(dim3 gridDim, dim3 blockDim, float* out, const float* in, size_t n,
+    void Expand(dim3 gridDim, dim3 blockDim, float* out, float* in, size_t n,
                 size_t total);
 
-    void Compress(dim3 gridDim, dim3 blockDim, float* out, const float* in, size_t n,
+    void Compress(dim3 gridDim, dim3 blockDim, float* out, float* in, size_t n,
                   size_t total);
+
+    void img2col(const dim3 gridDim, const dim3 blockDim, float* out, float* in, const size_t H_out,
+              const size_t W_out, const size_t C, const size_t K_h, const size_t K_w, const size_t stride,
+              const size_t pad, const size_t H_in, const size_t W_in, const size_t batch);
 #endif
 }
 

include/NeuZephyr/TensorOperations.cuh

@@ -798,8 +798,8 @@ namespace nz::data {
             for (auto j = 0; j < out.shape()[1]; j++) {
                 offsetC.push_back(i * out.shape().getStride(0) + j * out.shape().getStride(1));
                 offsetA.push_back(i * (lhs.shape().N() > 1 ? lhs.shape().getStride(0) : 0) + j * (lhs.shape().C() > 1
-                        ? lhs.shape().getStride(1)
-                        : 0));
+                    ? lhs.shape().getStride(1)
+                    : 0));
                 offsetB.push_back(i * (rhs.shape().N() > 1 ? rhs.shape().getStride(0) : 0) + j * (
                     rhs.shape().C() > 1 ? rhs.shape().getStride(1) : 0));
             }
@@ -869,8 +869,8 @@ namespace nz::data {
             for (auto j = 0; j < out.shape()[1]; j++) {
                 offsetC.push_back(i * out.shape().getStride(0) + j * out.shape().getStride(1));
                 offsetA.push_back(i * (lhs.shape().N() > 1 ? lhs.shape().getStride(0) : 0) + j * (lhs.shape().C() > 1
-                        ? lhs.shape().getStride(1)
-                        : 0));
+                    ? lhs.shape().getStride(1)
+                    : 0));
                 offsetB.push_back(i * (rhs.shape().N() > 1 ? rhs.shape().getStride(0) : 0) + j * (
                     rhs.shape().C() > 1 ? rhs.shape().getStride(1) : 0));
             }
@@ -939,8 +939,8 @@ namespace nz::data {
             for (auto j = 0; j < out.shape()[1]; j++) {
                 offsetC.push_back(i * out.shape().getStride(0) + j * out.shape().getStride(1));
                 offsetA.push_back(i * (lhs.shape().N() > 1 ? lhs.shape().getStride(0) : 0) + j * (lhs.shape().C() > 1
-                        ? lhs.shape().getStride(1)
-                        : 0));
+                    ? lhs.shape().getStride(1)
+                    : 0));
                 offsetB.push_back(i * (rhs.shape().N() > 1 ? rhs.shape().getStride(0) : 0) + j * (
                     rhs.shape().C() > 1 ? rhs.shape().getStride(1) : 0));
             }
@@ -1008,8 +1008,8 @@ namespace nz::data {
             for (auto j = 0; j < out.shape()[1]; j++) {
                 offsetC.push_back(i * out.shape().getStride(0) + j * out.shape().getStride(1));
                 offsetA.push_back(i * (lhs.shape().N() > 1 ? lhs.shape().getStride(0) : 0) + j * (lhs.shape().C() > 1
-                        ? lhs.shape().getStride(1)
-                        : 0));
+                    ? lhs.shape().getStride(1)
+                    : 0));
                 offsetB.push_back(i * (rhs.shape().N() > 1 ? rhs.shape().getStride(0) : 0) + j * (
                     rhs.shape().C() > 1 ? rhs.shape().getStride(1) : 0));
             }
@@ -1085,7 +1085,8 @@ namespace nz::data {
         return result;
     }
 
-    DL_API void iSoftmaxJacobian(float* out, float* in, size_t n, const std::vector<size_t>& offset_o, const std::vector<size_t>& offset_i);
+    DL_API void iSoftmaxJacobian(float* out, float* in, size_t n, const std::vector<size_t>& offset_o,
+                                 const std::vector<size_t>& offset_i);
 
     template <typename T>
     std::enable_if_t<is_valid_tensor_type<T>::value, T>
@@ -1103,5 +1104,21 @@ namespace nz::data {
         iSoftmaxJacobian(result.data(), in.data(), n, offset_o, offset_i);
         return result;
     }
+
+    DL_API void iImg2col(float* out, float* in, const size_t H_out,
+                         const size_t W_out, const size_t C, const size_t K_h, const size_t K_w, const size_t stride,
+                         const size_t pad, const size_t H_in, const size_t W_in, const size_t batch);
+
+    template <typename T>
+    std::enable_if_t<is_valid_tensor_type<T>::value, T>
+    tensorImg2col(const T& in, const size_t K_h, const size_t K_w, const size_t stride,
+                  const size_t pad) {
+        const size_t H_out = (in.shape().H() + 2 * pad - K_h) / stride + 1;
+        const size_t W_out = (in.shape().W() + 2 * pad - K_w) / stride + 1;
+        T result({in.shape()[0], 1, H_out * W_out, in.shape().C() * K_h * K_w});
+        iImg2col(result.data(), in.data(), H_out, W_out, in.shape().C(), K_h, K_w, stride, pad,
+                 in.shape().H(), in.shape().W(), in.shape()[0]);
+        return result;
+    }
 }
 #endif //TENSOROPERATIONS_CUH

src/OperationKernels.cu

@@ -1285,7 +1285,7 @@ namespace nz::krnl {
         out[idx] = in[idx % n];
     }
 
-    void Expand(const dim3 gridDim, const dim3 blockDim, float* out, const float* in, const size_t n,
+    void Expand(const dim3 gridDim, const dim3 blockDim, float* out, float* in, const size_t n,
                  const size_t total) {
         StreamManager<float>::Instance().submit(ExpandKernel, gridDim, blockDim, 0, out, in, n, total);
     }
@@ -1298,8 +1298,36 @@ namespace nz::krnl {
         atomicAdd(out + idx % n, in[idx]);
     }
 
-    void Compress(const dim3 gridDim, const dim3 blockDim, float* out, const float* in, const size_t n,
+    void Compress(const dim3 gridDim, const dim3 blockDim, float* out, float* in, const size_t n,
                   const size_t total) {
         StreamManager<float>::Instance().submit(CompressKernel, gridDim, blockDim, 0, out, in, n, total);
     }
+
+    __global__ void img2colKernel(float* out, const float* in, const size_t H_out, const size_t W_out, const size_t C,
+        const size_t K_h, const size_t K_w, const size_t stride, const size_t pad, const size_t H_in, const size_t W_in, const size_t batch) {
+        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+        if (idx >= H_out * W_out * C * K_h * K_w * batch) {
+            return;
+        }
+        const size_t fixedIdx = idx % (H_out * W_out * C * K_h * K_w);
+        const size_t currentBatch = idx / (H_out * W_out * C * K_h * K_w);
+        const size_t k = fixedIdx / (C * K_h * K_w);
+        const size_t m = fixedIdx % (C * K_h * K_w);
+        const size_t c = m / (K_h * K_w);
+        const long long h = (k / W_out) * stride - pad + (m % (K_h * K_w)) / K_w;
+        const long long w = (k % W_out) * stride - pad + m % K_w;
+        if (h >= 0 && h < H_in && w >= 0 && w < W_in) {
+            out[idx] = in[currentBatch * (C * H_in * W_in) + c * (H_in * W_in) + h * W_in + w];
+        }
+        else {
+            out[idx] = 0;
+        }
+    }
+
+    void img2col(const dim3 gridDim, const dim3 blockDim, float* out, float* in, const size_t H_out,
+              const size_t W_out, const size_t C, const size_t K_h, const size_t K_w, const size_t stride,
+              const size_t pad, const size_t H_in, const size_t W_in, const size_t batch) {
+        StreamManager<float>::Instance().submit(img2colKernel, gridDim, blockDim, 0, out, in, H_out, W_out, C,
+                                                K_h, K_w, stride, pad, H_in, W_in, batch);
+    }
 }

src/TensorOperations.cu

@@ -131,4 +131,12 @@ namespace nz::data {
         dim3 grid((n + block.x - 1) / block.x, (n + block.y - 1) / block.y);
         krnl::SoftmaxJacobian(grid, block, out, in, n, offset_o, offset_i);
     }
+
+    void iImg2col(float* out, float* in, const size_t H_out, const size_t W_out, const size_t C, const size_t K_h,
+        const size_t K_w, const size_t stride, const size_t pad, const size_t H_in, const size_t W_in,
+        const size_t batch) {
+        const dim3 block(BLOCKSIZE);
+        const dim3 grid((H_out * W_out * C * K_h * K_w * batch + BLOCKSIZE - 1) / BLOCKSIZE);
+        krnl::img2col(grid, block, out, in, H_out, W_out, C, K_h, K_w, stride, pad, H_in, W_in, batch);
+    }
 }

test/Test.cpp

@@ -2876,4 +2876,69 @@ TEST(Model, SGDOptimize) {
     expected.dataInject(input1Gard.begin(), input1Gard.end(), true);
 
     EXPECT_EQ(expected, *model.input1.output);
+}
+
+TEST(TensorBasic, img2colTest) {
+    const size_t n = 2;
+    const size_t c = 3;
+    const size_t h = 4;
+    const size_t w = 5;
+    const size_t k_h = 3;
+    const size_t k_w = 3;
+    const size_t stride = 1;
+    const size_t pad = 1;
+    const size_t H_out = (h + 2 * pad - k_h) / stride + 1;
+    const size_t W_out = (w + 2 * pad - k_w) / stride + 1;
+
+    std::vector<float> inputData({n*c*h*w});
+    std::vector<float> expectedData({n*H_out*W_out*k_h*k_w*c});
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_real_distribution<float> dist(0.1f, 0.9f);
+
+    for (auto& i : inputData) {
+        i = dist(gen);
+    }
+
+    for (size_t b = 0; b < n; ++b) {
+        for (size_t i = 0; i < H_out; ++i) {
+            for (size_t j = 0; j < W_out; ++j) {
+                const int h_start = static_cast<int>(i * stride) - pad;
+                const int w_start = static_cast<int>(j * stride) - pad;
+
+                for (size_t r = 0; r < k_h; ++r) {
+                    const int h_in = h_start + r;
+                    for (size_t s = 0; s < k_w; ++s) {
+                        const int w_in = w_start + s;
+                        for (size_t c_in = 0; c_in < c; ++c_in) {
+                            float val = 0.0f;
+                            if (h_in >= 0 && h_in < h && w_in >= 0 && w_in < w) {
+                                const size_t input_idx =
+                                    b * (c * h * w) +
+                                    c_in * (h * w) +
+                                    h_in * w +
+                                    w_in;
+                                val = inputData[input_idx];
+                            }
+                            const size_t expected_idx =
+                                b * (H_out * W_out * k_h * k_w * c) +
+                                (i * W_out + j) * (k_h * k_w * c) +
+                                c_in * (k_h * k_w) +
+                                r * k_w +
+                                s;
+                            expectedData[expected_idx] = val;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    Tensor input({n, c, h, w});
+    input.dataInject(inputData.begin(), inputData.end());
+    auto result = tensorImg2col(input, k_h, k_w, stride, pad);
+    Tensor expected({n, 1, H_out * W_out, k_h * k_w * c});
+    expected.dataInject(expectedData.begin(), expectedData.end());
+    EXPECT_EQ(expected, result);
 }