split_mlir_ir_log.py

#!/usr/bin/env python3
import re
from pathlib import Path
from typing import List

MARK = "// -----// IR Dump"
PAT  = re.compile(r'IR Dump \w+ (\w+)')

def main(infile: Path, outdir: Path):
    assert infile.exists(), f'{infile} not exists'
    assert not outdir.exists(), f'{outdir} exists'
    outdir.mkdir()
    index = 1
    filename =  outdir / 'tmp'
    file = None
    name = None
    scope = None
    for line in infile.open():
        if line.startswith(MARK):
            if file:
                file.close()
                filename.rename(outdir / f'{index:03}_{name}_@{scope}.mlir')
                index += 1
            file = filename.open('w')
            name = PAT.search(line)[1]
            scope = None
        else:
            if scope is None:
                if line.startswith('module'):
                    scope = 'module'
                elif line.startswith('func.func'):
                    scope = line[len('func.func @'):].split('(', 1)[0]
            file.write(line)
    file.close()
    filename.rename(outdir / f'{index:03}_{name}_@{scope}.mlir')

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(
        description=("split MLIR IRs produced by '-mlir-print-ir-after-all'"
                     " into each numbered files"),
    )
    parser.add_argument('infile', help='the log.mlir file')
    parser.add_argument('--outdir', nargs='?', help='output directory', default="irs")
    args = parser.parse_args()
    main(Path(args.infile), Path(args.outdir))