| 171 | } |
| 172 | |
| 173 | std::vector<ComputeTaskDescriptorPtr> Reshapex4(int id, ValueId input_id, |
| 174 | ValueId output_id, |
| 175 | const ReshapeAttributes& attr) { |
| 176 | auto desc = std::make_shared<ComputeTaskDescriptor>(); |
| 177 | desc->id = id; |
| 178 | desc->is_linkable = false; |
| 179 | desc->shader_source = GetReshapex4Code(); |
| 180 | |
| 181 | desc->input_buffers = { |
| 182 | {input_id, "device FLT4* const src_buffer"}, |
| 183 | }; |
| 184 | |
| 185 | desc->output_buffer = { |
| 186 | output_id, "device FLT4* dst_buffer", |
| 187 | [input_id, attr](const std::map<ValueId, BHWC>& buffers) { |
| 188 | int batch = buffers.find(input_id)->second.b; |
| 189 | return BHWC{batch, attr.new_shape.h, attr.new_shape.w, |
| 190 | attr.new_shape.c}; |
| 191 | }}; |
| 192 | |
| 193 | desc->uniform_buffers = { |
| 194 | {"constant uniforms& params", |
| 195 | [input_id, output_id](const std::map<ValueId, BHWC>& buffers) { |
| 196 | const auto& src_dim = buffers.find(input_id)->second; |
| 197 | const auto& dst_dim = buffers.find(output_id)->second; |
| 198 | std::vector<int32_t> uniform_params{ |
| 199 | // int4 src_size |
| 200 | src_dim.w, src_dim.h, IntegralDivideRoundUp(src_dim.c, 4), |
| 201 | src_dim.w * src_dim.h, |
| 202 | // int4 dst_size |
| 203 | dst_dim.w, dst_dim.h, IntegralDivideRoundUp(dst_dim.c, 4), |
| 204 | dst_dim.w * dst_dim.h, |
| 205 | // int2 plane_xz |
| 206 | src_dim.w * IntegralDivideRoundUp(src_dim.c, 4), |
| 207 | dst_dim.w * IntegralDivideRoundUp(dst_dim.c, 4), |
| 208 | 0, // dummy, for alignment |
| 209 | 0, // dummy, for alignment |
| 210 | 0, // dummy, for alignment |
| 211 | 0, // dummy, for alignment |
| 212 | 0, // dummy, for alignment |
| 213 | 0 // dummy, for alignment |
| 214 | }; |
| 215 | return VectorToUint8Vector(uniform_params); |
| 216 | }}, |
| 217 | }; |
| 218 | |
| 219 | desc->resize_function = [attr](const std::map<ValueId, BHWC>& buffers) { |
| 220 | const uint3 grid = uint3(attr.new_shape.w, attr.new_shape.h, |
| 221 | IntegralDivideRoundUp(attr.new_shape.c, 4)); |
| 222 | const uint3 groups_size = GetWorkGroupSizeForGrid(grid); |
| 223 | int groups_x = IntegralDivideRoundUp(grid.x, groups_size.x); |
| 224 | int groups_y = IntegralDivideRoundUp(grid.y, groups_size.y); |
| 225 | int groups_z = IntegralDivideRoundUp(grid.z, groups_size.z); |
| 226 | return std::make_pair(groups_size, uint3{groups_x, groups_y, groups_z}); |
| 227 | }; |
| 228 | |
| 229 | return {desc}; |
| 230 | } |
no test coverage detected