annotate_expected_remarks.py

import re
import subprocess
import sys

from black.trans import defaultdict

test_file = "/home/mlevental/dev_projects/triton/test/TritonGPU/amd/amd-range-analysis.mlir"

cmnd = [
    "/home/mlevental/dev_projects/llvm-project/cmake-build-debug/bin/triton-opt",
    "-split-input-file",
    "-allow-unregistered-dialect",
    "-test-tritonamdgpu-range-analysis",
    "-verify-diagnostics",
    test_file,
]

try:
    subprocess.check_output(cmnd, stderr=subprocess.PIPE)
except subprocess.CalledProcessError as e:
    stdout = e.output.decode(sys.getfilesystemencoding())
    stderr = e.stderr.decode(sys.getfilesystemencoding())
else:
    raise RuntimeError("expected failure")


f = stderr.splitlines()
diags = []
for ff in f:
    if "remark" not in ff:
        continue
    if "error: expected remark" in ff or "expected-remark" in ff:
        continue
    split, offset = map(int, re.findall(".*\.mlir:(\d+) offset :(\d+):", ff)[0])
    split -= 1
    diags.append((split + offset, re.findall(" unexpected remark: (.*)", ff)[0]))

g = open(test_file).readlines()

visited = defaultdict(lambda: 0)

for lineno, diag in diags[::-1]:
    # g.insert(lineno - 1 , diag)
    ws = len(g[lineno - 1]) - len(g[lineno - 1].lstrip())
    ws = " " * ws
    visited[lineno] += 1
    below = visited[lineno]
    g.insert(lineno - 1, f"{ws}// expected-remark@+{below} {{{{{diag.strip()}}}}}\n")


for gg in g:
    print(gg.rstrip())