运行时与 Client

pub trait Runtime: Sized + Send + Sync + 'static + core::fmt::Debug + Clone {
    type Server: ServerStorage;
    type Device: Device;

    fn client(device: &Self::Device) -> Client {
        Client::load::<Self::Server>(device.to_id())
    }

    fn can_read_tensor(shape: &Shape, strides: &Strides) -> bool;

    fn target_properties() -> TargetProperties;
    // ...
}

pub struct Client {
    device: DeviceHandle<dyn Server>,
    utilities: Arc<ServerUtilities>,
    stream_id: Option<StreamId>,
}

    /// Runs `task` on the device thread, blocking until it returns.
    fn submit_blocking<'a, R: Send, T: FnOnce(&mut dyn Any) -> R + Send + 'a>(
        &self,
        task: T,
    ) -> Result<R, CallError> {
        let state = self.state.service.clone();
        let current = StreamId::current();
        self.run_scoped(move || state.act_on(|s| current.executes(|| task(s.as_mut()))))
    }

    /// Asynchronously dispatches a task to the device thread.
    fn submit<T: FnOnce(&mut dyn Any) + Send + 'static>(&self, task: T) {
        self.submit_inner::<_, SEND_NO_FLUSH>(task)
            .expect("Can't have an error when submitting a task");
    }

        let stream_id = self.stream_id();
        let (handle_base, layouts) =
            self.utilities
                .layout_policy
                .apply(self.service_id(), stream_id, &descriptors);
        // ... 为每个张量配好 CopyDescriptor 和数据
        let (size, memory) = (handle_base.size(), handle_base.memory);
        self.device.submit(move |server| {
            // Cannot write on memory that wasn't initialized. The error is attached to the
            // buffer and is reported at the next sync point.
            if server.initialize_memory(memory, size, stream_id).is_ok() {
                server.write(descriptors, stream_id);
            }
        });

        layouts

  for round in 0..3 {
    let t = Instant::now();
    let h = client.create_from_slice(f32::as_bytes(&data));
    let t_create = t.elapsed();
    let t = Instant::now();
    let count = CubeCount::Static((n as u32).div_ceil(256), 1, 1);
    // SAFETY: the handle holds exactly n f32 values.
    let arg = unsafe { BufferArg::from_raw_parts(h.clone(), n) };
    scale::launch::<f32>(&client, count, CubeDim::new_1d(256), arg);
    let t_launch = t.elapsed();
    let t = Instant::now();
    let bytes = client.read_one(h).unwrap();
    let t_read = t.elapsed();
    // ...
  }

round 0: create 4.8595ms, launch 7.053792ms, read_one 132.533458ms (x[0] = 2)
round 1: create 2.840542ms, launch 31.416µs, read_one 6.989583ms (x[0] = 2)
round 2: create 2.421208ms, launch 17.5µs, read_one 3.219125ms (x[0] = 2)

            let (staging, binding) = self.auxiliary.reserve_staging(aligned_len).unwrap();

            self.tasks_count += 1;
            self.encoder.copy_buffer_to_buffer(
                &resource.buffer,
                resource.offset,
                &staging.buffer,
                0,
                aligned_len,
            );
        // ...
        // Flush all commands to the queue, so GPU gets started on copying to the
        // staging buffer. ...
        if let Err(err) = self.flush(stream_id) {
            return Box::pin(async move { Err(err) });
        }
        // ...
                let (sender, receiver) = cubecl_environment::future::channel::bounded(1);
                staging
                    .buffer
                    .slice(..)
                    .map_async(wgpu::MapMode::Read, move |v| {
                        // This might fail if the channel is closed (eg. the future is dropped).
                        // This is fine, just means results aren't needed anymore.
                        let _ = sender.try_send(v);
                    });
        // ...
        let poll = self.poll.start_polling();
        // ...
        Box::pin(async move {
            for receiver in callbacks.iter().flatten() {
                let mapped = receiver.recv().await;

fn pool_demo() {
  let client = cubecl::Device::default().client();
  usage(&client, "start");
  let data = vec![1u8; MIB];
  let handles: Vec<_> = (0..8).map(|_| client.create_from_slice(&data)).collect();
  usage(&client, "8 x 1 MiB created");
  pools(&client);
  drop(handles);
  usage(&client, "dropped");
  let again: Vec<_> = (0..8).map(|_| client.create_from_slice(&data)).collect();
  usage(&client, "8 x 1 MiB created again");
  drop(again);
  let big = client.create_from_slice(&vec![0u8; 64 * MIB]);
  usage(&client, "one 64 MiB created");
  pools(&client);
  drop(big);
  usage(&client, "64 MiB dropped");
  client.memory_cleanup().unwrap();
  usage(&client, "memory_cleanup");
  pools(&client);
}

pub enum Cleanup {
    /// On the way to a reservation: each pool gives back what its own policy
    /// says is unused, such as a page past its deallocation period or an
    /// outdated page that emptied.
    Periodic,
    /// Asked for: everything nothing holds goes back.
    Explicit,
}

burn start                         allocs=0   in_use=        0 reserved=        0
burn 8 x [512,512] f32             allocs=16  in_use=  8388992 reserved=  8515072
  report Some([SlicedPoolReport { page_size: 8388608, pages: 0, pages_peak: 0, largest_alloc: 0 }, SlicedPoolReport { page_size: 2097152, pages: 4, pages_peak: 4, largest_alloc: 1048576 }])
burn dropped                       allocs=8   in_use=      384 reserved=  8515072
burn memory_cleanup                allocs=0   in_use=        0 reserved=        0

2 launches, handles dropped        allocs=2   in_use=       16 reserved=  8421120
memory_cleanup                     allocs=0   in_use=        0 reserved=        0

  let _main = client.create_from_slice(&data);
  let other = std::thread::spawn({
    let client = client.clone();
    let data = data.clone();
    move || {
      let handle = client.create_from_slice(&data);
      let mine = client.memory_report(cubecl::MemoryScope::CurrentStream);
      println!("worker thread CurrentStream: {} stream(s)", mine.streams.len());
      handle
    }
  })
  .join()
  .unwrap();

worker thread CurrentStream: 1 stream(s)
StreamId { value: 0 }: allocs=1 in_use=1048576 reserved=2097152
StreamId { value: 1 }: allocs=1 in_use=1048576 reserved=2097152

start                              allocs=0   in_use=        0 reserved=        0
8 x 1 MiB created                  allocs=8   in_use=  8388608 reserved=  8388608
  dynamic Exclusive { max_alloc_size: 0 } pages=0 peak=0 largest=0 in_use=0
  dynamic Sliced { page_size: 8388608, max_slice_size: 65536 } pages=0 peak=0 largest=0 in_use=0
  dynamic Adaptive { page_size: 2097152, outdated_pages: 0 } pages=4 peak=4 largest=1048576 in_use=8388608
  persistent pages=0 in_use=0
dropped                            allocs=0   in_use=        0 reserved=  8388608
8 x 1 MiB created again            allocs=8   in_use=  8388608 reserved=  8388608
one 64 MiB created                 allocs=1   in_use= 67108864 reserved= 76546048
  dynamic Exclusive { max_alloc_size: 0 } pages=0 peak=0 largest=0 in_use=0
  dynamic Sliced { page_size: 8388608, max_slice_size: 65536 } pages=0 peak=0 largest=0 in_use=0
  dynamic Adaptive { page_size: 68157440, outdated_pages: 4 } pages=1 peak=1 largest=67108864 in_use=67108864
  dynamic Sliced { page_size: 2097152, max_slice_size: 2097152 } pages=4 peak=4 largest=1048576 in_use=0
  persistent pages=0 in_use=0
64 MiB dropped                     allocs=0   in_use=        0 reserved= 76546048
memory_cleanup                     allocs=0   in_use=        0 reserved=        0
  dynamic Exclusive { max_alloc_size: 0 } pages=0 peak=0 largest=0 in_use=0
  dynamic Sliced { page_size: 8388608, max_slice_size: 65536 } pages=0 peak=0 largest=0 in_use=0
  dynamic Adaptive { page_size: 68157440, outdated_pages: 0 } pages=0 peak=1 largest=67108864 in_use=0
  persistent pages=0 in_use=0