1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
// concat is a pure copy, so the kernels are keyed by element byte size
// (1/2/4/8) rather than logical type, matching the CUDA backend.
#define KERNEL_CONCAT(SUFFIX, T) \
kernel void kernel_concat_##SUFFIX( \
global const char * src0, \
ulong offset0, \
global const char * src1, \
ulong offset1, \
global char * dst, \
ulong offsetd, \
int ne00, \
int ne01, \
int ne02, \
int ne03, \
ulong nb00, \
ulong nb01, \
ulong nb02, \
ulong nb03, \
ulong nb10, \
ulong nb11, \
ulong nb12, \
ulong nb13, \
int ne0, \
ulong nb0, \
ulong nb1, \
ulong nb2, \
ulong nb3, \
int dim \
) { \
src0 = src0 + offset0; \
src1 = src1 + offset1; \
dst = dst + offsetd; \
\
const int i3 = get_group_id(2); \
const int i2 = get_group_id(1); \
const int i1 = get_group_id(0); \
\
int o[4] = {0, 0, 0, 0}; \
o[dim] = dim == 0 ? ne00 : (dim == 1 ? ne01 : (dim == 2 ? ne02 : ne03)); \
\
global const T * x; \
\
for (int i0 = get_local_id(0); i0 < ne0; i0 += get_local_size(0)) { \
if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { \
x = (global const T *)(src0 + (i3 )*nb03 + (i2 )*nb02 + (i1 )*nb01 + (i0 )*nb00); \
} else { \
x = (global const T *)(src1 + (i3 - o[3])*nb13 + (i2 - o[2])*nb12 + (i1 - o[1])*nb11 + (i0 - o[0])*nb10); \
} \
\
global T * y = (global T *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); \
\
*y = *x; \
} \
}
KERNEL_CONCAT(b1, char)
KERNEL_CONCAT(b2, short)
KERNEL_CONCAT(b4, int)
KERNEL_CONCAT(b8, long)
// packed variant for the common dim==0, small-ne0 case (4-byte elements only).
kernel void kernel_concat_b4_pack(
global const char * src0,
ulong offset0,
global const char * src1,
ulong offset1,
global char * dst,
ulong offsetd,
int ne00,
int ne01,
int ne02,
int ne03,
ulong nb00,
ulong nb01,
ulong nb02,
ulong nb03,
ulong nb10,
ulong nb11,
ulong nb12,
ulong nb13,
int ne0,
ulong nb0,
ulong nb1,
ulong nb2,
ulong nb3,
int dim,
int ne1,
int ne2,
int ne3
) {
src0 = src0 + offset0;
src1 = src1 + offset1;
dst = dst + offsetd;
int lsz = get_local_size(0);
int tpr = min(ne0, lsz); // threads per row
int rpw = lsz / tpr; // rows per workgroup
int lid = get_local_id(0);
int row = get_group_id(0)*rpw + lid / tpr;
int lane = lid - (lid / tpr) * tpr;
int nrows = ne1*ne2*ne3;
if (row >= nrows) {
return;
}
int i1 = row % ne1;
int t = row / ne1;
int i2 = t % ne2;
int i3 = t / ne2;
int o[4] = {0, 0, 0, 0};
o[dim] = dim == 0 ? ne00 : (dim == 1 ? ne01 : (dim == 2 ? ne02 : ne03));
for (int i0 = lane; i0 < ne0; i0 += tpr) {
global const int * x;
if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
x = (global const int *)(src0 + (i3 )*nb03 + (i2 )*nb02 + (i1 )*nb01 + (i0 )*nb00);
} else {
x = (global const int *)(src1 + (i3 - o[3])*nb13 + (i2 - o[2])*nb12 + (i1 - o[1])*nb11 + (i0 - o[0])*nb10);
}
global int * y = (global int *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
*y = *x;
}
}