fix flatten vulkan fp16p image pack1to4 and pack1to8

src/layer/arm/flatten_arm.cpp

@@ -33,10 +33,8 @@ Flatten_arm::Flatten_arm()
 
 int Flatten_arm::forward(const Mat& bottom_blob, Mat& top_blob, const Option& opt) const
 {
-    if (bottom_blob.elemsize / bottom_blob.elempack == 2u)
-    {
+    if (opt.use_bf16_storage)
         return forward_bf16s(bottom_blob, top_blob, opt);
-    }
 
     int dims = bottom_blob.dims;
 

src/layer/vulkan/flatten_vulkan.cpp

@@ -193,7 +193,7 @@ int Flatten_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCompute
         if (out_elempack == 1) out_elemsize = 4u;
     }
 
-    if (dims == 2 && elempack == 1)
+    if (dims == 2 && elempack == 1 && !(opt.use_fp16_packed && !opt.use_fp16_storage && out_elempack != 1))
     {
         top_blob = bottom_blob;
         top_blob.dims = 1;

src/layer/vulkan/shader/flatten_pack1to4.comp

@@ -67,13 +67,25 @@ void main()
 
     ivec4 i4 = gx * 4 + ivec4(0, 1, 2, 3);
 
-    int size = psc(w) * psc(h);
+    ivec4 v_offset;
 
-    ivec4 z4 = i4 / size;
-    ivec4 y4 = i4 % size / psc(w);
-    ivec4 x4 = i4 % size % psc(w);
+    if (psc(dims) == 2)
+    {
+        ivec4 y4 = i4 / psc(w);
+        ivec4 x4 = i4 % psc(w);
 
-    ivec4 v_offset = z4 * psc(cstep) + y4 * psc(w) + x4;
+        v_offset = y4 * psc(w) + x4;
+    }
+    else // if (psc(dims) == 3)
+    {
+        int size = psc(w) * psc(h);
+
+        ivec4 z4 = i4 / size;
+        ivec4 y4 = i4 % size / psc(w);
+        ivec4 x4 = i4 % size % psc(w);
+
+        v_offset = z4 * psc(cstep) + y4 * psc(w) + x4;
+    }
 
     buffer_cp1to4(top_blob_data, gx, bottom_blob_data, v_offset);
 }

src/layer/vulkan/shader/flatten_pack1to8.comp

@@ -69,18 +69,35 @@ void main()
     ivec4 i4 = gx * 8 + ivec4(0, 1, 2, 3);
     ivec4 ii4 = i4 + 4;
 
-    int size = psc(w) * psc(h);
-
-    ivec4 z4 = i4 / size;
-    ivec4 y4 = i4 % size / psc(w);
-    ivec4 x4 = i4 % size % psc(w);
-
-    ivec4 zz4 = ii4 / size;
-    ivec4 yy4 = ii4 % size / psc(w);
-    ivec4 xx4 = ii4 % size % psc(w);
-
-    ivec4 v_offset = z4 * psc(cstep) + y4 * psc(w) + x4;
-    ivec4 vv_offset = zz4 * psc(cstep) + yy4 * psc(w) + xx4;
+    ivec4 v_offset;
+    ivec4 vv_offset;
+
+    if (psc(dims) == 2)
+    {
+        ivec4 y4 = i4 / psc(w);
+        ivec4 x4 = i4 % psc(w);
+
+        ivec4 yy4 = ii4 / psc(w);
+        ivec4 xx4 = ii4 % psc(w);
+
+        v_offset = y4 * psc(w) + x4;
+        vv_offset = yy4 * psc(w) + xx4;
+    }
+    else // if (psc(dims) == 3)
+    {
+        int size = psc(w) * psc(h);
+
+        ivec4 z4 = i4 / size;
+        ivec4 y4 = i4 % size / psc(w);
+        ivec4 x4 = i4 % size % psc(w);
+
+        ivec4 zz4 = ii4 / size;
+        ivec4 yy4 = ii4 % size / psc(w);
+        ivec4 xx4 = ii4 % size % psc(w);
+
+        v_offset = z4 * psc(cstep) + y4 * psc(w) + x4;
+        vv_offset = zz4 * psc(cstep) + yy4 * psc(w) + xx4;
+    }
 
     buffer_cp1to8(top_blob_data, gx, bottom_blob_data, v_offset, vv_offset);
 }

tests/test_flatten.cpp

@@ -51,6 +51,7 @@ static int test_flatten_0()
         || test_flatten(RandomMat(6, 6, 15))
         || test_flatten(RandomMat(13, 13))
         || test_flatten(RandomMat(8, 12))
+        || test_flatten(RandomMat(8, 2))
         || test_flatten(RandomMat(32))
         || test_flatten(RandomMat(17))
         ;