migueldeicaza#123 add list devices. (migueldeicaza#172)

alexpantyukhin · migueldeicaza · commit 7cfa6a157408 · 2018-02-28T19:51:35.000-05:00
* add list devices.

* Adapting ListDevices method for existing method of handling errors.
diff --git a/TensorFlowSharp/Tensorflow.cs b/TensorFlowSharp/Tensorflow.cs
@@ -31,6 +31,7 @@
 using TF_Library = System.IntPtr;
 using TF_BufferPtr = System.IntPtr;
 using TF_Function = System.IntPtr;
+using TF_DeviceList = System.IntPtr;
 
 using size_t = System.UIntPtr;
 using System.Numerics;
@@ -2064,6 +2065,27 @@ public TFOutput this [int idx] {
 		}
 	}
 
+	public enum DeviceType
+	{
+		CPU, GPU, TPU
+	}
+
+	public class DeviceAttributes
+	{
+		public DeviceAttributes (string name, DeviceType deviceType, long memoryLimitBytes)
+		{
+			Name = name;
+			DeviceType = deviceType;
+			MemoryLimitBytes = memoryLimitBytes;
+		}
+		
+		public string Name { get; private set; }
+
+		public DeviceType DeviceType { get; private  set; }
+
+		public long MemoryLimitBytes { get; private  set; }
+	} 
+
 	/// <summary>
 	/// Contains options that are used to control how graph importing works.
 	/// </summary>
@@ -2321,6 +2343,48 @@ public TFSession (TFStatus status = null) : this (new TFGraph (), status)
 		[DllImport (NativeBinding.TensorFlowLibrary)]
 		static extern unsafe TF_Session TF_LoadSessionFromSavedModel (TF_SessionOptions session_options, LLBuffer* run_options, string export_dir, string [] tags, int tags_len, TF_Graph graph, LLBuffer* meta_graph_def, TF_Status status);
 
+		[DllImport (NativeBinding.TensorFlowLibrary)]
+		static extern unsafe TF_DeviceList TF_SessionListDevices (TF_Session session, TF_Status status);
+
+		[DllImport (NativeBinding.TensorFlowLibrary)]
+		static extern unsafe int TF_DeviceListCount (TF_DeviceList list);
+
+		[DllImport (NativeBinding.TensorFlowLibrary)]
+		static extern unsafe string TF_DeviceListName (TF_DeviceList list, int index, TF_Status status);
+
+		[DllImport (NativeBinding.TensorFlowLibrary)]
+		static extern unsafe string TF_DeviceListType (TF_DeviceList list, int index, TF_Status status);
+ 
+		[DllImport (NativeBinding.TensorFlowLibrary)]
+		static extern unsafe long TF_DeviceListMemoryBytes (TF_DeviceList list, int index, TF_Status status);
+
+		[DllImport (NativeBinding.TensorFlowLibrary)]
+		static extern unsafe void TF_DeleteDeviceList (TF_DeviceList list);
+
+		/// <summary>
+		/// Lists available devices in this session.
+		/// </summary>
+		public IEnumerable<DeviceAttributes> ListDevices(TFStatus status = null)
+		{
+			var cstatus = TFStatus.Setup (status);
+			var rawDeviceList = TF_SessionListDevices (this.Handle, cstatus.handle);
+			var size = TF_DeviceListCount (rawDeviceList);
+
+			var list = new List<DeviceAttributes> ();
+			for (var i = 0; i < size; i++) {
+				var name = TF_DeviceListName (rawDeviceList, i, cstatus.handle);
+				var deviceType = (DeviceType) Enum.Parse (typeof(DeviceType), TF_DeviceListType (rawDeviceList, i, cstatus.handle)); 
+				var memory = TF_DeviceListMemoryBytes (rawDeviceList, i, cstatus.handle);
+
+				list.Add (new DeviceAttributes (name, deviceType, memory));
+			}
+
+			// TODO: Fix deleting.
+			// TF_DeleteDeviceList (rawDeviceList);
+
+			return list;
+		}
+
 		/// <summary>
 		/// Creates a session and graph from a saved session model
 		/// </summary>
