wehub-resource-sync
63 lines
1.9 KiB
Metal
63 lines
1.9 KiB
Metal
// DS4 Metal concat kernel used by the graph.
|
|
|
|
struct ds4_metal_args_concat {
|
|
int32_t ne00;
|
|
int32_t ne01;
|
|
int32_t ne02;
|
|
int32_t ne03;
|
|
uint64_t nb00;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
int32_t ne10;
|
|
int32_t ne11;
|
|
int32_t ne12;
|
|
int32_t ne13;
|
|
uint64_t nb10;
|
|
uint64_t nb11;
|
|
uint64_t nb12;
|
|
uint64_t nb13;
|
|
int32_t ne0;
|
|
int32_t ne1;
|
|
int32_t ne2;
|
|
int32_t ne3;
|
|
uint64_t nb0;
|
|
uint64_t nb1;
|
|
uint64_t nb2;
|
|
uint64_t nb3;
|
|
int32_t dim;
|
|
};
|
|
|
|
// Concatenates two float tensors along one dimension. In DS4 this is a graph
|
|
// utility for assembling attention inputs with exactly the same tensor layout
|
|
// expected by the downstream kernels.
|
|
kernel void kernel_concat(
|
|
constant ds4_metal_args_concat & args,
|
|
device const char * src0,
|
|
device const char * src1,
|
|
device char * dst,
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort3 tpitg[[thread_position_in_threadgroup]],
|
|
ushort3 ntg[[threads_per_threadgroup]]) {
|
|
const int i3 = tgpig.z;
|
|
const int i2 = tgpig.y;
|
|
const int i1 = tgpig.x;
|
|
|
|
int o[4] = {0, 0, 0, 0};
|
|
o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03));
|
|
|
|
device const float * x;
|
|
|
|
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
|
if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
|
|
x = (device const float *)(src0 + (i3 )*args.nb03 + (i2 )*args.nb02 + (i1 )*args.nb01 + (i0 )*args.nb00);
|
|
} else {
|
|
x = (device const float *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10);
|
|
}
|
|
|
|
device float * y = (device float *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
|
|
|
*y = *x;
|
|
}
|
|
}
|