Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 15 additions & 22 deletions crates/onnx-runtime-ep-cuda/src/provider.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1093,7 +1093,10 @@ impl ExecutionProvider for CudaExecutionProvider {
Some(virtual_backing) => virtual_backing
.allocate_committed(size, alignment, committed_ranges)
.map_err(EpError::Memory)?,
None => self.memory().allocate(size, alignment).map_err(EpError::Memory)?,
None => self
.memory()
.allocate(size, alignment)
.map_err(EpError::Memory)?,
};
self.ep_allocations.fetch_add(1, Ordering::Relaxed);
// SAFETY: `ptr` is a fresh, unique, non-null device allocation of
Expand Down Expand Up @@ -1651,9 +1654,7 @@ impl ExecutionProvider for CudaExecutionProvider {
bytes: u64,
role: onnx_runtime_memory_governor::MemoryRole,
) -> Result<Option<onnx_runtime_memory_governor::MappedGrowthGrant>> {
if bytes == 0
|| !dynamic_lending_enabled()
|| self.memory().as_virtual_backing().is_none()
if bytes == 0 || !dynamic_lending_enabled() || self.memory().as_virtual_backing().is_none()
{
return Ok(None);
}
Expand Down Expand Up @@ -2151,12 +2152,10 @@ extern "C" __global__ void write_after_delay(unsigned int* out, long long spin)
let counting_allocator =
Arc::new(crate::device_allocator::CudaDeviceAllocator::new(context));

let error = provider
.with_memory(counting_allocator.clone())
.expect_err(
"with_memory must return Err once a VMM arena is already selected, not Ok(Self) \
let error = provider.with_memory(counting_allocator.clone()).expect_err(
"with_memory must return Err once a VMM arena is already selected, not Ok(Self) \
with the injection silently ignored",
);
);
let message = error.to_string();
assert!(
message.contains("already selected a VMM arena"),
Expand Down Expand Up @@ -2232,10 +2231,8 @@ extern "C" __global__ void write_after_delay(unsigned int* out, long long spin)
&self,
bytes: usize,
align: usize,
) -> std::result::Result<
std::ptr::NonNull<u8>,
onnx_runtime_memory_governor::MemoryError,
> {
) -> std::result::Result<std::ptr::NonNull<u8>, onnx_runtime_memory_governor::MemoryError>
{
self.base.allocate(bytes, align)
}

Expand Down Expand Up @@ -2299,9 +2296,7 @@ extern "C" __global__ void write_after_delay(unsigned int* out, long long spin)
unsafe { std::env::remove_var(crate::vmm_allocator::CUDA_VMM_ENV) };

let Ok(provider) = CudaExecutionProvider::new(0) else {
eprintln!(
"SKIPPED (no CUDA runtime): the offset-zero identity proof did NOT run."
);
eprintln!("SKIPPED (no CUDA runtime): the offset-zero identity proof did NOT run.");
panic!("CUDA test path did not run; report as a failed GPU test, not a pass");
};
assert!(
Expand Down Expand Up @@ -2339,10 +2334,8 @@ extern "C" __global__ void write_after_delay(unsigned int* out, long long spin)
&self,
bytes: usize,
align: usize,
) -> std::result::Result<
std::ptr::NonNull<u8>,
onnx_runtime_memory_governor::MemoryError,
> {
) -> std::result::Result<std::ptr::NonNull<u8>, onnx_runtime_memory_governor::MemoryError>
{
self.base.allocate(bytes, align)
}

Expand Down Expand Up @@ -2379,8 +2372,8 @@ extern "C" __global__ void write_after_delay(unsigned int* out, long long spin)
// Non-vacuous: prove the address collision this test exists to guard
// against is real before asserting on the rejection it causes.
let wrapper_addr = &wrapper as *const OffsetZeroWrapper as *const ();
let foreign_addr = &wrapper.foreign as *const crate::vmm_allocator::CudaVmmAllocator
as *const ();
let foreign_addr =
&wrapper.foreign as *const crate::vmm_allocator::CudaVmmAllocator as *const ();
assert!(
std::ptr::eq(wrapper_addr, foreign_addr),
"test setup did not reproduce the offset-zero address collision `#[repr(C)]` is \
Expand Down
4 changes: 3 additions & 1 deletion crates/onnx-runtime-memory-api/src/allocator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -226,7 +226,9 @@ pub trait DeviceAllocator: Send + Sync + Debug {
unsafe fn deallocate_with_unmapped(&self, ptr: NonNull<u8>, bytes: usize, align: usize) -> u64 {
let unmapped = match self.as_shared_mapping() {
// SAFETY: forwarded under this method's identical contract.
Some(shared_mapping) => unsafe { shared_mapping.release_shared_mapping(ptr, bytes, align) },
Some(shared_mapping) => unsafe {
shared_mapping.release_shared_mapping(ptr, bytes, align)
},
None => 0,
};
// SAFETY: forwarded under this method's identical contract.
Expand Down
41 changes: 33 additions & 8 deletions crates/onnx-runtime-memory-api/src/capability.rs
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,10 @@ pub trait VirtualBacking: Send + Sync {
/// allocator whose ranges can share physical granules (so committing them
/// together claims each shared granule only once) should override this
/// to do so atomically rather than one range at a time.
fn commit_allocation_ranges(&self, ranges: &[AllocationCommitRange]) -> Result<(), MemoryError> {
fn commit_allocation_ranges(
&self,
ranges: &[AllocationCommitRange],
) -> Result<(), MemoryError> {
for range in ranges {
self.commit_allocation_range(
range.ptr,
Expand Down Expand Up @@ -384,7 +387,10 @@ pub trait SharedMapping: Send + Sync {
fn device(&self) -> DeviceKey;

/// Create a pinned, read-only shared prefix of `bytes`.
fn create_shared_prefix(&self, bytes: usize) -> Result<Box<dyn SharedDevicePrefix>, MemoryError>;
fn create_shared_prefix(
&self,
bytes: usize,
) -> Result<Box<dyn SharedDevicePrefix>, MemoryError>;

/// Estimate the incremental **owned** physical bytes to admit one more
/// sharer of `prefix`.
Expand Down Expand Up @@ -511,7 +517,11 @@ mod tests {
unimplemented!("test double: never called")
}

fn mapped_bytes_for_allocation(&self, _bytes: usize, _align: usize) -> Result<u64, MemoryError> {
fn mapped_bytes_for_allocation(
&self,
_bytes: usize,
_align: usize,
) -> Result<u64, MemoryError> {
unimplemented!("test double: never called")
}

Expand Down Expand Up @@ -611,7 +621,11 @@ mod tests {
unimplemented!("test double: never called")
}

fn mapped_bytes_for_allocation(&self, _bytes: usize, _align: usize) -> Result<u64, MemoryError> {
fn mapped_bytes_for_allocation(
&self,
_bytes: usize,
_align: usize,
) -> Result<u64, MemoryError> {
unimplemented!("test double: never called")
}

Expand Down Expand Up @@ -796,11 +810,19 @@ mod tests {
DeviceKey::device(0)
}

fn create_shared_prefix(&self, bytes: usize) -> Result<Box<dyn SharedDevicePrefix>, MemoryError> {
Ok(Box::new(FakeSharedPrefix { bytes: bytes as u64 }))
fn create_shared_prefix(
&self,
bytes: usize,
) -> Result<Box<dyn SharedDevicePrefix>, MemoryError> {
Ok(Box::new(FakeSharedPrefix {
bytes: bytes as u64,
}))
}

fn incremental_owned_bytes_for_shared_prefix(&self, prefix: &dyn SharedDevicePrefix) -> u64 {
fn incremental_owned_bytes_for_shared_prefix(
&self,
prefix: &dyn SharedDevicePrefix,
) -> u64 {
prefix.committed_physical_bytes()
}

Expand Down Expand Up @@ -898,7 +920,10 @@ mod tests {

// SAFETY: `ptr` identifies a live allocation this capability mapped.
let first = unsafe { shared_mapping.release_shared_mapping(ptr, BYTES, 8) };
assert_eq!(first, BYTES as u64, "the first release must refund the full mapped amount");
assert_eq!(
first, BYTES as u64,
"the first release must refund the full mapped amount"
);

// SAFETY: calling the release operation again on the same `ptr` is
// exactly the scenario this test proves is safe: the bookkeeping was
Expand Down
6 changes: 5 additions & 1 deletion crates/onnx-runtime-session/src/executor/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -171,8 +171,12 @@ mod phase_profile {

/// Holds the planner on for one test and clears it on drop, so a leaked
/// gate cannot be observed by a later test that does not take the lock.
///
/// The field is never read on purpose: it is the [`globals_lock`] guard,
/// and holding it is the whole point. `dead_code` cannot see that a
/// `MutexGuard`'s value is its `Drop`, so it has to be told.
#[cfg(test)]
pub(super) struct ActivationPlanForTest(std::sync::MutexGuard<'static, ()>);
pub(super) struct ActivationPlanForTest(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>);

#[cfg(test)]
impl ActivationPlanForTest {
Expand Down
Loading