burn-router 0.22.0

Multi-backend router decorator for the Burn framework
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

use alloc::vec::Vec;

use burn_backend::{
    DeviceId,
    distributed::{CollectiveTensor, DistributedOps, ReduceOperation},
    tensor::{Device, FloatTensor},
};
use burn_ir::{AllReduceOpIr, DeviceIdIr, DistributedOperationIr, OperationIr, OperationOutput};

use crate::{BackendRouter, RouterChannel, RouterClient, get_client};

impl<R: RouterChannel> DistributedOps<Self> for BackendRouter<R> {
    fn all_reduce(
        tensor: FloatTensor<Self>,
        op: ReduceOperation,
        device_ids: Vec<DeviceId>,
    ) -> CollectiveTensor<Self> {
        let client = tensor.client.clone();
        let device_ids = device_ids.into_iter().map(DeviceIdIr::from).collect();
        let desc = AllReduceOpIr::create(tensor.into_ir(), op, device_ids, || {
            client.create_empty_handle()
        });

        let output = client
            .register(OperationIr::Distributed(DistributedOperationIr::AllReduce(
                desc,
            )))
            .output();

        CollectiveTensor::new(output)
    }

    fn sync_collective(device: &Device<Self>) {
        // Fire-and-forget, like `all_reduce`: register a `SyncCollective` op on the device's
        // stream instead of a blocking call. The interpreter resolves it through the normal op
        // stream like any other distributed op.
        let client = get_client::<R>(device);
        client.register_op(OperationIr::Distributed(
            DistributedOperationIr::SyncCollective,
        ));
    }
}