| 167 | } |
| 168 | |
| 169 | std::vector<ComputeTaskDescriptorPtr> ConcatX( |
| 170 | int id, std::vector<ValueId> input_ids, ValueId output_id, |
| 171 | const ConcatAttributes& attr, const std::vector<BHWC>& input_shapes) { |
| 172 | auto desc = std::make_shared<ComputeTaskDescriptor>(); |
| 173 | desc->id = id; |
| 174 | desc->is_linkable = false; |
| 175 | |
| 176 | std::string code = R"( |
| 177 | #include <metal_stdlib> |
| 178 | using namespace metal; |
| 179 | $0 |
| 180 | kernel void ComputeFunction( |
| 181 | $1 |
| 182 | uint3 gid[[thread_position_in_grid]]) { |
| 183 | if (int(gid.x) >= size.x || int(gid.y) >= size.y) { |
| 184 | return; |
| 185 | } |
| 186 | FLT4 value; |
| 187 | )"; |
| 188 | int output_width = 0; |
| 189 | for (int buffer_index = 0; buffer_index < input_shapes.size(); |
| 190 | buffer_index++) { |
| 191 | const auto& dims = input_shapes[buffer_index]; |
| 192 | output_width += dims.w; |
| 193 | |
| 194 | // Generated shader example: |
| 195 | // if (gid.x < 10) value = src_buffer0[(gid.y + gid.z * 3) * 4 + gid.x - 3]; |
| 196 | // else |
| 197 | if (buffer_index < input_shapes.size() - 1) { |
| 198 | code += "if (gid.x < " + std::to_string(output_width) + ")"; |
| 199 | } |
| 200 | code += "value = src_buffer" + std::to_string(buffer_index) + |
| 201 | "[(gid.y + gid.z * " + std::to_string(dims.h) + ") * " + |
| 202 | std::to_string(dims.w) + " + gid.x - " + |
| 203 | std::to_string(output_width - dims.w) + "];\n"; |
| 204 | if (buffer_index < input_shapes.size() - 1) { |
| 205 | code += "else "; |
| 206 | } |
| 207 | } |
| 208 | code += "const int linear_index = (gid.y + gid.z * " + |
| 209 | std::to_string(input_shapes[0].h) + ") * " + |
| 210 | std::to_string(output_width) + " + gid.x;"; |
| 211 | code += R"( |
| 212 | $2 |
| 213 | dst_buffer[linear_index] = value; |
| 214 | } |
| 215 | )"; |
| 216 | desc->shader_source = code; |
| 217 | |
| 218 | for (int i = 0; i < input_ids.size(); ++i) { |
| 219 | const std::string buffer_name = |
| 220 | "device FLT4* const src_buffer" + std::to_string(i); |
| 221 | desc->input_buffers.push_back({input_ids[i], buffer_name}); |
| 222 | } |
| 223 | |
| 224 | desc->output_buffer = { |
| 225 | output_id, "device FLT4* dst_buffer", |
| 226 | [input_ids, attr](const std::map<ValueId, BHWC>& buffers) { |
no test coverage detected