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

//! Private bridge from Matrix operations into an active ML tape.

use std::{cell::RefCell, marker::PhantomData, rc::Rc};

use crate::Result;

#[path = "autograd/matrix.gen.rs"]
mod matrix;

#[path = "autograd/node.gen.rs"]
mod node;

include!("autograd/exports.gen.rs");
#[path = "autograd/context.rs"]
mod context;

pub(crate) use context::GradientContext;
pub(crate) use node::GradNodeMatrix;

type Recorder = Rc<dyn Fn(GradNodeMatrix) -> Result<()>>;

thread_local! {
	static RECORDERS: RefCell<Vec<Recorder>> = const { RefCell::new(Vec::new()) };
}

/// Selection of one thread-local observer, removed exactly when dropped.
pub(crate) struct RecordingSelection {
	recorder: Recorder,
	_not_send_sync: PhantomData<Rc<()>>,
}

pub(crate) fn select(
	record: impl Fn(GradNodeMatrix) -> Result<()> + 'static,
) -> RecordingSelection {
	let recorder: Recorder = Rc::new(record);
	RECORDERS.with(|recorders| recorders.borrow_mut().push(recorder.clone()));
	RecordingSelection {
		recorder,
		_not_send_sync: PhantomData,
	}
}

pub(crate) fn record(node: GradNodeMatrix) -> Result<()> {
	let recorder = RECORDERS.with(|recorders| recorders.borrow().last().cloned());
	if let Some(recorder) = recorder {
		return recorder(node);
	}
	Ok(())
}

impl Drop for RecordingSelection {
	fn drop(&mut self) {
		RECORDERS.with(|recorders| {
			recorders
				.borrow_mut()
				.retain(|candidate| !Rc::ptr_eq(candidate, &self.recorder));
		});
	}
}