eltwise 4d (#4529)

src/layer/arm/eltwise_arm.cpp

@@ -60,9 +60,10 @@ int Eltwise_arm::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>&
     const Mat& bottom_blob = bottom_blobs[0];
     int w = bottom_blob.w;
     int h = bottom_blob.h;
+    int d = bottom_blob.d;
     int channels = bottom_blob.c;
     int elempack = bottom_blob.elempack;
-    int size = w * h * elempack;
+    int size = w * h * d * elempack;
 
     Mat& top_blob = top_blobs[0];
     top_blob.create_like(bottom_blob, opt.blob_allocator);
@@ -477,9 +478,10 @@ int Eltwise_arm::forward_bf16s(const std::vector<Mat>& bottom_blobs, std::vector
     const Mat& bottom_blob = bottom_blobs[0];
     int w = bottom_blob.w;
     int h = bottom_blob.h;
+    int d = bottom_blob.d;
     int channels = bottom_blob.c;
     int elempack = bottom_blob.elempack;
-    int size = w * h * elempack;
+    int size = w * h * d * elempack;
 
     Mat& top_blob = top_blobs[0];
     top_blob.create_like(bottom_blob, opt.blob_allocator);
@@ -705,7 +707,7 @@ int Eltwise_arm::forward_bf16s(const std::vector<Mat>& bottom_blobs, std::vector
         return 0;
     }
 
-    Mat top_blob_fp32(w, h, channels, (size_t)4u * elempack, elempack, opt.workspace_allocator);
+    Mat top_blob_fp32(w, h, d, channels, (size_t)4u * elempack, elempack, opt.workspace_allocator);
     if (top_blob_fp32.empty())
         return -100;
 

src/layer/arm/eltwise_arm_asimdhp.cpp

@@ -26,9 +26,10 @@ int Eltwise_arm::forward_fp16s(const std::vector<Mat>& bottom_blobs, std::vector
     const Mat& bottom_blob = bottom_blobs[0];
     int w = bottom_blob.w;
     int h = bottom_blob.h;
+    int d = bottom_blob.d;
     int channels = bottom_blob.c;
     int elempack = bottom_blob.elempack;
-    int size = w * h * elempack;
+    int size = w * h * d * elempack;
 
     Mat& top_blob = top_blobs[0];
     top_blob.create_like(bottom_blob, opt.blob_allocator);
@@ -287,7 +288,7 @@ int Eltwise_arm::forward_fp16s(const std::vector<Mat>& bottom_blobs, std::vector
         return 0;
     }
 
-    Mat top_blob_fp32(w, h, channels, (size_t)4u * elempack, elempack, opt.workspace_allocator);
+    Mat top_blob_fp32(w, h, d, channels, (size_t)4u * elempack, elempack, opt.workspace_allocator);
     if (top_blob_fp32.empty())
         return -100;
 
@@ -1097,9 +1098,10 @@ int Eltwise_arm::forward_fp16sa(const std::vector<Mat>& bottom_blobs, std::vecto
     const Mat& bottom_blob = bottom_blobs[0];
     int w = bottom_blob.w;
     int h = bottom_blob.h;
+    int d = bottom_blob.d;
     int channels = bottom_blob.c;
     int elempack = bottom_blob.elempack;
-    int size = w * h * elempack;
+    int size = w * h * d * elempack;
 
     Mat& top_blob = top_blobs[0];
     top_blob.create_like(bottom_blob, opt.blob_allocator);

src/layer/eltwise.cpp

@@ -35,8 +35,9 @@ int Eltwise::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top
     const Mat& bottom_blob = bottom_blobs[0];
     int w = bottom_blob.w;
     int h = bottom_blob.h;
+    int d = bottom_blob.d;
     int channels = bottom_blob.c;
-    int size = w * h;
+    int size = w * h * d;
 
     Mat& top_blob = top_blobs[0];
     top_blob.create_like(bottom_blob, opt.blob_allocator);

src/layer/loongarch/eltwise_loongarch.cpp

@@ -34,9 +34,10 @@ int Eltwise_loongarch::forward(const std::vector<Mat>& bottom_blobs, std::vector
     const Mat& bottom_blob = bottom_blobs[0];
     int w = bottom_blob.w;
     int h = bottom_blob.h;
+    int d = bottom_blob.d;
     int channels = bottom_blob.c;
     int elempack = bottom_blob.elempack;
-    int size = w * h * elempack;
+    int size = w * h * d * elempack;
 
     Mat& top_blob = top_blobs[0];
     top_blob.create_like(bottom_blob, opt.blob_allocator);

src/layer/mips/eltwise_mips.cpp

@@ -34,9 +34,10 @@ int Eltwise_mips::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>
     const Mat& bottom_blob = bottom_blobs[0];
     int w = bottom_blob.w;
     int h = bottom_blob.h;
+    int d = bottom_blob.d;
     int channels = bottom_blob.c;
     int elempack = bottom_blob.elempack;
-    int size = w * h * elempack;
+    int size = w * h * d * elempack;
 
     Mat& top_blob = top_blobs[0];
     top_blob.create_like(bottom_blob, opt.blob_allocator);

src/layer/vulkan/eltwise_vulkan.cpp

@@ -38,7 +38,7 @@ int Eltwise_vulkan::create_pipeline(const Option& opt)
     int elempack = 1;
     if (shape.dims == 1) elempack = opt.use_shader_pack8 && shape.w % 8 == 0 ? 8 : shape.w % 4 == 0 ? 4 : 1;
     if (shape.dims == 2) elempack = opt.use_shader_pack8 && shape.h % 8 == 0 ? 8 : shape.h % 4 == 0 ? 4 : 1;
-    if (shape.dims == 3) elempack = opt.use_shader_pack8 && shape.c % 8 == 0 ? 8 : shape.c % 4 == 0 ? 4 : 1;
+    if (shape.dims == 3 || shape.dims == 4) elempack = opt.use_shader_pack8 && shape.c % 8 == 0 ? 8 : shape.c % 4 == 0 ? 4 : 1;
 
     size_t elemsize;
     if (opt.use_fp16_storage)
@@ -57,14 +57,14 @@ int Eltwise_vulkan::create_pipeline(const Option& opt)
     Mat shape_packed;
     if (shape.dims == 1) shape_packed = Mat(shape.w / elempack, (void*)0, elemsize, elempack);
     if (shape.dims == 2) shape_packed = Mat(shape.w, shape.h / elempack, (void*)0, elemsize, elempack);
-    if (shape.dims == 3) shape_packed = Mat(shape.w, shape.h, shape.c / elempack, (void*)0, elemsize, elempack);
+    if (shape.dims == 3 || shape.dims == 4) shape_packed = Mat(shape.w, shape.h, shape.c / elempack, (void*)0, elemsize, elempack);
 
     std::vector<vk_specialization_type> specializations(2 + 5);
     specializations[0].i = op_type;
     specializations[1].i = coeffs.w == 0 ? 0 : 1;
     specializations[2 + 0].i = shape_packed.dims;
     specializations[2 + 1].i = shape_packed.w;
-    specializations[2 + 2].i = shape_packed.h;
+    specializations[2 + 2].i = shape_packed.h * shape_packed.d;
     specializations[2 + 3].i = shape_packed.c;
     specializations[2 + 4].i = shape_packed.cstep;
 
@@ -87,6 +87,12 @@ int Eltwise_vulkan::create_pipeline(const Option& opt)
         local_size_xyz.h = std::min(4, shape_packed.h);
         local_size_xyz.c = std::min(4, shape_packed.c);
     }
+    if (shape_packed.dims == 4)
+    {
+        local_size_xyz.w = std::min(4, shape_packed.w);
+        local_size_xyz.h = std::min(4, shape_packed.h * shape_packed.d);
+        local_size_xyz.c = std::min(4, shape_packed.c);
+    }
 
     // pack1
     if (shape.dims == 0 || elempack == 1)
@@ -164,7 +170,7 @@ int Eltwise_vulkan::forward(const std::vector<VkMat>& bottom_blobs, std::vector<
     std::vector<vk_constant_type> constants(5 + 2);
     constants[0].i = top_blob.dims;
     constants[1].i = top_blob.w;
-    constants[2].i = top_blob.h;
+    constants[2].i = top_blob.h * top_blob.d;
     constants[3].i = top_blob.c;
     constants[4].i = top_blob.cstep;
     constants[5].f = coeffs.w == 0 ? 1.f : coeffs[0];
@@ -186,7 +192,7 @@ int Eltwise_vulkan::forward(const std::vector<VkMat>& bottom_blobs, std::vector<
         std::vector<vk_constant_type> constants(5 + 2);
         constants[0].i = top_blob.dims;
         constants[1].i = top_blob.w;
-        constants[2].i = top_blob.h;
+        constants[2].i = top_blob.h * top_blob.d;
         constants[3].i = top_blob.c;
         constants[4].i = top_blob.cstep;
         constants[5].f = 1.f;
@@ -222,7 +228,7 @@ int Eltwise_vulkan::forward(const std::vector<VkImageMat>& bottom_blobs, std::ve
     std::vector<vk_constant_type> constants(5 + 2);
     constants[0].i = top_blob.dims;
     constants[1].i = top_blob.w;
-    constants[2].i = top_blob.h;
+    constants[2].i = top_blob.h * top_blob.d;
     constants[3].i = top_blob.c;
     constants[4].i = 0; //top_blob.cstep;
     constants[5].f = coeffs.w == 0 ? 1.f : coeffs[0];
@@ -244,7 +250,7 @@ int Eltwise_vulkan::forward(const std::vector<VkImageMat>& bottom_blobs, std::ve
         std::vector<vk_constant_type> constants(5 + 2);
         constants[0].i = top_blob.dims;
         constants[1].i = top_blob.w;
-        constants[2].i = top_blob.h;
+        constants[2].i = top_blob.h * top_blob.d;
         constants[3].i = top_blob.c;
         constants[4].i = 0; //top_blob.cstep;
         constants[5].f = 1.f;

src/layer/x86/eltwise_x86.cpp

@@ -36,9 +36,10 @@ int Eltwise_x86::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>&
     const Mat& bottom_blob = bottom_blobs[0];
     int w = bottom_blob.w;
     int h = bottom_blob.h;
+    int d = bottom_blob.d;
     int channels = bottom_blob.c;
     int elempack = bottom_blob.elempack;
-    int size = w * h * elempack;
+    int size = w * h * d * elempack;
 
     Mat& top_blob = top_blobs[0];
     top_blob.create_like(bottom_blob, opt.blob_allocator);

tests/test_eltwise.cpp

@@ -36,7 +36,7 @@ static int test_eltwise(const std::vector<ncnn::Mat>& a, int op_type, const ncnn
     int ret = test_layer<ncnn::Eltwise>("Eltwise", pd, weights, a);
     if (ret != 0)
     {
-        fprintf(stderr, "test_eltwise failed a[0].dims=%d a[0]=(%d %d %d) op_type=%d", a[0].dims, a[0].w, a[0].h, a[0].c, op_type);
+        fprintf(stderr, "test_eltwise failed a[0].dims=%d a[0]=(%d %d %d %d) op_type=%d", a[0].dims, a[0].w, a[0].h, a[0].d, a[0].c, op_type);
         fprintf(stderr, " coeffs=");
         print_float_array(coeffs);
         fprintf(stderr, "\n");
@@ -315,6 +315,45 @@ static int test_eltwise_11()
            || test_eltwise(a, 2, RandomMat(5));
 }
 
+static int test_eltwise_12()
+{
+    std::vector<ncnn::Mat> a(2);
+    a[0] = RandomMat(31, 5, 3, 12);
+    a[1] = RandomMat(31, 5, 3, 12);
+
+    std::vector<ncnn::Mat> b(3);
+    b[0] = RandomMat(32, 4, 5, 32);
+    b[1] = RandomMat(32, 4, 5, 32);
+    b[2] = RandomMat(32, 4, 5, 32);
+
+    std::vector<ncnn::Mat> c(4);
+    c[0] = RandomMat(33, 6, 7, 7);
+    c[1] = RandomMat(33, 6, 7, 7);
+    c[2] = RandomMat(33, 6, 7, 7);
+    c[3] = RandomMat(33, 6, 7, 7);
+
+    return 0
+           || test_eltwise(a, 0, ncnn::Mat())
+           || test_eltwise(a, 1, ncnn::Mat())
+           || test_eltwise(a, 2, ncnn::Mat())
+           || test_eltwise(b, 0, ncnn::Mat())
+           || test_eltwise(b, 1, ncnn::Mat())
+           || test_eltwise(b, 2, ncnn::Mat())
+           || test_eltwise(c, 0, ncnn::Mat())
+           || test_eltwise(c, 1, ncnn::Mat())
+           || test_eltwise(c, 2, ncnn::Mat())
+
+           || test_eltwise(a, 0, RandomMat(2))
+           || test_eltwise(a, 1, RandomMat(2))
+           || test_eltwise(a, 2, RandomMat(2))
+           || test_eltwise(b, 0, RandomMat(3))
+           || test_eltwise(b, 1, RandomMat(3))
+           || test_eltwise(b, 2, RandomMat(3))
+           || test_eltwise(c, 0, RandomMat(4))
+           || test_eltwise(c, 1, RandomMat(4))
+           || test_eltwise(c, 2, RandomMat(4));
+}
+
 int main()
 {
     SRAND(7767517);
@@ -331,5 +370,6 @@ int main()
            || test_eltwise_8()
            || test_eltwise_9()
            || test_eltwise_10()
-           || test_eltwise_11();
+           || test_eltwise_11()
+           || test_eltwise_12();
 }