hespas.estimator.cocossim_estimator

Classes

CocossimEstimator([hw_config])

class hespas.estimator.cocossim_estimator.CocossimEstimator(hw_config=None, **kwargs)

Bases: RooflineEstimator

cocossim_path = <hespas.estimator.config_option.ConfigOption object>

frequency = <hespas.estimator.config_option.ConfigOption object>

chips = <hespas.estimator.config_option.ConfigOption object>

conditions = {'cocossim_path': "'which perf_model' failed with: 'Command '['which', 'perf_model']' returned non-zero exit status 1.'"}

stablehlo_to_cocossim_dot_general(op_info, layer_id: int = 0) → List[str]: Convert a StableHLO dot_general to a Cocossim-compatible Matmul M K N. Only non-contract, non-batch dims form M (LHS) and N (RHS). Contract dims form K.

handle_dot_general_cocosim(op_info)

handle_fusion(op_info)

_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 = [<function RooflineEstimator.__setup_roofline_stats>]

_metadata_hooks = []

_module_metadata_hooks = []

_op_handlers = {'func.call': <function RooflineEstimator.handle_free_ops>, 'func.return': <function RooflineEstimator.handle_free_ops>, 'mhlo.bitcast': <function RooflineEstimator.handle_noflop_ops>, 'mhlo.copy': <function RooflineEstimator.handle_noflop_ops>, 'mhlo.fusion': <function CocossimEstimator.handle_fusion>, 'mhlo.ragged_dot': <function RooflineEstimator.handle_ragged_dot>, 'mhlo.return': <function RooflineEstimator.handle_free_ops>, 'stablehlo.abs': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.add': <function RooflineEstimator.handle_elementwise_binary>, 'stablehlo.and': <function RooflineEstimator.handle_elementwise_binary>, 'stablehlo.broadcast_in_dim': <function RooflineEstimator.handle_noflop_ops>, 'stablehlo.clamp': <function RooflineEstimator.handle_clamp>, 'stablehlo.compare': <function RooflineEstimator.handle_elementwise_binary>, 'stablehlo.complex': <function RooflineEstimator.handle_free_ops>, 'stablehlo.concatenate': <function RooflineEstimator.handle_concatenate>, 'stablehlo.constant': <function RooflineEstimator.handle_free_ops>, 'stablehlo.convert': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.convolution': <function RooflineEstimator.handle_convolution>, 'stablehlo.cosine': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.custom_call': <function RooflineEstimator.handle_custom_call>, 'stablehlo.divide': <function RooflineEstimator.handle_elementwise_binary>, 'stablehlo.dot': <function RooflineEstimator.handle_dot_general>, 'stablehlo.dot_general': <function CocossimEstimator.handle_dot_general_cocosim>, 'stablehlo.dynamic_slice': <function RooflineEstimator.handle_slice_ops>, 'stablehlo.dynamic_update_slice': <function RooflineEstimator.handle_dynamic_update_slice>, 'stablehlo.exponential': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.gather': <function RooflineEstimator.handle_gather>, 'stablehlo.get_tuple_element': <function RooflineEstimator.handle_free_ops>, 'stablehlo.imag': <function RooflineEstimator.handle_free_ops>, 'stablehlo.iota': <function RooflineEstimator.handle_free_ops>, 'stablehlo.is_finite': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.log': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.logistic': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.maximum': <function RooflineEstimator.handle_elementwise_binary>, 'stablehlo.minimum': <function RooflineEstimator.handle_elementwise_binary>, 'stablehlo.multiply': <function RooflineEstimator.handle_elementwise_binary>, 'stablehlo.negate': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.not': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.optimization_barrier': <function RooflineEstimator.handle_free_ops>, 'stablehlo.or': <function RooflineEstimator.handle_elementwise_binary>, 'stablehlo.pad': <function RooflineEstimator.handle_noflop_ops>, 'stablehlo.partition_id': <function RooflineEstimator.handle_free_ops>, 'stablehlo.power': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.real': <function RooflineEstimator.handle_free_ops>, 'stablehlo.reduce': <function RooflineEstimator.handle_reduce>, 'stablehlo.reduce_precision': <function RooflineEstimator.handle_noflop_ops>, 'stablehlo.reduce_window': <function RooflineEstimator.handle_reduce_window>, 'stablehlo.replica_id': <function RooflineEstimator.handle_free_ops>, 'stablehlo.reshape': <function RooflineEstimator.handle_noflop_ops>, 'stablehlo.return': <function RooflineEstimator.handle_free_ops>, 'stablehlo.reverse': <function RooflineEstimator.handle_noflop_ops>, 'stablehlo.round_nearest_even': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.rsqrt': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.scatter': <function RooflineEstimator.handle_scatter>, 'stablehlo.select': <function RooflineEstimator.handle_select>, 'stablehlo.select_and_scatter': <function RooflineEstimator.handle_select_and_scatter>, 'stablehlo.sign': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.sine': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.slice': <function RooflineEstimator.handle_slice_ops>, 'stablehlo.sort': <function RooflineEstimator.handle_sort>, 'stablehlo.sqrt': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.subtract': <function RooflineEstimator.handle_elementwise_binary>, 'stablehlo.tanh': <function RooflineEstimator.handle_unary_elemwise>, 'stablehlo.transpose': <function RooflineEstimator.handle_noflop_ops>, 'stablehlo.xor': <function RooflineEstimator.handle_elementwise_binary>}

_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>, <function RooflineEstimator.__get_total_bytes_flops>]

_post_op_hooks = [<function ComputeEstimator.__get__module_op_times>, <function RooflineEstimator.__get_per_op_stats>]

_post_run_hooks = [<function ComputeEstimator.__get_module_runtime>, <function ComputeEstimator.__module_run_end_time>, <function ComputeEstimator.__print_run_runtime>, <function RooflineEstimator.__get_module_bytes_flops>]

_pre_estimate_hooks = [<function ComputeEstimator.__setup_per_module_stat_tree>, <function ComputeEstimator.__total_estimate_start_time>, <function ComputeEstimator.__print_start_line>, <function RooflineEstimator.__setup_module_roofline_stats>]

_pre_op_hooks = [<function ComputeEstimator.__setup_per_op_stat_tree>, <function RooflineEstimator.__setup_per_op_roofline_stats>]

_pre_run_hooks = [<function ComputeEstimator.__module_run_start_time>]

bases_order = {'roofline': 0}

config_arguments = {'cache_dir': <hespas.estimator.config_option.ConfigOption object>, 'chips': <hespas.estimator.config_option.ConfigOption object>, 'cocossim_path': <hespas.estimator.config_option.ConfigOption object>, 'disable_cache': <hespas.estimator.config_option.ConfigOption object>, 'frequency': <hespas.estimator.config_option.ConfigOption object>, 'in_memory_only_cache': <hespas.estimator.config_option.ConfigOption object>, 'num_npus': <hespas.estimator.config_option.ConfigOption object>, 'roofline_error_on_unknown_type': <hespas.estimator.config_option.ConfigOption object>, 'roofline_hbm_power_ratio': <hespas.estimator.config_option.ConfigOption object>, 'roofline_memory_bandwidth': <hespas.estimator.config_option.ConfigOption object>, 'roofline_peak_flops': <hespas.estimator.config_option.ConfigOption object>, 'roofline_per_datatype_flops': <hespas.estimator.config_option.ConfigOption object>, 'roofline_tdp_W': <hespas.estimator.config_option.ConfigOption object>, 'roofline_warn_on_unknown_type': <hespas.estimator.config_option.ConfigOption object>, 'type': <hespas.estimator.config_option.ConfigOption object>}

config_options = {'cache_dir': <hespas.estimator.config_option.ConfigOption object>, 'chips': <hespas.estimator.config_option.ConfigOption object>, 'cocossim_path': <hespas.estimator.config_option.ConfigOption object>, 'disable_cache': <hespas.estimator.config_option.ConfigOption object>, 'frequency': <hespas.estimator.config_option.ConfigOption object>, 'in_memory_only_cache': <hespas.estimator.config_option.ConfigOption object>, 'num_npus': <hespas.estimator.config_option.ConfigOption object>, 'roofline': {'error_on_unknown_type': <hespas.estimator.config_option.ConfigOption object>, 'hbm_power_ratio': <hespas.estimator.config_option.ConfigOption object>, 'memory_bandwidth': <hespas.estimator.config_option.ConfigOption object>, 'peak_flops': <hespas.estimator.config_option.ConfigOption object>, 'per_datatype_flops': <hespas.estimator.config_option.ConfigOption object>, 'tdp_W': <hespas.estimator.config_option.ConfigOption object>, 'warn_on_unknown_type': <hespas.estimator.config_option.ConfigOption object>}, 'type': <hespas.estimator.config_option.ConfigOption object>}

display_name = 'cocossim'

display_name_map = {'roofline': <class 'hespas.estimator.roofline_estimator.RooflineEstimator'>}