oarust 0.8.6

// @generated by tool/gen/fn/generate.py; DO NOT EDIT.
// schema=ml_activation.json schema_version=1 generator_version=64
// schema_sha256=cfadfc3788fb5454d476c57f122b9c5913e4ba30cb3c3eb48ec41ae018f0ec89
//! ML activation Matrix operations — elementwise forward and backward passes.

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

/// Apply SiLU(gate) * up elementwise.
///
/// # Errors
///
/// Returns an error unless both inputs are equal-shape F32 on the same engine,
/// or runtime recording fails.
pub fn swiglu(gate: &Matrix, up: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::ml::SWIGLU;
	let operation = contract.name();
	let engine = gate.engine_handle();
	for matrix in &[gate, up] {
		if matrix.dtype() != DType::F32 {
			return Err(Error::invalid_argument(format!(
				"{operation} requires F32 matrices; found {}",
				matrix.dtype().token()
			)));
		}
		if !engine.same_as(matrix.engine_handle()) {
			return Err(Error::invalid_argument(format!(
				"{operation} inputs must belong to the same engine"
			)));
		}
	}
	if gate.shape() != up.shape() {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal input shapes"
		)));
	}
	let element_count = u32::try_from(gate.num_elements())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let output = Matrix::allocate(
		engine,
		gate.shape().to_vec(),
		gate.num_elements(),
		DType::F32,
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(gate.storage()),
			BufferBinding::read(up.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		engine.record_semantic(
			ComputeDispatch {
				kernel: KernelId::MlSwigluF32,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: KernelId::MlSwigluF32.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &[gate, up],
				outputs: &[&output],
				attributes: &[],
			},
		)?;
	}
	autograd::record_swiglu(gate, up, &output)?;
	Ok(output)
}

pub(in crate::ml) fn swiglu_backward(
	gate: &Matrix,
	up: &Matrix,
	output_gradient: &Matrix,
) -> Result<(Matrix, Matrix)> {
	let contract = crate::core::operation::ml::SWIGLU_BACKWARD;
	let operation = contract.name();
	let engine = gate.engine_handle();
	for matrix in &[gate, up, output_gradient] {
		if matrix.dtype() != DType::F32 {
			return Err(Error::invalid_argument(format!(
				"{operation} requires F32 matrices; found {}",
				matrix.dtype().token()
			)));
		}
		if !engine.same_as(matrix.engine_handle()) {
			return Err(Error::invalid_argument(format!(
				"{operation} inputs must belong to the same engine"
			)));
		}
	}
	if gate.shape() != up.shape() || gate.shape() != output_gradient.shape() {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal input and gradient shapes"
		)));
	}
	let element_count = u32::try_from(gate.num_elements())
		.map_err(|_| Error::invalid_argument(format!("{operation} element count exceeds u32")))?;
	let gate_gradient = Matrix::allocate(
		engine,
		gate.shape().to_vec(),
		gate.num_elements(),
		DType::F32,
	)?;
	let up_gradient = Matrix::allocate(
		engine,
		gate.shape().to_vec(),
		gate.num_elements(),
		DType::F32,
	)?;
	if element_count != 0 {
		let buffers = [
			BufferBinding::read(gate.storage()),
			BufferBinding::read(up.storage()),
			BufferBinding::read(output_gradient.storage()),
			BufferBinding::write(gate_gradient.storage()),
			BufferBinding::write(up_gradient.storage()),
		];
		let push_constants = [PushConstant::U32(element_count)];
		engine.record_semantic(
			ComputeDispatch {
				kernel: KernelId::MlSwigluBackwardF32,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: KernelId::MlSwigluBackwardF32.linear_workgroups(element_count),
			},
			SemanticDispatch {
				contract,
				inputs: &[gate, up, output_gradient],
				outputs: &[&gate_gradient, &up_gradient],
				attributes: &[],
			},
		)?;
	}
	Ok((gate_gradient, up_gradient))
}

/// Apply SiLU(gate) * up to concatenated halves of the final input axis.
///
/// # Errors
///
/// Returns an error unless `input` is F32 with final extent exactly twice the
/// nonzero `intermediate_size`, or runtime recording fails.
pub fn silu_mul(input: &Matrix, intermediate_size: usize) -> Result<Matrix> {
	let contract = crate::core::operation::ml::SILU_MUL;
	let operation = contract.name();
	let engine = input.engine_handle();
	if input.dtype() != DType::F32 {
		return Err(Error::invalid_argument(format!(
			"{operation} requires F32 matrices; found {}",
			input.dtype().token()
		)));
	}
	let Some(last_extent) = input.shape().last() else {
		return Err(Error::invalid_argument(format!(
			"{operation} input must have at least one dimension"
		)));
	};
	let concatenated = intermediate_size.checked_mul(2).ok_or_else(|| {
		Error::invalid_argument(format!("{operation} intermediate size overflows usize"))
	})?;
	if intermediate_size == 0 || *last_extent != concatenated {
		return Err(Error::invalid_argument(format!(
			"{operation} requires final input extent 2 * intermediate_size"
		)));
	}
	let output_count = input.num_elements() / 2;
	let output_count_u32 = u32::try_from(output_count).map_err(|_| {
		Error::invalid_argument(format!("{operation} output element count exceeds u32"))
	})?;
	let intermediate_u32 = u32::try_from(intermediate_size)
		.map_err(|_| Error::invalid_argument(format!("{operation} intermediate size exceeds u32")))?;
	let mut output_shape = input.shape().to_vec();
	*output_shape.last_mut().expect("validated nonempty shape") = intermediate_size;
	let output = Matrix::allocate(engine, output_shape, output_count, DType::F32)?;
	let buffers = [
		BufferBinding::read(input.storage()),
		BufferBinding::write(output.storage()),
	];
	let attributes = [OpAttribute::SignedInteger {
		name: "intermediate_size".into(),
		value: i64::try_from(intermediate_size)
			.map_err(|_| Error::invalid_argument(format!("{operation} intermediate size exceeds i64")))?,
	}];
	let push_constants = [
		PushConstant::U32(output_count_u32),
		PushConstant::U32(intermediate_u32),
	];
	engine.record_semantic(
		ComputeDispatch {
			kernel: KernelId::MlSiluMulF32,
			buffers: &buffers,
			push_constants: &push_constants,
			workgroups: KernelId::MlSiluMulF32.linear_workgroups(output_count_u32),
		},
		SemanticDispatch {
			contract,
			inputs: &[input],
			outputs: &[&output],
			attributes: &attributes,
		},
	)?;
	autograd::record_silu_mul(input, intermediate_size, &output)?;
	Ok(output)
}

pub(in crate::ml) fn silu_mul_backward(
	input: &Matrix,
	output_gradient: &Matrix,
	intermediate_size: usize,
) -> Result<Matrix> {
	let contract = crate::core::operation::ml::SILU_MUL_BACKWARD;
	let operation = contract.name();
	let engine = input.engine_handle();
	for matrix in &[input, output_gradient] {
		if matrix.dtype() != DType::F32 {
			return Err(Error::invalid_argument(format!(
				"{operation} requires F32 matrices; found {}",
				matrix.dtype().token()
			)));
		}
		if !engine.same_as(matrix.engine_handle()) {
			return Err(Error::invalid_argument(format!(
				"{operation} inputs must belong to the same engine"
			)));
		}
	}
	let Some(last_extent) = input.shape().last() else {
		return Err(Error::invalid_argument(format!(
			"{operation} input must have at least one dimension"
		)));
	};
	let concatenated = intermediate_size.checked_mul(2).ok_or_else(|| {
		Error::invalid_argument(format!("{operation} intermediate size overflows usize"))
	})?;
	let mut expected_gradient_shape = input.shape().to_vec();
	if intermediate_size == 0 || *last_extent != concatenated {
		return Err(Error::invalid_argument(format!(
			"{operation} requires final input extent 2 * intermediate_size"
		)));
	}
	*expected_gradient_shape
		.last_mut()
		.expect("validated nonempty shape") = intermediate_size;
	if output_gradient.shape() != expected_gradient_shape {
		return Err(Error::invalid_argument(format!(
			"{operation} output gradient shape must equal the half-width output shape"
		)));
	}
	let output_count = output_gradient.num_elements();
	let output_count_u32 = u32::try_from(output_count).map_err(|_| {
		Error::invalid_argument(format!("{operation} output element count exceeds u32"))
	})?;
	let intermediate_u32 = u32::try_from(intermediate_size)
		.map_err(|_| Error::invalid_argument(format!("{operation} intermediate size exceeds u32")))?;
	let input_gradient = Matrix::allocate(
		engine,
		input.shape().to_vec(),
		input.num_elements(),
		DType::F32,
	)?;
	let buffers = [
		BufferBinding::read(input.storage()),
		BufferBinding::read(output_gradient.storage()),
		BufferBinding::write(input_gradient.storage()),
	];
	let attributes = [OpAttribute::SignedInteger {
		name: "intermediate_size".into(),
		value: i64::try_from(intermediate_size)
			.map_err(|_| Error::invalid_argument(format!("{operation} intermediate size exceeds i64")))?,
	}];
	let push_constants = [
		PushConstant::U32(output_count_u32),
		PushConstant::U32(intermediate_u32),
	];
	engine.record_semantic(
		ComputeDispatch {
			kernel: KernelId::MlSiluMulBackwardF32,
			buffers: &buffers,
			push_constants: &push_constants,
			workgroups: KernelId::MlSiluMulBackwardF32.linear_workgroups(output_count_u32),
		},
		SemanticDispatch {
			contract,
			inputs: &[input, output_gradient],
			outputs: &[&input_gradient],
			attributes: &attributes,
		},
	)?;
	Ok(input_gradient)
}