oarust 0.8.6

GPU-first semantic computing engine written in Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

// @generated by tool/gen/fn/generate.py; DO NOT EDIT.
// schema=matrix_blas.json schema_version=1 generator_version=64
// schema_sha256=31a29676780c696e7f6a77f290398587a9e29e612903c5816503e35d3dda8687
//! Matrix BLAS operations, semantic recording and autograd attachment.

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

/// Multiply A by the transpose of B using the OA weight-layout convention.
///
/// `left` has shape `[M, K]`, `right` has shape `[N, K]`, and the
/// returned matrix has shape `[M, N]`. This matches OA's native
/// weight-layout convention and does not transpose storage.
///
/// The call records device work and returns without a host wait.
///
/// # Errors
///
/// Returns an error when either input is not a rank-two FP32 matrix,
/// their K extents differ, they belong to different engines, a dimension
/// exceeds the admitted shader ABI, or allocation or recording fails.
pub fn mat_mul_nt(left: &Matrix, right: &Matrix) -> Result<Matrix> {
	let contract = crate::core::operation::matrix::MAT_MUL_NT;
	let kernel = KernelId::MatrixMatMulNtTiledF32;
	let operation = contract.name();
	let [m, k] = left.shape() else {
		return Err(Error::invalid_argument(format!(
			"{operation} requires rank-two inputs; left is {:?}",
			left.shape()
		)));
	};
	let [n, right_k] = right.shape() else {
		return Err(Error::invalid_argument(format!(
			"{operation} requires rank-two inputs; right is {:?}",
			right.shape()
		)));
	};
	if k != right_k {
		return Err(Error::invalid_argument(format!(
			"{operation} requires equal K extents; left is {:?}, right is {:?}",
			left.shape(),
			right.shape()
		)));
	}
	if left.dtype() != DType::F32 || right.dtype() != DType::F32 {
		return Err(Error::invalid_argument(format!(
			"{operation} requires two F32 matrices; left is {}, right is {}",
			left.dtype().token(),
			right.dtype().token()
		)));
	}
	let engine = left.engine_handle();
	if !engine.same_as(right.engine_handle()) {
		return Err(Error::invalid_argument(format!(
			"{operation} inputs must belong to the same engine"
		)));
	}

	let output_count = m
		.checked_mul(*n)
		.ok_or_else(|| Error::invalid_argument(format!("{operation} output size overflows usize")))?;
	let m = u32::try_from(*m)
		.map_err(|_| Error::invalid_argument(format!("{operation} M extent exceeds u32")))?;
	let n = u32::try_from(*n)
		.map_err(|_| Error::invalid_argument(format!("{operation} N extent exceeds u32")))?;
	let k = u32::try_from(*k)
		.map_err(|_| Error::invalid_argument(format!("{operation} K extent exceeds u32")))?;
	for (label, count) in [
		("left", left.element_count()),
		("right", right.element_count()),
		("output", output_count),
	] {
		u32::try_from(count).map_err(|_| {
			Error::invalid_argument(format!("{operation} {label} element count exceeds u32"))
		})?;
	}

	let output = Matrix::allocate(
		engine,
		vec![m as usize, n as usize],
		output_count,
		DType::F32,
	)?;
	if m != 0 && n != 0 && k != 0 {
		let buffers = [
			BufferBinding::read(left.storage()),
			BufferBinding::read(right.storage()),
			BufferBinding::write(output.storage()),
		];
		let push_constants = [
			PushConstant::U32(m),
			PushConstant::U32(n),
			PushConstant::U32(k),
		];
		let inputs = [left, right];
		let outputs = [&output];
		engine.record_semantic(
			ComputeDispatch {
				kernel,
				buffers: &buffers,
				push_constants: &push_constants,
				workgroups: kernel.output_workgroups(m, n),
			},
			SemanticDispatch {
				contract,
				inputs: &inputs,
				outputs: &outputs,
				attributes: &[],
			},
		)?;
		autograd::record_mat_mul_nt(left, right, &output)?;
	}
	Ok(output)
}

/// Compute the mat_mul_nt adjoint using the admitted device operations.
///
/// This numerical provider does not inspect or mutate a gradient tape.
pub(crate) fn mat_mul_nt_backward(
	left: &Matrix,
	right: &Matrix,
	output_gradient: &Matrix,
) -> Result<(Matrix, Matrix)> {
	let [rows, inner] = left.shape() else {
		return Err(Error::internal("saved MatMulNt left rank changed"));
	};
	let [columns, right_inner] = right.shape() else {
		return Err(Error::internal("saved MatMulNt right rank changed"));
	};
	if inner != right_inner {
		return Err(Error::internal("saved MatMulNt inner extent changed"));
	}
	let output_gradient = output_gradient.reshape([1, *rows, *columns])?;
	let left_batched = left.reshape([1, *rows, *inner])?;
	let right_batched = right.reshape([1, *columns, *inner])?;
	let left_gradient =
		crate::ml::matrix::bmm(&output_gradient, &right_batched)?.reshape(left.shape().to_vec())?;
	let right_gradient =
		crate::ml::matrix::bmm_tn(&output_gradient, &left_batched)?.reshape(right.shape().to_vec())?;
	Ok((left_gradient, right_gradient))
}