oarust 0.8.6

//! ML-owned Matrix RNN operations and adjoints.

use crate::{
	DType, Error, Matrix, OpAttribute, Result,
	runtime::{BufferBinding, ComputeDispatch, KernelId, PushConstant, SemanticDispatch},
};

use crate::ml::{Parameter, autograd};

use super::super::validation::{shader_u32, validate_f32_same_engine};

use super::recurrent::{MAX_HIDDEN_SIZE, optional_bias};

/// Apply one Elman RNN cell update to a caller-supplied hidden state.
///
/// `gates_i` and `hidden` have shape `[batch, hidden]`, `weight_hh` has
/// `[hidden, hidden]`, and an optional recurrent bias has `[hidden]`.
///
/// # Errors
///
/// Returns an error unless all values are nonempty same-engine F32 matrices,
/// shapes agree, hidden size is at most 1024, or runtime recording fails.
pub fn rnn_cell(
	gates_i: &Matrix,
	hidden: &Matrix,
	weight_hh: &Matrix,
	bias_hh: Option<&Matrix>,
) -> Result<Matrix> {
	let bias_value = optional_bias(weight_hh, bias_hh, "oa::ml::matrix::rnn_cell")?;
	let has_bias = bias_hh.is_some();
	let (output, gates_h) = rnn_cell_dispatch(gates_i, hidden, weight_hh, &bias_value, has_bias)?;
	autograd::record_rnn_cell(
		gates_i,
		hidden,
		&output,
		gates_h,
		None,
		weight_hh.clone(),
		None,
		bias_value,
		has_bias,
	)?;
	Ok(output)
}

#[allow(
	clippy::too_many_arguments,
	reason = "module forwarding preserves optional parameter identity separately from Matrix values"
)]
pub(in crate::ml) fn rnn_cell_parameterized(
	gates_i: &Matrix,
	hidden: &Matrix,
	weight: (Parameter, Matrix, u64),
	bias: Option<(Parameter, u64)>,
	bias_value: &Matrix,
	has_bias: bool,
) -> Result<Matrix> {
	let (weight_parameter, weight_value, weight_version) = weight;
	let (output, gates_h) = rnn_cell_dispatch(gates_i, hidden, &weight_value, bias_value, has_bias)?;
	autograd::record_rnn_cell(
		gates_i,
		hidden,
		&output,
		gates_h,
		Some((weight_parameter, weight_version)),
		weight_value,
		bias,
		bias_value.clone(),
		has_bias,
	)?;
	Ok(output)
}

pub(in crate::ml) fn rnn_cell_backward(
	gates_i: &Matrix,
	gates_h: &Matrix,
	hidden: &Matrix,
	output_gradient: &Matrix,
	weight_hh: &Matrix,
	bias_hh: &Matrix,
	has_bias: bool,
) -> Result<(Matrix, Matrix, Matrix, Matrix)> {
	const OPERATION: &str = crate::core::operation::ml::RNN_CELL_BACKWARD.name();
	let [batch, hidden_size] = gates_i.shape() else {
		return Err(Error::invalid_argument(format!(
			"{OPERATION} input gates must have shape [B, H]"
		)));
	};
	if *batch == 0
		|| *hidden_size == 0
		|| *hidden_size > MAX_HIDDEN_SIZE
		|| gates_h.shape() != gates_i.shape()
		|| hidden.shape() != gates_i.shape()
		|| output_gradient.shape() != gates_i.shape()
		|| weight_hh.shape() != [*hidden_size, *hidden_size]
		|| bias_hh.shape() != [*hidden_size]
	{
		return Err(Error::invalid_argument(format!(
			"{OPERATION} saved values and output gradient do not match the RNN cell contract"
		)));
	}
	validate_f32_same_engine(
		OPERATION,
		&[
			gates_i,
			gates_h,
			hidden,
			output_gradient,
			weight_hh,
			bias_hh,
		],
	)?;
	let gates_i_gradient = Matrix::allocate(
		gates_i.engine_handle(),
		gates_i.shape().to_vec(),
		gates_i.num_elements(),
		DType::F32,
	)?;
	let gates_h_gradient = Matrix::allocate(
		gates_i.engine_handle(),
		gates_i.shape().to_vec(),
		gates_i.num_elements(),
		DType::F32,
	)?;
	let hidden_gradient = Matrix::allocate(
		gates_i.engine_handle(),
		hidden.shape().to_vec(),
		hidden.num_elements(),
		DType::F32,
	)?;
	let weight_gradient = Matrix::allocate(
		gates_i.engine_handle(),
		weight_hh.shape().to_vec(),
		weight_hh.num_elements(),
		DType::F32,
	)?;
	let bias_gradient = Matrix::allocate(
		gates_i.engine_handle(),
		bias_hh.shape().to_vec(),
		bias_hh.num_elements(),
		DType::F32,
	)?;
	let batch = shader_u32(*batch, "batch", OPERATION)?;
	let hidden_size = shader_u32(*hidden_size, "hidden size", OPERATION)?;
	let element_count = shader_u32(gates_i.num_elements(), "element count", OPERATION)?;
	let pointwise_buffers = [
		BufferBinding::read(gates_i.storage()),
		BufferBinding::read(gates_h.storage()),
		BufferBinding::read(output_gradient.storage()),
		BufferBinding::write(gates_i_gradient.storage()),
		BufferBinding::write(gates_h_gradient.storage()),
	];
	let pointwise_push_constants = [PushConstant::U32(element_count)];
	let input_buffers = [
		BufferBinding::read(weight_hh.storage()),
		BufferBinding::read(gates_h_gradient.storage()),
		BufferBinding::write(hidden_gradient.storage()),
	];
	let linear_push_constants = [
		PushConstant::U32(batch),
		PushConstant::U32(hidden_size),
		PushConstant::U32(hidden_size),
	];
	let parameter_buffers = [
		BufferBinding::read(hidden.storage()),
		BufferBinding::read(gates_h_gradient.storage()),
		BufferBinding::write(weight_gradient.storage()),
		BufferBinding::write(bias_gradient.storage()),
	];
	let dispatches = [
		ComputeDispatch {
			kernel: KernelId::MlRnnCellBackwardF32,
			buffers: &pointwise_buffers,
			push_constants: &pointwise_push_constants,
			workgroups: KernelId::MlRnnCellBackwardF32.linear_workgroups(element_count),
		},
		ComputeDispatch {
			kernel: KernelId::MlLinearBackwardF32,
			buffers: &input_buffers,
			push_constants: &linear_push_constants,
			workgroups: KernelId::MlLinearBackwardF32.output_workgroups(batch, hidden_size),
		},
		ComputeDispatch {
			kernel: KernelId::MlLinearParameterBackwardF32,
			buffers: &parameter_buffers,
			push_constants: &linear_push_constants,
			workgroups: [hidden_size, hidden_size.div_ceil(32), 1],
		},
	];
	let inputs = [
		gates_i,
		gates_h,
		hidden,
		output_gradient,
		weight_hh,
		bias_hh,
	];
	let outputs = [
		&gates_i_gradient,
		&hidden_gradient,
		&weight_gradient,
		&bias_gradient,
	];
	let attributes = [OpAttribute::Boolean {
		name: "has_bias".into(),
		value: has_bias,
	}];
	gates_i.engine_handle().record_split_semantic(
		&dispatches,
		SemanticDispatch {
			contract: crate::core::operation::ml::RNN_CELL_BACKWARD,
			inputs: &inputs,
			outputs: &outputs,
			attributes: &attributes,
		},
	)?;
	Ok((
		gates_i_gradient,
		hidden_gradient,
		weight_gradient,
		bias_gradient,
	))
}

/// Run a whole-sequence Elman recurrence from a zero hidden state.
///
/// `gates_i` has shape `[batch, sequence, hidden]`, `weight_hh` has
/// `[hidden, hidden]`, and an optional recurrent bias has `[hidden]`.
///
/// # Errors
///
/// Returns an error unless all values are nonempty same-engine F32 matrices,
/// shapes agree, hidden size is at most 1024, or runtime recording fails.
pub fn rnn_scan(gates_i: &Matrix, weight_hh: &Matrix, bias_hh: Option<&Matrix>) -> Result<Matrix> {
	let bias_value = optional_bias(weight_hh, bias_hh, "oa::ml::matrix::rnn_scan")?;
	let has_bias = bias_hh.is_some();
	let (output, hidden_previous) = rnn_scan_dispatch(gates_i, weight_hh, &bias_value, has_bias)?;
	autograd::record_rnn_scan(
		gates_i,
		&output,
		hidden_previous,
		None,
		weight_hh.clone(),
		None,
		bias_value,
		has_bias,
	)?;
	Ok(output)
}

pub(in crate::ml) fn rnn_scan_parameterized(
	gates_i: &Matrix,
	weight: (Parameter, Matrix, u64),
	bias: Option<(Parameter, u64)>,
	bias_value: &Matrix,
	has_bias: bool,
) -> Result<Matrix> {
	let (weight_parameter, weight_value, weight_version) = weight;
	let (output, hidden_previous) = rnn_scan_dispatch(gates_i, &weight_value, bias_value, has_bias)?;
	autograd::record_rnn_scan(
		gates_i,
		&output,
		hidden_previous,
		Some((weight_parameter, weight_version)),
		weight_value,
		bias,
		bias_value.clone(),
		has_bias,
	)?;
	Ok(output)
}

pub(in crate::ml) fn rnn_scan_backward(
	output_gradient: &Matrix,
	gates_i: &Matrix,
	hidden_previous: &Matrix,
	weight_hh: &Matrix,
	bias_hh: &Matrix,
	has_bias: bool,
) -> Result<(Matrix, Matrix, Matrix)> {
	const OPERATION: &str = crate::core::operation::ml::RNN_SCAN_BACKWARD.name();
	let [batch, sequence_length, hidden_size] = gates_i.shape() else {
		return Err(Error::invalid_argument(format!(
			"{OPERATION} input gates must have shape [B, S, H]"
		)));
	};
	if *batch == 0
		|| *sequence_length == 0
		|| *hidden_size == 0
		|| *hidden_size > MAX_HIDDEN_SIZE
		|| hidden_previous.shape() != gates_i.shape()
		|| output_gradient.shape() != gates_i.shape()
		|| weight_hh.shape() != [*hidden_size, *hidden_size]
		|| bias_hh.shape() != [*hidden_size]
	{
		return Err(Error::invalid_argument(format!(
			"{OPERATION} saved values and output gradient do not match the RNN scan contract"
		)));
	}
	validate_f32_same_engine(
		OPERATION,
		&[
			output_gradient,
			gates_i,
			hidden_previous,
			weight_hh,
			bias_hh,
		],
	)?;
	let gates_i_gradient = Matrix::allocate(
		gates_i.engine_handle(),
		gates_i.shape().to_vec(),
		gates_i.num_elements(),
		DType::F32,
	)?;
	let gates_h_gradient = Matrix::allocate(
		gates_i.engine_handle(),
		gates_i.shape().to_vec(),
		gates_i.num_elements(),
		DType::F32,
	)?;
	let weight_gradient = Matrix::allocate(
		gates_i.engine_handle(),
		weight_hh.shape().to_vec(),
		weight_hh.num_elements(),
		DType::F32,
	)?;
	let bias_gradient = Matrix::allocate(
		gates_i.engine_handle(),
		bias_hh.shape().to_vec(),
		bias_hh.num_elements(),
		DType::F32,
	)?;
	let rows = batch
		.checked_mul(*sequence_length)
		.ok_or_else(|| Error::invalid_argument(format!("{OPERATION} row count overflows usize")))?;
	let batch = shader_u32(*batch, "batch", OPERATION)?;
	let sequence_length = shader_u32(*sequence_length, "sequence length", OPERATION)?;
	let hidden_size = shader_u32(*hidden_size, "hidden size", OPERATION)?;
	let rows = shader_u32(rows, "flattened row count", OPERATION)?;
	let scan_buffers = [
		BufferBinding::read(output_gradient.storage()),
		BufferBinding::read(gates_i.storage()),
		BufferBinding::read(hidden_previous.storage()),
		BufferBinding::read(weight_hh.storage()),
		BufferBinding::read(bias_hh.storage()),
		BufferBinding::write(gates_i_gradient.storage()),
		BufferBinding::write(gates_h_gradient.storage()),
	];
	let scan_push_constants = [
		PushConstant::U32(hidden_size),
		PushConstant::U32(sequence_length),
		PushConstant::U32(batch),
		PushConstant::U32(u32::from(has_bias)),
	];
	let parameter_buffers = [
		BufferBinding::read(hidden_previous.storage()),
		BufferBinding::read(gates_h_gradient.storage()),
		BufferBinding::write(weight_gradient.storage()),
		BufferBinding::write(bias_gradient.storage()),
	];
	let parameter_push_constants = [
		PushConstant::U32(rows),
		PushConstant::U32(hidden_size),
		PushConstant::U32(hidden_size),
	];
	let dispatches = [
		ComputeDispatch {
			kernel: KernelId::MlRnnScanBackwardF32,
			buffers: &scan_buffers,
			push_constants: &scan_push_constants,
			workgroups: [batch, 1, 1],
		},
		ComputeDispatch {
			kernel: KernelId::MlLinearParameterBackwardF32,
			buffers: &parameter_buffers,
			push_constants: &parameter_push_constants,
			workgroups: [hidden_size, hidden_size.div_ceil(32), 1],
		},
	];
	let inputs = [
		output_gradient,
		gates_i,
		hidden_previous,
		weight_hh,
		bias_hh,
	];
	let outputs = [&gates_i_gradient, &weight_gradient, &bias_gradient];
	let attributes = [OpAttribute::Boolean {
		name: "has_bias".into(),
		value: has_bias,
	}];
	gates_i.engine_handle().record_split_semantic(
		&dispatches,
		SemanticDispatch {
			contract: crate::core::operation::ml::RNN_SCAN_BACKWARD,
			inputs: &inputs,
			outputs: &outputs,
			attributes: &attributes,
		},
	)?;
	Ok((gates_i_gradient, weight_gradient, bias_gradient))
}

fn rnn_cell_dispatch(
	gates_i: &Matrix,
	hidden: &Matrix,
	weight_hh: &Matrix,
	bias_hh: &Matrix,
	has_bias: bool,
) -> Result<(Matrix, Matrix)> {
	const OPERATION: &str = crate::core::operation::ml::RNN_CELL.name();
	let [batch, hidden_size] = gates_i.shape() else {
		return Err(Error::invalid_argument(format!(
			"{OPERATION} input gates must have shape [B, H]"
		)));
	};
	if *batch == 0
		|| *hidden_size == 0
		|| *hidden_size > MAX_HIDDEN_SIZE
		|| hidden.shape() != gates_i.shape()
		|| weight_hh.shape() != [*hidden_size, *hidden_size]
		|| bias_hh.shape() != [*hidden_size]
	{
		return Err(Error::invalid_argument(format!(
			"{OPERATION} requires gates and hidden [B, H], recurrent weight [H, H], bias storage [H], and H <= {MAX_HIDDEN_SIZE}"
		)));
	}
	validate_f32_same_engine(OPERATION, &[gates_i, hidden, weight_hh, bias_hh])?;
	let output = Matrix::allocate(
		gates_i.engine_handle(),
		gates_i.shape().to_vec(),
		gates_i.num_elements(),
		DType::F32,
	)?;
	let gates_h = Matrix::allocate(
		gates_i.engine_handle(),
		gates_i.shape().to_vec(),
		gates_i.num_elements(),
		DType::F32,
	)?;
	let batch = shader_u32(*batch, "batch", OPERATION)?;
	let hidden_size = shader_u32(*hidden_size, "hidden size", OPERATION)?;
	let buffers = [
		BufferBinding::read(hidden.storage()),
		BufferBinding::read(weight_hh.storage()),
		BufferBinding::read(bias_hh.storage()),
		BufferBinding::read(gates_i.storage()),
		BufferBinding::write(output.storage()),
		BufferBinding::write(gates_h.storage()),
	];
	let push_constants = [
		PushConstant::U32(hidden_size),
		PushConstant::U32(batch),
		PushConstant::U32(u32::from(has_bias)),
	];
	let attributes = [OpAttribute::Boolean {
		name: "has_bias".into(),
		value: has_bias,
	}];
	{
		let inputs: &[&Matrix] = &[gates_i, hidden, weight_hh, bias_hh];
		let outputs: &[&Matrix] = &[&output, &gates_h];
		let attributes: &[OpAttribute] = &attributes;
		let dispatch = ComputeDispatch {
			kernel: KernelId::MlRnnCellF32,
			buffers: &buffers,
			push_constants: &push_constants,
			workgroups: [batch, 1, 1],
		};
		let contract = dispatch.kernel.semantic_contract().ok_or_else(|| {
			Error::internal(format!(
				"lowering-only kernel {} cannot own one semantic ML operation",
				dispatch.kernel.report_name()
			))
		})?;
		let engine = inputs
			.first()
			.ok_or_else(|| Error::internal("semantic ML dispatch requires an input"))?
			.engine_handle();
		engine.record_semantic(
			dispatch,
			SemanticDispatch {
				contract,
				inputs,
				outputs,
				attributes,
			},
		)
	}?;
	Ok((output, gates_h))
}

fn rnn_scan_dispatch(
	gates_i: &Matrix,
	weight_hh: &Matrix,
	bias_hh: &Matrix,
	has_bias: bool,
) -> Result<(Matrix, Matrix)> {
	const OPERATION: &str = crate::core::operation::ml::RNN_SCAN.name();
	let [batch, sequence_length, hidden_size] = gates_i.shape() else {
		return Err(Error::invalid_argument(format!(
			"{OPERATION} input gates must have shape [B, S, H]"
		)));
	};
	if *batch == 0
		|| *sequence_length == 0
		|| *hidden_size == 0
		|| *hidden_size > MAX_HIDDEN_SIZE
		|| weight_hh.shape() != [*hidden_size, *hidden_size]
		|| bias_hh.shape() != [*hidden_size]
	{
		return Err(Error::invalid_argument(format!(
			"{OPERATION} requires nonempty gates [B, S, H], recurrent weight [H, H], bias storage [H], and H <= {MAX_HIDDEN_SIZE}"
		)));
	}
	validate_f32_same_engine(OPERATION, &[gates_i, weight_hh, bias_hh])?;
	let output = Matrix::allocate(
		gates_i.engine_handle(),
		gates_i.shape().to_vec(),
		gates_i.num_elements(),
		DType::F32,
	)?;
	let hidden_previous = Matrix::allocate(
		gates_i.engine_handle(),
		gates_i.shape().to_vec(),
		gates_i.num_elements(),
		DType::F32,
	)?;
	let batch = shader_u32(*batch, "batch", OPERATION)?;
	let sequence_length = shader_u32(*sequence_length, "sequence length", OPERATION)?;
	let hidden_size = shader_u32(*hidden_size, "hidden size", OPERATION)?;
	let buffers = [
		BufferBinding::read(gates_i.storage()),
		BufferBinding::read(weight_hh.storage()),
		BufferBinding::read(bias_hh.storage()),
		BufferBinding::write(output.storage()),
		BufferBinding::write(hidden_previous.storage()),
	];
	let push_constants = [
		PushConstant::U32(hidden_size),
		PushConstant::U32(sequence_length),
		PushConstant::U32(batch),
		PushConstant::U32(u32::from(has_bias)),
	];
	let attributes = [OpAttribute::Boolean {
		name: "has_bias".into(),
		value: has_bias,
	}];
	{
		let inputs: &[&Matrix] = &[gates_i, weight_hh, bias_hh];
		let outputs: &[&Matrix] = &[&output, &hidden_previous];
		let attributes: &[OpAttribute] = &attributes;
		let dispatch = ComputeDispatch {
			kernel: KernelId::MlRnnScanF32,
			buffers: &buffers,
			push_constants: &push_constants,
			workgroups: [batch, 1, 1],
		};
		let contract = dispatch.kernel.semantic_contract().ok_or_else(|| {
			Error::internal(format!(
				"lowering-only kernel {} cannot own one semantic ML operation",
				dispatch.kernel.report_name()
			))
		})?;
		let engine = inputs
			.first()
			.ok_or_else(|| Error::internal("semantic ML dispatch requires an input"))?
			.engine_handle();
		engine.record_semantic(
			dispatch,
			SemanticDispatch {
				contract,
				inputs,
				outputs,
				attributes,
			},
		)
	}?;
	Ok((output, hidden_previous))
}