Merge branch '1303' into temp

Author: yueyinqiu

RELEASENOTES.md

@@ -8,11 +8,14 @@ __Breaking Changes__:
 
 - `torchvision.dataset.MNIST` will try more mirrors.
     - The thrown exception might be changed when it fails to download `MNIST`, `FashionMNIST` or `KMNIST`.
+- `ObjectDisposedException` will now be thrown when trying to use the disposed dispose scopes.
+    - The constructor of dispose scopes is no longer `public`. Use `torch.NewDisposeScope` instead.
 
 __API Changes__:
 
 - #1291 `Tensor.grad()` and `Tensor.set_grad()` have been replaced by a new property `Tensor.grad`.
     - A potential memory leak caused by `set_grad` has been resolved.
+- `Include` method of dispose scopes has been removed. Use `Attach` instead.
 
 __Bug Fixes__:
 
@@ -25,6 +28,7 @@ __Bug Fixes__:
 
 - `TensorDataset` will now keep the aliases detached from dispose scopes, to avoid the unexpected disposal.
 - `DataLoaderEnumerator` has been completely rewritten to resolve the unexpected shuffler disposal, the ignorance of drop last and the incorrect count of worker.
+- #1303 Allow dispose scopes to be disposed out of LIFO order.
 
 # NuGet Version 0.102.4
 

src/TorchSharp/DisposeScope.cs

@@ -14,50 +14,46 @@ namespace TorchSharp
     /// </summary>
     public sealed class DisposeScope : IDisposable
     {
-        private readonly DisposeScopeManager _disposeScopeManager;
+        private DisposeScopeManager? _disposeScopeManager;
 
-        public DisposeScope(DisposeScopeManager disposeScopeManager)
+        internal DisposeScope(DisposeScopeManager disposeScopeManager)
         {
-            _disposeScopeManager = disposeScopeManager ?? throw new ArgumentNullException(nameof(disposeScopeManager));
-            if (disposeScopeManager.DisposeScopeStack.Count > 0) {
-                OuterScope = disposeScopeManager.DisposeScopeStack.Peek();
-            }
+            _disposeScopeManager = disposeScopeManager;
+            this.OuterScope = disposeScopeManager.CurrentDisposeScope;
         }
 
         /// <summary>
         /// The outer scope with relation to this scope.
         /// </summary>
-        internal DisposeScope? OuterScope { get; }
+        internal DisposeScope? OuterScope { get; set; }
 
         /// <summary>
         /// The disposables that are scheduled for disposing.
         /// </summary>
-        /// TODO: There is a ReferenceEqualityComparer coming in .NET 6, use that!
         internal HashSet<IDisposable> Disposables { get; private set; } =
             new HashSet<IDisposable>(ReferenceEqualityComparer<IDisposable>.Default);
 
         /// <summary>
         /// A view of the disposables in the scope - this list will not be kept in synch with the disposables
         /// in the scope.
         /// </summary>
-        public IReadOnlyList<IDisposable> DisposablesView => Disposables.ToList();
+        public IReadOnlyList<IDisposable> DisposablesView {
+            get {
+                if (this._disposeScopeManager is null)
+                    throw new ObjectDisposedException("The dispose scope has been disposed.");
+                return Disposables.ToArray();
+            }
+        }
 
         /// <summary>
         /// The number of disposables currently held in the scope
         /// </summary>
-        public int DisposablesCount => Disposables.Count;
-
-        /// <summary>
-        /// Includes a disposable in the scope - for tensors this is done automatically once the scope has been
-        /// created. Use this method to add additional disposables that should be disposed, but you typically
-        /// don't need to call this method.
-        /// </summary>
-        /// <param name="disposable">The disposable to keep in the scope</param>
-        /// <returns></returns>
-        public T Include<T>(T disposable) where T : IDisposable
-        {
-            Disposables.Add(disposable);
-            return disposable;
+        public int DisposablesCount {
+            get {
+                if (this._disposeScopeManager is null)
+                    throw new ObjectDisposedException("The dispose scope has been disposed.");
+                return Disposables.Count;
+            }
         }
 
         /// <summary>
@@ -156,6 +152,8 @@ public void MoveToOther(DisposeScope? scope, params IDisposable[] disposables) =
         /// </summary>
         public void MoveToOther(DisposeScope? scope, IEnumerable<IDisposable> disposables)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             foreach (var disposable in disposables) {
                 if (Disposables.Remove(disposable)) {
                     AddToOther(scope, disposable);
@@ -207,6 +205,8 @@ public T Detach<T>(T disposable) where T : IDisposable
         /// </summary>
         public void Detach(IEnumerable<IDisposable> disposables)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             foreach (var disposable in disposables) {
                 if (Disposables.Remove(disposable)) {
                     _disposeScopeManager.StatisticsInstance.DetachedFromScopeCount++;
@@ -219,6 +219,8 @@ public void Detach(IEnumerable<IDisposable> disposables)
 
         public void Attach(IDisposable disposable)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             if (disposable is torch.Tensor tensor) {
                 if (tensor.OwningDisposeScope == null && !tensor.IsInvalid) {
                     _disposeScopeManager.StatisticsInstance.DetachedFromScopeCount--;
@@ -240,6 +242,8 @@ public void Attach(IDisposable disposable)
         /// </summary>
         public void DisposeEverythingBut(IEnumerable<IDisposable> inKeep)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             // Avoiding multiple enumerations
             var oldList = Disposables;
             Disposables = inKeep.ToHashSet(ReferenceEqualityComparer<IDisposable>.Default);
@@ -315,8 +319,11 @@ public T DisposeEverythingBut<T>(T keep) where T : IDisposable
         /// </summary>
         public void Dispose()
         {
+            if (this._disposeScopeManager is null)
+                return;
             DisposeEverything();
             _disposeScopeManager.RemoveDisposeScope(this);
+            this._disposeScopeManager = null;
         }
 
         /// <summary>
@@ -328,6 +335,8 @@ public void Dispose()
         /// <param name="disposable">The disposable that was disposed</param>
         public void MarkAsDisposed(IDisposable disposable)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             _disposeScopeManager.StatisticsInstance.DisposedInScopeCount++;
             Disposables.Remove(disposable);
             if (disposable is torch.Tensor tensor) {
@@ -344,6 +353,8 @@ public void MarkAsDisposed(IDisposable disposable)
 
         private void AddToOther(DisposeScope? scope, IDisposable disposable)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             if (scope != null) {
                 scope.Disposables.Add(disposable);
             } else {

src/TorchSharp/DisposeScopeManager.cs

@@ -1,8 +1,6 @@
 // Copyright (c) .NET Foundation and Contributors.  All Rights Reserved.  See LICENSE in the project root for license information.
 
 using System;
-using System.Collections.Generic;
-using System.Diagnostics;
 
 #nullable enable
 namespace TorchSharp
@@ -15,43 +13,54 @@ namespace TorchSharp
     public class DisposeScopeManager
     {
         [ThreadStatic] private static DisposeScopeManager? _threadSingleton;
-        internal ThreadDisposeScopeStatistics StatisticsInstance { get; } = new ThreadDisposeScopeStatistics();
-
         internal static DisposeScopeManager ThreadSingleton => (_threadSingleton ??= new DisposeScopeManager());
-        internal Stack<DisposeScope> DisposeScopeStack { get; } = new();
 
-        public static ThreadDisposeScopeStatistics Statistics => ThreadSingleton.StatisticsInstance;
+        internal ThreadDisposeScopeStatistics StatisticsInstance { get; } = new ThreadDisposeScopeStatistics();
+        internal DisposeScope? CurrentDisposeScope { get; private set; } = null;
 
-        internal DisposeScope? RegisterOnCurrentDisposeScope(IDisposable disposable)
+        internal DisposeScope? RegisterOnCurrentDisposeScope(torch.Tensor tensor)
         {
-            if (DisposeScopeStack.Count == 0) {
+            if (this.CurrentDisposeScope is null) {
                 StatisticsInstance.CreatedOutsideScopeCount++;
                 return null;
             }
 
             StatisticsInstance.CreatedInScopeCount++;
-            var current = DisposeScopeStack.Peek();
-            current.Include(disposable);
-            return current;
-        }
-
-        internal static DisposeScope NewDisposeScope()
-        {
-            return ThreadSingleton.InnerNewDisposeScope();
+            this.CurrentDisposeScope.Disposables.Add(tensor);
+            return CurrentDisposeScope;
         }
 
         internal void RemoveDisposeScope(DisposeScope disposeScope)
         {
-            Debug.Assert(DisposeScopeStack.Count > 0);
-            Debug.Assert(DisposeScopeStack.Peek() == disposeScope);
-            DisposeScopeStack.Pop();
+            var scope = this.CurrentDisposeScope;
+            if (object.ReferenceEquals(scope, disposeScope)) {
+                this.CurrentDisposeScope = scope.OuterScope;
+                return;
+            }
+            if (scope is null) {
+                return;
+            }
+
+            for (; ; ) {
+                var outerScope = scope.OuterScope;
+                if (object.ReferenceEquals(outerScope, disposeScope)) {
+                    scope.OuterScope = outerScope.OuterScope;
+                    return;
+                }
+
+                if (outerScope is null) {
+                    return;
+                }
+                scope = outerScope;
+            }
         }
 
-        private DisposeScope InnerNewDisposeScope()
+        internal DisposeScope NewDisposeScope()
         {
-            var disposeScope = new DisposeScope(this);
-            DisposeScopeStack.Push(disposeScope);
-            return disposeScope;
+            this.CurrentDisposeScope = new DisposeScope(this);
+            return this.CurrentDisposeScope;
         }
+
+        public static ThreadDisposeScopeStatistics Statistics => ThreadSingleton.StatisticsInstance;
     }
 }

src/TorchSharp/Tensor/Tensor.cs

@@ -7374,7 +7374,8 @@ public static long max_int_value(ScalarType type)
         /// Creates a new dispose scope for the current thread. Any tensor created within the dispose scope will
         /// be automatically disposed once the dispose scope is disposed.
         /// </summary>
-        public static DisposeScope NewDisposeScope() => DisposeScopeManager.NewDisposeScope();
+        public static DisposeScope NewDisposeScope() =>
+            DisposeScopeManager.ThreadSingleton.NewDisposeScope();
 
         /// <summary>
         /// Creates a new dispose scope for the current thread, wrapping an expression.