cubecl-core 0.11.0

CubeCL core create
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

//! What a client reads about its stream's graph capture: [`Client::is_capturing`] is true from
//! `graph_prepare` until the capture ends, for every client of the capturing stream and for no
//! other stream.
//!
//! Not part of `testgen_all`: a capture holds its stream for the whole window, so a runtime calls
//! these from its graph tests, one capture at a time.

use crate::{self as cubecl};
use cubecl::prelude::*;
use cubecl_environment::stream::StreamId;
use cubecl_runtime::runtime::Runtime;

#[cube(launch)]
fn add_one(input: &[f32], output: &mut [f32]) {
    if ABSOLUTE_POS < output.len() {
        output[ABSOLUTE_POS] = input[ABSOLUTE_POS] + 1.0;
    }
}

/// A capture is visible to every client of the stream that prepared it, through the warmup run
/// and the recorded one, and to no other stream.
///
/// Code that has to decide the same way in the warmup and the recording reads it before every
/// launch, possibly through a client other than the one that prepared, so each client of the
/// stream has to see it; a neighbour that saw it would take the capture's path for nothing.
pub fn a_capture_is_seen_by_every_client_of_its_stream<R: Runtime>() {
    let client = R::client(&Default::default());
    let other_client = R::client(&Default::default());
    let owner = StreamId { value: 3_000_001 };
    let neighbour = StreamId { value: 3_000_002 };
    let capturing = |client: &Client, stream: StreamId| stream.executes(|| client.is_capturing());

    let n = 4usize;
    let (input, output) = owner.executes(|| {
        (
            client.create_from_slice(f32::as_bytes(&[1.0, 2.0, 3.0, 4.0])),
            client.empty(n * core::mem::size_of::<f32>()),
        )
    });
    let run = || {
        add_one::launch(
            &client,
            CubeCount::Static(1, 1, 1),
            CubeDim::new(&client, n),
            unsafe { BufferArg::from_raw_parts(input.clone(), n) },
            unsafe { BufferArg::from_raw_parts(output.clone(), n) },
        );
    };
    owner.executes(|| {
        run();
        client.read_one(output.clone()).unwrap();
    });

    assert!(!capturing(&client, owner), "no capture before prepare");
    owner
        .executes(|| client.graph_prepare())
        .expect("graph_prepare");
    assert!(capturing(&client, owner), "the warmup run is part of it");
    assert!(
        capturing(&other_client, owner),
        "every client of the stream sees it"
    );
    assert!(
        !capturing(&client, neighbour),
        "another stream isn't capturing"
    );

    owner.executes(|| {
        run();
        client.read_one(output.clone()).unwrap();
        client.start_capture().expect("start_capture");
    });
    assert!(
        capturing(&other_client, owner),
        "the recorded run is part of it"
    );
    assert!(!capturing(&other_client, neighbour));

    let graph = owner.executes(|| {
        run();
        client.stop_capture()
    });
    graph.expect("stop_capture");
    assert!(
        !capturing(&client, owner),
        "the capture ends with stop_capture"
    );
    assert!(!capturing(&other_client, owner));
}

/// A capture prepared but never opened ends at `stop_capture`, which is the only call left to a
/// caller whose warmup failed: the stream stops reporting it, and a whole capture then lands on
/// it, sealed and replayed.
pub fn a_prepared_capture_ends_at_stop_capture<R: Runtime>() {
    let client = R::client(&Default::default());
    let owner = StreamId { value: 3_000_003 };
    let capturing = || owner.executes(|| client.is_capturing());

    owner
        .executes(|| client.graph_prepare())
        .expect("graph_prepare");
    assert!(capturing(), "the warmup run is part of the capture");
    owner
        .executes(|| client.stop_capture())
        .expect_err("nothing recorded, so there is no graph to seal");
    assert!(!capturing(), "the capture ended with stop_capture");

    let n = 4usize;
    let (input, output) = owner.executes(|| {
        (
            client.create_from_slice(f32::as_bytes(&[1.0, 2.0, 3.0, 4.0])),
            client.empty(n * core::mem::size_of::<f32>()),
        )
    });
    let run = || {
        add_one::launch(
            &client,
            CubeCount::Static(1, 1, 1),
            CubeDim::new(&client, n),
            unsafe { BufferArg::from_raw_parts(input.clone(), n) },
            unsafe { BufferArg::from_raw_parts(output.clone(), n) },
        );
    };

    let out = owner.executes(|| {
        client
            .graph_prepare()
            .expect("the stream can be prepared again");
        run();
        client.read_one(output.clone()).unwrap();
        client.start_capture().expect("start_capture");
        run();
        let graph = client.stop_capture().expect("stop_capture");
        unsafe { graph.replay() }.expect("replay enqueues");
        client.read_one(output.clone()).unwrap()
    });
    assert_eq!(f32::from_bytes(&out), &[2.0, 3.0, 4.0, 5.0]);
    assert!(!capturing());
}