// Matrix Construct
//
// Copyright (C) Matrix Construct Developers, Authors & Contributors
// Copyright (C) 2016-2021 Jason Volk <jason@zemos.net>
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice is present in all copies. The
// full license for this software is available in the LICENSE file.

#pragma clang fp exceptions(ignore)
#pragma clang fp reassociate(on)
#pragma clang fp contract(fast)

#include <ircd/config.h>
#include <ircd/portable.h>
#include <clc/clc.h>
#include <ircd/simt/simt.h>
#include <ircd/gpt/vector.h>
#include <ircd/gpt/opts.h>
#include <ircd/gpt/ctrl.h>
#include <ircd/gpt/gpu.h>

//
// head
//

__kernel void
__attribute__((visibility("protected")))
ircd_gpt_alloc(__global const void *const restrict model,
               __global void *const restrict master,
               __constant const void *const opts,
               __global void *const restrict ctrl,
               __global void *const restrict frame0,
               __global void *const restrict frame1,
               __global void *const restrict frame2,
               __global void *const restrict frame3,
               __global void *const restrict frame4,
               __global void *const restrict frame5,
               __global void *const restrict frame6,
               __global void *const restrict frame7)
{
}

__kernel void
__attribute__((visibility("protected")))
ircd_gpt_enter(__global const void *const restrict model,
               __global void *const restrict state,
               __global void *const restrict master,
               __constant const struct ircd_gpt_opts *const opts,
               __global struct ircd_gpt_ctrl *const restrict ctrl)
{
	const ushort
	gi = get_global_id(0),
	li = get_local_id(0),
	ln = get_local_size(0),
	cycle = ctrl->clk.cycle;

	if(li == 0)
		ctrl->prof.entered = ircd_simt_cycles();
}

__kernel void
ircd_gpt_lm_embed(__global const struct ircd_gpt_ctrl *const ctrl,
                  __constant const struct ircd_gpt_opts *const opts,
                  __global ircd_gpt_vectorv *const restrict accum,
                  __global const ircd_gpt_vectorv *const restrict pos,
                  __global const ircd_gpt_vectorv *const restrict vocab)
{
	const ushort
	li = get_local_id(0),
	ln = get_local_size(0);

	const uint
	wo = get_global_offset(0);

	assume(ln == 192);
	assume(wo % ln == 0);

	const ushort
	wi = wo / ln + get_group_id(0);

	_ircd_gpt_lm_embed(ctrl, opts, accum, pos, vocab, wi, wi, li);
}

static void
_ircd_gpt_lm_embed(__global const struct ircd_gpt_ctrl *const ctrl,
                   __constant const struct ircd_gpt_opts *const opts,
                   __global ircd_gpt_vectorv *const restrict accum,
                   __global const ircd_gpt_vectorv *const restrict pos,
                   __global const ircd_gpt_vectorv *const restrict vocab,
                   const ushort out_idx,
                   const ushort tok_idx,
                   const ushort elem_idx)
{
	const ushort
	token = ctrl->token[tok_idx];

	const float4
	wpe = pos[tok_idx].elem[elem_idx],
	wte = vocab[token].elem[elem_idx],
	res = wte + wpe;

	accum[out_idx].elem[elem_idx] = res;
}

//
// Frontside
//

void
ircd_gpt_ffnn_fcon_tmul(__constant const struct ircd_gpt_opts *const opts,
                        __local ircd_gpt_ffnn_aperaturev *const restrict out,
                        __local const ircd_gpt_vectorv *const restrict in,
                        __global const ircd_gpt_ffnn_aperaturev *const restrict bias,
                        __global const ircd_gpt_ffnn_aperaturev *const restrict weight,
                        const uint li)
{
	const uint
	lanes = 4,
	segs = ircd_gpt_ffnn_segs,
	height = ircd_gpt_vector_elems / lanes;

	assume(height > 0);
	assume(height % lanes == 0);

	for(uint x = 0; x < segs; ++x)
		out->proj[x][li] = bias->proj[x][li];

	for(uint y = 0; y < height; ++y)
		for(uint k = 0; k < lanes; ++k)
			for(uint x = 0; x < segs; ++x)
			{
				const uint
				row = y * lanes + k;

				out->proj[x][li] += in->elem[y][k] * weight[row].proj[x][li];
			}
}

void
ircd_gpt_ffnn_fcon(__global const struct ircd_gpt_ctrl *const ctrl,
                   __constant const struct ircd_gpt_opts *const opts,
                   __local ircd_gpt_ffnn_aperaturev *const restrict out,
                   __local const ircd_gpt_vectorv *const restrict in,
                   __global const ircd_gpt_ffnn_aperaturev *const restrict bias,
                   __global const ircd_gpt_ffnn_aperaturev *const restrict weight,
                   const uint ln,
                   const uint li)
{
	const uint
	segs = ircd_gpt_ffnn_segs;

	// Fully connected
	ircd_gpt_ffnn_fcon_tmul
	(
		opts,
		out,
		in,
		bias,
		weight,
		li
	);

	for(uint i = 0; i < segs; ++i)
		ircd_gpt_ffnn_gelu(out, out, i * ln + li);
}

void
ircd_gpt_ffnn_proj_tmul(__constant const struct ircd_gpt_opts *const opts,
                        __local ircd_gpt_vectorv *const restrict out,
                        __local const ircd_gpt_ffnn_aperaturev *const restrict in,
                        __global const ircd_gpt_vectorv *const restrict bias,
                        __global const ircd_gpt_vectorv *const restrict weight,
                        const uint li)
{
	const uint
	lanes = 4,
	height = ircd_gpt_ffnn_fcon_elems / lanes;

	assume(height > 0);
	assume(height % lanes == 0);

	out->elem[li] = bias->elem[li];

	for(uint y = 0; y < height; ++y)
		for(uint k = 0; k < lanes; ++k)
		{
			const uint
			row = y * lanes + k;

			out->elem[li] += in->fcon[y][k] * weight[row].elem[li];
		}
}

void
ircd_gpt_ffnn(__global const struct ircd_gpt_ctrl *const ctrl,
              __constant const struct ircd_gpt_opts *const opts,
              __local ircd_gpt_vectorv *const restrict token,
              __local ircd_gpt_ffnn_aperaturev *const restrict buf,
              __local ircd_gpt_vectorv *const restrict tmp,
              __global const ircd_gpt_vectorv *const restrict norm_bias,
              __global const ircd_gpt_vectorv *const restrict norm_weight,
              __global const ircd_gpt_ffnn_aperaturev *const restrict fcon_bias,
              __global const ircd_gpt_ffnn_aperaturev *const restrict fcon_weight,
              __global const ircd_gpt_vectorv *const restrict proj_bias,
              __global const ircd_gpt_vectorv *const restrict proj_weight,
              const uint ln,
              const uint li)
{
	// Layer re-normalization
	ircd_gpt_norm(token, token, tmp, norm_bias, norm_weight, ln, li);

	// ln's writes are still pending but fcon reads results across threads.
	barrier(CLK_LOCAL_MEM_FENCE);

	// Fully connected
	ircd_gpt_ffnn_fcon
	(
		ctrl,
		opts,
		buf,
		token,
		fcon_bias,
		fcon_weight,
		ln,
		li
	);

	// fcon's writes are still pending but proj reads results across threads.
	barrier(CLK_LOCAL_MEM_FENCE);

	// Projection
	ircd_gpt_ffnn_proj_tmul
	(
		opts,
		token,
		buf,
		proj_bias,
		proj_weight,
		li
	);
}

static void
ircd_gpt_attn_self_samax(__global const struct ircd_gpt_ctrl *const ctrl,
                         __constant const struct ircd_gpt_opts *const opts,
                         __local float self[restrict][12],
                         const uint ln,
                         const uint li,
                         const uint wn,
                         const uint wi)
{
	struct ircd_math_samax samax =
	{
		.mu = -10000.0f,
		.sum = 0.0f,
	};

	__attribute__((opencl_unroll_hint))
	for(uint i = 0; i < wn; ++i)
		samax.mu = max(samax.mu, self[i][li]);

	__attribute__((opencl_unroll_hint))
	for(uint i = 0; i < wn; ++i)
		self[i][li] -= samax.mu;

	for(uint i = 0; i < wn; ++i)
		self[i][li] = native_exp(self[i][li]);

	__attribute__((opencl_unroll_hint))
	for(uint i = 0; i < wn; ++i)
		samax.sum += self[i][li];

	samax.sum += FLT_EPSILON;
	samax.lambda = 1.0f / samax.sum;

	__attribute__((opencl_unroll_hint))
	for(uint i = 0; i < wn; ++i)
		self[i][li] *= samax.lambda;
}

static void
ircd_gpt_attn_self_keys(__global const struct ircd_gpt_ctrl *const ctrl,
                        __constant const struct ircd_gpt_opts *const opts,
                        __local float self[restrict][ircd_gpt_attn_rank],
                        __global const ircd_gpt_attn_qkvv *const restrict token,
                        const uint ln,
                        const uint li,
                        const uint wi,
                        const uint kn,
                        const uint i)
{
	assume(i < wi);

	self[i][li] = 0.0f;

	__attribute__((opencl_unroll_hint))
	for(uint k = 0; k < kn; ++k)
	{
		float4
		qry = token[wi].qry.attn[li][k],
		key = token[i].key.attn[li][k],
		res = qry * key;

		self[i][li] += ircd_simt_hadd_f4(res);
	}

	self[i][li] /= 8.0f;
}

static void
ircd_gpt_attn_self_vals(__global const struct ircd_gpt_ctrl *const ctrl,
                        __constant const struct ircd_gpt_opts *const opts,
                        __local ircd_gpt_vectorv *const restrict out,
                        __local const float self[restrict][ircd_gpt_attn_rank],
                        __global const ircd_gpt_attn_qkvv *const restrict token,
                        const uint li,
                        const uint wi,
                        const uint ki,
                        const uint ti)
{
	out->attn[ti][ki] = 0.0f;

	__attribute__((opencl_unroll_hint))
	for(uint i = 0; i < wi; ++i)
	{
		const float4
		val = token[i].val.attn[ti][ki],
		attn = self[i][ti],
		res = attn * val;

		out->attn[ti][ki] += res;
	}
}

static void
ircd_gpt_attn_self(__global struct ircd_gpt_ctrl *const ctrl,
                   __constant const struct ircd_gpt_opts *const opts,
                   __local ircd_gpt_vectorv *const restrict out,
                   __local float self[restrict][ircd_gpt_attn_rank],
                   __global float attns[restrict][ircd_gpt_attn_rank],
                   __global const ircd_gpt_attn_qkvv *const restrict token,
                   const uint ln,
                   const uint li,
                   const uint wi)
{
	//assume(opts->attn_rank == sizeof(self[0]) / sizeof(float));
	assume(opts->attn_rank == ircd_gpt_attn_rank);
	assume(ctrl->count < ircd_gpt_context_tokens);
	assume(ctrl->tokens <= ircd_gpt_context_tokens);
	assume(ctrl->tokens > wi);
	assume(ctrl->tokens > 0);

	const uint
	wn = ctrl->tokens,
	kn = ln / opts->attn_rank,
	ki = li / opts->attn_rank,
	ti = li % opts->attn_rank;

	// Low-rank mask
	if(li < opts->attn_rank)
	{
		// Left attention
		uint i;
		for(i = 0; i < wi; ++i)
			ircd_gpt_attn_self_keys(ctrl, opts, self, token, ln, li, wi, kn, i);

		// Future mask
		__attribute__((opencl_unroll_hint))
		while(i < wn)
			self[i++][li] = -10000.0f;

		// Three-piece softmax
		ircd_gpt_attn_self_samax(ctrl, opts, self, ln, li, wn, wi);
	}

	// Propagate to full width for value dot prod.
	barrier(CLK_LOCAL_MEM_FENCE);
	ircd_gpt_attn_self_vals(ctrl, opts, out, self, token, li, wi, ki, ti);

	// Save softmax results for later analysis/observation.
	if(li < opts->attn_rank)
	{
		__attribute__((opencl_unroll_hint))
		for(uint i = 0; i < wn; ++i)
			attns[i][li] = self[i][li];
	}
}

static void
ircd_gpt_attn_proj_tmul(__constant const struct ircd_gpt_opts *const opts,
                        __local ircd_gpt_vectorv *const restrict out,
                        __local const ircd_gpt_vectorv *const restrict in,
                        __global const ircd_gpt_vectorv *const restrict bias,
                        __global const ircd_gpt_vectorv *const restrict weight,
                        const uint li)
{
	const uint
	lanes = 4,
	height = ircd_gpt_vector_elems / 4;

	assume(height > 0);
	assume(height % lanes == 0);

	out->elem[li] = bias->elem[li];

	for(uint y = 0; y < height; ++y)
		for(uint k = 0; k < lanes; ++k)
		{
			const uint
			row = y * lanes + k;

			const float4
			a = in->elem[y][k],
			b = weight[row].elem[li];

			out->elem[li] += a * b;
		}
}

__kernel void
__attribute__((visibility("protected")))
ircd_gpt_coil(__global struct ircd_gpt_ctrl *const ctrl,
              __constant const struct ircd_gpt_opts *const opts,
              __private const uint layer,
              __global ircd_gpt_vectorv *const restrict accum,
              __global float attns[restrict][ircd_gpt_attn_rank],
              __global const ircd_gpt_attn_qkvv *const restrict state,
              __global const ircd_gpt_vectorv *const restrict attn_proj_bias,
              __global const ircd_gpt_vectorv *const restrict attn_proj_weight,
              __global const ircd_gpt_vectorv *const restrict ffnn_norm_bias,
              __global const ircd_gpt_vectorv *const restrict ffnn_norm_weight,
              __global const ircd_gpt_ffnn_aperaturev *const restrict ffnn_fcon_bias,
              __global const ircd_gpt_ffnn_aperaturev *const restrict ffnn_fcon_weight,
              __global const ircd_gpt_vectorv *const restrict ffnn_proj_bias,
              __global const ircd_gpt_vectorv *const restrict ffnn_proj_weight)
{
	const uint
	li = get_local_id(0),
	ln = get_local_size(0),
	wo = get_global_offset(0),
	wi = wo / ln + get_group_id(0);

	assume(ln == 192);
	assume(wo % ln == 0);

	__local union
	{
		float
		attn_self[ircd_gpt_context_tokens][ircd_gpt_attn_rank];

		ircd_gpt_ffnn_aperaturev
		ffnn_fcon[2];

		ircd_gpt_vectorv
		vector[8];
	}
	buf;

	__local ircd_gpt_vectorv
	buf0, buf1,
	*const restrict attn_self = &buf1,
	*const restrict token = &buf0,
	*const restrict tmp = &buf1;

	// Self-attention backend; this computes the self-attention result now
	// that keys and values are globally visible across tokens.
	ircd_gpt_attn_self
	(
		ctrl,
		opts,
		attn_self,
		buf.attn_self,
		attns,
		state,
		ln,
		li,
		wi
	);

	barrier(CLK_LOCAL_MEM_FENCE);

	// Project result of self-attention.
	ircd_gpt_attn_proj_tmul
	(
		opts,
		token,
		attn_self,
		attn_proj_bias,
		attn_proj_weight,
		li
	);

	// Frontend accumulation
	{
		const float4
		attn = token->elem[li],
		resid = accum[wi].elem[li],
		result = resid + attn;

		token->elem[li] = result;
		accum[wi].elem[li] = result;
	}

	barrier(CLK_GLOBAL_MEM_FENCE);

	// Backend mlp
	ircd_gpt_ffnn
	(
		ctrl,
		opts,
		token,
		buf.ffnn_fcon,
		tmp,
		ffnn_norm_bias,
		ffnn_norm_weight,
		ffnn_fcon_bias,
		ffnn_fcon_weight,
		ffnn_proj_bias,
		ffnn_proj_weight,
		ln,
		li
	);

	// Backend accumulation
	{
		const float4
		ffnn = token->elem[li],
		resid = accum[wi].elem[li],
		result = resid + ffnn;

		accum[wi].elem[li] = result;
	}
}

static void
ircd_gpt_attn_fcon_tmul(__constant const struct ircd_gpt_opts *const opts,
                        __local ircd_gpt_attn_aperaturev *const restrict out,
                        __local const ircd_gpt_vectorv *const restrict in,
                        __global const ircd_gpt_attn_aperaturev *const restrict bias,
                        __global const ircd_gpt_attn_aperaturev *const restrict weight,
                        const uint ln,
                        const uint li)
{
	const uint
	lanes = 4,
	segs = ircd_gpt_attn_segs,
	height = ircd_gpt_vector_elems / lanes;

	assume(height > 0);
	assume(height % segs == 0);
	assume(height % lanes == 0);

	for(uint x = 0; x < segs; ++x)
		out->proj[x][li] = bias->proj[x][li];

	for(uint y = 0; y < height; ++y)
		for(uint k = 0; k < lanes; ++k)
			for(uint x = 0; x < segs; ++x)
			{
				const uint
				row = y * lanes + k;

				const float4
				a = in->elem[y][k],
				b = weight[row].proj[x][li];

				out->proj[x][li] += a * b;
			}
}

__kernel void
ircd_gpt_attn_fcon(__global const struct ircd_gpt_ctrl *const ctrl,
                   __constant const struct ircd_gpt_opts *const opts,
                   __private const uint layer,
                   __global ircd_gpt_attn_aperaturev *const restrict state,
                   __global const ircd_gpt_vectorv *const restrict accum,
                   __global const ircd_gpt_vectorv *const restrict norm_bias,
                   __global const ircd_gpt_vectorv *const restrict norm_weight,
                   __global const ircd_gpt_attn_aperaturev *const restrict fcon_bias,
                   __global const ircd_gpt_attn_aperaturev *const restrict fcon_weight)
{
	const uint
	li = get_local_id(0),
	ln = get_local_size(0),
	wo = get_global_offset(0),
	wi = wo / ln + get_group_id(0),
	segs = ircd_gpt_attn_segs;

	assume(ln == 192);
	assume(wo % ln == 0);

	__local ircd_gpt_attn_aperaturev
	attn;

	__local ircd_gpt_vectorv
	token, *const restrict tmp = attn.vector;

	token.elem[li] = accum[wi].elem[li];

	// Layer re-normalization
	ircd_gpt_norm(&token, &token, tmp, norm_bias, norm_weight, ln, li);

	// Ln's writes are still pending; fcon requires results across threads.
	barrier(CLK_LOCAL_MEM_FENCE);

	// Fully connected
	ircd_gpt_attn_fcon_tmul
	(
		opts,
		&attn,
		&token,
		fcon_bias,
		fcon_weight,
		ln,
		li
	);

	// Export queries, keys, and values.
	for(uint x = 0; x < segs; ++x)
		state[wi].proj[x][li] = attn.proj[x][li];
}

__kernel void
ircd_gpt_lm_norm(__global const struct ircd_gpt_ctrl *const ctrl,
                 __constant const struct ircd_gpt_opts *const opts,
                 __global ircd_gpt_vectorv *const restrict accum,
                 __global const ircd_gpt_vectorv *const restrict norm_bias,
                 __global const ircd_gpt_vectorv *const restrict norm_weight)
{
	const uint
	li = get_local_id(0),
	ln = get_local_size(0),
	wo = get_global_offset(0),
	wi = wo / ln + get_group_id(0);

	assume(ln == 192);
	assume(wo % ln == 0);

	__local ircd_gpt_vectorv
	tmp, token;

	token.elem[li] = accum[wi].elem[li];

	// Final re-normalization
	ircd_gpt_norm(&token, &token, &tmp, norm_bias, norm_weight, ln, li);

	accum[wi].elem[li] = token.elem[li];
}

__kernel void
ircd_gpt_lm_logit(__global const struct ircd_gpt_ctrl *const ctrl,
                  __constant const struct ircd_gpt_opts *const opts,
                  __global float *const restrict logit,
                  __global const ircd_gpt_vectorv *const restrict accum,
                  __global const ircd_gpt_vectorv *const restrict pos,
                  __global const ircd_gpt_vectorv *const restrict vocab)
{
	const uint
	gi = get_global_id(0),
	wi = ctrl->count - 1;

	assume(opts->embed_width == 192);
	assume(opts->logits <= 65536);

	if(gi >= opts->logits)
	{
		logit[gi] = -10000.0f;
		return;
	}

	float acc = 0.0f;
	for(uint j = 0; j < opts->embed_width; ++j)
	{
		const float4
		token = vocab[gi].elem[j],
		in = accum[wi].elem[j],
		wpe = pos[wi].elem[j],
		res = in * token + wpe;

		acc += ircd_simt_hadd_f4(res);
	}

	logit[gi] = acc;
}

__kernel void
__attribute__((reqd_work_group_size(256, 1, 1)))
ircd_gpt_lm_logsm(__global struct ircd_gpt_ctrl *const ctrl,
                  __constant const struct ircd_gpt_opts *const opts,
                  __global float logit[restrict 65536])
{
	const uint
	li = get_local_id(0),
	ln = get_local_size(0),
	//wo = get_global_offset(0),
	//wi = wo / ln + get_group_id(0),
	wn = 50432,
	tn = wn / ln,
	start = tn * li,
	stop = min(start + tn, opts->logits);

	__local float
	mu[256], sum[256], lambda[256];

	__local struct ircd_math_samax
	samax;

	assume(ln == 256);

	mu[li] = -10000.0f;
	__attribute__((opencl_unroll_hint))
	for(uint ti = start; ti < stop; ++ti)
		mu[li] = max(mu[li], logit[ti]);

	ircd_simt_reduce_max_flldr(mu, ln, li);

	if(li == 0)
		samax.mu = mu[li];

	ircd_simt_broadcast_flldr(mu, ln, li);

	sum[li] = 0.0f;
	for(uint ti = start; ti < stop; ++ti)
	{
		const float
		sub = logit[ti] - mu[li],
		res = native_exp(sub);

		sum[li] += res;
	}

	ircd_simt_reduce_add_flldr(sum, ln, li);

	if(li == 0)
		sum[li] += FLT_EPSILON,
		samax.sum = sum[li],
		samax.lambda = lambda[li] = 1.0f / sum[li];

	ircd_simt_broadcast_flldr(lambda, ln, li);

	for(uint ti = start; ti < stop; ++ti)
	{
		const float
		sub = logit[ti] - mu[li],
		res = lambda[li] * native_exp(sub);

		logit[ti] = res;
	}
}

void
ircd_gpt_lm_result_top(__local struct ircd_gpt_ctrl *const ctrl,
                       __constant const struct ircd_gpt_opts *const opts,
                       __local const ushort *const restrict idx,
                       __global const float *const restrict logsm,
                       const uint i)
{
	const ushort
	token = idx[i];

	const float
	samax = logsm[token];

	ctrl->top[i].token = token;
	ctrl->top[i].samax = samax;
}

void
ircd_gpt_lm_result_label_mean(__local struct ircd_gpt_ctrl *const ctrl,
                              __constant const struct ircd_gpt_opts *const opts,
                              __local struct ircd_math_mean *const mean,
                              const float last)
{
	const uint
	div = mean->div + 1,
	sum_sel = mean->div % 4;

	const float
	sum = mean->sum[0] + mean->sum[1] + mean->sum[2] + mean->sum[3] + last,
	res = sum / div;

	mean->sum[sum_sel] += last;
	mean->div = div;
	mean->last = last;
	mean->mean = res;
}

void
ircd_gpt_lm_result_label(__local struct ircd_gpt_ctrl *const ctrl,
                         __constant const struct ircd_gpt_opts *const opts,
                         __local struct ircd_gpt_ctrl_label *const label,
                         __global const float *const restrict logsm)
{
	const ushort
	token = label->logit.token;

	const float
	samax = logsm[token],
	loss = 0.0f - native_log(samax),
	ppl = (1.0f - samax) * native_log2(opts->logits);

	label->logit.samax = samax;
	ircd_gpt_lm_result_label_mean(ctrl, opts, &label->loss, loss);
	ircd_gpt_lm_result_label_mean(ctrl, opts, &label->ppl, ppl);
}

ushort
ircd_gpt_lm_result_select(__local struct ircd_gpt_ctrl *const ctrl,
                          __constant const struct ircd_gpt_opts *const opts,
                          __local const ushort *const restrict idx,
                          __global const float *const restrict logsm)
{
	const ulong
	ent_k = max(opts->top_k, 1U) - 1,
	rnd = ircd_simt_rand_xoshiro256pl(ctrl->rand);

	const float
	ent_p = min(max(opts->top_p, 0.0f), 1.0f),
	thresh = ent_p;

	float acc = 1.0f;
	ushort select = 0;
	for(; select < ent_k; ++select)
		if((acc -= logsm[idx[select]]) < thresh)
			break;

	const ushort
	token = idx[select];

	return token;
}

static ushort
ircd_gpt_lm_result(__local struct ircd_gpt_ctrl *const ctrl,
                   __constant const struct ircd_gpt_opts *const opts,
                   __local const ushort *const restrict idx,
                   __global const float *const restrict logsm)
{
	const ushort
	token = ircd_gpt_lm_result_select(ctrl, opts, idx, logsm);

	// Update the dynamic result label.
	ctrl->select.logit.token = token;
	ircd_gpt_lm_result_label(ctrl, opts, &ctrl->select, logsm);

	// Update the dynamic target label.
	ctrl->target.logit.token = ctrl->count < ctrl->tokens?
		ctrl->token[ctrl->count]:
		ctrl->select.logit.token;

	ircd_gpt_lm_result_label(ctrl, opts, &ctrl->target, logsm);

	const bool
	hit = ctrl->select.logit.token == ctrl->target.logit.token;

	// Update the token context.
	if(ctrl->count == ctrl->tokens)
		ctrl->token[ctrl->tokens++] = ctrl->select.logit.token;
	else
		ctrl->accept = -2;

	ctrl->miss += !hit;
	ctrl->hit += hit;
	ctrl->count++;
	return token;
}

static void
ircd_gpt_lm_result_attns(__local struct ircd_gpt_ctrl *const ctrl,
                         __constant const struct ircd_gpt_opts *const opts,
                         __global const float *const restrict attns,
                         const uint ln,
                         const uint li)
{
	const uint
	layer = li / opts->layers,
	head = li % opts->attn_rank,
	base = layer * opts->attn_self_elems;

	uint best = 0;
	float bestv = 10000.0f;
	for(uint i = 0; i < ctrl->count; ++i)
	{
		const uint
		bx = (((i + 1) * i) / 2) * opts->attn_rank,
		idx = base + bx + i * 12 + head;

		if(attns[idx] < bestv)
			bestv = attns[idx],
			best = i;
	}

	ctrl->attn[layer][head] = best;
}

__kernel void
__attribute__((visibility("protected")))
__attribute__((reqd_work_group_size(256, 1, 1)))
ircd_gpt_lm_select(__global struct ircd_gpt_ctrl *const restrict ctrl_,
                   __constant const struct ircd_gpt_opts *const opts,
                   __global const float logsm[restrict 65536],
                   __global const float *const restrict attns)
{
	const uint
	li = get_local_id(0),
	ln = get_local_size(0),
	logits_pad = ln - (opts->logits % ln),
	tn = (opts->logits + logits_pad) / ln,
	start = tn * li,
	stop = min(start + tn, opts->logits);

	__local ushort idx[256];
	__local struct ircd_gpt_ctrl ctrl;
	__private event_t event[1];

	assume(ln == 256);
	assume(start < stop);

	event[0] = async_work_group_copy
	(
		(__local char16 *)&ctrl,
		(__global const char16 *)ctrl_,
		sizeof(struct ircd_gpt_ctrl) / sizeof(char16),
		0
	);

	idx[li] = start;
	for(uint j = start + 1; j < stop; ++j)
		if(logsm[j] > logsm[idx[li]])
			idx[li] = j;

	ircd_simt_sort_idx16_flldr(idx, logsm, ln, li);
	wait_group_events(1, event);

	if(ctrl.count >= opts->buffer_tokens)
		return;

	if(li < opts->top_n)
		ircd_gpt_lm_result_top(&ctrl, opts, idx, logsm, li);

	if(li < opts->labels)
		ircd_gpt_lm_result_label(&ctrl, opts, ctrl.label + li, logsm);

	if(li < opts->layers * opts->attn_rank)
		ircd_gpt_lm_result_attns(&ctrl, opts, attns, ln, li);

	barrier(CLK_LOCAL_MEM_FENCE);

	if(li == 0)
		ircd_gpt_lm_result(&ctrl, opts, idx, logsm);

	barrier(CLK_LOCAL_MEM_FENCE);

	event[0] = async_work_group_copy
	(
		(__global char16 *)ctrl_,
		(__local const char16 *)&ctrl,
		sizeof(struct ircd_gpt_ctrl) / sizeof(char16),
		0
	);

	wait_group_events(1, event);
}

__kernel void
__attribute__((visibility("protected")))
__attribute__((reqd_work_group_size(256, 1, 1)))
ircd_gpt_leave(__global const void *const restrict model,
               __global void *const restrict state,
               __global void *const restrict master,
               __constant const struct ircd_gpt_opts *const opts,
               __global struct ircd_gpt_ctrl *const ctrl_,
               __global struct ircd_gpt_ctrl *const frame)
{
	const ushort
	li = get_local_id(0),
	ln = get_local_size(0);

	assume(ln == 256);

	__local struct ircd_gpt_ctrl _ctrl;
	__local struct ircd_gpt_ctrl *const ctrl = &_ctrl;

	if(li == 0)
		*ctrl = *ctrl_;

	barrier(CLK_LOCAL_MEM_FENCE);

	if(li == 0 && ctrl->accept < 0)
		ircd_gpt_accept(ctrl, opts);

	barrier(CLK_LOCAL_MEM_FENCE);

	const uint
	batch_size = opts->batch_size,
	samps = opts->training_steps + opts->validation_steps + opts->testing_steps,
	steps = samps / batch_size;

	const bool
	accepting = ctrl->accept >= 0,
	cycling = !accepting,
	sampling = accepting,
	stepping = sampling && (ctrl->clk.samp + 1) >= batch_size,
	epoching = stepping && (ctrl->clk.step + 1) >= steps;

	if(li == 0)
		ctrl->prof.finished = ircd_simt_cycles();

	if(li == 0)
	{
		// BARTS won't update the scalar after the copy. In this case we'll
		// be setting magic for all remaining frames :/
		#if defined(__R600__)
		ctrl->magic = 0xC7012C70UL;
		#endif

		*frame = *ctrl;
		frame->magic = 0xC7012C70UL;
	}

	if(li == 0 && !accepting)
	{
		ctrl->clk.cycle += cycling;
		ctrl->clk.samp += sampling;
		ctrl->clk.step += stepping;
		ctrl->clk.epoch += epoching;
	}

	if(li == 0 && accepting)
		ctrl->magic = 0xC7012C70UL;

	if(li == 0)
		*ctrl_ = *ctrl;
}

void
ircd_gpt_accept(__local struct ircd_gpt_ctrl *const ctrl,
                __constant const struct ircd_gpt_opts *const opts)
{
	const bool
	unlimited = opts->limit < 0,
	unprocessed = ctrl->accept == -2,
	acceptable = ctrl->accept < 0 && !unprocessed;

	const uint
	batch_size = opts->batch_size,
	samps = opts->training_steps + opts->validation_steps + opts->testing_steps,
	steps = samps / batch_size,
	unproc = ctrl->tokens - ctrl->count,
	limit = opts->limit;

	const int
	cycle_remain = opts->context_tokens - (ctrl->clk.cycle + 1), // cycle not yet incr
	token_remain = opts->context_tokens - ctrl->count, // but count already incr
	remain = min(cycle_remain, token_remain);

	int
	accept = ircd_gpt_accept_check(ctrl, opts),
	dispatch = accept < 0? min(abs(accept), (uint)remain): 0;

	if(opts->limit > 0 && ctrl->clk.cycle >= limit - 1)
	{
		accept = accept >= 0? accept: -1;
		dispatch = 0;
	}

	if(opts->limit == 0 && unprocessed)
	{
		accept = accept >= 0? accept: -1;
		dispatch = min((uint)remain, unproc);
	}

	if(opts->limit != 0 && !acceptable)
	{
		accept = accept >= 0? accept: -1;
		dispatch = max(dispatch, 1);
	}

	ctrl->accept = accept;
	ctrl->dispatch = min((uint)dispatch, opts->frames);
}

int
ircd_gpt_accept_check(__local struct ircd_gpt_ctrl *const ctrl,
                      __constant const struct ircd_gpt_opts *const opts)
{
	int best = opts->frames;
	for(uint i = 0; i < 4; ++i)
	{
		const int
		remain = ircd_gpt_accept_match(ctrl, opts, i);

		if(remain == 0)
			return i;

		if(remain < best)
			best = remain;
	}

	return -best;
}

uint
ircd_gpt_accept_match(__local struct ircd_gpt_ctrl *const ctrl,
                      __constant const struct ircd_gpt_opts *const opts,
                      const uint i)
{
	const uint
	len = ircd_gpt_accept_len(ctrl, opts, i),
	n = min(ctrl->count, len),
	maxlen = opts->frames;

	uint ret = len?: maxlen;
	for(uint j = 1; j <= n; ++j)
	{
		uint match = 0;
		for(; match < j; ++match)
		{
			const uint
			accept = opts->accept[i][match],
			token = ctrl->token[ctrl->count - j + match];

			if(token != accept)
				break;
		}

		if(match >= j)
			if(!(ret = len - match))
				break;
	}

	ret = max(ret, ctrl->tokens - ctrl->count);
	ret = min(ret, maxlen);
	return ret;
}

uint
ircd_gpt_accept_len(__local struct ircd_gpt_ctrl *const ctrl,
                    __constant const struct ircd_gpt_opts *const opts,
                    const uint i)
{
	uint len = 0;
	for(; len < 8; ++len)
		if(opts->accept[i][len] == (ushort)-1U)
			break;

	return len;
}

//
// backside
//

void
__attribute__((always_inline))
ircd_gpt_prop_elem(__global const struct ircd_gpt_ctrl *const ctrl,
                   __constant const struct ircd_gpt_opts *const opts,
                   __global float4 *const restrict param_,
                   __global float4 *const restrict exp_avg_,
                   __global float4 *const restrict exp_avg_sqr_)
{
	const uint
	li = get_local_id(0),
	ts = ctrl->clk.step;

	const float4
	param = param_[li],
	grad = ctrl->target.loss.mean,
	alpha[2] = { 1.0f - opts->beta[0], 1.0f - opts->beta[1], },
	exp_avg = ts? exp_avg_[li]: 0.0f,
	exp_avg_sqr = ts? exp_avg_sqr_[li]: 0.0f,
	exp_avg_mul = exp_avg * opts->beta[0],
	exp_avg_dot = exp_avg_mul + alpha[0] * grad,
	exp_avg_sqr_mul = exp_avg_sqr * opts->beta[1],
	exp_avg_sqr_dot = exp_avg_sqr_mul + alpha[1] * grad * grad,
	denom = native_sqrt(exp_avg_sqr_dot) + FLT_EPSILON,
	delta = opts->alpha * (exp_avg_dot / denom),
	update = param - delta;

	param_[li] = update;
	exp_avg_[li] = exp_avg_dot;
	exp_avg_sqr_[li] = exp_avg_sqr_dot;
}

//
// backpropagations
//

__kernel void
__attribute__((always_inline))
ircd_gpt_norm_prop(__global const struct ircd_gpt_ctrl *const ctrl,
                   __constant const struct ircd_gpt_opts *const opts,
                   __global ircd_gpt_vectorv *const restrict bias,
                   __global ircd_gpt_vectorv *const restrict bias_m0,
                   __global ircd_gpt_vectorv *const restrict bias_m1,
                   __global ircd_gpt_vectorv *const restrict weight,
                   __global ircd_gpt_vectorv *const restrict weight_m0,
                   __global ircd_gpt_vectorv *const restrict weight_m1)
{
	ircd_gpt_prop_elem
	(
		ctrl, opts,
		bias->elem,
		bias_m0->elem,
		bias_m1->elem
	);

	ircd_gpt_prop_elem
	(
		ctrl, opts,
		weight->elem,
		weight_m0->elem,
		weight_m1->elem
	);
}

__kernel void
ircd_gpt_coil_prop_attn(__global const struct ircd_gpt_ctrl *const ctrl,
                        __constant const struct ircd_gpt_opts *const opts,
                        __global ircd_gpt_vectorv *const restrict norm_bias,
                        __global ircd_gpt_vectorv *const restrict norm_bias_m0,
                        __global ircd_gpt_vectorv *const restrict norm_bias_m1,
                        __global ircd_gpt_vectorv *const restrict norm_weight,
                        __global ircd_gpt_vectorv *const restrict norm_weight_m0,
                        __global ircd_gpt_vectorv *const restrict norm_weight_m1,
                        __global ircd_gpt_attn_aperaturev *const restrict fcon_bias,
                        __global ircd_gpt_attn_aperaturev *const restrict fcon_bias_m0,
                        __global ircd_gpt_attn_aperaturev *const restrict fcon_bias_m1,
                        __global ircd_gpt_attn_aperaturev *const restrict fcon_weight,
                        __global ircd_gpt_attn_aperaturev *const restrict fcon_weight_m0,
                        __global ircd_gpt_attn_aperaturev *const restrict fcon_weight_m1,
                        __global ircd_gpt_vectorv *const restrict proj_bias,
                        __global ircd_gpt_vectorv *const restrict proj_bias_m0,
                        __global ircd_gpt_vectorv *const restrict proj_bias_m1,
                        __global ircd_gpt_vectorv *const restrict proj_weight,
                        __global ircd_gpt_vectorv *const restrict proj_weight_m0,
                        __global ircd_gpt_vectorv *const restrict proj_weight_m1)
{
	const uint
	fcon_height = opts->embed_elems,
	proj_height = opts->embed_elems,
	segs = 3;

	ircd_gpt_norm_prop
	(
		ctrl, opts,
		norm_bias,
		norm_bias_m0,
		norm_bias_m1,
		norm_weight,
		norm_weight_m0,
		norm_weight_m1
	);

	for(uint j = 0; j < segs; ++j)
		ircd_gpt_prop_elem
		(
			ctrl, opts,
			fcon_bias->proj[j],
			fcon_bias_m0->proj[j],
			fcon_bias_m1->proj[j]
		);

	for(uint i = 0; i < fcon_height; ++i)
		for(uint j = 0; j < segs; ++j)
			ircd_gpt_prop_elem
			(
				ctrl, opts,
				fcon_weight[i].proj[j],
				fcon_weight_m0[i].proj[j],
				fcon_weight_m1[i].proj[j]
			);

	ircd_gpt_prop_elem
	(
		ctrl, opts,
		proj_bias->elem,
		proj_bias_m0->elem,
		proj_bias_m1->elem
	);

	for(uint i = 0; i < proj_height; ++i)
		ircd_gpt_prop_elem
		(
			ctrl, opts,
			proj_weight[i].elem,
			proj_weight_m0[i].elem,
			proj_weight_m1[i].elem
		);
}

__kernel void
ircd_gpt_coil_prop_ffnn(__global const struct ircd_gpt_ctrl *const ctrl,
                        __constant const struct ircd_gpt_opts *const opts,
                        __global ircd_gpt_vectorv *const restrict norm_bias,
                        __global ircd_gpt_vectorv *const restrict norm_bias_m0,
                        __global ircd_gpt_vectorv *const restrict norm_bias_m1,
                        __global ircd_gpt_vectorv *const restrict norm_weight,
                        __global ircd_gpt_vectorv *const restrict norm_weight_m0,
                        __global ircd_gpt_vectorv *const restrict norm_weight_m1,
                        __global ircd_gpt_ffnn_aperaturev *const restrict fcon_bias,
                        __global ircd_gpt_ffnn_aperaturev *const restrict fcon_bias_m0,
                        __global ircd_gpt_ffnn_aperaturev *const restrict fcon_bias_m1,
                        __global ircd_gpt_ffnn_aperaturev *const restrict fcon_weight,
                        __global ircd_gpt_ffnn_aperaturev *const restrict fcon_weight_m0,
                        __global ircd_gpt_ffnn_aperaturev *const restrict fcon_weight_m1,
                        __global ircd_gpt_vectorv *const restrict proj_bias,
                        __global ircd_gpt_vectorv *const restrict proj_bias_m0,
                        __global ircd_gpt_vectorv *const restrict proj_bias_m1,
                        __global ircd_gpt_vectorv *const restrict proj_weight,
                        __global ircd_gpt_vectorv *const restrict proj_weight_m0,
                        __global ircd_gpt_vectorv *const restrict proj_weight_m1)
{
	const uint
	fcon_height = opts->embed_elems,
	proj_height = opts->ffnn_elems,
	segs = 4;

	ircd_gpt_norm_prop
	(
		ctrl, opts,
		norm_bias,
		norm_bias_m0,
		norm_bias_m1,
		norm_weight,
		norm_weight_m0,
		norm_weight_m1
	);

	for(uint j = 0; j < segs; ++j)
		ircd_gpt_prop_elem
		(
			ctrl, opts,
			fcon_bias->proj[j],
			fcon_bias_m0->proj[j],
			fcon_bias_m1->proj[j]
		);

	for(uint i = 0; i < fcon_height; ++i)
		for(uint j = 0; j < segs; ++j)
			ircd_gpt_prop_elem
			(
				ctrl, opts,
				fcon_weight[i].proj[j],
				fcon_weight_m0[i].proj[j],
				fcon_weight_m1[i].proj[j]
			);

	ircd_gpt_prop_elem
	(
		ctrl, opts,
		proj_bias->elem,
		proj_bias_m0->elem,
		proj_bias_m1->elem
	);

	for(uint i = 0; i < proj_height; ++i)
		ircd_gpt_prop_elem
		(
			ctrl, opts,
			proj_weight[i].elem,
			proj_weight_m0[i].elem,
			proj_weight_m1[i].elem
		);
}

__kernel void
ircd_gpt_lm_embed_prop(__global const struct ircd_gpt_ctrl *const ctrl,
                       __constant const struct ircd_gpt_opts *const opts,
                       __global ircd_gpt_vectorv *const restrict pos,
                       __global ircd_gpt_vectorv *const restrict pos_m0,
                       __global ircd_gpt_vectorv *const restrict pos_m1,
                       __global ircd_gpt_vectorv *const restrict token,
                       __global ircd_gpt_vectorv *const restrict token_m0,
                       __global ircd_gpt_vectorv *const restrict token_m1)
{
	const uint
	ln = get_local_size(0),
	wi = get_global_offset(0) / ln + get_group_id(0),
	wn = ctrl->count,
	cn = opts->context_tokens / wn,
	ci = cn * wi,
	tn = opts->logits / wn,
	ti = tn * wi;

	for(uint i = ci; i < ci + cn; ++i)
		ircd_gpt_prop_elem
		(
			ctrl, opts,
			pos[i].elem,
			pos_m0[i].elem,
			pos_m1[i].elem
		);

	for(uint i = ti; i < ti + tn; ++i)
		ircd_gpt_prop_elem
		(
			ctrl, opts,
			token[i].elem,
			token_m0[i].elem,
			token_m1[i].elem
		);
}

/// Gaussian Error Linear Unit
void
ircd_gpt_ffnn_gelu(__local ircd_gpt_ffnn_aperaturev *const out,
                   __local const ircd_gpt_ffnn_aperaturev *const in_,
                   const uint i)
{
	const float4
	in = in_->fcon[i];

	float4 a;
	a = 0.044715f;
	a *= in;
	a *= in;
	a += 1.0f;
	a *= 0.7978845608f;
	a *= in;

	a = tanh(a);
	a += 1.0f;
	a *= in;
	a *= 0.5f;

	out->fcon[i] = a;
}

void
ircd_gpt_norm(__local ircd_gpt_vectorv *const out,
              __local const ircd_gpt_vectorv *const in,
              __local ircd_gpt_vectorv *const restrict tmp,
              __global const ircd_gpt_vectorv *const restrict bias,
              __global const ircd_gpt_vectorv *const restrict weight,
              const uint ln,
              const uint li)
{
	// Layer re-normalization
	ircd_simt_math_norm_f4lldr(out->elem, in->elem, tmp->elem, ln, li);
	ircd_gpt_norm_fmad(out, out, bias, weight, li);
}

void
ircd_gpt_norm_fmad(__local ircd_gpt_vectorv *const out,
                   __local const ircd_gpt_vectorv *const in,
                   __global const ircd_gpt_vectorv *const restrict bias,
                   __global const ircd_gpt_vectorv *const restrict weight,
                   const uint i)
{
	out->elem[i] = in->elem[i] * weight->elem[i] + bias->elem[i];
}