1// Tencent is pleased to support the open source community by making ncnn available. 2// 3// Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved. 4// 5// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except 6// in compliance with the License. You may obtain a copy of the License at 7// 8// https://opensource.org/licenses/BSD-3-Clause 9// 10// Unless required by applicable law or agreed to in writing, software distributed 11// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR 12// CONDITIONS OF ANY KIND, either express or implied. See the License for the 13// specific language governing permissions and limitations under the License. 14 15#version 450 16 17#if NCNN_fp16_storage 18#extension GL_EXT_shader_16bit_storage: require 19#endif 20#if NCNN_fp16_arithmetic 21#extension GL_EXT_shader_explicit_arithmetic_types_float16: require 22#endif 23 24layout (constant_id = 0) const int region_type = 0; 25layout (constant_id = 1) const int pad_head = 0; 26layout (constant_id = 2) const int pad_tail = 0; 27 28#define shape_constant_id_offset 3 29layout (constant_id = shape_constant_id_offset + 0) const int dims = 0; 30layout (constant_id = shape_constant_id_offset + 1) const int w = 0; 31layout (constant_id = shape_constant_id_offset + 2) const int h = 0; 32layout (constant_id = shape_constant_id_offset + 3) const int c = 0; 33layout (constant_id = shape_constant_id_offset + 4) const int cstep = 0; 34 35layout (constant_id = shape_constant_id_offset + 5) const int outdims = 0; 36layout (constant_id = shape_constant_id_offset + 6) const int outw = 0; 37layout (constant_id = shape_constant_id_offset + 7) const int outh = 0; 38layout (constant_id = shape_constant_id_offset + 8) const int outc = 0; 39layout (constant_id = shape_constant_id_offset + 9) const int outcstep = 0; 40 41#if NCNN_image_shader 42layout (binding = 0) uniform unfp sampler3D bottom_blob; 43layout (binding = 1, r32f) writeonly uniform highp image3D square_workspace; 44#else 45layout (binding = 0) readonly buffer bottom_blob { sfp bottom_blob_data[]; }; 46layout (binding = 1) writeonly buffer square_workspace { float square_workspace_data[]; }; 47#endif 48 49layout (push_constant) uniform parameter 50{ 51 int dims; 52 int w; 53 int h; 54 int c; 55 int cstep; 56 57 int outdims; 58 int outw; 59 int outh; 60 int outc; 61 int outcstep; 62} p; 63 64void main() 65{ 66 int gx = int(gl_GlobalInvocationID.x); 67 int gy = int(gl_GlobalInvocationID.y); 68 int gz = int(gl_GlobalInvocationID.z); 69 70 if (gx >= psc(outw) || gy >= psc(outh) || gz >= psc(outc)) 71 return; 72 73 afp res; 74 75 if (region_type == 0) 76 { 77 int z = gz - pad_head; 78 79 if (z >= 0 && z < psc(c)) 80 { 81#if NCNN_image_shader 82 afp v = image3d_ld1(bottom_blob, ivec3(gx, gy, z)); 83#else 84 int v_offset = z * psc(cstep) + gy * psc(w) + gx; 85 afp v = buffer_ld1(bottom_blob_data, v_offset); 86#endif 87 res = v * v; 88 } 89 else 90 { 91 res = afp(0.f); 92 } 93 } 94 else 95 { 96 int x = gx - pad_head; 97 int y = gy - pad_head; 98 99 if (x >= 0 && x < psc(w) && y >= 0 && y < psc(h)) 100 { 101#if NCNN_image_shader 102 afp v = image3d_ld1(bottom_blob, ivec3(x, y, gz)); 103#else 104 int v_offset = gz * psc(cstep) + y * psc(w) + x; 105 afp v = buffer_ld1(bottom_blob_data, v_offset); 106#endif 107 res = v * v; 108 } 109 else 110 { 111 res = afp(0.f); 112 } 113 } 114 115#if NCNN_image_shader 116 imageStore(square_workspace, ivec3(gx, gy, gz), vec4(float(res))); 117#else 118 const int gi = gz * psc(outcstep) + gy * psc(outw) + gx; 119 120 square_workspace_data[gi] = float(res); 121#endif 122} 123