Incorrect Buffer copy_from and copy_to on CPU/DX and possibly Metal backends

[shiinamiyuki]

Describe the bug Some backends copy host data to staging buffer at submission time. If the host data is written by device then later upload to device again without syncing in between, the old data is uploaded.

To Reproduce The following code using rust frontend can reproduce the issue on CPU and DX on my machine

use luisa::prelude::*;
use luisa_compute as luisa;
use std::env::current_exe;

fn main() {
    luisa::init_logger_verbose();
    let args: Vec<String> = std::env::args().collect();
    assert!(
        args.len() <= 2,
        "Usage: {} <backend>. <backend>: cpu, cuda, dx, metal, remote",
        args[0]
    );

    let ctx = Context::new(current_exe().unwrap());
    let device = ctx.create_device(if args.len() == 2 {
        args[1].as_str()
    } else {
        "cpu"
    });
    let x = device.create_buffer::<f32>(1024);
    let y = device.create_buffer::<f32>(1024);
    let z = device.create_buffer::<f32>(1024);
    x.view(..).fill_fn(|i| i as f32);
    y.view(..).fill_fn(|i| 1000.0 * i as f32);

    let kernel = device.create_kernel_with_options::<fn(Buffer<f32>)>(
        KernelBuildOptions {
            name: Some("vecadd".into()),
            ..Default::default()
        },
        &track!(|buf_z| {
            // z is pass by arg
            let buf_x = &x; // x and y are captured
            let buf_y = &y;
            let tid = dispatch_id().x;
            let x = buf_x.read(tid);
            let y = buf_y.read(tid);
            buf_z.write(tid, x + y);
        }),
    );
    let mut z_data = vec![123.0f32; 1024];

    unsafe {
        let s = device.default_stream().scope();
        let z_data_ptr = z_data.as_mut_ptr();
        s.submit([
            z.copy_from_async(std::slice::from_raw_parts_mut(z_data_ptr, 1024)),
            kernel.dispatch_async([1024, 1, 1], &z),
            z.copy_to_async(std::slice::from_raw_parts_mut(z_data_ptr, 1024)),
            z.copy_from_async(std::slice::from_raw_parts_mut(z_data_ptr, 1024)),
            z.copy_to_buffer_async(&x)
        ]);
    }

    // this should produce the expected behavior
    // unsafe {
    //     let z_data_ptr = z_data.as_mut_ptr();

    //     z.copy_from(std::slice::from_raw_parts_mut(z_data_ptr, 1024));
    //     kernel.dispatch([1024, 1, 1], &z);
    //     z.copy_to(std::slice::from_raw_parts_mut(z_data_ptr, 1024));
    //     z.copy_from(std::slice::from_raw_parts_mut(z_data_ptr, 1024));
    //     z.copy_to_buffer(&x);
    // }

    println!("{:?}", &z_data[0..16]);
    let x_data = x.copy_to_vec();
    println!("{:?}", &x_data[0..16]);
}

Expected behavior The correct output should be

[0.0, 1001.0, 2002.0, 3003.0, 4004.0, 5005.0, 6006.0, 7007.0, 8008.0, 9009.0, 10010.0, 11011.0, 12012.0, 13013.0, 14014.0, 15015.0]
[0.0, 1001.0, 2002.0, 3003.0, 4004.0, 5005.0, 6006.0, 7007.0, 8008.0, 9009.0, 10010.0, 11011.0, 12012.0, 13013.0, 14014.0, 15015.0]

However on CPU and DX it is:

[0.0, 1001.0, 2002.0, 3003.0, 4004.0, 5005.0, 6006.0, 7007.0, 8008.0, 9009.0, 10010.0, 11011.0, 12012.0, 13013.0, 14014.0, 15015.0]
[123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0, 123.0]

instead

Desktop (please complete the following information):

• OS: Windows 11