layout_kernels.c File Reference

#include <stddef.h>
#include <string.h>

Go to the source code of this file.

Functions
void	ck_layout_head_to_token_f32 (const float *src, float *dst, int heads, int tokens, int head_dim)
void	ck_layout_token_to_head_f32 (const float *src, float *dst, int tokens, int heads, int head_dim)

Function Documentation

ck_layout_head_to_token_f32()

void ck_layout_head_to_token_f32	(	const float *	src,
		float *	dst,
		int	heads,
		int	tokens,
		int	head_dim
	)

Definition at line 19 of file layout_kernels.c.

{
    if (!src || !dst || heads <= 0 || tokens <= 0 || head_dim <= 0) return;
    for (int h = 0; h < heads; ++h) {
        for (int t = 0; t < tokens; ++t) {
            memcpy(dst + ((size_t)t * (size_t)heads + (size_t)h) * (size_t)head_dim,
                   src + ((size_t)h * (size_t)tokens + (size_t)t) * (size_t)head_dim,
                   (size_t)head_dim * sizeof(float));
        }
    }
}

ck_layout_token_to_head_f32()

void ck_layout_token_to_head_f32	(	const float *	src,
		float *	dst,
		int	tokens,
		int	heads,
		int	head_dim
	)

Definition at line 6 of file layout_kernels.c.

{
    if (!src || !dst || tokens <= 0 || heads <= 0 || head_dim <= 0) return;
    for (int t = 0; t < tokens; ++t) {
        for (int h = 0; h < heads; ++h) {
            memcpy(dst + ((size_t)h * (size_t)tokens + (size_t)t) * (size_t)head_dim,
                   src + ((size_t)t * (size_t)heads + (size_t)h) * (size_t)head_dim,
                   (size_t)head_dim * sizeof(float));
        }
    }
}