Adding reservation aware resource adapter

[nirandaperera]

This PR adds ReservationAwareResourceAdaptor, whose reservations are themselves RMM memory resources: hand one to cudf and every allocation made through it is charged against the reservation, turning "reserve before allocate" into something the type system and the allocator enforce rather than a convention.

• MemoryReservation is a shared_resource, so it can be passed anywhere a rmm::device_async_resource_ref is accepted. Allocating draws down its balance and throws rmm::out_of_memory when the grant is exhausted.
• The adaptor extends RmmResourceAdaptor's allocation tracking with a runtime-adjustable limit and a reserved-bytes counter (follows the limit logic in BufferResource)
• NOTE: Unspent balance is refunded only when the last reference dies. So, callers need to be cautious when over reserving!.

Usage

• General usage

    ReservationAwareResourceAdaptor adaptor{rmm::mr::cuda_memory_resource{}, 1MB};
    auto res = adaptor.reserve(2KB, AllowOverbooking::NO);

    {
        rmm::device_buffer buf1{1KB, stream, res};
        rmm::device_buffer buf2{1KB, stream, res};
        EXPECT_THROW((rmm::device_buffer buf3{1KB, stream, res}, rmm::out_of_memory); // res is exhausted
    }
   rmm::device_buffer buf4{2KB, stream, res}; // works again because buffers are released

• Create another reservation from existing reservation.

  void foo(MemoryReservation res, ...){
    ... 
    auto res2 = res.adapter().reserve(...); 
     ...
  }

Notes

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[nirandaperera]

/ok to test

[bdice]

I would use a method to do this, rather than making total_reserved_ a public member.

[nirandaperera]

Its private, but MemoryReservationImpl<> has been friended in ReservationAwareResourceAdaptorImpl.

[nirandaperera]

@bdice thank you for the review.
BTW I'd like to know your thoughts about the general direction of this idea, reserve before allocate. I think the worse case scenario would be, if there are N allocations, we could be creating N separate shared_resource instances for each reservation.

[wence-]

current_allocated is "expensive" since it has to acquire a mutex. So let's make this compare_exchange_strong and lose the while loop, no?

[nirandaperera]

Maybe not. I think the while loop serves 2 purposes.

• spurious failures
• total_reserved_ reserved could also be modified between L108-119.
  So, the latter might still be in effect. I think we can rework the logic to avoid this. Let me try

[nirandaperera]

I changed this to a more pessimistic approach on total_reserved_. We claim upfront, and if overbooking is not allowed, we rollback.

[nirandaperera]

/ok to test