ttir_builder.utils

ttir_builder.utils.shape_str(shape)

Converts shape tuple to string.

Parameters: shape (Union[Tuple[int, …], List[int]]) – Shape to convert to string
Returns: String representation of the shape (e.g., ‘32x32’ for shape (32, 32))
Return type: str

ttir_builder.utils.set_output_path(path)

Sets global output path.

Parameters: path (str) – Path to set as output directory

ttir_builder.utils.get_target_path(output_path, filename, target)

Gets target file path.

Parameters:
- output_path (str) – Base output directory
- filename (str) – Name of the file
- target (str) – Target subdirectory name
Returns: Full path to the target file
Return type: str

ttir_builder.utils.create_custom_pipeline_fn(pipeline, verify=True, print_ir=False)

Creates a custom pipeline function.

Parameters:
- pipeline (str) – Pipeline string specification
- verify (bool , optional) – Whether to enable verification (default: True)
- print_ir (Union[bool, str], optional) – If True or a path string, enables IR printing (default: False)
Returns: Function that runs the custom pipeline on a module
Return type: Callable

ttir_builder.utils.build_mlir_module(fn, inputs_shapes, inputs_types=None, mesh_shape=None, module_dump=False, base=None, output_root='.')

Define a MLIR module specified as a python function.

It will wrap fn in a MLIR FuncOp and then wrap that in a MLIR module, and finally tie arguments of that FuncOp to test function inputs. It will also pass a TTIRBuilder object as the last argument of test function.

Parameters:
- fn (Callable) – Python function to be converted to MLIR
- inputs_shapes (List[Shape]) – Shapes of the respective ranked tensor inputs of the test function.
- module_dump (bool) – Set to True to print out generated MLIR module.
- golden_dump (bool) – Set to True to dump golden info to flatbuffer file.
- inputs_types (List *[*dtype | TypeInfo ] | None)
- mesh_shape (Tuple *[*int , int ] | None)
- base (str | None)
- output_root (str)
Returns: MLIR module containing MLIR op graph defined by fn
Return type: Module

Example

>>> def test_add(in0: Operand, in1: Operand, builder: TTIRBuilder):
...     return builder.add(in0, in1)
...
>>> build_mlir_module(test_add, ((32, 32), (32, 32)))

This returns:

#any = #ttcore.operand_constraint<...>
module {
    func.func @test_add(
        %arg0: tensor<32x32xf32>,
        %arg1: tensor<32x32xf32>
    ) -> tensor<32x32xf32> {
        %0 = ttir.empty() : tensor<32x32xf32>
        %1 = "ttir.add"(%arg0, %arg1, %0) ...
        return %1 : tensor<32x32xf32>
    }
}

Check out: https://github.com/llvm/llvm-project/blob/main/mlir/test/python/dialects/tensor.py

ttir_builder.utils.run_pipeline(module, pipeline_fn=<nanobind.nb_func object>, pipeline_options=None, dump_to_file=True, output_file_name='test.mlir', system_desc_path=None, mesh_shape=None, argument_types_string=None)

Runs a pipeline over a module and optionally dumps to file.

Parameters:
- pipeline_fn (Callable) – Pipeline function to run. pipeline_fn(module, options)
- dump_to_file (bool) – Flag which indicates that generated TTNN module will be dumped to file.
- output_file_name (str) – Name of the output file.
- pipeline_options (List *[*str ] | None)
- system_desc_path (str | None)
- mesh_shape (Tuple *[*int , int ] | None)
- argument_types_string (str | None)
Return type: MLIR module containing MLIR op graph defined by module and pipeline_fn.

ttir_builder.utils.compile_to_flatbuffer(fn, inputs_shapes, inputs_types=None, system_desc_path='ttrt-artifacts/system_desc.ttsys', test_base='test', output_root='.', target='ttnn', mesh_shape=None, module_dump=True, argument_types_string=None, custom_pipeline=None, pipeline_options=None, print_ir=False)

Compiles a TTIRBuilder function fn to TTIR MLIR -> TT{Metal,NN} MLIR -> Flatbuffer.

This decorator is mainly a wrapper around the following functions, with each next function called on the output of the last:

build_mlir_module
run_pipeline
to_flatbuffer

The choice of TTNN vs. TTMetal is controlled by the target parameter.

Parameters:
- fn (Callable) – The TTIRBuilder function to compile. Must take builder : TTIRBuilder as a kwarg.
- inputs_shapes (List[Shape]) – Shapes of the respective ranked tensor inputs of the test function.
- inputs_types (Optional[List[torch.dtype]], optional) – The dtypes to use for the inputs to fn. Note that if supplied, len(inputs_shapes) == len(inputs_types) must be true. Default is None.
- test_base (str) – The string to be used as the base name for dumped files throughout the process. If None is provided, then the __name__ of fn will be used.
- output_root (str) – The path to dump all generated arguments under. If this path doesn’t exist, it will be created.
- target (str) – Either “ttnn” or “ttmetal”. This controls which backend to use.
- custom_pipeline (Union[Callable, str], optional) –
  
  Pipeline function to run. Can be either:
  - A Callable: custom_pipeline(module, options)
  - A str: “ttir-lower-to-layout,ttir-bufferization-pipeline”
- mesh_shape (Optional[Tuple[int, int]], optional) – A list that contains shape of the mesh to be applied on ttir to ttnn conversion path. Default is None.
- module_dump (bool , optional) – Set to True to print out generated TTIR MLIR module. Default is False.
- print_ir (Union[bool, str], optional) – Set to True to print IR to stdout. Set to dir path to print IR after each pass to its own file under that directory. Default is False.
- system_desc_path (str)
- argument_types_string (str | None)
- pipeline_options (List *[*str ] | None)

tt-mlir documentation