recurrent_state_kernels.c File Reference

#include "ckernel_engine.h"
#include "ck_threadpool.h"
#include <limits.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

Go to the source code of this file.

Functions
static void	ck_recurrent_conv_state_channel_range (int begin, int end, void *opaque)
static int	recurrent_conv_backward_extents (int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim, int *channels_out, int *total_len_out, size_t *elements_out)
void	recurrent_conv_state_update_backward (const float *d_conv_x, const float *d_state_out, float *d_state_in, float *d_q, float *d_k, float *d_v, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)
void	recurrent_conv_state_update_backward_workspace (const float *d_conv_x, const float *d_state_out, float *d_state_in, float *d_q, float *d_k, float *d_v, float *d_conv_total, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)
void	recurrent_conv_state_update_forward (const float *state_in, const float *q, const float *k, const float *v, float *conv_x, float *state_out, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)

Function Documentation

ck_recurrent_conv_state_channel_range()

static void ck_recurrent_conv_state_channel_range ( int  begin, int  end, void *  opaque  )

static

Definition at line 25 of file recurrent_state_kernels.c.

                                                                {
    const ck_recurrent_conv_state_args_t *args =
        (const ck_recurrent_conv_state_args_t *)opaque;
 
    const int first_seq = begin / args->channels;
    const int last_seq = (end - 1) / args->channels;
    for (int seq = first_seq; seq <= last_seq; ++seq) {
        int ch_begin = begin - seq * args->channels;
        int ch_end = end - seq * args->channels;
        if (ch_begin < 0) ch_begin = 0;
        if (ch_end > args->channels) ch_end = args->channels;
 
        const size_t state_seq_offset =
            (size_t)seq * (size_t)args->channels *
            (size_t)args->history_len;
        const size_t conv_seq_offset =
            (size_t)seq * (size_t)args->channels *
            (size_t)args->total_len;
        for (int ch = ch_begin; ch < ch_end; ++ch) {
            memcpy(
                args->conv_x + conv_seq_offset +
                    (size_t)ch * (size_t)args->total_len,
                args->state_in + state_seq_offset +
                    (size_t)ch * (size_t)args->history_len,
                (size_t)args->history_len * sizeof(float));
        }
 
        const int q_begin = ch_begin;
        const int q_end = ch_end < args->q_dim ? ch_end : args->q_dim;
        const int k_begin = ch_begin > args->q_dim ?
            ch_begin - args->q_dim : 0;
        const int k_end = ch_end < args->q_dim + args->k_dim ?
            ch_end - args->q_dim : args->k_dim;
        const int v_begin = ch_begin > args->q_dim + args->k_dim ?
            ch_begin - args->q_dim - args->k_dim : 0;
        const int v_end = ch_end - args->q_dim - args->k_dim < args->v_dim ?
            ch_end - args->q_dim - args->k_dim : args->v_dim;
 
        for (int tok = 0; tok < args->num_tokens; ++tok) {
            const int row = seq * args->num_tokens + tok;
            for (int col = q_begin; col < q_end; ++col) {
                args->conv_x[conv_seq_offset +
                    (size_t)col * (size_t)args->total_len +
                    (size_t)(args->history_len + tok)] =
                    args->q[(size_t)row * (size_t)args->q_dim + (size_t)col];
            }
            for (int col = k_begin; col < k_end; ++col) {
                args->conv_x[conv_seq_offset +
                    (size_t)(args->q_dim + col) * (size_t)args->total_len +
                    (size_t)(args->history_len + tok)] =
                    args->k[(size_t)row * (size_t)args->k_dim + (size_t)col];
            }
            for (int col = v_begin; col < v_end; ++col) {
                args->conv_x[conv_seq_offset +
                    (size_t)(args->q_dim + args->k_dim + col) *
                        (size_t)args->total_len +
                    (size_t)(args->history_len + tok)] =
                    args->v[(size_t)row * (size_t)args->v_dim + (size_t)col];
            }
        }
 
        for (int ch = ch_begin; ch < ch_end; ++ch) {
            memcpy(
                args->state_out + state_seq_offset +
                    (size_t)ch * (size_t)args->history_len,
                args->conv_x + conv_seq_offset +
                    (size_t)ch * (size_t)args->total_len +
                    (size_t)args->num_tokens,
                (size_t)args->history_len * sizeof(float));
        }
    }
}

References end.

Referenced by recurrent_conv_state_update_forward().

recurrent_conv_backward_extents()

static int recurrent_conv_backward_extents ( int  history_len, int  num_seqs, int  num_tokens, int  q_dim, int  k_dim, int  v_dim, int *  channels_out, int *  total_len_out, size_t *  elements_out  )

static

Definition at line 155 of file recurrent_state_kernels.c.

                                                                 {
    if (history_len < 0 || num_seqs <= 0 || num_tokens < 0 || q_dim < 0 ||
        k_dim < 0 || v_dim < 0 || q_dim > INT_MAX - k_dim ||
        q_dim + k_dim > INT_MAX - v_dim || history_len > INT_MAX - num_tokens) {
        return 0;
    }
    const int channels = q_dim + k_dim + v_dim;
    const int total_len = history_len + num_tokens;
    if (channels == 0 || total_len == 0) {
        return 0;
    }
    size_t elements = (size_t)num_seqs;
    if ((size_t)total_len > SIZE_MAX / elements) {
        return 0;
    }
    elements *= (size_t)total_len;
    if ((size_t)channels > SIZE_MAX / elements) {
        return 0;
    }
    *channels_out = channels;
    *total_len_out = total_len;
    *elements_out = elements * (size_t)channels;
    return 1;
}

Referenced by recurrent_conv_state_update_backward(), and recurrent_conv_state_update_backward_workspace().

recurrent_conv_state_update_backward()

void recurrent_conv_state_update_backward	(	const float *	d_conv_x,
		const float *	d_state_out,
		float *	d_state_in,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		int	history_len,
		int	num_seqs,
		int	num_tokens,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 257 of file recurrent_state_kernels.c.

                                                     {
    int channels = 0;
    int total_len = 0;
    size_t elements = 0;
    if (!recurrent_conv_backward_extents(
            history_len, num_seqs, num_tokens, q_dim, k_dim, v_dim,
            &channels, &total_len, &elements) ||
        elements > SIZE_MAX / sizeof(float)) {
        return;
    }
    (void)channels;
    (void)total_len;
    float *workspace = (float *)malloc(elements * sizeof(float));
    if (!workspace) {
        return;
    }
    recurrent_conv_state_update_backward_workspace(
        d_conv_x, d_state_out, d_state_in, d_q, d_k, d_v, workspace,
        history_len, num_seqs, num_tokens, q_dim, k_dim, v_dim);
    free(workspace);
}

References recurrent_conv_backward_extents(), and recurrent_conv_state_update_backward_workspace().

recurrent_conv_state_update_backward_workspace()

void recurrent_conv_state_update_backward_workspace	(	const float *	d_conv_x,
		const float *	d_state_out,
		float *	d_state_in,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		float *	d_conv_total,
		int	history_len,
		int	num_seqs,
		int	num_tokens,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 188 of file recurrent_state_kernels.c.

                                                                {
    int channels = 0;
    int total_len = 0;
    size_t elements = 0;
    if (!d_conv_x || !d_state_out || !d_state_in || !d_q || !d_k || !d_v ||
        !d_conv_total || !recurrent_conv_backward_extents(
            history_len, num_seqs, num_tokens, q_dim, k_dim, v_dim,
            &channels, &total_len, &elements)) {
        return;
    }
    if (elements > SIZE_MAX / sizeof(float)) {
        return;
    }
 
    memcpy(d_conv_total, d_conv_x, elements * sizeof(float));
 
    for (int seq = 0; seq < num_seqs; ++seq) {
        const float *d_state_out_seq = d_state_out + (size_t) seq * (size_t) channels * (size_t) history_len;
        float *d_conv_seq = d_conv_total + (size_t) seq * (size_t) channels * (size_t) total_len;
        for (int ch = 0; ch < channels; ++ch) {
            float *dst = d_conv_seq + (size_t) ch * (size_t) total_len + (size_t) num_tokens;
            const float *src = d_state_out_seq + (size_t) ch * (size_t) history_len;
            for (int idx = 0; idx < history_len; ++idx) {
                dst[idx] += src[idx];
            }
        }
    }
 
    for (int seq = 0; seq < num_seqs; ++seq) {
        const float *d_conv_seq = d_conv_total + (size_t) seq * (size_t) channels * (size_t) total_len;
        float *d_state_in_seq = d_state_in + (size_t) seq * (size_t) channels * (size_t) history_len;
 
        for (int ch = 0; ch < channels; ++ch) {
            memcpy(
                d_state_in_seq + (size_t) ch * (size_t) history_len,
                d_conv_seq + (size_t) ch * (size_t) total_len,
                (size_t) history_len * sizeof(float));
        }
 
        for (int tok = 0; tok < num_tokens; ++tok) {
            const int row = seq * num_tokens + tok;
            float *d_q_row = d_q + (size_t) row * (size_t) q_dim;
            float *d_k_row = d_k + (size_t) row * (size_t) k_dim;
            float *d_v_row = d_v + (size_t) row * (size_t) v_dim;
            for (int col = 0; col < q_dim; ++col) {
                d_q_row[col] = d_conv_seq[(size_t) col * (size_t) total_len + (size_t) (history_len + tok)];
            }
            for (int col = 0; col < k_dim; ++col) {
                d_k_row[col] = d_conv_seq[(size_t) (q_dim + col) * (size_t) total_len + (size_t) (history_len + tok)];
            }
            for (int col = 0; col < v_dim; ++col) {
                d_v_row[col] = d_conv_seq[(size_t) (q_dim + k_dim + col) * (size_t) total_len + (size_t) (history_len + tok)];
            }
        }
    }
}

References recurrent_conv_backward_extents().

Referenced by recurrent_conv_state_update_backward().

recurrent_conv_state_update_forward()

void recurrent_conv_state_update_forward	(	const float *	state_in,
		const float *	q,
		const float *	k,
		const float *	v,
		float *	conv_x,
		float *	state_out,
		int	history_len,
		int	num_seqs,
		int	num_tokens,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 100 of file recurrent_state_kernels.c.

                                                    {
    if (!state_in || !q || !k || !v || !conv_x || !state_out ||
        history_len < 0 || num_seqs <= 0 || num_tokens < 0 ||
        q_dim < 0 || k_dim < 0 || v_dim < 0 ||
        q_dim > INT_MAX - k_dim || q_dim + k_dim > INT_MAX - v_dim ||
        history_len > INT_MAX - num_tokens) {
        return;
    }
    const int channels = q_dim + k_dim + v_dim;
    const int total_len = history_len + num_tokens;
    if (channels <= 0 || total_len <= 0 || num_seqs > INT_MAX / channels ||
        (num_tokens > 0 && num_seqs > INT_MAX / num_tokens)) {
        return;
    }
 
    ck_recurrent_conv_state_args_t args = {
        state_in, q, k, v, conv_x, state_out,
        history_len, num_tokens, q_dim, k_dim, v_dim, channels, total_len,
    };
    const int jobs = num_seqs * channels;
    const size_t elements = (size_t)jobs * (size_t)num_tokens;
    const size_t minimum_elements_per_thread = 32768u;
    if (num_tokens <= 1 || elements < 2u * minimum_elements_per_thread) {
        ck_recurrent_conv_state_channel_range(0, jobs, &args);
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    size_t useful_threads =
        (elements + minimum_elements_per_thread - 1u) /
        minimum_elements_per_thread;
    if ((size_t)active > useful_threads) active = (int)useful_threads;
    if (active > jobs) active = jobs;
 
    if (active > 1 && ck_threadpool_thread_id(pool) <= 0) {
        ck_threadpool_parallel_for_n(
            pool, active, 0, jobs, 32,
            ck_recurrent_conv_state_channel_range, &args);
    } else {
        ck_recurrent_conv_state_channel_range(0, jobs, &args);
    }
}

References ck_recurrent_conv_state_channel_range(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and ck_threadpool_thread_id().

Referenced by ck_test_recurrent_conv_state_update().