xla::ComputationTracker Class Reference

Tracks computations for the XLA service. Registered with a xla::UserComputation instance and can be resolved from a handle for later use. More...

#include <computation_tracker.h>

Public Member Functions
ComputationHandle	NewComputation (const string &computation_name)
	Create a new UserComputation object and return the corresponding ComputationHandle for it.
StatusOr< std::unique_ptr< SessionModule > >	SnapshotComputation (const ComputationHandle &computation)
StatusOr< UserComputation * >	Resolve (const ComputationHandle &computation) const
StatusOr< std::unique_ptr< HloModule > >	BuildHloModule (const VersionedComputationHandle &entry_handle, const HloModuleConfig &config, bool include_unreachable_instructions=true) const
	Build a HLO module (which is basically a set of HLO instructions) using the specified computation as the entry. More...

Detailed Description

Tracks computations for the XLA service. Registered with a xla::UserComputation instance and can be resolved from a handle for later use.

Google Docs:

This class is also capable of serializing/deserializing computations that it tracks (and to serialize properly you need to serialize all referred-to computations as well).

Member Function Documentation

BuildHloModule()

StatusOr< std::unique_ptr< HloModule > > xla::ComputationTracker::BuildHloModule	(	const VersionedComputationHandle &	entry_handle,
		const HloModuleConfig &	config,
		bool	include_unreachable_instructions = true
	)		const

Build a HLO module (which is basically a set of HLO instructions) using the specified computation as the entry.

1. Get the UserComputation object by the entry handle first.
2. Build a list with entry computation and any embedded computations.
3. Traverse the list
   • Build HloComputation for each UserComputation (UserComputation's method xla::UserComputation::BuildHloComputation).
   • If the UserComputation is the entry computation, call HloModule's method xla::HloModule::AddEntryComputation() to add corresponding HloComputation to HloModule, otherwise call the method xla::HloModule::AddEmbeddedComputation().

• Header
• Implementation

Here is the call graph for this function:

Here is the caller graph for this function:

Resolve()

StatusOr< UserComputation * > xla::ComputationTracker::Resolve

(

const ComputationHandle &

computation

)

const

Google Docs:

Resolves a ComputationHandle to a UserComputation that is present in the map.

Here is the caller graph for this function:

SnapshotComputation()

StatusOr< std::unique_ptr< SessionModule > > xla::ComputationTracker::SnapshotComputation

(

const ComputationHandle &

computation

)

Google Docs:

Snapshots a computation (referenced by the provided handle) at its latest version, returning a module where it is the entry, and any referred-to computations are entrained as "embedded" (non-entry) computations.

Here is the call graph for this function:

Here is the caller graph for this function:

---

The documentation for this class was generated from the following files:

• tensorflow/tensorflow/compiler/xla/service/computation_tracker.h
• tensorflow/tensorflow/compiler/xla/service/computation_tracker.cc