use super::*;
use crate::encoder::RecordedBinding;
use crate::{
GrowthBucket, GrowthPoolKey, GrowthQuota, LegacyGrowthLimits, LegacyPoolGrowthAuthority,
MlxBufferPool,
};
use metal::foreign_types::ForeignType;
fn install(bytes: usize, roots: usize, cache: usize) -> RootAuthorityHandle {
NativeRootAuthority::install_early(RootCreationLimits {
max_attempted_requested_bytes: bytes,
max_root_attempts: roots,
max_raw_constructor_attempts: roots,
max_registered_devices: 1,
max_rope_cache_entries: cache,
})
.expect("fresh process authority")
}
fn history(a: &RootAuthorityHandle, bytes: usize, roots: usize) -> RootCreationSnapshot {
let s = a.snapshot().unwrap();
assert!(s.covered && s.unknown.is_none());
assert_eq!(
(
s.attempted_requested_bytes,
s.root_attempts,
s.raw_attempt_reservations,
s.successful_roots,
s.inflight
),
(bytes, roots, roots, roots, 0)
);
assert_eq!(s.registered_devices, 1);
s
}
fn known(b: &MlxBuffer, bytes: usize) -> usize {
match b.root_origin_snapshot() {
RootOriginSnapshot::KnownNative {
id,
requested_bytes,
} => {
assert_eq!(requested_bytes, bytes);
assert_ne!(id, 0);
id
}
other => panic!("expected known root, got {other:?}"),
}
}
#[test]
#[ignore = "separate supervised CPU-only 6 roots/416 requested bytes"]
fn actual_shared_factory_composition() {
let a = install(416, 6, 0);
let d = MlxDevice::new().unwrap();
a.register_device(&d).unwrap();
a.register_device(&d.clone()).unwrap();
let source = [0x5au8; 96];
let first = d.alloc_buffer(64, DType::U8, vec![64]).unwrap();
assert!(first.as_slice::<u8>().unwrap().iter().all(|x| *x == 0));
let mut ids = [0usize; 6];
ids[0] = known(&first, 64);
history(&a, 64, 1);
let mut second = unsafe { d.alloc_buffer_for_overwrite(64, DType::U8, vec![64]) }.unwrap();
second.as_mut_slice::<u8>().unwrap().fill(0x5a);
assert_eq!(second.as_slice::<u8>().unwrap(), &source[..64]);
ids[1] = known(&second, 64);
history(&a, 128, 2);
let third =
crate::weight::safetensors_to_metal_buffer(&d, &source[..64], DType::U8, vec![64]).unwrap();
assert_eq!(third.as_slice::<u8>().unwrap(), &source[..64]);
ids[2] = known(&third, 64);
history(&a, 192, 3);
let mut pool = MlxBufferPool::new();
let fourth = pool.alloc(&d, 64, DType::U8, vec![64]).unwrap();
assert!(fourth.as_slice::<u8>().unwrap().iter().all(|x| *x == 0));
ids[3] = known(&fourth, 64);
history(&a, 256, 4);
let quota = GrowthQuota {
bytes: 64,
roots: 1,
attempts: 1,
};
let growth = LegacyPoolGrowthAuthority::install(LegacyGrowthLimits {
quota,
max_pools: 1,
max_buckets_per_pool: 1,
max_rank: 1,
max_batch: 1,
})
.unwrap();
let key = GrowthPoolKey {
kind: 57,
worker_generation: 1,
};
let mut limited = growth
.bind(&d, key, quota, &[GrowthBucket { bytes: 64, quota }])
.unwrap();
let fifth = limited.alloc(&d, 64, DType::U8, vec![64]).unwrap();
assert!(fifth.as_slice::<u8>().unwrap().iter().all(|x| *x == 0));
ids[4] = known(&fifth, 64);
history(&a, 320, 5);
let local = limited.snapshot().unwrap();
assert_eq!(
(local.pool.bytes, local.pool.roots, local.pool.attempts),
(64, 1, 1)
);
let sixth = d.upload_shared_bytes(&source, DType::U8, vec![96]).unwrap();
assert_eq!(sixth.as_slice::<u8>().unwrap(), &source);
ids[5] = known(&sixth, 96);
history(&a, 416, 6);
for i in 0..ids.len() {
for j in 0..i {
assert_ne!(ids[i], ids[j]);
}
}
let view = first.slice_view(0, 16);
assert_eq!(known(&view, 64), ids[0]);
drop(view);
drop(first);
drop(second);
drop(third);
drop(sixth);
pool.release(fourth);
drop(pool);
drop(fifth);
unsafe { limited.reset(&d) }.unwrap();
limited.clear(&d).unwrap();
let before = history(&a, 416, 6);
assert_eq!(before.cache_slot_history, 0);
assert!(d.alloc_buffer(64, DType::U8, vec![64]).is_err());
let after = history(&a, 416, 6);
assert_eq!(after.denied, before.denied + 1);
println!("NATIVE57_ACTUAL_SHARED_COMPLETE roots=6 bytes=416 gpu=false");
}
#[test]
#[ignore = "separate supervised CPU-only one-page mapped lifetime control"]
fn actual_mapped_root_lineage_and_backing() {
use std::io::{Seek, SeekFrom, Write};
use std::os::unix::fs::{DirBuilderExt, MetadataExt, OpenOptionsExt};
extern "C" {
fn getpagesize() -> std::ffi::c_int;
}
let page = unsafe { getpagesize() };
assert_eq!(page, 16384);
let page = page as usize;
let a = install(page, 1, 0);
let dir = std::env::temp_dir().join(format!("native57-mapped-{}", std::process::id()));
std::fs::DirBuilder::new().mode(0o700).create(&dir).unwrap();
struct Directory(std::path::PathBuf);
impl Drop for Directory {
fn drop(&mut self) {
let _ = std::fs::remove_dir(&self.0);
}
}
let directory = Directory(dir.clone());
let path = dir.join("fixture");
let mut file = std::fs::OpenOptions::new()
.read(true)
.write(true)
.create_new(true)
.mode(0o600)
.open(&path)
.unwrap();
struct Cleanup(std::path::PathBuf);
impl Drop for Cleanup {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.0);
}
}
let cleanup = Cleanup(path.clone());
let block = [0x39u8; 256];
for _ in 0..page / 256 {
file.write_all(&block).unwrap();
}
file.sync_all().unwrap();
file.seek(SeekFrom::Start(0)).unwrap();
let before = file.metadata().unwrap();
assert_eq!(before.len(), page as u64);
drop(file);
let file = std::fs::File::open(&path).unwrap();
let after = file.metadata().unwrap();
assert_eq!(
(before.dev(), before.ino(), before.len()),
(after.dev(), after.ino(), after.len())
);
let escrow = Arc::new(unsafe { memmap2::MmapOptions::new().len(page).map(&file) }.unwrap());
let weak = Arc::downgrade(&escrow);
objc::rc::autoreleasepool(|| {
let d = MlxDevice::new().unwrap();
a.register_device(&d).unwrap();
let b = d
.map_file_buffer(escrow.clone(), 4096, 4096, DType::U8, vec![4096])
.unwrap();
assert_eq!(
(b.byte_offset(), b.data_byte_len(), b.byte_len()),
(4096, 4096, page)
);
known(&b, page);
assert!(b.as_slice::<u8>().unwrap().iter().all(|x| *x == 0x39));
let raw = b.metal_buffer().clone();
let recorded = RecordedBinding::Buffer {
metal_buffer: raw.clone(),
offset: b.byte_offset(),
};
drop(b);
if let Some(set) = d.residency_set() {
set.flush_pending();
}
assert!(Arc::strong_count(&escrow) > 1);
drop(raw);
assert!(Arc::strong_count(&escrow) > 1);
drop(recorded);
if let Some(set) = d.residency_set() {
set.flush_pending();
}
history(&a, page, 1);
});
assert_eq!(Arc::strong_count(&escrow), 1);
drop(escrow);
assert!(weak.upgrade().is_none());
drop(file);
std::fs::remove_file(&path).unwrap();
drop(cleanup);
std::fs::remove_dir(&dir).unwrap();
drop(directory);
println!(
"NATIVE57_ACTUAL_MAPPED_COMPLETE roots=1 bytes=16384 logical=4096 offset=4096 gpu=false"
);
}
#[test]
#[ignore = "separate supervised CPU-only already-charged raw import control"]
fn actual_raw_clone_import_is_unknown() {
let a = install(128, 2, 0);
let d = MlxDevice::new().unwrap();
a.register_device(&d).unwrap();
let source = [0x29u8; 64];
let b = d.upload_shared_bytes(&source, DType::U8, vec![64]).unwrap();
assert_eq!(b.as_slice::<u8>().unwrap(), &source);
known(&b, 64);
history(&a, 64, 1);
let raw = b.metal_buffer().clone();
let identity = raw.as_ptr();
let imported = MlxBuffer::from_raw(raw, DType::U8, vec![64]);
assert_eq!(imported.metal_buffer().as_ptr(), identity);
assert_eq!(
imported.root_origin_snapshot(),
RootOriginSnapshot::UnknownImport
);
let s = a.snapshot().unwrap();
assert!(!s.covered && s.unknown.is_some());
assert_eq!(
(
s.attempted_requested_bytes,
s.root_attempts,
s.raw_attempt_reservations,
s.successful_roots,
s.inflight
),
(64, 1, 1, 1, 0)
);
assert!(d.alloc_buffer(64, DType::U8, vec![64]).is_err());
let t = a.snapshot().unwrap();
assert!(!t.covered);
assert_eq!(t.unknown, s.unknown);
assert_eq!(
(t.raw_attempt_reservations, t.inflight),
(s.raw_attempt_reservations, 0)
);
assert_eq!(
(
t.attempted_requested_bytes,
t.root_attempts,
t.successful_roots
),
(64, 1, 1)
);
drop(imported);
drop(b);
println!("NATIVE57_ACTUAL_IMPORT_COMPLETE roots=1 bytes=64 covered=false gpu=false");
}
#[test]
#[ignore = "separate supervised zero-payload strict heap refusal control"]
fn actual_strict_heap_refuses_before_creation() {
let a = install(0, 0, 0);
let d = MlxDevice::new().unwrap();
a.register_device(&d).unwrap();
let authority = crate::ProcessPoolAuthority::install(crate::ProcessPoolLimits {
max_domains: 1,
max_reserved_heap_bytes: 65536,
max_child_creations: 1,
max_constructor_attempts: 1,
max_rank: 1,
max_batch: 1,
})
.unwrap();
let before = history(&a, 0, 0);
let result = authority.create_pool(
&d,
crate::BoundedPoolConfig {
requested_heap_bytes: 65536,
reserved_heap_bytes: 65536,
max_child_creations: 1,
},
);
assert!(matches!(result, Err(crate::BoundedPoolError::Unsupported)));
let after = history(&a, 0, 0);
assert_eq!(after.denied, before.denied + 1);
assert_eq!(authority.snapshot().unwrap().domains, 0);
println!("NATIVE57_ACTUAL_HEAP_COMPLETE roots=0 bytes=0 unsupported=true gpu=false");
}
#[test]
#[ignore = "separate supervised CPU-only 32-byte RoPE cold/hit/cache-limit control"]
fn actual_rope_cache_combined_reservation() {
let a = install(64, 2, 1);
let d = MlxDevice::new().unwrap();
a.register_device(&d).unwrap();
let inspect = |b: &MlxBuffer| {
let certificate = known(b, 32);
let values = b.as_slice::<f32>().unwrap();
assert_eq!(values.len(), 8);
for (pair, value) in values.iter().enumerate() {
let expected = 10.0f64.powf(-(pair as f64) / 2.0);
assert!(value.is_finite());
assert!((f64::from(*value) - expected).abs() <= 1e-6 * expected.max(1e-12));
}
(b.metal_buffer().as_ptr() as usize, certificate)
};
let identity =
crate::ops::native57_control_with_inv_freqs(d.metal_device(), 10000.0, 16, 8, inspect)
.unwrap();
let cold = history(&a, 32, 1);
assert_eq!(
(cold.cache_slot_history, cold.published_cache_entries),
(1, 1)
);
let hit =
crate::ops::native57_control_with_inv_freqs(d.metal_device(), 10000.0, 16, 8, inspect)
.unwrap();
assert_eq!(identity, hit);
let hit_state = history(&a, 32, 1);
assert_eq!(
(
hit_state.cache_slot_history,
hit_state.published_cache_entries,
hit_state.denied
),
(
cold.cache_slot_history,
cold.published_cache_entries,
cold.denied
)
);
assert!(
crate::ops::native57_control_with_inv_freqs(d.metal_device(), 20000.0, 16, 8, |_| ())
.is_err()
);
let denied = history(&a, 32, 1);
assert_eq!(
(denied.cache_slot_history, denied.published_cache_entries),
(1, 1)
);
assert_eq!(denied.denied, cold.denied + 1);
let old =
crate::ops::native57_control_with_inv_freqs(d.metal_device(), 10000.0, 16, 8, inspect)
.unwrap();
assert_eq!(old, identity);
let final_hit = history(&a, 32, 1);
assert_eq!(
(
final_hit.cache_slot_history,
final_hit.published_cache_entries,
final_hit.denied
),
(
denied.cache_slot_history,
denied.published_cache_entries,
denied.denied
)
);
println!("NATIVE57_ACTUAL_ROPE_COMPLETE roots=1 bytes=32 cache=1 gpu=false");
}