hespas.estimator.zigzag_onnx_estimator

Functions

get_hardware_performance_zigzag(workload, ...)

Classes

ZigzagONNXEstimator([hw_config])

hespas.estimator.zigzag_onnx_estimator._get_hardware_performance_zigzag(workload, accelerator, mapping, dump_folder)

hespas.estimator.zigzag_onnx_estimator.get_hardware_performance_zigzag(workload, accelerator, mapping, dump_folder)

hespas.estimator.zigzag_onnx_estimator._materialize_inline_zigzag_mapping(name: str) → str | None

hespas.estimator.zigzag_onnx_estimator._resolve_local_zigzag_mapping(name: str) → str | None

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

Bases: RooflineEstimator

static get_zigzag_mapping_path(name)

static get_zigzag_accel_path(name)

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

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

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

translate_mlir_to_onnx_model(op_info: Any) → onnx.ModelProto

Translate MLIR (or MLIRModule) to an ONNX-IR model.

static _normalize_mapping_for_onnx(mapping_path: str) → str

ZigZag’s ONNX Gemm/MatMul parser uses loop dims C/D/K. Most provided mappings use M/N (and implicitly K). Translate M->D, N->K in a temp mapping file to avoid per-workload overfitting.

handle_op(op_info)

_cache_hit_hooks = [<function Estimator.__count_cache_hits>, <function Estimator.__get_cached_module_times>, <function Estimator.__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_per_datatype_flops>, <function RooflineEstimator.__setup_roofline_stats>, <function RooflineEstimator.__setup_flops_per_element_map>, <function RooflineEstimator.__setup_seen_unknown_custom_kernel>]

_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_free_ops>, 'mhlo.copy': <function RooflineEstimator.handle_noflop_ops>, 'mhlo.fusion': <function RooflineEstimator.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 ZigzagONNXEstimator.handle_op>, '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 ZigzagONNXEstimator.handle_op>, '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_free_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_free_ops>, 'stablehlo.xor': <function RooflineEstimator.handle_elementwise_binary>}

_post_estimate_hooks = [<function Estimator.__setup_per_op_tree>, <function Estimator.__get_total_estimate_time>, <function Estimator.__get_total_runtime>, <function Estimator.__count_processed>, <function Estimator.__get_per_op_runtime>, <function RooflineEstimator.__get_total_bytes_flops>]

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

_post_run_hooks = [<function Estimator.__get_module_runtime>, <function Estimator.__module_run_end_time>, <function Estimator.__print_run_runtime>, <function RooflineEstimator.__get_module_bytes_flops>, <function RooflineEstimator.__add_kernel_launch_overhead>]

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

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

_pre_run_hooks = [<function Estimator.__module_run_start_time>]

bases_order = {'roofline': 0}

config_arguments = {'accelerator': ('accelerator', <hespas.estimator.config_option.ConfigOption object>), 'cache_dir': ('cache_dir', <hespas.estimator.config_option.ConfigOption object>), 'disable_cache': ('disable_cache', <hespas.estimator.config_option.ConfigOption object>), 'frequency': ('frequency', <hespas.estimator.config_option.ConfigOption object>), 'in_memory_only_cache': ('in_memory_only_cache', <hespas.estimator.config_option.ConfigOption object>), 'mapping': ('mapping', <hespas.estimator.config_option.ConfigOption object>), 'num_npus': ('num_npus', <hespas.estimator.config_option.ConfigOption object>), 'roofline_error_on_unknown_type': ('roofline', 'error_on_unknown_type', <hespas.estimator.config_option.ConfigOption object>), 'roofline_flops_per_element': ('roofline', 'flops_per_element', <hespas.estimator.config_option.ConfigOption object>), 'roofline_hbm_power_ratio': ('roofline', 'hbm_power_ratio', <hespas.estimator.config_option.ConfigOption object>), 'roofline_kernel_launch_overhead_s': ('roofline', 'kernel_launch_overhead_s', <hespas.estimator.config_option.ConfigOption object>), 'roofline_memory_bandwidth': ('roofline', 'memory_bandwidth', <hespas.estimator.config_option.ConfigOption object>), 'roofline_memory_compute_parallelism': ('roofline', 'memory_compute_parallelism', <hespas.estimator.config_option.ConfigOption object>), 'roofline_peak_flops': ('roofline', 'peak_flops', <hespas.estimator.config_option.ConfigOption object>), 'roofline_per_datatype_flops': ('roofline', 'per_datatype_flops', <hespas.estimator.config_option.ConfigOption object>), 'roofline_tdp_W': ('roofline', 'tdp_W', <hespas.estimator.config_option.ConfigOption object>), 'roofline_warn_on_unknown_type': ('roofline', 'warn_on_unknown_type', <hespas.estimator.config_option.ConfigOption object>), 'type': ('type', <hespas.estimator.config_option.ConfigOption object>)}

config_options = {'accelerator': <hespas.estimator.config_option.ConfigOption object>, 'cache_dir': <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>, 'mapping': <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>, 'flops_per_element': <hespas.estimator.config_option.ConfigOption object>, 'hbm_power_ratio': <hespas.estimator.config_option.ConfigOption object>, 'kernel_launch_overhead_s': <hespas.estimator.config_option.ConfigOption object>, 'memory_bandwidth': <hespas.estimator.config_option.ConfigOption object>, 'memory_compute_parallelism': <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 = 'zigzagonnx'

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