Package: pymlir
pyMLIR: Python Interface for the Multi-Level Intermediate RepresentationpyMLIR is a full Python interface to parse, process, and outputMLIRfiles according to the
syntax described in theMLIR documentation. pyMLIR
supports the basic dialects and can be extended with other dialects. It usesLarkto parse the MLIR syntax, and mirrors the classes into Python classes. Custom dialects can also be implemented with a
Python string-format-like syntax, or via direct parsing.Note that the tooldoes not depend on LLVM or MLIR. It can be installed and invoked directly from Python.InstructionsHow to install:pip install pymlirRequirements:Python 3.6 or newer, and the requirements insetup.pyorrequirements.txt. To manually install the
requirements, usepip install -r requirements.txtProblem parsing MLIR files?Run the file through LLVM'smlir-opt --mlir-print-op-genericto get the generic form
of the IR (instructions on how to build/install MLIR can be foundhere):$ mlir-opt file.mlir --mlir-print-op-generic > output.mlirFound other problems parsing files?Not all dialects and modes are supported. Feel free to send us an issue or
create a pull request! This is a community project and we welcome any contribution.Usage examplesParsing MLIR files into Pythonimportmlir# Read a file path, file handle (stream), or a stringast1=mlir.parse_path('/path/to/file.mlir')ast2=mlir.parse_file(open('/path/to/file.mlir','r'))ast3=mlir.parse_string('''module {func @toy_func(%tensor: tensor<2x3xf64>) -> tensor<3x2xf64> {%t_tensor = "toy.transpose"(%tensor) { inplace = true } : (tensor<2x3xf64>) -> tensor<3x2xf64>return %t_tensor : tensor<3x2xf64>}}''')Inspecting MLIR files in PythonMLIR files can be inspected by dumping their contents (which will print standard MLIR code), or by using the same tools
as you would with Python'sastmodule.importmlir# Dump valid MLIR filesm=mlir.parse_path('/path/to/file.mlir')print(m.dump())print('---')# Dump the AST directlyprint(m.dump_ast())print('---')# Or visit each node type by implementing visitor functionsclassMyVisitor(mlir.NodeVisitor):defvisit_Function(self,node:mlir.astnodes.Function):print('Function detected:',node.name.value)MyVisitor().visit(m)Transforming MLIR filesMLIR files can also be transformed with a Python-likeNodeTransformerobject.importmlirm=mlir.parse_path('/path/to/file.mlir')# Simple node transformer that removes all operations with a resultclassRemoveAllResultOps(mlir.NodeTransformer):defvisit_Operation(self,node:mlir.astnodes.Operation):# There are one or more outputs, return None to remove from ASTiflen(node.result_list)>0:returnNone# No outputs, no need to do anythingreturnself.generic_visit(node)m=RemoveAllResultOps().visit(m)# Write back to filewithopen('output.mlir','w')asfp:fp.write(m.dump())Using custom dialectsCustom dialects can be written and loaded as part of the pyMLIR parser.See full tutorial here.importmlirfromlarkimportUnexpectedCharactersfrom.mydialectimportdialect# Try to parse as-istry:m=mlir.parse_path('/path/to/matrixfile.mlir')exceptUnexpectedCharacters:# MyMatrix dialect not recognizedpass# Add dialect to the parserm=mlir.parse_path('/path/to/matrixfile.mlir',dialects=[dialect])# Print output backprint(m.dump_ast())MLIR from scratch with the builder APIpyMLIR has a Builder API that can create MLIR ASTs on the fly within Python code.importmlir.builderbuilder=mlir.builder.IRBuilder()mlirfile=builder.make_mlir_file()module=mlirfile.default_modulewithbuilder.goto_block(builder.make_block(module.region)):hello=builder.function("hello_world")block=builder.make_block(hello.region)builder.position_at_entry(block)x,y=builder.add_function_args(hello,[builder.F64,builder.F64],['a','b'])adder=builder.addf(x,y,builder.F64)builder.ret([adder],[builder.F64])print(mlirfile.dump())prints:module {
  func @hello_world(%a: f64, %b: f64) {
    %_pymlir_ssa = addf %a , %b : f64
    return %_pymlir_ssa : f64
  }
}See alsosaxpyfor a full example that registers and uses a dialect in the builder.Built-in dialect implementations and more examplesAll dialect implementations can be found in thedialectssubfolder. Additional uses
of the library, including a custom dialect implementation, can be found in thetestssubfolder.