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

Materialize inputs en masse when executing a 0-output lazy operation. #406

Merged
merged 2 commits into from
Aug 1, 2019
Merged

Materialize inputs en masse when executing a 0-output lazy operation. #406

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ca0d11c
Materialize inputs en-masse when executing a 0-output lazy operation.
bgogul Jul 31, 2019
9097f95
Use Array.compactMap
bgogul Aug 1, 2019
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
29 changes: 23 additions & 6 deletions Sources/TensorFlow/Core/LazyTensorOperation.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ class LazyTensorHandle: _AnyTensorHandle {
case .concrete(_): return nil
}
}

// Liveness tracking for LazyTensorOperations
//
static func isLive(_ op: LazyTensorOperation) -> Bool {
Expand Down Expand Up @@ -374,9 +374,26 @@ extension LazyTensorOperation: TFTensorOperation {
// If we want to stage this, we will need to add control dependencies.
// For the time-being, just build a TFE_Op and run it.
//
// Collect all the unmaterialized inputs.
var unmaterializedInputs = Array<LazyTensorOperation>()
unmaterializedInputs.reserveCapacity(inputs.count)
for input in inputs {
switch input {
case .single(let v):
if let lazyOperation = v.lazyTensorOperation {
unmaterializedInputs.append(lazyOperation)
}
case .list(let values):
unmaterializedInputs.append(
contentsOf: values.lazy.compactMap { $0.lazyTensorOperation }
)
}
}
// Materialize the inputs now.
LazyTensorOperation.materialize(targets: unmaterializedInputs)

// Build the TFEOp and execute.
let op = TFE_Op(name, outputCount)
// TODO(https://bugs.swift.org/browse/TF-604):
// Materialize inputs en masse and not one-by-one.
for input in inputs {
switch input {
case .single(let v):
Expand Down Expand Up @@ -800,7 +817,7 @@ extension LazyTensorOperation {
// Return materialized outputs if any.
if let outputs = outputs { return outputs }

materializeLiveTensors()
LazyTensorOperation.materialize(targets: [self])

// Our outputs should have been updated by now. Otherwise,
// something terrible happened!
Expand Down Expand Up @@ -838,8 +855,8 @@ extension LazyTensorOperation {
inputs = inputs.map { materializedAsNeeded(input: $0) }
}

private func materializeLiveTensors() {
let traceInfo = LazyTensorTraceBuilder.materializationTraceInfo(self)
static func materialize(targets: [LazyTensorOperation]) {
let traceInfo = LazyTensorTraceBuilder.materializationTraceInfo(targets)
debugLog("Extracted trace:\n\(traceInfo.trace)")

let function = TFFunction(trace: traceInfo.trace)
Expand Down