From 0e262aab3d48a078f9717f23b107c825e23605a8 Mon Sep 17 00:00:00 2001 From: behzad nouri Date: Thu, 8 Apr 2021 10:40:30 +0000 Subject: [PATCH] simplifies PinnedVec implementation (#16382) --- perf/src/cuda_runtime.rs | 128 +++++++++++---------------------------- 1 file changed, 35 insertions(+), 93 deletions(-) diff --git a/perf/src/cuda_runtime.rs b/perf/src/cuda_runtime.rs index 4d61c09db..45561af50 100644 --- a/perf/src/cuda_runtime.rs +++ b/perf/src/cuda_runtime.rs @@ -11,45 +11,39 @@ use rand::seq::SliceRandom; use rand::Rng; use rayon::prelude::*; use std::ops::{Index, IndexMut}; -use std::slice::SliceIndex; -use std::sync::{Arc, Weak}; +use std::slice::{Iter, IterMut, SliceIndex}; +use std::sync::Weak; use std::os::raw::c_int; const CUDA_SUCCESS: c_int = 0; -pub fn pin(_mem: &mut Vec) { +fn pin(_mem: &mut Vec) { if let Some(api) = perf_libs::api() { - unsafe { - use core::ffi::c_void; - use std::mem::size_of; + use std::ffi::c_void; + use std::mem::size_of; - let err = (api.cuda_host_register)( - _mem.as_mut_ptr() as *mut c_void, - _mem.capacity().saturating_mul(size_of::()), - 0, + let ptr = _mem.as_mut_ptr(); + let size = _mem.capacity().saturating_mul(size_of::()); + let err = unsafe { + (api.cuda_host_register)(ptr as *mut c_void, size, /*flags=*/ 0) + }; + if err != CUDA_SUCCESS { + panic!( + "cudaHostRegister error: {} ptr: {:?} bytes: {}", + err, ptr, size, ); - if err != CUDA_SUCCESS { - panic!( - "cudaHostRegister error: {} ptr: {:?} bytes: {}", - err, - _mem.as_ptr(), - _mem.capacity().saturating_mul(size_of::()), - ); - } } } } -pub fn unpin(_mem: *mut T) { +fn unpin(_mem: *mut T) { if let Some(api) = perf_libs::api() { - unsafe { - use core::ffi::c_void; + use std::ffi::c_void; - let err = (api.cuda_host_unregister)(_mem as *mut c_void); - if err != CUDA_SUCCESS { - panic!("cudaHostUnregister returned: {} ptr: {:?}", err, _mem); - } + let err = unsafe { (api.cuda_host_unregister)(_mem as *mut c_void) }; + if err != CUDA_SUCCESS { + panic!("cudaHostUnregister returned: {} ptr: {:?}", err, _mem); } } } @@ -57,12 +51,12 @@ pub fn unpin(_mem: *mut T) { // A vector wrapper where the underlying memory can be // page-pinned. Controlled by flags in case user only wants // to pin in certain circumstances. -#[derive(Debug)] +#[derive(Debug, Default)] pub struct PinnedVec { x: Vec, pinned: bool, pinnable: bool, - recycler: Option>>>, + recycler: Weak>>, } impl Reset for PinnedVec { @@ -74,21 +68,10 @@ impl Reset for PinnedVec { self.resize(size_hint, T::default()); } fn set_recycler(&mut self, recycler: Weak>) { - self.recycler = Some(recycler); + self.recycler = recycler; } fn unset_recycler(&mut self) { - self.recycler = None; - } -} - -impl Default for PinnedVec { - fn default() -> Self { - Self { - x: Vec::new(), - pinned: false, - pinnable: false, - recycler: None, - } + self.recycler = Weak::default(); } } @@ -99,29 +82,10 @@ impl From> for Vec { pinned_vec.pinned = false; } pinned_vec.pinnable = false; - pinned_vec.recycler = None; + pinned_vec.recycler = Weak::default(); std::mem::take(&mut pinned_vec.x) } } -pub struct PinnedIter<'a, T>(std::slice::Iter<'a, T>); - -pub struct PinnedIterMut<'a, T>(std::slice::IterMut<'a, T>); - -impl<'a, T: Clone + Default + Sized> Iterator for PinnedIter<'a, T> { - type Item = &'a T; - - fn next(&mut self) -> Option { - self.0.next() - } -} - -impl<'a, T: Clone + Default + Sized> Iterator for PinnedIterMut<'a, T> { - type Item = &'a mut T; - - fn next(&mut self) -> Option { - self.0.next() - } -} impl IntoIterator for PinnedVec { type Item = T; @@ -132,21 +96,12 @@ impl IntoIterator for PinnedVec { } } -impl<'a, T: Clone + Default + Sized> IntoIterator for &'a mut PinnedVec { - type Item = &'a T; - type IntoIter = PinnedIter<'a, T>; - - fn into_iter(self) -> Self::IntoIter { - PinnedIter(self.x.iter()) - } -} - impl<'a, T: Clone + Default + Sized> IntoIterator for &'a PinnedVec { type Item = &'a T; - type IntoIter = PinnedIter<'a, T>; + type IntoIter = Iter<'a, T>; fn into_iter(self) -> Self::IntoIter { - PinnedIter(self.x.iter()) + self.x.iter() } } @@ -167,12 +122,12 @@ impl> IndexMut for PinnedVec PinnedVec { - pub fn iter(&self) -> PinnedIter { - PinnedIter(self.x.iter()) + pub fn iter(&self) -> Iter<'_, T> { + self.x.iter() } - pub fn iter_mut(&mut self) -> PinnedIterMut { - PinnedIterMut(self.x.iter_mut()) + pub fn iter_mut(&mut self) -> IterMut<'_, T> { + self.x.iter_mut() } pub fn capacity(&self) -> usize { @@ -237,18 +192,12 @@ impl PinnedVec { x: source, pinned: false, pinnable: false, - recycler: None, + recycler: Weak::default(), } } pub fn with_capacity(capacity: usize) -> Self { - let x = Vec::with_capacity(capacity); - Self { - x, - pinned: false, - pinnable: false, - recycler: None, - } + Self::from_vec(Vec::with_capacity(capacity)) } pub fn is_empty(&self) -> bool { @@ -332,10 +281,6 @@ impl PinnedVec { self.pinned = true; } } - fn recycler_ref(&self) -> Option>> { - let r = self.recycler.as_ref()?; - r.upgrade() - } } impl Clone for PinnedVec { @@ -364,12 +309,9 @@ impl Clone for PinnedVec { impl Drop for PinnedVec { fn drop(&mut self) { - if let Some(strong) = self.recycler_ref() { - let mut vec = PinnedVec::default(); - std::mem::swap(&mut vec, self); - strong.recycle(vec); - } - if self.pinned { + if let Some(recycler) = self.recycler.upgrade() { + recycler.recycle(std::mem::take(self)); + } else if self.pinned { unpin(self.x.as_mut_ptr()); } }