hrx-rs 0.10.1

Owned native AMD GPU and NPU execution with Loom compilation
Documentation 
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

use super::super::{
    Access, Engine, NativeSession, Runtime,
    tests::{buffer, mock, runtime_without_workers},
};
use super::*;
use std::{
    sync::{atomic::AtomicUsize, mpsc},
    time::Duration,
};

#[test]
fn finalization_reserves_capacity_and_blocks_native_admission_until_publication() {
    let runtime = runtime_without_workers();
    let graph = mock(
        &runtime,
        Engine::Gpu,
        &buffer(&runtime),
        Access::Write,
        || Ok(()),
    );
    let completion = graph.submit().unwrap();
    let core = &runtime.inner.core;
    let run = next_action(&mut core.state.lock().unwrap(), None).unwrap();
    assert!(matches!(run, Action::Run(..)));
    perform(core, run);
    let finish = next_action(&mut core.state.lock().unwrap(), None).unwrap();
    assert!(matches!(finish, Action::Finish(..)));
    let ticket = {
        let mut state = core.state.lock().unwrap();
        assert!(state.pending.is_empty());
        assert_eq!(state.occupied(), 1);
        assert!(!completion.is_complete());
        let ticket = state.sequence();
        state.native.push_back(ticket);
        assert!(!state.native_ready(ticket));
        ticket
    };
    perform(core, finish);
    assert!(completion.is_complete());
    let mut state = core.state.lock().unwrap();
    assert!(state.native_ready(ticket));
    assert_eq!(state.occupied(), 1);
    state.native.pop_front();
    assert_eq!(state.occupied(), 0);
}

#[test]
fn native_waiter_runs_after_prior_submissions_before_later_compute() {
    let runtime = Runtime::new().unwrap();
    let order = Arc::new(AtomicUsize::new(0));
    let (entered, entering) = mpsc::channel();
    let (release, released) = mpsc::channel();
    let released = Mutex::new(released);
    let first = mock(&runtime, Engine::Gpu, &buffer(&runtime), Access::Write, {
        let order = order.clone();
        move || {
            assert_eq!(order.fetch_add(1, Ordering::SeqCst), 0);
            entered.send(()).unwrap();
            released
                .lock()
                .unwrap()
                .recv_timeout(Duration::from_secs(10))
                .unwrap();
            Ok(())
        }
    })
    .submit()
    .unwrap();
    entering.recv_timeout(Duration::from_secs(10)).unwrap();
    let (native_entered, native_entering) = mpsc::channel();
    let (native_release, native_released) = mpsc::channel();
    let native = std::thread::spawn({
        let runtime = runtime.clone();
        let order = order.clone();
        move || {
            // SAFETY: fixture performs no native GPU operations.
            let mut session = unsafe { NativeSession::new(&runtime, (), |_| Ok(())) };
            unsafe {
                session.run(|_| {
                    assert_eq!(order.fetch_add(1, Ordering::SeqCst), 1);
                    native_entered.send(()).unwrap();
                    native_released
                        .recv_timeout(Duration::from_secs(10))
                        .unwrap();
                })
            }
            .unwrap();
        }
    });
    let core = &runtime.inner.core;
    let state = core.state.lock().unwrap();
    let (state, timeout) = core
        .changed
        .wait_timeout_while(state, Duration::from_secs(10), |state| {
            state.native.is_empty()
        })
        .unwrap();
    assert!(!timeout.timed_out());
    assert_eq!(state.occupied(), 2);
    drop(state);
    let later = mock(&runtime, Engine::Gpu, &buffer(&runtime), Access::Write, {
        let order = order.clone();
        move || {
            assert_eq!(order.fetch_add(1, Ordering::SeqCst), 2);
            Ok(())
        }
    })
    .submit()
    .unwrap();
    release.send(()).unwrap();
    native_entering
        .recv_timeout(Duration::from_secs(10))
        .unwrap();
    assert!(first.is_complete());
    assert!(!later.is_complete());
    native_release.send(()).unwrap();
    native.join().unwrap();
    later.wait().unwrap();
    assert_eq!(order.load(Ordering::SeqCst), 3);
}