hespas.mlir_parser.mlir_module

MLIR Module wrapper class for handling split MLIR modules with metadata.

This module provides a convenient object-oriented interface for working with MLIR modules that have been split from a larger module, along with their associated metadata.

Classes

MLIRModule(*[, mlir_string, mlir_path, ...])

A wrapper class for MLIR module strings and their associated metadata.

Bases: object

A wrapper class for MLIR module strings and their associated metadata.

This class encapsulates an MLIR module (as a string or file path) along with metadata about the module, such as block type, operation count, input/output dimensions, and other relevant information used in the HESPAS workload analysis.

Each module is automatically assigned a unique incrementing index starting from 0.

Attributes:

mlir_string (str): The MLIR module as a string representation idx (int): Unique incrementing index for this module metadata (Dict[str, Any]): Dictionary containing module metadata

Example:

>>> # Create from MLIR string
>>> module = MLIRModule(
...     mlir_string="module { ... }",
...     block_type=NodeType.COMM_COLL_NODE,
...     parent_module="input.mlir",
... )
>>> print(module.op_count)
42
>>> print(module.idx)
0
>>>
>>> # Create from file path
>>> module = MLIRModule(
...     mlir_string="/path/to/module.mlir",  # File path
...     block_type=1
... )
>>> module.save_to_files("/output/dir")

_next_idx = 0

Initialize an MLIRModule instance.

Args:: mlir_string (str): Either the MLIR module as a string or a file path to an MLIR file block_type (str, optional): Type of the block ( COMM or COMP) parent_module (str, optional): File name of the module from which this was split **kwargs: Additional metadata fields

classmethod reset_index_counter(): Reset the index counter to 0. Useful for testing.

property comm_bytes

property mlir_string: str: Get the MLIR module as a string.

property analyzer

property block_type: NodeType: Get the block type (COMM or COMP).

property is_communication_block: bool: Check if this is a communication block.

property is_computation_block: bool: Check if this is a computation block.

property op_count: int: Get the number of operations in this module.

property op_count_expanded: int: Get the number of operations in this module.

property ops_list

property parent_module: str | None: Get the parent module file name from which this was split.

property module_path: Path | None: Get the module file path, or None if not set. This variable is set either if module is create from file. Or if module is create from string and written to file.

property module_file: str: Get the module file name.

property collective: str | None: Get the collective operation name (for COMM blocks).

property replica_groups: list[list[int]] | None

Get the parsed replica_groups for this collective (for COMM blocks).

Returns a list of groups, where each group is a list of device ids, or None if not a collective or if the collective uses a different attribute (e.g. collective_permute with source_target_pairs).

property input_dims: List[str]: Get the input dimensions as string representations (computed dynamically).

property output_dims: List[str]: Get the output dimensions as string representations (computed dynamically).

property metadata: Dict[str, Any]

Get metadata dictionary (computed dynamically for backward compatibility).

Note: This property exists for backward compatibility. It’s recommended to use individual properties directly instead of accessing metadata.

property input_dims_mlir: List

Get the input dimensions as MLIR RankedTensorType objects.

If the objects don’t exist (e.g., after unpickling), reconstruct them from the MLIR string.

Returns:: List: List of MLIR RankedTensorType objects

property output_dims_mlir: List

Get the output dimensions as MLIR RankedTensorType objects.

If the objects don’t exist (e.g., after unpickling), reconstruct them from the MLIR string.

Returns:: List: List of MLIR RankedTensorType objects

ensure_dir(output_dir)

save_mlir(output_path)

get_json_metadata()

save_json(output_path)

get_paths(output_dir, base_name=None)

save_to_files(output_dir: str | Path, base_name: str | None = None) → tuple

Save the MLIR module and metadata to separate files.

Args:

output_dir (Union[str, Path]): Directory to save files base_name (str, optional): Base name for the files (without extension). If not provided, uses “mini_module_{idx}”

Returns:

tuple: Paths to (mlir_file, json_file)

Example:

>>> module.save_to_files("/output")
... ("/output/mini_module_0.mlir", "/output/mini_module_0.json")

get_output_files(output_dir, base_name=None)

retrieve_mlir_from_file() → None

Retrieve and load the MLIR module content from the associated file.

This method reads the MLIR string from the file specified in the module_file attribute and returns it. It is useful for reloading the MLIR content after it has been cleared or modified.

Returns:: str: The MLIR module content as a string.

update_metadata(**kwargs) → None

Update metadata fields by setting instance variables.

Args:

**kwargs: Metadata fields to update

Example:

>>> module.update_metadata(op_count=50, custom_field="value")

get_metadata_copy() → Dict[str, Any]

Get a copy of the metadata dictionary (computed dynamically).

Returns: Dict[str, Any]: A copy of the metadata

property hash

classmethod get_next_idx() → int

Get the current value of the class-level index counter.

Returns:: int: Current value of _next_idx

classmethod set_next_idx(value: int) → None

Set the class-level index counter to a specific value.

Args:: value (int): New value for _next_idx

static get_abspath(path)

static get_abspath_dir(path)