xref: /llvm-project/mlir/utils/verify-canon/verify_canon.py (revision ca051dfe9f0996e7fdad8fde5817e5b6ce758ab9) 
1# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
2# See https://llvm.org/LICENSE.txt for license information.
3# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4
5# This script is a helper to verify canonicalization patterns using Alive2
6# https://alive2.llvm.org/ce/.
7# It performs the following steps:
8# - Filters out the provided test functions.
9# - Runs the canonicalization pass on the remaining functions.
10# - Lowers both the original and the canonicalized functions to LLVM IR.
11# - Prints the canonicalized and the original functions side-by-side in a format
12#   that can be copied into Alive2 for verification.
13# Example: `python verify_canon.py canonicalize.mlir -f func1 func2 func3`
14
15import subprocess
16import tempfile
17import sys
18from pathlib import Path
19from argparse import ArgumentParser
20
21
22def filter_funcs(ir, funcs):
23    if not funcs:
24        return ir
25
26    funcs_str = ",".join(funcs)
27    return subprocess.check_output(
28        ["mlir-opt", f"--symbol-privatize=exclude={funcs_str}", "--symbol-dce"],
29        input=ir,
30    )
31
32
33def add_func_prefix(src, prefix):
34    return src.replace("@", "@" + prefix)
35
36
37def merge_ir(chunks):
38    files = []
39    for chunk in chunks:
40        tmp = tempfile.NamedTemporaryFile(suffix=".ll")
41        tmp.write(chunk)
42        tmp.flush()
43        files.append(tmp)
44
45    return subprocess.check_output(["llvm-link", "-S"] + [f.name for f in files])
46
47
48if __name__ == "__main__":
49    parser = ArgumentParser()
50    parser.add_argument("file")
51    parser.add_argument("-f", "--func-names", nargs="+", default=[])
52    args = parser.parse_args()
53
54    file = args.file
55    funcs = args.func_names
56
57    orig_ir = Path(file).read_bytes()
58    orig_ir = filter_funcs(orig_ir, funcs)
59
60    to_llvm_args = ["--convert-to-llvm"]
61    orig_args = ["mlir-opt"] + to_llvm_args
62    canon_args = ["mlir-opt", "-canonicalize"] + to_llvm_args
63    translate_args = ["mlir-translate", "-mlir-to-llvmir"]
64
65    orig = subprocess.check_output(orig_args, input=orig_ir)
66    canonicalized = subprocess.check_output(canon_args, input=orig_ir)
67
68    orig = subprocess.check_output(translate_args, input=orig)
69    canonicalized = subprocess.check_output(translate_args, input=canonicalized)
70
71    enc = "utf-8"
72    orig = bytes(add_func_prefix(orig.decode(enc), "src_"), enc)
73    canonicalized = bytes(add_func_prefix(canonicalized.decode(enc), "tgt_"), enc)
74
75    res = merge_ir([orig, canonicalized])
76
77    print(res.decode(enc))
78