1411 lines
48 KiB
Rust
1411 lines
48 KiB
Rust
//! CommandEncoder dependency and submission rules from MTPLX's pinned runtime
|
|
//! (backend/metal/device.cpp). Native code only performs Metal API operations.
|
|
//! This is still test-bound; complete model/stream routing remains open.
|
|
use super::*;
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::ffi::CStr;
|
|
use std::ops::{Deref, DerefMut};
|
|
use std::sync::{Arc, Mutex, OnceLock};
|
|
|
|
#[derive(Clone, Copy)]
|
|
pub(super) enum Access {
|
|
Input(usize),
|
|
Output(usize),
|
|
Raw,
|
|
}
|
|
|
|
impl Binding<'_> {
|
|
/// Byte offset relative to the bound backing, like set_input/output_array.
|
|
/// Geometry validation belongs to the array; native binding checks the base.
|
|
pub(super) fn at_byte_offset(mut self, offset: u64) -> Self {
|
|
assert!(self.allocation.is_some());
|
|
self.raw.offset = offset;
|
|
self
|
|
}
|
|
|
|
pub(super) fn input(mut self, elements: usize) -> Self {
|
|
assert!(self.allocation.is_some());
|
|
self.access = Some(Access::Input(elements));
|
|
self
|
|
}
|
|
|
|
pub(super) fn output(mut self, elements: usize) -> Self {
|
|
assert!(self.allocation.is_some());
|
|
self.access = Some(Access::Output(elements));
|
|
self
|
|
}
|
|
|
|
pub(super) fn raw_buffer(mut self) -> Self {
|
|
assert!(self.allocation.is_some());
|
|
self.access = Some(Access::Raw);
|
|
self
|
|
}
|
|
}
|
|
|
|
thread_local! {
|
|
static ACTIVE: RefCell<Option<Rc<RefCell<Encoder>>>> = const { RefCell::new(None) };
|
|
}
|
|
|
|
/// One encoding-thread stream. The scope owns its queue across operations and
|
|
/// routes the existing operators/allocations, without adding dispatch boundaries.
|
|
pub(in super::super) struct Scope(Rc<RefCell<Encoder>>);
|
|
|
|
impl Scope {
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(in super::super) fn new() -> Result<Self, String> {
|
|
if ACTIVE.with(|active| active.borrow().is_some()) {
|
|
return Err("MTPLX encoder scope is already active".into());
|
|
}
|
|
let encoder = Rc::new(RefCell::new(Encoder::new(allocator::Allocator::new()?)?));
|
|
Self::activate(encoder)
|
|
}
|
|
|
|
pub(super) fn activate(encoder: Rc<RefCell<Encoder>>) -> Result<Self, String> {
|
|
if ACTIVE.with(|active| active.borrow().is_some()) {
|
|
return Err("MTPLX encoder scope is already active".into());
|
|
}
|
|
ACTIVE.with(|active| *active.borrow_mut() = Some(encoder.clone()));
|
|
Ok(Self(encoder))
|
|
}
|
|
|
|
pub(in super::super) fn synchronize(&self) -> Result<(), String> {
|
|
self.0.borrow_mut().synchronize()
|
|
}
|
|
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(in super::super) fn commits(&self) -> usize {
|
|
self.0.borrow().commits
|
|
}
|
|
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(in super::super) fn primitive_commits(&self) -> usize {
|
|
self.0.borrow().primitive_commits
|
|
}
|
|
}
|
|
|
|
impl Drop for Scope {
|
|
fn drop(&mut self) {
|
|
let owner = ACTIVE.with(|active| active.borrow_mut().take()).unwrap();
|
|
assert!(Rc::ptr_eq(&owner, &self.0));
|
|
// The final Rc drains Encoder after routing is removed. Never leave an
|
|
// expired thread-local owner, including during unwinding.
|
|
}
|
|
}
|
|
|
|
pub(super) fn with_active<T>(operation: impl FnOnce(&mut Encoder) -> T) -> Option<T> {
|
|
let owner = ACTIVE.with(|active| active.borrow().clone());
|
|
owner.map(|owner| operation(&mut owner.borrow_mut()))
|
|
}
|
|
|
|
pub(in super::super) fn active_allocator() -> Option<Arc<allocator::Allocator>> {
|
|
with_active(|encoder| encoder.allocator.clone())
|
|
}
|
|
|
|
/// Both CPU and GPU evaluators retain inputs and siblings by Data identity,
|
|
/// excluding the primary output (including an input donated to that output).
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(super) fn primitive_buffers<'a>(
|
|
inputs: &[&'a Buffer],
|
|
outputs: &'a [Buffer],
|
|
primary: usize,
|
|
) -> Result<Vec<&'a Buffer>, String> {
|
|
let selected = outputs
|
|
.get(primary)
|
|
.ok_or("invalid evaluator primary output")?;
|
|
Ok(data_buffers(
|
|
inputs.iter().copied().chain(outputs.iter()),
|
|
Some(selected),
|
|
))
|
|
}
|
|
|
|
pub(super) fn data_buffers<'a>(
|
|
buffers: impl Iterator<Item = &'a Buffer>,
|
|
primary: Option<&Buffer>,
|
|
) -> Vec<&'a Buffer> {
|
|
let mut seen = HashSet::new();
|
|
if let Some(primary) = primary {
|
|
seen.insert(primary.data_identity());
|
|
}
|
|
buffers
|
|
.filter(|buffer| seen.insert(buffer.data_identity()))
|
|
.collect()
|
|
}
|
|
|
|
pub(super) fn add_temporaries(buffers: &[Buffer]) {
|
|
with_active(|encoder| {
|
|
for buffer in buffers {
|
|
encoder.add_temporary(buffer);
|
|
}
|
|
});
|
|
}
|
|
|
|
/// Concatenate's disjoint slice writers share an output allocation. Keep their
|
|
/// dependency join outside the dispatch borrow, including on early return.
|
|
pub(super) struct ConcurrentScope(Option<Rc<RefCell<Encoder>>>);
|
|
|
|
pub(super) fn concurrent_scope() -> ConcurrentScope {
|
|
let owner = ACTIVE.with(|active| active.borrow().clone());
|
|
if let Some(owner) = &owner {
|
|
owner.borrow_mut().dependencies.concurrent = true;
|
|
}
|
|
ConcurrentScope(owner)
|
|
}
|
|
|
|
impl Drop for ConcurrentScope {
|
|
fn drop(&mut self) {
|
|
if let Some(owner) = &self.0 {
|
|
owner.borrow_mut().dependencies.end_concurrent();
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct Dependencies {
|
|
prev_inputs: HashSet<usize>,
|
|
prev_outputs: HashSet<usize>,
|
|
next_inputs: HashSet<usize>,
|
|
next_outputs: HashSet<usize>,
|
|
concurrent_outputs: HashSet<usize>,
|
|
all_inputs: HashSet<usize>,
|
|
all_outputs: HashSet<usize>,
|
|
needs_barrier: bool,
|
|
concurrent: bool,
|
|
ops: i32,
|
|
// The original adds array.data_size(), measured in ELEMENTS, not bytes.
|
|
data_size: usize,
|
|
}
|
|
|
|
impl Dependencies {
|
|
fn bind(&mut self, resource: usize, access: Access) {
|
|
match access {
|
|
Access::Raw => {
|
|
self.all_inputs.insert(resource);
|
|
self.all_outputs.insert(resource);
|
|
}
|
|
Access::Input(elements) | Access::Output(elements) => {
|
|
if self.all_inputs.insert(resource) {
|
|
self.data_size = self.data_size.wrapping_add(elements);
|
|
}
|
|
self.next_inputs.insert(resource);
|
|
self.needs_barrier |= self.prev_outputs.contains(&resource);
|
|
if matches!(access, Access::Output(_)) {
|
|
self.register_output(resource);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn register_output(&mut self, resource: usize) {
|
|
self.all_outputs.insert(resource);
|
|
if self.concurrent {
|
|
self.concurrent_outputs.insert(resource);
|
|
} else {
|
|
self.next_outputs.insert(resource);
|
|
self.needs_barrier |= self.prev_inputs.contains(&resource);
|
|
}
|
|
}
|
|
|
|
fn dispatched(&mut self) {
|
|
if self.needs_barrier {
|
|
self.needs_barrier = false;
|
|
std::mem::swap(&mut self.prev_inputs, &mut self.next_inputs);
|
|
std::mem::swap(&mut self.prev_outputs, &mut self.next_outputs);
|
|
} else {
|
|
self.prev_inputs.extend(self.next_inputs.iter());
|
|
self.prev_outputs.extend(self.next_outputs.iter());
|
|
}
|
|
self.next_inputs.clear();
|
|
self.next_outputs.clear();
|
|
self.ops += 1;
|
|
}
|
|
|
|
fn end_concurrent(&mut self) {
|
|
self.concurrent = false;
|
|
self.prev_outputs.extend(self.concurrent_outputs.drain());
|
|
}
|
|
|
|
fn end_encoding(&mut self) {
|
|
self.needs_barrier = false;
|
|
self.concurrent = false;
|
|
self.prev_inputs.clear();
|
|
self.next_inputs.clear();
|
|
self.prev_outputs.clear();
|
|
self.next_outputs.clear();
|
|
self.concurrent_outputs.clear();
|
|
self.all_inputs.clear();
|
|
// all_outputs is moved to the completion handler, like the reference.
|
|
assert!(self.all_outputs.is_empty());
|
|
}
|
|
|
|
fn needs_commit(&self, [ops, mb]: [i32; 2]) -> bool {
|
|
self.ops > ops || (self.data_size >> 20) > mb as usize
|
|
}
|
|
|
|
fn committed(&mut self) {
|
|
self.ops = 0;
|
|
self.data_size = 0;
|
|
}
|
|
}
|
|
|
|
struct Native(NonNull<c_void>);
|
|
|
|
impl Native {
|
|
fn new(pointer: *mut c_void, what: &str) -> Result<Self, String> {
|
|
NonNull::new(pointer)
|
|
.map(Self)
|
|
.ok_or_else(|| format!("failed to create MTPLX {what}"))
|
|
}
|
|
|
|
fn raw(&self) -> *mut c_void {
|
|
self.0.as_ptr()
|
|
}
|
|
}
|
|
|
|
impl Drop for Native {
|
|
fn drop(&mut self) {
|
|
unsafe { ds4_gpu_mtplx_object_free(self.raw()) };
|
|
}
|
|
}
|
|
|
|
struct Fence(Native);
|
|
// SAFETY: only immutable handles/ownership cross threads. Fence encoding is
|
|
// confined to Encoder's thread; completion callbacks only drop retained handles.
|
|
unsafe impl Send for Fence {}
|
|
unsafe impl Sync for Fence {}
|
|
|
|
type Producers = Arc<Mutex<HashMap<usize, Arc<Fence>>>>;
|
|
pub(super) type SharedError = Arc<Mutex<Option<String>>>;
|
|
|
|
pub(super) struct Event {
|
|
native: Native,
|
|
pub(super) error: Mutex<Option<SharedError>>,
|
|
}
|
|
|
|
// SAFETY: SharedEvent CPU operations are thread-safe native APIs; error pointer
|
|
// replacement is mutex protected. No mutable encoder crosses a thread boundary.
|
|
unsafe impl Send for Event {}
|
|
unsafe impl Sync for Event {}
|
|
|
|
impl Event {
|
|
pub(super) fn new() -> Result<Arc<Self>, String> {
|
|
Ok(Arc::new(Self {
|
|
native: Native::new(unsafe { ds4_gpu_mtplx_event_create() }, "shared event")?,
|
|
error: Mutex::new(None),
|
|
}))
|
|
}
|
|
|
|
pub(super) fn check_error(&self) -> Result<(), String> {
|
|
let error = self.error.lock().unwrap().clone();
|
|
if let Some(error) = error {
|
|
return error.lock().unwrap().take().map_or(Ok(()), Err);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn signal(&self, value: u64) {
|
|
unsafe { ds4_gpu_mtplx_event_signal(self.native.raw(), value) };
|
|
}
|
|
|
|
pub(super) fn wait(&self, value: u64) -> Result<(), String> {
|
|
self.check_error()?;
|
|
self.wait_raw(value);
|
|
self.check_error()
|
|
}
|
|
|
|
pub(super) fn wait_raw(&self, value: u64) {
|
|
unsafe { ds4_gpu_mtplx_event_wait(self.native.raw(), value) };
|
|
}
|
|
|
|
pub(super) fn is_signaled(&self, value: u64) -> bool {
|
|
unsafe { ds4_gpu_mtplx_event_value(self.native.raw()) >= value }
|
|
}
|
|
}
|
|
|
|
fn finish_submission(
|
|
error: &SharedError,
|
|
waited: &[Arc<Event>],
|
|
signaled: &[(Arc<Event>, u64)],
|
|
failure: Option<String>,
|
|
completion: Option<Box<dyn FnOnce() + Send>>,
|
|
) {
|
|
if let Some(completion) = completion {
|
|
completion();
|
|
}
|
|
for event in waited {
|
|
let source = event.error.lock().unwrap().clone();
|
|
// Snapshot first, then release the source lock before replacing ours.
|
|
// Different queues may complete on different Metal callback threads.
|
|
let message = source
|
|
.filter(|source| !Arc::ptr_eq(source, error))
|
|
.and_then(|source| source.lock().unwrap().clone());
|
|
if let Some(message) = message {
|
|
*error.lock().unwrap() = Some(message);
|
|
break;
|
|
}
|
|
}
|
|
if let Some(failure) = &failure {
|
|
error
|
|
.lock()
|
|
.unwrap()
|
|
.get_or_insert_with(|| format!("[METAL] Command buffer execution failed: {failure}."));
|
|
}
|
|
let has_error = error.lock().unwrap().is_some();
|
|
if has_error {
|
|
for (event, _) in signaled {
|
|
*event.error.lock().unwrap() = Some(error.clone());
|
|
}
|
|
}
|
|
// A GPU error skips Metal's encoded signal. Unblock downstream queues only
|
|
// after associating their events with the shared, consumable error state.
|
|
if failure.is_some() {
|
|
for (event, value) in signaled {
|
|
event.signal(*value);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn remove_completed(producers: &Producers, outputs: &HashSet<usize>, fence: &Arc<Fence>) {
|
|
let mut producers = producers.lock().unwrap();
|
|
for resource in outputs {
|
|
if producers
|
|
.get(resource)
|
|
.is_some_and(|f| Arc::ptr_eq(f, fence))
|
|
{
|
|
producers.remove(resource);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn on_completed(command: &Native, callback: impl FnOnce(Option<String>) + Send + 'static) {
|
|
type Callback = Box<dyn FnOnce(Option<String>) + Send>;
|
|
unsafe extern "C" fn completed(context: *mut c_void, error: *const c_char) {
|
|
let callback = unsafe { Box::from_raw(context.cast::<Callback>()) };
|
|
let error = (!error.is_null()).then(|| {
|
|
unsafe { CStr::from_ptr(error) }
|
|
.to_string_lossy()
|
|
.into_owned()
|
|
});
|
|
callback(error);
|
|
}
|
|
let callback: Callback = Box::new(callback);
|
|
unsafe {
|
|
ds4_gpu_mtplx_command_completed(
|
|
command.raw(),
|
|
Box::into_raw(Box::new(callback)).cast(),
|
|
completed,
|
|
)
|
|
};
|
|
}
|
|
|
|
pub(super) fn arch_suffix() -> Option<char> {
|
|
static ARCH: OnceLock<Option<char>> = OnceLock::new();
|
|
*ARCH.get_or_init(|| {
|
|
let mut suffix = unsafe { ds4_gpu_mtplx_arch_suffix() };
|
|
let arch = unsafe { libc::getenv(c"MLX_METAL_GPU_ARCH".as_ptr()) };
|
|
if !arch.is_null()
|
|
&& let Some(last) = unsafe { CStr::from_ptr(arch) }.to_bytes().last()
|
|
{
|
|
suffix = u32::from(*last);
|
|
}
|
|
char::from_u32(suffix)
|
|
})
|
|
}
|
|
|
|
fn limits() -> [i32; 2] {
|
|
static LIMITS: OnceLock<[i32; 2]> = OnceLock::new();
|
|
*LIMITS.get_or_init(|| {
|
|
let defaults = match arch_suffix() {
|
|
Some('p') => [20, 40],
|
|
Some('s' | 'd') => [50, 50],
|
|
_ => [40, 40],
|
|
};
|
|
std::array::from_fn(|i| {
|
|
let name = [c"MLX_MAX_OPS_PER_BUFFER", c"MLX_MAX_MB_PER_BUFFER"][i];
|
|
let value = unsafe { libc::getenv(name.as_ptr()) };
|
|
let default = defaults[i];
|
|
if value.is_null() {
|
|
default
|
|
} else {
|
|
unsafe { libc::atoi(value) }
|
|
}
|
|
})
|
|
})
|
|
}
|
|
|
|
pub(super) struct Encoder {
|
|
command: Option<Native>,
|
|
queue: Native,
|
|
encoding: Option<(Native, Arc<Fence>)>,
|
|
allocator: Arc<allocator::Allocator>,
|
|
attached: usize,
|
|
dependencies: Dependencies,
|
|
temporaries: Vec<(usize, BufferDataHold)>,
|
|
producers: Producers,
|
|
error: SharedError,
|
|
wait_events: Vec<Arc<Event>>,
|
|
signal_events: Vec<(Arc<Event>, u64)>,
|
|
limits: [i32; 2],
|
|
// Diagnostic counts, never used for scheduling.
|
|
barriers: usize,
|
|
waits: usize,
|
|
commits: usize,
|
|
primitive_commits: usize,
|
|
}
|
|
|
|
impl Encoder {
|
|
pub(super) fn new(allocator: Arc<allocator::Allocator>) -> Result<Self, String> {
|
|
let queue = Native::new(unsafe { ds4_gpu_mtplx_queue_create() }, "command queue")?;
|
|
let mut attached = 0;
|
|
unsafe {
|
|
allocator
|
|
.residency
|
|
.attach_new_sets(queue.raw(), &mut attached)
|
|
};
|
|
let command = Native::new(
|
|
unsafe { ds4_gpu_mtplx_command_create(queue.raw()) },
|
|
"command buffer",
|
|
)?;
|
|
Ok(Self {
|
|
queue,
|
|
command: Some(command),
|
|
encoding: None,
|
|
allocator,
|
|
attached,
|
|
dependencies: Dependencies::default(),
|
|
temporaries: Vec::new(),
|
|
producers: Arc::default(),
|
|
error: Arc::default(),
|
|
wait_events: Vec::new(),
|
|
signal_events: Vec::new(),
|
|
limits: limits(),
|
|
barriers: 0,
|
|
waits: 0,
|
|
commits: 0,
|
|
primitive_commits: 0,
|
|
})
|
|
}
|
|
|
|
fn check_error(&self) -> Result<(), String> {
|
|
self.error.lock().unwrap().take().map_or(Ok(()), Err)
|
|
}
|
|
|
|
fn get_encoder(&mut self) -> Result<*mut c_void, String> {
|
|
if self.encoding.is_none() {
|
|
self.check_error()?;
|
|
let command = self
|
|
.command
|
|
.as_ref()
|
|
.ok_or("MTPLX command buffer unavailable")?;
|
|
let encoder = Native::new(
|
|
unsafe { ds4_gpu_mtplx_encoder_create(command.raw()) },
|
|
"concurrent encoder",
|
|
)?;
|
|
let fence = match Native::new(unsafe { ds4_gpu_mtplx_fence_create() }, "fence") {
|
|
Ok(fence) => Arc::new(Fence(fence)),
|
|
Err(error) => {
|
|
unsafe { ds4_gpu_mtplx_encoder_end(encoder.raw()) };
|
|
return Err(error);
|
|
}
|
|
};
|
|
self.encoding = Some((encoder, fence));
|
|
}
|
|
Ok(self.encoding.as_ref().unwrap().0.raw())
|
|
}
|
|
|
|
pub(super) fn dispatch(
|
|
&mut self,
|
|
name: &str,
|
|
bindings: &[Binding<'_>],
|
|
constants: &[MetalConstant],
|
|
grid: [u32; 3],
|
|
threads: [u32; 3],
|
|
individual_threads: bool,
|
|
) -> Result<(), String> {
|
|
if bindings
|
|
.iter()
|
|
.any(|b| b.allocation.is_some() && b.access.is_none())
|
|
{
|
|
return Err("MTPLX encoder requires explicit input/output/raw binding roles".into());
|
|
}
|
|
let encoder = self.get_encoder()?;
|
|
for binding in bindings {
|
|
if let Some(access) = binding.access {
|
|
let resource =
|
|
unsafe { ds4_gpu_mtplx_tensor_allocation(binding.raw.tensor) } as usize;
|
|
self.dependencies.bind(resource, access);
|
|
}
|
|
}
|
|
let barrier = self.dependencies.needs_barrier;
|
|
let result = dispatch_geometry_on(
|
|
name,
|
|
bindings,
|
|
constants,
|
|
grid,
|
|
threads,
|
|
individual_threads,
|
|
Some((self.command.as_ref().unwrap().raw(), encoder, barrier)),
|
|
);
|
|
match result {
|
|
Ok(()) => {
|
|
self.barriers += usize::from(barrier);
|
|
self.dependencies.dispatched();
|
|
Ok(())
|
|
}
|
|
Err(error) => Err(error),
|
|
}
|
|
}
|
|
|
|
pub(super) fn barrier(&mut self) -> Result<(), String> {
|
|
let encoder = self.get_encoder()?;
|
|
unsafe { ds4_gpu_mtplx_encoder_barrier(encoder) };
|
|
self.barriers += 1;
|
|
// Explicit barriers do not change the reference's access bookkeeping.
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn register_output(&mut self, buffer: &Buffer) {
|
|
let resource = unsafe { ds4_gpu_mtplx_tensor_allocation(buffer.raw()) } as usize;
|
|
self.dependencies.register_output(resource);
|
|
}
|
|
|
|
pub(super) fn hold_until_completed(&self, hold: impl Send + 'static) -> Result<(), String> {
|
|
let command = self
|
|
.command
|
|
.as_ref()
|
|
.ok_or("MTPLX command buffer unavailable")?;
|
|
on_completed(command, move |_| drop(hold));
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
fn start_concurrent(&mut self) -> Concurrent<'_> {
|
|
self.dependencies.concurrent = true;
|
|
Concurrent(self)
|
|
}
|
|
|
|
fn add_temporary(&mut self, buffer: &Buffer) {
|
|
let resource = unsafe { ds4_gpu_mtplx_tensor_allocation(buffer.raw()) } as usize;
|
|
self.temporaries.push((resource, buffer.hold_data()));
|
|
}
|
|
|
|
/// Evaluator input/sibling Data is completion-owned, not backend scratch.
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(super) fn defer_data(&self, buffers: &[&Buffer]) -> bool {
|
|
let Some(command) = &self.command else {
|
|
return false;
|
|
};
|
|
let data = buffers
|
|
.iter()
|
|
.map(|buffer| buffer.hold_data())
|
|
.collect::<Vec<_>>();
|
|
on_completed(command, move |_| drop(data));
|
|
true
|
|
}
|
|
|
|
/// gpu::eval runs this once after the whole primitive, not per dispatch.
|
|
pub(super) fn eval_primitive(&mut self, buffers: &[&Buffer]) -> Result<(), String> {
|
|
let data = buffers
|
|
.iter()
|
|
.map(|buffer| buffer.hold_data())
|
|
.collect::<Vec<_>>();
|
|
if self.needs_commit() {
|
|
self.end_encoding();
|
|
let completion = scheduler::global().activity.start();
|
|
self.commit_with(Some(Box::new(move || {
|
|
// Reference task completion is announced before capture release.
|
|
drop(completion);
|
|
drop(data);
|
|
})))?;
|
|
self.primitive_commits += 1;
|
|
} else {
|
|
let command = self
|
|
.command
|
|
.as_ref()
|
|
.ok_or("MTPLX command buffer unavailable")?;
|
|
on_completed(command, move |_| drop(data));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// The fixed indexer tape has one open GPU stream. The general evaluator
|
|
/// must finalize all its open streams before using the same pressure gate.
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(super) fn check_backpressure(&mut self) -> Result<(), String> {
|
|
let activity = &scheduler::global().activity;
|
|
if activity.needs_throttle(self.allocator.stats().active, self.allocator.memory_limit()) {
|
|
self.end_encoding();
|
|
self.commit()?;
|
|
activity.wait_for_one();
|
|
while self.allocator.stats().active > self.allocator.memory_limit()
|
|
&& activity.count() > 0
|
|
{
|
|
activity.wait_for_one();
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn end_encoding(&mut self) {
|
|
let Some((encoder, fence)) = self.encoding.take() else {
|
|
return;
|
|
};
|
|
for (resource, _) in &self.temporaries {
|
|
self.dependencies.all_inputs.remove(resource);
|
|
self.dependencies.all_outputs.remove(resource);
|
|
}
|
|
let mut waiting_on = HashMap::new();
|
|
{
|
|
let mut producers = self.producers.lock().unwrap();
|
|
for resource in &self.dependencies.all_inputs {
|
|
if let Some(prior) = producers.get(resource)
|
|
&& let std::collections::hash_map::Entry::Vacant(entry) =
|
|
waiting_on.entry(Arc::as_ptr(prior) as usize)
|
|
{
|
|
unsafe { ds4_gpu_mtplx_encoder_fence(encoder.raw(), prior.0.raw(), 0) };
|
|
self.waits += 1;
|
|
entry.insert(prior.clone());
|
|
}
|
|
}
|
|
for resource in &self.dependencies.all_outputs {
|
|
producers.insert(*resource, fence.clone());
|
|
}
|
|
}
|
|
unsafe { ds4_gpu_mtplx_encoder_fence(encoder.raw(), fence.0.raw(), 1) };
|
|
let producers = self.producers.clone();
|
|
let outputs = std::mem::take(&mut self.dependencies.all_outputs);
|
|
let temporaries = std::mem::take(&mut self.temporaries);
|
|
on_completed(self.command.as_ref().unwrap(), move |_| {
|
|
remove_completed(&producers, &outputs, &fence);
|
|
drop((temporaries, waiting_on));
|
|
});
|
|
unsafe { ds4_gpu_mtplx_encoder_end(encoder.raw()) };
|
|
self.dependencies.end_encoding();
|
|
}
|
|
|
|
fn needs_commit(&self) -> bool {
|
|
self.dependencies.needs_commit(self.limits)
|
|
}
|
|
|
|
pub(super) fn wait_event(&mut self, event: Arc<Event>, value: u64) -> Result<(), String> {
|
|
self.end_encoding();
|
|
let command = self
|
|
.command
|
|
.as_ref()
|
|
.ok_or("MTPLX command buffer unavailable")?;
|
|
unsafe { ds4_gpu_mtplx_command_event(command.raw(), event.native.raw(), value, 0) };
|
|
self.wait_events.push(event);
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn signal_event(&mut self, event: Arc<Event>, value: u64) -> Result<(), String> {
|
|
self.end_encoding();
|
|
let command = self
|
|
.command
|
|
.as_ref()
|
|
.ok_or("MTPLX command buffer unavailable")?;
|
|
unsafe { ds4_gpu_mtplx_command_event(command.raw(), event.native.raw(), value, 1) };
|
|
self.signal_events.push((event, value));
|
|
Ok(())
|
|
}
|
|
|
|
fn commit(&mut self) -> Result<Native, String> {
|
|
self.commit_with(None)
|
|
}
|
|
|
|
fn commit_with(
|
|
&mut self,
|
|
completion: Option<Box<dyn FnOnce() + Send>>,
|
|
) -> Result<Native, String> {
|
|
// Reference callers end_encoding before commit; do not silently add an
|
|
// encoder boundary here, since that would change scheduling semantics.
|
|
if self.encoding.is_some() {
|
|
return Err("MTPLX commit requires end_encoding".into());
|
|
}
|
|
unsafe {
|
|
self.allocator
|
|
.residency
|
|
.attach_new_sets(self.queue.raw(), &mut self.attached)
|
|
};
|
|
let command = self
|
|
.command
|
|
.take()
|
|
.ok_or("MTPLX command buffer unavailable")?;
|
|
let error = self.error.clone();
|
|
let waited = std::mem::take(&mut self.wait_events);
|
|
let signaled = std::mem::take(&mut self.signal_events);
|
|
on_completed(&command, move |failure| {
|
|
finish_submission(&error, &waited, &signaled, failure, completion);
|
|
});
|
|
unsafe { ds4_gpu_mtplx_command_commit(command.raw()) };
|
|
self.commits += 1;
|
|
match Native::new(
|
|
unsafe { ds4_gpu_mtplx_command_create(self.queue.raw()) },
|
|
"command buffer",
|
|
) {
|
|
Ok(next) => self.command = Some(next),
|
|
Err(error) => {
|
|
// A submitted unretained buffer must drain even if allocating
|
|
// its replacement fails. Never reuse the submitted buffer.
|
|
unsafe { ds4_gpu_mtplx_command_wait(command.raw()) };
|
|
return Err(error);
|
|
}
|
|
}
|
|
self.dependencies.committed();
|
|
Ok(command)
|
|
}
|
|
|
|
pub(super) fn synchronize(&mut self) -> Result<(), String> {
|
|
self.drain()?;
|
|
self.check_error()
|
|
}
|
|
|
|
pub(super) fn finalize(&mut self) -> Result<(), String> {
|
|
self.end_encoding();
|
|
self.commit().map(drop)
|
|
}
|
|
|
|
fn drain(&mut self) -> Result<(), String> {
|
|
self.end_encoding();
|
|
let submitted = self.commit()?;
|
|
unsafe { ds4_gpu_mtplx_command_wait(submitted.raw()) };
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Drop for Encoder {
|
|
fn drop(&mut self) {
|
|
// Drain before queue, allocator and completion-owned state are released.
|
|
// Like the reference destructor, do not rethrow an asynchronous GPU error.
|
|
let _ = self.drain();
|
|
}
|
|
}
|
|
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
struct Concurrent<'a>(&'a mut Encoder);
|
|
impl Deref for Concurrent<'_> {
|
|
type Target = Encoder;
|
|
fn deref(&self) -> &Encoder {
|
|
self.0
|
|
}
|
|
}
|
|
impl DerefMut for Concurrent<'_> {
|
|
fn deref_mut(&mut self) -> &mut Encoder {
|
|
self.0
|
|
}
|
|
}
|
|
impl Drop for Concurrent<'_> {
|
|
fn drop(&mut self) {
|
|
self.0.dependencies.end_concurrent();
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn mtplx_encoder_barrier_epochs_and_commit_thresholds() {
|
|
let mut d = Dependencies::default();
|
|
d.bind(1, Access::Input(8));
|
|
d.bind(2, Access::Output(8));
|
|
assert!(!d.needs_barrier);
|
|
d.dispatched();
|
|
assert_eq!(d.data_size, 16);
|
|
d.bind(2, Access::Input(800)); // Same physical resource counts only once.
|
|
d.bind(3, Access::Output(8));
|
|
assert!(d.needs_barrier, "read after write");
|
|
d.dispatched();
|
|
assert_eq!(d.prev_inputs, HashSet::from([2, 3]));
|
|
assert_eq!(d.prev_outputs, HashSet::from([3]));
|
|
assert_eq!(d.data_size, 24);
|
|
d.bind(1, Access::Output(8));
|
|
assert!(!d.needs_barrier, "resource 1 is outside the current epoch");
|
|
d.dispatched();
|
|
d.bind(2, Access::Output(8));
|
|
assert!(d.needs_barrier, "write after read");
|
|
d.dispatched();
|
|
|
|
let mut raw = Dependencies::default();
|
|
raw.bind(9, Access::Raw);
|
|
raw.dispatched();
|
|
raw.bind(9, Access::Input(1024));
|
|
assert!(!raw.needs_barrier);
|
|
assert_eq!(raw.data_size, 0, "set_buffer already inserted all_inputs");
|
|
assert_eq!(raw.all_outputs, HashSet::from([9]));
|
|
|
|
let mut concurrent = Dependencies {
|
|
concurrent: true,
|
|
..Default::default()
|
|
};
|
|
for _ in 0..2 {
|
|
concurrent.bind(1, Access::Input(8));
|
|
concurrent.bind(2, Access::Output(8));
|
|
assert!(!concurrent.needs_barrier, "concurrent outputs are deferred");
|
|
concurrent.dispatched();
|
|
}
|
|
assert!(concurrent.prev_outputs.is_empty());
|
|
concurrent.end_concurrent();
|
|
concurrent.bind(2, Access::Input(8));
|
|
assert!(concurrent.needs_barrier);
|
|
concurrent.dispatched();
|
|
|
|
let mut threshold = Dependencies {
|
|
ops: 50,
|
|
data_size: (51 << 20) - 1,
|
|
..Default::default()
|
|
};
|
|
assert!(
|
|
!threshold.needs_commit([50, 50]),
|
|
"strict >, truncated Mi elements"
|
|
);
|
|
threshold.data_size += 1;
|
|
assert!(threshold.needs_commit([50, 50]));
|
|
threshold.data_size = 0;
|
|
threshold.ops += 1;
|
|
assert!(threshold.needs_commit([50, 50]));
|
|
threshold.end_encoding();
|
|
assert_eq!(
|
|
threshold.ops, 51,
|
|
"encoder end does not reset command counters"
|
|
);
|
|
threshold.committed();
|
|
assert_eq!((threshold.ops, threshold.data_size), (0, 0));
|
|
assert!(threshold.needs_commit([-1, -1]), "signed op threshold");
|
|
assert!(
|
|
!threshold.needs_commit([50, -1]),
|
|
"reference size_t conversion"
|
|
);
|
|
d.all_outputs.clear(); // Moved into the completion handler in the real encoder.
|
|
let counted = d.data_size;
|
|
let ops = d.ops;
|
|
d.end_encoding();
|
|
d.bind(1, Access::Input(8));
|
|
assert_eq!(d.data_size, counted + 8, "unique inputs reset per encoder");
|
|
assert_eq!(
|
|
d.ops, ops,
|
|
"command counters persist across encoder boundaries"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Apple Metal, no model weights"]
|
|
fn mtplx_encoder_commits_after_primitives_and_balances_completion() {
|
|
configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
crate::engine::reset_gpu_busy_stats();
|
|
crate::engine::gpu_canary_probe().unwrap();
|
|
assert_eq!(crate::engine::gpu_busy_stats().command_buffers, 0);
|
|
let scope = Scope::new().unwrap();
|
|
scope.0.borrow_mut().limits = [1, i32::MAX];
|
|
let input = Buffer::mtplx_bytes(16).unwrap();
|
|
let middle = Buffer::mtplx_bytes(16).unwrap();
|
|
let output = Buffer::mtplx_bytes(16).unwrap();
|
|
let expected = [0x80, 0x3f].repeat(8);
|
|
input.write(0, &expected).unwrap();
|
|
let copy = |from: &Buffer, to: &Buffer| {
|
|
qsa_copy_typed(
|
|
Some(from),
|
|
to,
|
|
[1, 1, 8],
|
|
[0, 0, 1],
|
|
[0, 0, 1],
|
|
"bfloat16",
|
|
8,
|
|
)
|
|
.unwrap();
|
|
};
|
|
copy(&input, &middle);
|
|
with_active(|encoder| encoder.eval_primitive(&[&input]))
|
|
.unwrap()
|
|
.unwrap();
|
|
assert_eq!(scope.commits(), 0, "the threshold is strict, not >=");
|
|
copy(&middle, &output);
|
|
assert_eq!(
|
|
scope.commits(),
|
|
0,
|
|
"dispatch must not commit a partial primitive"
|
|
);
|
|
with_active(|encoder| encoder.eval_primitive(&[&middle]))
|
|
.unwrap()
|
|
.unwrap();
|
|
assert_eq!(scope.primitive_commits(), 1);
|
|
assert_eq!(scope.commits(), 1);
|
|
scope.synchronize().unwrap();
|
|
let busy = crate::engine::gpu_busy_stats();
|
|
assert_eq!(busy.command_buffers, scope.commits() as u64);
|
|
assert!(busy.busy_seconds > 0.0);
|
|
assert!(busy.max_busy_seconds > 0.0);
|
|
assert_eq!(scheduler::global().activity.count(), 0);
|
|
assert!(input.unique_data());
|
|
assert!(middle.unique_data());
|
|
let mut actual = [0; 16];
|
|
output.read(0, &mut actual).unwrap();
|
|
assert_eq!(actual.as_slice(), expected);
|
|
|
|
// One primitive may encode several dependent kernels before the check.
|
|
copy(&input, &middle);
|
|
copy(&middle, &output);
|
|
copy(&output, &middle);
|
|
assert_eq!(scope.commits(), 2);
|
|
with_active(|encoder| encoder.eval_primitive(&[&input, &output]))
|
|
.unwrap()
|
|
.unwrap();
|
|
assert_eq!(scope.primitive_commits(), 2);
|
|
scope.synchronize().unwrap();
|
|
assert_eq!(
|
|
crate::engine::gpu_busy_stats().command_buffers,
|
|
scope.commits() as u64,
|
|
"record every MTPLX command once, without counting its later waits"
|
|
);
|
|
assert_eq!(scheduler::global().activity.count(), 0);
|
|
middle.read(0, &mut actual).unwrap();
|
|
assert_eq!(actual.as_slice(), expected);
|
|
|
|
// Fail before callback transfer, without submitting invalid GPU work.
|
|
let mut encoder = scope.0.borrow_mut();
|
|
let command = encoder.command.take();
|
|
encoder.dependencies.ops = 2;
|
|
assert_eq!(
|
|
encoder.eval_primitive(&[]),
|
|
Err("MTPLX command buffer unavailable".into())
|
|
);
|
|
assert_eq!(scheduler::global().activity.count(), 0);
|
|
encoder.command = command;
|
|
encoder.dependencies.committed();
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Apple Metal, no model weights"]
|
|
fn mtplx_encoder_scope_routes_operators_and_allocator() {
|
|
configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
let scope = Scope::new().unwrap();
|
|
assert!(
|
|
Scope::new().is_err(),
|
|
"nested scopes cannot replace the live stream"
|
|
);
|
|
let allocator = active_allocator().unwrap();
|
|
let input = Buffer::mtplx_bytes(3 * 64 * 4).unwrap();
|
|
input.write(0, &[0x80, 0x3f].repeat(3 * 64 * 2)).unwrap();
|
|
let output = Buffer::mtplx_bytes(3 * 64 * 2).unwrap();
|
|
let activated = gate_up_bf16(&input, &output, [3, 64]).unwrap();
|
|
{
|
|
let encoder = scope.0.borrow();
|
|
assert_eq!(encoder.dependencies.ops, 2);
|
|
assert_eq!(encoder.barriers, 1);
|
|
assert_eq!(
|
|
encoder.dependencies.data_size,
|
|
11 * 64,
|
|
"slice span (2*rows-1)*width plus two contiguous outputs, not full allocation bytes"
|
|
);
|
|
assert_eq!(encoder.commits, 0, "scope must not invent operator commits");
|
|
}
|
|
let bytes = [&input, &output, &activated]
|
|
.iter()
|
|
.map(|buffer| buffer.allocation_bytes() as usize)
|
|
.sum::<usize>();
|
|
assert_eq!(
|
|
allocator.stats().active,
|
|
bytes,
|
|
"all operator storage uses the pooled allocator"
|
|
);
|
|
scope.synchronize().unwrap();
|
|
drop((input, output, activated));
|
|
assert_eq!(allocator.stats().active, 0);
|
|
let failed_copy = || -> Result<(), String> {
|
|
let _concurrent = concurrent_scope();
|
|
assert!(scope.0.borrow().dependencies.concurrent);
|
|
Err("synthetic concatenate encoding failure".into())
|
|
};
|
|
assert!(failed_copy().is_err());
|
|
assert!(
|
|
!scope.0.borrow().dependencies.concurrent,
|
|
"an early return must restore the dependency join"
|
|
);
|
|
let source = Buffer::mtplx_bytes(32).unwrap();
|
|
let target = Buffer::mtplx_bytes(64).unwrap();
|
|
source.write(0, &[0x80, 0x3f].repeat(16)).unwrap();
|
|
assert!(qsa_dynamic_copy_rows([&source, &target], [8; 2], [1, 8], [None; 2], [7, 32]).is_err());
|
|
qsa_dynamic_copy_rows([&source, &target], [8; 2], [1, 8], [None; 2], [16, 32]).unwrap();
|
|
assert_eq!(
|
|
scope.0.borrow().dependencies.data_size,
|
|
48,
|
|
"copy counts bound array data_size, not the eight copied elements"
|
|
);
|
|
scope.synchronize().unwrap();
|
|
let mut copied = [0; 16];
|
|
target.read(0, &mut copied).unwrap();
|
|
assert_eq!(copied.as_slice(), [0x80, 0x3f].repeat(8));
|
|
drop((source, target));
|
|
assert!(
|
|
allocator.stats().cached > 0,
|
|
"completed storage returns to the same cache"
|
|
);
|
|
drop(scope);
|
|
assert!(
|
|
active_allocator().is_none(),
|
|
"scope teardown removes routing"
|
|
);
|
|
let outside = Buffer::mtplx_bytes(128).unwrap();
|
|
assert_eq!(
|
|
allocator.stats().active,
|
|
0,
|
|
"legacy work must not inherit the allocator"
|
|
);
|
|
drop(outside);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Apple Metal; empty original-kernel grids and nil bindings"]
|
|
fn mtplx_empty_bindings_preserve_nil_dependencies_without_dummy_storage() {
|
|
configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
{
|
|
let scope = Scope::new().unwrap();
|
|
let allocator = active_allocator().unwrap();
|
|
let real = Buffer::mtplx_bytes(16).unwrap();
|
|
let before = allocator.stats();
|
|
let empty = Buffer::mtplx_bytes(0).unwrap();
|
|
let another = Buffer::mtplx_bytes(0).unwrap();
|
|
let nthreads = 0_u32;
|
|
for binding in [tensor(0, &empty).output(0), tensor(0, &another).input(0)] {
|
|
scope
|
|
.0
|
|
.borrow_mut()
|
|
.dispatch(
|
|
"input_coherent",
|
|
&[binding, bytes(1, &nthreads)],
|
|
&[],
|
|
[0, 1, 1],
|
|
[1024, 1, 1],
|
|
false,
|
|
)
|
|
.unwrap();
|
|
}
|
|
{
|
|
let encoder = scope.0.borrow();
|
|
assert_eq!(encoder.dependencies.data_size, 0);
|
|
assert_eq!(encoder.dependencies.ops, 2);
|
|
assert_eq!(encoder.dependencies.all_inputs, HashSet::from([0]));
|
|
assert_eq!(encoder.dependencies.all_outputs, HashSet::from([0]));
|
|
assert_eq!(
|
|
encoder.barriers, 1,
|
|
"null resources participate in RAW dependencies"
|
|
);
|
|
}
|
|
assert_eq!(allocator.stats(), before);
|
|
// A zero-length slice can also retain a non-null allocation and point
|
|
// exactly one past its end. Do not substitute null Data for that view.
|
|
let tail = real.view(16, 0).unwrap();
|
|
assert_eq!(tail.data_identity(), real.data_identity());
|
|
assert_eq!(tail.allocation_bytes(), 16);
|
|
scope
|
|
.0
|
|
.borrow_mut()
|
|
.dispatch(
|
|
"input_coherent",
|
|
&[tensor(0, &tail).input(0), bytes(1, &nthreads)],
|
|
&[],
|
|
[0, 1, 1],
|
|
[1024, 1, 1],
|
|
false,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(scope.0.borrow().dependencies.data_size, 0);
|
|
drop((empty, another, tail));
|
|
scope.synchronize().unwrap();
|
|
assert_eq!(allocator.stats(), before);
|
|
}
|
|
// Also cover the native fallback owner path, which has an Objective-C
|
|
// resource set. nil is a valid Metal binding but not an NSSet element.
|
|
let empty = Buffer::mtplx_bytes(0).unwrap();
|
|
let nthreads = 0_u32;
|
|
let bindings = [tensor(0, &empty).input(0).raw, bytes(1, &nthreads).raw];
|
|
assert_ne!(
|
|
unsafe {
|
|
ds4_gpu_mtplx_dispatch(
|
|
c"input_coherent".as_ptr(),
|
|
bindings.as_ptr(),
|
|
2,
|
|
std::ptr::null(),
|
|
0,
|
|
[0, 1, 1].as_ptr(),
|
|
[1024, 1, 1].as_ptr(),
|
|
0,
|
|
std::ptr::null_mut(),
|
|
None,
|
|
)
|
|
},
|
|
0
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Apple Metal, no model weights"]
|
|
fn mtplx_encoder_untracked_dependencies_across_encoders_and_commands() {
|
|
configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
let allocator = allocator::Allocator::new().unwrap();
|
|
allocator.set_wired_limit(8 << 20).unwrap();
|
|
let mut enc = Encoder::new(allocator.clone()).unwrap();
|
|
assert_ne!(enc.queue.raw(), unsafe { ds4_gpu_mtplx_command_queue() });
|
|
assert_eq!(enc.attached, allocator.residency.num_sets());
|
|
const N: u32 = 4096;
|
|
let a = allocator.allocate((N * 2) as usize).unwrap().unwrap();
|
|
let x = allocator.allocate((N * 2) as usize).unwrap().unwrap();
|
|
let y = allocator.allocate((N * 2) as usize).unwrap().unwrap();
|
|
let expected = (0..N)
|
|
.flat_map(|i| (0x3f00_u16 + (i % 128) as u16).to_le_bytes())
|
|
.collect::<Vec<_>>();
|
|
a.write(0, &expected).unwrap();
|
|
let copy = |enc: &mut Encoder, source: &Buffer, output: &Buffer, count: u32| {
|
|
enc.dispatch(
|
|
"v_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, source).input(count as usize),
|
|
tensor(1, output).output(count as usize),
|
|
bytes(2, &count),
|
|
],
|
|
&[],
|
|
[count, 1, 1],
|
|
[count.min(1024), 1, 1],
|
|
true,
|
|
)
|
|
.unwrap();
|
|
};
|
|
assert!(
|
|
enc.dispatch(
|
|
"v_copybfloat16bfloat16",
|
|
&[tensor(0, &a), tensor(1, &x), bytes(2, &N)],
|
|
&[],
|
|
[N, 1, 1],
|
|
[1024, 1, 1],
|
|
true
|
|
)
|
|
.is_err()
|
|
);
|
|
assert!(
|
|
enc.encoding.is_none(),
|
|
"unclassified bindings cannot silently use serial semantics"
|
|
);
|
|
{
|
|
let mut concurrent = enc.start_concurrent();
|
|
for offset in [0, N] {
|
|
let input = a.view(u64::from(offset), u64::from(N)).unwrap();
|
|
let output = x.view(u64::from(offset), u64::from(N)).unwrap();
|
|
copy(&mut concurrent, &input, &output, N / 2);
|
|
}
|
|
}
|
|
assert_eq!(enc.barriers, 0);
|
|
copy(&mut enc, &x, &y, N);
|
|
assert_eq!(
|
|
enc.barriers, 1,
|
|
"join concurrent writers before their reader"
|
|
);
|
|
let previous = enc.dependencies.prev_outputs.clone();
|
|
enc.barrier().unwrap();
|
|
assert_eq!(
|
|
enc.dependencies.prev_outputs, previous,
|
|
"explicit barrier preserves bookkeeping"
|
|
);
|
|
assert_eq!(enc.barriers, 2);
|
|
enc.end_encoding();
|
|
assert_eq!(enc.waits, 0);
|
|
assert_eq!(enc.dependencies.ops, 3);
|
|
copy(&mut enc, &y, &x, N);
|
|
enc.end_encoding();
|
|
assert_eq!(enc.waits, 1, "two prior resources share one fence wait");
|
|
assert_eq!(enc.dependencies.ops, 4);
|
|
let submitted = enc.commit().unwrap(); // Deliberately no CPU wait.
|
|
assert_eq!(enc.dependencies.ops, 0);
|
|
assert!(!enc.needs_commit());
|
|
for _ in 0..6 {
|
|
copy(&mut enc, &x, &y, N);
|
|
copy(&mut enc, &y, &x, N);
|
|
enc.end_encoding();
|
|
drop(enc.commit().unwrap());
|
|
}
|
|
enc.synchronize().unwrap();
|
|
drop(submitted);
|
|
assert!(enc.producers.lock().unwrap().is_empty());
|
|
let mut actual = vec![0; expected.len()];
|
|
x.read(0, &mut actual).unwrap();
|
|
assert_eq!(actual, expected);
|
|
|
|
// Completion must not erase a newer writer's entry. Exercise the exact
|
|
// callback helper deterministically, independent of GPU completion timing.
|
|
let old = Arc::new(Fence(
|
|
Native::new(unsafe { ds4_gpu_mtplx_fence_create() }, "fence").unwrap(),
|
|
));
|
|
let new = Arc::new(Fence(
|
|
Native::new(unsafe { ds4_gpu_mtplx_fence_create() }, "fence").unwrap(),
|
|
));
|
|
let producers: Producers = Arc::default();
|
|
producers.lock().unwrap().insert(123, new.clone());
|
|
remove_completed(&producers, &HashSet::from([123]), &old);
|
|
assert!(Arc::ptr_eq(
|
|
producers.lock().unwrap().get(&123).unwrap(),
|
|
&new
|
|
));
|
|
remove_completed(&producers, &HashSet::from([123]), &new);
|
|
assert!(producers.lock().unwrap().is_empty());
|
|
|
|
// Temporaries do not enter the inter-encoder producer map. They stay alive
|
|
// through completion even after their descriptors and binding views drop.
|
|
let temporary = allocator.allocate((N * 2) as usize).unwrap().unwrap();
|
|
let resource = unsafe { ds4_gpu_mtplx_tensor_allocation(temporary.raw()) } as usize;
|
|
copy(&mut enc, &a, &temporary, N);
|
|
enc.add_temporary(&temporary);
|
|
drop(temporary);
|
|
enc.end_encoding();
|
|
assert!(!enc.producers.lock().unwrap().contains_key(&resource));
|
|
enc.synchronize().unwrap();
|
|
|
|
// Raw buffers match set_buffer: cross-encoder input AND output, without
|
|
// automatic intra-encoder barriers or data_size accounting.
|
|
enc.dispatch(
|
|
"v_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, &a).raw_buffer(),
|
|
tensor(1, &x).raw_buffer(),
|
|
bytes(2, &N),
|
|
],
|
|
&[],
|
|
[N, 1, 1],
|
|
[1024, 1, 1],
|
|
true,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(enc.dependencies.data_size, 0);
|
|
enc.end_encoding();
|
|
assert_eq!(enc.producers.lock().unwrap().len(), 2);
|
|
enc.synchronize().unwrap();
|
|
assert!(enc.producers.lock().unwrap().is_empty());
|
|
|
|
// Destruction must submit and wait with the allocator alive.
|
|
copy(&mut enc, &a, &y, N);
|
|
drop(enc);
|
|
y.read(0, &mut actual).unwrap();
|
|
assert_eq!(actual, expected);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Apple Metal, no model weights"]
|
|
fn mtplx_encoder_shared_events_and_error_completion_order() {
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
let allocator = allocator::Allocator::new().unwrap();
|
|
allocator.set_wired_limit(8 << 20).unwrap();
|
|
let input = allocator.allocate(128).unwrap().unwrap();
|
|
let middle = allocator.allocate(128).unwrap().unwrap();
|
|
let output = allocator.allocate(128).unwrap().unwrap();
|
|
let expected = [0x80, 0x3f].repeat(64);
|
|
input.write(0, &expected).unwrap();
|
|
let mut producer = Encoder::new(allocator.clone()).unwrap();
|
|
let mut consumer = Encoder::new(allocator).unwrap();
|
|
let ready = Event::new().unwrap();
|
|
let done = Event::new().unwrap();
|
|
assert!(!ready.is_signaled(1));
|
|
assert!(!done.is_signaled(2));
|
|
let copy = |enc: &mut Encoder, src: &Buffer, dst: &Buffer| {
|
|
enc.dispatch(
|
|
"v_copybfloat16bfloat16",
|
|
&[
|
|
tensor(0, src).input(64),
|
|
tensor(1, dst).output(64),
|
|
bytes(2, &64_u32),
|
|
],
|
|
&[],
|
|
[2, 1, 1],
|
|
[32, 1, 1],
|
|
false,
|
|
)
|
|
.unwrap();
|
|
};
|
|
copy(&mut producer, &input, &middle);
|
|
producer.signal_event(ready.clone(), 1).unwrap();
|
|
assert!(
|
|
producer.encoding.is_none(),
|
|
"signal ends the active encoder"
|
|
);
|
|
let completed = Arc::new(AtomicUsize::new(0));
|
|
let count = completed.clone();
|
|
drop(
|
|
producer
|
|
.commit_with(Some(Box::new(move || {
|
|
count.fetch_add(1, Ordering::SeqCst);
|
|
})))
|
|
.unwrap(),
|
|
);
|
|
consumer.wait_event(ready.clone(), 1).unwrap();
|
|
copy(&mut consumer, &middle, &output);
|
|
consumer.signal_event(done.clone(), 2).unwrap();
|
|
consumer.synchronize().unwrap();
|
|
producer.synchronize().unwrap();
|
|
assert_eq!(completed.load(Ordering::SeqCst), 1);
|
|
assert!(ready.is_signaled(1));
|
|
assert!(done.is_signaled(2));
|
|
done.wait(2).unwrap();
|
|
let mut actual = vec![0; 128];
|
|
output.read(0, &mut actual).unwrap();
|
|
assert_eq!(
|
|
actual, expected,
|
|
"untracked buffer crosses queues via shared event"
|
|
);
|
|
|
|
// Invoke the real completion policy with a synthetic status error; never
|
|
// deliberately hang/fault the GPU to exercise downstream unblocking.
|
|
let first = Event::new().unwrap();
|
|
let second = Event::new().unwrap();
|
|
let failed = Event::new().unwrap();
|
|
let error: SharedError = Arc::new(Mutex::new(Some("local".into())));
|
|
*first.error.lock().unwrap() = Some(Arc::new(Mutex::new(Some("first upstream".into()))));
|
|
*second.error.lock().unwrap() = Some(Arc::new(Mutex::new(Some("second upstream".into()))));
|
|
let observed = error.clone();
|
|
finish_submission(
|
|
&error,
|
|
&[first, second],
|
|
&[(failed.clone(), 3)],
|
|
Some("GPU failure".into()),
|
|
Some(Box::new(move || {
|
|
assert_eq!(observed.lock().unwrap().as_deref(), Some("local"))
|
|
})),
|
|
);
|
|
assert!(
|
|
failed.is_signaled(3),
|
|
"GPU error manually signals downstream event"
|
|
);
|
|
assert_eq!(error.lock().unwrap().as_deref(), Some("first upstream"));
|
|
assert_eq!(failed.wait(3).unwrap_err(), "first upstream");
|
|
assert!(
|
|
error.lock().unwrap().is_none(),
|
|
"event consumes the shared encoder error"
|
|
);
|
|
failed.check_error().unwrap();
|
|
|
|
// A self-referencing waited event is skipped, not recursively locked. A
|
|
// non-GPU upstream error poisons signals but does not signal them on CPU.
|
|
let self_event = Event::new().unwrap();
|
|
*self_event.error.lock().unwrap() = Some(error.clone());
|
|
let downstream = Event::new().unwrap();
|
|
*error.lock().unwrap() = Some("existing".into());
|
|
finish_submission(
|
|
&error,
|
|
&[self_event],
|
|
&[(downstream.clone(), 4)],
|
|
None,
|
|
None,
|
|
);
|
|
assert!(!downstream.is_signaled(4));
|
|
assert_eq!(downstream.check_error().unwrap_err(), "existing");
|
|
finish_submission(&error, &[], &[], Some("synthetic".into()), None);
|
|
assert_eq!(
|
|
error.lock().unwrap().as_deref(),
|
|
Some("[METAL] Command buffer execution failed: synthetic.")
|
|
);
|
|
let retained_error = producer.error.clone();
|
|
*retained_error.lock().unwrap() = Some("destructor must preserve this".into());
|
|
drop(producer);
|
|
assert_eq!(
|
|
retained_error.lock().unwrap().as_deref(),
|
|
Some("destructor must preserve this")
|
|
);
|
|
}
|