Adding more torch.cuda functionality.

Fixing the squeeze() issue.

Author: NiklasGustafsson

src/Native/LibTorchSharp/THSTensor.cpp

@@ -1230,6 +1230,11 @@ void THSTensor_dsplit_with_sizes(
     )
 }
 
+Tensor THSTensor_squeeze_no_dim(Tensor tensor)
+{
+    CATCH_TENSOR(tensor->squeeze());
+}
+
 Tensor THSTensor_squeeze(Tensor tensor, int64_t dim)
 {
     CATCH_TENSOR(tensor->squeeze(dim));

src/Native/LibTorchSharp/THSTensor.h

@@ -1045,6 +1045,7 @@ EXPORT_API(void) THSTensor_split_with_size(const Tensor tensor, Tensor* (*alloca
 EXPORT_API(void) THSTensor_split_with_sizes(const Tensor tensor, Tensor* (*allocator)(size_t length), const int64_t* sizes, const int length, const int64_t dim);
 
 EXPORT_API(Tensor) THSTensor_squeeze(Tensor tensor, int64_t dim);
+EXPORT_API(Tensor) THSTensor_squeeze_no_dim(Tensor tensor);
 
 EXPORT_API(Tensor) THSTensor_stack(const Tensor* tensor, const int length, const int64_t dim);
 

src/Native/LibTorchSharp/THSTorch.cpp

@@ -15,6 +15,17 @@ Generator THSGenerator_manual_seed(const int64_t seed)
     return THSGenerator_default_generator();
 }
 
+void THSCuda_manual_seed(const int64_t seed)
+{
+    CATCH(torch::cuda::manual_seed(seed);)
+}
+
+void THSCuda_manual_seed_all(const int64_t seed)
+{
+    CATCH(torch::cuda::manual_seed_all(seed);)
+}
+
+
 void THSGenerator_gen_manual_seed(const Generator generator, const int64_t seed)
 {
     generator->set_current_seed(seed);
@@ -69,6 +80,12 @@ int THSTorchCuda_device_count()
     return (int)torch::cuda::device_count();
 }
 
+EXPORT_API(void) THSTorchCuda_synchronize(const int64_t device_index)
+{
+    CATCH(torch::cuda::synchronize(device_index);)
+}
+
+
 const char * THSTorch_get_and_reset_last_err()
 {
     char *tmp = torch_last_err;

src/Native/LibTorchSharp/THSTorch.h

@@ -9,6 +9,8 @@
 
 // Sets manually the seed.
 EXPORT_API(void)      THSTorch_manual_seed(const int64_t seed);
+EXPORT_API(void) THSCuda_manual_seed(const int64_t seed);
+EXPORT_API(void) THSCuda_manual_seed_all(const int64_t seed);
 
 EXPORT_API(Generator) THSGenerator_manual_seed(const int64_t seed);
 EXPORT_API(void) THSGenerator_gen_manual_seed(const Generator gen, const int64_t seed);
@@ -24,6 +26,7 @@ EXPORT_API(void)      THSGenerator_dispose(const Generator generator);
 EXPORT_API(int) THSTorchCuda_is_available();
 EXPORT_API(int) THSTorchCuda_cudnn_is_available();
 EXPORT_API(int) THSTorchCuda_device_count();
+EXPORT_API(void) THSTorchCuda_synchronize(const int64_t device);
 
 // Returns the latest error. This is thread-local.
 EXPORT_API(const char *) THSTorch_get_and_reset_last_err();

src/TorchSharp/Tensor/Tensor.cs

@@ -1243,14 +1243,17 @@ public Tensor reshape(params long[] shape)
             [DllImport("LibTorchSharp")]
             static extern IntPtr THSTensor_squeeze(IntPtr tensor, long dimension);
 
+            [DllImport("LibTorchSharp")]
+            static extern IntPtr THSTensor_squeeze_no_dim(IntPtr tensor);
+
             /// <summary>
             /// Returns a tensor with all the dimensions of input of size 1 removed. When dim is given, a squeeze operation is done only in the given dimension.
             /// </summary>
             /// <param name="dim">If given, the input will be squeezed only in this dimension</param>
             /// <returns></returns>
-            public Tensor squeeze(long dim)
+            public Tensor squeeze(long? dim = null)
             {
-                var res = THSTensor_squeeze(handle, dim);
+                var res = dim.HasValue ? THSTensor_squeeze(handle, dim.Value) : THSTensor_squeeze_no_dim(handle);
                 if (res == IntPtr.Zero)
                     torch.CheckForErrors();
                 return new Tensor(res);
@@ -4239,9 +4242,9 @@ public Tensor slice(long dimension, long start, long finish, long step)
             ///  Returns a new tensor with a dimension of size one inserted at the specified position.
             ///  The returned tensor shares the same underlying data with this tensor.
             /// </summary>
-            public Tensor unsqueeze(long dimension)
+            public Tensor unsqueeze(long dim)
             {
-                var res = THSTensor_unsqueeze(handle, dimension);
+                var res = THSTensor_unsqueeze(handle, dim);
                 if (res == IntPtr.Zero) { torch.CheckForErrors(); }
                 return new Tensor(res);
             }
@@ -4253,9 +4256,9 @@ public Tensor unsqueeze(long dimension)
             ///  Returns a new tensor with a dimension of size one inserted at the specified position.
             ///  The returned tensor shares the same underlying data with this tensor.
             /// </summary>
-            public Tensor unsqueeze_(long dimension)
+            public Tensor unsqueeze_(long dim)
             {
-                var res = THSTensor_unsqueeze_(handle, dimension);
+                var res = THSTensor_unsqueeze_(handle, dim);
                 if (res == IntPtr.Zero) { torch.CheckForErrors(); }
                 return new Tensor(res);
             }

src/TorchSharp/Tensor/Tensor.torch.cs

@@ -37,6 +37,27 @@ public static Tensor cat(IList<Tensor> tensors, long dimension)
             }
         }
 
+        /// <summary>
+        /// Returns a tensor with all the dimensions of input of size 1 removed. When dim is given, a squeeze operation is done only in the given dimension.
+        /// </summary>
+        /// <param name="input"></param>
+        /// <param name="dim">If given, the input will be squeezed only in this dimension</param>
+        /// <returns></returns>
+        public static Tensor squeeze(Tensor input, long? dim = null) => input.squeeze(dim);
+
+
+        /// <summary>
+        ///  Returns a new tensor with a dimension of size one inserted at the specified position.
+        ///  The returned tensor shares the same underlying data with this tensor.
+        /// </summary>
+        public static Tensor unsqueeze(Tensor input, long dim) => input.unsqueeze(dim);
+
+        /// <summary>
+        ///  Returns a new tensor with a dimension of size one inserted at the specified position.
+        ///  The returned tensor shares the same underlying data with this tensor.
+        /// </summary>
+        public static Tensor unsqueeze_(Tensor input, long dim) => input.unsqueeze_(dim);
+
         [DllImport("LibTorchSharp")]
         extern static IntPtr THSTensor_stack(IntPtr tensor, int len, long dim);
 

src/TorchSharp/Torch.cs

@@ -287,6 +287,10 @@ internal static bool CallTorchCudaIsAvailable()
                 return THSTorchCuda_is_available();
             }
 
+            /// <summary>
+            /// Returns a bool indicating if CUDA is currently available.
+            /// </summary>
+            /// <returns></returns>
             public static bool is_available()
             {
                 TryInitializeDeviceType(DeviceType.CUDA);
@@ -305,12 +309,57 @@ public static bool is_cudnn_available()
             [DllImport("LibTorchSharp")]
             private static extern int THSTorchCuda_device_count();
 
+            /// <summary>
+            /// Returns the number of GPUs available.
+            /// </summary>
+            /// <returns></returns>
             public static int device_count()
             {
                 TryInitializeDeviceType(DeviceType.CUDA);
                 return THSTorchCuda_device_count();
             }
 
+            [DllImport("LibTorchSharp")]
+            private static extern void THSCuda_manual_seed(long seed);
+
+            /// <summary>
+            /// Sets the seed for generating random numbers for the current GPU.
+            /// It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.
+            /// </summary>
+            /// <param name="seed">The desired seed.</param>
+            public static void manual_seed(long seed)
+            {
+                TryInitializeDeviceType(DeviceType.CUDA);
+                THSCuda_manual_seed(seed);
+            }
+
+            [DllImport("LibTorchSharp")]
+            private static extern void THSCuda_manual_seed_all(long seed);
+
+            /// <summary>
+            /// Sets the seed for generating random numbers on all GPUs.
+            /// It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.
+            /// </summary>
+            /// <param name="seed"></param>
+            public static void manual_seed_all(long seed)
+            {
+                TryInitializeDeviceType(DeviceType.CUDA);
+                THSCuda_manual_seed_all(seed);
+            }
+
+            [DllImport("LibTorchSharp")]
+            private static extern void THSCuda_synchronize(long device_index);
+
+            /// <summary>
+            /// Waits for all kernels in all streams on a CUDA device to complete.
+            /// </summary>
+            /// <param name="seed">Device for which to synchronize.
+            /// It uses the current device, given by current_device(), if a device is not provided.</param>
+            public static void synchronize(long seed = -1L)
+            {
+                TryInitializeDeviceType(DeviceType.CUDA);
+                THSCuda_synchronize(seed);
+            }
         }
 
         [DllImport("LibTorchSharp")]

src/TorchSharp/TorchVision/Functional.cs

@@ -764,7 +764,7 @@ private static Tensor HSVtoRGB(Tensor h, Tensor s, Tensor v)
                 var q = torch.clamp((v * (1.0 - s * f)), 0.0, 1.0);
                 var t = torch.clamp((v * (1.0 - s * (1.0 - f))), 0.0, 1.0);
 
-                var iunsq = i.unsqueeze(dimension: -3);
+                var iunsq = i.unsqueeze(dim: -3);
                 var mask = iunsq == torch.arange(6, device: i.device).view(-1, 1, 1);
 
                 var a1 = torch.stack(new Tensor[] { v, q, p, p, t, v }, dimension: -3);

test/TorchSharpTest/TestTorchTensor.cs

@@ -4127,6 +4127,13 @@ public void SqueezeTest()
                 Assert.Equal(2.0f, res[1].ToSingle());
                 Assert.Equal(3.1f, res[2].ToSingle());
             }
+            // And all dims.
+            using (var res = Float32Tensor.from(data).expand(new long[] { 1, 1, 3 }).squeeze()) {
+                Assert.Equal(new long[] { 3 }, res.shape);
+                Assert.Equal(1.1f, res[0].ToSingle());
+                Assert.Equal(2.0f, res[1].ToSingle());
+                Assert.Equal(3.1f, res[2].ToSingle());
+            }
         }
 
         [Fact]