Make liveness tracking of LazyTensorOperations a thread local state.

Sources/TensorFlow/Core/LazyTensorContext.swift

@@ -0,0 +1,73 @@
+/// A class to keep track of runtime information about `LazyTensorOperation`
+/// instances created by the program. This will be manaaged as a thread local
+/// state.
+class LazyTensorOperationsTracker {
+    struct RefCounts {
+        let op: LazyTensorOperation
+        let liveRefCount: Int
+        let allRefCount: Int
+    }
+
+    private var refCounts: [ObjectIdentifier: RefCounts] = [:]
+
+    func incrementRefCount(_ op: LazyTensorOperation, isLive: Bool) {
+        let opID = ObjectIdentifier(op)
+        if let counts = refCounts[opID] {
+            refCounts[opID] = RefCounts(
+                op: op,
+                liveRefCount: counts.liveRefCount + (isLive ? 1 : 0),
+                allRefCount: counts.allRefCount + 1)
+        } else {
+            refCounts[opID] = RefCounts(
+                op: op, liveRefCount: isLive ? 1 : 0, allRefCount: 1)
+        }
+    }
+
+    func decrementRefCount(_ op: LazyTensorOperation, isLive: Bool) {
+        let opID = ObjectIdentifier(op)
+        if let counts = refCounts[opID] {
+            if counts.allRefCount > 1 {
+                refCounts[opID] = RefCounts(
+                    op: op,
+                    liveRefCount: counts.liveRefCount - (isLive ? 1 : 0),
+                    allRefCount: counts.allRefCount - 1)
+            } else {
+                refCounts.removeValue(forKey: opID)
+            }
+        }
+    }
+
+    func isLive(_ op: LazyTensorOperation) -> Bool {
+        let opID = ObjectIdentifier(op)
+        if let counts = refCounts[opID] {
+            return counts.liveRefCount > 0
+        }
+        return false
+    }
+
+    func forEachLiveOperation(
+        _ perform: (LazyTensorOperation) throws -> Void
+    ) rethrows -> Void {
+        for (_, counts) in refCounts where counts.liveRefCount > 0 {
+            try perform(counts.op)
+        }
+    }
+
+    func forEachOperation(
+        _ perform: (LazyTensorOperation) throws -> Void
+    ) rethrows -> Void {
+        for (_, counts) in refCounts { try perform(counts.op) }
+    }
+}
+
+struct LazyTensorContext {
+    var scopes = [LazyTensorOperationsTracker()]
+
+    static private var threadLocalContext: LazyTensorContext {
+        _ThreadLocalState.local.lazyTensorContext
+    }
+
+    static var operationsTracker: LazyTensorOperationsTracker {
+        return threadLocalContext.scopes.last!
+    }
+}

Sources/TensorFlow/Core/LazyTensorOperation.swift

@@ -50,80 +50,38 @@ class LazyTensorHandle: _AnyTensorHandle {
         precondition(
             index < op.outputCount, "Symbolic Tensor Index is out-of-bounds")
         handle = Handle.symbolic(op, index: index, isLive: false)
-        LazyTensorHandle.incrementRefCount(op, isLive: false)
+        LazyTensorContext.operationsTracker.incrementRefCount(op, isLive: false)
     }
 
     init(_lazyLive op: LazyTensorOperation, index: Int) {
         precondition(
             index < op.outputCount, "Symbolic Tensor Index is out-of-bounds")
         handle = Handle.symbolic(op, index: index, isLive: true)
-        LazyTensorHandle.incrementRefCount(op, isLive: true)
+        LazyTensorContext.operationsTracker.incrementRefCount(op, isLive: true)
     }
 
     deinit {
         if case let .symbolic(op, _, isLive) = handle {
-            LazyTensorHandle.decrementRefCount(op, isLive: isLive)
+            LazyTensorContext.operationsTracker.decrementRefCount(op, isLive: isLive)
         }
     }
-
+    
     // Liveness tracking for LazyTensorOperations
     //
-    struct LazyTensorOperationRefCounts {
-        let op: LazyTensorOperation
-        let liveRefCount: Int
-        let allRefCount: Int
-    }
-
-    private static var operationRefCounts: [
-        ObjectIdentifier: LazyTensorOperationRefCounts] = [:]
-
-    static func incrementRefCount(_ op: LazyTensorOperation, isLive: Bool) {
-        let opID = ObjectIdentifier(op)
-        if let counts = operationRefCounts[opID] {
-            operationRefCounts[opID] = LazyTensorOperationRefCounts(
-                op: op,
-                liveRefCount: counts.liveRefCount + (isLive ? 1 : 0),
-                allRefCount: counts.allRefCount + 1)
-        } else {
-            operationRefCounts[opID] = LazyTensorOperationRefCounts(
-                op: op, liveRefCount: isLive ? 1 : 0, allRefCount: 1)
-        }
-    }
-
-    static func decrementRefCount(_ op: LazyTensorOperation, isLive: Bool) {
-        let opID = ObjectIdentifier(op)
-        if let counts = operationRefCounts[opID] {
-            if counts.allRefCount > 1 {
-                operationRefCounts[opID] = LazyTensorOperationRefCounts(
-                    op: op,
-                    liveRefCount: counts.liveRefCount - (isLive ? 1 : 0),
-                    allRefCount: counts.allRefCount - 1)
-            } else {
-                operationRefCounts.removeValue(forKey: opID)
-            }
-        }
-    }
-
     static func isLive(_ op: LazyTensorOperation) -> Bool {
-        let opID = ObjectIdentifier(op)
-        if let counts = operationRefCounts[opID] {
-            return counts.liveRefCount > 0
-        }
-        return false
+        return LazyTensorContext.operationsTracker.isLive(op)
     }
 
     static func forEachLiveOperation(
         _ perform: (LazyTensorOperation) throws -> Void
     ) rethrows -> Void {
-        for (_, counts) in operationRefCounts where counts.liveRefCount > 0 {
-            try perform(counts.op)
-        }
+        try LazyTensorContext.operationsTracker.forEachLiveOperation(perform)
     }
 
     static func forEachOperation(
         _ perform: (LazyTensorOperation) throws -> Void
     ) rethrows -> Void {
-        for (_, counts) in operationRefCounts { try perform(counts.op) }
+        try LazyTensorContext.operationsTracker.forEachOperation(perform)
     }
 
     @usableFromInline

Sources/TensorFlow/Core/Runtime.swift

@@ -1193,10 +1193,14 @@ func _TFCOpSetAttrTypeArray(
     }
 }
 
-/// A class to keep around thread local state.
+/// A class to keep around thread local state:
+///  - DeviceScopes
+///  - LazyTensorContext
 class _ThreadLocalState {
     var deviceScopes = DeviceScopes()
 
+    var lazyTensorContext = LazyTensorContext()
+
     private static let key: pthread_key_t = {
         var key = pthread_key_t()
         pthread_key_create(&key) {