Feat: add back set_current, don't return a duration in Linker::complete

Author: RDambrosio016

crates/cust/CHANGELOG.md

@@ -27,6 +27,7 @@ any breaking changes, the API is the same.
 - Added `zeroed_async` to `DeviceBox`.
 - Added `drop_async` to `DeviceBox`.
 - Added `new_async` to `DeviceBox`.
+- `Linker::complete` now only returns the built cubin, and not the cubin and a duration.
 
 ## 0.2.2 - 12/5/21
 

crates/cust/src/context/mod.rs

@@ -576,6 +576,30 @@ impl CurrentContext {
         unsafe { cuda::cuCtxSetSharedMemConfig(transmute(cfg)).to_result() }
     }
 
+    /// Set the given context as the current context for this thread.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use cust::device::Device;
+    /// # use cust::context::{ Context, ContextFlags, CurrentContext };
+    /// # use std::error::Error;
+    /// #
+    /// # fn main () -> Result<(), Box<dyn Error>> {
+    /// # cust::init(cust::CudaFlags::empty())?;
+    /// # let device = Device::get_device(0)?;
+    /// let context = Context::new(device)?;
+    /// CurrentContext::set_current(&context)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn set_current<C: ContextHandle>(c: &C) -> CudaResult<()> {
+        unsafe {
+            cuda::cuCtxSetCurrent(c.get_inner()).to_result()?;
+            Ok(())
+        }
+    }
+
     /// Block to wait for a context's tasks to complete.
     pub fn synchronize() -> CudaResult<()> {
         unsafe {

crates/cust/src/link.rs

@@ -1,6 +1,7 @@
 //! Functions for linking together multiple PTX files into a module.
 
-use std::{mem::MaybeUninit, time::Duration};
+use std::mem::MaybeUninit;
+use std::ptr::null_mut;
 
 use crate::sys as cuda;
 
@@ -12,7 +13,6 @@ static UNNAMED: &str = "\0";
 #[derive(Debug)]
 pub struct Linker {
     raw: cuda::CUlinkState,
-    duration_ptr: *mut *mut f32,
 }
 
 impl Linker {
@@ -22,19 +22,11 @@ impl Linker {
         // Therefore we use box to alloc the memory for us, then into_raw it so we now have ownership
         // of the memory (and dont have any aliasing requirements attached either).
 
-        // technically it should be fine to alloc as Box<*mut f32> then dealloc with Box<Box<f32>> but
-        // in the future rust may make this guarantee untrue so just alloc as Box<Box<f32>> then cast.
-        let ptr = Box::into_raw(Box::new(Box::new(0.0f32))) as *mut *mut f32;
-
-        // cuda shouldnt be modifying this but trust the bindings in that it wants a *mut ptr.
-        let options = &mut [cuda::CUjit_option_enum::CU_JIT_WALL_TIME];
         unsafe {
             let mut raw = MaybeUninit::uninit();
-            cuda::cuLinkCreate_v2(1, options.as_mut_ptr(), ptr.cast(), raw.as_mut_ptr())
-                .to_result()?;
+            cuda::cuLinkCreate_v2(0, null_mut(), null_mut(), raw.as_mut_ptr()).to_result()?;
             Ok(Self {
                 raw: raw.assume_init(),
-                duration_ptr: ptr,
             })
         }
     }
@@ -121,7 +113,7 @@ impl Linker {
 
     /// Runs the linker to generate the final cubin bytes. Also returns a duration
     /// for how long it took to run the linker.
-    pub fn complete(self) -> CudaResult<(Vec<u8>, Duration)> {
+    pub fn complete(self) -> CudaResult<Vec<u8>> {
         let mut cubin = MaybeUninit::uninit();
         let mut size = MaybeUninit::uninit();
 
@@ -134,16 +126,7 @@ impl Linker {
             let mut vec = Vec::with_capacity(size);
             vec.extend_from_slice(slice);
 
-            // now that duration has been written to, retrieve it and deallocate it.
-            let duration = **self.duration_ptr;
-            // recreate a box from our pointer, which will take ownership
-            // of it and then drop it immediately. This is sound because
-            // complete consumes self.
-            Box::from_raw(self.duration_ptr as *mut Box<f32>);
-
-            // convert to nanos so we dont lose the decimal millisecs.
-            let duration = Duration::from_nanos((duration * 1e6) as u64);
-            Ok((vec, duration))
+            Ok(vec)
         }
     }
 }

crates/cust/src/module.rs

@@ -175,70 +175,72 @@ impl Module {
         }
     }
 
-    /// Creates a new module by loading a fatbin (fat binary) file.
-    ///
-    /// Fatbinary files are files that contain multiple ptx or cubin files. The driver will choose already-built
-    /// cubin if it is present, and otherwise JIT compile any PTX in the file to cubin.
-    ///
-    /// # Example
-    ///
-    /// ```
-    /// # use cust::*;
-    /// # use std::error::Error;
-    /// # fn main() -> Result<(), Box<dyn Error>> {
-    /// # let _ctx = quick_init()?;
-    /// use cust::module::Module;
-    /// let fatbin_bytes = std::fs::read("./resources/add.cubin")?;
-    /// assert!(fatbin_bytes.contains(&0));
-    /// let module = Module::from_cubin(&fatbin_bytes, &[])?;
-    /// # Ok(())
-    /// # }
-    /// ```
-    pub fn from_fatbin<T: AsRef<[u8]>>(
-        bytes: T,
-        options: &[ModuleJitOption],
-    ) -> CudaResult<Module> {
-        let mut bytes = bytes.as_ref().to_vec();
-        bytes.push(0);
-        // fatbins are just ELF files like cubins, and cuModuleLoadDataEx accepts ptx, cubin, and fatbin.
-        // We just make the distinction in case we want to do anything extra in the future. As well
-        // as keep things explicit to anyone reading the code.
-        Self::from_cubin(bytes, options)
-    }
+    // TODO(RDambrosio016): figure out why the heck cuda rejects cubins literally made by nvcc and loaded by fs::read
 
-    pub unsafe fn from_fatbin_unchecked<T: AsRef<[u8]>>(
-        bytes: T,
-        options: &[ModuleJitOption],
-    ) -> CudaResult<Module> {
-        Self::from_cubin_unchecked(bytes, options)
-    }
+    // /// Creates a new module by loading a fatbin (fat binary) file.
+    // ///
+    // /// Fatbinary files are files that contain multiple ptx or cubin files. The driver will choose already-built
+    // /// cubin if it is present, and otherwise JIT compile any PTX in the file to cubin.
+    // ///
+    // /// # Example
+    // ///
+    // /// ```
+    // /// # use cust::*;
+    // /// # use std::error::Error;
+    // /// # fn main() -> Result<(), Box<dyn Error>> {
+    // /// # let _ctx = quick_init()?;
+    // /// use cust::module::Module;
+    // /// let fatbin_bytes = std::fs::read("./resources/add.cubin")?;
+    // /// assert!(fatbin_bytes.contains(&0));
+    // /// let module = Module::from_cubin(&fatbin_bytes, &[])?;
+    // /// # Ok(())
+    // /// # }
+    // /// ```
+    // pub fn from_fatbin<T: AsRef<[u8]>>(
+    //     bytes: T,
+    //     options: &[ModuleJitOption],
+    // ) -> CudaResult<Module> {
+    //     let mut bytes = bytes.as_ref().to_vec();
+    //     bytes.push(0);
+    //     // fatbins are just ELF files like cubins, and cuModuleLoadDataEx accepts ptx, cubin, and fatbin.
+    //     // We just make the distinction in case we want to do anything extra in the future. As well
+    //     // as keep things explicit to anyone reading the code.
+    //     Self::from_cubin(bytes, options)
+    // }
 
-    pub fn from_cubin<T: AsRef<[u8]>>(bytes: T, options: &[ModuleJitOption]) -> CudaResult<Module> {
-        let bytes = bytes.as_ref();
-        goblin::elf::Elf::parse(bytes).expect("Cubin/Fatbin was not valid ELF!");
-        // SAFETY: we verified the bytes were valid ELF
-        unsafe { Self::from_cubin_unchecked(bytes, options) }
-    }
+    // pub unsafe fn from_fatbin_unchecked<T: AsRef<[u8]>>(
+    //     bytes: T,
+    //     options: &[ModuleJitOption],
+    // ) -> CudaResult<Module> {
+    //     Self::from_cubin_unchecked(bytes, options)
+    // }
 
-    pub unsafe fn from_cubin_unchecked<T: AsRef<[u8]>>(
-        bytes: T,
-        options: &[ModuleJitOption],
-    ) -> CudaResult<Module> {
-        let bytes = bytes.as_ref();
-        let mut module = Module {
-            inner: ptr::null_mut(),
-        };
-        let (mut options, mut option_values) = ModuleJitOption::into_raw(options);
-        cuda::cuModuleLoadDataEx(
-            &mut module.inner as *mut cuda::CUmodule,
-            bytes.as_ptr() as *const c_void,
-            options.len() as c_uint,
-            options.as_mut_ptr(),
-            option_values.as_mut_ptr(),
-        )
-        .to_result()?;
-        Ok(module)
-    }
+    // pub fn from_cubin<T: AsRef<[u8]>>(bytes: T, options: &[ModuleJitOption]) -> CudaResult<Module> {
+    //     let bytes = bytes.as_ref();
+    //     goblin::elf::Elf::parse(bytes).expect("Cubin/Fatbin was not valid ELF!");
+    //     // SAFETY: we verified the bytes were valid ELF
+    //     unsafe { Self::from_cubin_unchecked(bytes, options) }
+    // }
+
+    // pub unsafe fn from_cubin_unchecked<T: AsRef<[u8]>>(
+    //     bytes: T,
+    //     options: &[ModuleJitOption],
+    // ) -> CudaResult<Module> {
+    //     let bytes = bytes.as_ref();
+    //     let mut module = Module {
+    //         inner: ptr::null_mut(),
+    //     };
+    //     let (mut options, mut option_values) = ModuleJitOption::into_raw(options);
+    //     cuda::cuModuleLoadDataEx(
+    //         &mut module.inner as *mut cuda::CUmodule,
+    //         bytes.as_ptr() as *const c_void,
+    //         options.len() as c_uint,
+    //         options.as_mut_ptr(),
+    //         option_values.as_mut_ptr(),
+    //     )
+    //     .to_result()?;
+    //     Ok(module)
+    // }
 
     pub fn from_ptx_cstr(cstr: &CStr, options: &[ModuleJitOption]) -> CudaResult<Module> {
         unsafe {