ncnn

1
// yala is pleased to support the open source community by making ncnn available.
2
//
3
//
4
// Copyright (C) 2022 yala <zhaojunchao@loongson.cn>;<junchao82@qq.com>. All rights reserved.
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
#include "slice_loongarch.h"
16

17
namespace ncnn {
18

19
Slice_loongarch::Slice_loongarch()
20
{
21
#if __loongarch_sx
22
    support_packing = true;
23
#endif // __loongarch_sx
24
}
25

26
int Slice_loongarch::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_blobs, const Option& opt) const
27
{
28
    const Mat& bottom_blob = bottom_blobs[0];
29
    int dims = bottom_blob.dims;
30
    size_t elemsize = bottom_blob.elemsize;
31
    int elempack = bottom_blob.elempack;
32
    const int* slices_ptr = slices;
33
    const int* indices_ptr = indices;
34
    int positive_axis = axis < 0 ? dims + axis : axis;
35

36
    if (dims == 1) // positive_axis == 0
37
    {
38
        // slice vector
39
        int w = bottom_blob.w * elempack;
40
        int q = 0;
41
        for (size_t i = 0; i < top_blobs.size(); i++)
42
        {
43
            int slice;
44
            if (indices_ptr)
45
            {
46
                if (i == top_blobs.size() - 1)
47
                {
48
                    slice = w - q;
49
                }
50
                else
51
                {
52
                    int indice = indices_ptr[i];
53
                    int positive_indice = indice < 0 ? w + indice : indice;
54
                    slice = positive_indice - q;
55
                }
56
            }
57
            else
58
            {
59
                slice = slices_ptr[i];
60
                if (slice == -233)
61
                {
62
                    slice = static_cast<int>((w - q) / (top_blobs.size() - i));
63
                }
64
            }
65

66
            int out_elempack = 1;
67
#if __loongarch_sx
68
            if (opt.use_packing_layout)
69
                out_elempack = slice % 4 == 0 ? 4 : 1;
70
#endif
71
            size_t out_elemsize = elemsize / elempack * out_elempack;
72

73
            Mat& top_blob = top_blobs[i];
74
            top_blob.create(slice / out_elempack, out_elemsize, out_elempack, opt.blob_allocator);
75
            if (top_blob.empty())
76
                return -100;
77

78
            const float* ptr = (const float*)bottom_blob + q;
79
            float* outptr = top_blob;
80
            memcpy(outptr, ptr, top_blob.w * top_blob.elemsize);
81

82
            q += slice;
83
        }
84
    }
85

86
    if (dims == 2 && positive_axis == 0)
87
    {
88
        // slice image height
89
        int w = bottom_blob.w;
90
        int h = bottom_blob.h * elempack;
91

92
        int q = 0;
93
        for (size_t i = 0; i < top_blobs.size(); i++)
94
        {
95
            int slice;
96
            if (indices_ptr)
97
            {
98
                if (i == top_blobs.size() - 1)
99
                {
100
                    slice = h - q;
101
                }
102
                else
103
                {
104
                    int indice = indices_ptr[i];
105
                    int positive_indice = indice < 0 ? h + indice : indice;
106
                    slice = positive_indice - q;
107
                }
108
            }
109
            else
110
            {
111
                slice = slices_ptr[i];
112
                if (slice == -233)
113
                {
114
                    slice = static_cast<int>((h - q) / (top_blobs.size() - i));
115
                }
116
            }
117

118
            int out_elempack = 1;
119
#if __loongarch_sx
120
            if (opt.use_packing_layout)
121
                out_elempack = slice % 4 == 0 ? 4 : 1;
122
#endif
123
            size_t out_elemsize = elemsize / elempack * out_elempack;
124

125
            Mat& top_blob = top_blobs[i];
126
            top_blob.create(w, slice / out_elempack, out_elemsize, out_elempack, opt.blob_allocator);
127
            if (top_blob.empty())
128
                return -100;
129

130
            q += slice;
131
        }
132

133
        size_t out_elemsize = top_blobs[0].elemsize;
134
        int out_elempack = top_blobs[0].elempack;
135
        for (size_t i = 0; i < top_blobs.size(); i++)
136
        {
137
            out_elemsize = std::min(out_elemsize, top_blobs[i].elemsize);
138
            out_elempack = std::min(out_elempack, top_blobs[i].elempack);
139
        }
140

141
        Mat bottom_blob_unpacked = bottom_blob;
142
        if (elempack > out_elempack)
143
        {
144
            convert_packing(bottom_blob, bottom_blob_unpacked, out_elempack, opt);
145
            if (bottom_blob_unpacked.empty())
146
                return -100;
147
        }
148

149
        const float* ptr = bottom_blob_unpacked;
150
        for (size_t i = 0; i < top_blobs.size(); i++)
151
        {
152
            Mat& top_blob = top_blobs[i];
153

154
            if (out_elempack == 1 && top_blob.elempack == 4)
155
            {
156
                for (int j = 0; j < top_blob.h; j++)
157
                {
158
                    const float* r0 = ptr;
159
                    const float* r1 = ptr + w;
160
                    const float* r2 = ptr + w * 2;
161
                    const float* r3 = ptr + w * 3;
162

163
                    float* outptr0 = top_blob.row(j);
164

165
                    for (int j = 0; j < w; j++)
166
                    {
167
                        outptr0[0] = *r0++;
168
                        outptr0[1] = *r1++;
169
                        outptr0[2] = *r2++;
170
                        outptr0[3] = *r3++;
171

172
                        outptr0 += 4;
173
                    }
174

175
                    ptr += w * 4;
176
                }
177
            }
178
            else // if (out_elempack == 1 && top_blob.elempack == 1) if (out_elempack == 4 && top_blob.elempack == 4)
179
            {
180
                int size = w * top_blob.h;
181

182
                float* outptr = top_blob;
183
                memcpy(outptr, ptr, size * top_blob.elemsize);
184

185
                ptr += size * top_blob.elempack;
186
            }
187
        }
188
    }
189

190
    if (dims == 2 && positive_axis == 1)
191
    {
192
        // slice image width
193
        int w = bottom_blob.w;
194
        int h = bottom_blob.h;
195

196
        int q = 0;
197
        for (size_t i = 0; i < top_blobs.size(); i++)
198
        {
199
            int slice;
200
            if (indices_ptr)
201
            {
202
                if (i == top_blobs.size() - 1)
203
                {
204
                    slice = w - q;
205
                }
206
                else
207
                {
208
                    int indice = indices_ptr[i];
209
                    int positive_indice = indice < 0 ? w + indice : indice;
210
                    slice = positive_indice - q;
211
                }
212
            }
213
            else
214
            {
215
                slice = slices_ptr[i];
216
                if (slice == -233)
217
                {
218
                    slice = static_cast<int>((w - q) / (top_blobs.size() - i));
219
                }
220
            }
221

222
            Mat& top_blob = top_blobs[i];
223
            top_blob.create(slice, h, elemsize, elempack, opt.blob_allocator);
224
            if (top_blob.empty())
225
                return -100;
226

227
            q += slice;
228
        }
229

230
        #pragma omp parallel for num_threads(opt.num_threads)
231
        for (int j = 0; j < h; j++)
232
        {
233
            const float* ptr = bottom_blob.row(j);
234
            for (size_t i = 0; i < top_blobs.size(); i++)
235
            {
236
                Mat& top_blob = top_blobs[i];
237

238
                float* outptr = top_blob.row(j);
239
                memcpy(outptr, ptr, top_blob.w * elemsize);
240

241
                ptr += top_blob.w * elempack;
242
            }
243
        }
244
    }
245

246
    if ((dims == 3 || dims == 4) && positive_axis == 0)
247
    {
248
        // slice dim channel
249
        int w = bottom_blob.w;
250
        int h = bottom_blob.h;
251
        int d = bottom_blob.d;
252
        int channels = bottom_blob.c * elempack;
253

254
        int q = 0;
255
        for (size_t i = 0; i < top_blobs.size(); i++)
256
        {
257
            int slice;
258
            if (indices_ptr)
259
            {
260
                if (i == top_blobs.size() - 1)
261
                {
262
                    slice = channels - q;
263
                }
264
                else
265
                {
266
                    int indice = indices_ptr[i];
267
                    int positive_indice = indice < 0 ? channels + indice : indice;
268
                    slice = positive_indice - q;
269
                }
270
            }
271
            else
272
            {
273
                slice = slices_ptr[i];
274
                if (slice == -233)
275
                {
276
                    slice = static_cast<int>((channels - q) / (top_blobs.size() - i));
277
                }
278
            }
279

280
            int out_elempack = 1;
281
#if __loongarch_sx
282
            if (opt.use_packing_layout)
283
                out_elempack = slice % 4 == 0 ? 4 : 1;
284
#endif
285
            size_t out_elemsize = elemsize / elempack * out_elempack;
286

287
            Mat& top_blob = top_blobs[i];
288
            top_blob.create(w, h, d, slice / out_elempack, out_elemsize, out_elempack, opt.blob_allocator);
289
            if (top_blob.empty())
290
                return -100;
291

292
            top_blob.dims = dims;
293

294
            q += slice;
295
        }
296

297
        size_t out_elemsize = top_blobs[0].elemsize;
298
        int out_elempack = top_blobs[0].elempack;
299
        for (size_t i = 0; i < top_blobs.size(); i++)
300
        {
301
            out_elemsize = std::min(out_elemsize, top_blobs[i].elemsize);
302
            out_elempack = std::min(out_elempack, top_blobs[i].elempack);
303
        }
304

305
        Mat bottom_blob_unpacked = bottom_blob;
306
        if (elempack > out_elempack)
307
        {
308
            convert_packing(bottom_blob, bottom_blob_unpacked, out_elempack, opt);
309
            if (bottom_blob_unpacked.empty())
310
                return -100;
311
        }
312

313
        int p = 0;
314
        for (size_t i = 0; i < top_blobs.size(); i++)
315
        {
316
            Mat& top_blob = top_blobs[i];
317

318
            if (out_elempack == 1 && top_blob.elempack == 4)
319
            {
320
                int size = top_blob.w * top_blob.h * top_blob.d;
321

322
                for (int q = 0; q < top_blob.c; q++)
323
                {
324
                    const float* r0 = bottom_blob_unpacked.channel(p);
325
                    const float* r1 = bottom_blob_unpacked.channel(p + 1);
326
                    const float* r2 = bottom_blob_unpacked.channel(p + 2);
327
                    const float* r3 = bottom_blob_unpacked.channel(p + 3);
328

329
                    float* outptr0 = top_blob.channel(q);
330

331
                    for (int j = 0; j < size; j++)
332
                    {
333
                        outptr0[0] = *r0++;
334
                        outptr0[1] = *r1++;
335
                        outptr0[2] = *r2++;
336
                        outptr0[3] = *r3++;
337

338
                        outptr0 += 4;
339
                    }
340

341
                    p += 4;
342
                }
343
            }
344
            else // if (out_elempack == 1 && top_blob.elempack == 1) if (out_elempack == 4 && top_blob.elempack == 4)
345
            {
346
                int size = top_blob.total();
347

348
                const float* ptr = bottom_blob_unpacked.channel(p);
349
                float* outptr = top_blob;
350
                memcpy(outptr, ptr, size * top_blob.elemsize);
351

352
                p += top_blob.c;
353
            }
354
        }
355
    }
356

357
    if ((dims == 3 && positive_axis == 1) || (dims == 4 && positive_axis == 2))
358
    {
359
        // slice dim height
360
        int w = bottom_blob.w;
361
        int h = bottom_blob.h;
362
        int d = bottom_blob.d;
363
        int channels = bottom_blob.c;
364

365
        int q = 0;
366
        for (size_t i = 0; i < top_blobs.size(); i++)
367
        {
368
            int slice;
369
            if (indices_ptr)
370
            {
371
                if (i == top_blobs.size() - 1)
372
                {
373
                    slice = h - q;
374
                }
375
                else
376
                {
377
                    int indice = indices_ptr[i];
378
                    int positive_indice = indice < 0 ? h + indice : indice;
379
                    slice = positive_indice - q;
380
                }
381
            }
382
            else
383
            {
384
                slice = slices_ptr[i];
385
                if (slice == -233)
386
                {
387
                    slice = static_cast<int>((h - q) / (top_blobs.size() - i));
388
                }
389
            }
390

391
            Mat& top_blob = top_blobs[i];
392
            top_blob.create(w, slice, d, channels, elemsize, elempack, opt.blob_allocator);
393
            if (top_blob.empty())
394
                return -100;
395

396
            top_blob.dims = dims;
397

398
            q += slice;
399
        }
400

401
        #pragma omp parallel for num_threads(opt.num_threads)
402
        for (int p = 0; p < channels; p++)
403
        {
404
            const float* ptr = bottom_blob.channel(p);
405

406
            for (int j = 0; j < d; j++)
407
            {
408
                for (size_t i = 0; i < top_blobs.size(); i++)
409
                {
410
                    Mat& top_blob = top_blobs[i];
411

412
                    int size = top_blob.w * top_blob.h;
413

414
                    float* outptr = top_blob.channel(p).depth(j);
415
                    memcpy(outptr, ptr, size * elemsize);
416

417
                    ptr += size * elempack;
418
                }
419
            }
420
        }
421
    }
422

423
    if ((dims == 3 && positive_axis == 2) || (dims == 4 && positive_axis == 3))
424
    {
425
        // slice dim width
426
        int w = bottom_blob.w;
427
        int h = bottom_blob.h;
428
        int d = bottom_blob.d;
429
        int channels = bottom_blob.c;
430

431
        int q = 0;
432
        for (size_t i = 0; i < top_blobs.size(); i++)
433
        {
434
            int slice;
435
            if (indices_ptr)
436
            {
437
                if (i == top_blobs.size() - 1)
438
                {
439
                    slice = w - q;
440
                }
441
                else
442
                {
443
                    int indice = indices_ptr[i];
444
                    int positive_indice = indice < 0 ? w + indice : indice;
445
                    slice = positive_indice - q;
446
                }
447
            }
448
            else
449
            {
450
                slice = slices_ptr[i];
451
                if (slice == -233)
452
                {
453
                    slice = static_cast<int>((w - q) / (top_blobs.size() - i));
454
                }
455
            }
456

457
            Mat& top_blob = top_blobs[i];
458
            top_blob.create(slice, h, d, channels, elemsize, elempack, opt.blob_allocator);
459
            if (top_blob.empty())
460
                return -100;
461

462
            top_blob.dims = dims;
463

464
            q += slice;
465
        }
466

467
        #pragma omp parallel for num_threads(opt.num_threads)
468
        for (int p = 0; p < channels; p++)
469
        {
470
            const float* ptr = bottom_blob.channel(p);
471

472
            for (int j = 0; j < d; j++)
473
            {
474
                for (int k = 0; k < h; k++)
475
                {
476
                    for (size_t i = 0; i < top_blobs.size(); i++)
477
                    {
478
                        Mat& top_blob = top_blobs[i];
479

480
                        float* outptr = top_blob.channel(p).depth(j).row(k);
481
                        memcpy(outptr, ptr, top_blob.w * elemsize);
482

483
                        ptr += top_blob.w * elempack;
484
                    }
485
                }
486
            }
487
        }
488
    }
489

490
    if (dims == 4 && positive_axis == 1)
491
    {
492
        int w = bottom_blob.w;
493
        int h = bottom_blob.h;
494
        int d = bottom_blob.d;
495
        int channels = bottom_blob.c;
496

497
        int q = 0;
498
        for (size_t i = 0; i < top_blobs.size(); i++)
499
        {
500
            int slice;
501
            if (indices_ptr)
502
            {
503
                if (i == top_blobs.size() - 1)
504
                {
505
                    slice = d - q;
506
                }
507
                else
508
                {
509
                    int indice = indices_ptr[i];
510
                    int positive_indice = indice < 0 ? d + indice : indice;
511
                    slice = positive_indice - q;
512
                }
513
            }
514
            else
515
            {
516
                slice = slices_ptr[i];
517
                if (slice == -233)
518
                {
519
                    slice = static_cast<int>((d - q) / (top_blobs.size() - i));
520
                }
521
            }
522

523
            Mat& top_blob = top_blobs[i];
524
            top_blob.create(w, h, slice, channels, elemsize, elempack, opt.blob_allocator);
525
            if (top_blob.empty())
526
                return -100;
527

528
            q += slice;
529
        }
530

531
        #pragma omp parallel for num_threads(opt.num_threads)
532
        for (int p = 0; p < channels; p++)
533
        {
534
            const float* ptr = bottom_blob.channel(p);
535

536
            for (size_t i = 0; i < top_blobs.size(); i++)
537
            {
538
                Mat& top_blob = top_blobs[i];
539

540
                int size = top_blob.w * top_blob.h * top_blob.d;
541

542
                float* outptr = top_blob.channel(p);
543
                memcpy(outptr, ptr, size * elemsize);
544

545
                ptr += size * elempack;
546
            }
547
        }
548
    }
549

550
    return 0;
551
}
552

553
} // namespace ncnn
554