instancenorm image shader

src/layer/vulkan/instancenorm_vulkan.cpp

@@ -23,6 +23,7 @@ DEFINE_LAYER_CREATOR(InstanceNorm_vulkan)
 InstanceNorm_vulkan::InstanceNorm_vulkan()
 {
     support_vulkan = true;
+    support_image_storage = true;
 
     pipeline_instancenorm_reduce_sum4_fp16_to_fp32 = 0;
     pipeline_instancenorm_reduce_sum4_fp32[0] = 0;
@@ -76,12 +77,26 @@ int InstanceNorm_vulkan::create_pipeline(const Option& opt)
     Mat workspace_shape_packed(1, 1, channels / elempack, (void*)0, elemsize, elempack);
 
     {
-        Mat local_size_xyz(16, 1, std::min(4, channels / elempack), (void*)0);
-        if (workspace_shape_packed.dims != 0)
+        Mat local_size_xyz;
+        if (opt.use_image_storage)
         {
-            local_size_xyz.w = 16;
-            local_size_xyz.h = 1;
-            local_size_xyz.c = std::min(4, workspace_shape_packed.c);
+            local_size_xyz = Mat(4, 4, std::min(4, channels / elempack), (void*)0);
+            if (workspace_shape_packed.dims != 0)
+            {
+                local_size_xyz.w = 4;
+                local_size_xyz.h = 4;
+                local_size_xyz.c = std::min(4, workspace_shape_packed.c);
+            }
+        }
+        else
+        {
+            local_size_xyz = Mat(16, 1, std::min(4, channels / elempack), (void*)0);
+            if (workspace_shape_packed.dims != 0)
+            {
+                local_size_xyz.w = 16;
+                local_size_xyz.h = 1;
+                local_size_xyz.c = std::min(4, workspace_shape_packed.c);
+            }
         }
 
         // pack1
@@ -131,10 +146,11 @@ int InstanceNorm_vulkan::create_pipeline(const Option& opt)
     }
 
     {
-        std::vector<vk_specialization_type> specializations(0 + 3);
+        std::vector<vk_specialization_type> specializations(0 + 4);
         specializations[0].i = 0;// TODO resolve workspace_shape_packed.w;
-        specializations[1].i = workspace_shape_packed.c;
-        specializations[2].i = 0;// TODO resolve workspace_shape_packed.cstep;
+        specializations[1].i = 0;// TODO resolve workspace_shape_packed.h;
+        specializations[2].i = workspace_shape_packed.c;
+        specializations[3].i = 0;// TODO resolve workspace_shape_packed.cstep;
 
         Mat local_size_xyz(std::min(64, channels / elempack), 1, 1, (void*)0);
         if (workspace_shape_packed.dims != 0)
@@ -360,12 +376,26 @@ int InstanceNorm_vulkan::upload_model(VkTransfer& cmd, const Option& opt)
     Mat gamma_data_packed;
     convert_packing(gamma_data, gamma_data_packed, elempack);
 
-    cmd.record_upload(gamma_data_packed, gamma_data_gpu, opt);
+    if (opt.use_image_storage)
+    {
+        cmd.record_upload(gamma_data_packed, gamma_data_gpu_image, opt);
+    }
+    else
+    {
+        cmd.record_upload(gamma_data_packed, gamma_data_gpu, opt);
+    }
 
     Mat beta_data_packed;
     convert_packing(beta_data, beta_data_packed, elempack);
 
-    cmd.record_upload(beta_data_packed, beta_data_gpu, opt);
+    if (opt.use_image_storage)
+    {
+        cmd.record_upload(beta_data_packed, beta_data_gpu_image, opt);
+    }
+    else
+    {
+        cmd.record_upload(beta_data_packed, beta_data_gpu, opt);
+    }
 
     return 0;
 }
@@ -395,13 +425,15 @@ int InstanceNorm_vulkan::forward_inplace(VkMat& bottom_top_blob, VkCompute& cmd,
         bindings[0] = bottom_top_blob;
         bindings[1] = sum_workspace;
 
-        std::vector<vk_constant_type> constants(6);
+        std::vector<vk_constant_type> constants(8);
         constants[0].i = bottom_top_blob.w * bottom_top_blob.h;
-        constants[1].i = bottom_top_blob.c;
-        constants[2].i = bottom_top_blob.cstep;
-        constants[3].i = sum_workspace.w;
-        constants[4].i = sum_workspace.c;
-        constants[5].i = sum_workspace.cstep;
+        constants[1].i = 1;
+        constants[2].i = bottom_top_blob.c;
+        constants[3].i = bottom_top_blob.cstep;
+        constants[4].i = sum_workspace.w;
+        constants[5].i = 1;
+        constants[6].i = sum_workspace.c;
+        constants[7].i = sum_workspace.cstep;
 
         const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_sum4_fp16_to_fp32_pack8
                                  : elempack == 4 ? pipeline_instancenorm_reduce_sum4_fp16_to_fp32_pack4
@@ -426,13 +458,15 @@ int InstanceNorm_vulkan::forward_inplace(VkMat& bottom_top_blob, VkCompute& cmd,
         bindings[0] = sum_workspace;
         bindings[1] = sum_workspace_reduced;
 
-        std::vector<vk_constant_type> constants(6);
+        std::vector<vk_constant_type> constants(8);
         constants[0].i = sum_workspace.w;
-        constants[1].i = sum_workspace.c;
-        constants[2].i = sum_workspace.cstep;
-        constants[3].i = sum_workspace_reduced.w;
-        constants[4].i = sum_workspace_reduced.c;
-        constants[5].i = sum_workspace_reduced.cstep;
+        constants[1].i = 1;
+        constants[2].i = sum_workspace.c;
+        constants[3].i = sum_workspace.cstep;
+        constants[4].i = sum_workspace_reduced.w;
+        constants[5].i = 1;
+        constants[6].i = sum_workspace_reduced.c;
+        constants[7].i = sum_workspace_reduced.cstep;
 
         const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_sum4_fp32_pack8[pb%2]
                                  : elempack == 4 ? pipeline_instancenorm_reduce_sum4_fp32_pack4[pb%2]
@@ -451,11 +485,12 @@ int InstanceNorm_vulkan::forward_inplace(VkMat& bottom_top_blob, VkCompute& cmd,
         bindings[0] = sum_workspace;
         bindings[1] = mean_workspace;
 
-        std::vector<vk_constant_type> constants(4);
+        std::vector<vk_constant_type> constants(5);
         constants[0].i = sum_workspace.w;
-        constants[1].i = sum_workspace.c;
-        constants[2].i = sum_workspace.cstep;
-        constants[3].f = size;
+        constants[1].i = 1;
+        constants[2].i = sum_workspace.c;
+        constants[3].i = sum_workspace.cstep;
+        constants[4].f = size;
 
         const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_mean_pack8
                                  : elempack == 4 ? pipeline_instancenorm_reduce_mean_pack4
@@ -516,13 +551,15 @@ int InstanceNorm_vulkan::forward_inplace(VkMat& bottom_top_blob, VkCompute& cmd,
         bindings[0] = sqsum_workspace;
         bindings[1] = sqsum_workspace_reduced;
 
-        std::vector<vk_constant_type> constants(6);
+        std::vector<vk_constant_type> constants(8);
         constants[0].i = sqsum_workspace.w;
-        constants[1].i = sqsum_workspace.c;
-        constants[2].i = sqsum_workspace.cstep;
-        constants[3].i = sqsum_workspace_reduced.w;
-        constants[4].i = sqsum_workspace_reduced.c;
-        constants[5].i = sqsum_workspace_reduced.cstep;
+        constants[1].i = 1;
+        constants[2].i = sqsum_workspace.c;
+        constants[3].i = sqsum_workspace.cstep;
+        constants[4].i = sqsum_workspace_reduced.w;
+        constants[5].i = 1;
+        constants[6].i = sqsum_workspace_reduced.c;
+        constants[7].i = sqsum_workspace_reduced.cstep;
 
         const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_sum4_fp32_pack8[pb%2]
                                  : elempack == 4 ? pipeline_instancenorm_reduce_sum4_fp32_pack4[pb%2]
@@ -541,11 +578,12 @@ int InstanceNorm_vulkan::forward_inplace(VkMat& bottom_top_blob, VkCompute& cmd,
         bindings[0] = sqsum_workspace;
         bindings[1] = var_workspace;
 
-        std::vector<vk_constant_type> constants(4);
+        std::vector<vk_constant_type> constants(5);
         constants[0].i = sqsum_workspace.w;
-        constants[1].i = sqsum_workspace.c;
-        constants[2].i = sqsum_workspace.cstep;
-        constants[3].f = size;
+        constants[1].i = 1;
+        constants[2].i = sqsum_workspace.c;
+        constants[3].i = sqsum_workspace.cstep;
+        constants[4].f = size;
 
         const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_mean_pack8
                                  : elempack == 4 ? pipeline_instancenorm_reduce_mean_pack4
@@ -598,4 +636,243 @@ int InstanceNorm_vulkan::forward_inplace(VkMat& bottom_top_blob, VkCompute& cmd,
     return 0;
 }
 
+int InstanceNorm_vulkan::forward_inplace(VkImageMat& bottom_top_blob, VkCompute& cmd, const Option& opt) const
+{
+    int w = bottom_top_blob.w;
+    int h = bottom_top_blob.h;
+    int c = bottom_top_blob.c;
+    int size = w * h;
+    size_t elemsize = bottom_top_blob.elemsize;
+    int elempack = bottom_top_blob.elempack;
+
+    // mean
+    VkImageMat mean_workspace(c, elemsize, elempack, opt.workspace_vkallocator);
+    {
+        // reduce sum
+        VkImageMat sum_workspace;
+        {
+        int reduced_w = (bottom_top_blob.w + 1) / 2;
+        int reduced_h = (bottom_top_blob.h + 1) / 2;
+        int reduced_c = bottom_top_blob.c;
+
+        sum_workspace.create(reduced_w, reduced_h, reduced_c, 4u*elempack, elempack, opt.workspace_vkallocator);
+        {
+        std::vector<VkImageMat> bindings(2);
+        bindings[0] = bottom_top_blob;
+        bindings[1] = sum_workspace;
+
+        std::vector<vk_constant_type> constants(8);
+        constants[0].i = bottom_top_blob.w;
+        constants[1].i = bottom_top_blob.h;
+        constants[2].i = bottom_top_blob.c;
+        constants[3].i = 0;//bottom_top_blob.cstep;
+        constants[4].i = sum_workspace.w;
+        constants[5].i = sum_workspace.h;
+        constants[6].i = sum_workspace.c;
+        constants[7].i = 0;//sum_workspace.cstep;
+
+        const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_sum4_fp16_to_fp32_pack8
+                                 : elempack == 4 ? pipeline_instancenorm_reduce_sum4_fp16_to_fp32_pack4
+                                 : pipeline_instancenorm_reduce_sum4_fp16_to_fp32;
+
+        cmd.record_pipeline(pipeline, bindings, constants, sum_workspace);
+        }
+        }
+
+        int pb = 0;
+        while (sum_workspace.w > 2 || sum_workspace.h > 2)
+        {
+        int reduced_w = (sum_workspace.w + 1) / 2;
+        int reduced_h = (sum_workspace.h + 1) / 2;
+        int reduced_c = sum_workspace.c;
+
+        VkImageMat sum_workspace_reduced;
+        sum_workspace_reduced.create(reduced_w, reduced_h, reduced_c, 4u*elempack, elempack, opt.workspace_vkallocator);
+
+        {
+        std::vector<VkImageMat> bindings(2);
+        bindings[0] = sum_workspace;
+        bindings[1] = sum_workspace_reduced;
+
+        std::vector<vk_constant_type> constants(8);
+        constants[0].i = sum_workspace.w;
+        constants[1].i = sum_workspace.h;
+        constants[2].i = sum_workspace.c;
+        constants[3].i = 0;//sum_workspace.cstep;
+        constants[4].i = sum_workspace_reduced.w;
+        constants[5].i = sum_workspace_reduced.h;
+        constants[6].i = sum_workspace_reduced.c;
+        constants[7].i = 0;//sum_workspace_reduced.cstep;
+
+        const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_sum4_fp32_pack8[pb%2]
+                                 : elempack == 4 ? pipeline_instancenorm_reduce_sum4_fp32_pack4[pb%2]
+                                 : pipeline_instancenorm_reduce_sum4_fp32[pb%2];
+
+        cmd.record_pipeline(pipeline, bindings, constants, sum_workspace_reduced);
+
+        pb++;
+        }
+
+        sum_workspace = sum_workspace_reduced;
+        }
+
+        {
+        std::vector<VkImageMat> bindings(2);
+        bindings[0] = sum_workspace;
+        bindings[1] = mean_workspace;
+
+        std::vector<vk_constant_type> constants(5);
+        constants[0].i = sum_workspace.w;
+        constants[1].i = sum_workspace.h;
+        constants[2].i = sum_workspace.c;
+        constants[3].i = 0;//sum_workspace.cstep;
+        constants[4].f = size;
+
+        const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_mean_pack8
+                                 : elempack == 4 ? pipeline_instancenorm_reduce_mean_pack4
+                                 : pipeline_instancenorm_reduce_mean;
+
+        cmd.record_pipeline(pipeline, bindings, constants, mean_workspace);
+        }
+    }
+
+    // var
+    VkImageMat var_workspace(c, elemsize, elempack, opt.workspace_vkallocator);
+    {
+        // sub mean and square
+        VkImageMat square_workspace;
+        square_workspace.create(w, h, c, 4u*elempack, elempack, opt.workspace_vkallocator);
+        {
+        std::vector<VkImageMat> bindings(3);
+        bindings[0] = bottom_top_blob;
+        bindings[1] = mean_workspace;
+        bindings[2] = square_workspace;
+
+        std::vector<vk_constant_type> constants(10);
+        constants[0].i = bottom_top_blob.dims;
+        constants[1].i = bottom_top_blob.w;
+        constants[2].i = bottom_top_blob.h;
+        constants[3].i = bottom_top_blob.c;
+        constants[4].i = 0;//bottom_top_blob.cstep;
+        constants[5].i = square_workspace.dims;
+        constants[6].i = square_workspace.w;
+        constants[7].i = square_workspace.h;
+        constants[8].i = square_workspace.c;
+        constants[9].i = 0;//square_workspace.cstep;
+
+        const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_sub_mean_square_pack8
+                                 : elempack == 4 ? pipeline_instancenorm_sub_mean_square_pack4
+                                 : pipeline_instancenorm_sub_mean_square;
+
+        cmd.record_pipeline(pipeline, bindings, constants, square_workspace);
+        }
+
+        // reduce square
+        VkImageMat sqsum_workspace = square_workspace;
+
+        int pb = 0;
+        while (sqsum_workspace.w > 2 || sqsum_workspace.h > 2)
+        {
+        int reduced_w = (sqsum_workspace.w + 1) / 2;
+        int reduced_h = (sqsum_workspace.h + 1) / 2;
+        int reduced_c = sqsum_workspace.c;
+
+        VkImageMat sqsum_workspace_reduced;
+        sqsum_workspace_reduced.create(reduced_w, reduced_h, reduced_c, 4u*elempack, elempack, opt.workspace_vkallocator);
+
+        {
+        std::vector<VkImageMat> bindings(2);
+        bindings[0] = sqsum_workspace;
+        bindings[1] = sqsum_workspace_reduced;
+
+        std::vector<vk_constant_type> constants(8);
+        constants[0].i = sqsum_workspace.w;
+        constants[1].i = sqsum_workspace.h;
+        constants[2].i = sqsum_workspace.c;
+        constants[3].i = 0;//sqsum_workspace.cstep;
+        constants[4].i = sqsum_workspace_reduced.w;
+        constants[5].i = sqsum_workspace_reduced.h;
+        constants[6].i = sqsum_workspace_reduced.c;
+        constants[7].i = 0;//sqsum_workspace_reduced.cstep;
+
+        const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_sum4_fp32_pack8[pb%2]
+                                 : elempack == 4 ? pipeline_instancenorm_reduce_sum4_fp32_pack4[pb%2]
+                                 : pipeline_instancenorm_reduce_sum4_fp32[pb%2];
+
+        cmd.record_pipeline(pipeline, bindings, constants, sqsum_workspace_reduced);
+
+        pb++;
+        }
+
+        sqsum_workspace = sqsum_workspace_reduced;
+        }
+
+        {
+        std::vector<VkImageMat> bindings(2);
+        bindings[0] = sqsum_workspace;
+        bindings[1] = var_workspace;
+
+        std::vector<vk_constant_type> constants(5);
+        constants[0].i = sqsum_workspace.w;
+        constants[1].i = sqsum_workspace.h;
+        constants[2].i = sqsum_workspace.c;
+        constants[3].i = 0;//sqsum_workspace.cstep;
+        constants[4].f = size;
+
+        const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_reduce_mean_pack8
+                                 : elempack == 4 ? pipeline_instancenorm_reduce_mean_pack4
+                                 : pipeline_instancenorm_reduce_mean;
+
+        cmd.record_pipeline(pipeline, bindings, constants, var_workspace);
+        }
+    }
+
+    // coeffs
+    VkImageMat coeffs_workspace;
+    coeffs_workspace.create(c * 2, elemsize, elempack, opt.workspace_vkallocator);
+    {
+    std::vector<VkImageMat> bindings(5);
+    bindings[0] = coeffs_workspace;
+    bindings[1] = mean_workspace;
+    bindings[2] = var_workspace;
+    bindings[3] = gamma_data_gpu_image;
+    bindings[4] = beta_data_gpu_image;
+
+    std::vector<vk_constant_type> constants(0);
+
+    const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_coeffs_pack8
+                             : elempack == 4 ? pipeline_instancenorm_coeffs_pack4
+                             : pipeline_instancenorm_coeffs;
+
+    VkImageMat dispatcher;
+    dispatcher.w = c;
+    dispatcher.h = 1;
+    dispatcher.c = 1;
+    cmd.record_pipeline(pipeline, bindings, constants, dispatcher);
+    }
+
+    // norm
+    {
+    std::vector<VkImageMat> bindings(3);
+    bindings[0] = bottom_top_blob;
+    bindings[1] = bottom_top_blob;
+    bindings[2] = coeffs_workspace;
+
+    std::vector<vk_constant_type> constants(5);
+    constants[0].i = bottom_top_blob.dims;
+    constants[1].i = bottom_top_blob.w;
+    constants[2].i = bottom_top_blob.h;
+    constants[3].i = bottom_top_blob.c;
+    constants[4].i = 0;//bottom_top_blob.cstep;
+
+    const Pipeline* pipeline = elempack == 8 ? pipeline_instancenorm_norm_pack8
+                             : elempack == 4 ? pipeline_instancenorm_norm_pack4
+                             : pipeline_instancenorm_norm;
+
+    cmd.record_pipeline(pipeline, bindings, constants, bottom_top_blob);
+    }
+
+    return 0;
+}
+
 } // namespace ncnn

src/layer/vulkan/instancenorm_vulkan.h

@@ -31,10 +31,14 @@ public:
 
     using InstanceNorm::forward_inplace;
     virtual int forward_inplace(VkMat& bottom_top_blob, VkCompute& cmd, const Option& opt) const;
+    virtual int forward_inplace(VkImageMat& bottom_top_blob, VkCompute& cmd, const Option& opt) const;
 
 public:
     VkMat gamma_data_gpu;
     VkMat beta_data_gpu;
+    VkImageMat gamma_data_gpu_image;
+    VkImageMat beta_data_gpu_image;
+
     Pipeline* pipeline_instancenorm_reduce_sum4_fp16_to_fp32;
     Pipeline* pipeline_instancenorm_reduce_sum4_fp32[2];
     Pipeline* pipeline_instancenorm_reduce_mean;

src/layer/vulkan/shader/instancenorm_coeffs.comp

@@ -28,11 +28,19 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0, imfmtc1) writeonly uniform unfp image1D coeffs_blob;
+layout (binding = 1) uniform unfp sampler1D mean_blob;
+layout (binding = 2) uniform unfp sampler1D var_blob;
+layout (binding = 3) uniform unfp sampler1D gamma_blob;
+layout (binding = 4) uniform unfp sampler1D beta_blob;
+#else
 layout (binding = 0) writeonly buffer coeffs_blob { sfp coeffs_blob_data[]; };
 layout (binding = 1) readonly buffer mean_blob { sfp mean_data[]; };
 layout (binding = 2) readonly buffer var_blob { sfp var_data[]; };
 layout (binding = 3) readonly buffer gamma_blob { sfp gamma_data[]; };
 layout (binding = 4) readonly buffer beta_blob { sfp beta_data[]; };
+#endif
 
 void main()
 {
@@ -43,14 +51,26 @@ void main()
     if (gx >= w || gy >= 1 || gz >= 1)
         return;
 
+#if NCNN_image_shader
+    afp mean = image1d_ld1(mean_blob, gx);
+    afp var = image1d_ld1(var_blob, gx);
+    afp gamma = image1d_ld1(gamma_blob, gx);
+    afp beta = image1d_ld1(beta_blob, gx);
+#else
     afp mean = buffer_ld1(mean_data, gx);
     afp var = buffer_ld1(var_data, gx);
     afp gamma = buffer_ld1(gamma_data, gx);
     afp beta = buffer_ld1(beta_data, gx);
+#endif
 
     afp a = gamma / (sqrt(var + afp(eps)));
     afp b = - mean * a + beta;
 
+#if NCNN_image_shader
+    image1d_st1(coeffs_blob, gx*2, a);
+    image1d_st1(coeffs_blob, gx*2 +1, b);
+#else
     buffer_st1(coeffs_blob_data, gx*2, a);
     buffer_st1(coeffs_blob_data, gx*2 +1, b);
+#endif
 }

src/layer/vulkan/shader/instancenorm_coeffs_pack4.comp

@@ -28,11 +28,19 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0, imfmtc4) writeonly uniform unfp image1D coeffs_blob;
+layout (binding = 1) uniform unfp sampler1D mean_blob;
+layout (binding = 2) uniform unfp sampler1D var_blob;
+layout (binding = 3) uniform unfp sampler1D gamma_blob;
+layout (binding = 4) uniform unfp sampler1D beta_blob;
+#else
 layout (binding = 0) writeonly buffer coeffs_blob { sfpvec4 coeffs_blob_data[]; };
 layout (binding = 1) readonly buffer mean_blob { sfpvec4 mean_data[]; };
 layout (binding = 2) readonly buffer var_blob { sfpvec4 var_data[]; };
 layout (binding = 3) readonly buffer gamma_blob { sfpvec4 gamma_data[]; };
 layout (binding = 4) readonly buffer beta_blob { sfpvec4 beta_data[]; };
+#endif
 
 void main()
 {
@@ -43,14 +51,26 @@ void main()
     if (gx >= w || gy >= 1 || gz >= 1)
         return;
 
+#if NCNN_image_shader
+    afpvec4 mean = image1d_ld4(mean_blob, gx);
+    afpvec4 var = image1d_ld4(var_blob, gx);
+    afpvec4 gamma = image1d_ld4(gamma_blob, gx);
+    afpvec4 beta = image1d_ld4(beta_blob, gx);
+#else
     afpvec4 mean = buffer_ld4(mean_data, gx);
     afpvec4 var = buffer_ld4(var_data, gx);
     afpvec4 gamma = buffer_ld4(gamma_data, gx);
     afpvec4 beta = buffer_ld4(beta_data, gx);
+#endif
 
     afpvec4 a = gamma / (sqrt(var + afp(eps)));
     afpvec4 b = - mean * a + beta;
 
+#if NCNN_image_shader
+    image1d_st4(coeffs_blob, gx*2, a);
+    image1d_st4(coeffs_blob, gx*2 +1, b);
+#else
     buffer_st4(coeffs_blob_data, gx*2, a);
     buffer_st4(coeffs_blob_data, gx*2 +1, b);
+#endif
 }

src/layer/vulkan/shader/instancenorm_coeffs_pack8.comp

@@ -29,11 +29,19 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0, imfmtc4) writeonly uniform unfp image1D coeffs_blob;
+layout (binding = 1) uniform unfp sampler1D mean_blob;
+layout (binding = 2) uniform unfp sampler1D var_blob;
+layout (binding = 3) uniform unfp sampler1D gamma_blob;
+layout (binding = 4) uniform unfp sampler1D beta_blob;
+#else
 layout (binding = 0) writeonly buffer coeffs_blob { sfpvec8 coeffs_blob_data[]; };
 layout (binding = 1) readonly buffer mean_blob { sfpvec8 mean_data[]; };
 layout (binding = 2) readonly buffer var_blob { sfpvec8 var_data[]; };
 layout (binding = 3) readonly buffer gamma_blob { sfpvec8 gamma_data[]; };
 layout (binding = 4) readonly buffer beta_blob { sfpvec8 beta_data[]; };
+#endif
 
 void main()
 {
@@ -44,10 +52,17 @@ void main()
     if (gx >= w || gy >= 1 || gz >= 1)
         return;
 
+#if NCNN_image_shader
+    afpvec8 mean = image1d_ld8(mean_blob, gx);
+    afpvec8 var = image1d_ld8(var_blob, gx);
+    afpvec8 gamma = image1d_ld8(gamma_blob, gx);
+    afpvec8 beta = image1d_ld8(beta_blob, gx);
+#else
     afpvec8 mean = buffer_ld8(mean_data, gx);
     afpvec8 var = buffer_ld8(var_data, gx);
     afpvec8 gamma = buffer_ld8(gamma_data, gx);
     afpvec8 beta = buffer_ld8(beta_data, gx);
+#endif
 
     afpvec8 a;
     afpvec8 b;
@@ -56,6 +71,11 @@ void main()
     b[0] = - mean[0] * a[0] + beta[0];
     b[1] = - mean[1] * a[1] + beta[1];
 
+#if NCNN_image_shader
+    image1d_st8(coeffs_blob, gx*2, a);
+    image1d_st8(coeffs_blob, gx*2 +1, b);
+#else
     buffer_st8(coeffs_blob_data, gx*2, a);
     buffer_st8(coeffs_blob_data, gx*2 +1, b);
+#endif
 }

src/layer/vulkan/shader/instancenorm_norm.comp

@@ -32,8 +32,14 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform unfp sampler3D bottom_blob;
+layout (binding = 1, imfmtc1) writeonly uniform unfp image3D top_blob;
+layout (binding = 2) uniform unfp sampler1D coeffs_blob;
+#else
 layout (binding = 0) buffer bottom_top_blob { sfp bottom_top_blob_data[]; };
 layout (binding = 1) readonly buffer coeffs_blob { sfp coeffs_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
@@ -53,14 +59,25 @@ void main()
     if (gx >= psc(w) || gy >= psc(h) || gz >= psc(c))
         return;
 
+#if NCNN_image_shader
+    afp v = image3d_ld1(bottom_blob, ivec3(gx, gy, gz));
+
+    afp a = image1d_ld1(coeffs_blob, gz*2);
+    afp b = image1d_ld1(coeffs_blob, gz*2 +1);
+#else
     const int gi = gz * psc(cstep) + gy * psc(w) + gx;
 
     afp v = buffer_ld1(bottom_top_blob_data, gi);
 
     afp a = buffer_ld1(coeffs_blob_data, gz*2);
     afp b = buffer_ld1(coeffs_blob_data, gz*2 +1);
+#endif
 
     v = v * a + b;
 
+#if NCNN_image_shader
+    image3d_st1(top_blob, ivec3(gx, gy, gz), v);
+#else
     buffer_st1(bottom_top_blob_data, gi, v);
+#endif
 }

src/layer/vulkan/shader/instancenorm_norm_pack4.comp

@@ -32,8 +32,14 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform unfp sampler3D bottom_blob;
+layout (binding = 1, imfmtc4) writeonly uniform unfp image3D top_blob;
+layout (binding = 2) uniform unfp sampler1D coeffs_blob;
+#else
 layout (binding = 0) buffer bottom_top_blob { sfpvec4 bottom_top_blob_data[]; };
 layout (binding = 1) readonly buffer coeffs_blob { sfpvec4 coeffs_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
@@ -53,14 +59,25 @@ void main()
     if (gx >= psc(w) || gy >= psc(h) || gz >= psc(c))
         return;
 
+#if NCNN_image_shader
+    afpvec4 v = image3d_ld4(bottom_blob, ivec3(gx, gy, gz));
+
+    afpvec4 a = image1d_ld4(coeffs_blob, gz*2);
+    afpvec4 b = image1d_ld4(coeffs_blob, gz*2 +1);
+#else
     const int gi = gz * psc(cstep) + gy * psc(w) + gx;
 
     afpvec4 v = buffer_ld4(bottom_top_blob_data, gi);
 
     afpvec4 a = buffer_ld4(coeffs_blob_data, gz*2);
     afpvec4 b = buffer_ld4(coeffs_blob_data, gz*2 +1);
+#endif
 
     v = v * a + b;
 
+#if NCNN_image_shader
+    image3d_st4(top_blob, ivec3(gx, gy, gz), v);
+#else
     buffer_st4(bottom_top_blob_data, gi, v);
+#endif
 }

src/layer/vulkan/shader/instancenorm_norm_pack8.comp

@@ -33,8 +33,14 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform unfp sampler3D bottom_blob;
+layout (binding = 1, imfmtc4) writeonly uniform unfp image3D top_blob;
+layout (binding = 2) uniform unfp sampler1D coeffs_blob;
+#else
 layout (binding = 0) buffer bottom_top_blob { sfpvec8 bottom_top_blob_data[]; };
 layout (binding = 1) readonly buffer coeffs_blob { sfpvec8 coeffs_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
@@ -54,15 +60,26 @@ void main()
     if (gx >= psc(w) || gy >= psc(h) || gz >= psc(c))
         return;
 
+#if NCNN_image_shader
+    afpvec8 v = image3d_ld8(bottom_blob, ivec3(gx, gy, gz));
+
+    afpvec8 a = image1d_ld8(coeffs_blob, gz*2);
+    afpvec8 b = image1d_ld8(coeffs_blob, gz*2 +1);
+#else
     const int gi = gz * psc(cstep) + gy * psc(w) + gx;
 
     afpvec8 v = buffer_ld8(bottom_top_blob_data, gi);
 
     afpvec8 a = buffer_ld8(coeffs_blob_data, gz*2);
     afpvec8 b = buffer_ld8(coeffs_blob_data, gz*2 +1);
+#endif
 
     v[0] = v[0] * a[0] + b[0];
     v[1] = v[1] * a[1] + b[1];
 
+#if NCNN_image_shader
+    image3d_st8(top_blob, ivec3(gx, gy, gz), v);
+#else
     buffer_st8(bottom_top_blob_data, gi, v);
+#endif
 }

src/layer/vulkan/shader/instancenorm_reduce_mean.comp

@@ -22,20 +22,27 @@
 #endif
 
 #define shape_constant_id_offset 0
-layout (constant_id = shape_constant_id_offset + 0) const int size = 0;
-layout (constant_id = shape_constant_id_offset + 1) const int c = 0;
-layout (constant_id = shape_constant_id_offset + 2) const int cstep = 0;
+layout (constant_id = shape_constant_id_offset + 0) const int w = 0;
+layout (constant_id = shape_constant_id_offset + 1) const int h = 0;
+layout (constant_id = shape_constant_id_offset + 2) const int c = 0;
+layout (constant_id = shape_constant_id_offset + 3) const int cstep = 0;
 
 layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform highp sampler3D bottom_top_blob;
+layout (binding = 1, imfmtc1) writeonly uniform unfp image1D mean_blob;
+#else
 layout (binding = 0) readonly buffer bottom_top_blob { float bottom_top_blob_data[]; };
 layout (binding = 1) writeonly buffer mean_blob { sfp mean_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
-    int size;
+    int w;
+    int h;
     int c;
     int cstep;
     float area;
@@ -52,15 +59,29 @@ void main()
 
     afp sum = afp(0.f);
 
+#if NCNN_image_shader
+    for (int i = 0; i < p.h; i++)
+    {
+        for (int j = 0; j < p.w; j++)
+        {
+            sum += afp(texelFetch(bottom_top_blob, ivec3(j, i, gx), 0).r);
+        }
+    }
+#else
     int v_offset = gx * psc(cstep);
 
-    for (int i = 0; i < p.size; i++)
+    for (int i = 0; i < p.w; i++)
     {
         sum += afp(bottom_top_blob_data[v_offset]);
         v_offset += 1;
     }
+#endif
 
     afp mean = sum / afp(p.area);
 
+#if NCNN_image_shader
+    image1d_st1(mean_blob, gx, mean);
+#else
     buffer_st1(mean_data, gx, mean);
+#endif
 }

src/layer/vulkan/shader/instancenorm_reduce_mean_pack4.comp

@@ -22,20 +22,27 @@
 #endif
 
 #define shape_constant_id_offset 0
-layout (constant_id = shape_constant_id_offset + 0) const int size = 0;
-layout (constant_id = shape_constant_id_offset + 1) const int c = 0;
-layout (constant_id = shape_constant_id_offset + 2) const int cstep = 0;
+layout (constant_id = shape_constant_id_offset + 0) const int w = 0;
+layout (constant_id = shape_constant_id_offset + 1) const int h = 0;
+layout (constant_id = shape_constant_id_offset + 2) const int c = 0;
+layout (constant_id = shape_constant_id_offset + 3) const int cstep = 0;
 
 layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform highp sampler3D bottom_top_blob;
+layout (binding = 1, imfmtc4) writeonly uniform unfp image1D mean_blob;
+#else
 layout (binding = 0) readonly buffer bottom_top_blob { vec4 bottom_top_blob_data[]; };
 layout (binding = 1) writeonly buffer mean_blob { sfpvec4 mean_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
-    int size;
+    int w;
+    int h;
     int c;
     int cstep;
     float area;
@@ -52,15 +59,29 @@ void main()
 
     afpvec4 sum = afpvec4(0.f);
 
+#if NCNN_image_shader
+    for (int i = 0; i < p.h; i++)
+    {
+        for (int j = 0; j < p.w; j++)
+        {
+            sum += afpvec4(texelFetch(bottom_top_blob, ivec3(j, i, gx), 0));
+        }
+    }
+#else
     int v_offset = gx * psc(cstep);
 
-    for (int i = 0; i < p.size; i++)
+    for (int i = 0; i < p.w; i++)
     {
         sum += afpvec4(bottom_top_blob_data[v_offset]);
         v_offset += 1;
     }
+#endif
 
     afpvec4 mean = sum / afp(p.area);
 
+#if NCNN_image_shader
+    image1d_st4(mean_blob, gx, mean);
+#else
     buffer_st4(mean_data, gx, mean);
+#endif
 }

src/layer/vulkan/shader/instancenorm_reduce_mean_pack8.comp

@@ -23,20 +23,27 @@ struct sfpvec8 { f16vec4 abcd; f16vec4 efgh; };
 #endif
 
 #define shape_constant_id_offset 0
-layout (constant_id = shape_constant_id_offset + 0) const int size = 0;
-layout (constant_id = shape_constant_id_offset + 1) const int c = 0;
-layout (constant_id = shape_constant_id_offset + 2) const int cstep = 0;
+layout (constant_id = shape_constant_id_offset + 0) const int w = 0;
+layout (constant_id = shape_constant_id_offset + 1) const int h = 0;
+layout (constant_id = shape_constant_id_offset + 2) const int c = 0;
+layout (constant_id = shape_constant_id_offset + 3) const int cstep = 0;
 
 layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform highp sampler3D bottom_top_blob;
+layout (binding = 1, imfmtc4) writeonly uniform unfp image1D mean_blob;
+#else
 layout (binding = 0) readonly buffer bottom_top_blob { mat2x4 bottom_top_blob_data[]; };
 layout (binding = 1) writeonly buffer mean_blob { sfpvec8 mean_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
-    int size;
+    int w;
+    int h;
     int c;
     int cstep;
     float area;
@@ -53,17 +60,31 @@ void main()
 
     afpvec8 sum = afpvec8(afpvec4(0.f), afpvec4(0.f));
 
+#if NCNN_image_shader
+    for (int i = 0; i < p.h; i++)
+    {
+        for (int j = 0; j < p.w; j++)
+        {
+            sum += afpvec8(texelFetch(bottom_top_blob, ivec3(j * 2, i, gx), 0), texelFetch(bottom_top_blob, ivec3(j * 2 + 1, i, gx), 0));
+        }
+    }
+#else
     int v_offset = gx * psc(cstep);
 
-    for (int i = 0; i < p.size; i++)
+    for (int i = 0; i < p.w; i++)
     {
         sum += afpvec8(bottom_top_blob_data[v_offset]);
         v_offset += 1;
     }
+#endif
 
     afpvec8 mean;
     mean[0] = sum[0] / afp(p.area);
     mean[1] = sum[1] / afp(p.area);
 
+#if NCNN_image_shader
+    image1d_st8(mean_blob, gx, mean);
+#else
     buffer_st8(mean_data, gx, mean);
+#endif
 }

src/layer/vulkan/shader/instancenorm_reduce_sum4_fp16_to_fp32.comp

@@ -27,16 +27,23 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform unfp sampler3D bottom_top_blob;
+layout (binding = 1, r32f) writeonly uniform highp image3D sum_blob;
+#else
 layout (binding = 0) readonly buffer bottom_top_blob { sfp bottom_top_blob_data[]; };
 layout (binding = 1) writeonly buffer sum_blob { float sum_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
-    int size;
+    int w;
+    int h;
     int c;
     int cstep;
 
-    int outsize;
+    int outw;
+    int outh;
     int outc;
     int outcstep;
 } p;
@@ -47,37 +54,75 @@ void main()
     int gy = int(gl_GlobalInvocationID.y);
     int gz = int(gl_GlobalInvocationID.z);
 
-    if (gx >= p.outsize || gy >= 1 || gz >= p.outc)
+    if (gx >= p.outw || gy >= p.outh || gz >= p.outc)
         return;
 
-    int gi = gz * p.outcstep + gx;
+    float sum;
+
+#if NCNN_image_shader
+    int sx = gx * 2;
+    int sy = gy * 2;
 
+    if (sy == p.h - 1)
+    {
+        if (sx == p.w - 1)
+        {
+            float v0 = texelFetch(bottom_top_blob, ivec3(sx, sy, gz), 0).r;
+
+            sum = v0;
+        }
+        else
+        {
+            float v0 = texelFetch(bottom_top_blob, ivec3(sx, sy, gz), 0).r;
+            float v1 = texelFetch(bottom_top_blob, ivec3(sx + 1, sy, gz), 0).r;
+
+            sum = v0 + v1;
+        }
+    }
+    else
+    {
+        if (sx == p.w - 1)
+        {
+            float v0 = texelFetch(bottom_top_blob, ivec3(sx, sy, gz), 0).r;
+            float v2 = texelFetch(bottom_top_blob, ivec3(sx, sy + 1, gz), 0).r;
+
+            sum = v0 + v2;
+        }
+        else
+        {
+            float v0 = texelFetch(bottom_top_blob, ivec3(sx, sy, gz), 0).r;
+            float v1 = texelFetch(bottom_top_blob, ivec3(sx + 1, sy, gz), 0).r;
+            float v2 = texelFetch(bottom_top_blob, ivec3(sx, sy + 1, gz), 0).r;
+            float v3 = texelFetch(bottom_top_blob, ivec3(sx + 1, sy + 1, gz), 0).r;
+
+            sum = v0 + v1 + v2 + v3;
+        }
+    }
+#else
     int sx = gx * 4;
 
     int v_offset = gz * p.cstep + sx;
 
-    if (sx == p.size - 1)
+    if (sx == p.w - 1)
     {
-        sum_blob_data[gi] = float(buffer_ld1(bottom_top_blob_data, v_offset));
+        float v0 = float(buffer_ld1(bottom_top_blob_data, v_offset));
+
+        sum = v0;
     }
-    else if (sx == p.size - 2)
+    else if (sx == p.w - 2)
     {
         float v0 = float(buffer_ld1(bottom_top_blob_data, v_offset));
         float v1 = float(buffer_ld1(bottom_top_blob_data, v_offset + 1));
 
-        float sum = v0 + v1;
-
-        sum_blob_data[gi] = sum;
+        sum = v0 + v1;
     }
-    else if (sx == p.size - 3)
+    else if (sx == p.w - 3)
     {
         float v0 = float(buffer_ld1(bottom_top_blob_data, v_offset));
         float v1 = float(buffer_ld1(bottom_top_blob_data, v_offset + 1));
         float v2 = float(buffer_ld1(bottom_top_blob_data, v_offset + 2));
 
-        float sum = v0 + v1 + v2;
-
-        sum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2;
     }
     else
     {
@@ -86,8 +131,15 @@ void main()
         float v2 = float(buffer_ld1(bottom_top_blob_data, v_offset + 2));
         float v3 = float(buffer_ld1(bottom_top_blob_data, v_offset + 3));
 
-        float sum = v0 + v1 + v2 + v3;
-
-        sum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2 + v3;
     }
+#endif
+
+#if NCNN_image_shader
+    imageStore(sum_blob, ivec3(gx, gy, gz), vec4(sum));
+#else
+    int gi = gz * p.outcstep + gx;
+
+    sum_blob_data[gi] = sum;
+#endif
 }

src/layer/vulkan/shader/instancenorm_reduce_sum4_fp16_to_fp32_pack4.comp

@@ -31,16 +31,23 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform unfp sampler3D bottom_top_blob;
+layout (binding = 1, rgba32f) writeonly uniform highp image3D sum_blob;
+#else
 layout (binding = 0) readonly buffer bottom_top_blob { sfpvec4 bottom_top_blob_data[]; };
 layout (binding = 1) writeonly buffer sum_blob { vec4 sum_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
-    int size;
+    int w;
+    int h;
     int c;
     int cstep;
 
-    int outsize;
+    int outw;
+    int outh;
     int outc;
     int outcstep;
 } p;
@@ -51,37 +58,73 @@ void main()
     int gy = int(gl_GlobalInvocationID.y);
     int gz = int(gl_GlobalInvocationID.z);
 
-    if (gx >= p.outsize || gy >= 1 || gz >= p.outc)
+    if (gx >= p.outw || gy >= p.outh || gz >= p.outc)
         return;
 
-    int gi = gz * p.outcstep + gx;
+    vec4 sum;
 
+#if NCNN_image_shader
+    int sx = gx * 2;
+    int sy = gy * 2;
+
+    if (sy == p.h - 1)
+    {
+        if (sx == p.w - 1)
+        {
+            vec4 v0 = texelFetch(bottom_top_blob, ivec3(sx, sy, gz), 0);
+
+            sum = v0;
+        }
+        else
+        {
+            vec4 v0 = texelFetch(bottom_top_blob, ivec3(sx, sy, gz), 0);
+            vec4 v1 = texelFetch(bottom_top_blob, ivec3(sx + 1, sy, gz), 0);
+
+            sum = v0 + v1;
+        }
+    }
+    else
+    {
+        if (sx == p.w - 1)
+        {
+            vec4 v0 = texelFetch(bottom_top_blob, ivec3(sx, sy, gz), 0);
+            vec4 v2 = texelFetch(bottom_top_blob, ivec3(sx, sy + 1, gz), 0);
+
+            sum = v0 + v2;
+        }
+        else
+        {
+            vec4 v0 = texelFetch(bottom_top_blob, ivec3(sx, sy, gz), 0);
+            vec4 v1 = texelFetch(bottom_top_blob, ivec3(sx + 1, sy, gz), 0);
+            vec4 v2 = texelFetch(bottom_top_blob, ivec3(sx, sy + 1, gz), 0);
+            vec4 v3 = texelFetch(bottom_top_blob, ivec3(sx + 1, sy + 1, gz), 0);
+
+            sum = v0 + v1 + v2 + v3;
+        }
+    }
+#else
     int sx = gx * 4;
 
     int v_offset = gz * p.cstep + sx;
 
-    if (sx == p.size - 1)
+    if (sx == p.w - 1)
     {
-        sum_blob_data[gi] = vec4(buffer_ld4(bottom_top_blob_data, v_offset));
+        sum = vec4(buffer_ld4(bottom_top_blob_data, v_offset));
     }
-    else if (sx == p.size - 2)
+    else if (sx == p.w - 2)
     {
         vec4 v0 = vec4(buffer_ld4(bottom_top_blob_data, v_offset));
         vec4 v1 = vec4(buffer_ld4(bottom_top_blob_data, v_offset + 1));
 
-        vec4 sum = v0 + v1;
-
-        sum_blob_data[gi] = sum;
+        sum = v0 + v1;
     }
-    else if (sx == p.size - 3)
+    else if (sx == p.w - 3)
     {
         vec4 v0 = vec4(buffer_ld4(bottom_top_blob_data, v_offset));
         vec4 v1 = vec4(buffer_ld4(bottom_top_blob_data, v_offset + 1));
         vec4 v2 = vec4(buffer_ld4(bottom_top_blob_data, v_offset + 2));
 
-        vec4 sum = v0 + v1 + v2;
-
-        sum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2;
     }
     else
     {
@@ -90,8 +133,15 @@ void main()
         vec4 v2 = vec4(buffer_ld4(bottom_top_blob_data, v_offset + 2));
         vec4 v3 = vec4(buffer_ld4(bottom_top_blob_data, v_offset + 3));
 
-        vec4 sum = v0 + v1 + v2 + v3;
-
-        sum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2 + v3;
     }
+#endif
+
+#if NCNN_image_shader
+    imageStore(sum_blob, ivec3(gx, gy, gz), sum);
+#else
+    int gi = gz * p.outcstep + gx;
+
+    sum_blob_data[gi] = sum;
+#endif
 }

src/layer/vulkan/shader/instancenorm_reduce_sum4_fp16_to_fp32_pack8.comp

@@ -26,16 +26,23 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform unfp sampler3D bottom_top_blob;
+layout (binding = 1, rgba32f) writeonly uniform highp image3D sum_blob;
+#else
 layout (binding = 0) readonly buffer bottom_top_blob { sfpvec8 bottom_top_blob_data[]; };
 layout (binding = 1) writeonly buffer sum_blob { mat2x4 sum_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
-    int size;
+    int w;
+    int h;
     int c;
     int cstep;
 
-    int outsize;
+    int outw;
+    int outh;
     int outc;
     int outcstep;
 } p;
@@ -46,37 +53,73 @@ void main()
     int gy = int(gl_GlobalInvocationID.y);
     int gz = int(gl_GlobalInvocationID.z);
 
-    if (gx >= p.outsize || gy >= 1 || gz >= p.outc)
+    if (gx >= p.outw || gy >= p.outh || gz >= p.outc)
         return;
 
-    int gi = gz * p.outcstep + gx;
+    mat2x4 sum;
 
+#if NCNN_image_shader
+    int sx = gx * 2;
+    int sy = gy * 2;
+
+    if (sy == p.h - 1)
+    {
+        if (sx == p.w - 1)
+        {
+            mat2x4 v0 = mat2x4(texelFetch(bottom_top_blob, ivec3(sx * 2, sy, gz), 0), texelFetch(bottom_top_blob, ivec3(sx * 2 + 1, sy, gz), 0));
+
+            sum = v0;
+        }
+        else
+        {
+            mat2x4 v0 = mat2x4(texelFetch(bottom_top_blob, ivec3(sx * 2, sy, gz), 0), texelFetch(bottom_top_blob, ivec3(sx * 2 + 1, sy, gz), 0));
+            mat2x4 v1 = mat2x4(texelFetch(bottom_top_blob, ivec3((sx + 1) * 2, sy, gz), 0), texelFetch(bottom_top_blob, ivec3((sx + 1) * 2 + 1, sy, gz), 0));
+
+            sum = v0 + v1;
+        }
+    }
+    else
+    {
+        if (sx == p.w - 1)
+        {
+            mat2x4 v0 = mat2x4(texelFetch(bottom_top_blob, ivec3(sx * 2, sy, gz), 0), texelFetch(bottom_top_blob, ivec3(sx * 2 + 1, sy, gz), 0));
+            mat2x4 v2 = mat2x4(texelFetch(bottom_top_blob, ivec3(sx * 2, sy + 1, gz), 0), texelFetch(bottom_top_blob, ivec3(sx * 2 + 1, sy + 1, gz), 0));
+
+            sum = v0 + v2;
+        }
+        else
+        {
+            mat2x4 v0 = mat2x4(texelFetch(bottom_top_blob, ivec3(sx * 2, sy, gz), 0), texelFetch(bottom_top_blob, ivec3(sx * 2 + 1, sy, gz), 0));
+            mat2x4 v1 = mat2x4(texelFetch(bottom_top_blob, ivec3((sx + 1) * 2, sy, gz), 0), texelFetch(bottom_top_blob, ivec3((sx + 1) * 2 + 1, sy, gz), 0));
+            mat2x4 v2 = mat2x4(texelFetch(bottom_top_blob, ivec3(sx * 2, sy + 1, gz), 0), texelFetch(bottom_top_blob, ivec3(sx * 2 + 1, sy + 1, gz), 0));
+            mat2x4 v3 = mat2x4(texelFetch(bottom_top_blob, ivec3((sx + 1) * 2, sy + 1, gz), 0), texelFetch(bottom_top_blob, ivec3((sx + 1) * 2 + 1, sy + 1, gz), 0));
+
+            sum = v0 + v1 + v2 + v3;
+        }
+    }
+#else
     int sx = gx * 4;
 
     int v_offset = gz * p.cstep + sx;
 
-    if (sx == p.size - 1)
+    if (sx == p.w - 1)
     {
-        sum_blob_data[gi] = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset));
+        sum = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset));
     }
-    else if (sx == p.size - 2)
+    else if (sx == p.w - 2)
     {
         mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset));
         mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset + 1));
 
-        mat2x4 sum = v0 + v1;
-
-        sum_blob_data[gi] = sum;
+        sum = v0 + v1;
     }
-    else if (sx == p.size - 3)
+    else if (sx == p.w - 3)
     {
         mat2x4 v0 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset));
         mat2x4 v1 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset + 1));
         mat2x4 v2 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset + 2));
 
-        mat2x4 sum = v0 + v1 + v2;
-
-        sum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2;
     }
     else
     {
@@ -85,8 +128,16 @@ void main()
         mat2x4 v2 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset + 2));
         mat2x4 v3 = mat2x4(buffer_ld8(bottom_top_blob_data, v_offset + 3));
 
-        mat2x4 sum = v0 + v1 + v2 + v3;
-
-        sum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2 + v3;
     }
+#endif
+
+#if NCNN_image_shader
+    imageStore(sum_blob, ivec3(gx * 2, gy, gz), sum[0]);
+    imageStore(sum_blob, ivec3(gx * 2 + 1, gy, gz), sum[1]);
+#else
+    int gi = gz * p.outcstep + gx;
+
+    sum_blob_data[gi] = sum;
+#endif
 }

src/layer/vulkan/shader/instancenorm_reduce_sum4_fp32.comp

@@ -25,16 +25,23 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform highp sampler3D square_blob;
+layout (binding = 1, r32f) writeonly uniform highp image3D sqsum_blob;
+#else
 layout (binding = 0) readonly buffer square_blob { float square_blob_data[]; };
 layout (binding = 1) writeonly buffer sqsum_blob { float sqsum_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
-    int size;
+    int w;
+    int h;
     int c;
     int cstep;
 
-    int outsize;
+    int outw;
+    int outh;
     int outc;
     int outcstep;
 } p;
@@ -45,37 +52,75 @@ void main()
     int gy = int(gl_GlobalInvocationID.y);
     int gz = int(gl_GlobalInvocationID.z);
 
-    if (gx >= p.outsize || gy >= 1 || gz >= p.outc)
+    if (gx >= p.outw || gy >= p.outh || gz >= p.outc)
         return;
 
-    int gi = gz * p.outcstep + gx;
+    float sum;
+
+#if NCNN_image_shader
+    int sx = gx * 2;
+    int sy = gy * 2;
 
+    if (sy == p.h - 1)
+    {
+        if (sx == p.w - 1)
+        {
+            float v0 = texelFetch(square_blob, ivec3(sx, sy, gz), 0).r;
+
+            sum = v0;
+        }
+        else
+        {
+            float v0 = texelFetch(square_blob, ivec3(sx, sy, gz), 0).r;
+            float v1 = texelFetch(square_blob, ivec3(sx + 1, sy, gz), 0).r;
+
+            sum = v0 + v1;
+        }
+    }
+    else
+    {
+        if (sx == p.w - 1)
+        {
+            float v0 = texelFetch(square_blob, ivec3(sx, sy, gz), 0).r;
+            float v2 = texelFetch(square_blob, ivec3(sx, sy + 1, gz), 0).r;
+
+            sum = v0 + v2;
+        }
+        else
+        {
+            float v0 = texelFetch(square_blob, ivec3(sx, sy, gz), 0).r;
+            float v1 = texelFetch(square_blob, ivec3(sx + 1, sy, gz), 0).r;
+            float v2 = texelFetch(square_blob, ivec3(sx, sy + 1, gz), 0).r;
+            float v3 = texelFetch(square_blob, ivec3(sx + 1, sy + 1, gz), 0).r;
+
+            sum = v0 + v1 + v2 + v3;
+        }
+    }
+#else
     int sx = gx * 4;
 
     int v_offset = gz * p.cstep + sx;
 
-    if (sx == p.size - 1)
+    if (sx == p.w - 1)
     {
-        sqsum_blob_data[gi] = square_blob_data[v_offset];
+        float v0 = square_blob_data[v_offset];
+
+        sum = v0;
     }
-    else if (sx == p.size - 2)
+    else if (sx == p.w - 2)
     {
         float v0 = square_blob_data[v_offset];
         float v1 = square_blob_data[v_offset + 1];
 
-        float sum = v0 + v1;
-
-        sqsum_blob_data[gi] = sum;
+        sum = v0 + v1;
     }
-    else if (sx == p.size - 2)
+    else if (sx == p.w - 2)
     {
         float v0 = square_blob_data[v_offset];
         float v1 = square_blob_data[v_offset + 1];
         float v2 = square_blob_data[v_offset + 2];
 
-        float sum = v0 + v1 + v2;
-
-        sqsum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2;
     }
     else
     {
@@ -84,8 +129,15 @@ void main()
         float v2 = square_blob_data[v_offset + 2];
         float v3 = square_blob_data[v_offset + 3];
 
-        float sum = v0 + v1 + v2 + v3;
-
-        sqsum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2 + v3;
     }
+#endif
+
+#if NCNN_image_shader
+    imageStore(sqsum_blob, ivec3(gx, gy, gz), vec4(sum));
+#else
+    int gi = gz * p.outcstep + gx;
+
+    sqsum_blob_data[gi] = sum;
+#endif
 }

src/layer/vulkan/shader/instancenorm_reduce_sum4_fp32_pack4.comp

@@ -25,16 +25,23 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform highp sampler3D square_blob;
+layout (binding = 1, rgba32f) writeonly uniform highp image3D sqsum_blob;
+#else
 layout (binding = 0) readonly buffer square_blob { vec4 square_blob_data[]; };
 layout (binding = 1) writeonly buffer sqsum_blob { vec4 sqsum_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
-    int size;
+    int w;
+    int h;
     int c;
     int cstep;
 
-    int outsize;
+    int outw;
+    int outh;
     int outc;
     int outcstep;
 } p;
@@ -45,37 +52,73 @@ void main()
     int gy = int(gl_GlobalInvocationID.y);
     int gz = int(gl_GlobalInvocationID.z);
 
-    if (gx >= p.outsize || gy >= 1 || gz >= p.outc)
+    if (gx >= p.outw || gy >= p.outh || gz >= p.outc)
         return;
 
-    int gi = gz * p.outcstep + gx;
+    vec4 sum;
 
+#if NCNN_image_shader
+    int sx = gx * 2;
+    int sy = gy * 2;
+
+    if (sy == p.h - 1)
+    {
+        if (sx == p.w - 1)
+        {
+            vec4 v0 = texelFetch(square_blob, ivec3(sx, sy, gz), 0);
+
+            sum = v0;
+        }
+        else
+        {
+            vec4 v0 = texelFetch(square_blob, ivec3(sx, sy, gz), 0);
+            vec4 v1 = texelFetch(square_blob, ivec3(sx + 1, sy, gz), 0);
+
+            sum = v0 + v1;
+        }
+    }
+    else
+    {
+        if (sx == p.w - 1)
+        {
+            vec4 v0 = texelFetch(square_blob, ivec3(sx, sy, gz), 0);
+            vec4 v2 = texelFetch(square_blob, ivec3(sx, sy + 1, gz), 0);
+
+            sum = v0 + v2;
+        }
+        else
+        {
+            vec4 v0 = texelFetch(square_blob, ivec3(sx, sy, gz), 0);
+            vec4 v1 = texelFetch(square_blob, ivec3(sx + 1, sy, gz), 0);
+            vec4 v2 = texelFetch(square_blob, ivec3(sx, sy + 1, gz), 0);
+            vec4 v3 = texelFetch(square_blob, ivec3(sx + 1, sy + 1, gz), 0);
+
+            sum = v0 + v1 + v2 + v3;
+        }
+    }
+#else
     int sx = gx * 4;
 
     int v_offset = gz * p.cstep + sx;
 
-    if (sx == p.size - 1)
+    if (sx == p.w - 1)
     {
-        sqsum_blob_data[gi] = square_blob_data[v_offset];
+        sum = square_blob_data[v_offset];
     }
-    else if (sx == p.size - 2)
+    else if (sx == p.w - 2)
     {
         vec4 v0 = square_blob_data[v_offset];
         vec4 v1 = square_blob_data[v_offset + 1];
 
-        vec4 sum = v0 + v1;
-
-        sqsum_blob_data[gi] = sum;
+        sum = v0 + v1;
     }
-    else if (sx == p.size - 3)
+    else if (sx == p.w - 3)
     {
         vec4 v0 = square_blob_data[v_offset];
         vec4 v1 = square_blob_data[v_offset + 1];
         vec4 v2 = square_blob_data[v_offset + 2];
 
-        vec4 sum = v0 + v1 + v2;
-
-        sqsum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2;
     }
     else
     {
@@ -84,8 +127,15 @@ void main()
         vec4 v2 = square_blob_data[v_offset + 2];
         vec4 v3 = square_blob_data[v_offset + 3];
 
-        vec4 sum = v0 + v1 + v2 + v3;
-
-        sqsum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2 + v3;
     }
+#endif
+
+#if NCNN_image_shader
+    imageStore(sqsum_blob, ivec3(gx, gy, gz), sum);
+#else
+    int gi = gz * p.outcstep + gx;
+
+    sqsum_blob_data[gi] = sum;
+#endif
 }

src/layer/vulkan/shader/instancenorm_reduce_sum4_fp32_pack8.comp

@@ -26,16 +26,23 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform highp sampler3D square_blob;
+layout (binding = 1, rgba32f) writeonly uniform highp image3D sqsum_blob;
+#else
 layout (binding = 0) readonly buffer square_blob { mat2x4 square_blob_data[]; };
 layout (binding = 1) writeonly buffer sqsum_blob { mat2x4 sqsum_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
-    int size;
+    int w;
+    int h;
     int c;
     int cstep;
 
-    int outsize;
+    int outw;
+    int outh;
     int outc;
     int outcstep;
 } p;
@@ -46,37 +53,73 @@ void main()
     int gy = int(gl_GlobalInvocationID.y);
     int gz = int(gl_GlobalInvocationID.z);
 
-    if (gx >= p.outsize || gy >= 1 || gz >= p.outc)
+    if (gx >= p.outw || gy >= p.outh || gz >= p.outc)
         return;
 
-    int gi = gz * p.outcstep + gx;
+    mat2x4 sum;
 
+#if NCNN_image_shader
+    int sx = gx * 2;
+    int sy = gy * 2;
+
+    if (sy == p.h - 1)
+    {
+        if (sx == p.w - 1)
+        {
+            mat2x4 v0 = mat2x4(texelFetch(square_blob, ivec3(sx * 2, sy, gz), 0), texelFetch(square_blob, ivec3(sx * 2 + 1, sy, gz), 0));
+
+            sum = v0;
+        }
+        else
+        {
+            mat2x4 v0 = mat2x4(texelFetch(square_blob, ivec3(sx * 2, sy, gz), 0), texelFetch(square_blob, ivec3(sx * 2 + 1, sy, gz), 0));
+            mat2x4 v1 = mat2x4(texelFetch(square_blob, ivec3((sx + 1) * 2, sy, gz), 0), texelFetch(square_blob, ivec3((sx + 1) * 2 + 1, sy, gz), 0));
+
+            sum = v0 + v1;
+        }
+    }
+    else
+    {
+        if (sx == p.w - 1)
+        {
+            mat2x4 v0 = mat2x4(texelFetch(square_blob, ivec3(sx * 2, sy, gz), 0), texelFetch(square_blob, ivec3(sx * 2 + 1, sy, gz), 0));
+            mat2x4 v2 = mat2x4(texelFetch(square_blob, ivec3(sx * 2, sy + 1, gz), 0), texelFetch(square_blob, ivec3(sx * 2 + 1, sy + 1, gz), 0));
+
+            sum = v0 + v2;
+        }
+        else
+        {
+            mat2x4 v0 = mat2x4(texelFetch(square_blob, ivec3(sx * 2, sy, gz), 0), texelFetch(square_blob, ivec3(sx * 2 + 1, sy, gz), 0));
+            mat2x4 v1 = mat2x4(texelFetch(square_blob, ivec3((sx + 1) * 2, sy, gz), 0), texelFetch(square_blob, ivec3((sx + 1) * 2 + 1, sy, gz), 0));
+            mat2x4 v2 = mat2x4(texelFetch(square_blob, ivec3(sx * 2, sy + 1, gz), 0), texelFetch(square_blob, ivec3(sx * 2 + 1, sy + 1, gz), 0));
+            mat2x4 v3 = mat2x4(texelFetch(square_blob, ivec3((sx + 1) * 2, sy + 1, gz), 0), texelFetch(square_blob, ivec3((sx + 1) * 2 + 1, sy + 1, gz), 0));
+
+            sum = v0 + v1 + v2 + v3;
+        }
+    }
+#else
     int sx = gx * 4;
 
     int v_offset = gz * p.cstep + sx;
 
-    if (sx == p.size - 1)
+    if (sx == p.w - 1)
     {
-        sqsum_blob_data[gi] = square_blob_data[v_offset];
+        sum = square_blob_data[v_offset];
     }
-    else if (sx == p.size - 2)
+    else if (sx == p.w - 2)
     {
         mat2x4 v0 = square_blob_data[v_offset];
         mat2x4 v1 = square_blob_data[v_offset + 1];
 
-        mat2x4 sum = v0 + v1;
-
-        sqsum_blob_data[gi] = sum;
+        sum = v0 + v1;
     }
-    else if (sx == p.size - 3)
+    else if (sx == p.w - 3)
     {
         mat2x4 v0 = square_blob_data[v_offset];
         mat2x4 v1 = square_blob_data[v_offset + 1];
         mat2x4 v2 = square_blob_data[v_offset + 2];
 
-        mat2x4 sum = v0 + v1 + v2;
-
-        sqsum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2;
     }
     else
     {
@@ -85,8 +128,16 @@ void main()
         mat2x4 v2 = square_blob_data[v_offset + 2];
         mat2x4 v3 = square_blob_data[v_offset + 3];
 
-        mat2x4 sum = v0 + v1 + v2 + v3;
-
-        sqsum_blob_data[gi] = sum;
+        sum = v0 + v1 + v2 + v3;
     }
+#endif
+
+#if NCNN_image_shader
+    imageStore(sqsum_blob, ivec3(gx * 2, gy, gz), sum[0]);
+    imageStore(sqsum_blob, ivec3(gx * 2 + 1, gy, gz), sum[1]);
+#else
+    int gi = gz * p.outcstep + gx;
+
+    sqsum_blob_data[gi] = sum;
+#endif
 }

src/layer/vulkan/shader/instancenorm_sub_mean_square.comp

@@ -38,9 +38,15 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform unfp sampler3D bottom_top_blob;
+layout (binding = 1) uniform unfp sampler1D mean_blob;
+layout (binding = 2, r32f) writeonly uniform highp image3D square_blob;
+#else
 layout (binding = 0) readonly buffer bottom_top_blob { sfp bottom_top_blob_data[]; };
 layout (binding = 1) readonly buffer mean_blob { sfp mean_data[]; };
 layout (binding = 2) writeonly buffer square_blob { float square_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
@@ -66,15 +72,24 @@ void main()
     if (gx >= psc(outw) || gy >= psc(outh) || gz >= psc(outc))
         return;
 
-    const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;
-
+#if NCNN_image_shader
+    afp v = image3d_ld1(bottom_top_blob, ivec3(gx, gy, gz));
+    afp mean = image1d_ld1(mean_blob, gz);
+#else
     int v_offset = gz * psc(cstep) + gy * psc(w) + gx;
 
     afp v = buffer_ld1(bottom_top_blob_data, v_offset);
     afp mean = buffer_ld1(mean_data, gz);
+#endif
 
     v = v - mean;
     v = v * v;
 
+#if NCNN_image_shader
+    imageStore(square_blob, ivec3(gx, gy, gz), vec4(float(v)));
+#else
+    const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;
+
     square_blob_data[gi] = float(v);
+#endif
 }

src/layer/vulkan/shader/instancenorm_sub_mean_square_pack4.comp

@@ -38,9 +38,15 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform unfp sampler3D bottom_top_blob;
+layout (binding = 1) uniform unfp sampler1D mean_blob;
+layout (binding = 2, rgba32f) writeonly uniform highp image3D square_blob;
+#else
 layout (binding = 0) readonly buffer bottom_top_blob { sfpvec4 bottom_top_blob_data[]; };
 layout (binding = 1) readonly buffer mean_blob { sfpvec4 mean_data[]; };
 layout (binding = 2) writeonly buffer square_blob { vec4 square_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
@@ -66,15 +72,24 @@ void main()
     if (gx >= psc(outw) || gy >= psc(outh) || gz >= psc(outc))
         return;
 
-    const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;
-
+#if NCNN_image_shader
+    afpvec4 v = image3d_ld4(bottom_top_blob, ivec3(gx, gy, gz));
+    afpvec4 mean = image1d_ld4(mean_blob, gz);
+#else
     int v_offset = gz * psc(cstep) + gy * psc(w) + gx;
 
     afpvec4 v = buffer_ld4(bottom_top_blob_data, v_offset);
     afpvec4 mean = buffer_ld4(mean_data, gz);
+#endif
 
     v = v - mean;
     v = v * v;
 
+#if NCNN_image_shader
+    imageStore(square_blob, ivec3(gx, gy, gz), vec4(v));
+#else
+    const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;
+
     square_blob_data[gi] = vec4(v);
+#endif
 }

src/layer/vulkan/shader/instancenorm_sub_mean_square_pack8.comp

@@ -39,9 +39,15 @@ layout (local_size_x_id = 233) in;
 layout (local_size_y_id = 234) in;
 layout (local_size_z_id = 235) in;
 
+#if NCNN_image_shader
+layout (binding = 0) uniform unfp sampler3D bottom_top_blob;
+layout (binding = 1) uniform unfp sampler1D mean_blob;
+layout (binding = 2, rgba32f) writeonly uniform highp image3D square_blob;
+#else
 layout (binding = 0) readonly buffer bottom_top_blob { sfpvec8 bottom_top_blob_data[]; };
 layout (binding = 1) readonly buffer mean_blob { sfpvec8 mean_data[]; };
 layout (binding = 2) writeonly buffer square_blob { mat2x4 square_blob_data[]; };
+#endif
 
 layout (push_constant) uniform parameter
 {
@@ -67,17 +73,27 @@ void main()
     if (gx >= psc(outw) || gy >= psc(outh) || gz >= psc(outc))
         return;
 
-    const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;
-
+#if NCNN_image_shader
+    afpvec8 v = image3d_ld8(bottom_top_blob, ivec3(gx, gy, gz));
+    afpvec8 mean = image1d_ld8(mean_blob, gz);
+#else
     int v_offset = gz * psc(cstep) + gy * psc(w) + gx;
 
     afpvec8 v = buffer_ld8(bottom_top_blob_data, v_offset);
     afpvec8 mean = buffer_ld8(mean_data, gz);
+#endif
 
     v[0] = v[0] - mean[0];
     v[1] = v[1] - mean[1];
     v[0] = v[0] * v[0];
     v[1] = v[1] * v[1];
 
+#if NCNN_image_shader
+    imageStore(square_blob, ivec3(gx * 2, gy, gz), vec4(v[0]));
+    imageStore(square_blob, ivec3(gx * 2 + 1, gy, gz), vec4(v[1]));
+#else
+    const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;
+
     square_blob_data[gi] = mat2x4(v);
+#endif
 }