diff --git a/src/layer/vulkan/convolution_vulkan.cpp b/src/layer/vulkan/convolution_vulkan.cpp index db6926e43..795c705b0 100644 --- a/src/layer/vulkan/convolution_vulkan.cpp +++ b/src/layer/vulkan/convolution_vulkan.cpp @@ -178,12 +178,31 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) if (opt.use_winograd_convolution && (opt.use_winograd23_convolution || opt.use_winograd43_convolution) && is_conv3x3s1d1 && num_input >= 16 && num_output >= 16) { - bool use_cooperative_matrix_16_8_8 = vkdev->info.support_cooperative_matrix_16_8_8() && opt.use_cooperative_matrix && !opt.use_shader_pack8 && opt.use_fp16_storage && num_input % 8 == 0 && num_output % 8 == 0; - bool use_cooperative_matrix_16_16_16 = vkdev->info.support_cooperative_matrix_16_16_16() && opt.use_cooperative_matrix && !opt.use_shader_pack8 && opt.use_fp16_storage && num_input % 16 == 0 && num_output % 16 == 0; - if (vkdev->info.subgroup_size() != 32 && (!vkdev->info.support_subgroup_size_control() || vkdev->info.min_subgroup_size() > 32 || vkdev->info.max_subgroup_size() < 32)) + use_cooperative_matrix = vkdev->info.support_cooperative_matrix() && opt.use_cooperative_matrix && !opt.use_shader_pack8 && (opt.use_fp16_storage || opt.use_fp16_packed); + + if (use_cooperative_matrix) { - use_cooperative_matrix_16_8_8 = false; - use_cooperative_matrix_16_16_16 = false; + int size = 1024; + // f43 and f23 share the same size parameter, set zero for dynamic dispatch + // if (out_shape.dims != 0) + // { + // int block_x = (out_shape.w + 3) / 4; + // int block_y = (out_shape.h + 3) / 4; + // int block_x = (out_shape.w + 1) / 2; + // int block_y = (out_shape.h + 1) / 2; + // size = block_x * block_y; + // } + + vkdev->info.get_optimal_cooperative_matrix_mnk(size, num_output, num_input, VK_COMPONENT_TYPE_FLOAT16_KHR, opt.use_fp16_arithmetic ? VK_COMPONENT_TYPE_FLOAT16_KHR : VK_COMPONENT_TYPE_FLOAT32_KHR, VK_SCOPE_SUBGROUP_KHR, coopmat_M, coopmat_N, coopmat_K); + + // assert coopmat_M != 0 && coopmat_N != 0 && coopmat_K != 0 + + UNROLL_SG_M = std::min((size + coopmat_M - 1) / coopmat_M, 2); + UNROLL_SG_N = std::min((num_output + coopmat_N - 1) / coopmat_N, 2); + UNROLL_SG_K = std::min((num_input + coopmat_K - 1) / coopmat_K, 2); + + UNROLL_WG_M = std::min((size + coopmat_M * UNROLL_SG_M - 1) / (coopmat_M * UNROLL_SG_M), 2); + UNROLL_WG_N = std::min((num_output + coopmat_N * UNROLL_SG_N - 1) / (coopmat_N * UNROLL_SG_N), 2); } // winograd43 transform kernel @@ -237,54 +256,123 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) } } - if (use_cooperative_matrix_16_8_8) + if (use_cooperative_matrix) { - // src = 36-inch-outch - // dst = 8b-8a-inch/8a-outch/8b-36 - weight_winograd43_data_packed.create(num_input / 8, num_output / 8, 36, (size_t)4 * 8 * 8, 8 * 8); - + // from 36-inch-outch to inch-outch-36 + Mat weight_data_tm_r2(num_input, num_output, 36); for (int k = 0; k < 36; k++) { - float* g00 = weight_winograd43_data_packed.channel(k); + float* g00 = weight_data_tm_r2.channel(k); - for (int q = 0; q + (8 - 1) < num_output; q += 8) + for (int q = 0; q < num_output; q++) { - for (int p = 0; p + (8 - 1) < num_input; p += 8) + for (int p = 0; p < num_input; p++) { - for (int i = 0; i < 8; i++) - { - for (int j = 0; j < 8; j++) - { - const float* k00 = weight_data_tm.channel(q + j).row(p + i); - g00[0] = k00[k]; - g00++; - } - } + *g00++ = weight_data_tm[(q * num_input + p) * 36 + k]; } } } - } - else if (use_cooperative_matrix_16_16_16) - { - // src = 36-inch-outch - // dst = 16b-16a-inch/16a-outch/16b-36 - weight_winograd43_data_packed.create(num_input / 16, num_output / 16, 36, (size_t)4 * 16 * 16, 16 * 16); - for (int k = 0; k < 36; k++) - { - float* g00 = weight_winograd43_data_packed.channel(k); + // +-N-+ + // K | + // +SG_UN + // | | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UN +---+ + // | | | + // | +- -+ + // | | | + // v +---+ - for (int q = 0; q + (16 - 1) < num_output; q += 16) + // +-N-+ + // K | + // +SG_UN + // | | + // ^ +---+ + // | | | + // WG_UN+- -+ + // | | | + // v +---+ + + const int blocks_n = (num_output + coopmat_N * UNROLL_SG_N * UNROLL_WG_N - 1) / (coopmat_N * UNROLL_SG_N * UNROLL_WG_N); + // const int blocks_k = (num_input + coopmat_K * UNROLL_SG_K - 1) / (coopmat_K * UNROLL_SG_K); + const int kk = (num_input + coopmat_K - 1) / coopmat_K; + + weight_winograd43_data_packed.create(coopmat_N * coopmat_K * UNROLL_SG_N * UNROLL_WG_N * kk, blocks_n, 36); + for (int b = 0; b < 36; b++) + { + for (int bn = 0; bn < blocks_n; bn++) { - for (int p = 0; p + (16 - 1) < num_input; p += 16) + float* p = weight_winograd43_data_packed.channel(b).row(bn); + + int k = 0; + for (; k + UNROLL_SG_K - 1 < kk; k += UNROLL_SG_K) { - for (int i = 0; i < 16; i++) + // const int ki = k * coopmat_K; + + for (int wn = 0; wn < UNROLL_WG_N; wn++) { - for (int j = 0; j < 16; j++) + for (int zk = 0; zk < UNROLL_SG_K; zk++) { - const float* k00 = weight_data_tm.channel(q + j).row(p + i); - g00[0] = k00[k]; - g00++; + for (int zn = 0; zn < UNROLL_SG_N; zn++) + { + for (int i = 0; i < coopmat_K; i++) + { + for (int j = 0; j < coopmat_N; j++) + { + const int gni = ((bn * UNROLL_WG_N + wn) * UNROLL_SG_N + zn) * coopmat_N + j; + const int gki = (k + zk) * coopmat_K + i; + + if (gni < num_output && gki < num_input) + { + *p++ = weight_data_tm_r2.channel(b)[gni * num_input + gki]; + } + else + { + *p++ = 0.f; + } + } + } + } + } + } + } + for (; k < kk; k++) + { + // const int ki = k * coopmat_K; + + for (int wn = 0; wn < UNROLL_WG_N; wn++) + { + // for (int zk = 0; zk < UNROLL_SG_K; zk++) + { + for (int zn = 0; zn < UNROLL_SG_N; zn++) + { + for (int i = 0; i < coopmat_K; i++) + { + for (int j = 0; j < coopmat_N; j++) + { + const int gni = ((bn * UNROLL_WG_N + wn) * UNROLL_SG_N + zn) * coopmat_N + j; + // const int gki = (k + zk) * coopmat_K + i; + const int gki = k * coopmat_K + i; + + if (gni < num_output && gki < num_input) + { + *p++ = weight_data_tm_r2.channel(b)[gni * num_input + gki]; + } + else + { + *p++ = 0.f; + } + } + } + } } } } @@ -334,26 +422,26 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) if (out_shape.dims != 0) { - int block_x = (out_shape.w + 3) / 4; - int block_y = (out_shape.h + 3) / 4; + block_x = (out_shape.w + 3) / 4; + block_y = (out_shape.h + 3) / 4; - shape_winograd_input_transformed = Mat(block_x * block_y, shape.c, 36, (void*)0); - shape_winograd_gemm = Mat(block_x * block_y, out_shape.c, 36, (void*)0); + shape_winograd_input_transformed = Mat(block_x * block_y, 1, shape.c * 36, (void*)0); + shape_winograd_gemm = Mat(block_x * block_y, 1, out_shape.c * 36, (void*)0); } - if (shape_winograd_input_transformed.dims == 3) shape_winograd_input_transformed_packed = Mat(shape_winograd_input_transformed.w, shape_winograd_input_transformed.h / elempack, 36, (void*)0, elemsize, elempack); + if (shape_winograd_input_transformed.dims == 3) shape_winograd_input_transformed_packed = Mat(shape_winograd_input_transformed.w, 1, shape_winograd_input_transformed.h / elempack * 36, (void*)0, elemsize, elempack); - if (shape_winograd_gemm.dims == 3) shape_winograd_gemm_packed = Mat(shape_winograd_gemm.w, shape_winograd_gemm.h / out_elempack, 36, (void*)0, out_elemsize, out_elempack); + if (shape_winograd_gemm.dims == 3) shape_winograd_gemm_packed = Mat(shape_winograd_gemm.w, 1, shape_winograd_gemm.h / out_elempack * 36, (void*)0, out_elemsize, out_elempack); { - std::vector specializations(0 + 7); - specializations[0 + 0].i = shape_bordered_packed.w; - specializations[0 + 1].i = shape_bordered_packed.h; - specializations[0 + 2].i = shape_bordered_packed.c; - specializations[0 + 3].i = shape_bordered_packed.cstep; - specializations[0 + 4].i = shape_winograd_input_transformed_packed.cstep; - specializations[0 + 5].i = block_x; - specializations[0 + 6].i = block_y; + std::vector specializations(1 + 6); + specializations[0].i = num_input / elempack; + specializations[1 + 0].i = shape_bordered_packed.w; + specializations[1 + 1].i = shape_bordered_packed.h; + specializations[1 + 2].i = shape_bordered_packed.cstep; + specializations[1 + 3].i = shape_winograd_input_transformed_packed.cstep; + specializations[1 + 4].i = block_x; + specializations[1 + 5].i = block_y; int shader_type_index = -1; if (elempack == 1) shader_type_index = LayerShaderType::convolution_3x3s1d1_winograd43_transform_input; @@ -365,14 +453,45 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) pipeline_convolution_3x3s1d1_winograd43_transform_input->create(shader_type_index, opt, specializations); } + if (use_cooperative_matrix) { - std::vector specializations(1 + 5); + Mat weight_winograd43_data_packed_fp16 = Mat(weight_winograd43_data_packed.w, weight_winograd43_data_packed.h, weight_winograd43_data_packed.c, (void*)0, 2u, 1); + + std::vector specializations(14 + 3); + specializations[0].u32 = 36; //batch + specializations[1].u32 = coopmat_M; + specializations[2].u32 = coopmat_N; + specializations[3].u32 = coopmat_K; + specializations[4].u32 = UNROLL_SG_M; + specializations[5].u32 = UNROLL_SG_N; + specializations[6].u32 = UNROLL_SG_K; + specializations[7].u32 = UNROLL_WG_M; + specializations[8].u32 = UNROLL_WG_N; + specializations[9].u32 = num_input; + specializations[10].u32 = num_output; + specializations[11].u32 = elempack; + specializations[12].u32 = out_elempack; + specializations[13].u32 = weight_winograd43_data_packed_fp16.cstep; + specializations[14 + 0].u32 = shape_winograd_input_transformed_packed.w; + specializations[14 + 1].u32 = shape_winograd_input_transformed_packed.cstep; + specializations[14 + 2].u32 = shape_winograd_gemm_packed.cstep; + + const int subgroup_size = vkdev->info.subgroup_size(); + + pipeline_convolution_3x3s1d1_winograd43_gemm = new Pipeline(vkdev); + pipeline_convolution_3x3s1d1_winograd43_gemm->set_subgroup_size(subgroup_size); + pipeline_convolution_3x3s1d1_winograd43_gemm->set_local_size_xyz(subgroup_size * UNROLL_WG_M * UNROLL_WG_N, 1, 1); + pipeline_convolution_3x3s1d1_winograd43_gemm->create(LayerShaderType::convolution_winograd_gemm_cm, opt, specializations); + } + else + { + std::vector specializations(3 + 3); specializations[0].i = 36; - specializations[1 + 0].i = shape_winograd_input_transformed_packed.h; - specializations[1 + 1].i = shape_winograd_input_transformed_packed.cstep; - specializations[1 + 2].i = shape_winograd_gemm_packed.w; - specializations[1 + 3].i = shape_winograd_gemm_packed.h; - specializations[1 + 4].i = shape_winograd_gemm_packed.cstep; + specializations[1].i = num_input / elempack; + specializations[2].i = num_output / out_elempack; + specializations[3 + 0].i = shape_winograd_input_transformed_packed.cstep; + specializations[3 + 1].i = shape_winograd_gemm_packed.w; + specializations[3 + 2].i = shape_winograd_gemm_packed.cstep; int shader_type_index = -1; if (elempack == 1 && out_elempack == 1) shader_type_index = LayerShaderType::convolution_3x3s1d1_winograd_gemm; @@ -385,27 +504,8 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) if (elempack == 4 && out_elempack == 8) shader_type_index = LayerShaderType::convolution_pack4to8_3x3s1d1_winograd_gemm; if (elempack == 8 && out_elempack == 4) shader_type_index = LayerShaderType::convolution_pack8to4_3x3s1d1_winograd_gemm; - if (use_cooperative_matrix_16_8_8) - { - shader_type_index = LayerShaderType::convolution_pack4_3x3s1d1_winograd_gemm_cm_16_8_8; - } - else if (use_cooperative_matrix_16_16_16) - { - shader_type_index = LayerShaderType::convolution_pack4_3x3s1d1_winograd_gemm_cm_16_16_16; - } - pipeline_convolution_3x3s1d1_winograd43_gemm = new Pipeline(vkdev); - if (use_cooperative_matrix_16_8_8) - { - pipeline_convolution_3x3s1d1_winograd43_gemm->set_subgroup_size(32); - pipeline_convolution_3x3s1d1_winograd43_gemm->set_local_size_xyz(32, 1, 1); - } - else if (use_cooperative_matrix_16_16_16) - { - pipeline_convolution_3x3s1d1_winograd43_gemm->set_subgroup_size(32); - pipeline_convolution_3x3s1d1_winograd43_gemm->set_local_size_xyz(32, 1, 1); - } - else if (opt.use_shader_local_memory) + if (opt.use_shader_local_memory) { pipeline_convolution_3x3s1d1_winograd43_gemm->set_local_size_xyz(8, 8, 1); } @@ -417,18 +517,18 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) } { - std::vector specializations(4 + 7); + std::vector specializations(5 + 6); specializations[0].i = bias_term; specializations[1].i = activation_type; specializations[2].f = activation_params.w >= 1 ? activation_params[0] : 0.f; specializations[3].f = activation_params.w == 2 ? activation_params[1] : 0.f; - specializations[4 + 0].i = shape_winograd_gemm_packed.h; - specializations[4 + 1].i = shape_winograd_gemm_packed.cstep; - specializations[4 + 2].i = block_x; - specializations[4 + 3].i = block_y; - specializations[4 + 4].i = out_shape_packed.w; - specializations[4 + 5].i = out_shape_packed.h; - specializations[4 + 6].i = out_shape_packed.cstep; + specializations[4].i = num_output / out_elempack; + specializations[5 + 0].i = shape_winograd_gemm_packed.cstep; + specializations[5 + 1].i = block_x; + specializations[5 + 2].i = block_y; + specializations[5 + 3].i = out_shape_packed.w; + specializations[5 + 4].i = out_shape_packed.h; + specializations[5 + 5].i = out_shape_packed.cstep; int shader_type_index = -1; if (out_elempack == 1) shader_type_index = LayerShaderType::convolution_3x3s1d1_winograd43_transform_output; @@ -490,54 +590,123 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) } } - if (use_cooperative_matrix_16_8_8) + if (use_cooperative_matrix) { - // src = 16-inch-outch - // dst = 8b-8a-inch/8a-outch/8b-16 - weight_winograd23_data_packed.create(num_input / 8, num_output / 8, 16, (size_t)4 * 8 * 8, 8 * 8); - + // from 16-inch-outch to inch-outch-16 + Mat weight_data_tm_r2(num_input, num_output, 16); for (int k = 0; k < 16; k++) { - float* g00 = weight_winograd23_data_packed.channel(k); + float* g00 = weight_data_tm_r2.channel(k); - for (int q = 0; q + (8 - 1) < num_output; q += 8) + for (int q = 0; q < num_output; q++) { - for (int p = 0; p + (8 - 1) < num_input; p += 8) + for (int p = 0; p < num_input; p++) { - for (int i = 0; i < 8; i++) - { - for (int j = 0; j < 8; j++) - { - const float* k00 = weight_data_tm.channel(q + j).row(p + i); - g00[0] = k00[k]; - g00++; - } - } + *g00++ = weight_data_tm[(q * num_input + p) * 16 + k]; } } } - } - else if (use_cooperative_matrix_16_16_16) - { - // src = 16-inch-outch - // dst = 16b-16a-inch/16a-outch/16b-16 - weight_winograd23_data_packed.create(num_input / 16, num_output / 16, 16, (size_t)4 * 16 * 16, 16 * 16); - for (int k = 0; k < 16; k++) - { - float* g00 = weight_winograd23_data_packed.channel(k); + // +-N-+ + // K | + // +SG_UN + // | | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UN +---+ + // | | | + // | +- -+ + // | | | + // v +---+ - for (int q = 0; q + (16 - 1) < num_output; q += 16) + // +-N-+ + // K | + // +SG_UN + // | | + // ^ +---+ + // | | | + // WG_UN+- -+ + // | | | + // v +---+ + + const int blocks_n = (num_output + coopmat_N * UNROLL_SG_N * UNROLL_WG_N - 1) / (coopmat_N * UNROLL_SG_N * UNROLL_WG_N); + // const int blocks_k = (num_input + coopmat_K * UNROLL_SG_K - 1) / (coopmat_K * UNROLL_SG_K); + const int kk = (num_input + coopmat_K - 1) / coopmat_K; + + weight_winograd23_data_packed.create(coopmat_N * coopmat_K * UNROLL_SG_N * UNROLL_WG_N * kk, blocks_n, 16); + for (int b = 0; b < 16; b++) + { + for (int bn = 0; bn < blocks_n; bn++) { - for (int p = 0; p + (16 - 1) < num_input; p += 16) + float* p = weight_winograd23_data_packed.channel(b).row(bn); + + int k = 0; + for (; k + UNROLL_SG_K - 1 < kk; k += UNROLL_SG_K) { - for (int i = 0; i < 16; i++) + // const int ki = k * coopmat_K; + + for (int wn = 0; wn < UNROLL_WG_N; wn++) { - for (int j = 0; j < 16; j++) + for (int zk = 0; zk < UNROLL_SG_K; zk++) { - const float* k00 = weight_data_tm.channel(q + j).row(p + i); - g00[0] = k00[k]; - g00++; + for (int zn = 0; zn < UNROLL_SG_N; zn++) + { + for (int i = 0; i < coopmat_K; i++) + { + for (int j = 0; j < coopmat_N; j++) + { + const int gni = ((bn * UNROLL_WG_N + wn) * UNROLL_SG_N + zn) * coopmat_N + j; + const int gki = (k + zk) * coopmat_K + i; + + if (gni < num_output && gki < num_input) + { + *p++ = weight_data_tm_r2.channel(b)[gni * num_input + gki]; + } + else + { + *p++ = 0.f; + } + } + } + } + } + } + } + for (; k < kk; k++) + { + // const int ki = k * coopmat_K; + + for (int wn = 0; wn < UNROLL_WG_N; wn++) + { + // for (int zk = 0; zk < UNROLL_SG_K; zk++) + { + for (int zn = 0; zn < UNROLL_SG_N; zn++) + { + for (int i = 0; i < coopmat_K; i++) + { + for (int j = 0; j < coopmat_N; j++) + { + const int gni = ((bn * UNROLL_WG_N + wn) * UNROLL_SG_N + zn) * coopmat_N + j; + // const int gki = (k + zk) * coopmat_K + i; + const int gki = k * coopmat_K + i; + + if (gni < num_output && gki < num_input) + { + *p++ = weight_data_tm_r2.channel(b)[gni * num_input + gki]; + } + else + { + *p++ = 0.f; + } + } + } + } } } } @@ -587,26 +756,26 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) if (out_shape.dims != 0) { - int block_x = (out_shape.w + 1) / 2; - int block_y = (out_shape.h + 1) / 2; + block_x = (out_shape.w + 1) / 2; + block_y = (out_shape.h + 1) / 2; - shape_winograd_input_transformed = Mat(block_x * block_y, shape.c, 16, (void*)0); - shape_winograd_gemm = Mat(block_x * block_y, out_shape.c, 16, (void*)0); + shape_winograd_input_transformed = Mat(block_x * block_y, 1, shape.c * 16, (void*)0); + shape_winograd_gemm = Mat(block_x * block_y, 1, out_shape.c * 16, (void*)0); } - if (shape_winograd_input_transformed.dims == 3) shape_winograd_input_transformed_packed = Mat(shape_winograd_input_transformed.w, shape_winograd_input_transformed.h / elempack, 16, (void*)0, elemsize, elempack); + if (shape_winograd_input_transformed.dims == 3) shape_winograd_input_transformed_packed = Mat(shape_winograd_input_transformed.w, 1, shape_winograd_input_transformed.h / elempack * 16, (void*)0, elemsize, elempack); - if (shape_winograd_gemm.dims == 3) shape_winograd_gemm_packed = Mat(shape_winograd_gemm.w, shape_winograd_gemm.h / out_elempack, 16, (void*)0, out_elemsize, out_elempack); + if (shape_winograd_gemm.dims == 3) shape_winograd_gemm_packed = Mat(shape_winograd_gemm.w, 1, shape_winograd_gemm.h / out_elempack * 16, (void*)0, out_elemsize, out_elempack); { - std::vector specializations(0 + 7); - specializations[0 + 0].i = shape_bordered_packed.w; - specializations[0 + 1].i = shape_bordered_packed.h; - specializations[0 + 2].i = shape_bordered_packed.c; - specializations[0 + 3].i = shape_bordered_packed.cstep; - specializations[0 + 4].i = shape_winograd_input_transformed_packed.cstep; - specializations[0 + 5].i = block_x; - specializations[0 + 6].i = block_y; + std::vector specializations(1 + 6); + specializations[0].i = num_input / elempack; + specializations[1 + 0].i = shape_bordered_packed.w; + specializations[1 + 1].i = shape_bordered_packed.h; + specializations[1 + 2].i = shape_bordered_packed.cstep; + specializations[1 + 3].i = shape_winograd_input_transformed_packed.cstep; + specializations[1 + 4].i = block_x; + specializations[1 + 5].i = block_y; int shader_type_index = -1; if (elempack == 1) shader_type_index = LayerShaderType::convolution_3x3s1d1_winograd23_transform_input; @@ -618,14 +787,45 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) pipeline_convolution_3x3s1d1_winograd23_transform_input->create(shader_type_index, opt, specializations); } + if (use_cooperative_matrix) { - std::vector specializations(1 + 5); + Mat weight_winograd23_data_packed_fp16 = Mat(weight_winograd23_data_packed.w, weight_winograd23_data_packed.h, weight_winograd23_data_packed.c, (void*)0, 2u, 1); + + std::vector specializations(14 + 3); + specializations[0].u32 = 16; //batch + specializations[1].u32 = coopmat_M; + specializations[2].u32 = coopmat_N; + specializations[3].u32 = coopmat_K; + specializations[4].u32 = UNROLL_SG_M; + specializations[5].u32 = UNROLL_SG_N; + specializations[6].u32 = UNROLL_SG_K; + specializations[7].u32 = UNROLL_WG_M; + specializations[8].u32 = UNROLL_WG_N; + specializations[9].u32 = num_input; + specializations[10].u32 = num_output; + specializations[11].u32 = elempack; + specializations[12].u32 = out_elempack; + specializations[13].u32 = weight_winograd23_data_packed_fp16.cstep; + specializations[14 + 0].u32 = shape_winograd_input_transformed_packed.w; + specializations[14 + 1].u32 = shape_winograd_input_transformed_packed.cstep; + specializations[14 + 2].u32 = shape_winograd_gemm_packed.cstep; + + const int subgroup_size = vkdev->info.subgroup_size(); + + pipeline_convolution_3x3s1d1_winograd23_gemm = new Pipeline(vkdev); + pipeline_convolution_3x3s1d1_winograd23_gemm->set_subgroup_size(subgroup_size); + pipeline_convolution_3x3s1d1_winograd23_gemm->set_local_size_xyz(subgroup_size * UNROLL_WG_M * UNROLL_WG_N, 1, 1); + pipeline_convolution_3x3s1d1_winograd23_gemm->create(LayerShaderType::convolution_winograd_gemm_cm, opt, specializations); + } + else + { + std::vector specializations(3 + 3); specializations[0].i = 16; - specializations[1 + 0].i = shape_winograd_input_transformed_packed.h; - specializations[1 + 1].i = shape_winograd_input_transformed_packed.cstep; - specializations[1 + 2].i = shape_winograd_gemm_packed.w; - specializations[1 + 3].i = shape_winograd_gemm_packed.h; - specializations[1 + 4].i = shape_winograd_gemm_packed.cstep; + specializations[1].i = num_input / elempack; + specializations[2].i = num_output / out_elempack; + specializations[3 + 0].i = shape_winograd_input_transformed_packed.cstep; + specializations[3 + 1].i = shape_winograd_gemm_packed.w; + specializations[3 + 2].i = shape_winograd_gemm_packed.cstep; int shader_type_index = -1; if (elempack == 1 && out_elempack == 1) shader_type_index = LayerShaderType::convolution_3x3s1d1_winograd_gemm; @@ -638,27 +838,8 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) if (elempack == 4 && out_elempack == 8) shader_type_index = LayerShaderType::convolution_pack4to8_3x3s1d1_winograd_gemm; if (elempack == 8 && out_elempack == 4) shader_type_index = LayerShaderType::convolution_pack8to4_3x3s1d1_winograd_gemm; - if (use_cooperative_matrix_16_8_8) - { - shader_type_index = LayerShaderType::convolution_pack4_3x3s1d1_winograd_gemm_cm_16_8_8; - } - else if (use_cooperative_matrix_16_16_16) - { - shader_type_index = LayerShaderType::convolution_pack4_3x3s1d1_winograd_gemm_cm_16_16_16; - } - pipeline_convolution_3x3s1d1_winograd23_gemm = new Pipeline(vkdev); - if (use_cooperative_matrix_16_8_8) - { - pipeline_convolution_3x3s1d1_winograd23_gemm->set_subgroup_size(32); - pipeline_convolution_3x3s1d1_winograd23_gemm->set_local_size_xyz(32, 1, 1); - } - else if (use_cooperative_matrix_16_16_16) - { - pipeline_convolution_3x3s1d1_winograd23_gemm->set_subgroup_size(32); - pipeline_convolution_3x3s1d1_winograd23_gemm->set_local_size_xyz(32, 1, 1); - } - else if (opt.use_shader_local_memory) + if (opt.use_shader_local_memory) { pipeline_convolution_3x3s1d1_winograd23_gemm->set_local_size_xyz(8, 8, 1); } @@ -670,18 +851,18 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) } { - std::vector specializations(4 + 7); + std::vector specializations(5 + 6); specializations[0].i = bias_term; specializations[1].i = activation_type; specializations[2].f = activation_params.w >= 1 ? activation_params[0] : 0.f; specializations[3].f = activation_params.w == 2 ? activation_params[1] : 0.f; - specializations[4 + 0].i = shape_winograd_gemm_packed.h; - specializations[4 + 1].i = shape_winograd_gemm_packed.cstep; - specializations[4 + 2].i = block_x; - specializations[4 + 3].i = block_y; - specializations[4 + 4].i = out_shape_packed.w; - specializations[4 + 5].i = out_shape_packed.h; - specializations[4 + 6].i = out_shape_packed.cstep; + specializations[4].i = num_output / out_elempack; + specializations[5 + 0].i = shape_winograd_gemm_packed.cstep; + specializations[5 + 1].i = block_x; + specializations[5 + 2].i = block_y; + specializations[5 + 3].i = out_shape_packed.w; + specializations[5 + 4].i = out_shape_packed.h; + specializations[5 + 5].i = out_shape_packed.cstep; int shader_type_index = -1; if (out_elempack == 1) shader_type_index = LayerShaderType::convolution_3x3s1d1_winograd23_transform_output; @@ -1029,96 +1210,27 @@ int Convolution_vulkan::create_pipeline(const Option& _opt) } else { - // src = kw-kh-inch-outch - // dst = pa-pb-kw-kh-inch/pa-outch/pb - bool use_cooperative_matrix_16_8_8 = vkdev->info.support_cooperative_matrix_16_8_8() && opt.use_cooperative_matrix && is_conv1x1s1d1 && !opt.use_shader_pack8 && opt.use_fp16_storage && num_input % 8 == 0 && num_output % 8 == 0; - bool use_cooperative_matrix_16_16_16 = vkdev->info.support_cooperative_matrix_16_16_16() && opt.use_cooperative_matrix && is_conv1x1s1d1 && !opt.use_shader_pack8 && opt.use_fp16_storage && num_input % 16 == 0 && num_output % 16 == 0; - if (vkdev->info.subgroup_size() != 32 && (!vkdev->info.support_subgroup_size_control() || vkdev->info.min_subgroup_size() > 32 || vkdev->info.max_subgroup_size() < 32)) - { - use_cooperative_matrix_16_8_8 = false; - use_cooperative_matrix_16_16_16 = false; - } + Mat weight_data_r2 = weight_data.reshape(maxk, num_input, num_output); - if (use_cooperative_matrix_16_8_8) - { - // dst = 8b-8a-inch/8a-outch/8b - Mat weight_data_r2 = weight_data.reshape(maxk, num_input, num_output); - - weight_data_packed.create(maxk, num_input / 8, num_output / 8, (size_t)4 * 8 * 8, 8 * 8); - - for (int q = 0; q + 7 < num_output; q += 8) - { - float* g00 = weight_data_packed.channel(q / 8); + weight_data_packed.create(maxk, num_input / elempack, num_output / out_elempack, (size_t)4 * elempack * out_elempack, elempack * out_elempack); - for (int p = 0; p + 7 < num_input; p += 8) - { - for (int k = 0; k < maxk; k++) - { - for (int i = 0; i < 8; i++) - { - for (int j = 0; j < 8; j++) - { - const float* k00 = weight_data_r2.channel(q + j).row(p + i); - g00[0] = k00[k]; - g00++; - } - } - } - } - } - } - else if (use_cooperative_matrix_16_16_16) + for (int q = 0; q + (out_elempack - 1) < num_output; q += out_elempack) { - // dst = 16b-16a-inch/16a-outch/16b - Mat weight_data_r2 = weight_data.reshape(maxk, num_input, num_output); + float* g00 = weight_data_packed.channel(q / out_elempack); - weight_data_packed.create(maxk, num_input / 16, num_output / 16, (size_t)4 * 16 * 16, 16 * 16); - - for (int q = 0; q + 15 < num_output; q += 16) + for (int p = 0; p + (elempack - 1) < num_input; p += elempack) { - float* g00 = weight_data_packed.channel(q / 16); - - for (int p = 0; p + 15 < num_input; p += 16) + for (int k = 0; k < maxk; k++) { - for (int k = 0; k < maxk; k++) + for (int i = 0; i < out_elempack; i++) { - for (int i = 0; i < 16; i++) - { - for (int j = 0; j < 16; j++) - { - const float* k00 = weight_data_r2.channel(q + j).row(p + i); - g00[0] = k00[k]; - g00++; - } - } - } - } - } - } - else - { - Mat weight_data_r2 = weight_data.reshape(maxk, num_input, num_output); + const Mat k0 = weight_data_r2.channel(q + i); - weight_data_packed.create(maxk, num_input / elempack, num_output / out_elempack, (size_t)4 * elempack * out_elempack, elempack * out_elempack); - - for (int q = 0; q + (out_elempack - 1) < num_output; q += out_elempack) - { - float* g00 = weight_data_packed.channel(q / out_elempack); - - for (int p = 0; p + (elempack - 1) < num_input; p += elempack) - { - for (int k = 0; k < maxk; k++) - { - for (int i = 0; i < out_elempack; i++) + for (int j = 0; j < elempack; j++) { - const Mat k0 = weight_data_r2.channel(q + i); - - for (int j = 0; j < elempack; j++) - { - const float* k00 = k0.row(p + j); - g00[0] = k00[k]; - g00++; - } + const float* k00 = k0.row(p + j); + g00[0] = k00[k]; + g00++; } } } @@ -1579,14 +1691,6 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom if (opt.use_winograd_convolution && (opt.use_winograd23_convolution || opt.use_winograd43_convolution) && is_conv3x3s1d1 && num_input >= 16 && num_output >= 16) { - bool use_cooperative_matrix_16_8_8 = vkdev->info.support_cooperative_matrix_16_8_8() && opt.use_cooperative_matrix && !opt.use_shader_pack8 && opt.use_fp16_storage && num_input % 8 == 0 && num_output % 8 == 0; - bool use_cooperative_matrix_16_16_16 = vkdev->info.support_cooperative_matrix_16_16_16() && opt.use_cooperative_matrix && !opt.use_shader_pack8 && opt.use_fp16_storage && num_input % 16 == 0 && num_output % 16 == 0; - if (vkdev->info.subgroup_size() != 32 && (!vkdev->info.support_subgroup_size_control() || vkdev->info.min_subgroup_size() > 32 || vkdev->info.max_subgroup_size() < 32)) - { - use_cooperative_matrix_16_8_8 = false; - use_cooperative_matrix_16_16_16 = false; - } - bool pre_winograd43 = opt.use_winograd43_convolution; if (opt.use_winograd23_convolution) { @@ -1595,9 +1699,7 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom if (vkdev->info.type() != 0 && (w <= 12 && h <= 12)) pre_winograd43 = false; - if (use_cooperative_matrix_16_8_8 && (w <= 18 && h <= 18)) - pre_winograd43 = false; - else if (use_cooperative_matrix_16_16_16 && (w <= 18 && h <= 18)) + if (use_cooperative_matrix && (w <= 18 && h <= 18)) pre_winograd43 = false; } @@ -1610,7 +1712,7 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom // transform input VkMat bottom_tm_blob; { - bottom_tm_blob.create(block_x * block_y, channels, 36, elemsize, elempack, opt.workspace_vkallocator); + bottom_tm_blob.create(block_x * block_y, 1, channels * 36, elemsize, elempack, opt.workspace_vkallocator); if (bottom_tm_blob.empty()) return -100; @@ -1618,19 +1720,18 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom bindings[0] = bottom_blob_bordered; bindings[1] = bottom_tm_blob; - std::vector constants(7); + std::vector constants(6); constants[0].i = bottom_blob_bordered.w; constants[1].i = bottom_blob_bordered.h; - constants[2].i = bottom_blob_bordered.c; - constants[3].i = bottom_blob_bordered.cstep; - constants[4].i = bottom_tm_blob.cstep; - constants[5].i = block_x; - constants[6].i = block_y; + constants[2].i = bottom_blob_bordered.cstep; + constants[3].i = bottom_tm_blob.cstep; + constants[4].i = block_x; + constants[5].i = block_y; VkMat dispatcher; dispatcher.w = block_x; dispatcher.h = block_y; - dispatcher.c = bottom_tm_blob.h; + dispatcher.c = channels; cmd.record_pipeline(pipeline_convolution_3x3s1d1_winograd43_transform_input, bindings, constants, dispatcher); } @@ -1638,41 +1739,53 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom // gemm VkMat top_tm_blob; { - top_tm_blob.create(block_x * block_y, num_output / out_elempack, 36, out_elemsize, out_elempack, opt.workspace_vkallocator); + top_tm_blob.create(block_x * block_y, 1, num_output / out_elempack * 36, out_elemsize, out_elempack, opt.workspace_vkallocator); if (top_tm_blob.empty()) return -100; - std::vector bindings(3); - bindings[0] = bottom_tm_blob; - bindings[1] = top_tm_blob; - bindings[2] = weight_data_gpu_tm_winograd43; + if (use_cooperative_matrix) + { + std::vector bindings(3); + bindings[0] = bottom_tm_blob; + bindings[1] = top_tm_blob; + bindings[2] = weight_data_gpu_tm_winograd43; - std::vector constants(5); - constants[0].i = bottom_tm_blob.h; - constants[1].i = bottom_tm_blob.cstep; - constants[2].i = top_tm_blob.w; - constants[3].i = top_tm_blob.h; - constants[4].i = top_tm_blob.cstep; + std::vector constants(3); + constants[0].i = bottom_tm_blob.w; + constants[1].i = bottom_tm_blob.cstep; + constants[2].i = top_tm_blob.cstep; - VkMat dispatcher; - dispatcher.w = (top_tm_blob.w + 3) / 4; - dispatcher.h = top_tm_blob.h; - dispatcher.c = 36; + const int blocks_x = (bottom_tm_blob.w + coopmat_M * UNROLL_SG_M * UNROLL_WG_M - 1) / (coopmat_M * UNROLL_SG_M * UNROLL_WG_M); + const int blocks_y = (num_output + coopmat_N * UNROLL_SG_N * UNROLL_WG_N - 1) / (coopmat_N * UNROLL_SG_N * UNROLL_WG_N); - if (use_cooperative_matrix_16_8_8) - { - dispatcher.w = ((top_tm_blob.w + 15) / 16 + 1) / 2 * 32; - dispatcher.h = ((top_tm_blob.h + 1) / 2 + 3) / 4; + const int subgroup_size = vkdev->info.subgroup_size(); + + VkMat dispatcher; + dispatcher.w = (blocks_x * blocks_y) * (subgroup_size * UNROLL_WG_M * UNROLL_WG_N); + dispatcher.h = 1; dispatcher.c = 36; + + cmd.record_pipeline(pipeline_convolution_3x3s1d1_winograd43_gemm, bindings, constants, dispatcher); } - else if (use_cooperative_matrix_16_16_16) + else { - dispatcher.w = ((top_tm_blob.w + 15) / 16 + 1) / 2 * 32; - dispatcher.h = ((top_tm_blob.h + 3) / 4 + 1) / 2; + std::vector bindings(3); + bindings[0] = bottom_tm_blob; + bindings[1] = top_tm_blob; + bindings[2] = weight_data_gpu_tm_winograd43; + + std::vector constants(3); + constants[0].i = bottom_tm_blob.cstep; + constants[1].i = top_tm_blob.w; + constants[2].i = top_tm_blob.cstep; + + VkMat dispatcher; + dispatcher.w = (top_tm_blob.w + 3) / 4; + dispatcher.h = num_output / out_elempack; dispatcher.c = 36; - } - cmd.record_pipeline(pipeline_convolution_3x3s1d1_winograd43_gemm, bindings, constants, dispatcher); + cmd.record_pipeline(pipeline_convolution_3x3s1d1_winograd43_gemm, bindings, constants, dispatcher); + } } // transform output @@ -1686,14 +1799,13 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom bindings[1] = top_blob; bindings[2] = bias_data_gpu; - std::vector constants(7); - constants[0].i = top_tm_blob.h; - constants[1].i = top_tm_blob.cstep; - constants[2].i = block_x; - constants[3].i = block_y; - constants[4].i = top_blob.w; - constants[5].i = top_blob.h; - constants[6].i = top_blob.cstep; + std::vector constants(6); + constants[0].i = top_tm_blob.cstep; + constants[1].i = block_x; + constants[2].i = block_y; + constants[3].i = top_blob.w; + constants[4].i = top_blob.h; + constants[5].i = top_blob.cstep; VkMat dispatcher; dispatcher.w = block_x; @@ -1712,7 +1824,7 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom // transform input VkMat bottom_tm_blob; { - bottom_tm_blob.create(block_x * block_y, channels, 16, elemsize, elempack, opt.workspace_vkallocator); + bottom_tm_blob.create(block_x * block_y, 1, channels * 16, elemsize, elempack, opt.workspace_vkallocator); if (bottom_tm_blob.empty()) return -100; @@ -1720,19 +1832,18 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom bindings[0] = bottom_blob_bordered; bindings[1] = bottom_tm_blob; - std::vector constants(7); + std::vector constants(6); constants[0].i = bottom_blob_bordered.w; constants[1].i = bottom_blob_bordered.h; - constants[2].i = bottom_blob_bordered.c; - constants[3].i = bottom_blob_bordered.cstep; - constants[4].i = bottom_tm_blob.cstep; - constants[5].i = block_x; - constants[6].i = block_y; + constants[2].i = bottom_blob_bordered.cstep; + constants[3].i = bottom_tm_blob.cstep; + constants[4].i = block_x; + constants[5].i = block_y; VkMat dispatcher; dispatcher.w = block_x; dispatcher.h = block_y; - dispatcher.c = bottom_tm_blob.h; + dispatcher.c = channels; cmd.record_pipeline(pipeline_convolution_3x3s1d1_winograd23_transform_input, bindings, constants, dispatcher); } @@ -1740,41 +1851,53 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom // gemm VkMat top_tm_blob; { - top_tm_blob.create(block_x * block_y, num_output / out_elempack, 16, out_elemsize, out_elempack, opt.workspace_vkallocator); + top_tm_blob.create(block_x * block_y, 1, num_output / out_elempack * 16, out_elemsize, out_elempack, opt.workspace_vkallocator); if (top_tm_blob.empty()) return -100; - std::vector bindings(3); - bindings[0] = bottom_tm_blob; - bindings[1] = top_tm_blob; - bindings[2] = weight_data_gpu_tm_winograd23; + if (use_cooperative_matrix) + { + std::vector bindings(3); + bindings[0] = bottom_tm_blob; + bindings[1] = top_tm_blob; + bindings[2] = weight_data_gpu_tm_winograd23; - std::vector constants(5); - constants[0].i = bottom_tm_blob.h; - constants[1].i = bottom_tm_blob.cstep; - constants[2].i = top_tm_blob.w; - constants[3].i = top_tm_blob.h; - constants[4].i = top_tm_blob.cstep; + std::vector constants(3); + constants[0].i = bottom_tm_blob.w; + constants[1].i = bottom_tm_blob.cstep; + constants[2].i = top_tm_blob.cstep; - VkMat dispatcher; - dispatcher.w = (top_tm_blob.w + 3) / 4; - dispatcher.h = top_tm_blob.h; - dispatcher.c = 16; + const int blocks_x = (bottom_tm_blob.w + coopmat_M * UNROLL_SG_M * UNROLL_WG_M - 1) / (coopmat_M * UNROLL_SG_M * UNROLL_WG_M); + const int blocks_y = (num_output + coopmat_N * UNROLL_SG_N * UNROLL_WG_N - 1) / (coopmat_N * UNROLL_SG_N * UNROLL_WG_N); - if (use_cooperative_matrix_16_8_8) - { - dispatcher.w = ((top_tm_blob.w + 15) / 16 + 1) / 2 * 32; - dispatcher.h = ((top_tm_blob.h + 1) / 2 + 3) / 4; + const int subgroup_size = vkdev->info.subgroup_size(); + + VkMat dispatcher; + dispatcher.w = (blocks_x * blocks_y) * (subgroup_size * UNROLL_WG_M * UNROLL_WG_N); + dispatcher.h = 1; dispatcher.c = 16; + + cmd.record_pipeline(pipeline_convolution_3x3s1d1_winograd23_gemm, bindings, constants, dispatcher); } - else if (use_cooperative_matrix_16_16_16) + else { - dispatcher.w = ((top_tm_blob.w + 15) / 16 + 1) / 2 * 32; - dispatcher.h = ((top_tm_blob.h + 3) / 4 + 1) / 2; + std::vector bindings(3); + bindings[0] = bottom_tm_blob; + bindings[1] = top_tm_blob; + bindings[2] = weight_data_gpu_tm_winograd23; + + std::vector constants(3); + constants[0].i = bottom_tm_blob.cstep; + constants[1].i = top_tm_blob.w; + constants[2].i = top_tm_blob.cstep; + + VkMat dispatcher; + dispatcher.w = (top_tm_blob.w + 3) / 4; + dispatcher.h = num_output / out_elempack; dispatcher.c = 16; - } - cmd.record_pipeline(pipeline_convolution_3x3s1d1_winograd23_gemm, bindings, constants, dispatcher); + cmd.record_pipeline(pipeline_convolution_3x3s1d1_winograd23_gemm, bindings, constants, dispatcher); + } } // transform output @@ -1788,14 +1911,13 @@ int Convolution_vulkan::forward(const VkMat& bottom_blob, VkMat& top_blob, VkCom bindings[1] = top_blob; bindings[2] = bias_data_gpu; - std::vector constants(7); - constants[0].i = top_tm_blob.h; - constants[1].i = top_tm_blob.cstep; - constants[2].i = block_x; - constants[3].i = block_y; - constants[4].i = top_blob.w; - constants[5].i = top_blob.h; - constants[6].i = top_blob.cstep; + std::vector constants(6); + constants[0].i = top_tm_blob.cstep; + constants[1].i = block_x; + constants[2].i = block_y; + constants[3].i = top_blob.w; + constants[4].i = top_blob.h; + constants[5].i = top_blob.cstep; VkMat dispatcher; dispatcher.w = block_x; diff --git a/src/layer/vulkan/shader/convolution_1x1s1d1_cm.comp b/src/layer/vulkan/shader/convolution_1x1s1d1_cm.comp index 4d46f572f..ac003a7cb 100644 --- a/src/layer/vulkan/shader/convolution_1x1s1d1_cm.comp +++ b/src/layer/vulkan/shader/convolution_1x1s1d1_cm.comp @@ -331,8 +331,6 @@ void main() { if (elempack == 1) { - const uint cstepd4 = psc(cstep) / 4; - const uint Md4_K_USGM_USGK = Md4 * K * UNROLL_SG_M * UNROLL_SG_K; const uint Md4_K_USGM_USGK_d_subgroupsize = (Md4_K_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); [[unroll]] for (uint q = 0; q < Md4_K_USGM_USGK_d_subgroupsize; q++) @@ -347,8 +345,6 @@ void main() } else // if (elempack == 4) { - const uint inchd4 = inch / 4; - const uint Kd4_M_USGM_USGK = Kd4 * M * UNROLL_SG_M * UNROLL_SG_K; const uint Kd4_M_USGM_USGK_d_subgroupsize = (Kd4_M_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); [[unroll]] for (uint q = 0; q < Kd4_M_USGM_USGK_d_subgroupsize; q++) @@ -365,11 +361,7 @@ void main() // load weight { - const uint Nd4_K_USGN = Nd4 * K * UNROLL_SG_N; const uint Nd4_K_USGN_USGK = Nd4 * K * UNROLL_SG_N * UNROLL_SG_K; - - const uint w_offset = ((wgni * kk) * UNROLL_WG_N) * Nd4_K_USGN + ((k / UNROLL_SG_K) * UNROLL_WG_N + sgni) * Nd4_K_USGN_USGK; - const uint Nd4_K_USGN_USGK_d_subgroupsize = (Nd4_K_USGN_USGK + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); [[unroll]] for (uint q = 0; q < Nd4_K_USGN_USGK_d_subgroupsize; q++) { @@ -537,29 +529,29 @@ void main() { if (elempack == 1) { - #if ncnn_VK_KHR_cooperative_matrix +#if ncnn_VK_KHR_cooperative_matrix coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, gl_CooperativeMatrixLayoutColumnMajor); - #elif ncnn_VK_NV_cooperative_matrix +#elif ncnn_VK_NV_cooperative_matrix coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, true); - #endif +#endif } else // if (elempack == 4) { - #if ncnn_VK_KHR_cooperative_matrix +#if ncnn_VK_KHR_cooperative_matrix coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, gl_CooperativeMatrixLayoutRowMajor); - #elif ncnn_VK_NV_cooperative_matrix +#elif ncnn_VK_NV_cooperative_matrix coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, false); - #endif +#endif } } [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) { - #if ncnn_VK_KHR_cooperative_matrix +#if ncnn_VK_KHR_cooperative_matrix coopMatLoad(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, gl_CooperativeMatrixLayoutRowMajor); - #elif ncnn_VK_NV_cooperative_matrix +#elif ncnn_VK_NV_cooperative_matrix coopMatLoadNV(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, false); - #endif +#endif } // sum += k * v @@ -567,11 +559,11 @@ void main() { [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) { - #if ncnn_VK_KHR_cooperative_matrix +#if ncnn_VK_KHR_cooperative_matrix sum[zn][zm] = coopMatMulAdd(A[zm], B[zn], sum[zn][zm]); - #elif ncnn_VK_NV_cooperative_matrix +#elif ncnn_VK_NV_cooperative_matrix sum[zn][zm] = coopMatMulAddNV(A[zm], B[zn], sum[zn][zm]); - #endif +#endif } } } @@ -585,8 +577,6 @@ void main() { if (elempack == 1) { - const uint cstepd4 = psc(cstep) / 4; - const uint Md4_K_USGM_USGK = Md4 * K * UNROLL_SG_M * UNROLL_SG_K; const uint Md4_K_USGM_USGK_d_subgroupsize = (Md4_K_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); [[unroll]] for (uint q = 0; q < Md4_K_USGM_USGK_d_subgroupsize; q++) @@ -601,8 +591,6 @@ void main() } else // if (elempack == 4) { - const uint inchd4 = inch / 4; - const uint Kd4_M_USGM_USGK = Kd4 * M * UNROLL_SG_M * UNROLL_SG_K; const uint Kd4_M_USGM_USGK_d_subgroupsize = (Kd4_M_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); [[unroll]] for (uint q = 0; q < Kd4_M_USGM_USGK_d_subgroupsize; q++) @@ -619,11 +607,7 @@ void main() // load weight { - const uint Nd4_K_USGN = Nd4 * K * UNROLL_SG_N; const uint Nd4_K_USGN_USGK = Nd4 * K * UNROLL_SG_N * UNROLL_SG_K; - - const uint w_offset = ((wgni * kk) * UNROLL_WG_N) * Nd4_K_USGN + ((k / UNROLL_SG_K) * UNROLL_WG_N + sgni) * Nd4_K_USGN_USGK; - const uint Nd4_K_USGN_USGK_d_subgroupsize = (Nd4_K_USGN_USGK + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); [[unroll]] for (uint q = 0; q < Nd4_K_USGN_USGK_d_subgroupsize; q++) { diff --git a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd23_transform_input.comp b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd23_transform_input.comp index f4d984bd9..acbb45da3 100644 --- a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd23_transform_input.comp +++ b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd23_transform_input.comp @@ -3,16 +3,17 @@ #version 450 -#define shape_constant_id_offset 0 +layout (constant_id = 0) const int c = 0; + +#define shape_constant_id_offset 1 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 (constant_id = shape_constant_id_offset + 2) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outcstep = 0; -layout (constant_id = shape_constant_id_offset + 5) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 6) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 4) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 5) const int block_y = 0; layout (binding = 0) readonly buffer bottom_blob { sfp bottom_blob_data[]; }; layout (binding = 1) writeonly buffer bottom_tm_blob { sfp bottom_tm_blob_data[]; }; @@ -21,7 +22,6 @@ layout (push_constant) uniform parameter { int w; int h; - int c; int cstep; int outcstep; @@ -36,7 +36,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 4x4 @@ -115,22 +115,22 @@ void main() v33 = m33 - m31; // store 16 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - buffer_st1(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep), v00); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep), v01); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep), v02); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep), v03); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep), v10); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep), v11); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep), v12); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep), v13); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep), v20); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep), v21); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep), v22); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep), v23); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep), v30); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep), v31); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep), v32); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep), v33); + int v_tm_offset = gz * psc(outcstep) + gy * psc(block_x) + gx; + + buffer_st1(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep) * c, v00); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep) * c, v01); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep) * c, v02); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep) * c, v03); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep) * c, v10); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep) * c, v11); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep) * c, v12); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep) * c, v13); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep) * c, v20); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep) * c, v21); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep) * c, v22); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep) * c, v23); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep) * c, v30); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep) * c, v31); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep) * c, v32); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep) * c, v33); } diff --git a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd23_transform_output.comp b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd23_transform_output.comp index 5daa98842..5e970635d 100644 --- a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd23_transform_output.comp +++ b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd23_transform_output.comp @@ -10,17 +10,17 @@ layout (constant_id = 0) const int bias_term = 0; layout (constant_id = 1) const int activation_type = 0; layout (constant_id = 2) const float activation_param_0 = 0; layout (constant_id = 3) const float activation_param_1 = 0; +layout (constant_id = 4) const int c = 0; -#define shape_constant_id_offset 4 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; +#define shape_constant_id_offset 5 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 2) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 3) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 1) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 2) const int block_y = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 5) const int outh = 0; -layout (constant_id = shape_constant_id_offset + 6) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 4) const int outh = 0; +layout (constant_id = shape_constant_id_offset + 5) const int outcstep = 0; layout (binding = 0) readonly buffer top_tm_blob { sfp top_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_blob { sfp top_blob_data[]; }; @@ -28,7 +28,6 @@ layout (binding = 2) readonly buffer bias_blob { sfp bias_data[]; }; layout (push_constant) uniform parameter { - int c; int cstep; int block_x; @@ -45,28 +44,28 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 16 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - afp v00 = buffer_ld1(top_tm_blob_data, v_tm_offset + 0 * psc(cstep)); - afp v01 = buffer_ld1(top_tm_blob_data, v_tm_offset + 1 * psc(cstep)); - afp v02 = buffer_ld1(top_tm_blob_data, v_tm_offset + 2 * psc(cstep)); - afp v03 = buffer_ld1(top_tm_blob_data, v_tm_offset + 3 * psc(cstep)); - afp v10 = buffer_ld1(top_tm_blob_data, v_tm_offset + 4 * psc(cstep)); - afp v11 = buffer_ld1(top_tm_blob_data, v_tm_offset + 5 * psc(cstep)); - afp v12 = buffer_ld1(top_tm_blob_data, v_tm_offset + 6 * psc(cstep)); - afp v13 = buffer_ld1(top_tm_blob_data, v_tm_offset + 7 * psc(cstep)); - afp v20 = buffer_ld1(top_tm_blob_data, v_tm_offset + 8 * psc(cstep)); - afp v21 = buffer_ld1(top_tm_blob_data, v_tm_offset + 9 * psc(cstep)); - afp v22 = buffer_ld1(top_tm_blob_data, v_tm_offset + 10 * psc(cstep)); - afp v23 = buffer_ld1(top_tm_blob_data, v_tm_offset + 11 * psc(cstep)); - afp v30 = buffer_ld1(top_tm_blob_data, v_tm_offset + 12 * psc(cstep)); - afp v31 = buffer_ld1(top_tm_blob_data, v_tm_offset + 13 * psc(cstep)); - afp v32 = buffer_ld1(top_tm_blob_data, v_tm_offset + 14 * psc(cstep)); - afp v33 = buffer_ld1(top_tm_blob_data, v_tm_offset + 15 * psc(cstep)); + int v_tm_offset = gz * psc(cstep) + gy * psc(block_x) + gx; + + afp v00 = buffer_ld1(top_tm_blob_data, v_tm_offset + 0 * psc(cstep) * c); + afp v01 = buffer_ld1(top_tm_blob_data, v_tm_offset + 1 * psc(cstep) * c); + afp v02 = buffer_ld1(top_tm_blob_data, v_tm_offset + 2 * psc(cstep) * c); + afp v03 = buffer_ld1(top_tm_blob_data, v_tm_offset + 3 * psc(cstep) * c); + afp v10 = buffer_ld1(top_tm_blob_data, v_tm_offset + 4 * psc(cstep) * c); + afp v11 = buffer_ld1(top_tm_blob_data, v_tm_offset + 5 * psc(cstep) * c); + afp v12 = buffer_ld1(top_tm_blob_data, v_tm_offset + 6 * psc(cstep) * c); + afp v13 = buffer_ld1(top_tm_blob_data, v_tm_offset + 7 * psc(cstep) * c); + afp v20 = buffer_ld1(top_tm_blob_data, v_tm_offset + 8 * psc(cstep) * c); + afp v21 = buffer_ld1(top_tm_blob_data, v_tm_offset + 9 * psc(cstep) * c); + afp v22 = buffer_ld1(top_tm_blob_data, v_tm_offset + 10 * psc(cstep) * c); + afp v23 = buffer_ld1(top_tm_blob_data, v_tm_offset + 11 * psc(cstep) * c); + afp v30 = buffer_ld1(top_tm_blob_data, v_tm_offset + 12 * psc(cstep) * c); + afp v31 = buffer_ld1(top_tm_blob_data, v_tm_offset + 13 * psc(cstep) * c); + afp v32 = buffer_ld1(top_tm_blob_data, v_tm_offset + 14 * psc(cstep) * c); + afp v33 = buffer_ld1(top_tm_blob_data, v_tm_offset + 15 * psc(cstep) * c); // const float itm[2][4] = { // {1.0f, 1.0f, 1.0f, 0.0f}, diff --git a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd43_transform_input.comp b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd43_transform_input.comp index 449a9600b..62478297e 100644 --- a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd43_transform_input.comp +++ b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd43_transform_input.comp @@ -3,16 +3,17 @@ #version 450 -#define shape_constant_id_offset 0 +layout (constant_id = 0) const int c = 0; + +#define shape_constant_id_offset 1 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 (constant_id = shape_constant_id_offset + 2) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outcstep = 0; -layout (constant_id = shape_constant_id_offset + 5) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 6) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 4) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 5) const int block_y = 0; layout (binding = 0) readonly buffer bottom_blob { sfp bottom_blob_data[]; }; layout (binding = 1) writeonly buffer bottom_tm_blob { sfp bottom_tm_blob_data[]; }; @@ -21,7 +22,6 @@ layout (push_constant) uniform parameter { int w; int h; - int c; int cstep; int outcstep; @@ -36,7 +36,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 6x6 @@ -187,42 +187,42 @@ void main() v55 = m51 - m53 * afp(2.5) + m55; // store 36 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - buffer_st1(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep), v00); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep), v01); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep), v02); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep), v03); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep), v04); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep), v05); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep), v10); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep), v11); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep), v12); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep), v13); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep), v14); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep), v15); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep), v20); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep), v21); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep), v22); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep), v23); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 16 * psc(outcstep), v24); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 17 * psc(outcstep), v25); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 18 * psc(outcstep), v30); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 19 * psc(outcstep), v31); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 20 * psc(outcstep), v32); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 21 * psc(outcstep), v33); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 22 * psc(outcstep), v34); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 23 * psc(outcstep), v35); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 24 * psc(outcstep), v40); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 25 * psc(outcstep), v41); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 26 * psc(outcstep), v42); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 27 * psc(outcstep), v43); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 28 * psc(outcstep), v44); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 29 * psc(outcstep), v45); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 30 * psc(outcstep), v50); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 31 * psc(outcstep), v51); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 32 * psc(outcstep), v52); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 33 * psc(outcstep), v53); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 34 * psc(outcstep), v54); - buffer_st1(bottom_tm_blob_data, v_tm_offset + 35 * psc(outcstep), v55); + int v_tm_offset = gz * psc(outcstep) + gy * psc(block_x) + gx; + + buffer_st1(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep) * c, v00); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep) * c, v01); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep) * c, v02); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep) * c, v03); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep) * c, v04); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep) * c, v05); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep) * c, v10); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep) * c, v11); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep) * c, v12); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep) * c, v13); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep) * c, v14); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep) * c, v15); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep) * c, v20); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep) * c, v21); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep) * c, v22); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep) * c, v23); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 16 * psc(outcstep) * c, v24); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 17 * psc(outcstep) * c, v25); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 18 * psc(outcstep) * c, v30); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 19 * psc(outcstep) * c, v31); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 20 * psc(outcstep) * c, v32); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 21 * psc(outcstep) * c, v33); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 22 * psc(outcstep) * c, v34); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 23 * psc(outcstep) * c, v35); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 24 * psc(outcstep) * c, v40); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 25 * psc(outcstep) * c, v41); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 26 * psc(outcstep) * c, v42); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 27 * psc(outcstep) * c, v43); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 28 * psc(outcstep) * c, v44); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 29 * psc(outcstep) * c, v45); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 30 * psc(outcstep) * c, v50); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 31 * psc(outcstep) * c, v51); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 32 * psc(outcstep) * c, v52); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 33 * psc(outcstep) * c, v53); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 34 * psc(outcstep) * c, v54); + buffer_st1(bottom_tm_blob_data, v_tm_offset + 35 * psc(outcstep) * c, v55); } diff --git a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd43_transform_output.comp b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd43_transform_output.comp index 2531afd55..914632cc3 100644 --- a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd43_transform_output.comp +++ b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd43_transform_output.comp @@ -10,17 +10,17 @@ layout (constant_id = 0) const int bias_term = 0; layout (constant_id = 1) const int activation_type = 0; layout (constant_id = 2) const float activation_param_0 = 0; layout (constant_id = 3) const float activation_param_1 = 0; +layout (constant_id = 4) const int c = 0; -#define shape_constant_id_offset 4 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; +#define shape_constant_id_offset 5 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 2) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 3) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 1) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 2) const int block_y = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 5) const int outh = 0; -layout (constant_id = shape_constant_id_offset + 6) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 4) const int outh = 0; +layout (constant_id = shape_constant_id_offset + 5) const int outcstep = 0; layout (binding = 0) readonly buffer top_tm_blob { sfp top_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_blob { sfp top_blob_data[]; }; @@ -28,7 +28,6 @@ layout (binding = 2) readonly buffer bias_blob { sfp bias_data[]; }; layout (push_constant) uniform parameter { - int c; int cstep; int block_x; @@ -45,48 +44,48 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 36 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - afp v00 = buffer_ld1(top_tm_blob_data, v_tm_offset + 0 * psc(cstep)); - afp v01 = buffer_ld1(top_tm_blob_data, v_tm_offset + 1 * psc(cstep)); - afp v02 = buffer_ld1(top_tm_blob_data, v_tm_offset + 2 * psc(cstep)); - afp v03 = buffer_ld1(top_tm_blob_data, v_tm_offset + 3 * psc(cstep)); - afp v04 = buffer_ld1(top_tm_blob_data, v_tm_offset + 4 * psc(cstep)); - afp v05 = buffer_ld1(top_tm_blob_data, v_tm_offset + 5 * psc(cstep)); - afp v10 = buffer_ld1(top_tm_blob_data, v_tm_offset + 6 * psc(cstep)); - afp v11 = buffer_ld1(top_tm_blob_data, v_tm_offset + 7 * psc(cstep)); - afp v12 = buffer_ld1(top_tm_blob_data, v_tm_offset + 8 * psc(cstep)); - afp v13 = buffer_ld1(top_tm_blob_data, v_tm_offset + 9 * psc(cstep)); - afp v14 = buffer_ld1(top_tm_blob_data, v_tm_offset + 10 * psc(cstep)); - afp v15 = buffer_ld1(top_tm_blob_data, v_tm_offset + 11 * psc(cstep)); - afp v20 = buffer_ld1(top_tm_blob_data, v_tm_offset + 12 * psc(cstep)); - afp v21 = buffer_ld1(top_tm_blob_data, v_tm_offset + 13 * psc(cstep)); - afp v22 = buffer_ld1(top_tm_blob_data, v_tm_offset + 14 * psc(cstep)); - afp v23 = buffer_ld1(top_tm_blob_data, v_tm_offset + 15 * psc(cstep)); - afp v24 = buffer_ld1(top_tm_blob_data, v_tm_offset + 16 * psc(cstep)); - afp v25 = buffer_ld1(top_tm_blob_data, v_tm_offset + 17 * psc(cstep)); - afp v30 = buffer_ld1(top_tm_blob_data, v_tm_offset + 18 * psc(cstep)); - afp v31 = buffer_ld1(top_tm_blob_data, v_tm_offset + 19 * psc(cstep)); - afp v32 = buffer_ld1(top_tm_blob_data, v_tm_offset + 20 * psc(cstep)); - afp v33 = buffer_ld1(top_tm_blob_data, v_tm_offset + 21 * psc(cstep)); - afp v34 = buffer_ld1(top_tm_blob_data, v_tm_offset + 22 * psc(cstep)); - afp v35 = buffer_ld1(top_tm_blob_data, v_tm_offset + 23 * psc(cstep)); - afp v40 = buffer_ld1(top_tm_blob_data, v_tm_offset + 24 * psc(cstep)); - afp v41 = buffer_ld1(top_tm_blob_data, v_tm_offset + 25 * psc(cstep)); - afp v42 = buffer_ld1(top_tm_blob_data, v_tm_offset + 26 * psc(cstep)); - afp v43 = buffer_ld1(top_tm_blob_data, v_tm_offset + 27 * psc(cstep)); - afp v44 = buffer_ld1(top_tm_blob_data, v_tm_offset + 28 * psc(cstep)); - afp v45 = buffer_ld1(top_tm_blob_data, v_tm_offset + 29 * psc(cstep)); - afp v50 = buffer_ld1(top_tm_blob_data, v_tm_offset + 30 * psc(cstep)); - afp v51 = buffer_ld1(top_tm_blob_data, v_tm_offset + 31 * psc(cstep)); - afp v52 = buffer_ld1(top_tm_blob_data, v_tm_offset + 32 * psc(cstep)); - afp v53 = buffer_ld1(top_tm_blob_data, v_tm_offset + 33 * psc(cstep)); - afp v54 = buffer_ld1(top_tm_blob_data, v_tm_offset + 34 * psc(cstep)); - afp v55 = buffer_ld1(top_tm_blob_data, v_tm_offset + 35 * psc(cstep)); + int v_tm_offset = gz * psc(cstep) + gy * psc(block_x) + gx; + + afp v00 = buffer_ld1(top_tm_blob_data, v_tm_offset + 0 * psc(cstep) * c); + afp v01 = buffer_ld1(top_tm_blob_data, v_tm_offset + 1 * psc(cstep) * c); + afp v02 = buffer_ld1(top_tm_blob_data, v_tm_offset + 2 * psc(cstep) * c); + afp v03 = buffer_ld1(top_tm_blob_data, v_tm_offset + 3 * psc(cstep) * c); + afp v04 = buffer_ld1(top_tm_blob_data, v_tm_offset + 4 * psc(cstep) * c); + afp v05 = buffer_ld1(top_tm_blob_data, v_tm_offset + 5 * psc(cstep) * c); + afp v10 = buffer_ld1(top_tm_blob_data, v_tm_offset + 6 * psc(cstep) * c); + afp v11 = buffer_ld1(top_tm_blob_data, v_tm_offset + 7 * psc(cstep) * c); + afp v12 = buffer_ld1(top_tm_blob_data, v_tm_offset + 8 * psc(cstep) * c); + afp v13 = buffer_ld1(top_tm_blob_data, v_tm_offset + 9 * psc(cstep) * c); + afp v14 = buffer_ld1(top_tm_blob_data, v_tm_offset + 10 * psc(cstep) * c); + afp v15 = buffer_ld1(top_tm_blob_data, v_tm_offset + 11 * psc(cstep) * c); + afp v20 = buffer_ld1(top_tm_blob_data, v_tm_offset + 12 * psc(cstep) * c); + afp v21 = buffer_ld1(top_tm_blob_data, v_tm_offset + 13 * psc(cstep) * c); + afp v22 = buffer_ld1(top_tm_blob_data, v_tm_offset + 14 * psc(cstep) * c); + afp v23 = buffer_ld1(top_tm_blob_data, v_tm_offset + 15 * psc(cstep) * c); + afp v24 = buffer_ld1(top_tm_blob_data, v_tm_offset + 16 * psc(cstep) * c); + afp v25 = buffer_ld1(top_tm_blob_data, v_tm_offset + 17 * psc(cstep) * c); + afp v30 = buffer_ld1(top_tm_blob_data, v_tm_offset + 18 * psc(cstep) * c); + afp v31 = buffer_ld1(top_tm_blob_data, v_tm_offset + 19 * psc(cstep) * c); + afp v32 = buffer_ld1(top_tm_blob_data, v_tm_offset + 20 * psc(cstep) * c); + afp v33 = buffer_ld1(top_tm_blob_data, v_tm_offset + 21 * psc(cstep) * c); + afp v34 = buffer_ld1(top_tm_blob_data, v_tm_offset + 22 * psc(cstep) * c); + afp v35 = buffer_ld1(top_tm_blob_data, v_tm_offset + 23 * psc(cstep) * c); + afp v40 = buffer_ld1(top_tm_blob_data, v_tm_offset + 24 * psc(cstep) * c); + afp v41 = buffer_ld1(top_tm_blob_data, v_tm_offset + 25 * psc(cstep) * c); + afp v42 = buffer_ld1(top_tm_blob_data, v_tm_offset + 26 * psc(cstep) * c); + afp v43 = buffer_ld1(top_tm_blob_data, v_tm_offset + 27 * psc(cstep) * c); + afp v44 = buffer_ld1(top_tm_blob_data, v_tm_offset + 28 * psc(cstep) * c); + afp v45 = buffer_ld1(top_tm_blob_data, v_tm_offset + 29 * psc(cstep) * c); + afp v50 = buffer_ld1(top_tm_blob_data, v_tm_offset + 30 * psc(cstep) * c); + afp v51 = buffer_ld1(top_tm_blob_data, v_tm_offset + 31 * psc(cstep) * c); + afp v52 = buffer_ld1(top_tm_blob_data, v_tm_offset + 32 * psc(cstep) * c); + afp v53 = buffer_ld1(top_tm_blob_data, v_tm_offset + 33 * psc(cstep) * c); + afp v54 = buffer_ld1(top_tm_blob_data, v_tm_offset + 34 * psc(cstep) * c); + afp v55 = buffer_ld1(top_tm_blob_data, v_tm_offset + 35 * psc(cstep) * c); #define sq2 1.41421356237 #define sq2_m2 1.41421356237*2 diff --git a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd_gemm.comp b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd_gemm.comp index 5a06f65a6..b1f869086 100644 --- a/src/layer/vulkan/shader/convolution_3x3s1d1_winograd_gemm.comp +++ b/src/layer/vulkan/shader/convolution_3x3s1d1_winograd_gemm.comp @@ -6,14 +6,13 @@ #define LOCAL_MEMORY_UNROLL_INCH 8 layout (constant_id = 0) const int batch = 1; +layout (constant_id = 1) const int c = 0; +layout (constant_id = 2) const int outc = 0; -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +#define shape_constant_id_offset 3 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; +layout (constant_id = shape_constant_id_offset + 1) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 2) const int outcstep = 0; layout (binding = 0) readonly buffer bottom_tm_blob { sfp bottom_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_tm_blob { sfp top_tm_blob_data[]; }; @@ -21,11 +20,8 @@ layout (binding = 2) readonly buffer weight_tm_blob { sfp weight_tm_data[]; }; layout (push_constant) uniform parameter { - int c; int cstep; - int outw; - int outc; int outcstep; } p; @@ -41,7 +37,7 @@ void main() int gz = int(gl_GlobalInvocationID.z); #if !NCNN_shader_local_memory - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; #endif @@ -50,21 +46,21 @@ void main() afp sum2 = afp(0.f); afp sum3 = afp(0.f); - int v_offset = gz * psc(cstep) + gx; - int w_offset = gz * psc(c) * psc(outc) + gy * psc(c); + int v_offset = gz * c * psc(cstep) + gx; + int w_offset = gz * c * outc + gy * c; #if NCNN_shader_local_memory const int lx = int(gl_LocalInvocationID.x); const int ly = int(gl_LocalInvocationID.y); int z = 0; - for (; z + (LOCAL_MEMORY_UNROLL_INCH - 1) < psc(c); z += LOCAL_MEMORY_UNROLL_INCH) + for (; z + (LOCAL_MEMORY_UNROLL_INCH - 1) < c; z += LOCAL_MEMORY_UNROLL_INCH) { if (ly < 4) { for (int z4 = 0; z4 < LOCAL_MEMORY_UNROLL_INCH; z4++) { - tmp_v[lx][z4][ly] = sfp2lfp(buffer_ld1(bottom_tm_blob_data, v_offset + z4 * psc(outw) + ly)); + tmp_v[lx][z4][ly] = sfp2lfp(buffer_ld1(bottom_tm_blob_data, v_offset + z4 * psc(cstep) + ly)); } } @@ -93,21 +89,21 @@ void main() sum3 += v3 * k; } - v_offset += LOCAL_MEMORY_UNROLL_INCH * psc(outw); + v_offset += LOCAL_MEMORY_UNROLL_INCH * psc(cstep); w_offset += LOCAL_MEMORY_UNROLL_INCH; barrier(); } - if (z < psc(c)) + if (z < c) { - const int remain = psc(c) - z; + const int remain = c - z; if (ly < 4) { for (int z4 = 0; z4 < remain; z4++) { - tmp_v[lx][z4][ly] = sfp2lfp(buffer_ld1(bottom_tm_blob_data, v_offset + z4 * psc(outw) + ly)); + tmp_v[lx][z4][ly] = sfp2lfp(buffer_ld1(bottom_tm_blob_data, v_offset + z4 * psc(cstep) + ly)); } } @@ -137,7 +133,7 @@ void main() } } #else - for (int z = 0; z < psc(c); z++) + for (int z = 0; z < c; z++) { afp v0 = buffer_ld1(bottom_tm_blob_data, v_offset + 0); afp v1 = buffer_ld1(bottom_tm_blob_data, v_offset + 1); @@ -151,17 +147,17 @@ void main() sum2 += v2 * k; sum3 += v3 * k; - v_offset += psc(outw); + v_offset += psc(cstep); w_offset += 1; } #endif #if NCNN_shader_local_memory - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; #endif - int gi = gz * psc(outcstep) + gy * psc(outw) + gx; + int gi = gz * outc * psc(outcstep) + gy * psc(outcstep) + gx; buffer_st1(top_tm_blob_data, gi + 0, sum0); if (gx + 1 < psc(outw)) buffer_st1(top_tm_blob_data, gi + 1, sum1); diff --git a/src/layer/vulkan/shader/convolution_gemm_cm.comp b/src/layer/vulkan/shader/convolution_gemm_cm.comp index b2613639b..322826259 100644 --- a/src/layer/vulkan/shader/convolution_gemm_cm.comp +++ b/src/layer/vulkan/shader/convolution_gemm_cm.comp @@ -385,8 +385,6 @@ void main() { if (elempack == 1) { - const uint cstepd4 = psc(cstep) / 4; - const uint Md4_K_USGM_USGK = Md4 * K * UNROLL_SG_M * UNROLL_SG_K; const uint Md4_K_USGM_USGK_d_subgroupsize = (Md4_K_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); [[unroll]] for (uint q = 0; q < Md4_K_USGM_USGK_d_subgroupsize; q++) @@ -401,8 +399,6 @@ void main() } else // if (elempack == 4) { - const uint inchd4 = inch / 4; - const uint Kd4_M_USGM_USGK = Kd4 * M * UNROLL_SG_M * UNROLL_SG_K; const uint Kd4_M_USGM_USGK_d_subgroupsize = (Kd4_M_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); [[unroll]] for (uint q = 0; q < Kd4_M_USGM_USGK_d_subgroupsize; q++) @@ -419,11 +415,7 @@ void main() // load weight { - const uint Nd4_K_USGN = Nd4 * K * UNROLL_SG_N; const uint Nd4_K_USGN_USGK = Nd4 * K * UNROLL_SG_N * UNROLL_SG_K; - - const uint w_offset = ((wgni * kk) * UNROLL_WG_N) * Nd4_K_USGN + ((k / UNROLL_SG_K) * UNROLL_WG_N + sgni) * Nd4_K_USGN_USGK; - const uint Nd4_K_USGN_USGK_d_subgroupsize = (Nd4_K_USGN_USGK + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); [[unroll]] for (uint q = 0; q < Nd4_K_USGN_USGK_d_subgroupsize; q++) { @@ -629,29 +621,29 @@ void main() { if (elempack == 1) { - #if ncnn_VK_KHR_cooperative_matrix +#if ncnn_VK_KHR_cooperative_matrix coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, gl_CooperativeMatrixLayoutColumnMajor); - #elif ncnn_VK_NV_cooperative_matrix +#elif ncnn_VK_NV_cooperative_matrix coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, true); - #endif +#endif } else // if (elempack == 4) { - #if ncnn_VK_KHR_cooperative_matrix +#if ncnn_VK_KHR_cooperative_matrix coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, gl_CooperativeMatrixLayoutRowMajor); - #elif ncnn_VK_NV_cooperative_matrix +#elif ncnn_VK_NV_cooperative_matrix coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, false); - #endif +#endif } } [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) { - #if ncnn_VK_KHR_cooperative_matrix +#if ncnn_VK_KHR_cooperative_matrix coopMatLoad(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, gl_CooperativeMatrixLayoutRowMajor); - #elif ncnn_VK_NV_cooperative_matrix +#elif ncnn_VK_NV_cooperative_matrix coopMatLoadNV(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, false); - #endif +#endif } // sum += k * v @@ -659,11 +651,11 @@ void main() { [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) { - #if ncnn_VK_KHR_cooperative_matrix +#if ncnn_VK_KHR_cooperative_matrix sum[zn][zm] = coopMatMulAdd(A[zm], B[zn], sum[zn][zm]); - #elif ncnn_VK_NV_cooperative_matrix +#elif ncnn_VK_NV_cooperative_matrix sum[zn][zm] = coopMatMulAddNV(A[zm], B[zn], sum[zn][zm]); - #endif +#endif } } } @@ -691,8 +683,6 @@ void main() } else // if (elempack == 4) { - const uint inchd4 = inch / 4; - const uint Kd4_M_USGM_USGK = Kd4 * M * UNROLL_SG_M * UNROLL_SG_K; const uint Kd4_M_USGM_USGK_d_subgroupsize = (Kd4_M_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); [[unroll]] for (uint q = 0; q < Kd4_M_USGM_USGK_d_subgroupsize; q++) @@ -709,11 +699,7 @@ void main() // load weight { - const uint Nd4_K_USGN = Nd4 * K * UNROLL_SG_N; const uint Nd4_K_USGN_USGK = Nd4 * K * UNROLL_SG_N * UNROLL_SG_K; - - const uint w_offset = ((wgni * kk) * UNROLL_WG_N) * Nd4_K_USGN + ((k / UNROLL_SG_K) * UNROLL_WG_N + sgni) * Nd4_K_USGN_USGK; - const uint Nd4_K_USGN_USGK_d_subgroupsize = (Nd4_K_USGN_USGK + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); [[unroll]] for (uint q = 0; q < Nd4_K_USGN_USGK_d_subgroupsize; q++) { diff --git a/src/layer/vulkan/shader/convolution_pack1to4_3x3s1d1_winograd_gemm.comp b/src/layer/vulkan/shader/convolution_pack1to4_3x3s1d1_winograd_gemm.comp index 8da2d9baf..ab2c3706d 100644 --- a/src/layer/vulkan/shader/convolution_pack1to4_3x3s1d1_winograd_gemm.comp +++ b/src/layer/vulkan/shader/convolution_pack1to4_3x3s1d1_winograd_gemm.comp @@ -6,14 +6,13 @@ #define LOCAL_MEMORY_UNROLL_INCH 8 layout (constant_id = 0) const int batch = 1; +layout (constant_id = 1) const int c = 0; +layout (constant_id = 2) const int outc = 0; -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +#define shape_constant_id_offset 3 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; +layout (constant_id = shape_constant_id_offset + 1) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 2) const int outcstep = 0; layout (binding = 0) readonly buffer bottom_tm_blob { sfp bottom_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_tm_blob { sfpvec4 top_tm_blob_data[]; }; @@ -21,11 +20,8 @@ layout (binding = 2) readonly buffer weight_tm_blob { sfpvec4 weight_tm_data[]; layout (push_constant) uniform parameter { - int c; int cstep; - int outw; - int outc; int outcstep; } p; @@ -41,7 +37,7 @@ void main() int gz = int(gl_GlobalInvocationID.z); #if !NCNN_shader_local_memory - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; #endif @@ -50,21 +46,21 @@ void main() afpvec4 sum2 = afpvec4(0.f); afpvec4 sum3 = afpvec4(0.f); - int v_offset = gz * psc(cstep) + gx; - int w_offset = gz * psc(c) * psc(outc) + gy * psc(c); + int v_offset = gz * c * psc(cstep) + gx; + int w_offset = gz * c * outc + gy * c; #if NCNN_shader_local_memory const int lx = int(gl_LocalInvocationID.x); const int ly = int(gl_LocalInvocationID.y); int z = 0; - for (; z + (LOCAL_MEMORY_UNROLL_INCH - 1) < psc(c); z += LOCAL_MEMORY_UNROLL_INCH) + for (; z + (LOCAL_MEMORY_UNROLL_INCH - 1) < c; z += LOCAL_MEMORY_UNROLL_INCH) { if (ly < 4) { for (int z4 = 0; z4 < LOCAL_MEMORY_UNROLL_INCH; z4++) { - tmp_v[lx][z4][ly] = sfp2lfp(buffer_ld1(bottom_tm_blob_data, v_offset + z4 * psc(outw) + ly)); + tmp_v[lx][z4][ly] = sfp2lfp(buffer_ld1(bottom_tm_blob_data, v_offset + z4 * psc(cstep) + ly)); } } @@ -93,21 +89,21 @@ void main() sum3 += v3 * k; } - v_offset += LOCAL_MEMORY_UNROLL_INCH * psc(outw); + v_offset += LOCAL_MEMORY_UNROLL_INCH * psc(cstep); w_offset += LOCAL_MEMORY_UNROLL_INCH; barrier(); } - if (z < psc(c)) + if (z < c) { - const int remain = psc(c) - z; + const int remain = c - z; if (ly < 4) { for (int z4 = 0; z4 < remain; z4++) { - tmp_v[lx][z4][ly] = sfp2lfp(buffer_ld1(bottom_tm_blob_data, v_offset + z4 * psc(outw) + ly)); + tmp_v[lx][z4][ly] = sfp2lfp(buffer_ld1(bottom_tm_blob_data, v_offset + z4 * psc(cstep) + ly)); } } @@ -137,7 +133,7 @@ void main() } } #else - for (int z = 0; z < psc(c); z++) + for (int z = 0; z < c; z++) { afp v0 = buffer_ld1(bottom_tm_blob_data, v_offset + 0); afp v1 = buffer_ld1(bottom_tm_blob_data, v_offset + 1); @@ -151,17 +147,17 @@ void main() sum2 += v2 * k; sum3 += v3 * k; - v_offset += psc(outw); + v_offset += psc(cstep); w_offset += 1; } #endif #if NCNN_shader_local_memory - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; #endif - int gi = gz * psc(outcstep) + gy * psc(outw) + gx; + int gi = gz * outc * psc(outcstep) + gy * psc(outcstep) + gx; buffer_st4(top_tm_blob_data, gi + 0, sum0); if (gx + 1 < psc(outw)) buffer_st4(top_tm_blob_data, gi + 1, sum1); diff --git a/src/layer/vulkan/shader/convolution_pack1to8_3x3s1d1_winograd_gemm.comp b/src/layer/vulkan/shader/convolution_pack1to8_3x3s1d1_winograd_gemm.comp index 2e1e82c63..1bac0df03 100644 --- a/src/layer/vulkan/shader/convolution_pack1to8_3x3s1d1_winograd_gemm.comp +++ b/src/layer/vulkan/shader/convolution_pack1to8_3x3s1d1_winograd_gemm.comp @@ -4,14 +4,13 @@ #version 450 layout (constant_id = 0) const int batch = 1; +layout (constant_id = 1) const int c = 0; +layout (constant_id = 2) const int outc = 0; -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +#define shape_constant_id_offset 3 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; +layout (constant_id = shape_constant_id_offset + 1) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 2) const int outcstep = 0; layout (binding = 0) readonly buffer bottom_tm_blob { sfp bottom_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_tm_blob { sfpvec8 top_tm_blob_data[]; }; @@ -19,11 +18,8 @@ layout (binding = 2) readonly buffer weight_tm_blob { sfpvec8 weight_tm_data[]; layout (push_constant) uniform parameter { - int c; int cstep; - int outw; - int outc; int outcstep; } p; @@ -33,7 +29,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; afpvec8 sum0 = afpvec8(afpvec4(0.f), afpvec4(0.f)); @@ -41,10 +37,10 @@ void main() afpvec8 sum2 = afpvec8(afpvec4(0.f), afpvec4(0.f)); afpvec8 sum3 = afpvec8(afpvec4(0.f), afpvec4(0.f)); - int v_offset = gz * psc(cstep) + gx; - int w_offset = gz * psc(c) * psc(outc) + gy * psc(c); + int v_offset = gz * c * psc(cstep) + gx; + int w_offset = gz * c * outc + gy * c; - for (int z = 0; z < psc(c); z++) + for (int z = 0; z < c; z++) { afp v0 = buffer_ld1(bottom_tm_blob_data, v_offset + 0); afp v1 = buffer_ld1(bottom_tm_blob_data, v_offset + 1); @@ -66,11 +62,11 @@ void main() sum3[0] += v3 * k[0]; sum3[1] += v3 * k[1]; - v_offset += psc(outw); + v_offset += psc(cstep); w_offset += 1; } - int gi = gz * psc(outcstep) + gy * psc(outw) + gx; + int gi = gz * outc * psc(outcstep) + gy * psc(outcstep) + gx; buffer_st8(top_tm_blob_data, gi + 0, sum0); if (gx + 1 < psc(outw)) buffer_st8(top_tm_blob_data, gi + 1, sum1); diff --git a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd23_transform_input.comp b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd23_transform_input.comp index 205c8b561..da2b496a2 100644 --- a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd23_transform_input.comp +++ b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd23_transform_input.comp @@ -3,16 +3,17 @@ #version 450 -#define shape_constant_id_offset 0 +layout (constant_id = 0) const int c = 0; + +#define shape_constant_id_offset 1 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 (constant_id = shape_constant_id_offset + 2) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outcstep = 0; -layout (constant_id = shape_constant_id_offset + 5) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 6) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 4) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 5) const int block_y = 0; layout (binding = 0) readonly buffer bottom_blob { sfpvec4 bottom_blob_data[]; }; layout (binding = 1) writeonly buffer bottom_tm_blob { sfpvec4 bottom_tm_blob_data[]; }; @@ -21,7 +22,6 @@ layout (push_constant) uniform parameter { int w; int h; - int c; int cstep; int outcstep; @@ -36,7 +36,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 4x4 @@ -115,22 +115,22 @@ void main() v33 = m33 - m31; // store 16 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - buffer_st4(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep), v00); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep), v01); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep), v02); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep), v03); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep), v10); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep), v11); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep), v12); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep), v13); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep), v20); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep), v21); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep), v22); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep), v23); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep), v30); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep), v31); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep), v32); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep), v33); + int v_tm_offset = gz * psc(outcstep) + gy * psc(block_x) + gx; + + buffer_st4(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep) * c, v00); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep) * c, v01); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep) * c, v02); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep) * c, v03); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep) * c, v10); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep) * c, v11); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep) * c, v12); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep) * c, v13); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep) * c, v20); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep) * c, v21); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep) * c, v22); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep) * c, v23); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep) * c, v30); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep) * c, v31); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep) * c, v32); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep) * c, v33); } diff --git a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd23_transform_output.comp b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd23_transform_output.comp index 35a98cc32..b4c6a3b74 100644 --- a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd23_transform_output.comp +++ b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd23_transform_output.comp @@ -10,17 +10,17 @@ layout (constant_id = 0) const int bias_term = 0; layout (constant_id = 1) const int activation_type = 0; layout (constant_id = 2) const float activation_param_0 = 0; layout (constant_id = 3) const float activation_param_1 = 0; +layout (constant_id = 4) const int c = 0; -#define shape_constant_id_offset 4 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; +#define shape_constant_id_offset 5 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 2) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 3) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 1) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 2) const int block_y = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 5) const int outh = 0; -layout (constant_id = shape_constant_id_offset + 6) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 4) const int outh = 0; +layout (constant_id = shape_constant_id_offset + 5) const int outcstep = 0; layout (binding = 0) readonly buffer top_tm_blob { sfpvec4 top_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_blob { sfpvec4 top_blob_data[]; }; @@ -28,7 +28,6 @@ layout (binding = 2) readonly buffer bias_blob { sfpvec4 bias_data[]; }; layout (push_constant) uniform parameter { - int c; int cstep; int block_x; @@ -45,28 +44,28 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 16 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - afpvec4 v00 = buffer_ld4(top_tm_blob_data, v_tm_offset + 0 * psc(cstep)); - afpvec4 v01 = buffer_ld4(top_tm_blob_data, v_tm_offset + 1 * psc(cstep)); - afpvec4 v02 = buffer_ld4(top_tm_blob_data, v_tm_offset + 2 * psc(cstep)); - afpvec4 v03 = buffer_ld4(top_tm_blob_data, v_tm_offset + 3 * psc(cstep)); - afpvec4 v10 = buffer_ld4(top_tm_blob_data, v_tm_offset + 4 * psc(cstep)); - afpvec4 v11 = buffer_ld4(top_tm_blob_data, v_tm_offset + 5 * psc(cstep)); - afpvec4 v12 = buffer_ld4(top_tm_blob_data, v_tm_offset + 6 * psc(cstep)); - afpvec4 v13 = buffer_ld4(top_tm_blob_data, v_tm_offset + 7 * psc(cstep)); - afpvec4 v20 = buffer_ld4(top_tm_blob_data, v_tm_offset + 8 * psc(cstep)); - afpvec4 v21 = buffer_ld4(top_tm_blob_data, v_tm_offset + 9 * psc(cstep)); - afpvec4 v22 = buffer_ld4(top_tm_blob_data, v_tm_offset + 10 * psc(cstep)); - afpvec4 v23 = buffer_ld4(top_tm_blob_data, v_tm_offset + 11 * psc(cstep)); - afpvec4 v30 = buffer_ld4(top_tm_blob_data, v_tm_offset + 12 * psc(cstep)); - afpvec4 v31 = buffer_ld4(top_tm_blob_data, v_tm_offset + 13 * psc(cstep)); - afpvec4 v32 = buffer_ld4(top_tm_blob_data, v_tm_offset + 14 * psc(cstep)); - afpvec4 v33 = buffer_ld4(top_tm_blob_data, v_tm_offset + 15 * psc(cstep)); + int v_tm_offset = gz * psc(cstep) + gy * psc(block_x) + gx; + + afpvec4 v00 = buffer_ld4(top_tm_blob_data, v_tm_offset + 0 * psc(cstep) * c); + afpvec4 v01 = buffer_ld4(top_tm_blob_data, v_tm_offset + 1 * psc(cstep) * c); + afpvec4 v02 = buffer_ld4(top_tm_blob_data, v_tm_offset + 2 * psc(cstep) * c); + afpvec4 v03 = buffer_ld4(top_tm_blob_data, v_tm_offset + 3 * psc(cstep) * c); + afpvec4 v10 = buffer_ld4(top_tm_blob_data, v_tm_offset + 4 * psc(cstep) * c); + afpvec4 v11 = buffer_ld4(top_tm_blob_data, v_tm_offset + 5 * psc(cstep) * c); + afpvec4 v12 = buffer_ld4(top_tm_blob_data, v_tm_offset + 6 * psc(cstep) * c); + afpvec4 v13 = buffer_ld4(top_tm_blob_data, v_tm_offset + 7 * psc(cstep) * c); + afpvec4 v20 = buffer_ld4(top_tm_blob_data, v_tm_offset + 8 * psc(cstep) * c); + afpvec4 v21 = buffer_ld4(top_tm_blob_data, v_tm_offset + 9 * psc(cstep) * c); + afpvec4 v22 = buffer_ld4(top_tm_blob_data, v_tm_offset + 10 * psc(cstep) * c); + afpvec4 v23 = buffer_ld4(top_tm_blob_data, v_tm_offset + 11 * psc(cstep) * c); + afpvec4 v30 = buffer_ld4(top_tm_blob_data, v_tm_offset + 12 * psc(cstep) * c); + afpvec4 v31 = buffer_ld4(top_tm_blob_data, v_tm_offset + 13 * psc(cstep) * c); + afpvec4 v32 = buffer_ld4(top_tm_blob_data, v_tm_offset + 14 * psc(cstep) * c); + afpvec4 v33 = buffer_ld4(top_tm_blob_data, v_tm_offset + 15 * psc(cstep) * c); // const float itm[2][4] = { // {1.0f, 1.0f, 1.0f, 0.0f}, diff --git a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd43_transform_input.comp b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd43_transform_input.comp index 82a28bd36..8ad930d84 100644 --- a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd43_transform_input.comp +++ b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd43_transform_input.comp @@ -3,16 +3,17 @@ #version 450 -#define shape_constant_id_offset 0 +layout (constant_id = 0) const int c = 0; + +#define shape_constant_id_offset 1 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 (constant_id = shape_constant_id_offset + 2) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outcstep = 0; -layout (constant_id = shape_constant_id_offset + 5) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 6) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 4) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 5) const int block_y = 0; layout (binding = 0) readonly buffer bottom_blob { sfpvec4 bottom_blob_data[]; }; layout (binding = 1) writeonly buffer bottom_tm_blob { sfpvec4 bottom_tm_blob_data[]; }; @@ -21,7 +22,6 @@ layout (push_constant) uniform parameter { int w; int h; - int c; int cstep; int outcstep; @@ -36,7 +36,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 6x6 @@ -187,42 +187,42 @@ void main() v55 = m51 - m53 * afp(2.5) + m55; // store 36 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - buffer_st4(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep), v00); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep), v01); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep), v02); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep), v03); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep), v04); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep), v05); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep), v10); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep), v11); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep), v12); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep), v13); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep), v14); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep), v15); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep), v20); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep), v21); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep), v22); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep), v23); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 16 * psc(outcstep), v24); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 17 * psc(outcstep), v25); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 18 * psc(outcstep), v30); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 19 * psc(outcstep), v31); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 20 * psc(outcstep), v32); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 21 * psc(outcstep), v33); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 22 * psc(outcstep), v34); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 23 * psc(outcstep), v35); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 24 * psc(outcstep), v40); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 25 * psc(outcstep), v41); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 26 * psc(outcstep), v42); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 27 * psc(outcstep), v43); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 28 * psc(outcstep), v44); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 29 * psc(outcstep), v45); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 30 * psc(outcstep), v50); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 31 * psc(outcstep), v51); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 32 * psc(outcstep), v52); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 33 * psc(outcstep), v53); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 34 * psc(outcstep), v54); - buffer_st4(bottom_tm_blob_data, v_tm_offset + 35 * psc(outcstep), v55); + int v_tm_offset = gz * psc(outcstep) + gy * psc(block_x) + gx; + + buffer_st4(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep) * c, v00); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep) * c, v01); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep) * c, v02); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep) * c, v03); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep) * c, v04); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep) * c, v05); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep) * c, v10); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep) * c, v11); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep) * c, v12); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep) * c, v13); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep) * c, v14); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep) * c, v15); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep) * c, v20); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep) * c, v21); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep) * c, v22); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep) * c, v23); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 16 * psc(outcstep) * c, v24); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 17 * psc(outcstep) * c, v25); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 18 * psc(outcstep) * c, v30); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 19 * psc(outcstep) * c, v31); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 20 * psc(outcstep) * c, v32); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 21 * psc(outcstep) * c, v33); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 22 * psc(outcstep) * c, v34); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 23 * psc(outcstep) * c, v35); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 24 * psc(outcstep) * c, v40); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 25 * psc(outcstep) * c, v41); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 26 * psc(outcstep) * c, v42); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 27 * psc(outcstep) * c, v43); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 28 * psc(outcstep) * c, v44); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 29 * psc(outcstep) * c, v45); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 30 * psc(outcstep) * c, v50); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 31 * psc(outcstep) * c, v51); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 32 * psc(outcstep) * c, v52); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 33 * psc(outcstep) * c, v53); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 34 * psc(outcstep) * c, v54); + buffer_st4(bottom_tm_blob_data, v_tm_offset + 35 * psc(outcstep) * c, v55); } diff --git a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd43_transform_output.comp b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd43_transform_output.comp index 23cdbf853..cae99013d 100644 --- a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd43_transform_output.comp +++ b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd43_transform_output.comp @@ -10,17 +10,17 @@ layout (constant_id = 0) const int bias_term = 0; layout (constant_id = 1) const int activation_type = 0; layout (constant_id = 2) const float activation_param_0 = 0; layout (constant_id = 3) const float activation_param_1 = 0; +layout (constant_id = 4) const int c = 0; -#define shape_constant_id_offset 4 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; +#define shape_constant_id_offset 5 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 2) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 3) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 1) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 2) const int block_y = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 5) const int outh = 0; -layout (constant_id = shape_constant_id_offset + 6) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 4) const int outh = 0; +layout (constant_id = shape_constant_id_offset + 5) const int outcstep = 0; layout (binding = 0) readonly buffer top_tm_blob { sfpvec4 top_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_blob { sfpvec4 top_blob_data[]; }; @@ -28,7 +28,6 @@ layout (binding = 2) readonly buffer bias_blob { sfpvec4 bias_data[]; }; layout (push_constant) uniform parameter { - int c; int cstep; int block_x; @@ -45,48 +44,48 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 36 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - afpvec4 v00 = buffer_ld4(top_tm_blob_data, v_tm_offset + 0 * psc(cstep)); - afpvec4 v01 = buffer_ld4(top_tm_blob_data, v_tm_offset + 1 * psc(cstep)); - afpvec4 v02 = buffer_ld4(top_tm_blob_data, v_tm_offset + 2 * psc(cstep)); - afpvec4 v03 = buffer_ld4(top_tm_blob_data, v_tm_offset + 3 * psc(cstep)); - afpvec4 v04 = buffer_ld4(top_tm_blob_data, v_tm_offset + 4 * psc(cstep)); - afpvec4 v05 = buffer_ld4(top_tm_blob_data, v_tm_offset + 5 * psc(cstep)); - afpvec4 v10 = buffer_ld4(top_tm_blob_data, v_tm_offset + 6 * psc(cstep)); - afpvec4 v11 = buffer_ld4(top_tm_blob_data, v_tm_offset + 7 * psc(cstep)); - afpvec4 v12 = buffer_ld4(top_tm_blob_data, v_tm_offset + 8 * psc(cstep)); - afpvec4 v13 = buffer_ld4(top_tm_blob_data, v_tm_offset + 9 * psc(cstep)); - afpvec4 v14 = buffer_ld4(top_tm_blob_data, v_tm_offset + 10 * psc(cstep)); - afpvec4 v15 = buffer_ld4(top_tm_blob_data, v_tm_offset + 11 * psc(cstep)); - afpvec4 v20 = buffer_ld4(top_tm_blob_data, v_tm_offset + 12 * psc(cstep)); - afpvec4 v21 = buffer_ld4(top_tm_blob_data, v_tm_offset + 13 * psc(cstep)); - afpvec4 v22 = buffer_ld4(top_tm_blob_data, v_tm_offset + 14 * psc(cstep)); - afpvec4 v23 = buffer_ld4(top_tm_blob_data, v_tm_offset + 15 * psc(cstep)); - afpvec4 v24 = buffer_ld4(top_tm_blob_data, v_tm_offset + 16 * psc(cstep)); - afpvec4 v25 = buffer_ld4(top_tm_blob_data, v_tm_offset + 17 * psc(cstep)); - afpvec4 v30 = buffer_ld4(top_tm_blob_data, v_tm_offset + 18 * psc(cstep)); - afpvec4 v31 = buffer_ld4(top_tm_blob_data, v_tm_offset + 19 * psc(cstep)); - afpvec4 v32 = buffer_ld4(top_tm_blob_data, v_tm_offset + 20 * psc(cstep)); - afpvec4 v33 = buffer_ld4(top_tm_blob_data, v_tm_offset + 21 * psc(cstep)); - afpvec4 v34 = buffer_ld4(top_tm_blob_data, v_tm_offset + 22 * psc(cstep)); - afpvec4 v35 = buffer_ld4(top_tm_blob_data, v_tm_offset + 23 * psc(cstep)); - afpvec4 v40 = buffer_ld4(top_tm_blob_data, v_tm_offset + 24 * psc(cstep)); - afpvec4 v41 = buffer_ld4(top_tm_blob_data, v_tm_offset + 25 * psc(cstep)); - afpvec4 v42 = buffer_ld4(top_tm_blob_data, v_tm_offset + 26 * psc(cstep)); - afpvec4 v43 = buffer_ld4(top_tm_blob_data, v_tm_offset + 27 * psc(cstep)); - afpvec4 v44 = buffer_ld4(top_tm_blob_data, v_tm_offset + 28 * psc(cstep)); - afpvec4 v45 = buffer_ld4(top_tm_blob_data, v_tm_offset + 29 * psc(cstep)); - afpvec4 v50 = buffer_ld4(top_tm_blob_data, v_tm_offset + 30 * psc(cstep)); - afpvec4 v51 = buffer_ld4(top_tm_blob_data, v_tm_offset + 31 * psc(cstep)); - afpvec4 v52 = buffer_ld4(top_tm_blob_data, v_tm_offset + 32 * psc(cstep)); - afpvec4 v53 = buffer_ld4(top_tm_blob_data, v_tm_offset + 33 * psc(cstep)); - afpvec4 v54 = buffer_ld4(top_tm_blob_data, v_tm_offset + 34 * psc(cstep)); - afpvec4 v55 = buffer_ld4(top_tm_blob_data, v_tm_offset + 35 * psc(cstep)); + int v_tm_offset = gz * psc(cstep) + gy * psc(block_x) + gx; + + afpvec4 v00 = buffer_ld4(top_tm_blob_data, v_tm_offset + 0 * psc(cstep) * c); + afpvec4 v01 = buffer_ld4(top_tm_blob_data, v_tm_offset + 1 * psc(cstep) * c); + afpvec4 v02 = buffer_ld4(top_tm_blob_data, v_tm_offset + 2 * psc(cstep) * c); + afpvec4 v03 = buffer_ld4(top_tm_blob_data, v_tm_offset + 3 * psc(cstep) * c); + afpvec4 v04 = buffer_ld4(top_tm_blob_data, v_tm_offset + 4 * psc(cstep) * c); + afpvec4 v05 = buffer_ld4(top_tm_blob_data, v_tm_offset + 5 * psc(cstep) * c); + afpvec4 v10 = buffer_ld4(top_tm_blob_data, v_tm_offset + 6 * psc(cstep) * c); + afpvec4 v11 = buffer_ld4(top_tm_blob_data, v_tm_offset + 7 * psc(cstep) * c); + afpvec4 v12 = buffer_ld4(top_tm_blob_data, v_tm_offset + 8 * psc(cstep) * c); + afpvec4 v13 = buffer_ld4(top_tm_blob_data, v_tm_offset + 9 * psc(cstep) * c); + afpvec4 v14 = buffer_ld4(top_tm_blob_data, v_tm_offset + 10 * psc(cstep) * c); + afpvec4 v15 = buffer_ld4(top_tm_blob_data, v_tm_offset + 11 * psc(cstep) * c); + afpvec4 v20 = buffer_ld4(top_tm_blob_data, v_tm_offset + 12 * psc(cstep) * c); + afpvec4 v21 = buffer_ld4(top_tm_blob_data, v_tm_offset + 13 * psc(cstep) * c); + afpvec4 v22 = buffer_ld4(top_tm_blob_data, v_tm_offset + 14 * psc(cstep) * c); + afpvec4 v23 = buffer_ld4(top_tm_blob_data, v_tm_offset + 15 * psc(cstep) * c); + afpvec4 v24 = buffer_ld4(top_tm_blob_data, v_tm_offset + 16 * psc(cstep) * c); + afpvec4 v25 = buffer_ld4(top_tm_blob_data, v_tm_offset + 17 * psc(cstep) * c); + afpvec4 v30 = buffer_ld4(top_tm_blob_data, v_tm_offset + 18 * psc(cstep) * c); + afpvec4 v31 = buffer_ld4(top_tm_blob_data, v_tm_offset + 19 * psc(cstep) * c); + afpvec4 v32 = buffer_ld4(top_tm_blob_data, v_tm_offset + 20 * psc(cstep) * c); + afpvec4 v33 = buffer_ld4(top_tm_blob_data, v_tm_offset + 21 * psc(cstep) * c); + afpvec4 v34 = buffer_ld4(top_tm_blob_data, v_tm_offset + 22 * psc(cstep) * c); + afpvec4 v35 = buffer_ld4(top_tm_blob_data, v_tm_offset + 23 * psc(cstep) * c); + afpvec4 v40 = buffer_ld4(top_tm_blob_data, v_tm_offset + 24 * psc(cstep) * c); + afpvec4 v41 = buffer_ld4(top_tm_blob_data, v_tm_offset + 25 * psc(cstep) * c); + afpvec4 v42 = buffer_ld4(top_tm_blob_data, v_tm_offset + 26 * psc(cstep) * c); + afpvec4 v43 = buffer_ld4(top_tm_blob_data, v_tm_offset + 27 * psc(cstep) * c); + afpvec4 v44 = buffer_ld4(top_tm_blob_data, v_tm_offset + 28 * psc(cstep) * c); + afpvec4 v45 = buffer_ld4(top_tm_blob_data, v_tm_offset + 29 * psc(cstep) * c); + afpvec4 v50 = buffer_ld4(top_tm_blob_data, v_tm_offset + 30 * psc(cstep) * c); + afpvec4 v51 = buffer_ld4(top_tm_blob_data, v_tm_offset + 31 * psc(cstep) * c); + afpvec4 v52 = buffer_ld4(top_tm_blob_data, v_tm_offset + 32 * psc(cstep) * c); + afpvec4 v53 = buffer_ld4(top_tm_blob_data, v_tm_offset + 33 * psc(cstep) * c); + afpvec4 v54 = buffer_ld4(top_tm_blob_data, v_tm_offset + 34 * psc(cstep) * c); + afpvec4 v55 = buffer_ld4(top_tm_blob_data, v_tm_offset + 35 * psc(cstep) * c); #define sq2 1.41421356237 #define sq2_m2 1.41421356237*2 diff --git a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm.comp b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm.comp index 005f04c73..e0f35f1ef 100644 --- a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm.comp +++ b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm.comp @@ -6,14 +6,13 @@ #define LOCAL_MEMORY_UNROLL_INCH 8 layout (constant_id = 0) const int batch = 1; +layout (constant_id = 1) const int c = 0; +layout (constant_id = 2) const int outc = 0; -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +#define shape_constant_id_offset 3 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; +layout (constant_id = shape_constant_id_offset + 1) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 2) const int outcstep = 0; layout (binding = 0) readonly buffer bottom_tm_blob { sfpvec4 bottom_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_tm_blob { sfpvec4 top_tm_blob_data[]; }; @@ -21,11 +20,8 @@ layout (binding = 2) readonly buffer weight_tm_blob { sfpvec4 weight_tm_data[]; layout (push_constant) uniform parameter { - int c; int cstep; - int outw; - int outc; int outcstep; } p; @@ -41,7 +37,7 @@ void main() int gz = int(gl_GlobalInvocationID.z); #if !NCNN_shader_local_memory - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; #endif @@ -50,21 +46,21 @@ void main() afpvec4 sum2 = afpvec4(0.f); afpvec4 sum3 = afpvec4(0.f); - int v_offset = gz * psc(cstep) + gx; - int w_offset = (gz * psc(c) * psc(outc) + gy * psc(c)) * 4; + int v_offset = gz * c * psc(cstep) + gx; + int w_offset = (gz * c * outc + gy * c) * 4; #if NCNN_shader_local_memory const int lx = int(gl_LocalInvocationID.x); const int ly = int(gl_LocalInvocationID.y); int z = 0; - for (; z + (LOCAL_MEMORY_UNROLL_INCH - 1) < psc(c); z += LOCAL_MEMORY_UNROLL_INCH) + for (; z + (LOCAL_MEMORY_UNROLL_INCH - 1) < c; z += LOCAL_MEMORY_UNROLL_INCH) { if (ly < 4) { for (int z4 = 0; z4 < LOCAL_MEMORY_UNROLL_INCH; z4++) { - tmp_v[lx][z4][ly] = sfp2lfpvec4(bottom_tm_blob_data[v_offset + z4 * psc(outw) + ly]); + tmp_v[lx][z4][ly] = sfp2lfpvec4(bottom_tm_blob_data[v_offset + z4 * psc(cstep) + ly]); } } @@ -98,21 +94,21 @@ void main() sum3 += v3 * k; } - v_offset += LOCAL_MEMORY_UNROLL_INCH * psc(outw); + v_offset += LOCAL_MEMORY_UNROLL_INCH * psc(cstep); w_offset += LOCAL_MEMORY_UNROLL_INCH * 4; barrier(); } - if (z < psc(c)) + if (z < c) { - const int remain = psc(c) - z; + const int remain = c - z; if (ly < 4) { for (int z4 = 0; z4 < remain; z4++) { - tmp_v[lx][z4][ly] = sfp2lfpvec4(bottom_tm_blob_data[v_offset + z4 * psc(outw) + ly]); + tmp_v[lx][z4][ly] = sfp2lfpvec4(bottom_tm_blob_data[v_offset + z4 * psc(cstep) + ly]); } } @@ -147,7 +143,7 @@ void main() } } #else - for (int z = 0; z < psc(c); z++) + for (int z = 0; z < c; z++) { afpvec4 v0 = buffer_ld4(bottom_tm_blob_data, v_offset + 0); afpvec4 v1 = buffer_ld4(bottom_tm_blob_data, v_offset + 1); @@ -166,17 +162,17 @@ void main() sum2 += v2 * k; sum3 += v3 * k; - v_offset += psc(outw); + v_offset += psc(cstep); w_offset += 4; } #endif #if NCNN_shader_local_memory - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; #endif - int gi = gz * psc(outcstep) + gy * psc(outw) + gx; + int gi = gz * outc * psc(outcstep) + gy * psc(outcstep) + gx; buffer_st4(top_tm_blob_data, gi + 0, sum0); if (gx + 1 < psc(outw)) buffer_st4(top_tm_blob_data, gi + 1, sum1); diff --git a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm_cm_16_16_16.comp b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm_cm_16_16_16.comp deleted file mode 100644 index 5e1ed3439..000000000 --- a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm_cm_16_16_16.comp +++ /dev/null @@ -1,258 +0,0 @@ -// Copyright 2023 Tencent -// SPDX-License-Identifier: BSD-3-Clause - -#version 450 - -#extension GL_KHR_memory_scope_semantics: require -#extension GL_EXT_shader_explicit_arithmetic_types: require -#extension GL_EXT_shader_explicit_arithmetic_types_float16: require -#if ncnn_VK_KHR_cooperative_matrix -#extension GL_KHR_cooperative_matrix: require -#elif ncnn_VK_NV_cooperative_matrix -#extension GL_NV_cooperative_matrix: require -#endif - -layout (constant_id = 0) const int batch = 1; - -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; - -layout (binding = 0) readonly buffer bottom_tm_blob { uvec2 bottom_tm_blob_data[]; }; -layout (binding = 1) writeonly buffer top_tm_blob { uvec2 top_tm_blob_data[]; }; -layout (binding = 2) readonly buffer weight_tm_blob { uvec2 weight_tm_data[]; }; - -layout (push_constant) uniform parameter -{ - int c; - int cstep; - - int outw; - int outc; - int outcstep; -} p; - -#define UNROLL_INCH 2 - -shared uvec2 tmp_v0[UNROLL_INCH * 16*4]; -shared uvec2 tmp_v1[UNROLL_INCH * 16*4]; -shared uvec2 tmp_k0[UNROLL_INCH * 16*4]; -shared uvec2 tmp_k1[UNROLL_INCH * 16*4]; - -void main() -{ - int gx = int(gl_GlobalInvocationID.x) / 32 * 2 * 16; - int gy = int(gl_GlobalInvocationID.y) * 2 * 4; - int gz = int(gl_GlobalInvocationID.z); - - const int lx = int(gl_LocalInvocationID.x); - - const int lxd16 = lx / 16; // 0 1 - const int lxm16 = lx % 16; // 0 1 2 3 .... 15 - -#if ncnn_VK_KHR_cooperative_matrix - coopmat sum0 = coopmat(0.f); - coopmat sum1 = coopmat(0.f); - coopmat sum2 = coopmat(0.f); - coopmat sum3 = coopmat(0.f); -#elif ncnn_VK_NV_cooperative_matrix -#if NCNN_fp16_arithmetic - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> sum0 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 16>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> sum1 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 16>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> sum2 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 16>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> sum3 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 16>(0.f); -#else - fcoopmatNV<32, gl_ScopeSubgroup, 16, 16> sum0 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 16>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 16> sum1 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 16>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 16> sum2 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 16>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 16> sum3 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 16>(0.f); -#endif -#endif - - const int N = psc(c) / 4; - - int z = 0; - for (; z + (UNROLL_INCH - 1) < N; z += UNROLL_INCH) - { - { - for (int j = 0; j < 4; j++) - { - const int tmp_i = lxd16*16*4 + lxm16 * 4 + j; - - const int v_offset = gz * psc(cstep) + ((z + lxd16) * 4 + j) * psc(outw) + (gx + lxm16); - - tmp_v0[tmp_i] = (gx + lxm16) < psc(outw) ? bottom_tm_blob_data[v_offset] : uvec2(0); - tmp_v1[tmp_i] = (gx + lxm16 + 16) < psc(outw) ? bottom_tm_blob_data[v_offset + 16] : uvec2(0); - - const int w_offset = gz * psc(outc) * psc(c) * 4 + gy * psc(c) * 4 + (z + lxd16) * 4 * 16 + (lxm16 * 4 + j); - - tmp_k0[tmp_i] = weight_tm_data[w_offset]; - tmp_k1[tmp_i] = weight_tm_data[w_offset + psc(c) * 16]; - } - } - - barrier(); - - for (int z4 = 0; z4 < UNROLL_INCH; z4++) - { -#if ncnn_VK_KHR_cooperative_matrix - coopmat A0; - coopmat A1; - coopMatLoad(A0, tmp_v0, z4*16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(A1, tmp_v1, z4*16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - - coopmat B0; - coopmat B1; - coopMatLoad(B0, tmp_k0, z4*16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(B1, tmp_k1, z4*16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - - // sum += v * k - sum0 = coopMatMulAdd(A0, B0, sum0); - sum1 = coopMatMulAdd(A1, B0, sum1); - sum2 = coopMatMulAdd(A0, B1, sum2); - sum3 = coopMatMulAdd(A1, B1, sum3); -#elif ncnn_VK_NV_cooperative_matrix - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> A0; - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> A1; - coopMatLoadNV(A0, tmp_v0, z4*16*4, 4, false); - coopMatLoadNV(A1, tmp_v1, z4*16*4, 4, false); - - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> B0; - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> B1; - coopMatLoadNV(B0, tmp_k0, z4*16*4, 4, false); - coopMatLoadNV(B1, tmp_k1, z4*16*4, 4, false); - - // sum += v * k - sum0 = coopMatMulAddNV(A0, B0, sum0); - sum1 = coopMatMulAddNV(A1, B0, sum1); - sum2 = coopMatMulAddNV(A0, B1, sum2); - sum3 = coopMatMulAddNV(A1, B1, sum3); -#endif - } - - barrier(); - } - - if (z < N) - { - const int remain = N - z; - - if (lxd16 == 0) - { - for (int j = 0; j < 4; j++) - { - const int tmp_i = lxd16*16*4 + lxm16 * 4 + j; - - const int v_offset = gz * psc(cstep) + ((z + lxd16) * 4 + j) * psc(outw) + (gx + lxm16); - - tmp_v0[tmp_i] = (gx + lxm16) < psc(outw) ? bottom_tm_blob_data[v_offset] : uvec2(0); - tmp_v1[tmp_i] = (gx + lxm16 + 16) < psc(outw) ? bottom_tm_blob_data[v_offset + 16] : uvec2(0); - - const int w_offset = gz * psc(outc) * psc(c) * 4 + gy * psc(c) * 4 + (z + lxd16) * 4 * 16 + (lxm16 * 4 + j); - - tmp_k0[tmp_i] = weight_tm_data[w_offset]; - tmp_k1[tmp_i] = weight_tm_data[w_offset + psc(c) * 16]; - } - } - - barrier(); - - for (int z4 = 0; z4 < remain; z4++) - { -#if ncnn_VK_KHR_cooperative_matrix - coopmat A0; - coopmat A1; - coopMatLoad(A0, tmp_v0, z4*16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(A1, tmp_v1, z4*16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - - coopmat B0; - coopmat B1; - coopMatLoad(B0, tmp_k0, z4*16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(B1, tmp_k1, z4*16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - - // sum += v * k - sum0 = coopMatMulAdd(A0, B0, sum0); - sum1 = coopMatMulAdd(A1, B0, sum1); - sum2 = coopMatMulAdd(A0, B1, sum2); - sum3 = coopMatMulAdd(A1, B1, sum3); -#elif ncnn_VK_NV_cooperative_matrix - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> A0; - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> A1; - coopMatLoadNV(A0, tmp_v0, z4*16*4, 4, false); - coopMatLoadNV(A1, tmp_v1, z4*16*4, 4, false); - - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> B0; - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> B1; - coopMatLoadNV(B0, tmp_k0, z4*16*4, 4, false); - coopMatLoadNV(B1, tmp_k1, z4*16*4, 4, false); - - // sum += v * k - sum0 = coopMatMulAddNV(A0, B0, sum0); - sum1 = coopMatMulAddNV(A1, B0, sum1); - sum2 = coopMatMulAddNV(A0, B1, sum2); - sum3 = coopMatMulAddNV(A1, B1, sum3); -#endif - } - - barrier(); - } - - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) - return; - -#if ncnn_VK_KHR_cooperative_matrix -#if NCNN_fp16_arithmetic - coopMatStore(sum0, tmp_v0, 0, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum1, tmp_v1, 0, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum2, tmp_v0, 16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum3, tmp_v1, 16*4, 4, gl_CooperativeMatrixLayoutRowMajor); -#else - coopmat sum0_fp16 = coopmat(sum0); - coopmat sum1_fp16 = coopmat(sum1); - coopmat sum2_fp16 = coopmat(sum2); - coopmat sum3_fp16 = coopmat(sum3); - - coopMatStore(sum0_fp16, tmp_v0, 0, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum1_fp16, tmp_v1, 0, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum2_fp16, tmp_v0, 16*4, 4, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum3_fp16, tmp_v1, 16*4, 4, gl_CooperativeMatrixLayoutRowMajor); -#endif -#elif ncnn_VK_NV_cooperative_matrix -#if NCNN_fp16_arithmetic - coopMatStoreNV(sum0, tmp_v0, 0, 4, false); - coopMatStoreNV(sum1, tmp_v1, 0, 4, false); - coopMatStoreNV(sum2, tmp_v0, 16*4, 4, false); - coopMatStoreNV(sum3, tmp_v1, 16*4, 4, false); -#else - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> sum0_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 16>(sum0); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> sum1_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 16>(sum1); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> sum2_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 16>(sum2); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 16> sum3_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 16>(sum3); - - coopMatStoreNV(sum0_fp16, tmp_v0, 0, 4, false); - coopMatStoreNV(sum1_fp16, tmp_v1, 0, 4, false); - coopMatStoreNV(sum2_fp16, tmp_v0, 16*4, 4, false); - coopMatStoreNV(sum3_fp16, tmp_v1, 16*4, 4, false); -#endif -#endif - - barrier(); - - { - for (int j = 0; j < 4; j++) - { - const int tmp_vi = lxm16 * 4 + j + lxd16*16*4; - const int gi = gz * psc(outcstep) + (gy + lxd16 * 4 + j) * psc(outw) + (gx + lxm16); - - if (gy + lxd16 * 4 + j < psc(outc)) - { - if (gx + lxm16 < psc(outw)) top_tm_blob_data[gi] = tmp_v0[tmp_vi]; - if (gx + lxm16 + 16 < psc(outw)) top_tm_blob_data[gi + 16] = tmp_v1[tmp_vi]; - } - } - } -} diff --git a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm_cm_16_8_8.comp b/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm_cm_16_8_8.comp deleted file mode 100644 index 69a0ceb8e..000000000 --- a/src/layer/vulkan/shader/convolution_pack4_3x3s1d1_winograd_gemm_cm_16_8_8.comp +++ /dev/null @@ -1,343 +0,0 @@ -// Copyright 2023 Tencent -// SPDX-License-Identifier: BSD-3-Clause - -#version 450 - -#extension GL_KHR_memory_scope_semantics: require -#extension GL_EXT_shader_explicit_arithmetic_types: require -#extension GL_EXT_shader_explicit_arithmetic_types_float16: require -#if ncnn_VK_KHR_cooperative_matrix -#extension GL_KHR_cooperative_matrix: require -#elif ncnn_VK_NV_cooperative_matrix -#extension GL_NV_cooperative_matrix: require -#endif - -layout (constant_id = 0) const int batch = 1; - -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; - -layout (binding = 0) readonly buffer bottom_tm_blob { uvec2 bottom_tm_blob_data[]; }; -layout (binding = 1) writeonly buffer top_tm_blob { uvec2 top_tm_blob_data[]; }; -layout (binding = 2) readonly buffer weight_tm_blob { uvec2 weight_tm_data[]; }; - -layout (push_constant) uniform parameter -{ - int c; - int cstep; - - int outw; - int outc; - int outcstep; -} p; - -#define UNROLL_INCH 4 - -shared uvec2 tmp_v0[UNROLL_INCH * 16*2]; -shared uvec2 tmp_v1[UNROLL_INCH * 16*2]; -shared uvec2 tmp_k0[UNROLL_INCH * 8*2]; -shared uvec2 tmp_k1[UNROLL_INCH * 8*2]; -shared uvec2 tmp_k2[UNROLL_INCH * 8*2]; -shared uvec2 tmp_k3[UNROLL_INCH * 8*2]; - -void main() -{ - int gx = int(gl_GlobalInvocationID.x) / 32 * 2 * 16; - int gy = int(gl_GlobalInvocationID.y) * 2 * 4; - int gz = int(gl_GlobalInvocationID.z); - - const int lx = int(gl_LocalInvocationID.x); - - const int lxd8 = lx / 8; // 0 1 2 3 - const int lxm8 = lx % 8; // 0 1 2 3 .... 7 - -#if ncnn_VK_KHR_cooperative_matrix - coopmat sum0 = coopmat(0.f); - coopmat sum1 = coopmat(0.f); - coopmat sum2 = coopmat(0.f); - coopmat sum3 = coopmat(0.f); - coopmat sum4 = coopmat(0.f); - coopmat sum5 = coopmat(0.f); - coopmat sum6 = coopmat(0.f); - coopmat sum7 = coopmat(0.f); -#elif ncnn_VK_NV_cooperative_matrix -#if NCNN_fp16_arithmetic - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum0 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum1 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum2 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum3 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum4 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum5 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum6 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum7 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(0.f); -#else - fcoopmatNV<32, gl_ScopeSubgroup, 16, 8> sum0 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 8> sum1 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 8> sum2 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 8> sum3 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 8> sum4 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 8> sum5 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 8> sum6 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 8>(0.f); - fcoopmatNV<32, gl_ScopeSubgroup, 16, 8> sum7 = fcoopmatNV<32, gl_ScopeSubgroup, 16, 8>(0.f); -#endif -#endif - - const int N = psc(c) / 2; - - int z = 0; - for (; z + (UNROLL_INCH - 1) < N; z += UNROLL_INCH) - { - { - for (int j = 0; j < 2; j++) - { - const int tmp_vi = lxd8*16*2 + lxm8 * 2 + j; - - int v_offset = gz * psc(cstep) + ((z + lxd8) * 2 + j) * psc(outw) + (gx + lxm8); - - tmp_v0[tmp_vi] = (gx + lxm8) < psc(outw) ? bottom_tm_blob_data[v_offset] : uvec2(0); - tmp_v0[tmp_vi + 16] = (gx + lxm8 + 8) < psc(outw) ? bottom_tm_blob_data[v_offset + 8] : uvec2(0); - tmp_v1[tmp_vi] = (gx + lxm8 + 16) < psc(outw) ? bottom_tm_blob_data[v_offset + 16] : uvec2(0); - tmp_v1[tmp_vi + 16] = (gx + lxm8 + 24) < psc(outw) ? bottom_tm_blob_data[v_offset + 24] : uvec2(0); - - const int tmp_ki = lxd8*8*2 + lxm8 * 2 + j; - - int w_offset = gz * psc(outc) * psc(c) * 4 + gy * psc(c) * 4 + (z + lxd8) * 2 * 8 + (lxm8 * 2 + j); - - tmp_k0[tmp_ki] = weight_tm_data[w_offset]; - tmp_k1[tmp_ki] = weight_tm_data[w_offset + psc(c) * 8]; - tmp_k2[tmp_ki] = weight_tm_data[w_offset + psc(c) * 16]; - tmp_k3[tmp_ki] = weight_tm_data[w_offset + psc(c) * 24]; - } - } - - barrier(); - - for (int z4 = 0; z4 < UNROLL_INCH; z4++) - { -#if ncnn_VK_KHR_cooperative_matrix - coopmat A0; - coopmat A1; - coopMatLoad(A0, tmp_v0, z4*16*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(A1, tmp_v1, z4*16*2, 2, gl_CooperativeMatrixLayoutRowMajor); - - coopmat B0; - coopmat B1; - coopmat B2; - coopmat B3; - coopMatLoad(B0, tmp_k0, z4*8*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(B1, tmp_k1, z4*8*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(B2, tmp_k2, z4*8*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(B3, tmp_k3, z4*8*2, 2, gl_CooperativeMatrixLayoutRowMajor); - - // sum += v * k - sum0 = coopMatMulAdd(A0, B0, sum0); - sum1 = coopMatMulAdd(A1, B0, sum1); - sum2 = coopMatMulAdd(A0, B1, sum2); - sum3 = coopMatMulAdd(A1, B1, sum3); - sum4 = coopMatMulAdd(A0, B2, sum4); - sum5 = coopMatMulAdd(A1, B2, sum5); - sum6 = coopMatMulAdd(A0, B3, sum6); - sum7 = coopMatMulAdd(A1, B3, sum7); -#elif ncnn_VK_NV_cooperative_matrix - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> A0; - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> A1; - coopMatLoadNV(A0, tmp_v0, z4*16*2, 2, false); - coopMatLoadNV(A1, tmp_v1, z4*16*2, 2, false); - - fcoopmatNV<16, gl_ScopeSubgroup, 8, 8> B0; - fcoopmatNV<16, gl_ScopeSubgroup, 8, 8> B1; - fcoopmatNV<16, gl_ScopeSubgroup, 8, 8> B2; - fcoopmatNV<16, gl_ScopeSubgroup, 8, 8> B3; - coopMatLoadNV(B0, tmp_k0, z4*8*2, 2, false); - coopMatLoadNV(B1, tmp_k1, z4*8*2, 2, false); - coopMatLoadNV(B2, tmp_k2, z4*8*2, 2, false); - coopMatLoadNV(B3, tmp_k3, z4*8*2, 2, false); - - // sum += v * k - sum0 = coopMatMulAddNV(A0, B0, sum0); - sum1 = coopMatMulAddNV(A1, B0, sum1); - sum2 = coopMatMulAddNV(A0, B1, sum2); - sum3 = coopMatMulAddNV(A1, B1, sum3); - sum4 = coopMatMulAddNV(A0, B2, sum4); - sum5 = coopMatMulAddNV(A1, B2, sum5); - sum6 = coopMatMulAddNV(A0, B3, sum6); - sum7 = coopMatMulAddNV(A1, B3, sum7); -#endif - } - - barrier(); - } - - if (z < N) - { - const int remain = N - z; - - if (lxd8 < remain) - { - for (int j = 0; j < 2; j++) - { - const int tmp_vi = lxd8*16*2 + lxm8 * 2 + j; - - int v_offset = gz * psc(cstep) + ((z + lxd8) * 2 + j) * psc(outw) + (gx + lxm8); - - tmp_v0[tmp_vi] = (gx + lxm8) < psc(outw) ? bottom_tm_blob_data[v_offset] : uvec2(0); - tmp_v0[tmp_vi + 16] = (gx + lxm8 + 8) < psc(outw) ? bottom_tm_blob_data[v_offset + 8] : uvec2(0); - tmp_v1[tmp_vi] = (gx + lxm8 + 16) < psc(outw) ? bottom_tm_blob_data[v_offset + 16] : uvec2(0); - tmp_v1[tmp_vi + 16] = (gx + lxm8 + 24) < psc(outw) ? bottom_tm_blob_data[v_offset + 24] : uvec2(0); - - const int tmp_ki = lxd8*8*2 + lxm8 * 2 + j; - - int w_offset = gz * psc(outc) * psc(c) * 4 + gy * psc(c) * 4 + (z + lxd8) * 2 * 8 + (lxm8 * 2 + j); - - tmp_k0[tmp_ki] = weight_tm_data[w_offset]; - tmp_k1[tmp_ki] = weight_tm_data[w_offset + psc(c) * 8]; - tmp_k2[tmp_ki] = weight_tm_data[w_offset + psc(c) * 16]; - tmp_k3[tmp_ki] = weight_tm_data[w_offset + psc(c) * 24]; - } - } - - barrier(); - - for (int z4 = 0; z4 < remain; z4++) - { -#if ncnn_VK_KHR_cooperative_matrix - coopmat A0; - coopmat A1; - coopMatLoad(A0, tmp_v0, z4*16*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(A1, tmp_v1, z4*16*2, 2, gl_CooperativeMatrixLayoutRowMajor); - - coopmat B0; - coopmat B1; - coopmat B2; - coopmat B3; - coopMatLoad(B0, tmp_k0, z4*8*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(B1, tmp_k1, z4*8*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(B2, tmp_k2, z4*8*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(B3, tmp_k3, z4*8*2, 2, gl_CooperativeMatrixLayoutRowMajor); - - // sum += v * k - sum0 = coopMatMulAdd(A0, B0, sum0); - sum1 = coopMatMulAdd(A1, B0, sum1); - sum2 = coopMatMulAdd(A0, B1, sum2); - sum3 = coopMatMulAdd(A1, B1, sum3); - sum4 = coopMatMulAdd(A0, B2, sum4); - sum5 = coopMatMulAdd(A1, B2, sum5); - sum6 = coopMatMulAdd(A0, B3, sum6); - sum7 = coopMatMulAdd(A1, B3, sum7); -#elif ncnn_VK_NV_cooperative_matrix - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> A0; - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> A1; - coopMatLoadNV(A0, tmp_v0, z4*16*2, 2, false); - coopMatLoadNV(A1, tmp_v1, z4*16*2, 2, false); - - fcoopmatNV<16, gl_ScopeSubgroup, 8, 8> B0; - fcoopmatNV<16, gl_ScopeSubgroup, 8, 8> B1; - fcoopmatNV<16, gl_ScopeSubgroup, 8, 8> B2; - fcoopmatNV<16, gl_ScopeSubgroup, 8, 8> B3; - coopMatLoadNV(B0, tmp_k0, z4*8*2, 2, false); - coopMatLoadNV(B1, tmp_k1, z4*8*2, 2, false); - coopMatLoadNV(B2, tmp_k2, z4*8*2, 2, false); - coopMatLoadNV(B3, tmp_k3, z4*8*2, 2, false); - - // sum += v * k - sum0 = coopMatMulAddNV(A0, B0, sum0); - sum1 = coopMatMulAddNV(A1, B0, sum1); - sum2 = coopMatMulAddNV(A0, B1, sum2); - sum3 = coopMatMulAddNV(A1, B1, sum3); - sum4 = coopMatMulAddNV(A0, B2, sum4); - sum5 = coopMatMulAddNV(A1, B2, sum5); - sum6 = coopMatMulAddNV(A0, B3, sum6); - sum7 = coopMatMulAddNV(A1, B3, sum7); -#endif - } - - barrier(); - } - - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) - return; - -#if ncnn_VK_KHR_cooperative_matrix -#if NCNN_fp16_arithmetic - coopMatStore(sum0, tmp_v0, 0, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum1, tmp_v1, 0, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum2, tmp_v0, 16*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum3, tmp_v1, 16*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum4, tmp_v0, 16*4, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum5, tmp_v1, 16*4, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum6, tmp_v0, 16*6, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum7, tmp_v1, 16*6, 2, gl_CooperativeMatrixLayoutRowMajor); -#else - coopmat sum0_fp16 = coopmat(sum0); - coopmat sum1_fp16 = coopmat(sum1); - coopmat sum2_fp16 = coopmat(sum2); - coopmat sum3_fp16 = coopmat(sum3); - coopmat sum4_fp16 = coopmat(sum4); - coopmat sum5_fp16 = coopmat(sum5); - coopmat sum6_fp16 = coopmat(sum6); - coopmat sum7_fp16 = coopmat(sum7); - - coopMatStore(sum0_fp16, tmp_v0, 0, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum1_fp16, tmp_v1, 0, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum2_fp16, tmp_v0, 16*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum3_fp16, tmp_v1, 16*2, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum4_fp16, tmp_v0, 16*4, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum5_fp16, tmp_v1, 16*4, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum6_fp16, tmp_v0, 16*6, 2, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(sum7_fp16, tmp_v1, 16*6, 2, gl_CooperativeMatrixLayoutRowMajor); -#endif -#elif ncnn_VK_NV_cooperative_matrix -#if NCNN_fp16_arithmetic - coopMatStoreNV(sum0, tmp_v0, 0, 2, false); - coopMatStoreNV(sum1, tmp_v1, 0, 2, false); - coopMatStoreNV(sum2, tmp_v0, 16*2, 2, false); - coopMatStoreNV(sum3, tmp_v1, 16*2, 2, false); - coopMatStoreNV(sum4, tmp_v0, 16*4, 2, false); - coopMatStoreNV(sum5, tmp_v1, 16*4, 2, false); - coopMatStoreNV(sum6, tmp_v0, 16*6, 2, false); - coopMatStoreNV(sum7, tmp_v1, 16*6, 2, false); -#else - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum0_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(sum0); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum1_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(sum1); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum2_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(sum2); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum3_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(sum3); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum4_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(sum4); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum5_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(sum5); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum6_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(sum6); - fcoopmatNV<16, gl_ScopeSubgroup, 16, 8> sum7_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, 16, 8>(sum7); - - coopMatStoreNV(sum0_fp16, tmp_v0, 0, 2, false); - coopMatStoreNV(sum1_fp16, tmp_v1, 0, 2, false); - coopMatStoreNV(sum2_fp16, tmp_v0, 16*2, 2, false); - coopMatStoreNV(sum3_fp16, tmp_v1, 16*2, 2, false); - coopMatStoreNV(sum4_fp16, tmp_v0, 16*4, 2, false); - coopMatStoreNV(sum5_fp16, tmp_v1, 16*4, 2, false); - coopMatStoreNV(sum6_fp16, tmp_v0, 16*6, 2, false); - coopMatStoreNV(sum7_fp16, tmp_v1, 16*6, 2, false); -#endif -#endif - - barrier(); - - const int lxd16 = lx / 16; // 0 1 - const int lxm16 = lx % 16; // 0 1 2 3 .... 15 - - { - for (int j = 0; j < 4; j++) - { - const int tmp_vi = lxm16 * 2 + lxd16 + j*16*2; - const int gi = gz * psc(outcstep) + (gy + lxd16 + j*2) * psc(outw) + (gx + lxm16); - - if (gy + j * 2 + lxd16 < psc(outc)) - { - if (gx + lxm16 < psc(outw)) top_tm_blob_data[gi] = tmp_v0[tmp_vi]; - if (gx + lxm16 + 16 < psc(outw)) top_tm_blob_data[gi + 16] = tmp_v1[tmp_vi]; - } - } - } -} diff --git a/src/layer/vulkan/shader/convolution_pack4to1_3x3s1d1_winograd_gemm.comp b/src/layer/vulkan/shader/convolution_pack4to1_3x3s1d1_winograd_gemm.comp index d425fe501..017c77578 100644 --- a/src/layer/vulkan/shader/convolution_pack4to1_3x3s1d1_winograd_gemm.comp +++ b/src/layer/vulkan/shader/convolution_pack4to1_3x3s1d1_winograd_gemm.comp @@ -6,14 +6,13 @@ #define LOCAL_MEMORY_UNROLL_INCH 8 layout (constant_id = 0) const int batch = 1; +layout (constant_id = 1) const int c = 0; +layout (constant_id = 2) const int outc = 0; -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +#define shape_constant_id_offset 3 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; +layout (constant_id = shape_constant_id_offset + 1) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 2) const int outcstep = 0; layout (binding = 0) readonly buffer bottom_tm_blob { sfpvec4 bottom_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_tm_blob { sfp top_tm_blob_data[]; }; @@ -21,11 +20,8 @@ layout (binding = 2) readonly buffer weight_tm_blob { sfpvec4 weight_tm_data[]; layout (push_constant) uniform parameter { - int c; int cstep; - int outw; - int outc; int outcstep; } p; @@ -41,7 +37,7 @@ void main() int gz = int(gl_GlobalInvocationID.z); #if !NCNN_shader_local_memory - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; #endif @@ -50,21 +46,21 @@ void main() afp sum2 = afp(0.f); afp sum3 = afp(0.f); - int v_offset = gz * psc(cstep) + gx; - int w_offset = gz * psc(c) * psc(outc) + gy * psc(c); + int v_offset = gz * c * psc(cstep) + gx; + int w_offset = gz * c * outc + gy * c; #if NCNN_shader_local_memory const int lx = int(gl_LocalInvocationID.x); const int ly = int(gl_LocalInvocationID.y); int z = 0; - for (; z + (LOCAL_MEMORY_UNROLL_INCH - 1) < psc(c); z += LOCAL_MEMORY_UNROLL_INCH) + for (; z + (LOCAL_MEMORY_UNROLL_INCH - 1) < c; z += LOCAL_MEMORY_UNROLL_INCH) { if (ly < 4) { for (int z4 = 0; z4 < LOCAL_MEMORY_UNROLL_INCH; z4++) { - tmp_v[lx][z4][ly] = sfp2lfpvec4(bottom_tm_blob_data[v_offset + z4 * psc(outw) + ly]); + tmp_v[lx][z4][ly] = sfp2lfpvec4(bottom_tm_blob_data[v_offset + z4 * psc(cstep) + ly]); } } @@ -93,21 +89,21 @@ void main() sum3 += dot(v3, k); } - v_offset += LOCAL_MEMORY_UNROLL_INCH * psc(outw); + v_offset += LOCAL_MEMORY_UNROLL_INCH * psc(cstep); w_offset += LOCAL_MEMORY_UNROLL_INCH; barrier(); } - if (z < psc(c)) + if (z < c) { - const int remain = psc(c) - z; + const int remain = c - z; if (ly < 4) { for (int z4 = 0; z4 < remain; z4++) { - tmp_v[lx][z4][ly] = sfp2lfpvec4(bottom_tm_blob_data[v_offset + z4 * psc(outw) + ly]); + tmp_v[lx][z4][ly] = sfp2lfpvec4(bottom_tm_blob_data[v_offset + z4 * psc(cstep) + ly]); } } @@ -137,7 +133,7 @@ void main() } } #else - for (int z = 0; z < psc(c); z++) + for (int z = 0; z < c; z++) { afpvec4 v0 = buffer_ld4(bottom_tm_blob_data, v_offset + 0); afpvec4 v1 = buffer_ld4(bottom_tm_blob_data, v_offset + 1); @@ -151,17 +147,17 @@ void main() sum2 += dot(v2, k); sum3 += dot(v3, k); - v_offset += psc(outw); + v_offset += psc(cstep); w_offset += 1; } #endif #if NCNN_shader_local_memory - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; #endif - int gi = gz * psc(outcstep) + gy * psc(outw) + gx; + int gi = gz * outc * psc(outcstep) + gy * psc(outcstep) + gx; buffer_st1(top_tm_blob_data, gi + 0, sum0); if (gx + 1 < psc(outw)) buffer_st1(top_tm_blob_data, gi + 1, sum1); diff --git a/src/layer/vulkan/shader/convolution_pack4to8_3x3s1d1_winograd_gemm.comp b/src/layer/vulkan/shader/convolution_pack4to8_3x3s1d1_winograd_gemm.comp index c377a1e84..b42eb4ac2 100644 --- a/src/layer/vulkan/shader/convolution_pack4to8_3x3s1d1_winograd_gemm.comp +++ b/src/layer/vulkan/shader/convolution_pack4to8_3x3s1d1_winograd_gemm.comp @@ -4,14 +4,13 @@ #version 450 layout (constant_id = 0) const int batch = 1; +layout (constant_id = 1) const int c = 0; +layout (constant_id = 2) const int outc = 0; -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +#define shape_constant_id_offset 3 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; +layout (constant_id = shape_constant_id_offset + 1) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 2) const int outcstep = 0; layout (binding = 0) readonly buffer bottom_tm_blob { sfpvec4 bottom_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_tm_blob { sfpvec8 top_tm_blob_data[]; }; @@ -19,11 +18,8 @@ layout (binding = 2) readonly buffer weight_tm_blob { sfpvec4 weight_tm_data[]; layout (push_constant) uniform parameter { - int c; int cstep; - int outw; - int outc; int outcstep; } p; @@ -33,7 +29,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; afpvec8 sum0 = afpvec8(afpvec4(0.f), afpvec4(0.f)); @@ -41,10 +37,10 @@ void main() afpvec8 sum2 = afpvec8(afpvec4(0.f), afpvec4(0.f)); afpvec8 sum3 = afpvec8(afpvec4(0.f), afpvec4(0.f)); - int v_offset = gz * psc(cstep) + gx; - int w_offset = (gz * psc(c) * psc(outc) + gy * psc(c)) * 8; + int v_offset = gz * c * psc(cstep) + gx; + int w_offset = (gz * c * outc + gy * c) * 8; - for (int z = 0; z < psc(c); z++) + for (int z = 0; z < c; z++) { afpvec4 v0 = buffer_ld4(bottom_tm_blob_data, v_offset + 0); afpvec4 v1 = buffer_ld4(bottom_tm_blob_data, v_offset + 1); @@ -97,11 +93,11 @@ void main() sum3[1].b += dot(v3, k6); sum3[1].a += dot(v3, k7); - v_offset += psc(outw); + v_offset += psc(cstep); w_offset += 8; } - int gi = gz * psc(outcstep) + gy * psc(outw) + gx; + int gi = gz * outc * psc(outcstep) + gy * psc(outcstep) + gx; buffer_st8(top_tm_blob_data, gi + 0, sum0); if (gx + 1 < psc(outw)) buffer_st8(top_tm_blob_data, gi + 1, sum1); diff --git a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd23_transform_input.comp b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd23_transform_input.comp index e02c4a706..b7afb5ca9 100644 --- a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd23_transform_input.comp +++ b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd23_transform_input.comp @@ -3,16 +3,17 @@ #version 450 -#define shape_constant_id_offset 0 +layout (constant_id = 0) const int c = 0; + +#define shape_constant_id_offset 1 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 (constant_id = shape_constant_id_offset + 2) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outcstep = 0; -layout (constant_id = shape_constant_id_offset + 5) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 6) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 4) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 5) const int block_y = 0; layout (binding = 0) readonly buffer bottom_blob { sfpvec8 bottom_blob_data[]; }; layout (binding = 1) writeonly buffer bottom_tm_blob { sfpvec8 bottom_tm_blob_data[]; }; @@ -21,7 +22,6 @@ layout (push_constant) uniform parameter { int w; int h; - int c; int cstep; int outcstep; @@ -36,7 +36,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 4x4 @@ -115,22 +115,22 @@ void main() v33 = m33 - m31; // store 16 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - buffer_st8(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep), v00); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep), v01); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep), v02); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep), v03); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep), v10); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep), v11); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep), v12); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep), v13); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep), v20); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep), v21); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep), v22); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep), v23); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep), v30); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep), v31); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep), v32); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep), v33); + int v_tm_offset = gz * psc(outcstep) + gy * psc(block_x) + gx; + + buffer_st8(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep) * c, v00); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep) * c, v01); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep) * c, v02); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep) * c, v03); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep) * c, v10); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep) * c, v11); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep) * c, v12); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep) * c, v13); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep) * c, v20); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep) * c, v21); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep) * c, v22); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep) * c, v23); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep) * c, v30); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep) * c, v31); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep) * c, v32); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep) * c, v33); } diff --git a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd23_transform_output.comp b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd23_transform_output.comp index 7749464e0..8d3184761 100644 --- a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd23_transform_output.comp +++ b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd23_transform_output.comp @@ -10,17 +10,17 @@ layout (constant_id = 0) const int bias_term = 0; layout (constant_id = 1) const int activation_type = 0; layout (constant_id = 2) const float activation_param_0 = 0; layout (constant_id = 3) const float activation_param_1 = 0; +layout (constant_id = 4) const int c = 0; -#define shape_constant_id_offset 4 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; +#define shape_constant_id_offset 5 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 2) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 3) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 1) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 2) const int block_y = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 5) const int outh = 0; -layout (constant_id = shape_constant_id_offset + 6) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 4) const int outh = 0; +layout (constant_id = shape_constant_id_offset + 5) const int outcstep = 0; layout (binding = 0) readonly buffer top_tm_blob { sfpvec8 top_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_blob { sfpvec8 top_blob_data[]; }; @@ -28,7 +28,6 @@ layout (binding = 2) readonly buffer bias_blob { sfpvec8 bias_data[]; }; layout (push_constant) uniform parameter { - int c; int cstep; int block_x; @@ -45,28 +44,28 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= psc(c)) + if (gx >= psc(block_x) || gy >= psc(block_y) || gz >= c) return; // load 16 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - afpvec8 v00 = buffer_ld8(top_tm_blob_data, v_tm_offset + 0 * psc(cstep)); - afpvec8 v01 = buffer_ld8(top_tm_blob_data, v_tm_offset + 1 * psc(cstep)); - afpvec8 v02 = buffer_ld8(top_tm_blob_data, v_tm_offset + 2 * psc(cstep)); - afpvec8 v03 = buffer_ld8(top_tm_blob_data, v_tm_offset + 3 * psc(cstep)); - afpvec8 v10 = buffer_ld8(top_tm_blob_data, v_tm_offset + 4 * psc(cstep)); - afpvec8 v11 = buffer_ld8(top_tm_blob_data, v_tm_offset + 5 * psc(cstep)); - afpvec8 v12 = buffer_ld8(top_tm_blob_data, v_tm_offset + 6 * psc(cstep)); - afpvec8 v13 = buffer_ld8(top_tm_blob_data, v_tm_offset + 7 * psc(cstep)); - afpvec8 v20 = buffer_ld8(top_tm_blob_data, v_tm_offset + 8 * psc(cstep)); - afpvec8 v21 = buffer_ld8(top_tm_blob_data, v_tm_offset + 9 * psc(cstep)); - afpvec8 v22 = buffer_ld8(top_tm_blob_data, v_tm_offset + 10 * psc(cstep)); - afpvec8 v23 = buffer_ld8(top_tm_blob_data, v_tm_offset + 11 * psc(cstep)); - afpvec8 v30 = buffer_ld8(top_tm_blob_data, v_tm_offset + 12 * psc(cstep)); - afpvec8 v31 = buffer_ld8(top_tm_blob_data, v_tm_offset + 13 * psc(cstep)); - afpvec8 v32 = buffer_ld8(top_tm_blob_data, v_tm_offset + 14 * psc(cstep)); - afpvec8 v33 = buffer_ld8(top_tm_blob_data, v_tm_offset + 15 * psc(cstep)); + int v_tm_offset = gz * psc(cstep) + gy * psc(block_x) + gx; + + afpvec8 v00 = buffer_ld8(top_tm_blob_data, v_tm_offset + 0 * psc(cstep) * c); + afpvec8 v01 = buffer_ld8(top_tm_blob_data, v_tm_offset + 1 * psc(cstep) * c); + afpvec8 v02 = buffer_ld8(top_tm_blob_data, v_tm_offset + 2 * psc(cstep) * c); + afpvec8 v03 = buffer_ld8(top_tm_blob_data, v_tm_offset + 3 * psc(cstep) * c); + afpvec8 v10 = buffer_ld8(top_tm_blob_data, v_tm_offset + 4 * psc(cstep) * c); + afpvec8 v11 = buffer_ld8(top_tm_blob_data, v_tm_offset + 5 * psc(cstep) * c); + afpvec8 v12 = buffer_ld8(top_tm_blob_data, v_tm_offset + 6 * psc(cstep) * c); + afpvec8 v13 = buffer_ld8(top_tm_blob_data, v_tm_offset + 7 * psc(cstep) * c); + afpvec8 v20 = buffer_ld8(top_tm_blob_data, v_tm_offset + 8 * psc(cstep) * c); + afpvec8 v21 = buffer_ld8(top_tm_blob_data, v_tm_offset + 9 * psc(cstep) * c); + afpvec8 v22 = buffer_ld8(top_tm_blob_data, v_tm_offset + 10 * psc(cstep) * c); + afpvec8 v23 = buffer_ld8(top_tm_blob_data, v_tm_offset + 11 * psc(cstep) * c); + afpvec8 v30 = buffer_ld8(top_tm_blob_data, v_tm_offset + 12 * psc(cstep) * c); + afpvec8 v31 = buffer_ld8(top_tm_blob_data, v_tm_offset + 13 * psc(cstep) * c); + afpvec8 v32 = buffer_ld8(top_tm_blob_data, v_tm_offset + 14 * psc(cstep) * c); + afpvec8 v33 = buffer_ld8(top_tm_blob_data, v_tm_offset + 15 * psc(cstep) * c); // const float itm[2][4] = { // {1.0f, 1.0f, 1.0f, 0.0f}, diff --git a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd43_transform_input.comp b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd43_transform_input.comp index 76236386a..3e86574d0 100644 --- a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd43_transform_input.comp +++ b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd43_transform_input.comp @@ -3,16 +3,17 @@ #version 450 -#define shape_constant_id_offset 0 +layout (constant_id = 0) const int c = 0; + +#define shape_constant_id_offset 1 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 (constant_id = shape_constant_id_offset + 2) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outcstep = 0; -layout (constant_id = shape_constant_id_offset + 5) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 6) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 4) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 5) const int block_y = 0; layout (binding = 0) readonly buffer bottom_blob { sfpvec8 bottom_blob_data[]; }; layout (binding = 1) writeonly buffer bottom_tm_blob { sfpvec8 bottom_tm_blob_data[]; }; @@ -21,7 +22,6 @@ layout (push_constant) uniform parameter { int w; int h; - int c; int cstep; int outcstep; @@ -36,7 +36,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= p.block_x || gy >= p.block_y || gz >= psc(c)) + if (gx >= p.block_x || gy >= p.block_y || gz >= c) return; // load 6x6 @@ -187,42 +187,42 @@ void main() v55 = m51 - m53 * afp(2.5) + m55; // store 36 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - buffer_st8(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep), v00); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep), v01); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep), v02); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep), v03); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep), v04); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep), v05); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep), v10); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep), v11); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep), v12); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep), v13); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep), v14); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep), v15); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep), v20); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep), v21); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep), v22); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep), v23); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 16 * psc(outcstep), v24); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 17 * psc(outcstep), v25); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 18 * psc(outcstep), v30); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 19 * psc(outcstep), v31); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 20 * psc(outcstep), v32); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 21 * psc(outcstep), v33); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 22 * psc(outcstep), v34); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 23 * psc(outcstep), v35); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 24 * psc(outcstep), v40); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 25 * psc(outcstep), v41); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 26 * psc(outcstep), v42); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 27 * psc(outcstep), v43); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 28 * psc(outcstep), v44); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 29 * psc(outcstep), v45); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 30 * psc(outcstep), v50); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 31 * psc(outcstep), v51); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 32 * psc(outcstep), v52); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 33 * psc(outcstep), v53); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 34 * psc(outcstep), v54); - buffer_st8(bottom_tm_blob_data, v_tm_offset + 35 * psc(outcstep), v55); + int v_tm_offset = gz * psc(outcstep) + gy * psc(block_x) + gx; + + buffer_st8(bottom_tm_blob_data, v_tm_offset + 0 * psc(outcstep) * c, v00); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 1 * psc(outcstep) * c, v01); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 2 * psc(outcstep) * c, v02); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 3 * psc(outcstep) * c, v03); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 4 * psc(outcstep) * c, v04); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 5 * psc(outcstep) * c, v05); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 6 * psc(outcstep) * c, v10); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 7 * psc(outcstep) * c, v11); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 8 * psc(outcstep) * c, v12); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 9 * psc(outcstep) * c, v13); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 10 * psc(outcstep) * c, v14); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 11 * psc(outcstep) * c, v15); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 12 * psc(outcstep) * c, v20); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 13 * psc(outcstep) * c, v21); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 14 * psc(outcstep) * c, v22); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 15 * psc(outcstep) * c, v23); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 16 * psc(outcstep) * c, v24); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 17 * psc(outcstep) * c, v25); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 18 * psc(outcstep) * c, v30); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 19 * psc(outcstep) * c, v31); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 20 * psc(outcstep) * c, v32); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 21 * psc(outcstep) * c, v33); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 22 * psc(outcstep) * c, v34); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 23 * psc(outcstep) * c, v35); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 24 * psc(outcstep) * c, v40); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 25 * psc(outcstep) * c, v41); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 26 * psc(outcstep) * c, v42); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 27 * psc(outcstep) * c, v43); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 28 * psc(outcstep) * c, v44); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 29 * psc(outcstep) * c, v45); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 30 * psc(outcstep) * c, v50); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 31 * psc(outcstep) * c, v51); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 32 * psc(outcstep) * c, v52); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 33 * psc(outcstep) * c, v53); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 34 * psc(outcstep) * c, v54); + buffer_st8(bottom_tm_blob_data, v_tm_offset + 35 * psc(outcstep) * c, v55); } diff --git a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd43_transform_output.comp b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd43_transform_output.comp index 12a7f9e31..1a2cd0775 100644 --- a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd43_transform_output.comp +++ b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd43_transform_output.comp @@ -10,17 +10,17 @@ layout (constant_id = 0) const int bias_term = 0; layout (constant_id = 1) const int activation_type = 0; layout (constant_id = 2) const float activation_param_0 = 0; layout (constant_id = 3) const float activation_param_1 = 0; +layout (constant_id = 4) const int c = 0; -#define shape_constant_id_offset 4 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; +#define shape_constant_id_offset 5 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; -layout (constant_id = shape_constant_id_offset + 2) const int block_x = 0; -layout (constant_id = shape_constant_id_offset + 3) const int block_y = 0; +layout (constant_id = shape_constant_id_offset + 1) const int block_x = 0; +layout (constant_id = shape_constant_id_offset + 2) const int block_y = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 5) const int outh = 0; -layout (constant_id = shape_constant_id_offset + 6) const int outcstep = 0; +layout (constant_id = shape_constant_id_offset + 3) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 4) const int outh = 0; +layout (constant_id = shape_constant_id_offset + 5) const int outcstep = 0; layout (binding = 0) readonly buffer top_tm_blob { sfpvec8 top_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_blob { sfpvec8 top_blob_data[]; }; @@ -28,7 +28,6 @@ layout (binding = 2) readonly buffer bias_blob { sfpvec8 bias_data[]; }; layout (push_constant) uniform parameter { - int c; int cstep; int block_x; @@ -45,48 +44,48 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= p.block_x || gy >= p.block_y || gz >= psc(c)) + if (gx >= p.block_x || gy >= p.block_y || gz >= c) return; // load 36 - int v_tm_offset = gz * psc(block_x) * psc(block_y) + gy * psc(block_x) + gx; - - afpvec8 v00 = buffer_ld8(top_tm_blob_data, v_tm_offset + 0 * psc(cstep)); - afpvec8 v01 = buffer_ld8(top_tm_blob_data, v_tm_offset + 1 * psc(cstep)); - afpvec8 v02 = buffer_ld8(top_tm_blob_data, v_tm_offset + 2 * psc(cstep)); - afpvec8 v03 = buffer_ld8(top_tm_blob_data, v_tm_offset + 3 * psc(cstep)); - afpvec8 v04 = buffer_ld8(top_tm_blob_data, v_tm_offset + 4 * psc(cstep)); - afpvec8 v05 = buffer_ld8(top_tm_blob_data, v_tm_offset + 5 * psc(cstep)); - afpvec8 v10 = buffer_ld8(top_tm_blob_data, v_tm_offset + 6 * psc(cstep)); - afpvec8 v11 = buffer_ld8(top_tm_blob_data, v_tm_offset + 7 * psc(cstep)); - afpvec8 v12 = buffer_ld8(top_tm_blob_data, v_tm_offset + 8 * psc(cstep)); - afpvec8 v13 = buffer_ld8(top_tm_blob_data, v_tm_offset + 9 * psc(cstep)); - afpvec8 v14 = buffer_ld8(top_tm_blob_data, v_tm_offset + 10 * psc(cstep)); - afpvec8 v15 = buffer_ld8(top_tm_blob_data, v_tm_offset + 11 * psc(cstep)); - afpvec8 v20 = buffer_ld8(top_tm_blob_data, v_tm_offset + 12 * psc(cstep)); - afpvec8 v21 = buffer_ld8(top_tm_blob_data, v_tm_offset + 13 * psc(cstep)); - afpvec8 v22 = buffer_ld8(top_tm_blob_data, v_tm_offset + 14 * psc(cstep)); - afpvec8 v23 = buffer_ld8(top_tm_blob_data, v_tm_offset + 15 * psc(cstep)); - afpvec8 v24 = buffer_ld8(top_tm_blob_data, v_tm_offset + 16 * psc(cstep)); - afpvec8 v25 = buffer_ld8(top_tm_blob_data, v_tm_offset + 17 * psc(cstep)); - afpvec8 v30 = buffer_ld8(top_tm_blob_data, v_tm_offset + 18 * psc(cstep)); - afpvec8 v31 = buffer_ld8(top_tm_blob_data, v_tm_offset + 19 * psc(cstep)); - afpvec8 v32 = buffer_ld8(top_tm_blob_data, v_tm_offset + 20 * psc(cstep)); - afpvec8 v33 = buffer_ld8(top_tm_blob_data, v_tm_offset + 21 * psc(cstep)); - afpvec8 v34 = buffer_ld8(top_tm_blob_data, v_tm_offset + 22 * psc(cstep)); - afpvec8 v35 = buffer_ld8(top_tm_blob_data, v_tm_offset + 23 * psc(cstep)); - afpvec8 v40 = buffer_ld8(top_tm_blob_data, v_tm_offset + 24 * psc(cstep)); - afpvec8 v41 = buffer_ld8(top_tm_blob_data, v_tm_offset + 25 * psc(cstep)); - afpvec8 v42 = buffer_ld8(top_tm_blob_data, v_tm_offset + 26 * psc(cstep)); - afpvec8 v43 = buffer_ld8(top_tm_blob_data, v_tm_offset + 27 * psc(cstep)); - afpvec8 v44 = buffer_ld8(top_tm_blob_data, v_tm_offset + 28 * psc(cstep)); - afpvec8 v45 = buffer_ld8(top_tm_blob_data, v_tm_offset + 29 * psc(cstep)); - afpvec8 v50 = buffer_ld8(top_tm_blob_data, v_tm_offset + 30 * psc(cstep)); - afpvec8 v51 = buffer_ld8(top_tm_blob_data, v_tm_offset + 31 * psc(cstep)); - afpvec8 v52 = buffer_ld8(top_tm_blob_data, v_tm_offset + 32 * psc(cstep)); - afpvec8 v53 = buffer_ld8(top_tm_blob_data, v_tm_offset + 33 * psc(cstep)); - afpvec8 v54 = buffer_ld8(top_tm_blob_data, v_tm_offset + 34 * psc(cstep)); - afpvec8 v55 = buffer_ld8(top_tm_blob_data, v_tm_offset + 35 * psc(cstep)); + int v_tm_offset = gz * psc(cstep) + gy * psc(block_x) + gx; + + afpvec8 v00 = buffer_ld8(top_tm_blob_data, v_tm_offset + 0 * psc(cstep) * c); + afpvec8 v01 = buffer_ld8(top_tm_blob_data, v_tm_offset + 1 * psc(cstep) * c); + afpvec8 v02 = buffer_ld8(top_tm_blob_data, v_tm_offset + 2 * psc(cstep) * c); + afpvec8 v03 = buffer_ld8(top_tm_blob_data, v_tm_offset + 3 * psc(cstep) * c); + afpvec8 v04 = buffer_ld8(top_tm_blob_data, v_tm_offset + 4 * psc(cstep) * c); + afpvec8 v05 = buffer_ld8(top_tm_blob_data, v_tm_offset + 5 * psc(cstep) * c); + afpvec8 v10 = buffer_ld8(top_tm_blob_data, v_tm_offset + 6 * psc(cstep) * c); + afpvec8 v11 = buffer_ld8(top_tm_blob_data, v_tm_offset + 7 * psc(cstep) * c); + afpvec8 v12 = buffer_ld8(top_tm_blob_data, v_tm_offset + 8 * psc(cstep) * c); + afpvec8 v13 = buffer_ld8(top_tm_blob_data, v_tm_offset + 9 * psc(cstep) * c); + afpvec8 v14 = buffer_ld8(top_tm_blob_data, v_tm_offset + 10 * psc(cstep) * c); + afpvec8 v15 = buffer_ld8(top_tm_blob_data, v_tm_offset + 11 * psc(cstep) * c); + afpvec8 v20 = buffer_ld8(top_tm_blob_data, v_tm_offset + 12 * psc(cstep) * c); + afpvec8 v21 = buffer_ld8(top_tm_blob_data, v_tm_offset + 13 * psc(cstep) * c); + afpvec8 v22 = buffer_ld8(top_tm_blob_data, v_tm_offset + 14 * psc(cstep) * c); + afpvec8 v23 = buffer_ld8(top_tm_blob_data, v_tm_offset + 15 * psc(cstep) * c); + afpvec8 v24 = buffer_ld8(top_tm_blob_data, v_tm_offset + 16 * psc(cstep) * c); + afpvec8 v25 = buffer_ld8(top_tm_blob_data, v_tm_offset + 17 * psc(cstep) * c); + afpvec8 v30 = buffer_ld8(top_tm_blob_data, v_tm_offset + 18 * psc(cstep) * c); + afpvec8 v31 = buffer_ld8(top_tm_blob_data, v_tm_offset + 19 * psc(cstep) * c); + afpvec8 v32 = buffer_ld8(top_tm_blob_data, v_tm_offset + 20 * psc(cstep) * c); + afpvec8 v33 = buffer_ld8(top_tm_blob_data, v_tm_offset + 21 * psc(cstep) * c); + afpvec8 v34 = buffer_ld8(top_tm_blob_data, v_tm_offset + 22 * psc(cstep) * c); + afpvec8 v35 = buffer_ld8(top_tm_blob_data, v_tm_offset + 23 * psc(cstep) * c); + afpvec8 v40 = buffer_ld8(top_tm_blob_data, v_tm_offset + 24 * psc(cstep) * c); + afpvec8 v41 = buffer_ld8(top_tm_blob_data, v_tm_offset + 25 * psc(cstep) * c); + afpvec8 v42 = buffer_ld8(top_tm_blob_data, v_tm_offset + 26 * psc(cstep) * c); + afpvec8 v43 = buffer_ld8(top_tm_blob_data, v_tm_offset + 27 * psc(cstep) * c); + afpvec8 v44 = buffer_ld8(top_tm_blob_data, v_tm_offset + 28 * psc(cstep) * c); + afpvec8 v45 = buffer_ld8(top_tm_blob_data, v_tm_offset + 29 * psc(cstep) * c); + afpvec8 v50 = buffer_ld8(top_tm_blob_data, v_tm_offset + 30 * psc(cstep) * c); + afpvec8 v51 = buffer_ld8(top_tm_blob_data, v_tm_offset + 31 * psc(cstep) * c); + afpvec8 v52 = buffer_ld8(top_tm_blob_data, v_tm_offset + 32 * psc(cstep) * c); + afpvec8 v53 = buffer_ld8(top_tm_blob_data, v_tm_offset + 33 * psc(cstep) * c); + afpvec8 v54 = buffer_ld8(top_tm_blob_data, v_tm_offset + 34 * psc(cstep) * c); + afpvec8 v55 = buffer_ld8(top_tm_blob_data, v_tm_offset + 35 * psc(cstep) * c); #define sq2 1.41421356237 #define sq2_m2 1.41421356237*2 diff --git a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd_gemm.comp b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd_gemm.comp index e9436d3ba..ae19570b2 100644 --- a/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd_gemm.comp +++ b/src/layer/vulkan/shader/convolution_pack8_3x3s1d1_winograd_gemm.comp @@ -4,14 +4,13 @@ #version 450 layout (constant_id = 0) const int batch = 1; +layout (constant_id = 1) const int c = 0; +layout (constant_id = 2) const int outc = 0; -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +#define shape_constant_id_offset 3 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; +layout (constant_id = shape_constant_id_offset + 1) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 2) const int outcstep = 0; layout (binding = 0) readonly buffer bottom_tm_blob { sfpvec8 bottom_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_tm_blob { sfpvec8 top_tm_blob_data[]; }; @@ -19,11 +18,8 @@ layout (binding = 2) readonly buffer weight_tm_blob { sfpvec8 weight_tm_data[]; layout (push_constant) uniform parameter { - int c; int cstep; - int outw; - int outc; int outcstep; } p; @@ -33,7 +29,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; afpvec8 sum0 = afpvec8(afpvec4(0.f), afpvec4(0.f)); @@ -41,10 +37,10 @@ void main() afpvec8 sum2 = afpvec8(afpvec4(0.f), afpvec4(0.f)); afpvec8 sum3 = afpvec8(afpvec4(0.f), afpvec4(0.f)); - int v_offset = gz * psc(cstep) + gx; - int w_offset = (gz * psc(c) * psc(outc) + gy * psc(c)) * 8; + int v_offset = gz * c * psc(cstep) + gx; + int w_offset = (gz * c * outc + gy * c) * 8; - for (int z = 0; z < psc(c); z++) + for (int z = 0; z < c; z++) { afpvec8 v0 = buffer_ld8(bottom_tm_blob_data, v_offset + 0); afpvec8 v1 = buffer_ld8(bottom_tm_blob_data, v_offset + 1); @@ -97,11 +93,11 @@ void main() sum3[1].b += dot(v3[0], k6[0]) + dot(v3[1], k6[1]); sum3[1].a += dot(v3[0], k7[0]) + dot(v3[1], k7[1]); - v_offset += psc(outw); + v_offset += psc(cstep); w_offset += 8; } - int gi = gz * psc(outcstep) + gy * psc(outw) + gx; + int gi = gz * outc * psc(outcstep) + gy * psc(outcstep) + gx; buffer_st8(top_tm_blob_data, gi + 0, sum0); if (gx + 1 < psc(outw)) buffer_st8(top_tm_blob_data, gi + 1, sum1); diff --git a/src/layer/vulkan/shader/convolution_pack8to1_3x3s1d1_winograd_gemm.comp b/src/layer/vulkan/shader/convolution_pack8to1_3x3s1d1_winograd_gemm.comp index e381bc8b1..4e407891a 100644 --- a/src/layer/vulkan/shader/convolution_pack8to1_3x3s1d1_winograd_gemm.comp +++ b/src/layer/vulkan/shader/convolution_pack8to1_3x3s1d1_winograd_gemm.comp @@ -4,14 +4,13 @@ #version 450 layout (constant_id = 0) const int batch = 1; +layout (constant_id = 1) const int c = 0; +layout (constant_id = 2) const int outc = 0; -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +#define shape_constant_id_offset 3 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; +layout (constant_id = shape_constant_id_offset + 1) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 2) const int outcstep = 0; layout (binding = 0) readonly buffer bottom_tm_blob { sfpvec8 bottom_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_tm_blob { sfp top_tm_blob_data[]; }; @@ -19,11 +18,8 @@ layout (binding = 2) readonly buffer weight_tm_blob { sfpvec8 weight_tm_data[]; layout (push_constant) uniform parameter { - int c; int cstep; - int outw; - int outc; int outcstep; } p; @@ -33,7 +29,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; afp sum0 = afp(0.f); @@ -41,10 +37,10 @@ void main() afp sum2 = afp(0.f); afp sum3 = afp(0.f); - int v_offset = gz * psc(cstep) + gx; - int w_offset = gz * psc(c) * psc(outc) + gy * psc(c); + int v_offset = gz * c * psc(cstep) + gx; + int w_offset = gz * c * outc + gy * c; - for (int z = 0; z < psc(c); z++) + for (int z = 0; z < c; z++) { afpvec8 v0 = buffer_ld8(bottom_tm_blob_data, v_offset + 0); afpvec8 v1 = buffer_ld8(bottom_tm_blob_data, v_offset + 1); @@ -59,11 +55,11 @@ void main() sum2 += dot(v2[0], k[0]) + dot(v2[1], k[1]); sum3 += dot(v3[0], k[0]) + dot(v3[1], k[1]); - v_offset += psc(outw); + v_offset += psc(cstep); w_offset += 1; } - int gi = gz * psc(outcstep) + gy * psc(outw) + gx; + int gi = gz * outc * psc(outcstep) + gy * psc(outcstep) + gx; buffer_st1(top_tm_blob_data, gi + 0, sum0); if (gx + 1 < psc(outw)) buffer_st1(top_tm_blob_data, gi + 1, sum1); diff --git a/src/layer/vulkan/shader/convolution_pack8to4_3x3s1d1_winograd_gemm.comp b/src/layer/vulkan/shader/convolution_pack8to4_3x3s1d1_winograd_gemm.comp index 7f3dc739d..409f80c9d 100644 --- a/src/layer/vulkan/shader/convolution_pack8to4_3x3s1d1_winograd_gemm.comp +++ b/src/layer/vulkan/shader/convolution_pack8to4_3x3s1d1_winograd_gemm.comp @@ -4,14 +4,13 @@ #version 450 layout (constant_id = 0) const int batch = 1; +layout (constant_id = 1) const int c = 0; +layout (constant_id = 2) const int outc = 0; -#define shape_constant_id_offset 1 -layout (constant_id = shape_constant_id_offset + 0) const int c = 0; -layout (constant_id = shape_constant_id_offset + 1) const int cstep = 0; - -layout (constant_id = shape_constant_id_offset + 2) const int outw = 0; -layout (constant_id = shape_constant_id_offset + 3) const int outc = 0; -layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0; +#define shape_constant_id_offset 3 +layout (constant_id = shape_constant_id_offset + 0) const int cstep = 0; +layout (constant_id = shape_constant_id_offset + 1) const int outw = 0; +layout (constant_id = shape_constant_id_offset + 2) const int outcstep = 0; layout (binding = 0) readonly buffer bottom_tm_blob { sfpvec8 bottom_tm_blob_data[]; }; layout (binding = 1) writeonly buffer top_tm_blob { sfpvec4 top_tm_blob_data[]; }; @@ -19,11 +18,8 @@ layout (binding = 2) readonly buffer weight_tm_blob { sfpvec8 weight_tm_data[]; layout (push_constant) uniform parameter { - int c; int cstep; - int outw; - int outc; int outcstep; } p; @@ -33,7 +29,7 @@ void main() int gy = int(gl_GlobalInvocationID.y); int gz = int(gl_GlobalInvocationID.z); - if (gx >= psc(outw) || gy >= psc(outc) || gz >= batch) + if (gx >= psc(outw) || gy >= outc || gz >= batch) return; afpvec4 sum0 = afpvec4(0.f); @@ -41,10 +37,10 @@ void main() afpvec4 sum2 = afpvec4(0.f); afpvec4 sum3 = afpvec4(0.f); - int v_offset = gz * psc(cstep) + gx; - int w_offset = (gz * psc(c) * psc(outc) + gy * psc(c)) * 4; + int v_offset = gz * c * psc(cstep) + gx; + int w_offset = (gz * c * outc + gy * c) * 4; - for (int z = 0; z < psc(c); z++) + for (int z = 0; z < c; z++) { afpvec8 v0 = buffer_ld8(bottom_tm_blob_data, v_offset + 0); afpvec8 v1 = buffer_ld8(bottom_tm_blob_data, v_offset + 1); @@ -77,11 +73,11 @@ void main() sum3.b += dot(v3[0], k2[0]) + dot(v3[1], k2[1]); sum3.a += dot(v3[0], k3[0]) + dot(v3[1], k3[1]); - v_offset += psc(outw); + v_offset += psc(cstep); w_offset += 4; } - int gi = gz * psc(outcstep) + gy * psc(outw) + gx; + int gi = gz * outc * psc(outcstep) + gy * psc(outcstep) + gx; buffer_st4(top_tm_blob_data, gi + 0, sum0); if (gx + 1 < psc(outw)) buffer_st4(top_tm_blob_data, gi + 1, sum1); diff --git a/src/layer/vulkan/shader/convolution_winograd_gemm_cm.comp b/src/layer/vulkan/shader/convolution_winograd_gemm_cm.comp new file mode 100644 index 000000000..657cb173f --- /dev/null +++ b/src/layer/vulkan/shader/convolution_winograd_gemm_cm.comp @@ -0,0 +1,1183 @@ +// Copyright 2025 Tencent +// SPDX-License-Identifier: BSD-3-Clause + +#version 450 + +#extension GL_GOOGLE_include_directive: require +#include "vulkan_activation.comp" + +#extension GL_EXT_control_flow_attributes: require + +#extension GL_KHR_shader_subgroup_basic: require + +#extension GL_KHR_memory_scope_semantics: require +#extension GL_EXT_shader_explicit_arithmetic_types: require +#extension GL_EXT_shader_explicit_arithmetic_types_float16: require +#if ncnn_VK_KHR_cooperative_matrix +#extension GL_KHR_cooperative_matrix: require +#elif ncnn_VK_NV_cooperative_matrix +#extension GL_NV_cooperative_matrix: require +#endif + +layout (constant_id = 0) const uint batch = 1; +layout (constant_id = 1) const uint M = 1; +layout (constant_id = 2) const uint N = 1; +layout (constant_id = 3) const uint K = 1; +layout (constant_id = 4) const uint UNROLL_SG_M = 2; +layout (constant_id = 5) const uint UNROLL_SG_N = 2; +layout (constant_id = 6) const uint UNROLL_SG_K = 2; +layout (constant_id = 7) const uint UNROLL_WG_M = 2; +layout (constant_id = 8) const uint UNROLL_WG_N = 2; +layout (constant_id = 9) const uint inch = 1; +layout (constant_id = 10) const uint outch = 1; +layout (constant_id = 11) const uint elempack = 1; +layout (constant_id = 12) const uint out_elempack = 1; +layout (constant_id = 13) const uint wbstep = 0; + +#define shape_constant_id_offset 14 +layout (constant_id = shape_constant_id_offset + 0) const uint size = 0; +layout (constant_id = shape_constant_id_offset + 1) const uint cstep = 0; +layout (constant_id = shape_constant_id_offset + 2) const uint outcstep = 0; + +layout (binding = 0) readonly buffer bottom_tm_blob { uvec2 bottom_tm_blob_data[]; }; +layout (binding = 1) writeonly buffer top_tm_blob { uvec2 top_tm_blob_data[]; }; +layout (binding = 2) readonly buffer weight_tm_blob { uvec2 weight_tm_data[]; }; + +layout (push_constant) uniform parameter +{ + uint size; + uint cstep; + uint outcstep; +} p; + +shared uvec2 tmp_v[UNROLL_WG_M][UNROLL_SG_M * UNROLL_SG_K * M * K / 4]; + +shared uvec2 tmp_k[UNROLL_WG_N][UNROLL_SG_N * UNROLL_SG_K * K * N / 4]; + +shared uvec2 tmp_o[UNROLL_WG_N * UNROLL_WG_M][UNROLL_SG_N * UNROLL_SG_M * M * N / 4]; + +void main() +{ + // assert gl_WorkGroupSize.x == gl_SubgroupSize + // but neither gl_SubgroupSize nor gl_WorkGroupSize.x is a constant + const uint local_size = ncnn_subgroupSize * UNROLL_WG_M * UNROLL_WG_N; + + // [ WG_UN * WG_UM * [ SG_UN * SG_UM * subgroup ] ] + + // <----WG_UN----> + // +---N--+-SG_UN+------+------+ + // | | | |XXXXXX| + // M | XXXX<----coopmat + // | | | |XXXXXX| + // +-- --SG0-- --+-- --SG2-- --+ + // | | | | | + // SG_UM | | + // | | | | | + // ^ +------+--WORKGROUP--+------+ + // | | | | | | + // | | | | + // | | | | | | + // WG_UM+-- --SG1-- --+-- --SG3-- --+ + // | | | | | | + // | | | | + // | | | | | | + // v +------+------+------+------+ + // + + const uint gz = gl_GlobalInvocationID.z; + if (gz >= batch) + return; + + const uint wgi = gl_WorkGroupID.x; + const uint sgi = gl_SubgroupID; + + const uint wgmm = (psc(size) + M * UNROLL_SG_M * UNROLL_WG_M - 1) / (M * UNROLL_SG_M * UNROLL_WG_M); + const uint wgnn = (outch + N * UNROLL_SG_N * UNROLL_WG_N - 1) / (N * UNROLL_SG_N * UNROLL_WG_N); + + const uint wgmi = wgi / wgnn; + const uint wgni = wgi % wgnn; + + const uint sgmi = sgi / UNROLL_WG_N; + const uint sgni = sgi % UNROLL_WG_N; + +// const uint mm = (psc(size) + M - 1) / M; +// const uint nn = (outch + N - 1) / N; + const uint kk = (inch + K - 1) / K; + + if (wgmi >= wgmm) + return; + + const uint li = gl_LocalInvocationID.x; + const uint si = gl_SubgroupInvocationID; + + const uint Md4 = M / 4; + const uint Nd4 = N / 4; + const uint Kd4 = K / 4; + + const uint ni = (wgni * UNROLL_WG_N + sgni) * UNROLL_SG_N; + const uint mi = (wgmi * UNROLL_WG_M + sgmi) * UNROLL_SG_M; + +#if ncnn_VK_KHR_cooperative_matrix + coopmat sum[UNROLL_SG_N][UNROLL_SG_M]; +#elif ncnn_VK_NV_cooperative_matrix +#if NCNN_fp16_arithmetic + fcoopmatNV<16, gl_ScopeSubgroup, M, N> sum[UNROLL_SG_N][UNROLL_SG_M]; +#else + fcoopmatNV<32, gl_ScopeSubgroup, M, N> sum[UNROLL_SG_N][UNROLL_SG_M]; +#endif +#endif + + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { +#if ncnn_VK_KHR_cooperative_matrix + sum[zn][zm] = coopmat(0.f); +#elif ncnn_VK_NV_cooperative_matrix +#if NCNN_fp16_arithmetic + sum[zn][zm] = fcoopmatNV<16, gl_ScopeSubgroup, M, N>(0.f); +#else + sum[zn][zm] = fcoopmatNV<32, gl_ScopeSubgroup, M, N>(0.f); +#endif +#endif + } + } + + uint k = 0; + + if (kk >= UNROLL_SG_K * 2) + { + // local stack and shared memory ping-pong + + // prefetch + uvec2 prefetch_tmp_v[(UNROLL_SG_M * UNROLL_SG_K * M * K / 4 + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N)]; + uvec2 prefetch_tmp_k[(UNROLL_SG_N * UNROLL_SG_K * K * N / 4 + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M)]; + + // prefetch the very first + { + const uint ki = 0; + + // load bottom_blob + { + if (elempack == 1) + { + // +-M-+ + // K | + // +SG_UM + // | | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UM +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + const uint cstepd4 = psc(cstep) / 4; + + const uint Md4_K_USGM_USGK = Md4 * K * UNROLL_SG_M * UNROLL_SG_K; + const uint Md4_K_USGM_USGK_d_subgroupsize = (Md4_K_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Md4_K_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Md4_K_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Md4_K_USGM_USGK) + { + const uint zk = siq / (Md4 * K * UNROLL_SG_M); + const uint zmij = siq % (Md4 * K * UNROLL_SG_M); + const uint zm = zmij / (Md4 * K); + const uint ij = zmij % (Md4 * K); + const uint i = ij / Md4; + const uint j = ij % Md4; + + const uint gk = ki + zk * K + i; + const uint gm = (mi + zm) * Md4 + j; + + uvec2 v = gk < inch && gm < cstepd4 ? bottom_tm_blob_data[(gz * inch + gk) * cstepd4 + gm] : uvec2(0); + + prefetch_tmp_v[q] = v; + } + } + } + else // if (elempack == 4) + { + // +-K-+ + // M | + // +- -+ + // SG_UM | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UM +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + const uint inchd4 = inch / 4; + + const uint Kd4_M_USGM_USGK = Kd4 * M * UNROLL_SG_M * UNROLL_SG_K; + const uint Kd4_M_USGM_USGK_d_subgroupsize = (Kd4_M_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Kd4_M_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Kd4_M_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Kd4_M_USGM_USGK) + { + const uint zk = siq / (Kd4 * M * UNROLL_SG_M); + const uint zmij = siq % (Kd4 * M * UNROLL_SG_M); + const uint zmi = zmij / Kd4; + const uint j = zmij % Kd4; + + const uint gm = mi * M + zmi; + const uint gk = ki / 4 + zk * Kd4 + j; + + uvec2 v = gk < inchd4 && gm < psc(cstep) ? bottom_tm_blob_data[(gz * inchd4 + gk) * psc(cstep) + gm] : uvec2(0); + + prefetch_tmp_v[q] = v; + } + } + } + } + + // load weight + { + // +-N-+ + // K | + // +SG_UN + // | | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UN +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + // weight_tm_data coopmat_N * coopmat_K * UNROLL_SG_N * UNROLL_WG_N * kk, blocks_n + + const uint Nd4_K_USGN = Nd4 * K * UNROLL_SG_N; + const uint Nd4_K_USGN_USGK = Nd4 * K * UNROLL_SG_N * UNROLL_SG_K; + + const uint w_offset = gz * wbstep / 4 + ((wgni * kk) * UNROLL_WG_N) * Nd4_K_USGN + ((k / UNROLL_SG_K) * UNROLL_WG_N + sgni) * Nd4_K_USGN_USGK; + + const uint Nd4_K_USGN_USGK_d_subgroupsize = (Nd4_K_USGN_USGK + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); + [[unroll]] for (uint q = 0; q < Nd4_K_USGN_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_M + sgmi) * ncnn_subgroupSize + si; + + if (Nd4_K_USGN_USGK % (ncnn_subgroupSize * UNROLL_WG_M) == 0 || siq < Nd4_K_USGN_USGK) + { + prefetch_tmp_k[q] = weight_tm_data[w_offset + siq]; + } + } + } + } + + k += UNROLL_SG_K; + + for (; k + UNROLL_SG_K - 1 < kk; k += UNROLL_SG_K) + { + barrier(); + + // copy prefetch to shared memory + { + // load bottom_blob + { + if (elempack == 1) + { + const uint Md4_K_USGM_USGK = Md4 * K * UNROLL_SG_M * UNROLL_SG_K; + const uint Md4_K_USGM_USGK_d_subgroupsize = (Md4_K_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Md4_K_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Md4_K_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Md4_K_USGM_USGK) + { + tmp_v[sgmi][siq] = prefetch_tmp_v[q]; + } + } + } + else // if (elempack == 4) + { + const uint Kd4_M_USGM_USGK = Kd4 * M * UNROLL_SG_M * UNROLL_SG_K; + const uint Kd4_M_USGM_USGK_d_subgroupsize = (Kd4_M_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Kd4_M_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Kd4_M_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Kd4_M_USGM_USGK) + { + tmp_v[sgmi][siq] = prefetch_tmp_v[q]; + } + } + } + } + + // load weight + { + const uint Nd4_K_USGN_USGK = Nd4 * K * UNROLL_SG_N * UNROLL_SG_K; + const uint Nd4_K_USGN_USGK_d_subgroupsize = (Nd4_K_USGN_USGK + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); + [[unroll]] for (uint q = 0; q < Nd4_K_USGN_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_M + sgmi) * ncnn_subgroupSize + si; + + if (Nd4_K_USGN_USGK % (ncnn_subgroupSize * UNROLL_WG_M) == 0 || siq < Nd4_K_USGN_USGK) + { + tmp_k[sgni][siq] = prefetch_tmp_k[q]; + } + } + } + } + + barrier(); + + // prefetch the next + { + const uint ki = k * K; + + // load bottom_blob + { + if (elempack == 1) + { + // +-M-+ + // K | + // +SG_UM + // | | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UM +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + const uint cstepd4 = psc(cstep) / 4; + + const uint Md4_K_USGM_USGK = Md4 * K * UNROLL_SG_M * UNROLL_SG_K; + const uint Md4_K_USGM_USGK_d_subgroupsize = (Md4_K_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Md4_K_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Md4_K_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Md4_K_USGM_USGK) + { + const uint zk = siq / (Md4 * K * UNROLL_SG_M); + const uint zmij = siq % (Md4 * K * UNROLL_SG_M); + const uint zm = zmij / (Md4 * K); + const uint ij = zmij % (Md4 * K); + const uint i = ij / Md4; + const uint j = ij % Md4; + + const uint gk = ki + zk * K + i; + const uint gm = (mi + zm) * Md4 + j; + + uvec2 v = gk < inch && gm < cstepd4 ? bottom_tm_blob_data[(gz * inch + gk) * cstepd4 + gm] : uvec2(0); + + prefetch_tmp_v[q] = v; + } + } + } + else // if (elempack == 4) + { + // +-K-+ + // M | + // +- -+ + // SG_UM | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UM +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + const uint inchd4 = inch / 4; + + const uint Kd4_M_USGM_USGK = Kd4 * M * UNROLL_SG_M * UNROLL_SG_K; + const uint Kd4_M_USGM_USGK_d_subgroupsize = (Kd4_M_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Kd4_M_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Kd4_M_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Kd4_M_USGM_USGK) + { + const uint zk = siq / (Kd4 * M * UNROLL_SG_M); + const uint zmij = siq % (Kd4 * M * UNROLL_SG_M); + const uint zmi = zmij / Kd4; + const uint j = zmij % Kd4; + + const uint gm = mi * M + zmi; + const uint gk = ki / 4 + zk * Kd4 + j; + + uvec2 v = gk < inchd4 && gm < psc(cstep) ? bottom_tm_blob_data[(gz * inchd4 + gk) * psc(cstep) + gm] : uvec2(0); + + prefetch_tmp_v[q] = v; + } + } + } + } + + // load weight + { + // +-N-+ + // K | + // +SG_UN + // | | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UN +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + // weight_tm_data coopmat_N * coopmat_K * UNROLL_SG_N * UNROLL_WG_N * kk, blocks_n + + const uint Nd4_K_USGN = Nd4 * K * UNROLL_SG_N; + const uint Nd4_K_USGN_USGK = Nd4 * K * UNROLL_SG_N * UNROLL_SG_K; + + const uint w_offset = gz * wbstep / 4 + ((wgni * kk) * UNROLL_WG_N) * Nd4_K_USGN + ((k / UNROLL_SG_K) * UNROLL_WG_N + sgni) * Nd4_K_USGN_USGK; + + const uint Nd4_K_USGN_USGK_d_subgroupsize = (Nd4_K_USGN_USGK + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); + [[unroll]] for (uint q = 0; q < Nd4_K_USGN_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_M + sgmi) * ncnn_subgroupSize + si; + + if (Nd4_K_USGN_USGK % (ncnn_subgroupSize * UNROLL_WG_M) == 0 || siq < Nd4_K_USGN_USGK) + { + prefetch_tmp_k[q] = weight_tm_data[w_offset + siq]; + } + } + } + } + +#if ncnn_VK_KHR_cooperative_matrix + coopmat A[UNROLL_SG_M]; + coopmat B[UNROLL_SG_N]; +#elif ncnn_VK_NV_cooperative_matrix + fcoopmatNV<16, gl_ScopeSubgroup, M, K> A[UNROLL_SG_M]; + fcoopmatNV<16, gl_ScopeSubgroup, K, N> B[UNROLL_SG_N]; +#endif + + [[unroll]] for (uint zk = 0; zk < UNROLL_SG_K; zk++) + { + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { + if (elempack == 1) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, gl_CooperativeMatrixLayoutColumnMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, true); +#endif + } + else // if (elempack == 4) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, gl_CooperativeMatrixLayoutRowMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, false); +#endif + } + } + + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, gl_CooperativeMatrixLayoutRowMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, false); +#endif + } + + // sum += k * v + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { +#if ncnn_VK_KHR_cooperative_matrix + sum[zn][zm] = coopMatMulAdd(A[zm], B[zn], sum[zn][zm]); +#elif ncnn_VK_NV_cooperative_matrix + sum[zn][zm] = coopMatMulAddNV(A[zm], B[zn], sum[zn][zm]); +#endif + } + } + } + } + + barrier(); + + // the last copy prefetch to shared memory + { + // load bottom_blob + { + if (elempack == 1) + { + const uint Md4_K_USGM_USGK = Md4 * K * UNROLL_SG_M * UNROLL_SG_K; + const uint Md4_K_USGM_USGK_d_subgroupsize = (Md4_K_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Md4_K_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Md4_K_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Md4_K_USGM_USGK) + { + tmp_v[sgmi][siq] = prefetch_tmp_v[q]; + } + } + } + else // if (elempack == 4) + { + const uint Kd4_M_USGM_USGK = Kd4 * M * UNROLL_SG_M * UNROLL_SG_K; + const uint Kd4_M_USGM_USGK_d_subgroupsize = (Kd4_M_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Kd4_M_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Kd4_M_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Kd4_M_USGM_USGK) + { + tmp_v[sgmi][siq] = prefetch_tmp_v[q]; + } + } + } + } + + // load weight + { + const uint Nd4_K_USGN_USGK = Nd4 * K * UNROLL_SG_N * UNROLL_SG_K; + const uint Nd4_K_USGN_USGK_d_subgroupsize = (Nd4_K_USGN_USGK + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); + [[unroll]] for (uint q = 0; q < Nd4_K_USGN_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_M + sgmi) * ncnn_subgroupSize + si; + + if (Nd4_K_USGN_USGK % (ncnn_subgroupSize * UNROLL_WG_M) == 0 || siq < Nd4_K_USGN_USGK) + { + tmp_k[sgni][siq] = prefetch_tmp_k[q]; + } + } + } + } + + barrier(); + +#if ncnn_VK_KHR_cooperative_matrix + coopmat A[UNROLL_SG_M]; + coopmat B[UNROLL_SG_N]; +#elif ncnn_VK_NV_cooperative_matrix + fcoopmatNV<16, gl_ScopeSubgroup, M, K> A[UNROLL_SG_M]; + fcoopmatNV<16, gl_ScopeSubgroup, K, N> B[UNROLL_SG_N]; +#endif + + [[unroll]] for (uint zk = 0; zk < UNROLL_SG_K; zk++) + { + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { + if (elempack == 1) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, gl_CooperativeMatrixLayoutColumnMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, true); +#endif + } + else // if (elempack == 4) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, gl_CooperativeMatrixLayoutRowMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, false); +#endif + } + } + + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, gl_CooperativeMatrixLayoutRowMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, false); +#endif + } + + // sum += k * v + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { +#if ncnn_VK_KHR_cooperative_matrix + sum[zn][zm] = coopMatMulAdd(A[zm], B[zn], sum[zn][zm]); +#elif ncnn_VK_NV_cooperative_matrix + sum[zn][zm] = coopMatMulAddNV(A[zm], B[zn], sum[zn][zm]); +#endif + } + } + } + } + else if (kk >= UNROLL_SG_K) + { + // no ping-pong version + + const uint ki = 0; + + // load bottom_blob + { + if (elempack == 1) + { + // +-M-+ + // K | + // +SG_UM + // | | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UM +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + const uint cstepd4 = psc(cstep) / 4; + + const uint Md4_K_USGM_USGK = Md4 * K * UNROLL_SG_M * UNROLL_SG_K; + const uint Md4_K_USGM_USGK_d_subgroupsize = (Md4_K_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Md4_K_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Md4_K_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Md4_K_USGM_USGK) + { + const uint zk = siq / (Md4 * K * UNROLL_SG_M); + const uint zmij = siq % (Md4 * K * UNROLL_SG_M); + const uint zm = zmij / (Md4 * K); + const uint ij = zmij % (Md4 * K); + const uint i = ij / Md4; + const uint j = ij % Md4; + + const uint gk = ki + zk * K + i; + const uint gm = (mi + zm) * Md4 + j; + + uvec2 v = gk < inch && gm < cstepd4 ? bottom_tm_blob_data[(gz * inch + gk) * cstepd4 + gm] : uvec2(0); + + tmp_v[sgmi][siq] = v; + } + } + } + else // if (elempack == 4) + { + // +-K-+ + // M | + // +- -+ + // SG_UM | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UM +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + const uint inchd4 = inch / 4; + + const uint Kd4_M_USGM_USGK = Kd4 * M * UNROLL_SG_M * UNROLL_SG_K; + const uint Kd4_M_USGM_USGK_d_subgroupsize = (Kd4_M_USGM_USGK + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Kd4_M_USGM_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Kd4_M_USGM_USGK % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Kd4_M_USGM_USGK) + { + const uint zk = siq / (Kd4 * M * UNROLL_SG_M); + const uint zmij = siq % (Kd4 * M * UNROLL_SG_M); + const uint zmi = zmij / Kd4; + const uint j = zmij % Kd4; + + const uint gm = mi * M + zmi; + const uint gk = ki / 4 + zk * Kd4 + j; + + uvec2 v = gk < inchd4 && gm < psc(cstep) ? bottom_tm_blob_data[(gz * inchd4 + gk) * psc(cstep) + gm] : uvec2(0); + + tmp_v[sgmi][siq] = v; + } + } + } + } + + // load weight + { + // +-N-+ + // K | + // +SG_UN + // | | + // ^ +---+ + // | | | + // SG_UK+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UN +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + // weight_tm_data coopmat_N * coopmat_K * UNROLL_SG_N * UNROLL_WG_N * kk, blocks_n + + const uint Nd4_K_USGN = Nd4 * K * UNROLL_SG_N; + const uint Nd4_K_USGN_USGK = Nd4 * K * UNROLL_SG_N * UNROLL_SG_K; + + const uint w_offset = gz * wbstep / 4 + ((wgni * kk) * UNROLL_WG_N) * Nd4_K_USGN + ((k / UNROLL_SG_K) * UNROLL_WG_N + sgni) * Nd4_K_USGN_USGK; + + const uint Nd4_K_USGN_USGK_d_subgroupsize = (Nd4_K_USGN_USGK + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); + [[unroll]] for (uint q = 0; q < Nd4_K_USGN_USGK_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_M + sgmi) * ncnn_subgroupSize + si; + + if (Nd4_K_USGN_USGK % (ncnn_subgroupSize * UNROLL_WG_M) == 0 || siq < Nd4_K_USGN_USGK) + { + tmp_k[sgni][siq] = weight_tm_data[w_offset + siq]; + } + } + } + + barrier(); + +#if ncnn_VK_KHR_cooperative_matrix + coopmat A[UNROLL_SG_M]; + coopmat B[UNROLL_SG_N]; +#elif ncnn_VK_NV_cooperative_matrix + fcoopmatNV<16, gl_ScopeSubgroup, M, K> A[UNROLL_SG_M]; + fcoopmatNV<16, gl_ScopeSubgroup, K, N> B[UNROLL_SG_N]; +#endif + + [[unroll]] for (uint zk = 0; zk < UNROLL_SG_K; zk++) + { + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { + if (elempack == 1) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, gl_CooperativeMatrixLayoutColumnMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Md4 * K), Md4, true); +#endif + } + else // if (elempack == 4) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, gl_CooperativeMatrixLayoutRowMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(A[zm], tmp_v[sgmi], (zk * UNROLL_SG_M + zm) * (Kd4 * M), Kd4, false); +#endif + } + } + + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, gl_CooperativeMatrixLayoutRowMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(B[zn], tmp_k[sgni], (zk * UNROLL_SG_N + zn) * (Nd4 * K), Nd4, false); +#endif + } + + // sum += k * v + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { +#if ncnn_VK_KHR_cooperative_matrix + sum[zn][zm] = coopMatMulAdd(A[zm], B[zn], sum[zn][zm]); +#elif ncnn_VK_NV_cooperative_matrix + sum[zn][zm] = coopMatMulAddNV(A[zm], B[zn], sum[zn][zm]); +#endif + } + } + } + + k += UNROLL_SG_K; + } + + for (; k < kk; k++) + { + const uint ki = k * K; + + barrier(); + + // load bottom_tm_blob + { + if (elempack == 1) + { + // +-M-+ + // K | + // +SG_UM + // | | + // ^ +---+ + // | | | + // WG_UM+- -+ + // | | | + // v +---+ + + const uint cstepd4 = psc(cstep) / 4; + + const uint Md4_K_USGM = Md4 * K * UNROLL_SG_M; + const uint Md4_K_USGM_d_subgroupsize = (Md4_K_USGM + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Md4_K_USGM_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Md4_K_USGM % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Md4_K_USGM) + { + const uint zm = siq / (Md4 * K); + const uint ij = siq % (Md4 * K); + const uint i = ij / Md4; + const uint j = ij % Md4; + + const uint gk = ki + i; + const uint gm = (mi + zm) * Md4 + j; + + uvec2 v = gk < inch && gm < cstepd4 ? bottom_tm_blob_data[(gz * inch + gk) * cstepd4 + gm] : uvec2(0); + + tmp_v[sgmi][siq] = v; + } + } + } + else // if (elempack == 4) + { + // +-K-+ + // M | + // +- -+ + // SG_UM | + // ^ +---+ + // | | | + // WG_UM+- -+ + // | | | + // v +---+ + + const uint inchd4 = inch / 4; + + const uint Kd4_M_USGM = Kd4 * M * UNROLL_SG_M; + const uint Kd4_M_USGM_d_subgroupsize = (Kd4_M_USGM + (ncnn_subgroupSize * UNROLL_WG_N - 1)) / (ncnn_subgroupSize * UNROLL_WG_N); + [[unroll]] for (uint q = 0; q < Kd4_M_USGM_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_N + sgni) * ncnn_subgroupSize + si; + + if (Kd4_M_USGM % (ncnn_subgroupSize * UNROLL_WG_N) == 0 || siq < Kd4_M_USGM) + { + const uint zmi = siq / Kd4; + const uint j = siq % Kd4; + + const uint gm = mi * M + zmi; + const uint gk = ki / 4 + j; + + uvec2 v = gk < inchd4 && gm < psc(size) ? bottom_tm_blob_data[(gz * inchd4 + gk) * psc(cstep) + gm] : uvec2(0); + + tmp_v[sgmi][siq] = v; + } + } + } + } + + // load weight + { + // +-N-+ + // K | + // +SG_UN + // | | + // ^ +---+ + // | | | + // WG_UN+- -+ + // | | | + // v +---+ + + // weight_tm_data coopmat_N * coopmat_K * UNROLL_SG_N * UNROLL_WG_N * kk, blocks_n + + const uint Nd4_K_USGN = Nd4 * K * UNROLL_SG_N; + + const uint w_offset = gz * wbstep / 4 + ((wgni * kk + k) * UNROLL_WG_N + sgni) * Nd4_K_USGN; + + const uint Nd4_K_USGN_d_subgroupsize = (Nd4_K_USGN + (ncnn_subgroupSize * UNROLL_WG_M - 1)) / (ncnn_subgroupSize * UNROLL_WG_M); + [[unroll]] for (uint q = 0; q < Nd4_K_USGN_d_subgroupsize; q++) + { + const uint siq = (q * UNROLL_WG_M + sgmi) * ncnn_subgroupSize + si; + + if (Nd4_K_USGN % (ncnn_subgroupSize * UNROLL_WG_M) == 0 || siq < Nd4_K_USGN) + { + tmp_k[sgni][siq] = weight_tm_data[w_offset + siq]; + } + } + } + + barrier(); + +#if ncnn_VK_KHR_cooperative_matrix + coopmat A[UNROLL_SG_M]; + coopmat B[UNROLL_SG_N]; +#elif ncnn_VK_NV_cooperative_matrix + fcoopmatNV<16, gl_ScopeSubgroup, M, K> A[UNROLL_SG_M]; + fcoopmatNV<16, gl_ScopeSubgroup, K, N> B[UNROLL_SG_N]; +#endif + + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { + if (elempack == 1) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(A[zm], tmp_v[sgmi], zm * (Md4 * K), Md4, gl_CooperativeMatrixLayoutColumnMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(A[zm], tmp_v[sgmi], zm * (Md4 * K), Md4, true); +#endif + } + else // if (elempack == 4) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(A[zm], tmp_v[sgmi], zm * (Kd4 * M), Kd4, gl_CooperativeMatrixLayoutRowMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(A[zm], tmp_v[sgmi], zm * (Kd4 * M), Kd4, false); +#endif + } + } + + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { +#if ncnn_VK_KHR_cooperative_matrix + coopMatLoad(B[zn], tmp_k[sgni], zn * (Nd4 * K), Nd4, gl_CooperativeMatrixLayoutRowMajor); +#elif ncnn_VK_NV_cooperative_matrix + coopMatLoadNV(B[zn], tmp_k[sgni], zn * (Nd4 * K), Nd4, false); +#endif + } + + // sum += k * v + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { +#if ncnn_VK_KHR_cooperative_matrix + sum[zn][zm] = coopMatMulAdd(A[zm], B[zn], sum[zn][zm]); +#elif ncnn_VK_NV_cooperative_matrix + sum[zn][zm] = coopMatMulAddNV(A[zm], B[zn], sum[zn][zm]); +#endif + } + } + } + + [[unroll]] for (uint zn = 0; zn < UNROLL_SG_N; zn++) + { + [[unroll]] for (uint zm = 0; zm < UNROLL_SG_M; zm++) + { + if (out_elempack == 1) + { +#if ncnn_VK_KHR_cooperative_matrix +#if NCNN_fp16_arithmetic + coopMatStore(sum[zn][zm], tmp_o[sgi], (zn * UNROLL_SG_M + zm) * (Md4 * N), Md4, gl_CooperativeMatrixLayoutColumnMajor); +#else + coopmat sum_fp16 = coopmat(sum[zn][zm]); + coopMatStore(sum_fp16, tmp_o[sgi], (zn * UNROLL_SG_M + zm) * (Md4 * N), Md4, gl_CooperativeMatrixLayoutColumnMajor); +#endif +#elif ncnn_VK_NV_cooperative_matrix +#if NCNN_fp16_arithmetic + coopMatStoreNV(sum[zn][zm], tmp_o[sgi], (zn * UNROLL_SG_M + zm) * (Md4 * N), Md4, true); +#else + fcoopmatNV<16, gl_ScopeSubgroup, M, N> sum_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, M, N>(sum[zn][zm]); + coopMatStoreNV(sum_fp16, tmp_o[sgi], (zn * UNROLL_SG_M + zm) * (Md4 * N), Md4, true); +#endif +#endif + } + else // if (out_elempack == 4) + { +#if ncnn_VK_KHR_cooperative_matrix +#if NCNN_fp16_arithmetic + coopMatStore(sum[zn][zm], tmp_o[sgi], (zn * UNROLL_SG_M + zm) * (Nd4 * M), Nd4, gl_CooperativeMatrixLayoutRowMajor); +#else + coopmat sum_fp16 = coopmat(sum[zn][zm]); + coopMatStore(sum_fp16, tmp_o[sgi], (zn * UNROLL_SG_M + zm) * (Nd4 * M), Nd4, gl_CooperativeMatrixLayoutRowMajor); +#endif +#elif ncnn_VK_NV_cooperative_matrix +#if NCNN_fp16_arithmetic + coopMatStoreNV(sum[zn][zm], tmp_o[sgi], (zn * UNROLL_SG_M + zm) * (Nd4 * M), Nd4, false); +#else + fcoopmatNV<16, gl_ScopeSubgroup, M, N> sum_fp16 = fcoopmatNV<16, gl_ScopeSubgroup, M, N>(sum[zn][zm]); + coopMatStoreNV(sum_fp16, tmp_o[sgi], (zn * UNROLL_SG_M + zm) * (Nd4 * M), Nd4, false); +#endif +#endif + } + } + } + + barrier(); + + // store top_tm_blob + { + if (out_elempack == 1) + { + // +-M-+ + // N | + // +SG_UM + // | | + // ^ +---+ + // | | | + // SG_UN+- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UM +- -+ + // | | | + // | +- -+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // | +---+ + // | | | + // | +- -+ + // | | | + // WG_UN +---+ + // | | | + // | +- -+ + // | | | + // | +---+ + // | | | + // | +- -+ + // | | | + // v +---+ + + const uint outcstepd4 = psc(outcstep) / 4; + + const uint Md4_N_USGM_USGN = Md4 * N * UNROLL_SG_M * UNROLL_SG_N; + const uint Md4_N_USGM_USGN_d_subgroupsize = (Md4_N_USGM_USGN + ncnn_subgroupSize - 1) / ncnn_subgroupSize; + [[unroll]] for (uint q = 0; q < Md4_N_USGM_USGN_d_subgroupsize; q++) + { + const uint siq = si + q * ncnn_subgroupSize; + + if (Md4_N_USGM_USGN % ncnn_subgroupSize == 0 || siq < Md4_N_USGM_USGN) + { + const uint zn = siq / (Md4 * N * UNROLL_SG_M); + const uint zmij = siq % (Md4 * N * UNROLL_SG_M); + const uint zm = zmij / (Md4 * N); + const uint ij = zmij % (Md4 * N); + const uint i = ij / Md4; + const uint j = ij % Md4; + + const uint gn = (ni + zn) * N + i; + const uint gm = (mi + zm) * Md4 + j; + + if (gn < outch && gm < outcstepd4) + { + top_tm_blob_data[(gz * outch + gn) * outcstepd4 + gm] = tmp_o[sgi][siq]; + } + } + } + } + else // if (out_elempack == 4) + { + // +-N-+ + // M | + // +---+ + // SG_UM | + // ^ +---+ + // | | | + // SG_UN+---+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // WG_UM +---+ + // | | | + // | +---+ + // | | | + // ^ v +---+ + // | | | + // | +- -+ + // | | | + // | +---+ + // | | | + // | +---+ + // | | | + // WG_UN +---+ + // | | | + // | +- -+ + // | | | + // | +---+ + // | | | + // | +---+ + // | | | + // v +---+ + + const uint outchd4 = outch / 4; + + const uint Nd4_M_USGM_USGN = Nd4 * M * UNROLL_SG_M * UNROLL_SG_N; + const uint Nd4_M_USGM_USGN_d_subgroupsize = (Nd4_M_USGM_USGN + ncnn_subgroupSize - 1) / ncnn_subgroupSize; + [[unroll]] for (uint q = 0; q < Nd4_M_USGM_USGN_d_subgroupsize; q++) + { + const uint siq = si + q * ncnn_subgroupSize; + + if (Nd4_M_USGM_USGN % ncnn_subgroupSize == 0 || siq < Nd4_M_USGM_USGN) + { + const uint zn = siq / (Nd4 * M * UNROLL_SG_M); + const uint zmij = siq % (Nd4 * M * UNROLL_SG_M); + const uint zmi = zmij / Nd4; + const uint j = zmij % Nd4; + + const uint gn = (ni + zn) * Nd4 + j; + const uint gm = mi * M + zmi; + + if (gn < outchd4 && gm < psc(size)) + { + top_tm_blob_data[(gz * outchd4 + gn) * psc(outcstep) + gm] = tmp_o[sgi][siq]; + } + } + } + } + } +}