Remove rust bindings for pinned memory

jeremyfelder · jeremyfelder · commit 2ca04c8e15c0 · 2024-06-19T17:32:14.000+03:00
diff --git a/wrappers/rust/icicle-core/src/msm/mod.rs b/wrappers/rust/icicle-core/src/msm/mod.rs
@@ -388,11 +388,6 @@ macro_rules! impl_msm_tests {
             check_msm::<$curve>()
         }
 
-        #[test]
-        fn test_msm_pinned() {
-            check_msm_pinned_memory::<$curve>()
-        }
-
         #[test]
         fn test_msm_batch() {
             check_msm_batch::<$curve>()
diff --git a/wrappers/rust/icicle-core/src/msm/tests.rs b/wrappers/rust/icicle-core/src/msm/tests.rs
@@ -2,7 +2,7 @@ use crate::curve::{Affine, Curve, Projective};
 use crate::msm::{msm, precompute_points, MSMConfig, MSM};
 use crate::traits::{FieldImpl, GenerateRandom};
 use icicle_cuda_runtime::device::{get_device_count, set_device, warmup};
-use icicle_cuda_runtime::memory::{CudaHostAllocFlags, CudaHostRegisterFlags, DeviceVec, HostOrDeviceSlice, HostSlice};
+use icicle_cuda_runtime::memory::{DeviceVec, HostSlice};
 use icicle_cuda_runtime::stream::CudaStream;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 
@@ -105,88 +105,6 @@ where
         });
 }
 
-pub fn check_msm_pinned_memory<C: Curve + MSM<C> + 'static>()
-where
-    <C::ScalarField as FieldImpl>::Config: GenerateRandom<C::ScalarField>,
-    C::ScalarField: ArkConvertible<ArkEquivalent = <C::ArkSWConfig as ArkCurveConfig>::ScalarField>,
-    C::BaseField: ArkConvertible<ArkEquivalent = <C::ArkSWConfig as ArkCurveConfig>::BaseField>,
-{
-    let largest_size = 1 << 16;
-    // let test_sizes = [1 << 10, largest_size];
-    let test_size = largest_size;
-    let mut msm_results = DeviceVec::<Projective<C>>::cuda_malloc(1).unwrap();
-    let random_points = generate_random_affine_points_with_zeroes(largest_size, 2);
-    let points: &HostSlice<Affine<C>> = HostSlice::from_slice(&random_points);
-
-    // TODO: implement cudaHostRegister via HostSlice.pin
-    // let pin = points.is_pinnable();
-    // if pin {
-    //     points.pin(CudaHostRegisterFlags::DEFAULT).unwrap();
-    //     let flags = points.get_memory_flags().unwrap();
-    //     println!("Flags of registered pin: {:?}", flags);
-    //     unsafe {
-    //         println!("points address Rust after pin: {:?}", points.as_ptr());
-    //     }
-    //     points.unpin();
-    //     unsafe {
-    //         println!("points address Rust after unpin: {:?}", points.as_ptr());
-    //     }
-    // }
-
-    let scalars = <<C as Curve>::ScalarField as FieldImpl>::Config::generate_random(largest_size);
-
-    let mut scalars_d = DeviceVec::<<C as Curve>::ScalarField>::cuda_malloc(test_size).unwrap();
-    let stream = CudaStream::create().unwrap();
-    scalars_d
-        .copy_from_host_async(HostSlice::from_slice(&scalars[..test_size]), &stream)
-        .unwrap();
-
-    let mut cfg = MSMConfig::default();
-    cfg.ctx
-        .stream = &stream;
-    cfg.is_async = true;
-    msm(&scalars_d[..], points, &cfg, &mut msm_results[..]).unwrap();
-
-    let mut msm_host_result = vec![Projective::<C>::zero(); 1];
-    msm_results
-        .copy_to_host(HostSlice::from_mut_slice(&mut msm_host_result[..]))
-        .unwrap();
-    stream
-        .synchronize()
-        .unwrap();
-
-    let msm_res_affine: ark_ec::short_weierstrass::Affine<<C as Curve>::ArkSWConfig> = msm_host_result[0]
-        .to_ark()
-        .into();
-    assert!(msm_res_affine.is_on_curve());
-
-    points.allocate_pinned(points.len(), CudaHostAllocFlags::DEFAULT).unwrap();
-    // let allocated_pinned_points = HostSlice::allocate_pinned(points.len(), CudaHostAllocFlags::DEFAULT).unwrap();
-    // allocated_pinned_points
-    //     .as_mut_slice()
-    //     .clone_from_slice(points.as_slice());
-
-    msm(&scalars_d[..], points, &cfg, &mut msm_results[..]).unwrap();
-
-    let mut msm_host_result = vec![Projective::<C>::zero(); 1];
-    msm_results
-        .copy_to_host(HostSlice::from_mut_slice(&mut msm_host_result[..]))
-        .unwrap();
-    stream
-        .synchronize()
-        .unwrap();
-
-    let msm_res_affine: ark_ec::short_weierstrass::Affine<<C as Curve>::ArkSWConfig> = msm_host_result[0]
-        .to_ark()
-        .into();
-    assert!(msm_res_affine.is_on_curve());
-    points.free_pinned();
-
-    stream
-        .destroy()
-        .unwrap();
-}
-
 pub fn check_msm_batch<C: Curve + MSM<C>>()
 where
     <C::ScalarField as FieldImpl>::Config: GenerateRandom<C::ScalarField>,
diff --git a/wrappers/rust/icicle-cuda-runtime/src/device.rs b/wrappers/rust/icicle-cuda-runtime/src/device.rs
@@ -1,7 +1,7 @@
 use crate::{
     bindings::{
-        cudaDeviceAttr, cudaDeviceGetAttribute, cudaFreeAsync, cudaGetDevice, cudaGetDeviceCount, cudaMallocAsync,
-        cudaMemGetInfo, cudaPointerAttributes, cudaPointerGetAttributes, cudaSetDevice,
+        cudaFreeAsync, cudaGetDevice, cudaGetDeviceCount, cudaMallocAsync, cudaMemGetInfo, cudaPointerAttributes,
+        cudaPointerGetAttributes, cudaSetDevice,
     },
     error::{CudaResult, CudaResultWrap},
     stream::CudaStream,
@@ -32,11 +32,6 @@ pub fn get_device_from_pointer(ptr: *const ::std::os::raw::c_void) -> CudaResult
     }
 }
 
-pub fn get_device_attribute(attr: cudaDeviceAttr, device: usize) -> CudaResult<usize> {
-    let mut res: i32 = 0;
-    unsafe { cudaDeviceGetAttribute(&mut res, attr, device as i32) }.wrap_value(res as usize)
-}
-
 pub fn check_device(device_id: usize) {
     match device_id == get_device().unwrap() {
         true => (),
diff --git a/wrappers/rust/icicle-cuda-runtime/src/memory.rs b/wrappers/rust/icicle-cuda-runtime/src/memory.rs
@@ -1,34 +1,17 @@
 use crate::bindings::{
-    cudaDeviceAttr, cudaFree, cudaFreeHost, cudaHostAlloc, cudaHostAllocDefault, cudaHostAllocPortable,
-    cudaHostGetFlags, cudaHostRegister, cudaHostRegisterDefault, cudaHostRegisterPortable, cudaHostUnregister,
-    cudaMalloc, cudaMallocAsync, cudaMemPool_t, cudaMemcpy, cudaMemcpyAsync, cudaMemcpyKind,
+    cudaFree, cudaMalloc, cudaMallocAsync, cudaMemPool_t, cudaMemcpy, cudaMemcpyAsync, cudaMemcpyKind,
 };
-use crate::device::{check_device, get_device_attribute, get_device_from_pointer};
+use crate::device::{check_device, get_device_from_pointer};
 use crate::error::{CudaError, CudaResult, CudaResultWrap};
 use crate::stream::CudaStream;
-use bitflags::bitflags;
 use std::mem::{size_of, ManuallyDrop, MaybeUninit};
 use std::ops::{
     Deref, DerefMut, Index, IndexMut, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive,
 };
-use std::os::raw::{c_uint, c_void};
+use std::os::raw::c_void;
 use std::slice::from_raw_parts_mut;
 use std::slice::SliceIndex;
 
-bitflags! {
-    pub struct CudaHostAllocFlags: u32 {
-        const DEFAULT = cudaHostAllocDefault;
-        const PORTABLE = cudaHostAllocPortable;
-    }
-}
-
-bitflags! {
-    pub struct CudaHostRegisterFlags: u32 {
-        const DEFAULT = cudaHostRegisterDefault;
-        const PORTABLE = cudaHostRegisterPortable;
-    }
-}
-
 #[derive(Debug)]
 pub struct HostSlice<T>([T]);
 pub struct DeviceVec<T>(ManuallyDrop<Box<[T]>>);
@@ -134,78 +117,6 @@ impl<T> HostSlice<T> {
         self.0
             .iter_mut()
     }
-
-    // TODO: @jeremy Fix the issue where ptr pinned by cudaHostRegister cannot be used in primitives
-    // pub fn is_pinnable(&self) -> bool {
-    //     let pinnable = get_device_attribute(cudaDeviceAttr::cudaDevAttrHostRegisterSupported, 0).unwrap();
-    //     let lockable =
-    //         get_device_attribute(cudaDeviceAttr::cudaDevAttrPageableMemoryAccessUsesHostPageTables, 0).unwrap();
-
-    //     pinnable == 1 && lockable == 0
-    // }
-
-    // pub fn pin(&self, flags: CudaHostRegisterFlags) -> CudaResult<()> {
-    //     if self.is_pinnable() {
-    //         unsafe {
-    //             let ptr = self.as_ptr() as *mut c_void;
-    //             let flags_to_set = flags.bits();
-    //             cudaHostRegister(ptr, self.len(), flags_to_set as c_uint).wrap()
-    //         }
-    //     } else {
-    //         Ok(())
-    //     }
-    // }
-
-    // pub fn unpin(&self) -> CudaResult<()> {
-    //     unsafe {
-    //         let mut flags = 0;
-    //         let ptr = self.as_ptr() as *mut c_void;
-    //         cudaHostGetFlags(&mut flags, ptr).wrap()?;
-    //         cudaHostUnregister(ptr).wrap()
-    //     }
-    // }
-
-    pub fn allocate_pinned(&self, count: usize, flags: CudaHostAllocFlags) -> CudaResult<()> {
-        let size = count
-            .checked_mul(size_of::<T>())
-            .unwrap_or(0);
-        if size == 0 {
-            return Err(CudaError::cudaErrorMemoryAllocation); //TODO: only CUDA backend should return CudaError
-        }
-
-        // let mut pinned_host_ptr = MaybeUninit::<*mut c_void>::uninit();
-
-        // unsafe {
-        //     cudaHostAlloc(pinned_host_ptr.as_mut_ptr(), size, flags.bits).wrap()?;
-        //     let pinned_host_slice = from_raw_parts_mut(pinned_host_ptr.assume_init() as *mut T, count);
-        //     Ok(Self::from_mut_slice(pinned_host_slice))
-        // }
-
-        unsafe {
-            let p_host = self.as_ptr() as *mut *mut c_void;
-            cudaHostAlloc(p_host, size, flags.bits()).wrap()?;
-        }
-        
-        Ok(())
-    }
-
-    pub fn free_pinned(&self) -> CudaResult<()> {
-        unsafe {
-            let mut flags: u32 = 0;
-            let ptr = self.as_ptr() as *mut c_void;
-            cudaHostGetFlags(&mut flags, ptr).wrap()?;
-            cudaFreeHost(ptr).wrap()
-        }
-    }
-
-    pub fn get_memory_flags(&self) -> CudaResult<u32> {
-        unsafe {
-            let mut flags: u32 = 1234;
-            let ptr = self.as_ptr() as *mut c_void;
-            cudaHostGetFlags(&mut flags, ptr).wrap()?;
-            Ok(flags)
-        }
-    }
 }
 
 impl<T> DeviceSlice<T> {
@@ -513,47 +424,3 @@ impl<T> Drop for DeviceVec<T> {
 
 #[allow(non_camel_case_types)]
 pub type CudaMemPool = cudaMemPool_t;
-
-pub(crate) mod tests {
-    use crate::error::CudaError;
-    use crate::memory::{CudaHostAllocFlags, HostOrDeviceSlice, HostSlice};
-
-    // TODO: @jeremy Fix the issue where ptr pinned by cudaHostRegister cannot be used in primitives
-    // #[test]
-    // fn test_pin_memory() {
-    //     let data = vec![1, 2, 3, 4, 5, 7, 8, 9];
-    //     let data_host_slice = HostSlice::from_slice(&data);
-
-    //     data_host_slice
-    //         .pin(CudaHostRegisterFlags::DEFAULT)
-    //         .expect("Registering host mem failed");
-    //     let err = data_host_slice
-    //         .pin(CudaHostRegisterFlags::DEFAULT)
-    //         .expect_err("Registering already registered memory succeeded");
-    //     assert_eq!(err, CudaError::cudaErrorHostMemoryAlreadyRegistered);
-
-    //     data_host_slice
-    //         .unpin()
-    //         .expect("Unregistering pinned memory failed");
-    //     let err = data_host_slice
-    //         .unpin()
-    //         .expect_err("Unregistering non-registered pinned memory succeeded");
-    //     assert_eq!(err, CudaError::cudaErrorInvalidValue);
-    // }
-
-    // #[test]
-    // fn test_allocated_pinned_memory() {
-    //     let data = vec![1, 2, 3, 4, 5, 7, 8, 9];
-    //     let data_host_slice = HostSlice::from_slice(&data);
-    //     let newly_allocated_pinned_host_slice: &HostSlice<i32> =
-    //         HostSlice::allocate_pinned(data_host_slice.len(), CudaHostAllocFlags::DEFAULT)
-    //             .expect("Allocating new pinned memory failed");
-    //     newly_allocated_pinned_host_slice
-    //         .free_pinned()
-    //         .expect("Freeing pinned memory failed");
-    //     let err = newly_allocated_pinned_host_slice
-    //         .free_pinned()
-    //         .expect_err("Freeing non-pinned memory succeeded");
-    //     assert_eq!(err, CudaError::cudaErrorInvalidValue);
-    // }
-}
