From 8e9ab44463288039058195b082ea68709d28d133 Mon Sep 17 00:00:00 2001
From: Anders Smedegaard Pedersen <smedegaard@users.noreply.github.com>
Date: Sat, 2 Nov 2024 19:00:20 +0100
Subject: [PATCH] 7 hipdevicetotalmem (#8)

* 'add device_total_mem()'
---
 src/runtime/safe.rs | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/runtime/safe.rs b/src/runtime/safe.rs
index e042a5d..01c874c 100644
--- a/src/runtime/safe.rs
+++ b/src/runtime/safe.rs
@@ -95,7 +95,7 @@ pub fn set_device(device: Device) -> Result<Device> {
 /// # Returns
 /// * `Result<Version>` - On success, returns a `Version` struct containing the major and minor version
 ///   numbers of the device's compute capability. On failure, returns an error indicating what went wrong.
-pub fn get_device_compute_capability(device: Device) -> Result<Version> {
+pub fn device_compute_capability(device: Device) -> Result<Version> {
     unsafe {
         let mut major: i32 = -1;
         let mut minor: i32 = -1;
@@ -105,15 +105,35 @@ pub fn get_device_compute_capability(device: Device) -> Result<Version> {
     }
 }
 
+pub fn device_total_mem(device: Device) -> Result<usize> {
+    unsafe {
+        let mut size: usize = 0;
+        let code = sys::hipDeviceTotalMem(&mut size, device.id);
+        (size, code).to_result()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_device_total_mem() {
+        let device = Device::new(0);
+        let result = device_total_mem(device);
+        assert!(result.is_ok());
+        let size = result.unwrap();
+        assert!(size > 0);
+        println!("Total memory in bytes: {}", size);
+    }
+
     #[test]
     fn test_get_device_compute_capability() {
         let device = Device::new(0);
-        let result = get_device_compute_capability(device);
+        let result = device_compute_capability(device);
+        assert!(result.is_ok());
         let version = result.unwrap();
+        assert!(version.major > 0);
         println!("Compute Capability: {}.{}", version.major, version.minor);
     }