hespas.estimator.estimator

Functions

register_cache_hit_hook(func)	Decorator that registers the current method as a hook to call when the given module and hash hit in the cache.
register_cache_miss_hook(func)	Decorator that registers the current method as a hook to call when the given module and hash mises in the cache.
register_default_op_handler(func)	Decorator that registers the current method as a stablehlo operator handler if no op handler can be found.
register_init_hook(func)	Decorator that registers the current method as something to be called at object initialisation.
register_metadata_hook(func)	Decorator that registers the current method as a hook to call to generate metadata for all the modules that have been estimated.
register_module_metadata_hook(func)	Decorator that registers the current method as a hook to call to generate metadata from the given ModuleResult to be used in the EstimatorResult.
register_op_handler(op_names)	Decorator that registers the current method as a stablehlo operator handler.
register_post_estimate_hook(func)	Decorator that registers the current method as a hook to call after estimating the module given as the first argument, and the result of estimation as its second argument.
register_post_op_hook(func)	Decorator that registers the current method as a hook to call after estimating the operator given as the first argument, and the result of the estion as its second argument.
register_post_run_hook(func)	Decorator that registers the current method as a hook to call after estimating the module given as the first argument, and the result of the estimation run as its second argument.
register_pre_estimate_hook(func)	Decorator that registers the current method as a hook to call before estimating the module given as the first argument.
register_pre_op_hook(func)	Decorator that registers the current method as a hook to call before estimating the operator given as the first argument.
register_pre_run_hook(func)	Decorator that registers the current method as a hook to call after estimating the module given as the first argument.

Classes

Estimator([hw_config])

Base class for all Estimators.

hespas.estimator.estimator.register_op_handler(op_names: str | list[str]) → Callable[[object, OpInfo], OpResult]

Decorator that registers the current method as a stablehlo operator handler. The wrapped method should take an OpInfo as the argument, and return an OpResult.

Parameters:

op_names – Either a single string naming a StableHLO operator or a list of strings of StableHLO operator names

Returns:

Annotated method (will be used in class construction to determine the op handlers)

hespas.estimator.estimator.register_default_op_handler(func: Callable[[object, OpInfo], OpResult]) → Callable[[object, OpInfo], OpResult]

Decorator that registers the current method as a stablehlo operator handler if no op handler can be found. The wrapped method should take an OpInfo as the argument, and return an OpResult.

Returns:

Annotated method (will be used in class construction to determine the op handlers)

hespas.estimator.estimator.register_pre_op_hook(func: Callable[[object, OpInfo], Never]) → Callable[[object, OpInfo], Never]

Decorator that registers the current method as a hook to call before estimating the operator given as the first argument. It should not return anything, but may manipulate object fields.

Returns:

Annotated method (will be used in class construction to create the list of pre-op hooks)

hespas.estimator.estimator.register_post_op_hook(func: Callable[[object, OpInfo, OpResult], Never]) → Callable[[object, OpInfo, OpResult], Never]

Decorator that registers the current method as a hook to call after estimating the operator given as the first argument, and the result of the estion as its second argument. It should not return anything, but may manipulate object fields.

Returns:

Annotated method (will be used in class construction to create the list of post-op hooks)

hespas.estimator.estimator.register_pre_run_hook(func: Callable[[object, MLIRModule], Never]) → Callable[[object, MLIRModule], Never]

Decorator that registers the current method as a hook to call after estimating the module given as the first argument. It should not return anything, but may manipulate object fields. This hook is only called when an estimator is actually run for this module, and not when a cached result is returned.

Returns:

Annotated method (will be used in class construction to create the list of pre-run hooks)

hespas.estimator.estimator.register_post_run_hook(func: Callable[[object, MLIRModule, ModuleResult], Never]) → Callable[[object, MLIRModule, ModuleResult], Never]

Decorator that registers the current method as a hook to call after estimating the module given as the first argument, and the result of the estimation run as its second argument. It should not return anything, but may manipulate object fields. This hook is only called when an estimator is actually run for this module, and not when a cached result is returned.

Returns:

Annotated method (will be used in class construction to create the list of post-run hooks)

hespas.estimator.estimator.register_pre_estimate_hook(func: Callable[[object, MLIRModule], Never]) → Callable[[object, MLIRModule], Never]

Decorator that registers the current method as a hook to call before estimating the module given as the first argument. It should not return anything, but may manipulate object fields. This hook is called for every module that is estimated, included cached results.

Returns:

Annotated method (will be used in class construction to create the list of pre-estimate hooks)

hespas.estimator.estimator.register_post_estimate_hook(func: Callable[[object, MLIRModule, EstimatorResult], Never]) → Callable[[object, MLIRModule, EstimatorResult], Never]

Decorator that registers the current method as a hook to call after estimating the module given as the first argument, and the result of estimation as its second argument. It should not return anything, but may manipulate object fields. This hook is called for every module that is estimated, included cached results.

Returns:

Annotated method (will be used in class construction to create the list of post-estimate hooks)

hespas.estimator.estimator.register_module_metadata_hook(func: Callable[[object, ModuleResult], dict[str, Any]]) → Callable[[object, ModuleResult], dict[str, Any]]

Decorator that registers the current method as a hook to call to generate metadata from the given ModuleResult to be used in the EstimatorResult. This should return a (possibly empty) dict of metadata and it will be merged along with other registered metadata hooks. The statistics metadata should be specific to the module within the ModuleResult.

Returns:

Annotated method (will be used in class construction to create the list of metadata hooks)

hespas.estimator.estimator.register_cache_hit_hook(func: Callable[[object, MLIRModule, str], Never]) → Callable[[object, MLIRModule, str], Never]

Decorator that registers the current method as a hook to call when the given module and hash hit in the cache. It should not return anything, but may manipulate object fields.

Returns:

Annotated method (will be used in class construction to create the list of cache hit hooks)

hespas.estimator.estimator.register_cache_miss_hook(func: Callable[[object, MLIRModule, str], Never]) → Callable[[object, MLIRModule, str], Never]

Decorator that registers the current method as a hook to call when the given module and hash mises in the cache. It should not return anything, but may manipulate object fields.

Returns:

Annotated method (will be used in class construction to create the list of cache miss hooks)

hespas.estimator.estimator.register_metadata_hook(func: Callable[[object], dict[str, Any]]) → Callable[[object], dict[str, Any]]

Decorator that registers the current method as a hook to call to generate metadata for all the modules that have been estimated. This should return a (possibly empty) dict of metadata and it will be merged along with other registered metadata hooks. The statistics metadata should be for all modules that have been estimated.

Returns:

Annotated method (will be used in class construction to create the list of metadata hooks)

hespas.estimator.estimator.register_init_hook(func: Callable[[object, OpInfo], OpResult]) → Callable[[object, OpInfo], OpResult]

Decorator that registers the current method as something to be called at object initialisation.

Returns:

Annotated method (will be used in class construction tp create the list of init hooks)

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

Bases: object

Base class for all Estimators. Handles configuration and running of the estimation. Deriving classes should decend from ComputeEstimator, which contains additional functionality relevant to most estimators, with this class being a direct implementation of the API itself.

When an Estimator is called to estimate the runtime of an MLIR module, there are two classes of behaviour that the Estimator will follow depending on the method the new class registers.

The @register_op_handler decorator registers a specific method to handle one or more StableHLO operator names. This method should take a single OpInfo argument, and return an OpResult representing the runtime estimate of this specific operator, in the form

@register_op_handler(['stablehlo.dot_general', 'stablehlo.dot'])
def handle_dot(self, op_info: OpInfo) -> OpResult:
    ...

If a class registers one or more op handlers, or derives from a class that does; the Estimator will use the “per-op” behaviour. This takes a list of operators from an MLIR module, runs registered estimators for each operator found, and returns a result based on summing the runtime of all operators. Deriving classes dot need, and should not, specifically override any method to achieve this; but rather just register op handling functions using the decorator. Deriving classes can also register a default op handler if no handler is registered for the current op by using @register_default_op_handler. If no op handlers are registered, the Estimator will instead use the “whole module” behaviour. This calls the method _get_whole_module_result(), which the deriving class should override with its own implementation, which should take a single MLIR module as its argument, and return a ModuleResult representing the runtime estimate of the module, in the form

def _get_whole_module_result(self, module: MLIRModule) -> ModuleResult:
    ...

These are the only methods that are required to implement Estimator behaviour in the Compute API, as most of the behaviour is contained within the base class. Within API are additional methods for deriving classes to further implement behaviour and configuration

To define configuration options for the new Estimator, instead of directly implementing configuring parsing, new Estimators should use ConfigOptions within the class definition, like so

class NewEstimator(Estimator):
    new_config_option = ConfigOption(str, description="A new configuration option with no default")
    new_default_config_option = ConfigOption(int, description="A new configuration option with a default", default=2)

The first argument should be a “validation” function which returns the validated/converted value (like int()), then all other arguments are optional. If no ‘default’ is specified, the Estimator will throw an error if this configuration option is not given unless the argument ‘optional’ is set to True.

The arguments for a ConfigOption are:

• validation_func: Function to validate and return a validated value

• default: Default value (Optional)

• optional: Whether this value is optional. If a default is specified, this will default to True

• post_validate_func: Function to be called on the value after validation occurs (Optional)

• description: Human readable description of the config option. Use in the argparse argument parser generation

• top_level: Whether descending classes have this in a subdictionary in configuration or at a top level this does not change the field name within the instantiated class, but the arguments to the call to instantiate the class and the configuration dict that is fed in.

During initialisation, each of these options will be assigned to ‘self.new_config_option’ and ‘self.new_default_config_option’ in the instantiated object. Classes that derive from this new estimator will also have these configuration options automatically added. The Estimator class itself also generates arguments to the class initialisation itself from these, as well as parsing of the configuration dict that can be passed in, and argparse arguments if the estimator is used directly as a module.

There are additional methods and decorators to help with implementing Estimator functionality:

__post_init__()

Estimators should not override __init__, but rather if needed define a __post_init__ method that takes no arguments to perform additional initialisation. This will be called Estimator’s __init__ configuration is complete and the cache is loaded if enabled

_add_event()

Estimators should not override _add_event(), but derived classes should call it to add performance counter events to the current performance counter event stream, to be later written out with output_events()

There are a set of ‘hooks’ that are called a specific points in the estimation process, usually prior or subsequent to an estimation. As these registered methods are instance methods of the instantiated Estimator, these have access to the object’s fields and can be used to implement things like statistics or the generation of performance counter events. Multiple methods can be registered to occur when these hooks are called, and registered methods will be inherited from parent classes. The hooks will be called in the order they are registered, and in the order the parent classes are inheriented. As an example:

class OtherEstimator(Estimator):
    ...
    @register_pre_run_hook
    def third_pre_run_hook(...):

class BaseEstimator(Estimator):
    ...
    @register_pre_run_hook
    def new_pre_run_hook(...):
        ...

    @register_pre_run_hook
    def second_pre_run_hook(...):

class ChildEstimator(BaseEstimator, OtherEstimator):
    ...
    @register_pre_run_hook
    def child_pre_run_hook(...):

When the pre-run hooks are called for BaseEstimator, the order will be:

1. new_pre_run_hook()

2. second_pre_run_hook()

And in OtherEstimator:

1. third_pre_run_hook()

And finally ChildEstimator:

1. new_pre_run_hook()

2. second_pre_run_hook()

3. third_pre_run_hook()

4. child_pre_run_hook()

The current list of hooks are:

@register_pre_op_hook

This registers a hook that will be run before an operator estimation is made. This will only be called for estimators that use the per-op behaviour. The registered method should take an OpInfo as its only argument, and not return anything.

@register_post_op_hook

This registers a hook that will be run after an operator estimation is made. This will only be called for estimators that use the per-op behaviour. The registered method should take an OpInfo as its first argument, and an OpResult which represents the completed estimation of the operator given as its first argument as its second argument. It should not return.

@register_pre_run_hook

This registers a hook that will be run before a module estimation is made, but not if the estimation is cached. The registered method should take a MLIRModule as its only argument, and not return anything.

@register_post_run_hook

This registers a hook that will be run after a module estimation is made, but not if the estimation is cached. The registered method should take a MLIRModule as its first argument, and an ModuleResult as its second argument, which represents the completed estimation of the MLIR module given as its first argument. It should not return.

@register_pre_estimate_hook

This registers a hook that will be run before a module estimation is made, regardless of caching. The registered method should take a MLIRModule as its only argument, and not return anything.

@register_post_estimate_hook

This registers a hook that will be run after module estimation is made, regardless of caching. The registered method should take a MLIRModule as its first argument, and an EstimatorResult as its second argument, which represents the completed estimation of the MLIR module given as its first argument. It should not return.

@register_module_metadata_hook

This registers a hook that will be run when the per-module metadata are placed into the EstimatorResult. This only occurs when an estimation is actually run, and not if it is cached. This specific hook runs after the pre-run hooks, but before the post-run hooks. It should take a ModuleResult as its only argument, and return a dict containing a set of metadata to place with the EstimatorResult that is generated for this run. This dict will be merged with all other registered hooks, in the order the hooks are registered.

@register_cache_hit_hook

This registers a hook that will be run on cache hits. The registered method should take a MLIRModule as its first argument, and the cache hash as a string for the module given as its first argument. It should not return.

@register_cache_miss_hook

This registers a hook that will be run on cache misses. The registered method should take a MLIRModule as its first argument, and the cache hash as a string for the module given as its first argument. It should not return.

@register_metadata_hook

This registers a hook that will be run when the public interface ‘get_metadata()’ is called. It should not take any arguments, and return a dict of metadata for all the estimations that have occured. This dict will be merged with all other registered hooks, in the order the hooks are registered.

There are a set of public interface functions implemented as the common interface to interact with Estimators. Derived classes should avoid overriding these and keep them as the public interface, and prefer registering handlers and hooks to modify behaviour instead. The current set of public interfaces are:

get_estimate()

This is the main public interface to run an estimation. This takes an MLIRModule as its only argument to run, and triggers the entire estimation process. It results an EstimatorResult representing the estimation of its argument.

write_cache()

This is the public interface for writing the current cache state to the configured file. Does nothing if the cache is configured to be disabled. It should not be overriden by decending classes, but should be called if the estimator wants to specifically ensure the current cache state is saved.

get_metadata()

This is the public interface for accessing the metadata dict for the set of estimations this instantiated Estimator has performed. The specific content is not defined, but acts as a generic container. All registered metadata_hook functions are call when this interface is used, and their returned dicts are merged in order to produce the dict this method returns.

get_statistics()

This is the public interface for accessing the statistics tree for the set of estimations this instantiated Estimator has performed.

output_events()

This is the public interface to write the performance event counters that have been collected during the set of estimations that this instantiated Estimator has performed to a file. The specification of this file can be found in the CounterEventStream definition file.

compute_estimator = False

conditions = None

allow_multiprocess = False

cache_dir = <hespas.estimator.config_option.ConfigOption object>

disable_cache = <hespas.estimator.config_option.ConfigOption object>

in_memory_only_cache = <hespas.estimator.config_option.ConfigOption object>

num_npus = <hespas.estimator.config_option.ConfigOption object>

__init__(hw_config=None, **kwargs)

Common initialisation for all Estimators. Should not be overriden, handles common configuration.

Parameters:

• hw_config – Dictionary containing the configuration for this estimator. Optional, and will be overriden by anything in kwargs

• kwargs – Dictionary of arguments to apply to this estimator’s config. The legal arguments are defined in the class’s config_arguments attribute

__setup_config_dict(config_dict: dict[str, Any], **kwargs) → dict[str, Any]

Method that combines both the arguments given to the Estimator as well as its pass hw_config into a final dictionary that used for configuration. Checks for config argument existance, but not validity. Should not be overriden nor called manually.

Parameters:

• config_dict – hw_config dictionary passed to __init__()

• kwargs – Override arguments passed to __init__()

Returns:

Elaborated configuration dictionary with overrides applied for use in _configure()

__configure(config_dict=None, config_options=None, subdict_path=None)

Method that configures the current estimator. Each configuration option is checked for validity, and existence; and ensures that all required options are set. Since the configuration options are inherited, this method recursively evaluates and validates the sub-dictionaries within the current configuration dictionary. Each configuration option is set as an attribute. This method should not be overriden or manually called.

Parameters:

• config_dict – The current input dictionary to explore. Any values that are themselves dictionaries will be explored recursively

• config_options – The current set of valid options for this level.

• subdict_path – The current path of subdictionaries that have been traversed. Used for error printing.

__get_cache_writer() → Never

Sets up the cache writer executor

__flush_outstanding_writes() → Never

Collects the outstanding cache futures and flushes them

__close_executor() → Never

Ensure the executor for cache writes is shut down

__load_cache() → Never

Loads the specified cache file to populate the cache, if it exists.

__write_cache(hash: str | None = None) → Never

Write the current cache state to the files within the directory specified as ‘cache_dir’. Should not be called externally or by descending classes - use the ‘write_cache()’ public interface instead.

Parameters:

hash – The specific cache element to write, or the entire cache if the param is None

__get_hash(module: MLIRModule) → str

__create_new_cache_entry(module: MLIRModule, result: EstimatorResult, hash: str | None = None, write: bool = True) → EstimatorResult

Creates a new cache entry with the current module, runtime estimate, and metadata. Write out (by default) Should not be overidden or called manually.

Parameters:

• module – Module that this estimate was generated for

• result – The estimation result generated from the Estimator

• hash – Hash for this combination of hardware configuration and result. Performance optimisation, will be generated if not specified.

Returns:

The estimation result that was cached

__get_op_estimate_by_name(op_name: str, op_info: OpInfo) → OpResult

This method takes a operator name and anop_info object and generates an estimation result for the Estimators that use per-op estimation. If there is no handler found for the specific operator, the default handler is called. The default handler method can be overriden by using the @register_default_op_handler decorator. This should not be called directly, but rather Estimators should use the @register_op_handler decorator with one or more operator names to register handling methods.

Parameters:

• op_name – The operator name to use to index the registered op_handlers

• op_info – The operator to estimate the time for.

Returns:

The estimation result for this operator

__get_op_estimate(op_info: OpInfo) → OpResult

This method takes an op_info object and generates an estimation result for the Estimators that use per-op estimation. If there is no handler found for the specific operator, the default handler is called. The default handler method can be overriden by using the @register_default_op_handler decorator. This should not be called directly, but rather Estimators should use the @register_op_handler decorator with one or more operator names to register handling methods.

Parameters:

op_info – The operator to estimate the time for. A specific handler for that op will be called, or the default

Returns:

The estimation result for this operator

__get_op_estimates(op_info_list: [<class 'hespas.utils.op_info.OpInfo'>]) → [<class 'hespas.estimator.result.OpResult'>]

This method takes a list of operators and calls _get_op_estimate() on each one. Users should only use the @register_op_handler decorator and not directly call this method.

Parameters:

op_info_list – List of operator objects to produce estimations for

Returns:

A list of estimations for the input operator list

__get_op_estimate_list(op_info_list: [<class 'hespas.utils.op_info.OpInfo'>]) → [<class 'hespas.estimator.result.OpResult'>]

This method takes a list of operators and calls _get_op_estimate() on each one. It also calls pre-/post- hooks that are registered for operator estimations. It will return a list of estimation results for each operator. Users should only use the @register_op_handler decorator and not directly call this method.

Parameters:

op_info_list – List of operator objects to produce estimations for

Returns:

A list of estimations for the input operator list

__get_per_op_module_result(module: MLIRModule) → ModuleResult

This method takes a module, splits it into a list of operators, and then calls _get_op_estimate_list() to generate a list of estimation results for each operator. It should not be called or overriden by users, rather Estimators should register handling methods using the @register_op_handler decorator. The list of operator estimations is directly use to instantiate the module result.

Parameters:

module – Module to generate a list of operators from

Returns:

An estimation result for the whole module generated from the list of operator estimations

__get_estimator_run_result(module: MLIRModule) → EstimatorResult

This method takes a module, and runs per-op runtime estimations if any op handlers are defined, otherwise runs the whole module estimate. It also runs the pre-/post- hooks that are registers for each actual (non-cached) estimation run. This will also call get_metadata() and get_statistics() with the resultant ModuleResult in order to populate the metadata and statistics fields of the EstimatorResult. This method should not be directly called or overriden by users, rather either _get_whole_module_result() should be overriden or a set of op handling methods should be registered with @register_op_handler.

Parameters:

module – Module to generate the estimate for

Returns:

A estimation result for the input module

__get_estimator_cached_result(module: MLIRModule, write: bool = True) → EstimatorResult

This method takes a module and checks the cache to see if a matching module estimate has already been made. If it has, return that, otherwise do a new estimation run for the current module and write the result to the cache. It also record both cache hits and misses.This method should not be called directly or overridden by users.

Parameters:

• module – Module to look for in the cache, or if not found, generate a new estimation for

• write – Write the cache to file each time the cache is written to

Returns:

An completed estimation result

__get_estimator_result(module: MLIRModule, write: bool = True) → EstimatorResult

This is the root method of starting an estimate. It will run the pre-/post- hooks that are run for every module, as well getting the estimation result for this module. If caching is enabled, it will attempt to get the result from the cache, otherwise it will run a new estimation every time. This method should not be overriden or directly called by users, rather _get_whole_module_estimate() should be overridden to specify an estimation that takes the entire module, or per-operator estimation handling method should be registered with @register_op_handler.

Parameters:

• module – The module to generate an estimation for

• write – Write the cache to file each time the cache is written to (ignored if the cache is disabled)

__init_hooks() → Never

Runs the init hooks that have been registered. Should not be called directly or overriden, rather users should use the @register_init_hook decorator

__pre_op_hooks(op_info: OpInfo) → Never

Runs the pre-operator estimation hooks that have been registered. Should not be called directly or overriden, rather users should use the @register_pre_op_hook decorator

Parameters:

op_info – The operator that will be passed into the registered hooks

__post_op_hooks(op_info: OpInfo, result: OpResult) → Never

Runs the post-operator estimation hooks that have been registered. Should not be called directly or overriden, rather users should use the @register_post_op_hook decorator

Parameters:

• op_info – The operator that will be passed into the registered hooks

• result – The result of the operator estimation

__pre_run_hooks(module: MLIRModule) → Never

Runs the pre-run estimation hooks that have been registered. Should not be called directly or overriden, rather users should use the @register_pre_run_hook decorator

Parameters:

module – The module that will be passed into the registered hooks

__post_run_hooks(module: MLIRModule, result: ModuleResult) → Never

Runs the post-run estimation hooks that have been registered. Should not be called directly or overriden, rather users should use the @register_post_run_hook decorator

Parameters:

• module – The module that will be passed into the registered hooks

• result – The estimation result of the run

__pre_estimate_hooks(module: MLIRModule) → Never

Runs the pre-estimate hooks that have been registered. Should not be called directly or overriden, rather users should use the @register_pre_estimate_hook decorator

Parameters:

module – The module that will be passed into the registered hooks

__post_estimate_hooks(module: MLIRModule, result: EstimatorResult) → Never

Runs the post-estimate hooks that have been registered. Should not be called directly or overriden, rather users should use the @register_post_estimate_hook decorator

Parameters:

• module – The module that will be passed into the registered hooks

• result – The result of the estimation

__cache_hit_hooks(module: MLIRModule, hash: str) → Never

Runs the cache hit hooks that have been registered. Should not be called directly or overriden, rather users should use the @register_cache_hit_hook decorator

Parameters:

• module – The module that will be passed into the registered hooks

• hash – The hash for this hardware config and module

__cache_miss_hooks(module: MLIRModule, hash: str) → Never

Runs the cache miss hooks that have been registered. Should not be called directly or overriden, rather users should use the @register_cache_miss_hook decorator

Parameters:

• module – The module that will be passed into the registered hooks

• hash – The hash for this hardware config and module

__get_module_metadata(module_result: ModuleResult) → dict[str, Any]

Runs the per-module metadata generatation hooks that have been registered and merges the results from all of the hooks. Should not be called directly or overriden, rather users should use the @register_module_metadata_hook decorator

Parameters:

module_result – The module estimation result that will be passed into the registered hooks

Returns:

A dict containing the metadata for this module’s estimation

__get_module_statistics(module_result: ModuleResult) → dict[str, Any]

Returns the per-module stats, found in the ‘per_module’ subtree, for this module.

Parameters:

module_result – The module estimation result to get the module name from

Returns:

A dict containing the statistics for this module’s estimation

__get_metadata() → dict[str, Any]

Runs the metadata generatation hooks that have been registered and merges the results from all of the hooks. Returns a dict containing the metadata for all modules that have been estimated. Should not be called directly or should not be overriden, rather users should use the @register_metadata_hook decorator

Returns:

A dict containing the metadata for all of the modules that have been estimated

_get_whole_module_result(module: MLIRModule) → ModuleResult

This method takes a module and returns an estimation. This should be overriden by users in an Estimator that generates an estimation using the whole module (like the XLA estimator) to implement the estimation functionality. Besides taking a module as its input and returning an estimation result, the implemention is left to the user.

Parameters:

module – Module to generate the estimate for

Returns:

An estimation result for the whole module

_add_event(event_name: str, timestamp: int | float, timeunit: str = '', metadata: dict[str, Any] | None = None) → Never

Interface to add a performance counter event to the current event stream. Should be called by descendent classes, but not overriden.

Parameters:

• event_name – String specifying the name of the event

• timestamp – Timestamp value of the event

• timeunit – Optional unit for the timestamp

• metadata – Optional metadata for this event

get_metadata() → dict[str, Any]

Returns a dict containing the metadata for all modules that have been estimated. This is the public interface and should be called directly, but should not be overriden, rather users should use the @register_metadata_hook decorator

Returns:

A dict containing the metadata for all of the modules that have been estimated

get_statistics() → ContainerStatistic

Returns the root ContainerStatistic with all statistics collected by this estimator.

Returns:

A ContainerStatistic containing the statistics for all of the modules that have been estimated

write_cache() → Never

Writes the current state of the cache to the configured file if the cache has not been disabled, otherwise does nothing. This is the public interface for writing out the cache file and should be called directly, but should not be overriden.

output_events(file_path: str | Path) → Never

Interface to write out the current set of performance counter events to file. Should be called by externally, but not overriden.

Parameters:

file_path – Path to the file to write the event counter file to

get_estimate(module: MLIRModule) → EstimatorResult

Public facing interface to run an estimation on the given module. Should be called directly as the primary interface to run an estimation. Should not be overriden.

Parameters:

module – The module to run an estimation for

Returns:

An EstimatorResult containing the runtime, metadata, and statistics

classmethod run(hw_config: dict[str, Any], mlir_file: str | Path) → EstimatorResult

Runs the estimator with the specified hardware config for a single MLIR file and returns the result.

Parameters:

• hw_config – Hardware configuration dict for this estimator

• mlir_file – MLIR module file to estimate

Returns:

An EstimatorResult containing the results of estimation

_cache_hit_hooks = []

_cache_miss_hooks = []

_init_hooks = []

_metadata_hooks = []

_module_metadata_hooks = []

_op_handlers = {}

_post_estimate_hooks = []

_post_op_hooks = []

_post_run_hooks = []

_pre_estimate_hooks = []

_pre_op_hooks = []

_pre_run_hooks = []

bases_order = {}

config_arguments = {'cache_dir': <hespas.estimator.config_option.ConfigOption object>, 'disable_cache': <hespas.estimator.config_option.ConfigOption object>, 'in_memory_only_cache': <hespas.estimator.config_option.ConfigOption object>, 'num_npus': <hespas.estimator.config_option.ConfigOption object>}

config_options = {'cache_dir': <hespas.estimator.config_option.ConfigOption object>, 'disable_cache': <hespas.estimator.config_option.ConfigOption object>, 'in_memory_only_cache': <hespas.estimator.config_option.ConfigOption object>, 'num_npus': <hespas.estimator.config_option.ConfigOption object>}

display_name = ''

display_name_map = {}

classmethod main(args: dict[str, Any] | None = None) → Never

Universal commandline interface for running an Estimator

Parameters:

args – An optional dict of args to use in arg parsing.