adding the MPS version of running_compute_processes for >= Volta.

Author: devnexen

nvml-wrapper/src/device.rs

@@ -624,6 +624,48 @@ impl<'nvml> Device<'nvml> {
         }
     }
 
+    /**
+    Gets information about processes with a compute context running on this `Device`.
+    Note that processes list can differ between the accounting call and the list gathering
+
+    # Errors
+
+    * `Uninitialized`, if the library has not been successfully initialized
+    * `InvalidArg`, if this `Device` is invalid
+    * `GpuLost`, if this `Device` has fallen off the bus or is otherwise inaccessible
+    * `Unknown`, on any unexpected error
+
+    # Device Support
+
+    Supports Volta or newer fully supported devices.
+    */
+    #[doc(alias = "nvmlDeviceGetiMPSComputeRunningProcesses_v3")]
+    pub fn mps_running_compute_processes(&self) -> Result<Vec<ProcessInfo>, NvmlError> {
+        let sym = nvml_sym(
+            self.nvml
+                .lib
+                .nvmlDeviceGetMPSComputeRunningProcesses_v3
+                .as_ref(),
+        )?;
+
+        unsafe {
+            let mut hints: c_uint = 0;
+
+            nvml_try(sym(self.device, &mut hints, ptr::null_mut()))?;
+
+            let mut hints_shr = match hints.checked_shr(1) {
+                Some(h) => h,
+                None => hints,
+            };
+            let mut processes: Vec<nvmlProcessInfo_t> = Vec::with_capacity(hints_shr as usize);
+
+            nvml_try(sym(self.device, &mut hints_shr, processes.as_mut_ptr()))?;
+
+            processes.truncate(hints_shr as usize);
+            Ok(processes.into_iter().map(ProcessInfo::from).collect())
+        }
+    }
+
     /**
     Gets the number of processes with a compute context running on this `Device`.
 
@@ -6633,6 +6675,12 @@ mod test {
         test_with_device(3, &nvml, |device| device.running_compute_processes_v2())
     }
 
+    #[test]
+    fn mps_running_compute_processes() {
+        let nvml = nvml();
+        test_with_device(3, &nvml, |device| device.mps_running_compute_processes())
+    }
+
     #[cfg(target_os = "linux")]
     #[test]
     fn cpu_affinity() {