cubecl_core/frontend/container/
peer.rs

1//! Buffers in memory shared with the other devices of a runtime.
2
3use crate::{
4    frontend::container::{buffer_len::expand_buffer_length_native, slice},
5    prelude::*,
6    unexpanded,
7};
8use core::marker::PhantomData;
9use cubecl_runtime::server::Handle;
10use serde::{Deserialize, Serialize};
11
12/// A buffer in memory a client [shared with its peers](cubecl_runtime::client::Client::share_with_peers),
13/// as a kernel parameter: `&Peer<[T]>` reads it, `&mut Peer<[T]>` writes it.
14///
15/// Kernels on other devices may read and write the buffer while this one runs, wherever the
16/// memory lives: on the launching device, where its owner's kernels bind it, or on another one,
17/// which kernels there bind directly. Inside the kernel it is a slice like any other, with what
18/// sharing changes:
19///
20/// - The compiler promises nothing about who writes it or what it aliases, so a read is never
21///   served from a cache another device's writes go around. The kernel's other buffers keep
22///   every promise.
23/// - An atomic on it synchronizes at
24///   [`SyncScope::System`](crate::ir::dialect::synchronization::SyncScope::System), whatever
25///   scope the call names: a device-scope atomic may compile to an instruction that is not
26///   atomic across the link between the devices.
27///
28/// Nothing orders the accesses of different devices. Kernels hand data off through an atomic flag
29/// in a peer buffer, with [`sync_system()`](crate::prelude::sync_system()) between the data and the flag on both sides.
30pub struct Peer<T: ?Sized> {
31    _buffer: PhantomData<T>,
32}
33
34impl<C: CubePrimitive> CubeType for Peer<[C]> {
35    type ExpandType = SliceExpand<C>;
36}
37
38/// A `Peer` buffer reads as a slice inside a kernel. Only the expansion runs: these exist for
39/// the kernel's body to type-check as Rust.
40impl<C: CubePrimitive> core::ops::Deref for Peer<[C]> {
41    type Target = [C];
42
43    fn deref(&self) -> &[C] {
44        unexpanded!()
45    }
46}
47
48impl<C: CubePrimitive> core::ops::DerefMut for Peer<[C]> {
49    fn deref_mut(&mut self) -> &mut [C] {
50        unexpanded!()
51    }
52}
53
54/// The compilation argument of a [`Peer`] buffer: it never aliases another binding, since its
55/// memory is shared and an alias would be shared by accident.
56#[derive(Clone, PartialEq, Eq, Hash, Debug, Serialize, Deserialize)]
57pub struct PeerCompilationArg;
58
59/// A shared buffer passed to a [`Peer`] parameter.
60pub struct PeerArg {
61    pub(crate) handle: BufferBinding,
62}
63
64impl PeerArg {
65    /// `handle`, of `length` elements, for a [`Peer`] parameter.
66    ///
67    /// # Safety
68    ///
69    /// Specifying the wrong length may lead to out-of-bounds reads and writes.
70    ///
71    /// # Panics
72    ///
73    /// When `handle` was never [shared](cubecl_runtime::client::Client::share_with_peers): the
74    /// memory of a buffer only its device knows cannot be bound as shared.
75    pub unsafe fn from_raw_parts(handle: &Handle, length: usize) -> Self {
76        assert!(
77            handle.peer.is_some(),
78            "a `Peer` parameter takes a handle shared with `share_with_peers`"
79        );
80        Self {
81            handle: unsafe { BufferBinding::from_raw_parts(handle.clone(), length) },
82        }
83    }
84}
85
86impl<C: CubePrimitive> LaunchArg for Peer<[C]> {
87    type RuntimeArg = PeerArg;
88    type CompilationArg = PeerCompilationArg;
89
90    fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
91        let elem_size = launcher.with_scope(|scope| C::__expand_size(scope));
92        launcher.register_peer(arg, elem_size);
93        PeerCompilationArg
94    }
95
96    fn expand(_arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> SliceExpand<C> {
97        let buffer = builder.peer_buffer(C::__expand_as_type(&builder.scope));
98        let scope = &builder.scope;
99        let len = expand_buffer_length_native(scope, buffer);
100        let slice_var =
101            slice::from_raw_parts::<C>(scope, buffer, 0usize.into_expand(scope), len.into());
102        slice_var.expand.into()
103    }
104}
cubecl_core/frontend/container/peer.rs

cubecl_core/frontend/container/
peer.rs
