| 117 | } // namespace |
| 118 | |
| 119 | std::vector<ComputeTaskDescriptorPtr> Reshape(int id, ValueId input_id, |
| 120 | ValueId output_id, |
| 121 | const ReshapeAttributes& attr) { |
| 122 | auto desc = std::make_shared<ComputeTaskDescriptor>(); |
| 123 | desc->id = id; |
| 124 | desc->is_linkable = false; |
| 125 | desc->shader_source = GetReshapeCode(); |
| 126 | |
| 127 | desc->input_buffers = { |
| 128 | {input_id, "device FLT4* const src_buffer"}, |
| 129 | }; |
| 130 | |
| 131 | desc->output_buffer = { |
| 132 | output_id, "device FLT4* dst_buffer", |
| 133 | [input_id, attr](const std::map<ValueId, BHWC>& buffers) { |
| 134 | int batch = buffers.find(input_id)->second.b; |
| 135 | return BHWC{batch, attr.new_shape.h, attr.new_shape.w, |
| 136 | attr.new_shape.c}; |
| 137 | }}; |
| 138 | |
| 139 | desc->uniform_buffers = { |
| 140 | {"constant uniforms& params", |
| 141 | [input_id, output_id](const std::map<ValueId, BHWC>& buffers) { |
| 142 | const auto& src_dim = buffers.find(input_id)->second; |
| 143 | const auto& dst_dim = buffers.find(output_id)->second; |
| 144 | std::vector<int> uniform_params{ |
| 145 | // int4 src_size |
| 146 | src_dim.w, |
| 147 | src_dim.h, |
| 148 | src_dim.c, |
| 149 | src_dim.c * src_dim.w, |
| 150 | // int4 dst_size |
| 151 | dst_dim.w, |
| 152 | dst_dim.h, |
| 153 | dst_dim.c, |
| 154 | dst_dim.c * dst_dim.w, |
| 155 | }; |
| 156 | return VectorToUint8Vector(uniform_params); |
| 157 | }}, |
| 158 | }; |
| 159 | |
| 160 | desc->resize_function = [attr](const std::map<ValueId, BHWC>& buffers) { |
| 161 | const uint3 grid = uint3(attr.new_shape.w, attr.new_shape.h, |
| 162 | IntegralDivideRoundUp(attr.new_shape.c, 4)); |
| 163 | const uint3 groups_size = GetWorkGroupSizeForGrid(grid); |
| 164 | int groups_x = IntegralDivideRoundUp(grid.x, groups_size.x); |
| 165 | int groups_y = IntegralDivideRoundUp(grid.y, groups_size.y); |
| 166 | int groups_z = IntegralDivideRoundUp(grid.z, groups_size.z); |
| 167 | return std::make_pair(groups_size, uint3{groups_x, groups_y, groups_z}); |
| 168 | }; |
| 169 | |
| 170 | return {desc}; |
| 171 | } |
| 172 | |
| 173 | std::vector<ComputeTaskDescriptorPtr> Reshapex4(int id, ValueId input_id, |
| 174 | ValueId output_id, |
no test coverage detected