DS4Server/src/engine/metal/qwen_mtplx/submission.rs

//! Residency at actual submission boundaries. This is not yet the reference
//! dependency/fence encoder; dispatch access roles and that port remain open.
use super::super::*;
use super::allocator::Allocator;
use std::cell::Cell;
use std::rc::Rc;
use std::sync::Arc;

pub(super) struct Submission {
    allocator: Arc<Allocator>,
    queue: NonNull<c_void>,
    attached: Cell<usize>,
    #[cfg(test)]
    commits: Cell<usize>,
}

impl Submission {
    pub(super) fn new(allocator: Arc<Allocator>) -> Result<Rc<Self>, String> {
        let queue = NonNull::new(unsafe { ds4_gpu_mtplx_retain_command_queue() })
            .ok_or("Metal command queue is not initialized")?;
        let owner = Rc::new(Self {
            allocator,
            queue,
            attached: Cell::new(0),
            #[cfg(test)]
            commits: Cell::new(0),
        });
        owner.attach(); // Reference CommandEncoder attaches existing sets at creation.
        Ok(owner)
    }

    fn attach(&self) {
        let mut attached = self.attached.get();
        unsafe {
            self.allocator
                .residency
                .attach_new_sets(self.queue.as_ptr(), &mut attached)
        };
        self.attached.set(attached);
    }

    pub(super) fn begin(self: &Rc<Self>) -> Result<SubmissionCommands, String> {
        if unsafe { ds4_gpu_mtplx_command_queue() } != self.queue.as_ptr() {
            return Err("Metal submission owner belongs to a different queue".into());
        }
        unsafe extern "C" fn before_commit(context: *mut c_void, queue: *mut c_void) {
            // The scope owns an Rc at this stable address. The native hook is
            // thread-local and synchronous; Submission/Rc cannot cross threads.
            let owner = unsafe { &*context.cast::<Submission>() };
            assert_eq!(queue, owner.queue.as_ptr());
            owner.attach();
            #[cfg(test)]
            owner.commits.set(owner.commits.get() + 1);
        }
        let context = Rc::as_ptr(self).cast_mut().cast();
        check(
            unsafe { ds4_gpu_mtplx_submission_hook(context, Some(before_commit)) },
            "registering Metal submission owner",
        )?;
        match Commands::begin() {
            Ok(commands) => Ok(SubmissionCommands {
                commands: Some(commands),
                owner: self.clone(),
            }),
            Err(error) => {
                unsafe { ds4_gpu_mtplx_submission_hook(context, None) };
                Err(error)
            }
        }
    }
}

impl Drop for Submission {
    fn drop(&mut self) {
        unsafe { ds4_gpu_mtplx_queue_free(self.queue.as_ptr()) };
    }
}

pub(super) struct SubmissionCommands {
    commands: Option<Commands>,
    owner: Rc<Submission>,
}

impl SubmissionCommands {
    pub(super) fn flush(&mut self) -> Result<(), String> {
        self.commands.as_mut().unwrap().flush()
    }

    pub(super) fn eval_boundary(&mut self) -> Result<(), String> {
        self.commands.as_mut().unwrap().eval_boundary()
    }

    pub(super) fn finish(mut self) -> Result<(), String> {
        self.commands.take().unwrap().finish()
    }
}

impl Drop for SubmissionCommands {
    fn drop(&mut self) {
        // Drop/finish drains native work while the hook and its owner still live.
        drop(self.commands.take());
        let context = Rc::as_ptr(&self.owner).cast_mut().cast();
        assert_ne!(unsafe { ds4_gpu_mtplx_submission_hook(context, None) }, 0);
    }
}

#[test]
fn mtplx_all_native_command_commits_pass_the_submission_hook() {
    let source = include_str!("../../../../native/metal/ds4_metal.m");
    assert_eq!(source.matches("[cb commit]").count(), 1);
    assert!(!source.contains("[g_batch_cb commit]"));
    for line in source.lines().filter(|line| line.contains(" commit]")) {
        assert!(
            line.contains("[cb commit]")
                || line.contains("[g_model_residency_set commit]")
                || line.contains("[residency_set commit]")
                || line.contains("id<MTLResidencySet>)set commit]"),
            "unreviewed native commit bypass: {line}"
        );
    }
    let helper = source
        .split_once("static void ds4_gpu_commit_command_buffer(")
        .unwrap()
        .1
        .split_once("static void ds4_gpu_stream_expert_cache_note_owned_created")
        .unwrap()
        .0;
    assert!(
        helper
            .find("g_mtplx_before_commit(g_mtplx_submission_context")
            .unwrap()
            < helper.find("[cb commit]").unwrap()
    );
    assert!(helper.contains("cb.commandQueue == g_queue"));
}

#[test]
#[ignore = "requires Apple Metal, no model weights"]
fn mtplx_submission_attaches_at_flush_eval_finish_and_drop() {
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let allocator = Allocator::new().unwrap();
    allocator.residency.set_max_per_set(8 << 20);
    allocator.set_wired_limit(128 << 20).unwrap();
    let owner = Submission::new(allocator.clone()).unwrap();
    assert_eq!(owner.attached.get(), allocator.residency.num_sets());
    let input = allocator.allocate(16).unwrap().unwrap();
    let output = allocator.allocate(16).unwrap().unwrap();
    let expected = [0x80, 0x3f].repeat(8);
    input.write(0, &expected).unwrap();
    let copy = || {
        assert!(
            super::qsa_dynamic_copy_rows([&input, &output], [8, 8], [1, 8], [None, None], [8; 2])
                .unwrap()
                .is_empty()
        )
    };
    let mut commands = owner.begin().unwrap();
    assert!(
        owner.begin().is_err(),
        "nested scopes must not replace the outer hook"
    );
    copy();
    let buffers = (0..8)
        .map(|_| allocator.allocate(5 << 20).unwrap().unwrap())
        .collect::<Vec<_>>();
    assert!(owner.attached.get() < allocator.residency.num_sets());
    commands.flush().unwrap();
    assert_eq!(owner.commits.get(), 1);
    assert_eq!(owner.attached.get(), allocator.residency.num_sets());
    let big = allocator.allocate(32 << 20).unwrap().unwrap();
    assert!(owner.attached.get() < allocator.residency.num_sets());
    commands.eval_boundary().unwrap();
    assert_eq!(owner.commits.get(), 2);
    assert_eq!(owner.attached.get(), allocator.residency.num_sets());
    commands.finish().unwrap();
    assert_eq!(owner.commits.get(), 3);
    let mut actual = [0; 16];
    output.read(0, &mut actual).unwrap();
    assert_eq!(actual.as_slice(), expected);
    drop((buffers, big));

    // Drop submits with the owner still registered. Reuse the same queue cursor.
    let commands = owner.begin().unwrap();
    copy();
    drop(commands);
    assert_eq!(owner.commits.get(), 4);
    let commands = owner.begin().unwrap();
    assert_eq!(unsafe { ds4_gpu_end_commands() }, 1);
    assert_eq!(
        owner.commits.get(),
        5,
        "native-internal submission bypasses no hook"
    );
    copy(); // With no open batch, the native dispatch commits synchronously.
    assert_eq!(owner.commits.get(), 6);
    drop(commands);
    assert_eq!(
        owner.commits.get(),
        6,
        "no duplicate commit during scope cleanup"
    );

    // The legacy scope neither inherits the owner nor calls an expired context.
    let commands = Commands::begin().unwrap();
    copy();
    commands.finish().unwrap();
    assert_eq!(owner.commits.get(), 6);
    output.read(0, &mut actual).unwrap();
    assert_eq!(actual.as_slice(), expected);
}