165 lines · python
1#!/usr/bin/env python32 3"""4This script reads the input from stdin, extracts all lines starting with5"# FDATA: " (or a given prefix instead of "FDATA"), parses the directives,6replaces symbol names ("#name#") with either symbol values or with offsets from7respective anchor symbols, and prints the resulting file to stdout.8"""9 10import argparse11import os12import platform13import shutil14import subprocess15import sys16import re17 18parser = argparse.ArgumentParser()19parser.add_argument("input")20parser.add_argument("objfile", help="Object file to extract symbol values from")21parser.add_argument("output")22parser.add_argument("prefix", nargs="?", default="FDATA", help="Custom FDATA prefix")23parser.add_argument(24 "--nmtool",25 default="llvm-nm" if platform.system() == "Windows" else "nm",26 help="Path to nm tool",27)28parser.add_argument("--no-lbr", action="store_true")29parser.add_argument("--no-redefine", action="store_true")30 31args = parser.parse_args()32 33# Regexes to extract FDATA lines from input and parse FDATA and pre-aggregated34# profile data35prefix_pat = re.compile(f"^(#|//) {args.prefix}: (.*)")36 37# FDATA records:38# <is symbol?> <closest elf symbol or DSO name> <relative FROM address>39# <is symbol?> <closest elf symbol or DSO name> <relative TO address>40# <number of mispredictions> <number of branches>41fdata_pat = re.compile(r"([01].*) (?P<mispred>\d+) (?P<exec>\d+)")42 43# Pre-aggregated profile:44# {T|R|S|E|B|F|f|r} <start> [<end>] [<ft_end>] <count> [<mispred_count>]45# <loc>: [<id>:]<offset>46preagg_pat = re.compile(r"(?P<type>[TRSBFfr]) (?P<offsets_count>.*)")47 48# No-LBR profile:49# <is symbol?> <closest elf symbol or DSO name> <relative address> <count>50nolbr_pat = re.compile(r"([01].*) (?P<count>\d+)")51 52# Replacement symbol: #symname#53replace_pat = re.compile(r"#(?P<symname>[^#]+)#")54 55# Read input and construct the representation of fdata expressions56# as (src_tuple, dst_tuple, mispred_count, exec_count) tuples, where src and dst57# are represented as (is_sym, anchor, offset) tuples58exprs = []59with open(args.input, "r") as f:60 for line in f.readlines():61 prefix_match = prefix_pat.match(line)62 if not prefix_match:63 continue64 profile_line = prefix_match.group(2)65 fdata_match = fdata_pat.match(profile_line)66 preagg_match = preagg_pat.match(profile_line)67 nolbr_match = nolbr_pat.match(profile_line)68 if fdata_match:69 src_dst, mispred, execnt = fdata_match.groups()70 # Split by whitespaces not preceded by a backslash (negative lookbehind)71 chunks = re.split(r"(?<!\\) +", src_dst)72 # Check if the number of records separated by non-escaped whitespace73 # exactly matches the format.74 assert (75 len(chunks) == 676 ), f"ERROR: wrong format/whitespaces must be escaped:\n{line}"77 exprs.append(("FDATA", (*chunks, mispred, execnt)))78 elif nolbr_match:79 loc, count = nolbr_match.groups()80 # Split by whitespaces not preceded by a backslash (negative lookbehind)81 chunks = re.split(r"(?<!\\) +", loc)82 # Check if the number of records separated by non-escaped whitespace83 # exactly matches the format.84 assert (85 len(chunks) == 386 ), f"ERROR: wrong format/whitespaces must be escaped:\n{line}"87 exprs.append(("NOLBR", (*chunks, count)))88 elif preagg_match:89 exprs.append(("PREAGG", preagg_match.groups()))90 else:91 exit("ERROR: unexpected input:\n%s" % line)92 93# Read nm output: <symbol value> <symbol type> <symbol name>94# Ignore .exe on Windows host.95is_llvm_nm = os.path.basename(os.path.realpath(shutil.which(args.nmtool))).startswith(96 "llvm-nm"97)98nm_output = subprocess.run(99 [100 args.nmtool,101 "--defined-only",102 "--special-syms" if is_llvm_nm else "--synthetic",103 args.objfile,104 ],105 text=True,106 capture_output=True,107).stdout108# Populate symbol map109symbols = {}110for symline in nm_output.splitlines():111 symval, _, symname = symline.split(maxsplit=2)112 if symname in symbols and args.no_redefine:113 continue114 symbols[symname] = symval115 116 117def evaluate_symbol(issym, anchor, offsym):118 sym_match = replace_pat.match(offsym)119 if not sym_match:120 # No need to evaluate symbol value, return as is121 return f"{issym} {anchor} {offsym}"122 symname = sym_match.group("symname")123 assert symname in symbols, f"ERROR: symbol {symname} is not defined in binary"124 # Evaluate to an absolute offset if issym is false125 if issym == "0":126 return f"{issym} {anchor} {symbols[symname]}"127 # Evaluate symbol against its anchor if issym is true128 assert anchor in symbols, f"ERROR: symbol {anchor} is not defined in binary"129 anchor_value = int(symbols[anchor], 16)130 symbol_value = int(symbols[symname], 16)131 sym_offset = symbol_value - anchor_value132 return f'{issym} {anchor} {format(sym_offset, "x")}'133 134 135def replace_symbol(matchobj):136 """137 Expects matchobj to only capture one group which contains the symbol name.138 """139 symname = matchobj.group("symname")140 assert symname in symbols, f"ERROR: symbol {symname} is not defined in binary"141 return symbols[symname]142 143 144with open(args.output, "w", newline="\n") as f:145 if args.no_lbr:146 print("no_lbr", file=f)147 for etype, expr in exprs:148 if etype == "FDATA":149 issym1, anchor1, offsym1, issym2, anchor2, offsym2, execnt, mispred = expr150 print(151 evaluate_symbol(issym1, anchor1, offsym1),152 evaluate_symbol(issym2, anchor2, offsym2),153 execnt,154 mispred,155 file=f,156 )157 elif etype == "NOLBR":158 issym, anchor, offsym, count = expr159 print(evaluate_symbol(issym, anchor, offsym), count, file=f)160 elif etype == "PREAGG":161 # Replace all symbols enclosed in ##162 print(expr[0], re.sub(replace_pat, replace_symbol, expr[1]), file=f)163 else:164 exit("ERROR: unhandled expression type:\n%s" % etype)165