Skip to content

Commit

Permalink
simplifies PinnedVec implementation (solana-labs#16382)
Browse files Browse the repository at this point in the history
  • Loading branch information
@behzadnouri
behzadnouri authored Apr 8, 2021
1 parent 3645092 commit 0e262aa
Showing 1 changed file with 36 additions and 94 deletions.
130 changes: 36 additions & 94 deletions perf/src/cuda_runtime.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,58 +11,52 @@ 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;

let err = (api.cuda_host_register)(
_mem.as_mut_ptr() as *mut c_void,
_mem.capacity().saturating_mul(size_of::<T>()),
0,
use std::ffi::c_void;
use std::mem::size_of;

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);
}
}
}

// 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> {
Expand All @@ -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();
}
}

Expand All @@ -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;
Expand All @@ -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()
}
}

Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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> {
Expand Down Expand Up @@ -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());
}
}
Expand Down

0 comments on commit 0e262aa

Please sign in to comment.