Skip to content

Instantly share code, notes, and snippets.

@scturtle
Last active February 1, 2024 03:45
Show Gist options
  • Save scturtle/e17000f356e4d4c199e696f6d0259b05 to your computer and use it in GitHub Desktop.
Save scturtle/e17000f356e4d4c199e696f6d0259b05 to your computer and use it in GitHub Desktop.
#!/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))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment