Skip to main content

ruccl/tensor_device/
communicator.rs

1use super::{TensorDevice, TensorDeviceError, TensorElement};
2use crate::in_process::{Communicator, collective::InProcessCollective};
3use crate::rank::{communicator::RankCommunicator, device_collective::DeviceCollective};
4use ruda_tensor::Backend;
5
6impl<B: Backend> RankCommunicator<TensorDevice<B>> {
7    /// Select the wire element type from the typed tensor buffer contract.
8    pub fn tensor_collective<T: TensorElement>(&self) -> DeviceCollective<'_, T, TensorDevice<B>> {
9        self.device_collective(T::ELEMENT_TYPE)
10    }
11}
12
13impl<B: Backend> Communicator<TensorDevice<B>> {
14    pub fn tensor_collective<T: TensorElement>(
15        &self,
16    ) -> InProcessCollective<'_, T, TensorDevice<B>, TensorDeviceError> {
17        self.collective()
18    }
19}