feat: introduce useful util functions for cc and multi-gpu ppcie (#90)

Adds util functions for cc and multi-gpu ppcie

Author: jorgeantonio21

nvml-wrapper/src/device.rs

@@ -785,6 +785,188 @@ impl<'nvml> Device<'nvml> {
         }
     }
 
+    /**
+    Checks simultaneously if confidential compute is enabled, if the device is in a production environment,
+    and if the device is accepting client requests.
+    # Errors
+    * `Uninitialized`, if the library has not been successfully initialized
+    * `NotSupported`, if this query is not supported by the device
+    * `InvalidArg`, if confidential compute state is invalid
+    */
+    pub fn check_confidential_compute_status(&self) -> Result<bool, NvmlError> {
+        let cc_state_sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeState.as_ref())?;
+        let cc_gpus_ready_sym = nvml_sym(
+            self.nvml
+                .lib
+                .nvmlSystemGetConfComputeGpusReadyState
+                .as_ref(),
+        )?;
+
+        unsafe {
+            let mut state: nvmlConfComputeSystemState_t = mem::zeroed();
+            nvml_try(cc_state_sym(&mut state))?;
+
+            let is_cc_enabled = state.ccFeature == NVML_CC_SYSTEM_FEATURE_ENABLED;
+            let is_prod_environment = state.environment == NVML_CC_SYSTEM_ENVIRONMENT_PROD;
+
+            let mut cc_gpus_ready: std::os::raw::c_uint = 0;
+            nvml_try(cc_gpus_ready_sym(&mut cc_gpus_ready))?;
+            let is_accepting_client_requests =
+                cc_gpus_ready == NVML_CC_ACCEPTING_CLIENT_REQUESTS_TRUE;
+
+            Ok(is_cc_enabled && is_prod_environment && is_accepting_client_requests)
+        }
+    }
+
+    /**
+    Gets the confidential compute state for this `Device`.
+    # Errors
+    * `Uninitialized`, if the library has not been successfully initialized
+    * `InvalidArg`, if device is invalid or memory is NULL
+    * `NotSupported`, if this query is not supported by the device
+    */
+    #[doc(alias = "nvmlDeviceGetConfComputeGpusReadyState")]
+    pub fn get_confidential_compute_state(&self) -> Result<bool, NvmlError> {
+        let sym = nvml_sym(
+            self.nvml
+                .lib
+                .nvmlSystemGetConfComputeGpusReadyState
+                .as_ref(),
+        )?;
+
+        unsafe {
+            let mut is_accepting_work: u32 = 0;
+            nvml_try(sym(&mut is_accepting_work))?;
+            Ok(is_accepting_work == NVML_CC_ACCEPTING_CLIENT_REQUESTS_TRUE)
+        }
+    }
+
+    /**
+    Sets the confidential compute state for this `Device`.
+    # Errors
+    * `Uninitialized`, if the library has not been successfully initialized
+    * `InvalidArg`, if device is invalid or memory is NULL
+    * `NotSupported`, if this query is not supported by the device
+    */
+    #[doc(alias = "nvmlDeviceSetConfComputeState")]
+    pub fn set_confidential_compute_state(&self, is_accepting_work: bool) -> Result<(), NvmlError> {
+        let sym = nvml_sym(
+            self.nvml
+                .lib
+                .nvmlSystemSetConfComputeGpusReadyState
+                .as_ref(),
+        )?;
+
+        unsafe {
+            nvml_try(sym(is_accepting_work as u32))?;
+            Ok(())
+        }
+    }
+
+    /**
+    Gets the confidential compute state for this `Device`.
+    # Errors
+
+    * `Uninitialized`, if the library has not been successfully initialized
+    * `InvalidArg`, if device is invalid or counters is NULL
+    * `NotSupported`, if the device does not support this feature
+    * `GpuLost`, if the target GPU has fallen off the bus or is otherwise inaccessible
+    * `ArgumentVersionMismatch`, if the provided version is invalid/unsupported
+    * `Unknown`, on any unexpected error
+    */
+    #[doc(alias = "nvmlDeviceSetConfComputeSettings")]
+    pub fn is_cc_enabled(&self) -> Result<bool, NvmlError> {
+        let sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeSettings.as_ref())?;
+
+        unsafe {
+            let mut settings: nvmlSystemConfComputeSettings_t = mem::zeroed();
+            nvml_try(sym(&mut settings))?;
+            Ok(settings.ccFeature == NVML_CC_SYSTEM_FEATURE_ENABLED)
+        }
+    }
+
+    /**
+    Gets the confidential compute state for this `Device`.
+    # Errors
+
+    * `Uninitialized`, if the library has not been successfully initialized
+    * `InvalidArg`, if device is invalid or counters is NULL
+    * `NotSupported`, if the device does not support this feature
+    * `GpuLost`, if the target GPU has fallen off the bus or is otherwise inaccessible
+    * `ArgumentVersionMismatch`, if the provided version is invalid/unsupported
+    * `Unknown`, on any unexpected error
+    */
+    #[doc(alias = "nvmlSystemGetConfComputeSettings")]
+    pub fn is_multi_gpu_protected_pcie_enabled(&self) -> Result<bool, NvmlError> {
+        let sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeSettings.as_ref())?;
+
+        unsafe {
+            let mut settings: nvmlSystemConfComputeSettings_t = mem::zeroed();
+            nvml_try(sym(&mut settings))?;
+            Ok(settings.multiGpuMode == NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE)
+        }
+    }
+
+    /**
+    Gets the confidential compute state for this `Device`.
+    # Errors
+
+    * `Uninitialized`, if the library has not been successfully initialized
+    * `InvalidArg`, if device is invalid or counters is NULL
+    * `NotSupported`, if the device does not support this feature
+    * `GpuLost`, if the target GPU has fallen off the bus or is otherwise inaccessible
+    * `ArgumentVersionMismatch`, if the provided version is invalid/unsupported
+    * `Unknown`, on any unexpected error
+    */
+    #[doc(alias = "nvmlSystemGetConfComputeSettings")]
+    pub fn is_cc_dev_mode_enabled(&self) -> Result<bool, NvmlError> {
+        let sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeSettings.as_ref())?;
+
+        unsafe {
+            let mut settings: nvmlSystemConfComputeSettings_t = mem::zeroed();
+            nvml_try(sym(&mut settings))?;
+            Ok(settings.devToolsMode == NVML_CC_SYSTEM_DEVTOOLS_MODE_ON)
+        }
+    }
+
+    /**
+    Gets the confidential compute capabilities for this `Device`.
+    # Errors
+    * `Uninitialized`, if the library has not been successfully initialized
+    * `InvalidArg`, if device is invalid or memory is NULL
+    * `NotSupported`, if this query is not supported by the device
+    */
+    pub fn get_confidential_compute_capabilities(
+        &self,
+    ) -> Result<ConfidentialComputeCapabilities, NvmlError> {
+        let sym = nvml_sym(self.nvml.lib.nvmlSystemGetConfComputeCapabilities.as_ref())?;
+
+        unsafe {
+            let mut capabilities: nvmlConfComputeSystemCaps_t = mem::zeroed();
+            nvml_try(sym(&mut capabilities))?;
+
+            let cpu_caps = match capabilities.cpuCaps {
+                NVML_CC_SYSTEM_CPU_CAPS_NONE => ConfidentialComputeCpuCapabilities::None,
+                NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV => ConfidentialComputeCpuCapabilities::AmdSev,
+                NVML_CC_SYSTEM_CPU_CAPS_INTEL_TDX => ConfidentialComputeCpuCapabilities::IntelTdx,
+                _ => return Err(NvmlError::Unknown),
+            };
+
+            let gpus_caps = match capabilities.gpusCaps {
+                NVML_CC_SYSTEM_GPUS_CC_CAPABLE => ConfidentialComputeGpuCapabilities::Capable,
+                NVML_CC_SYSTEM_GPUS_CC_NOT_CAPABLE => {
+                    ConfidentialComputeGpuCapabilities::NotCapable
+                }
+                _ => return Err(NvmlError::Unknown),
+            };
+
+            Ok(ConfidentialComputeCapabilities {
+                cpu_caps,
+                gpus_caps,
+            })
+        }
+    }
+
     /**
     Fetches the confidential compute attestation report for this [`Device`].
 

nvml-wrapper/src/structs/device.rs

@@ -4,6 +4,38 @@ use crate::enum_wrappers::device::OperationMode;
 #[cfg(feature = "serde")]
 use serde_derive::{Deserialize, Serialize};
 
+/// Returned from `Device.get_confidential_compute_capabilities()`
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct ConfidentialComputeCapabilities {
+    /// The CPU capabilities.
+    pub cpu_caps: ConfidentialComputeCpuCapabilities,
+    /// The GPU capabilities.
+    pub gpus_caps: ConfidentialComputeGpuCapabilities,
+}
+
+/// The possible CPU capabilities for confidential compute (either None, AMD SEV or Intel TDX)
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum ConfidentialComputeCpuCapabilities {
+    /// No CPU capabilities.
+    None,
+    /// AMD SEV confidential compute capabilities.
+    AmdSev,
+    /// Intel TDX confidential compute capabilities.
+    IntelTdx,
+}
+
+/// The possible GPU capabilities for confidential compute (either not capable or capable)
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum ConfidentialComputeGpuCapabilities {
+    /// Capable.
+    Capable,
+    /// Not capable.
+    NotCapable,
+}
+
 /// Returned from `Device.confidential_compute_gpu_attestation_report_bytes()`
 #[derive(Debug, Clone, Eq, PartialEq, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]