wehub-resource-sync
53 lines
1.6 KiB
Metal
53 lines
1.6 KiB
Metal
// DS4 Metal repeat kernel used for HC embedding expansion.
|
|
|
|
struct ds4_metal_args_repeat {
|
|
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 ne0;
|
|
int32_t ne1;
|
|
int32_t ne2;
|
|
int32_t ne3;
|
|
uint64_t nb0;
|
|
uint64_t nb1;
|
|
uint64_t nb2;
|
|
uint64_t nb3;
|
|
};
|
|
|
|
// Repeats a source row into the HC channel dimension. DS4 uses this when the
|
|
// token embedding has to become an HC activation block before layer 0.
|
|
template<typename T>
|
|
kernel void kernel_repeat(
|
|
constant ds4_metal_args_repeat & args,
|
|
device const char * src0,
|
|
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;
|
|
|
|
const int i03 = i3%args.ne03;
|
|
const int i02 = i2%args.ne02;
|
|
const int i01 = i1%args.ne01;
|
|
|
|
device const char * src0_ptr = src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01;
|
|
device char * dst_ptr = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
|
|
|
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
|
const int i00 = i0%args.ne00;
|
|
*((device T *)(dst_ptr + i0*args.nb0)) = *((device T *)(src0_ptr + i00*args.nb00));
|
|
}
|
|
}
|
|
|
|
typedef decltype(kernel_repeat<float>) kernel_repeat_t;
|
|
|
|
// Host-visible F32 repeat used for HC expansion of embeddings.
|
|
template [[host_name("kernel_repeat_f32")]] kernel kernel_repeat_t kernel_repeat<float>;
|