oarust 0.8.6

// @generated by tool/gen/fn/generate.py; DO NOT EDIT.
// schema=ml_training.json schema_version=1 generator_version=64
// schema_sha256=6678961ee57f24310158676c7b4fe0821e21ca48db7edcb20d2a749a164976ae
//! Rotary position embedding Matrix operations.

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

fn validate_rope(
	input: &Matrix,
	num_heads: usize,
	head_dim: usize,
	theta_base: f32,
	operation: &'static str,
) -> Result<(u32, u32, u32, u32)> {
	let [tokens, width] = input.shape() else {
		return Err(Error::invalid_argument(format!(
			"{operation} input must have shape [tokens, num_heads * head_dim]"
		)));
	};
	let expected_width = num_heads
		.checked_mul(head_dim)
		.ok_or_else(|| Error::invalid_argument(format!("{operation} width overflows usize")))?;
	if *tokens == 0
		|| num_heads == 0
		|| head_dim == 0
		|| !head_dim.is_multiple_of(2)
		|| *width != expected_width
	{
		return Err(Error::invalid_argument(format!(
			"{operation} requires nonempty [tokens, num_heads * even head_dim] input"
		)));
	}
	if !theta_base.is_finite() || theta_base <= 0.0 {
		return Err(Error::invalid_argument(format!(
			"{operation} theta base must be finite and positive"
		)));
	}
	validate_f32_same_engine(operation, &[input])?;
	let pair_count = tokens
		.checked_mul(num_heads)
		.and_then(|count| count.checked_mul(head_dim / 2))
		.ok_or_else(|| Error::invalid_argument(format!("{operation} pair count overflows usize")))?;
	Ok((
		shader_u32(*tokens, "token count", operation)?,
		shader_u32(num_heads, "head count", operation)?,
		shader_u32(head_dim, "head dimension", operation)?,
		shader_u32(pair_count, "rotary pair count", operation)?,
	))
}

/// Apply split-half rotary position embeddings to `[tokens, heads * head_dim]`.
///
/// `position_offset` selects the first absolute position, making the same
/// operation usable for full sequences and future KV-cache extension.
///
/// # Errors
///
/// Returns an error unless the input is nonempty F32, `head_dim` is even, the
/// final width matches `num_heads * head_dim`, and `theta_base` is finite and
/// positive.
pub fn rope(
	input: &Matrix,
	num_heads: usize,
	head_dim: usize,
	theta_base: f32,
	position_offset: u32,
) -> Result<Matrix> {
	const OPERATION: &str = crate::core::operation::ml::ROPE.name();
	let (tokens, num_heads_u32, head_dim_u32, pair_count) =
		validate_rope(input, num_heads, head_dim, theta_base, OPERATION)?;
	let output = Matrix::allocate(
		input.engine_handle(),
		input.shape().to_vec(),
		input.num_elements(),
		DType::F32,
	)?;
	let buffers = [
		BufferBinding::read(input.storage()),
		BufferBinding::write(output.storage()),
	];
	let push_constants = [
		PushConstant::U32(tokens),
		PushConstant::U32(num_heads_u32),
		PushConstant::U32(head_dim_u32),
		PushConstant::F32(theta_base),
		PushConstant::U32(position_offset),
	];
	let attributes = [
		OpAttribute::UnsignedInteger {
			name: "num_heads".into(),
			value: u64::from(num_heads_u32),
		},
		OpAttribute::UnsignedInteger {
			name: "head_dim".into(),
			value: u64::from(head_dim_u32),
		},
		OpAttribute::Float {
			name: "theta_base".into(),
			value: f64::from(theta_base),
		},
		OpAttribute::UnsignedInteger {
			name: "position_offset".into(),
			value: u64::from(position_offset),
		},
	];
	let kernel = KernelId::MlRopeF32;
	{
		let inputs: &[&Matrix] = &[input];
		let outputs: &[&Matrix] = &[&output];
		let attributes: &[OpAttribute] = &attributes;
		let dispatch = ComputeDispatch {
			kernel,
			buffers: &buffers,
			push_constants: &push_constants,
			workgroups: kernel.linear_workgroups(pair_count),
		};
		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,
			},
		)
	}?;
	autograd::record_rope(
		input,
		num_heads,
		head_dim,
		theta_base,
		position_offset,
		&output,
	)?;
	Ok(output)
}

pub(in crate::ml) fn rope_backward(
	output_gradient: &Matrix,
	num_heads: usize,
	head_dim: usize,
	theta_base: f32,
	position_offset: u32,
) -> Result<Matrix> {
	const OPERATION: &str = crate::core::operation::ml::ROPE_BACKWARD.name();
	let (tokens, num_heads, head_dim, pair_count) =
		validate_rope(output_gradient, num_heads, head_dim, theta_base, OPERATION)?;
	let input_gradient = Matrix::allocate(
		output_gradient.engine_handle(),
		output_gradient.shape().to_vec(),
		output_gradient.num_elements(),
		DType::F32,
	)?;
	let buffers = [
		BufferBinding::read(output_gradient.storage()),
		BufferBinding::write(input_gradient.storage()),
	];
	let push_constants = [
		PushConstant::U32(tokens),
		PushConstant::U32(num_heads),
		PushConstant::U32(head_dim),
		PushConstant::F32(theta_base),
		PushConstant::U32(position_offset),
	];
	let attributes = [
		OpAttribute::UnsignedInteger {
			name: "num_heads".into(),
			value: u64::from(num_heads),
		},
		OpAttribute::UnsignedInteger {
			name: "head_dim".into(),
			value: u64::from(head_dim),
		},
		OpAttribute::Float {
			name: "theta_base".into(),
			value: f64::from(theta_base),
		},
		OpAttribute::UnsignedInteger {
			name: "position_offset".into(),
			value: u64::from(position_offset),
		},
	];
	let kernel = KernelId::MlRopeBackwardF32;
	{
		let inputs: &[&Matrix] = &[output_gradient];
		let outputs: &[&Matrix] = &[&input_gradient];
		let attributes: &[OpAttribute] = &attributes;
		let dispatch = ComputeDispatch {
			kernel,
			buffers: &buffers,
			push_constants: &push_constants,
			workgroups: kernel.linear_workgroups(pair_count),
		};
		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(input_gradient)
}