API Reference¶

class hloenv.HloEnv(*args, **kwargs)¶

The class representing the HloEnv. Each HloEnv instance can be loaded with a single HloModule, and used to run Passes/Pipelines on that module, as well as extract features, evaluate and obtain the graph features for that HloModule.

Overloaded function.

__init__(self: hloenv.HloEnv, hlo_filepath: str, platform: str) -> None

Creates a HloEnv and loads in a HloModule from a specified filepath.

Parameters:	hlo_filepath (str) – The path of the HloModule text file platform (str) – The platform we wish to run the HloModule on. Currently only ‘gpu’ is supported

__init__(self: hloenv.HloEnv, hlo_data: str, format: str, platform: str) -> None

Creates a HloEnv and loads in a HloModule from its string representation.

Parameters:	hlo_data (str) – The HloModule string platform (str) – The platform we wish to run the HloModule on. Currently only ‘gpu’ is supported

__init__(self: hloenv.HloEnv, alt_hlo_module: hloenv.HloModule, platform: str) -> None

Creates a HloEnv and loads in a HloModule from an existing HloModule object.

Parameters:	alt_hlo_module (`HloModule`) – The HloModule object platform (str) – The platform we wish to run the HloModule on. Currently only ‘gpu’ is supported

apply_alternatives(self: hloenv.HloEnv, decisions: numpy.ndarray[numpy.uint64]) → None¶

Applies the specified decisions to the alternative nodes in the HloModule graph, and prunes the resulting graph.

Parameters:	decisions (ndarray) – 2D array of (node_uid, decision) pairs where decision = i indicates that we wish to select the ith alternative at the alternative node corresponding to node_uid.

clone_hlo(self: hloenv.HloEnv) → hloenv.HloModule¶: Clones the currently loaded HloModule and returns it.

evaluate(self: hloenv.HloEnv, times: int = 20, do_not_prep_for_eval: bool = False) → hloenv.EvaluationResult¶

Evaluates the HloModule loaded into the environment N times and returns both the output and the duration of each evaluation.

Parameters:	times (int) – The number of evaluations to perform do_not_prep_for_eval (bool, optional) – Whether to prepare the HloModule for evaluation. This can result in changes to the HloModule (e.g. insertion of Copy instructions), so set this to True if the HloModule has already gone through this process. Defaults to false.
Returns:	The structure containing the durations and output of the evaluation
Return type:	`EvaluationResult`

export_hlo_to_str(self: hloenv.HloEnv) → str¶: Exports the currently loaded HloModule into its text representation. This text can be loaded by the HloEnv.

get_hlo_graph(self: hloenv.HloEnv, debug: bool = False, inline_fused_comp: bool = False, do_hash_verification: bool = False) → hloenv.HloGraph¶: Converts the HloModule into a HloGraph object that describes the features and structure of the HloModule.

get_hlo_module(self: hloenv.HloEnv) → hloenv.HloModule¶: Get the HloModule object loaded in the HloEnv

get_hlo_module_hash(self: hloenv.HloEnv) → int¶: Returns the HloDagHash of the class:HloModule loaded in the environment.

has_equal_output(self: hloenv.HloEnv, first_module: hloenv.HloModule, second_module: hloenv.HloModule, times: int = 1) → bool¶

has_equal_output_as(self: hloenv.HloEnv, other_module: hloenv.HloModule, times: int = 1) → bool¶

load_hlo(*args, **kwargs)¶

Overloaded function.

load_hlo(self: hloenv.HloEnv, hlo_module: hloenv.HloModule) -> None

Loads in a new HloModule from an existing HloModule object.

Args:

hlo_module (HloModule): The HloModule to be loaded in.
load_hlo(self: hloenv.HloEnv, hlo_data: str, format: str = ‘path’) -> None

Loads in a new HloModule from text data.

Args:

hlo_data (str): The HloModule data to be loaded in in the form of a filepath or raw Hlo text string. format (str, optional): The format of the Hlo data. Defaults to “path”.

optimize_hlo_module(self: hloenv.HloEnv) → None¶: Runs the original Xla Optimization Pipeline on the HloModule to obtain a baseline reference against XLA.

prepare_for_eval(self: hloenv.HloEnv) → None¶: Prepare the HloModule for IR emitting and evaluation. This step is automatically run during HloEnv.evaluate, unless the do_not_prep_for_eval parameter is set to True, hence in most cases you will not have to run this function.

run(self: hloenv.HloEnv, pass_pipeline: hloenv::PassInterface) → bool¶

Runs the specified Pass or Pipeline on the HloModule loaded in the environment.

Parameters:	pass_pipeline (Union[`Pass`, `Pipeline`]) – The Pass or Pipeline we wish to run on the `HloModule` loaded in the environment.
Returns:	True if alternatives were generated (i.e. the pass or one of the passes in the pipeline is an AltPipeline), False otherwise. Note that if this returns True, it indicates that an apply_alternatives call must be run to pick the decisions at each alternative before the HloModule can be evaluated.
Return type:	bool

class hloenv.HloGraph(self: hloenv.HloGraph, arg0: xla::HloModule) → None¶

The graph representation of a HloModule the describes its structure and individual instructions features. See Playing with HLO Graph Features for more details.

get_graph_load_errors(self: hloenv.HloGraph) → int¶

hash(self: hloenv.HloGraph) → int¶

class hloenv.HloModule(self: hloenv.HloModule, input: str, format: str = 'path') → None¶

The class representing an XLA HloModule. Each HloModule can be loaded into the HloEnv, where we can evaluate it, obtain its hash, or run specific Pass and Pipeline on it

Creates a HloEnv and loads in a HloModule from a specified filepath.

Parameters:	input (str) – The Hlo text input in the form of a string or filepath. format (str, optional) – The format of the input. Can be either “path” for a filepath to a Hlo text file, or “txt” for the raw Hlo text string. Defaults to “path”.

clone(self: hloenv.HloModule) → hloenv.HloModule¶: Clones the HloModule.

computation_count¶: The number of computations in the HloModule.

config¶: The config options of the HloModule

cost_analysis(self: hloenv.HloModule) → Dict[str, float]¶: Returns a dictionary containing the flops, transcendentals, bytes accessed and optimal seconds of the Hlo Module.

extract_fusions_as_module(self: hloenv.HloModule, arg0: int) → List[hloenv.HloModule]¶

extract_instructions_as_module(self: hloenv.HloModule, arg0: int) → List[Tuple[str, hloenv.HloModule]]¶

extract_random_submodule(self: hloenv.HloModule, arg0: int, arg1: int) → hloenv.HloModule¶

hash(self: hloenv.HloModule) → int¶: The DAGHash of the HloModule. This DAGHash is a custom hash implementation that differs from Tensorflow’s existing HloModule hash implementation to better account for the structure and parameters of the Hlo Instructions.

instruction_count¶: The number of instructions in the HloModule.

is_bef_enabled(self: hloenv.HloModule) → bool¶: Returns whethe Binary Executable Format (BEF) is enabled for the executable.

to_string(self: hloenv.HloModule) → str¶: Converts the HloModule to a string representation. This string representation can also used to initialize a new HloEnv or loaded into an existing one.

class hloenv.EvaluationResult¶

A structure containing the duration and output of the HloModule evaluation.

compute_durations¶: The duration in nanoseconds of the computation, without data transfer, as measured on the device.

durations¶: The default duration in nanoseconds. This returns the execution duration as measured within the Tensorflow evaluation code, starting from the point when the executable has been enqueued on the compute stream till the completion of the executable.

full_durations¶: The full duration in nanoseconds as measured within HloEnv.evaluate(). This captures the entire execution process including processes such as enqueueing the computation on the compute stream, and is hence more subject to timing noise.

output¶: The output of the HloModule.