Skip to content
This repository was archived by the owner on Jul 1, 2023. It is now read-only.

Make useLazyTensor a thread-local state. #387

Merged
merged 2 commits into from
Jul 25, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Sources/TensorFlow/Core/LazyTensorTrace.swift
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ class LazyTensorTraceBuilder {

/// Returns a trace obtained by tracing the given function.
static func trace<In: TensorGroup, Out: TensorGroup>(_ fn: (In) -> Out) -> LazyTensorTrace {
precondition(_RuntimeConfig.useLazyTensor, "Lazy tensor is not enabled for tracing.")
precondition(_ThreadLocalState.useLazyTensor, "Lazy tensor is not enabled for tracing.")

// Set up inputs for running `fn`.
let inputOps = In._typeList.map { Self.makePlaceholder(dataType: $0) }
Expand Down
16 changes: 15 additions & 1 deletion Sources/TensorFlow/Core/Runtime.swift
Original file line number Diff line number Diff line change
Expand Up @@ -693,7 +693,7 @@ extension _ExecutionContext {
static func makeOp(
_ name: String, _ outputCount: Int
) -> TFTensorOperation {
return _RuntimeConfig.useLazyTensor
return _ThreadLocalState.useLazyTensor
? LazyTensorOperation(name, outputCount)
: TFE_Op(name, outputCount)
}
Expand Down Expand Up @@ -1201,6 +1201,20 @@ class _ThreadLocalState {

var lazyTensorContext = LazyTensorContext()

static var useLazyTensor: Bool {
get {
_ThreadLocalState.local.lazyTensorEnabled ?? _RuntimeConfig.useLazyTensor
}
set {
_ThreadLocalState.local.lazyTensorEnabled = newValue
}
}

/// When true, use lazy evaluation. If this is not set, we should use the
/// value of `_RuntimeConfig.useLazyTensor` to determine if lazy evaluation
/// is enabled.
private var lazyTensorEnabled: Bool? = nil
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this optional?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This may not always be set. If this is not set, we should use the global _RuntimeConfig.useLazyTensor property to determine if lazy tensor is enabled. I added a comment to explain this.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, isn't using the global _RuntimeConfig.useLazyTensor still not thread-safe? Would it make sense to make this flag the only one?


private static let key: pthread_key_t = {
var key = pthread_key_t()
pthread_key_create(&key) {
Expand Down