simplifies PinnedVec implementation (#16382)

This commit is contained in:
behzad nouri 2021-04-08 10:40:30 +00:00 committed by GitHub
parent 3645092a52
commit 0e262aab3d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 35 additions and 93 deletions

View File

@ -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<T>(_mem: &mut Vec<T>) {
fn pin<T>(_mem: &mut Vec<T>) {
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::<T>()),
0,
let ptr = _mem.as_mut_ptr();
let size = _mem.capacity().saturating_mul(size_of::<T>());
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::<T>()),
);
}
}
}
}
pub fn unpin<T>(_mem: *mut T) {
fn unpin<T>(_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<T>(_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<T: Default + Clone + Sized> {
x: Vec<T>,
pinned: bool,
pinnable: bool,
recycler: Option<Weak<RecyclerX<PinnedVec<T>>>>,
recycler: Weak<RecyclerX<PinnedVec<T>>>,
}
impl<T: Default + Clone + Sized> Reset for PinnedVec<T> {
@ -74,21 +68,10 @@ impl<T: Default + Clone + Sized> Reset for PinnedVec<T> {
self.resize(size_hint, T::default());
}
fn set_recycler(&mut self, recycler: Weak<RecyclerX<Self>>) {
self.recycler = Some(recycler);
self.recycler = recycler;
}
fn unset_recycler(&mut self) {
self.recycler = None;
}
}
impl<T: Clone + Default + Sized> Default for PinnedVec<T> {
fn default() -> Self {
Self {
x: Vec::new(),
pinned: false,
pinnable: false,
recycler: None,
}
self.recycler = Weak::default();
}
}
@ -99,29 +82,10 @@ impl<T: Clone + Default + Sized> From<PinnedVec<T>> for Vec<T> {
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::Item> {
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::Item> {
self.0.next()
}
}
impl<T: Clone + Default + Sized> IntoIterator for PinnedVec<T> {
type Item = T;
@ -132,21 +96,12 @@ impl<T: Clone + Default + Sized> IntoIterator for PinnedVec<T> {
}
}
impl<'a, T: Clone + Default + Sized> IntoIterator for &'a mut PinnedVec<T> {
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<T> {
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<T: Clone + Default + Sized, I: SliceIndex<[T]>> IndexMut<I> for PinnedVec<T
}
impl<T: Clone + Default + Sized> PinnedVec<T> {
pub fn iter(&self) -> PinnedIter<T> {
PinnedIter(self.x.iter())
pub fn iter(&self) -> Iter<'_, T> {
self.x.iter()
}
pub fn iter_mut(&mut self) -> PinnedIterMut<T> {
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<T: Clone + Default + Sized> PinnedVec<T> {
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<T: Clone + Default + Sized> PinnedVec<T> {
self.pinned = true;
}
}
fn recycler_ref(&self) -> Option<Arc<RecyclerX<Self>>> {
let r = self.recycler.as_ref()?;
r.upgrade()
}
}
impl<T: Clone + Default + Sized> Clone for PinnedVec<T> {
@ -364,12 +309,9 @@ impl<T: Clone + Default + Sized> Clone for PinnedVec<T> {
impl<T: Sized + Default + Clone> Drop for PinnedVec<T> {
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());
}
}