[BUG]: cuda-async allows use after free without unsafe

[soooch]

Operating System

N/A

GPU and GPU Architecture

N/A

Error Description

It looks like DeviceFuture::drop now invokes DeviceFuture::release_in_flight_result. This prevents the scenario outlined in #242. However, in current Rust, it's unfortunately always safe to leak any value via core::mem::forget. So, we can simply forget the DeviceFuture then drop any formerly borrowed inputs.

Actually, the scenario demonstrated in the repro below is fairly similar to that found in the Leakpocalypse which is the finding that led to the modern version of the scoped threads API. You might also find without.boats' writing on this issue in The Scoped Task trilemma instructive.

Luckily, the Borrowing in this API is not the fully general form that without.boats discusses in that post. At least for Tensor, the resource that matters (the storage field) is already wrapped in an Arc. Taking ownership of the Storage by holding a ref for the stream lifetime of the dependent operation is one avenue that could be explored.

Minimal Reproduction Example

use core::future::{Future, IntoFuture};
use core::task;

use cutile::prelude::*;
use my_module::add;

#[cutile::module]
mod my_module {

    use cutile::core::*;

    // for example, but could be something more expensive
    #[cutile::entry(print_ir = true)]
    fn add<const S: [i32; 1]>(
        z: &mut Tensor<f32, S>,
        x: &Tensor<f32, { [-1] }>,
        y: &Tensor<f32, { [-1] }>,
    ) {
        let tile_x = x.load_like(z);
        let tile_y = y.load_like(z);
        z.store(tile_x + tile_y);
    }
}

#[tokio::main]
async fn main() -> Result<(), cuda_async::error::DeviceError> {
    let len = 2usize.pow(5);

    let x_buf = api::arange(len).await?;
    let mut context = task::Context::from_waker(&task::Waker::noop());
    let mut fut = Box::pin(
        add(
            api::zeros::<f32>(&[len]).partition([2]),
            &x_buf,
            api::ones(&[len]),
        )
        .into_future(),
    );

    // polling once for example.
    // could poll arbitrary number of times (until ready).
    // can wait however long u want between each poll.
    let _ = fut.as_mut().poll(&mut context);
    core::mem::forget(fut);
    drop(x_buf);

    Ok(())
}