hespas.estimator.xla_estimator

Classes

XLAEstimator([hw_config])

class hespas.estimator.xla_estimator.XLAEstimator(hw_config=None, **kwargs)

Bases: ComputeEstimator

xla_translate_path = <hespas.estimator.config_option.ConfigOption object>
hlo_runner_main_path = <hespas.estimator.config_option.ConfigOption object>
sample = <hespas.estimator.config_option.ConfigOption object>
translate = <hespas.estimator.config_option.ConfigOption object>
conditions = {'hlo_runner_main_path': "'which hlo_runner_main' failed with: 'Command '['which', 'hlo_runner_main']' returned non-zero exit status 1.'", 'xla_translate_path': "'xla-translate -h' failed with: '[Errno 2] No such file or directory: 'xla-translate''"}
__get_hlo_file_path(module)
get_mlir_to_hlo_args(module)

Translates MLIR to HLO using the XLA translation tool.

get_exec_args(module)
generic_compile(module)

Compiles an MLIR module string into an IREE VMFB (VM Flatbuffer) file.

Args:

module_str (str): The MLIR module representation as a string. output_vmfb_file (str): The path to the output VMFB file.

Uses:

  • Calls ireec.compile_str to compile the module.

  • Uses the target backend and extra arguments from self.hw_config[‘compile’].

  • Do not include any trailing or leading white spaces in the configuration json

    for the compilation parameters.

  • Depending on the specific GPU, the “sm” version needs to be changed

Returns:

True if the compilation was successful, False otherwise.

xla_runtime_estimate(module)

Measures runtime performance of an XLA compiled model in seconds. :return: Estimated runtime in seconds or None if execution fails.

_get_whole_module_result(module)

Compile the MLIR module using IREE and execute it on the target CPU to obtain a runtime estimate in microseconds. :return: Estimated runtime in microseconds or None if execution fails.

_parse_benchmark_output(output)

Extract runtime estimates from the benchmark output. :param output: The output string from the benchmark execution. :return: Runtime estimate in seconds or None if parsing fails.

_cache_hit_hooks = [<function ComputeEstimator.__count_cache_hits>, <function ComputeEstimator.__get_cached_module_times>, <function ComputeEstimator.__print_cached_runtime>]
_cache_miss_hooks = []
_default_op_handler(op_info: OpInfo) → OpResult

This is a default for descended classes that an exception is thrown if the operator isn’t known. This can be overriden through the @register_default_op_handler decorator, but not directly. This should not be directly overidden or called manually.

Parameters:

op_info – The operator to estimate the time of. This method will not estimate, and just throw an exception

Returns:

The result of the estimator for this operator (to match the type of operator estimators - will not return)

Raises:

InvalidOpError – Raises an InvalidOpError for any unknown operator

_init_hooks = []
_metadata_hooks = []
_module_metadata_hooks = []
_op_handlers = {}
_post_estimate_hooks = [<function ComputeEstimator.__setup_per_op_tree>, <function ComputeEstimator.__get_total_estimate_time>, <function ComputeEstimator.__get_total_runtime>, <function ComputeEstimator.__count_processed>, <function ComputeEstimator.__get_per_op_runtime>]
_post_op_hooks = [<function ComputeEstimator.__get__module_op_times>]
_post_run_hooks = [<function ComputeEstimator.__get_module_runtime>, <function ComputeEstimator.__module_run_end_time>, <function ComputeEstimator.__print_run_runtime>]
_pre_estimate_hooks = [<function ComputeEstimator.__setup_per_module_stat_tree>, <function ComputeEstimator.__total_estimate_start_time>, <function ComputeEstimator.__print_start_line>]
_pre_op_hooks = [<function ComputeEstimator.__setup_per_op_stat_tree>]
_pre_run_hooks = [<function ComputeEstimator.__module_run_start_time>]
bases_order = {'compute': 0}
config_arguments = {'cache_dir': <hespas.estimator.config_option.ConfigOption object>, 'disable_cache': <hespas.estimator.config_option.ConfigOption object>, 'hlo_runner_main_path': <hespas.estimator.config_option.ConfigOption object>, 'in_memory_only_cache': <hespas.estimator.config_option.ConfigOption object>, 'num_npus': <hespas.estimator.config_option.ConfigOption object>, 'sample': <hespas.estimator.config_option.ConfigOption object>, 'translate': <hespas.estimator.config_option.ConfigOption object>, 'type': <hespas.estimator.config_option.ConfigOption object>, 'xla_translate_path': <hespas.estimator.config_option.ConfigOption object>}
config_options = {'cache_dir': <hespas.estimator.config_option.ConfigOption object>, 'disable_cache': <hespas.estimator.config_option.ConfigOption object>, 'hlo_runner_main_path': <hespas.estimator.config_option.ConfigOption object>, 'in_memory_only_cache': <hespas.estimator.config_option.ConfigOption object>, 'num_npus': <hespas.estimator.config_option.ConfigOption object>, 'sample': <hespas.estimator.config_option.ConfigOption object>, 'translate': <hespas.estimator.config_option.ConfigOption object>, 'type': <hespas.estimator.config_option.ConfigOption object>, 'xla_translate_path': <hespas.estimator.config_option.ConfigOption object>}
display_name = 'xla'
display_name_map = {'compute': <class 'hespas.estimator.compute_estimator.ComputeEstimator'>}