|
|
import os |
|
|
import json |
|
|
from tqdm import tqdm |
|
|
import itertools as it |
|
|
import sympy as sp |
|
|
|
|
|
from .disassemble import DisassemblerARM32 |
|
|
from .parser import sympy_to_prefix, isint |
|
|
|
|
|
def match_constants(exprconst, asmconst, constsym, eps=1e-5): |
|
|
def _close(a, b): |
|
|
return abs(a - b) <= eps |
|
|
mapping = {} |
|
|
mapped = set() |
|
|
|
|
|
for ec in exprconst: |
|
|
ecf = float(exprconst[ec]) |
|
|
ecsym = constsym[ec] |
|
|
if abs(ecf) < eps: |
|
|
continue |
|
|
for ac in asmconst: |
|
|
acf = asmconst[ac] |
|
|
acsym = constsym[ac] |
|
|
if _close(acf, ecf): |
|
|
mapping[ecsym] = acsym |
|
|
mapped.add(ec) |
|
|
break |
|
|
if _close(acf, 1/ecf): |
|
|
mapping[ecsym] = 1/acsym |
|
|
mapped.add(ec) |
|
|
break |
|
|
if _close(acf, -ecf): |
|
|
mapping[ecsym] = -acsym |
|
|
mapped.add(ec) |
|
|
break |
|
|
return mapping, mapped |
|
|
|
|
|
def replace_naming(pref): |
|
|
ret = [] |
|
|
for p in pref: |
|
|
if p == "x0": |
|
|
ret.append("x") |
|
|
elif p[0] == "c" and isint(p[1:]): |
|
|
|
|
|
ret.append("k"+p[1:]) |
|
|
else: |
|
|
ret.append(p) |
|
|
return ret |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
import argparse |
|
|
parser = argparse.ArgumentParser("Pre-process assembly to replace constants and dump") |
|
|
parser.add_argument("--list", required=True) |
|
|
parser.add_argument("--prefix", required=True) |
|
|
args = parser.parse_args() |
|
|
|
|
|
with open(args.list, "r") as f: |
|
|
mdllist = list(f) |
|
|
opts = ["O0", "O1", "O2", "O3"] |
|
|
|
|
|
asmf = open(args.prefix + ".asm", "w") |
|
|
eqnf = open(args.prefix + ".eqn", "w") |
|
|
constf = open(args.prefix + ".const.jsonl", "w") |
|
|
|
|
|
basedir = os.path.dirname(args.list) |
|
|
for mdl in tqdm(mdllist): |
|
|
mdl = mdl.strip() |
|
|
mdlname = os.path.basename(mdl) |
|
|
with open(os.path.join(basedir, mdl, "expressions.json")) as f: |
|
|
expressions = json.load(f) |
|
|
yexpr = expressions["expressions"]["y"] |
|
|
exprconsts = {c: float(expressions["constants"][c]) for c in expressions["constants"]} |
|
|
if len(exprconsts) > 4: |
|
|
continue |
|
|
yexpr = sp.parse_expr(yexpr) |
|
|
exprconstsym = {c: sp.Symbol(c) for c in expressions["constants"]} |
|
|
|
|
|
for opt in opts: |
|
|
funcname = f"{mdlname}_run" |
|
|
binf = os.path.join(basedir, mdl, opt, f"c_bin.elf") |
|
|
D = DisassemblerARM32(binf) |
|
|
diss = D.disassemble(funcname) |
|
|
constants = D.constants |
|
|
if len(constants) > 3: |
|
|
continue |
|
|
|
|
|
exprconstsym.update({c: sp.Symbol(f"c{c}") for c in constants}) |
|
|
mapping, mapped = match_constants(exprconsts, constants, exprconstsym) |
|
|
if len(mapped) != len(constants): |
|
|
continue |
|
|
|
|
|
exprsubs = yexpr.subs(mapping) |
|
|
exprprefix = replace_naming(sympy_to_prefix(exprsubs)) |
|
|
|
|
|
asmf.write(diss + "\n") |
|
|
eqnf.write(" ".join(exprprefix) + "\n") |
|
|
constf.write(json.dumps(constants) + "\n") |
|
|
|
|
|
asmf.close() |
|
|
eqnf.close() |
|
|
constf.close() |
|
|
|
