Implement more robust memory pool lookup (#243) #335

Quba1 · 2024-12-02T14:18:13Z

Fixes #243

I have implemented a more robust way to find big enough memory pool for reserve/alloc. I think comments explain quite well what happens in the code.

Minimal example of the solution

pub fn find_mem_pool(pool: &[usize], size: usize) -> usize {
    let ind = pool
        .binary_search(&size)
        .or_else(|found_index| {
            if found_index < pool.len() {
                Ok(found_index)
            } else {
                Err(found_index)
            }
        })
        .unwrap_or_else(|_| panic!("Cannot find pool big enough to reserve {size} bytes"));

    ind
}

#[cfg(test)]
mod tests {
    use crate::find_mem_pool;

    #[test]
    fn partitioning() {
        let v = [1, 3, 3, 5, 7];

        assert_eq!(find_mem_pool(&v, 0), 0);
        assert_eq!(find_mem_pool(&v, 1), 0);
        assert_eq!(find_mem_pool(&v, 2), 1);
        assert_eq!(find_mem_pool(&v, 3), 2);
        assert_eq!(find_mem_pool(&v, 5), 3);
        assert_eq!(find_mem_pool(&v, 6), 4);
        assert_eq!(find_mem_pool(&v, 7), 4);

        // this panics
        find_mem_pool(&v, 8);
    }
}

I have also added a simple test to check if the code panics when too much memory is requested.

To avoid code repetition in reserve() and alloc() I've extracted the solution into a separate function. But because alloc and reserve require &mut Storage the finder function needs to be implemented for directly for pools (to have two mutable borrows for pools and storage). But as Vec<DynamicPool> is a foreign type I had to implement the function via trait, hence the addition of PoolFind trait.

I hope this PR helps you. If something needs improvements or corrections, please let me know.

nathanielsimard · 2024-12-02T18:11:39Z

crates/cubecl-runtime/src/memory_management/memory_manage.rs

+        let pool_ind = self
+            // This finds a (not necessarily first) pool that matches the size exactly
+            // or returns error with index where a matching element could be inserted while maintaining sorted order
+            .binary_search_by_key(&size, |p| p.max_alloc_size())


size won't be equal to any of the max_alloc_sizes, we should find a bucket where size >= min_slice_size && size <= max_alloc_size. We don't need binary search since we only have < 10 buckets, max 20. Having the trait is great though.

Sorry, I didn't know about the size >= min_slice_size condition. But I also can't find min_slice_size defined anywhere.

Is the Vec<DynamicPools> sorted by min_slice_size too? Because then binary search would still be possible.

Binary search was what partition_point() was doing and this implementation only extends that technique. But with such a small array size I do agree that linear search will not impact the performance and will allow for the size >= min_slice_size condition - I'll change the implementation.

Hiya - give I wrote this I'll barge in - there is no min_slice_size, that's implicit by the fact that that should have been allocated in the previous pool. Stuff being sorted is therefore a condition anyway - and binary search seems fine? I agree it's a small performance win but... idk, seems clearer & faster. I just didn't quite handle the error condition of the search correctly.

The trait also seems like it could just be a struct? Doubt you'd want multiple possible implementations of this.

In case there's only one condition, then the most important consideration for using binary search is that if there are two pools with same max_alloc_size then it's not predictable which one will be chosen. The question is if that's an acceptable behaviour.

The trait also seems like it could just be a struct? Doubt you'd want multiple possible implementations of this.

Struct or newtype would require changing the type of pools in MemoryManagement and I didn't feel comfortable doing that. So I created the trait, which I admit is a little bit excessive, but necessary for borrow checker.

Implement more robust memory pool lookup (tracel-ai#243)

2e89116

nathanielsimard reviewed Dec 2, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement more robust memory pool lookup (#243) #335

Implement more robust memory pool lookup (#243) #335

Quba1 commented Dec 2, 2024

nathanielsimard Dec 2, 2024

Quba1 Dec 3, 2024

ArthurBrussee Dec 3, 2024

Quba1 Dec 3, 2024

Implement more robust memory pool lookup (#243) #335

Are you sure you want to change the base?

Implement more robust memory pool lookup (#243) #335

Conversation

Quba1 commented Dec 2, 2024

nathanielsimard Dec 2, 2024

Choose a reason for hiding this comment

Quba1 Dec 3, 2024

Choose a reason for hiding this comment

ArthurBrussee Dec 3, 2024

Choose a reason for hiding this comment

Quba1 Dec 3, 2024

Choose a reason for hiding this comment