Finish some todos (#951)

Author: nathanielsimard

crates/cubecl-hip/src/compute/command.rs

@@ -7,14 +7,15 @@ use cubecl_core::{
 };
 use cubecl_hip_sys::{
     HIP_SUCCESS, hipMemcpyKind_hipMemcpyDeviceToHost, hipMemcpyKind_hipMemcpyHostToDevice,
+    ihipStream_t,
 };
 use cubecl_runtime::{
     id::KernelId,
     logging::ServerLogger,
     memory_management::{MemoryAllocationMode, MemoryHandle},
     stream::ResolvedStreams,
 };
-use std::sync::Arc;
+use std::{ffi::c_void, sync::Arc};
 
 use crate::{
     compute::{
@@ -31,7 +32,7 @@ use crate::{
 /// registration, and task execution.
 pub struct Command<'a> {
     ctx: &'a mut HipContext,
-    streams: ResolvedStreams<'a, HipStreamBackend>,
+    pub(crate) streams: ResolvedStreams<'a, HipStreamBackend>,
 }
 
 impl<'a> Command<'a> {
@@ -274,58 +275,13 @@ impl<'a> Command<'a> {
             return Err(IoError::UnsupportedStrides);
         }
 
-        let rank = shape.len();
         let resource = self.resource(binding)?;
         let stream = match stream_id {
             Some(id) => self.streams.get(&id),
             None => self.streams.current(),
         };
 
-        if rank <= 1 {
-            unsafe {
-                let status = cubecl_hip_sys::hipMemcpyDtoHAsync(
-                    bytes.as_mut_ptr() as *mut _,
-                    resource.ptr,
-                    bytes.len(),
-                    stream.sys,
-                );
-
-                if status != HIP_SUCCESS {
-                    return Err(IoError::Unknown(format!("HIP memcpy failed: {}", status)));
-                }
-            }
-            return Ok(());
-        }
-
-        let dim_x = shape[rank - 1];
-        let width_bytes = dim_x * elem_size;
-        let dim_y: usize = shape.iter().rev().skip(1).product();
-        let pitch = strides[rank - 2] * elem_size;
-
-        unsafe {
-            let status = cubecl_hip_sys::hipMemcpy2DAsync(
-                bytes.as_mut_ptr() as *mut _,
-                width_bytes,
-                resource.ptr,
-                pitch,
-                width_bytes,
-                dim_y,
-                hipMemcpyKind_hipMemcpyDeviceToHost,
-                stream.sys,
-            );
-
-            // Fallback, sometimes the copy doesn't work.
-            if status != HIP_SUCCESS {
-                let status = cubecl_hip_sys::hipMemcpyDtoHAsync(
-                    bytes.as_mut_ptr() as *mut _,
-                    resource.ptr,
-                    bytes.len(),
-                    stream.sys,
-                );
-                assert_eq!(status, HIP_SUCCESS, "Should send data to device");
-            }
-        }
-        Ok(())
+        unsafe { write_to_cpu(shape, strides, elem_size, bytes, resource.ptr, stream.sys) }
     }
 
     /// Writes data from the host to GPU memory as specified by the copy descriptor.
@@ -469,3 +425,61 @@ impl<'a> Command<'a> {
         };
     }
 }
+
+pub(crate) unsafe fn write_to_cpu(
+    shape: &[usize],
+    strides: &[usize],
+    elem_size: usize,
+    bytes: &mut Bytes,
+    resource_ptr: *mut c_void,
+    stream: *mut ihipStream_t,
+) -> Result<(), IoError> {
+    let rank = shape.len();
+
+    if rank <= 1 {
+        let status = unsafe {
+            cubecl_hip_sys::hipMemcpyDtoHAsync(
+                bytes.as_mut_ptr() as *mut _,
+                resource_ptr,
+                bytes.len(),
+                stream,
+            )
+        };
+
+        if status != HIP_SUCCESS {
+            return Err(IoError::Unknown(format!("HIP memcpy failed: {}", status)));
+        }
+        return Ok(());
+    }
+
+    let dim_x = shape[rank - 1];
+    let width_bytes = dim_x * elem_size;
+    let dim_y: usize = shape.iter().rev().skip(1).product();
+    let pitch = strides[rank - 2] * elem_size;
+
+    unsafe {
+        let status = cubecl_hip_sys::hipMemcpy2DAsync(
+            bytes.as_mut_ptr() as *mut _,
+            width_bytes,
+            resource_ptr,
+            pitch,
+            width_bytes,
+            dim_y,
+            hipMemcpyKind_hipMemcpyDeviceToHost,
+            stream,
+        );
+
+        // Fallback, sometimes the copy doesn't work.
+        if status != HIP_SUCCESS {
+            let status = cubecl_hip_sys::hipMemcpyDtoHAsync(
+                bytes.as_mut_ptr() as *mut _,
+                resource_ptr,
+                bytes.len(),
+                stream,
+            );
+            assert_eq!(status, HIP_SUCCESS, "Should send data to device");
+        }
+    }
+
+    Ok(())
+}

crates/cubecl-hip/src/compute/server.rs

@@ -1,7 +1,9 @@
 use super::storage::gpu::GpuResource;
 use super::storage::gpu::GpuStorage;
 use crate::compute::command::Command;
+use crate::compute::command::write_to_cpu;
 use crate::compute::context::HipContext;
+use crate::compute::fence::Fence;
 use crate::compute::stream::HipStreamBackend;
 use crate::runtime::HipCompiler;
 use cubecl_common::bytes::Bytes;
@@ -238,7 +240,17 @@ impl ComputeServer for HipServer {
 }
 
 impl ServerCommunication for HipServer {
-    const SERVER_COMM_ENABLED: bool = false;
+    const SERVER_COMM_ENABLED: bool = true;
+
+    fn copy(
+        server_src: &mut Self,
+        server_dst: &mut Self,
+        src: CopyDescriptor<'_>,
+        stream_id_src: StreamId,
+        stream_id_dst: StreamId,
+    ) -> Result<Allocation, IoError> {
+        Self::change_server_serialized(server_src, server_dst, src, stream_id_src, stream_id_dst)
+    }
 }
 
 impl HipServer {
@@ -277,6 +289,77 @@ impl HipServer {
 
         Command::new(&mut self.ctx, streams)
     }
+
+    fn change_server_serialized(
+        server_src: &mut Self,
+        server_dst: &mut Self,
+        src: CopyDescriptor<'_>,
+        stream_id_src: StreamId,
+        stream_id_dst: StreamId,
+    ) -> Result<Allocation, IoError> {
+        let shape = src.shape.to_vec();
+        let strides = src.strides.to_vec();
+        let elem_size = src.elem_size;
+        let binding = src.binding.clone();
+        let num_bytes = shape.iter().product::<usize>() * elem_size;
+
+        // We start by creating a command on the destination server.
+        //
+        // Here we allocate the necessary bytes using pinned memory managed by the destination
+        // server along a new GPU handle. This way, the bytes could be reused later by that server,
+        // and the lifetime of that handle is aligned with the execution order of the destination server,
+        // removing the need to keep the bytes handle alive using synchronization, which would be the
+        // case if we allocated the bytes using the source server.
+        let mut command_dst = server_dst.command_no_inputs(stream_id_dst);
+        let handle = command_dst.reserve(binding.size())?;
+        let mut bytes = command_dst.reserve_cpu(num_bytes, true, None);
+        let copy_desc = handle.copy_descriptor(&shape, &strides, elem_size);
+
+        // We need to free the command before creating another one.
+        core::mem::drop(command_dst);
+
+        // We create a command on the source server to retrieve the correct resource from the
+        // source memory pools. We also make sure the current stream is aligned with the stream of
+        // the binding, where the data was first allocated.
+        //
+        // We use the source stream to copy the data from the source server into the allocated
+        // bytes. This ensures that the source binding follows the correct execution order, meaning
+        // that we don't have to keep the source handle alive using synchronization, which would be
+        // the case if we performed the copy on the destination server.
+        let mut command_src = server_src.command(stream_id_src, [&src.binding].into_iter());
+        let resource_src = command_src.resource(binding.clone())?;
+        let stream_src = command_src.streams.current().sys;
+
+        unsafe {
+            write_to_cpu(
+                &shape,
+                &strides,
+                elem_size,
+                &mut bytes,
+                resource_src.ptr,
+                stream_src,
+            )?;
+        }
+        let fence_src = Fence::new(stream_src);
+
+        // We need to free the command before creating another one.
+        core::mem::drop(command_src);
+
+        // Finally, we recreate a new command on the destination server to write the data stored in
+        // pinned memory into the destination server. Here we need to wait for the initial copy
+        // made by the source server using an event. The synchronization is done lazily on the
+        // destination stream, which is very efficient.
+        let mut command_dst = server_dst.command_no_inputs(stream_id_dst);
+        let stream_dst = command_dst.streams.current().sys;
+
+        fence_src.wait_async(stream_dst);
+        command_dst.write_to_gpu(copy_desc, &bytes)?;
+
+        // We drop the last command.
+        core::mem::drop(command_dst);
+
+        Ok(Allocation { handle, strides })
+    }
 }
 
 pub(crate) fn contiguous_strides(shape: &[usize]) -> Vec<usize> {

crates/cubecl-runtime/src/memory_management/memory_pool/exclusive_pool.rs

@@ -1,5 +1,5 @@
 use crate::{
-    memory_management::MemoryUsage,
+    memory_management::{BytesFormat, MemoryUsage},
     server::IoError,
     storage::{ComputeStorage, StorageHandle, StorageUtilization},
 };
@@ -24,7 +24,23 @@ pub struct ExclusiveMemoryPool {
 
 impl core::fmt::Display for ExclusiveMemoryPool {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-        f.write_str("ExclusiveMemoryPool")
+        f.write_fmt(format_args!(
+            " - Exclusive Pool max_alloc_size={}\n",
+            BytesFormat::new(self.max_alloc_size)
+        ))?;
+
+        for page in self.pages.iter() {
+            let is_free = page.slice.is_free();
+            let size = BytesFormat::new(page.slice.effective_size());
+
+            f.write_fmt(format_args!("   - Page {size} is_free={is_free}\n"))?;
+        }
+
+        if !self.pages.is_empty() {
+            f.write_fmt(format_args!("\n{}\n", self.get_memory_usage()))?;
+        }
+
+        Ok(())
     }
 }
 

crates/cubecl-runtime/src/memory_management/memory_pool/persistent_pool.rs

@@ -136,16 +136,22 @@ impl MemoryPool for PersistentPool {
         explicit: bool,
     ) {
         if explicit {
-            self.slices.retain(|_, slice| {
+            let mut removed = Vec::new();
+            self.slices.retain(|id, slice| {
                 if slice.is_free() {
                     storage.dealloc(slice.storage.id);
+                    removed.push((*id, slice.effective_size()));
                     false
                 } else {
                     true
                 }
-                // TODO: Remove the slice id from the sizes map.
             });
 
+            for (id, size) in removed {
+                let ids = self.sizes.get_mut(&size).expect("The size should match");
+                ids.retain(|id_| *id_ != id);
+            }
+
             storage.flush();
         }
     }