Save inference parity implementation and evaluation harness

src/engine/metal/qwen_mtplx/submission.rs

@@ -0,0 +1,211 @@
+//! 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);
+}