Add REMEND python module

pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "remend"
+version = "1.0"
+authors = [{name="Meet Udeshi", email="m.udeshi@nyu.edu"}]
+description = "Neural Decompilation for Reverse Engineering Math Equations from Binary Executables"
+readme = "README.md"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+requires-python = ">=3.9"
+dependencies = [
+    "networkx",
+    "capstone",
+    "Levenshtein",
+    "tqdm",
+    "numpy",
+    "sympy",
+    "fairseq",
+    "torch",
+    "matplotlib",
+    "tokenizers"
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["remend"]

remend/bpe.py

@@ -0,0 +1,64 @@
+from tokenizers import pre_tokenizers, Tokenizer
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+from tokenizers.pre_tokenizers import Whitespace, PreTokenizer
+import random
+import os
+from tqdm import tqdm
+import itertools as it
+
+class ImmPreTokenizer:
+    def pre_tokenize(self, pretok):
+        pretok.split(self.hex_imm_split)
+    def hex_imm_split(self, i, norm_str):
+        tok = str(norm_str)
+        if tok[:2] == "0x" or tok.isdigit():
+            return [norm_str[i:i+1] for i in range(len(tok))]
+        else:
+            return [norm_str]
+
+def get_asm_tok(files, save):
+    asm_tok = Tokenizer(BPE(unk_token="@@UNK@@"))
+    asm_tok.pre_tokenizer = pre_tokenizers.Sequence([Whitespace(), PreTokenizer.custom(ImmPreTokenizer())])
+    asm_train = BpeTrainer(special_tokens=["@@UNK@@"])
+
+    asm_tok.train(files, asm_train)
+    asm_tok.pre_tokenizer = Whitespace() # Hack to save, careful to restore ImmPreTokenizer
+    asm_tok.save(save)
+    asm_tok.pre_tokenizer = pre_tokenizers.Sequence([Whitespace(), PreTokenizer.custom(ImmPreTokenizer())])
+
+    return asm_tok
+
+def load_asm_tok(load):
+    asm_tok = Tokenizer.from_file(load)
+    asm_tok.pre_tokenizer = pre_tokenizers.Sequence([Whitespace(), PreTokenizer.custom(ImmPreTokenizer())])
+    return asm_tok
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Train the tokenizer and tokenize the asm")
+    parser.add_argument("-i", "--indir", required=True, help="output directory")
+    parser.add_argument("-o", "--outdir", default="tokenized", help="output directory")
+    args = parser.parse_args()
+
+    os.makedirs(args.outdir, exist_ok=True)
+    injoin = lambda p: os.path.join(args.indir, p)
+    pjoin = lambda p: os.path.join(args.outdir, p)
+    max_asm_toks = 0
+
+    asm_tok = get_asm_tok([injoin("train.asm"), injoin("valid.asm")], pjoin("asm_tokens.json"))
+    for split in ["train", "valid", "test"]:
+        asmfile = split + ".asm"
+        with open(injoin(asmfile), "r") as asmf, open(pjoin(asmfile), "w") as asmtokf:
+            for asm in tqdm(asmf, desc=f"Tokenizing {split}"):
+                asm = asm.strip()
+                asm_enc = asm_tok.encode(asm)
+                max_asm_toks = max(max_asm_toks, len(asm_enc.tokens))
+                asm_seq = " ".join(asm_enc.tokens)
+                asmtokf.write(asm_seq + "\n")
+
+    print("Maximum tokens:", max_asm_toks)
+
+# After this, run command:
+# fairseq-preprocess -s asm -t eqn --trainpref {OUTDIR}/train --validpref {OUTDIR}/valid --testpref {OUTDIR}/test --destdir {OUTDIR}

remend/bpe_apply.py

@@ -0,0 +1,25 @@
+from tqdm import tqdm
+
+from .bpe import load_asm_tok
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Tokenize using existing tokenizer")
+    parser.add_argument("-t", "--tokenizer", required=True, help="existing tokenizer")
+    parser.add_argument("-i", "--input", required=True, help="input file")
+    parser.add_argument("-o", "--output", required=True, help="output file")
+    args = parser.parse_args()
+
+    max_asm_toks = 0
+    asm_tok = load_asm_tok(args.tokenizer)
+
+    with open(args.input, "r") as asmf, open(args.output, "w") as asmtokf:
+        for asm in tqdm(asmf, desc=f"Tokenizing"):
+            asm = asm.strip()
+            asm_enc = asm_tok.encode(asm)
+            max_asm_toks = max(max_asm_toks, len(asm_enc.tokens))
+            asm_seq = " ".join(asm_enc.tokens)
+            asmtokf.write(asm_seq + "\n")
+
+    print("Maximum tokens:", max_asm_toks)
+

remend/change_eqn_format.py

@@ -0,0 +1,79 @@
+from .parser import isint, OPERATORS
+
+def prefix_to_brackets(eqn):
+    stack = []
+    lastop = []
+    intunit = []
+    N = len(eqn)
+    i = 0
+    while i < N:
+        # print("Stack", stack)
+        val = eqn[i]
+        if val.startswith("INT"):
+            intunit.append(val)
+            i += 1
+            while i < N and isint(eqn[i]):
+                intunit.append(eqn[i])
+                i += 1
+            stack.append(" ".join(intunit))
+            intunit = []
+            i -= 1
+        elif val in OPERATORS:
+            _, numops = OPERATORS[val]
+            lastop.append((len(stack), numops))
+            stack.append(val)
+        else:
+            stack.append(val)
+
+        while len(lastop) > 0  and len(stack) > lastop[-1][0] + lastop[-1][1]:
+            # Combine op
+            # print(lastop[-1], stack[lastop[-1][0]:])
+            op = " ".join(stack[lastop[-1][0]:])
+            del stack[lastop[-1][0]:]
+            lastop.pop()
+            stack.append(f"( {op} )")
+        i += 1
+    assert(len(stack) == 1)
+    return stack[0]
+
+def prefix_to_postfix(eqn):
+    if eqn[0].startswith("INT"):
+        intunit = [eqn[0]]
+        for i, val in enumerate(eqn[1:]):
+            if not isint(val):
+                break
+            intunit.append(val)
+        return intunit, eqn[i+1:]
+    elif eqn[0] in OPERATORS:
+        _, numops = OPERATORS[eqn[0]]
+        remeqn = eqn[1:]
+        ops = []
+        for i in range(numops):
+            op, remeqn = prefix_to_postfix(remeqn)
+            ops.extend(op)
+        ops.append(eqn[0]) # Restructured to postfix
+        return ops, remeqn
+    else:
+        return [eqn[0]], eqn[1:]
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Change equation format from prefix to other")
+    parser.add_argument("--eqn", required=True)
+    parser.add_argument("--out", required=True)
+    args = parser.parse_args()
+
+    with open(args.eqn, "r") as inf, open(args.out, "w") as outf:
+        for eqn in inf:
+            postfix, _ = prefix_to_postfix(eqn.strip().split(" "))
+            outf.write(" ".join(postfix) + "\n")
+
+    # eqn = "div mul x add INT+ 5 add mul INT+ 3 x mul pow x INT+ 2 add INT- 5 add mul INT- 3 x mul x mul add INT+ 1 mul k0 pow x INT+ 3 add INT+ 4 x add INT+ 5 mul INT+ 3 x"
+    # eqn = "div add mul INT+ 3 x pow x INT- 4 mul sub x k0 add mul INT+ 5 x k1"
+    # print(" ".join(prefix_to_postfix(eqn.split(" "))[0]))
+    # postfix = "x INT+ 3 mul INT+ 5 add x INT+ 4 add INT+ 3 x pow k0 mul INT+ 1 add mul x mul x INT- 3 mul add INT- 5 add INT+ 2 x pow mul x INT+ 3 mul add INT+ 5 add x mul div"
+    # print(prefix_to_brackets(eqn.split(" ")))
+
+    # (div (mul x (add INT+ 5 (add (mul INT+ 3 x) (mul (pow x INT+ 2) (add INT- 5 (add (mul INT- 3 x) (mul x (mul (add INT+ 1 (mul k0 (pow x INT+ 3))) (add INT+ 4 x))))))))) (add INT+ 5 (mul INT+ 3 x)))
+    # ( div ( mul x ( add INT+ 5 ( add ( mul INT+ 3 x ) ( mul ( pow x INT+ 2 ) ( add INT- 5 ( add ( mul INT- 3 x ) ( mul x ( mul ( add INT+ 1 ( mul k0 ( pow x INT+ 3 ) ) ) ( add INT+ 4 x ) ) ) ) ) ) ) ) ) ( add INT+ 5 ( mul INT+ 3 x ) ) )
+    

remend/check_generated.py

@@ -0,0 +1,143 @@
+import sympy as sp
+import sys
+import re
+from tqdm import tqdm
+from Levenshtein import distance
+import networkx as nx
+from networkx import graph_edit_distance
+
+from .parser import parse_prefix_to_sympy, parse_postfix_to_sympy, isint
+
+def percent(a, n):
+    return f"{a/n*100:0.1f}%"
+
+def do_simplify_match(orig_expr, gen_expr):
+    orig_simp = sp.simplify(orig_expr)
+    gen_simp = sp.simplify(gen_expr)
+    if orig_simp == gen_simp:
+        return True
+    return False
+
+def do_structure_match(orig_toks, gen_toks):
+    def _isconst(t):
+        return re.match(r"c[0-9]+", t)
+    def _isvar(t):
+        return re.match(r"x[0-9]+", t)
+    if len(orig_toks) != len(gen_toks):
+        return False
+    for orig, gen in zip(orig_toks, gen_toks):
+        if (_isconst(orig) and _isconst(gen)) \
+                or (_isvar(orig) and _isvar(gen)) \
+                or (isint(orig) and isint(gen)) \
+                or (orig.startswith("INT") and gen.startswith("INT")) \
+                or (orig == gen):
+            continue
+        # Mismatched
+        return False
+    return True
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Check generated expressions")
+    parser.add_argument("-g", required=True, help="Generated expressions file")
+    parser.add_argument("-r", required=True, help="Results file")
+    parser.add_argument("--simplify", action="store_true", default=False)
+    parser.add_argument("--postfix", action="store_true", default=False)
+    args = parser.parse_args()
+    
+
+    orig_list = []
+    gen_list = []
+    with open(args.g, 'r') as f:
+        for line in tqdm(f, desc="Reading file"):
+            comps = line.strip().split("\t")
+            if line[0] == 'T':
+                num = int(comps[0][2:])
+                tokens = comps[1].split(" ")
+                orig_list.append((num, tokens))
+            elif line[0] == 'H':
+                num = int(comps[0][2:])
+                tokens = comps[2].split(" ")
+                gen_list.append((num, tokens))
+
+    N = len(orig_list)
+    gen_errors = []
+    parsed = []
+    exact_match = []
+    structure_match = []
+    simplify_match = []
+
+    orig_exprs = {}
+    gen_exprs = {}
+
+    all_aed = []
+    # all_ged = []
+
+    results = []
+
+    for (orig_num, orig_toks), (gen_num, gen_toks) in tqdm(zip(orig_list, gen_list), desc="Parsing expressions", total=N):
+        assert orig_num == gen_num
+        aed = distance(orig_toks, gen_toks) / (len(orig_toks) + len(gen_toks))
+        all_aed.append(aed)
+        res = {"id": gen_num, "aed": aed, "matched": False, "parsed": False}
+
+        if aed == 0:
+            parsed.append(orig_num)
+            exact_match.append(orig_num)
+            structure_match.append(orig_num)
+            res["parsed"] = True
+            res["matched"] = "Exact"
+            results.append(res)
+            continue
+
+        if do_structure_match(orig_toks, gen_toks):
+            structure_match.append(orig_num)
+            res["matched"] = "Structure"
+
+        if "<<unk>>" in orig_toks:
+            # Why this happened?
+            res["parsed"] = False
+            res["matched"] = False
+            results.append(res)
+            continue
+
+        if args.postfix:
+            orig_expr = parse_postfix_to_sympy(orig_toks)
+        else:
+            orig_expr = parse_prefix_to_sympy(orig_toks)
+        try:
+            if args.postfix:
+                gen_expr = parse_postfix_to_sympy(gen_toks)
+            else:
+                gen_expr = parse_prefix_to_sympy(gen_toks)
+            res["parsed"] = True
+        except: # Exception as e:
+            gen_errors.append(gen_num)
+            results.append(res)
+            continue
+
+        parsed.append(gen_num)
+        orig_exprs[gen_num] = orig_expr
+        gen_exprs[gen_num] = gen_expr
+
+        if orig_expr == gen_expr:
+            exact_match.append(gen_num)
+            res["matched"] = "Exact"
+        elif args.simplify and do_simplify_match(orig_expr, gen_expr):
+            simplify_match.append(gen_num)
+            res["matched"] = "Simplify"
+        results.append(res)
+
+    with open(args.r, "w") as resf:
+        for res in results:
+            resf.write("{id} {aed} {parsed} {matched}\n".format(**res))
+        resf.write("\n")
+        print("Total", N, file=resf)
+        print("Parse error", len(gen_errors), percent(len(gen_errors), N), file=resf)
+        print("Exact match", len(exact_match), percent(len(exact_match), N), file=resf)
+        print("Structure match", len(structure_match), percent(len(structure_match), N), file=resf)
+        if args.simplify:
+            print("Simplify match", len(simplify_match), percent(len(simplify_match), N), file=resf)
+        print("Avg SED", sum(all_aed) / len(all_aed), max(all_aed), file=resf)
+        # print("Avg GED", sum(all_ged) / len(all_ged), max(all_ged), file=resf)
+

remend/compile_dataset.py

@@ -0,0 +1,185 @@
+from tqdm import tqdm
+import random
+import sympy as sp
+import json
+import subprocess as sproc
+from os.path import realpath, dirname, join as pjoin
+from os import makedirs
+import multiprocessing as mp
+from time import sleep
+import logging
+
+from .implementation import Implementor
+from .parser import parse_prefix_to_sympy, sympy_to_prefix, constant_fold
+from .disassemble import DisassemblerARM32, DisassemblerAArch64, DisassemblerX64
+from .util import DecodeError, timeout, sympy_expr_ok
+
+SCRIPT = pjoin(dirname(realpath(__file__)), "compile_eqn.sh")
+
+QUEUE_END = "QUEUE_END_SENTINEL"
+
+def compile_c(code, elf, arch="arm32", src="/tmp/myfunc.c", opt=0):
+    with open(src, "w") as f:
+        f.write(code)
+    ret = sproc.run(["bash", "-e", SCRIPT, arch+"-c", src, elf, f"-O{opt}"], capture_output=True)
+    if ret.returncode != 0:
+        raise DecodeError("compile failed")
+
+def compile_fortran(code, elf, arch="arm32", src="/tmp/myfunc.f95", opt=0):
+    with open(src, "w") as f:
+        f.write(code)
+    ret = sproc.run(["bash", "-e", SCRIPT, arch+"-fortran", src, elf, f"-O{opt}"], capture_output=True)
+    if ret.returncode != 0:
+        raise DecodeError("compile failed")
+
+class EquationCompiler:
+    def __init__(self, q, arch, impl, opt, outdir, prefix, dtype="double"):
+        if "fortran" in impl:
+            self.compiler = compile_fortran
+        else:
+            self.compiler = compile_c
+
+        if arch == "arm32":
+            self.disassembler = DisassemblerARM32
+        elif arch == "aarch64":
+            self.disassembler = DisassemblerAArch64
+        elif arch == "x64":
+            self.disassembler = DisassemblerX64
+        else:
+            raise DecodeError("arch not supported: " + arch)
+
+        self.q = q
+        self.impl = impl
+        self.opt = opt
+        self.outdir = outdir
+        self.prefix = prefix
+        self.dtype = dtype
+        self.arch = arch
+
+    def run(self):
+        outdir = pjoin(self.outdir, f"O{self.opt}", self.impl)
+        makedirs(outdir, exist_ok=True)
+        outfiles = {
+            "asm": open(pjoin(outdir, self.prefix + ".asm"), "w"),
+            "eqn": open(pjoin(outdir, self.prefix + ".eqn"), "w"),
+            "src": open(pjoin(outdir, self.prefix + ".src"), "w"),
+            "const": open(pjoin(outdir, self.prefix + ".const.jsonl"), "w"),
+            "err": open(pjoin(outdir, self.prefix + ".error"), "w")
+        }
+        l = 0
+        tmpsrc = f"/tmp/myfunc_{self.impl}_{self.opt}_{self.prefix}"
+        if "fortran" in self.impl:
+            tmpsrc += ".f95"
+            func = "myfunc_"
+        else:
+            tmpsrc += ".c"
+            func = "myfunc"
+        tmpelf = f"/tmp/myfunc_{self.arch}_{self.impl}_{self.opt}_{self.prefix}.elf"
+
+        while True:
+            data = self.q.get()
+            if data == QUEUE_END:
+                # Queue is closed, break from inf loop
+                break
+            n, expr, expr_const, pref = data
+            impl = Implementor(expr, constants=expr_const, dtype=self.dtype)
+            try:
+                code = impl.implement(self.impl)
+                self.compiler(code, tmpelf, arch=self.arch, src=tmpsrc, opt=self.opt)
+                disasm = self.disassembler(tmpelf, expr_constants=expr_const,
+                                           match_constants=True)
+                asm = disasm.disassemble(func)
+                if len(disasm.constants) < len(expr_const):
+                    print(n, "constants not identified", disasm.constants, expr_const,
+                          file=outfiles["err"])
+                    continue
+            except DecodeError as e:
+                print(n, "impl error", e, expr, expr_const, pref, file=outfiles["err"])
+                continue
+
+            outfiles["asm"].write(asm + "\n")
+            outfiles["eqn"].write(pref + "\n")
+            outfiles["src"].write(f"==== pick={n} line={l} ====\n" + code + "\n")
+            outfiles["const"].write(json.dumps(expr_const) + "\n")
+            l += 1
+
+        for f in outfiles:
+            outfiles[f].close()
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Compile prefix to asm->eqn dataset")
+    parser.add_argument("-f", "--file", required=True, help="Input file")
+    parser.add_argument("--outdir", required=True, help="Output directory")
+    parser.add_argument("--prefix", required=True, help="File prefix")
+    parser.add_argument("--impl", nargs="+", required=True,
+                        choices=["dag_c", "cse_c", "dag_fortran", "cse_fortran"])
+    parser.add_argument("--pick", type=float, required=True,
+                        help="Ratio of samples to pick (0 to 1)")
+    parser.add_argument("--start", type=int, default=0, help="Start from index")
+    parser.add_argument("--count", type=int, default=0, help="Process only these many")
+    parser.add_argument("--seed", type=int, default=1225)
+    parser.add_argument("--min-tokens", help="Minimum tokens in equations", type=int, default=5)
+    parser.add_argument("--min-ops", help="Minimum ops in equations", type=int, default=5)
+    parser.add_argument("--dtype", help="Implementation datatype", type=str,
+                        choices=["double", "float"], default="double")
+    parser.add_argument("--arch", help="Target architecture", type=str,
+                        choices=["arm32", "aarch64", "x64"], default="arm32")
+    parser.add_argument("-O", "--opt", nargs="+", type=int, choices=[0, 1, 2, 3], default=[0],
+                        help="Optimization level (s)")
+    
+    # Dont show warnings
+    logging.getLogger("cle").setLevel(logging.ERROR)
+
+    args = parser.parse_args()
+    random.seed(args.seed)
+
+    eqcompilers = [EquationCompiler(mp.Queue(), args.arch, impl, opt, args.outdir, args.prefix, dtype=args.dtype)
+                        for impl in args.impl
+                        for opt in args.opt]
+    pool = [mp.Process(target=eqc.run, args=()) for eqc in eqcompilers]
+    for proc in pool:
+        proc.start()
+
+    count = 0
+    prefixf = open(args.file, "r")
+    for n, line in tqdm(enumerate(prefixf), desc="Parsing file"):
+        # Skip for start lines and with some probability
+        if n < args.start or random.random() > args.pick:
+            continue
+        comps = line.strip().split("\t")
+        pref = comps[0][comps[0].find("Y'")+3:]
+        prefl = pref.split(" ")
+        # pref = comps[1].split(" ")
+        if len(prefl) < args.min_tokens:
+            continue
+        try:
+            expr = parse_prefix_to_sympy(prefl)
+            with timeout(10):
+                expr = sp.simplify(expr)
+            if not sympy_expr_ok(expr):
+                # Simplified is bad
+                continue
+            expr, expr_const = constant_fold(expr)
+            pref = " ".join(sympy_to_prefix(expr))
+        except:
+            continue
+
+        if sp.count_ops(expr) < args.min_ops:
+            continue
+
+        for eqc in eqcompilers:
+            # Poll on this queue to get empty
+            while eqc.q.qsize() > 5:
+                sleep(1)
+            eqc.q.put((n, expr, expr_const, pref))
+        count += 1
+        if args.count > 0 and count >= args.count:
+            break
+
+    # Close queues
+    for eqc in eqcompilers:
+        eqc.q.put(QUEUE_END)
+    for proc in pool:
+        proc.join()

remend/compile_eqn.sh

@@ -0,0 +1,44 @@
+#!bin/bash
+
+MODE=$1
+SRC=$2
+ELF=$3
+OPT=$4
+
+if [ ! -f "$SRC" ]
+then
+    echo "Please provide source file"
+    exit 1
+fi
+
+if [ "$ELF" == "" ]
+then
+    echo "Please provide elf file path"
+    exit 1
+fi
+
+
+if [ "$MODE" == "arm32-c" ]
+then
+    arm-linux-gnueabihf-gcc $OPT $SRC -lm -o $ELF
+elif [ "$MODE" == "arm32-fortran" ]
+then
+    arm-linux-gnueabihf-gfortran -std=gnu $OPT $SRC -o $ELF
+elif [ "$MODE" == "aarch64-c" ]
+then
+    aarch64-linux-gnu-gcc $OPT $SRC -lm -o $ELF
+elif [ "$MODE" == "aarch64-fortran" ]
+then
+    aarch64-linux-gnu-gfortran -std=gnu $OPT $SRC -o $ELF
+elif [ "$MODE" == "x64-c" ]
+then
+    gcc $OPT $SRC -lm -o $ELF
+elif [ "$MODE" == "x64-fortran" ]
+then
+    gfortran -std=gnu $OPT $SRC -o $ELF
+else
+    echo "Incorrect mode: $MODE. Choose from: {arm32,aarch64,x64}-{c,fortran}"
+    exit 1
+fi
+
+# arm-linux-gnueabihf-objdump --no-show-raw-insn --no-addresses -d $1.elf | sed -n -e 's/\s;\s.*$//' -e "/myfunc>:$/,/^$/p" | sed '1d;$d' | tr '\n' ' '

remend/convert_generated.py

@@ -0,0 +1,24 @@
+from .parser import parse_prefix_to_sympy
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Parse result prefix to equation")
+    parser.add_argument("--input", required=True, help="Input result file")
+    args = parser.parse_args()
+    
+    res_list = []
+
+    with open(args.input, 'r') as f:
+        for line in f:
+            comps = line.strip().split("\t")
+            if line[0] == 'H':
+                num = int(comps[0][2:])
+                tokens = comps[2].split(" ")
+                res_list.append((num, tokens))
+
+    for n, toks in res_list:
+        try:
+            ex = parse_prefix_to_sympy(toks)
+            print(n, ex)
+        except Exception as e:
+            print(n, "could not parse:", str(e))

remend/deduplicate_split.py

@@ -0,0 +1,111 @@
+import sys
+import random
+import os
+import re
+from tqdm import tqdm
+
+def filter_poly(asm, eqn):
+    rejects = {"ln", "exp", "sin", "cos", "sqrt", "tan", "asin", "acos", "atan", "E", "pi", "cot"}
+    return any(t in rejects for t in asm.strip().split(" ")) \
+            or any(t in rejects for t in eqn.strip().split(" "))
+
+def filter_bigint(asm, eqn):
+    if re.search(r"CONST=[0-9]{4,}", asm):
+        return True
+    return False
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Deduplicate ASM and split files into train/test/valid")
+    parser.add_argument("--inprefix", required=True, help="Prefix of input files")
+    parser.add_argument("--outdir", required=True)
+    parser.add_argument("--split", type=float, default=0.05)
+    parser.add_argument("--seed", type=int, default=1225)
+    parser.add_argument("--filter", choices=["poly", "bigint"], default=None)
+    parser.add_argument("--no-separate-eqn", action="store_true")
+    
+    args = parser.parse_args()
+    
+    eq_mapped = {}
+    combined_ds = []
+    asm_hash = set()
+    removed = 0
+
+    with open(args.inprefix + ".asm", "r") as asmf, \
+            open(args.inprefix + ".eqn", "r") as eqnf, \
+            open(args.inprefix + ".const.jsonl", "r") as constf:
+        for i, (asm, eqn, const) in tqdm(enumerate(zip(asmf, eqnf, constf)),
+                                         desc="Read files", leave=False):
+            h = hash(asm)
+            if h in asm_hash:
+                # Skip this repeated line
+                removed += 1
+                continue
+
+            if re.search(r"[0-9]\.[0-9]", eqn):
+                # Float not represented, remove
+                removed += 1
+                continue
+
+            if args.filter == "poly" and filter_poly(asm, eqn):
+                removed += 1
+                continue
+            if args.filter == "bigint" and filter_bigint(asm, eqn):
+                removed += 1
+                continue
+
+            asm_hash.add(h)
+            if args.no_separate_eqn:
+                combined_ds.append((i, asm, eqn, const))
+            else:
+                if eqn not in eq_mapped:
+                    eq_mapped[eqn] = []
+                eq_mapped[eqn].append((i, asm, const))
+
+    print("Removed", removed)
+
+    if args.no_separate_eqn:
+        dataset = combined_ds
+    else:
+        dataset = list(eq_mapped.keys())
+
+    random.seed(args.seed)
+    random.shuffle(dataset)
+
+    N = len(dataset)
+    Ntest = int(N * args.split)
+    
+    splits = {
+        "train": dataset[:N-2*Ntest],
+        "valid": dataset[N-2*Ntest:N-Ntest],
+        "test": dataset[N-Ntest:]
+    }
+    splitidxs = {s: [] for s in splits}
+
+    idxf = open(os.path.join(args.outdir, "splits.txt"), "w")
+    for s in splits:
+        asmfn = os.path.join(args.outdir, f"{s}.asm")
+        eqnfn = os.path.join(args.outdir, f"{s}.eqn")
+        constfn = os.path.join(args.outdir, f"{s}.const.jsonl")
+        with open(asmfn, "w") as asmf, open(eqnfn, "w") as eqnf, \
+                open(constfn, "w") as constf:
+            if args.no_separate_eqn:
+                for i, asm, eqn, const in splits[s]:
+                    asmf.write(asm)
+                    eqnf.write(eqn)
+                    constf.write(const)
+                    splitidxs[s].append(i)
+            else:
+                for eqn in splits[s]:
+                    for i, asm, const in eq_mapped[eqn]:
+                        asmf.write(asm)
+                        eqnf.write(eqn)
+                        constf.write(const)
+                        splitidxs[s].append(i)
+        print("Split", s, len(splitidxs[s]))
+        idxf.write(f"==== {s} ====\n")
+        for j, i in enumerate(splitidxs[s]):
+            idxf.write(f"{j}: {i}\n")
+        idxf.write("\n")
+    idxf.close()
+

remend/disassemble.py

@@ -0,0 +1,553 @@
+from capstone import *
+from capstone.arm import *
+from capstone.arm64 import *
+from capstone.x86 import *
+import cle
+import struct
+from math import e as CONST_E, pi as CONST_PI
+import sympy as sp
+
+from .util import DecodeError
+
+def int2fp32(v):
+    if type(v) == int:
+        v = struct.unpack("<f", v.to_bytes(4, "little"))
+        v = v[0]
+    return v
+def int2fp64(v):
+    if type(v) == int:
+        v = struct.unpack("<d", v.to_bytes(8, "little"))
+        v = v[0]
+    return v
+
+def align4(v):
+    return v & (0xFFFFFFFC)
+
+class DisassemblerBase:
+    def __init__(self, expr_constants={}, match_constants=False):
+        self.loader = None # Load in child class
+        self.reg_values = {}
+        self.constidx = 0
+        self.constants = {}
+        self.constaddrs = set()
+        self.expr_constants = expr_constants
+        self.match_constants = match_constants
+
+    def get_function_bytes(self, funcname):
+        func = self.loader.find_symbol(funcname)
+        if not func:
+            raise DecodeError(f"Function {funcname} not found in binary")
+        faddr = func.rebased_addr
+        if (not isinstance(self, DisassemblerX64)) and faddr % 2 == 1:
+            # Unaligned address, aligning
+            faddr = faddr - 1
+        fbytes = self.loader.memory.load(faddr, func.size)
+        self.funcrange = faddr, faddr + func.size
+        return faddr, fbytes
+
+    def find_constant(self, constants, value):
+        for ec in constants:
+            if abs(value - constants[ec]) < 1e-5:
+                return ec, ""
+            elif abs(1/value - constants[ec]) < 1e-5:
+                return ec, "1/"
+            elif abs(-value - constants[ec]) < 1e-5:
+                return ec, "-"
+            elif abs(-1/value - constants[ec]) < 1e-5:
+                return ec, "-1/"
+        return False
+
+    def add_constant(self, value, addr=0, size=0):
+        # Don't map known constants like e, pi, 0
+        if value == 0:
+            cname = "CONST=0"
+        elif abs(value - CONST_E) < 1e-7:
+            cname = "CONST=E"
+        elif abs(value - CONST_PI) < 1e-7:
+            cname = "CONST=pi"
+        elif self.match_constants and \
+                (ecmatch := self.find_constant(self.expr_constants, value)):
+            # Gives the name and expression of the matched constant
+            ecname, ecxpr = ecmatch
+            # print(value, ecname, ecxpr, self.expr_constants[ecname])
+            cname = f"{ecxpr}CSYM{ecname[1:]}"
+            self.constants[ecname] = value
+        elif size > 0 and addr in self.constaddrs and \
+                (smatch := self.find_constant(self.constants, value)):
+            sname, sxpr = smatch
+            cname = f"{sxpr}CSYM{sname}"
+        else:
+            rep = sp.nsimplify(value, [sp.E, sp.pi], tolerance=1e-7)
+            if isinstance(rep, sp.Integer) or \
+                    (isinstance(rep, sp.Rational) and rep.q <= 16):
+                cname = f"CONST={rep}"
+            elif not self.match_constants:
+                cname = f"CSYM{self.constidx}"
+                self.constants[self.constidx] = value
+                self.constidx += 1
+            else:
+                raise DecodeError(f"Cannot represent unmatched float {value}")
+
+        if size > 0:
+            self.constaddrs |= {addr+i for i in range(size)}
+        return cname
+
+    def disassemble(self, function):
+        raise NotImplementedError("Call disassemble on child classes, not base")
+
+
+class DisassemblerARM32(DisassemblerBase):
+    def __init__(self, binpath, expr_constants={}, match_constants=False):
+        super().__init__(expr_constants=expr_constants, match_constants=match_constants)
+        self.md = Cs(CS_ARCH_ARM, CS_MODE_THUMB)
+        self.md.detail = True
+        self.loader = cle.Loader(binpath)
+
+    def check_mov_imm(self, insn):
+        if insn.id not in {ARM_INS_MOV, ARM_INS_MOVW,
+                           ARM_INS_MOVT, ARM_INS_ADR}:
+            return False
+        ops = list(insn.operands)
+        if len(ops) != 2:
+            return False
+        if ops[0].type != ARM_OP_REG or ops[1].type != ARM_OP_IMM:
+            return False
+        imm = ops[1].value.imm
+        if imm < 0:
+            imm = 2**32 + imm # 2's complement
+        if insn.id == ARM_INS_ADR:
+            # Add PC value
+            imm += insn.address + 4
+        return ops[0].value.reg, imm
+
+    def check_float_store(self, insn):
+        if insn.id not in {ARM_INS_STR, ARM_INS_STRD}:
+            return False
+        ops = list(insn.operands)
+        if insn.id == ARM_INS_STRD:
+            dest = ops[0].value.reg
+            dest2 = ops[1].value.reg
+            if dest not in self.reg_values or dest2 not in self.reg_values:
+                return False
+            fval = int2fp64((self.reg_values[dest2]<<32) + self.reg_values[dest])
+        else:
+            dest = ops[0].value.reg
+            if dest not in self.reg_values:
+                return False
+            fval = int2fp32(self.reg_values[dest])
+        if abs(fval) < 1e-3 or abs(fval) > 100:
+            return False
+        return fval
+
+    def check_ldrd(self, insn):
+        if insn.id != ARM_INS_LDRD:
+            return False
+        ops = insn.op_str.split(", ")
+        if len(ops) != 3:
+            return False
+        mem = ops[2] # format: [<reg> + #<offset>]
+        if mem[0] != "[" or mem[-1] != "]":
+            return False
+        memcomps = mem[1:-1].split(" ")
+        if memcomps[0] == "pc":
+            base = align4(insn.address + 4)
+        else:
+            basereg = ARM_REG_R0 + int(memcomps[0][1:]) # Shitty hack, may malfunction
+            if basereg not in self.reg_values:
+                return False
+            base = align4(self.reg_values[basereg])
+        if len(memcomps) == 3:
+            offset = int(memcomps[2][1:])
+        else:
+            offset = 0
+        addr = base + offset
+        fhex = self.loader.memory.load(addr, 8)
+        fval = struct.unpack("d", fhex)[0]
+        return fval, addr, 8
+
+    def check_vldr(self, insn):
+        if insn.id != ARM_INS_VLDR:
+            return False
+        ops = list(insn.operands)
+        dest = ops[0]
+        if ops[1].type != ARM_OP_MEM:
+            return False
+        mem = ops[1].value.mem
+        if mem.base == ARM_REG_PC:
+            # Align4(PC) + Imm
+            # For whatever reason, in Thumb PC=addr+4
+            addr = align4(insn.address + 4) + mem.disp
+        elif mem.base in self.reg_values:
+            addr = align4(self.reg_values[mem.base]) + mem.disp
+        else:
+            return False
+        if addr < self.loader.min_addr or addr + 8 > self.loader.max_addr:
+            # Out of bounds
+            return False
+        if dest.value.reg >= ARM_REG_D0 and dest.value.reg <= ARM_REG_D31:
+            size = 8
+            fhex = self.loader.memory.load(addr, 8)
+            fval = struct.unpack("d", fhex)[0]
+        else:
+            size = 4
+            fhex = self.loader.memory.load(addr, 4)
+            fval = struct.unpack("f", fhex)[0]
+        return fval, addr, size
+
+    def check_vmov(self, insn):
+        # fconsts/d == vmov.f32/f64 (old/new names)
+        if insn.id not in {ARM_INS_FCONSTS, ARM_INS_FCONSTD}:
+            return False
+        ops = list(insn.operands)
+        if len(ops) != 2 or ops[1].type != ARM_OP_FP:
+            return False
+        fval = ops[1].value.fp
+        destname = insn.reg_name(ops[0].value.reg)
+        asm = f"{insn.mnemonic} {destname}, {fval}"
+        return asm, fval
+        
+    def check_branch_symbol(self, insn):
+        if insn.id not in {ARM_INS_B, ARM_INS_BL, ARM_INS_BLX}:
+            return False
+        ops = list(insn.operands)
+        if len(ops) != 1 or ops[0].type != ARM_OP_IMM:
+            return False
+        addr = ops[0].value.imm
+        if addr > self.funcrange[0] and addr < self.funcrange[1]:
+            # Self-branch
+            func = f"SELF+{hex(addr - self.funcrange[0])}"
+        else:
+            func = self.loader.find_plt_stub_name(addr)
+            if func is None:
+                # Some tail call optimized PLT stubs have extra instructions
+                # that are not identified by CLE, so check with offset of 4 also.
+                func = self.loader.find_plt_stub_name(addr + 4)
+            if func is None:
+                return False
+        asm = f"{insn.mnemonic} <{func}>"
+        return asm
+
+    def get_function_bytes(self, funcname):
+        func = self.loader.find_symbol(funcname)
+        if not func:
+            raise DecodeError(f"Function {funcname} not found in binary")
+        faddr = func.rebased_addr
+        if faddr % 2 == 1:
+            # Unaligned address, aligning
+            faddr = faddr - 1
+        fbytes = self.loader.memory.load(faddr, func.size)
+        self.funcrange = faddr, faddr + func.size
+        return faddr, fbytes
+
+    def disassemble(self, funcname):
+        funcaddr, funcbytes = self.get_function_bytes(funcname) 
+        disassm = []
+
+        for insn in self.md.disasm(funcbytes, funcaddr):
+            if insn.address in self.constaddrs:
+                # Skip if this is a constant value and not instruction
+                continue
+
+            cname = None
+            asm = None
+
+            if vldr := self.check_vldr(insn):
+                fval, faddr, fsize = vldr
+                cname = self.add_constant(fval, faddr, fsize)
+            elif ldrd := self.check_ldrd(insn):
+                fval, faddr, fsize = ldrd
+                cname = self.add_constant(fval, faddr, fsize)
+            elif strfloat := self.check_float_store(insn):
+                fval = strfloat
+                cname = self.add_constant(fval)
+            elif vmovfloat := self.check_vmov(insn):
+                asm, fval = vmovfloat
+                cname = self.add_constant(fval)
+            elif branch := self.check_branch_symbol(insn):
+                asm = branch
+
+            # Maintain values of immediate moves.
+            # Needs to be done after processing current instruction.
+            if movimm := self.check_mov_imm(insn):
+                reg, imm = movimm
+                if insn.id == ARM_INS_MOVT:
+                    if reg not in self.reg_values:
+                        self.reg_values[reg] = 0
+                    self.reg_values[reg] += imm << 16
+                else:
+                    self.reg_values[reg] = imm
+            else:
+                reads, writes = insn.regs_access()
+                for r in writes:
+                    # Remove this reg if written to
+                    if r in self.reg_values:
+                        del self.reg_values[r]
+
+            if not asm:
+                asm = f"{insn.mnemonic} {insn.op_str}"
+            if cname:
+                asm += f", {cname}"
+            disassm.append(asm)
+
+        fulldiss = "; ".join(disassm)
+        return fulldiss
+
+class DisassemblerAArch64(DisassemblerBase):
+    def __init__(self, binpath, expr_constants={}, match_constants=False):
+        super().__init__(expr_constants=expr_constants, match_constants=match_constants)
+        self.md = Cs(CS_ARCH_ARM64, CS_MODE_ARM)
+        self.md.detail = True
+        self.loader = cle.Loader(binpath)
+
+    def reg_size_type(self, reg):
+        # Bit width and datatype of register
+        if reg >= ARM64_REG_W0 and reg <= ARM64_REG_W30:
+            return 32, int
+        elif reg >= ARM64_REG_X0 and reg <= ARM64_REG_X30:
+            return 64, int
+        elif reg >= ARM64_REG_S0 and reg <= ARM64_REG_S31:
+            return 32, float
+        elif reg >= ARM64_REG_D0 and reg <= ARM64_REG_D31:
+            return 64, float
+        return 0, None
+
+    def check_mov_imm(self, insn):
+        if insn.id not in {ARM64_INS_ADRP, ARM64_INS_ADR, ARM64_INS_MOV, ARM64_INS_MOVK}:
+            return False
+
+        ops = insn.operands
+        if len(ops) != 2:
+            return False
+        if ops[0].type != ARM64_OP_REG or ops[1].type != ARM64_OP_IMM:
+            return False
+
+        imm = ops[1].value.imm
+        if ops[1].shift.type == 1: # LSL
+            imm <<= ops[1].shift.value
+            mask = 0xFFFF << ops[1].shift.value
+
+        if insn.id == ARM64_INS_ADRP:
+            # imm -= 0x400000 # Subtract global offset for some reason
+            # imm = ((insn.address + 4) & (~4095)) + imm
+            # Really confused about this, maybe I can use the imm directly
+            pass
+        elif insn.id == ARM64_INS_ADR:
+            imm -= 0x400000 # Subtract global offset for some reason
+            imm += insn.address + 4
+        elif insn.id == ARM64_INS_MOVK:
+            # load previous reg value
+            if ops[0].value.reg in self.reg_values:
+                curr = self.reg_values[ops[0].value.reg]
+                imm = (imm & mask) | (curr & (~mask))
+            
+        return ops[0].value.reg, imm
+
+    def check_fmov(self, insn):
+        if insn.id != ARM64_INS_FMOV:
+            return False
+        ops = insn.operands
+        if len(ops) != 2: # or ops[1].type != ARM64_OP_FP:
+            return False
+
+        destsize, _ = self.reg_size_type(ops[0].value.reg)
+        destname = insn.reg_name(ops[0].value.reg)
+        if ops[1].type == ARM64_OP_FP:
+            fval = ops[1].value.fp
+            asm = f"{insn.mnemonic} {destname}, {fval}"
+        elif ops[1].type == ARM64_OP_REG:
+            reg = ops[1].value.reg
+            if reg not in self.reg_values:
+                return False
+            # TODO datatype
+            fhex = self.reg_values[reg]
+            if destsize == 64:
+                if fhex < 0:
+                    fhex += 2**64
+                fval = int2fp64(fhex)
+            elif destsize == 32:
+                if fhex < 0:
+                    fhex += 2**32
+                fval = int2fp32(fhex)
+            else:
+                return False
+
+            if abs(fval) < 1e-5 or abs(fval) > 1e5:
+                return False
+            asm = f"{insn.mnemonic} {insn.op_str}"
+        return asm, fval
+
+    def check_ldr(self, insn):
+        if insn.id != ARM64_INS_LDR:
+            return False
+        ops = insn.op_str[:-1].split(", ")
+        destsize, desttype = self.reg_size_type(insn.operands[0].value.reg)
+        if len(ops) < 2 or desttype != float:
+            return False
+        reg = ops[1]
+        if reg[0] != "[" or "sp" in reg:
+            return False
+        basereg = ARM64_REG_X0 + int(reg[2:]) # Shitty hack, may malfunction
+        if basereg not in self.reg_values:
+            return False
+        base = align4(self.reg_values[basereg])
+        if len(ops) == 3:
+            offset = ops[2][1:]
+            if offset.startswith("0x"):
+                offset = int(offset[2:], base=16)
+            else:
+                offset = int(offset)
+        else:
+            offset = 0
+        addr = base + offset
+        if destsize == 64:
+            fhex = self.loader.memory.load(addr, 8)
+            fval = struct.unpack("d", fhex)[0]
+            return fval, addr, 8
+        elif destsize == 32:
+            fhex = self.loader.memory.load(addr, 4)
+            fval = struct.unpack("f", fhex)[0]
+            return fval, addr, 4
+        else:
+            return False
+
+
+    def check_branch_symbol(self, insn):
+        if insn.id not in {ARM64_INS_BL, ARM64_INS_B}:
+            return False
+        ops = insn.operands
+        if len(ops) != 1 or ops[0].type != ARM_OP_IMM:
+            return False
+        addr = ops[0].value.imm
+        if addr > self.funcrange[0] and addr < self.funcrange[1]:
+            # Self-branch
+            func = f"SELF+{hex(addr - self.funcrange[0])}"
+        else:
+            func = self.loader.find_plt_stub_name(addr)
+            if func is None:
+                # Some tail call optimized PLT stubs have extra instructions
+                # that are not identified by CLE, so check with offset of 4 also.
+                func = self.loader.find_plt_stub_name(addr + 4)
+            if func is None:
+                return False
+        asm = f"{insn.mnemonic} <{func}>"
+        return asm
+
+    def disassemble(self, funcname):
+        funcaddr, funcbytes = self.get_function_bytes(funcname) 
+        disassm = []
+
+        for insn in self.md.disasm(funcbytes, funcaddr):
+            if insn.address in self.constaddrs:
+                # Skip if this is a constant value and not instruction
+                continue
+
+            cname = None
+            asm = None
+            # Maintain values of immediate moves
+            if movimm := self.check_mov_imm(insn):
+                reg, imm = movimm
+                self.reg_values[reg] = imm
+            else:
+                reads, writes = insn.regs_access()
+                for r in writes:
+                    # Remove this reg if written to
+                    if r in self.reg_values:
+                        del self.reg_values[r]
+
+            if fmov := self.check_fmov(insn):
+                asm, fval = fmov
+                cname = self.add_constant(fval)
+            elif ldr := self.check_ldr(insn):
+                fval, faddr, fsize = ldr
+                cname = self.add_constant(fval, faddr, fsize)
+            elif branch := self.check_branch_symbol(insn):
+                asm = branch
+
+            if not asm:
+                asm = f"{insn.mnemonic} {insn.op_str}"
+            if cname:
+                asm += f", {cname}"
+            disassm.append(asm)
+        
+        fulldiss = "; ".join(disassm)
+        return fulldiss
+
+class DisassemblerX64(DisassemblerBase):
+    def __init__(self, binpath, expr_constants={}, match_constants=False):
+        super().__init__(expr_constants=expr_constants, match_constants=match_constants)
+        self.md = Cs(CS_ARCH_X86, CS_MODE_64)
+        self.md.detail = True
+        self.loader = cle.Loader(binpath)
+
+    def check_call_symbol(self, insn):
+        if insn.id != X86_INS_CALL:
+            return False
+        ops = insn.operands
+        # TODO check this ARM_OP
+        if len(ops) != 1 or ops[0].type != ARM_OP_IMM:
+            return False
+        addr = ops[0].value.imm
+        func = self.loader.find_plt_stub_name(addr)
+        if func is None:
+            return False
+        asm = f"{insn.mnemonic} <{func}>"
+        return asm
+
+    def check_fload(self, insn):
+        # Cannot rely on ID because any instruction
+        # can access memory.
+        ops = insn.operands
+        memops = [op for op in ops 
+                  if (op.type == X86_OP_MEM and
+                      op.value.mem.base == X86_REG_RIP)]
+        if len(memops) != 1:
+            return False
+        mem, size = memops[0].value.mem, memops[0].size
+        if size > 8:
+            return False
+        addr = insn.address + insn.size + mem.disp
+        fhex = self.loader.memory.load(addr, size)
+        fval = struct.unpack("f" if size == 4 else "d", fhex)[0]
+        return fval, addr, size
+
+    def disassemble(self, funcname):
+        funcaddr, funcbytes = self.get_function_bytes(funcname) 
+        disassm = []
+
+        for insn in self.md.disasm(funcbytes, funcaddr):
+            asm = None
+            cname = None
+            if fload := self.check_fload(insn):
+                fval, faddr, fsize = fload
+                cname = self.add_constant(fval, faddr, fsize)
+            elif call := self.check_call_symbol(insn):
+                asm = call
+
+            if not asm:
+                asm = f"{insn.mnemonic} {insn.op_str}"
+            if cname:
+                asm += f", {cname}"
+            disassm.append(asm)
+        
+        fulldiss = "; ".join(disassm)
+        return fulldiss
+
+
+# Regular
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Pre-process assembly to replace constants and dump")
+    parser.add_argument("--bin", required=True)
+    parser.add_argument("--func", required=True)
+    parser.add_argument("--arch", required=True)
+    args = parser.parse_args()
+    
+    if args.arch == "arm32":
+        D = DisassemblerARM32(args.bin)
+    elif args.arch == "aarch64":
+        D = DisassemblerAArch64(args.bin)
+    elif args.arch == "x64":
+        D = DisassemblerX64(args.bin)
+    diss = D.disassemble(args.func)
+    print(diss)
+    print(D.constants)

remend/edit_model.py

@@ -0,0 +1,16 @@
+removed = ["encoder.layers.0.in_proj_weight", "encoder.layers.0.in_proj_bias", "encoder.layers.0.out_proj_weight", "encoder.layers.0.out_proj_bias", "encoder.layers.0.fc1_weight", "encoder.layers.0.fc1_bias", "encoder.layers.0.fc2_weight", "encoder.layers.0.fc2_bias", "encoder.layers.1.in_proj_weight", "encoder.layers.1.in_proj_bias", "encoder.layers.1.out_proj_weight", "encoder.layers.1.out_proj_bias", "encoder.layers.1.fc1_weight", "encoder.layers.1.fc1_bias", "encoder.layers.1.fc2_weight", "encoder.layers.1.fc2_bias", "encoder.layers.2.in_proj_weight", "encoder.layers.2.in_proj_bias", "encoder.layers.2.out_proj_weight", "encoder.layers.2.out_proj_bias", "encoder.layers.2.fc1_weight", "encoder.layers.2.fc1_bias", "encoder.layers.2.fc2_weight", "encoder.layers.2.fc2_bias", "encoder.layers.3.in_proj_weight", "encoder.layers.3.in_proj_bias", "encoder.layers.3.out_proj_weight", "encoder.layers.3.out_proj_bias", "encoder.layers.3.fc1_weight", "encoder.layers.3.fc1_bias", "encoder.layers.3.fc2_weight", "encoder.layers.3.fc2_bias", "encoder.layers.4.in_proj_weight", "encoder.layers.4.in_proj_bias", "encoder.layers.4.out_proj_weight", "encoder.layers.4.out_proj_bias", "encoder.layers.4.fc1_weight", "encoder.layers.4.fc1_bias", "encoder.layers.4.fc2_weight", "encoder.layers.4.fc2_bias", "encoder.layers.5.in_proj_weight", "encoder.layers.5.in_proj_bias", "encoder.layers.5.out_proj_weight", "encoder.layers.5.out_proj_bias", "encoder.layers.5.fc1_weight", "encoder.layers.5.fc1_bias", "encoder.layers.5.fc2_weight", "encoder.layers.5.fc2_bias"]
+
+if __name__ == "__main__":
+    import argparse
+    import torch
+
+    parser = argparse.ArgumentParser("Edit the checkpoint to remove extra dict weights")
+    parser.add_argument("-c", "--checkpoint", required=True, help="Input checkpoint")
+    parser.add_argument("-e", "--edited", required=True, help="Edited checkpoint")
+    args = parser.parse_args()
+    
+    sd = torch.load(args.checkpoint, weights_only=False)
+    for k in removed:
+        if k in sd['model']:
+            del sd['model'][k]
+    torch.save(sd, args.edited)

remend/eval_generated.py

@@ -0,0 +1,100 @@
+import sympy as sp
+import numpy as np
+import warnings
+from sympy.abc import x
+import sys
+import json
+from tqdm import tqdm
+
+from .parser import parse_prefix_to_sympy, isint
+
+# Ignore sympy lambda warnings.
+warnings.simplefilter("ignore")
+
+def percent(a, n):
+    return f"{a/n*100:0.1f}%"
+
+def do_eval_match(orig_expr, gen_expr):
+    try:
+        origl = sp.lambdify(x, orig_expr)
+        genl = sp.lambdify(x, gen_expr)
+        count = 0
+
+        for v in np.arange(0.2, 1, 0.01):
+            o = origl(v)
+            g = genl(v)
+            if o == float('nan') or o == float('inf'):
+                continue
+            if g == float('nan') or g == float('inf'):
+                continue
+            # if type(o) != np.float64 or type(g) != np.float64:
+            #     print(orig_expr, o, gen_expr, g)
+            #     return False
+            if abs((o-g)/o) > 1e-5:
+                return False
+            count += 1
+    except:
+        return False
+    return count >= 5
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Check generated expressions")
+    parser.add_argument("-g", required=True, help="Generated expressions file")
+    parser.add_argument("-c", required=True, help="Constants file")
+    parser.add_argument("-e", required=True, help="Equations file")
+    parser.add_argument("-r", required=True, help="Results file")
+    args = parser.parse_args()
+
+    gens = []
+    with open(args.g, 'r') as genf, open(args.c) as constf, open(args.e) as eqnf:
+        for line in tqdm(genf, desc="Reading file"):
+            comps = line.strip().split("\t")
+            if line[0] == 'H':
+                num = int(comps[0][2:])
+                tokens = comps[2].split(" ")
+                eqn = next(eqnf)
+                const = next(constf)
+                const = json.loads(const.strip())
+                gens.append((num, tokens, eqn.strip(), const))
+
+    parsed = []
+    matched = []
+    results = []
+
+    for n, toks, eqn, const in tqdm(gens, desc="Evaluating expressions"):
+        res = {"id": n, "parsed": False, "matched": False, "orig": "", "gen": ""}
+        if "<<unk>>" in toks:
+            # Not parsed
+            results.append(res)
+            continue
+        try:
+            gen_expr = parse_prefix_to_sympy(toks)
+        except Exception as e:
+            # Not parsed
+            results.append(res)
+            continue
+
+        res["parsed"] = True
+        parsed.append(n)
+
+        gen_expr = gen_expr.subs([(sp.Symbol("k"+c), const[c]) for c in const])
+        orig_expr = sp.parse_expr(eqn, local_dict={"x0":x})
+        res["orig"] = str(orig_expr)
+        res["gen"] = str(gen_expr)
+
+        if not do_eval_match(orig_expr, gen_expr):
+            results.append(res)
+            continue
+        res["matched"] = True
+        matched.append(n)
+        results.append(res)
+
+    with open(args.r, "w") as resf:
+        for res in results:
+            resf.write("{id} {parsed} {matched} \"{orig}\" \"{gen}\"\n".format(**res))
+        resf.write("\n")
+        N = len(gens)
+        print("Total", N, file=resf)
+        print("Parsed", len(parsed), percent(len(parsed), N), file=resf)
+        print("Matched", len(matched), percent(len(matched), N), file=resf)

remend/experiment.py

@@ -0,0 +1,75 @@
+import sympy as sp
+import random
+
+from .parser import parse_prefix_to_sympy
+
+isconst = lambda e: not any(c.is_symbol for c in e.atoms())
+def constfold(expr):
+    q = [expr]
+    cidx = 0
+    subsmap = {}
+    constmap = {}
+
+    while len(q) > 0:
+        curr_expr = q.pop(0)
+        if isinstance(curr_expr, sp.Number) or isconst(e):
+            const_expr = curr_expr.evalf()
+            rep = sp.nsimplify(const_expr, [sp.E, sp.pi])
+            if isinstance(rep, sp.Integer) or \
+                    (isinstance(rep, sp.Rational) and rep.q <= 16):
+                subsmap[curr_expr] = rep
+            else:
+                subsmap[curr_expr] = sp.Symbol(f"k{cidx}")
+                constmap[f"k{cidx}"] = float(const_expr)
+                cidx += 1
+        else:
+            for child in curr_expr.args:
+                q.append(child)
+
+    return expr.subs(subsmap), constmap
+
+def replace_const(expr):
+    cidx = 0
+    subsmap = {}
+    constmap = {}
+
+    for c in sp.preorder_traversal(expr):
+        if isinstance(c, sp.Float):
+            rep = sp.nsimplify(c)
+            if isinstance(rep, sp.Integer) or \
+                    (isinstance(rep, sp.Rational) and rep.q <= 16):
+                subsmap[c] = rep
+            else:
+                subsmap[c] = sp.Symbol(f"c{cidx}")
+                constmap[f"c{cidx}"] = float(c)
+                cidx += 1
+    return expr.subs(subsmap), constmap
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Random experiments")
+    parser.add_argument("-f", required=True)
+    parser.add_argument("-p", type=float, default=0.1)
+    parser.add_argument("-n", type=int, default=20)
+    args = parser.parse_args()
+
+    random.seed(1225)
+    
+    count = 0
+    with open(args.f, "r") as f:
+        for line in f:
+            if random.random() > args.p:
+                continue
+
+            prefl = line.strip().split(" ")
+
+            orig = parse_prefix_to_sympy(prefl)
+            # simp = sp.simplify(expr)
+            expr = constfold(orig)
+            expr, consts = replace_const(expr)
+            print(orig, expr, consts)
+
+            count += 1
+            if count == args.n:
+                break

remend/find_duplicates.py

@@ -0,0 +1,69 @@
+import sys
+from tqdm import tqdm
+from Levenshtein import distance
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Find duplicates in the dataset ASM")
+    parser.add_argument("--train", required=True)
+    # parser.add_argument("--valid", required=True)
+    parser.add_argument("--test", required=True)
+    parser.add_argument("--result", required=False)
+    parser.add_argument("--distance", action="store_true", default=False)
+    args = parser.parse_args()
+    
+    train = []
+    train_hash = {}
+    # valid = []
+    test = []
+    with open(args.train, "r") as tf:
+        for idx, line in tqdm(enumerate(tf), desc="Read train", leave=False):
+            train_hash[hash(line)] = idx
+            comps = line.strip().split(" ")
+            train.append(comps)
+    # with open(args.valid, "r") as tf:
+    #     for line in tqdm(tf, desc="Read valid", leave=False):
+    #         valid.append(line.strip().split(" "))
+    with open(args.test, "r") as tf:
+        for line in tqdm(tf, desc="Read test", leave=False):
+            test.append(line)
+
+    selfcheck = args.train == args.test
+    if args.result:
+        rf = open(args.result, "w")
+        searchdist = args.distance
+    else:
+        searchdist = False # Dont compute if no result file
+        rf = None
+
+    def reswrite(s):
+        if rf:
+            rf.write(s)
+
+    exact = 0
+    for i, testline in tqdm(enumerate(test), desc="Test", total=len(test)):
+        testl = testline.strip().split(" ")
+        htest = hash(testline)
+        if htest in train_hash:
+            # Found exact match
+            j = train_hash[htest]
+            if not selfcheck or j != i:
+                exact += 1
+                reswrite(f"{i} {j} 0 0.0\n")
+                continue
+
+        # If not, then search
+        if searchdist:
+            minavgdist, mindist, minj = 100, 100, -1 
+            for j, trainl in enumerate(train):
+                if abs(len(trainl) - len(testl)) > 10:
+                    dist = abs(len(trainl) - len(testl)) * 2 # HACK to speed it up
+                else:
+                    dist = distance(trainl, testl)
+                avgdist = dist / (len(trainl) + len(testl))
+                if mindist > dist:
+                    minavgdist, mindist, minj = avgdist, dist, j
+
+            reswrite(f"{i} {minj} {mindist} {minavgdist}\n")
+
+    print("Exact duplicates:", exact)

remend/implementation.py

@@ -0,0 +1,210 @@
+import sympy as sp
+from sympy.codegen import ast
+import itertools as it
+import networkx as nx
+
+from .parser import OPERATORS, sympy_to_dag
+from .util import DecodeError
+
+def isnum(s):
+    try:
+        float(s)
+        return True
+    except ValueError:
+        return False
+
+class Implementor:
+    def __init__(self, expr, constants={}, dtype="double"):
+        self.expr = expr
+        self.constants = constants
+        self.cdtype = dtype
+        self.cpf = "lf" if dtype == "double" else "f"
+        self.fdtype = "double precision" if dtype == "double" else "real"
+
+    def implement(self, impl):
+        if impl == "dag_c":
+            return self.dag_to_c_impl()
+        elif impl == "cse_c":
+            return self.sympy_cse_c_impl()
+        elif impl == "dag_fortran":
+            return self.dag_to_fortran_impl()
+        elif impl == "cse_fortran":
+            return self.sympy_cse_fortran_impl()
+
+
+    def op_c_impl(self, f, children):
+        if f == "add":
+            return " + ".join(children);
+        elif f == "mul":
+            return " * ".join(children);
+        elif f == "pow":
+            assert len(children) == 2
+            if self.cdtype == "double":
+                return f"pow({children[0]}, {children[1]})"
+            else:
+                return f"powf({children[0]}, {children[1]})"
+        elif f == "ln":
+            assert len(children) == 1
+            if self.cdtype == "double":
+                return f"log({children[0]})"
+            else:
+                return f"logf({children[0]})"
+        else:
+            if f in OPERATORS and OPERATORS[f][1] == 1:
+                assert len(children) == 1
+                if self.cdtype == "double":
+                    return f"{f}({children[0]})"
+                else:
+                    return f"{f}f({children[0]})"
+            else:
+                raise DecodeError(f"C impl: operation {f} not handled")
+
+    def op_f_impl(self, f, children):
+        if f == "add":
+            j = ")+(".join(children)
+            return "(" + j + ")"
+        elif f == "mul":
+            j = ")*(".join(children)
+            return "(" + j + ")"
+        elif f == "pow":
+            assert len(children) == 2
+            return f"({children[0]})**({children[1]})"
+        elif f == "ln":
+            assert len(children) == 1
+            return f"log({children[0]})"
+        else:
+            if f in OPERATORS and OPERATORS[f][1] == 1:
+                assert len(children) == 1
+                return f"{f}({children[0]})"
+            else:
+                raise DecodeError(f"F impl: operation {f} not handled")
+
+    def full_c_code(self, body):
+        pre = f"#include <stdio.h>\n#include <math.h>\n{self.cdtype} myfunc({self.cdtype} x) {{"
+        post = f"}}\nint main() {{ {self.cdtype} x; scanf(\"%{self.cpf}\", &x); printf(\"%{self.cpf}\", myfunc(x)); }}"
+        return f"{pre}\n{body}\n{post}"
+
+    def full_f_code(self, body):
+        pre = "function myfunc(x) result(y)\nimplicit none\n" + \
+              f"{self.fdtype}, intent(in) :: x\n{self.fdtype} :: y, E, pi\n"
+        post = "end function myfunc\nprogram main\nimplicit none\n" + \
+               f"{self.fdtype} :: x\n{self.fdtype} :: myfunc\n" + \
+               "read(*, *) x\nprint *, \"y is:\", myfunc(x)\nend program main"
+        return f"{pre}\n{body}\n{post}"
+
+    def dag_to_c_impl(self):
+        dag = sympy_to_dag(self.expr, csuf="F" if self.cdtype == "float" else "")
+        cstr = ""
+        added_pi, added_E = False, False
+        for c in self.constants:
+            cstr += f"{self.cdtype} {c} = {self.constants[c]};\n"
+        varidx = it.count()
+        for node in reversed(list(nx.topological_sort(dag))):
+            label = dag.nodes[node]["label"]
+            children = [dag.nodes[n]["var"] for n in dag.adj[node]]
+            if len(children) == 0:
+                if label == "pi":
+                    if self.cdtype == "float" and not added_pi:
+                        cstr += "const float pi = 3.14159265F;\n"
+                        added_pi = True
+                    else:
+                        label = "M_PI"
+                elif label == "E":
+                    if self.cdtype == "float" and not added_E:
+                        cstr += "const float E = 2.71828183F;\n"
+                        added_E = True
+                    else:
+                        label = "M_E"
+                dag.nodes[node]["var"] = label
+                continue
+            varname = f"t{next(varidx)}"
+            cexpr = self.op_c_impl(label, children)
+            dag.nodes[node]["var"] = varname
+            cstr += f"{self.cdtype} {varname} = {cexpr};\n"
+            retname = varname
+        cstr += f"return {retname};\n"
+        return self.full_c_code(cstr)
+
+    def dag_to_fortran_impl(self):
+        csuf = "" if self.fdtype == "real" else "d0"
+        dag = sympy_to_dag(self.expr, csuf=csuf)
+        varstr = ""
+        fstr = "parameter E = 2.71828183\nparameter pi = 3.14159265\n"
+        for c in self.constants:
+            varstr += f"{self.fdtype} :: {c}\n"
+            fstr += f"parameter {c} = {self.constants[c]}{csuf}\n"
+        varidx = it.count()
+        allvars = []
+        for node in reversed(list(nx.topological_sort(dag))):
+            label = dag.nodes[node]["label"]
+            children = [dag.nodes[n]["var"] for n in dag.adj[node]]
+            if len(children) == 0:
+                dag.nodes[node]["var"] = label
+                continue
+            varname = f"t{next(varidx)}"
+            fexpr = self.op_f_impl(label, children)
+            dag.nodes[node]["var"] = varname
+            fstr += f"{varname} = {fexpr}\n"
+            retname = varname
+            varstr += f"{self.fdtype} :: {varname}\n"
+        fstr += f"y = {retname};\n"
+        fstr = varstr + "\n" + fstr
+        return self.full_f_code(fstr)
+
+    def sympy_cse_c_impl(self):
+        if self.cdtype == "float":
+            extraargs = {
+                "type_aliases": {ast.real: ast.float32},
+                "math_macros": {},
+            }
+        else:
+            extraargs = {}
+        cstr = ""
+        for c in self.constants:
+            cstr += f"{self.cdtype} {c} = {self.constants[c]};\n"
+        xvars, xpr = sp.cse(self.expr)
+        for vname, vxpr in xvars:
+            code = sp.ccode(vxpr, assign_to=vname.name, **extraargs)
+            cstr += f"{self.cdtype} {vname.name}; {code};\n"
+        assert len(xpr) == 1
+        code = sp.ccode(xpr[0], assign_to="y", **extraargs)
+        cstr += f"{self.cdtype} y; {code}; return y;\n"
+        return self.full_c_code(cstr)
+
+    def sympy_cse_fortran_impl(self):
+        csuf = "" if self.fdtype == "real" else "d0"
+        varstr = ""
+        fstr = ""
+        for c in self.constants:
+            varstr += f"{self.fdtype} :: {c}\n"
+            fstr += f"parameter {c} = {self.constants[c]}{csuf}\n"
+        xvars, xpr = sp.cse(self.expr)
+        for vname, vxpr in xvars:
+            varstr += f"{self.fdtype} :: {vname.name}\n"
+            fstr += sp.fcode(vxpr, assign_to=vname.name, standard=95, source_format="free") + "\n"
+        assert len(xpr) == 1
+        fstr += sp.fcode(xpr[0], assign_to="y", standard=95, source_format="free") + "\n"
+        fstr = varstr + "\n" + fstr
+        if self.fdtype == "real":
+            # Hack to fix sympy generation
+            fstr = fstr.replace("d0", "")
+        return self.full_f_code(fstr)
+
+
+
+# For testing only
+if __name__ == "__main__":
+    from .parser import parse_prefix_to_sympy, sympy_to_dag
+
+    prefs = "add mul div INT+ 1 INT+ 5 x mul div INT+ 1 INT+ 5 mul x tan pow x INT+ 2".split(" ")
+    exp = parse_prefix_to_sympy(prefs)
+    impl = Implementor(exp, dtype="float")
+
+    print("DAG C:")
+    print(impl.dag_to_c_impl())
+    print("DAG Fortran:")
+    print(impl.dag_to_fortran_impl())
+    print("CSE C:")
+    print(impl.sympy_cse_c_impl())
+    print("CSE Fortran:")
+    print(impl.sympy_cse_fortran_impl())

remend/parser.py

@@ -0,0 +1,449 @@
+import sympy as sp
+import networkx as nx
+import itertools as it
+import sys
+
+from .util import DecodeError, sympy_expr_ok
+
+OPERATORS = {
+    # Elementary functions
+    'add': (lambda a,b: a+b, 2),
+    'sub': (lambda a,b: a-b, 2),
+    'mul': (lambda a,b: a*b, 2),
+    'div': (lambda a,b: a/b, 2),
+    'pow': (lambda a,b: a**b, 2),
+    # 'inv': (lambda a: 1/a, 1),
+    # 'pow2': (lambda a: a**2, 1),
+    # 'pow3': (lambda a: a**3, 1),
+    # 'pow4': (lambda a: a**4, 1),
+    # 'pow5': (lambda a: a**5, 1),
+    'sqrt': (lambda a: sp.sqrt(a), 1),
+    'exp': (lambda a: sp.exp(a), 1),
+    'ln': (lambda a: sp.ln(a), 1),
+    # 'abs': (lambda a: sp.abs(a), 1),
+    # 'sign': (lambda a: sp.sign(a), 1),
+    # Trigonometric Functions
+    'sin': (lambda a: sp.sin(a), 1),
+    'cos': (lambda a: sp.cos(a), 1),
+    'tan': (lambda a: sp.tan(a), 1),
+    'cot': (lambda a: sp.cot(a), 1),
+    'sec': (lambda a: sp.sec(a), 1),
+    'csc': (lambda a: sp.csc(a), 1),
+    # Trigonometric Inverses
+    'asin': (lambda a: sp.asin(a), 1),
+    'acos': (lambda a: sp.acos(a), 1),
+    'atan': (lambda a: sp.atan(a), 1),
+    'acot': (lambda a: sp.acot(a), 1),
+    'asec': (lambda a: sp.asec(a), 1),
+    'acsc': (lambda a: sp.acsc(a), 1),
+    # Hyperbolic
+    # 'sinh': (lambda a: sp.sinh(a), 1),
+    # 'cosh': (lambda a: sp.cosh(a), 1),
+    # 'tanh': (lambda a: sp.tanh(a), 1),
+}
+
+CONSTANTS = {
+    'E': sp.E,
+    'pi': sp.pi,
+    '0': 0,
+    '1': 1,
+    '2': 2,
+    '3': 3,
+    '4': 4,
+    '5': 5,
+    '6': 6,
+    '7': 7,
+    '8': 8,
+    '9': 9,
+}
+
+VARIABLES = {
+    'x': sp.Symbol('x'),
+    'x0': sp.Symbol('x0'),
+    'x1': sp.Symbol('x1'),
+
+    'c0': sp.Symbol('c0'),
+    'c1': sp.Symbol('c1'),
+    'c2': sp.Symbol('c2'),
+    'c3': sp.Symbol('c3'),
+    'c4': sp.Symbol('c4'),
+    'c5': sp.Symbol('c5'),
+    'c6': sp.Symbol('c6'),
+    'c7': sp.Symbol('c7'),
+    'c8': sp.Symbol('c8'),
+    'c9': sp.Symbol('c9'),
+    'c10': sp.Symbol('c10'),
+
+    'k0': sp.Symbol('k0'),
+    'k1': sp.Symbol('k1'),
+    'k2': sp.Symbol('k2'),
+    'k3': sp.Symbol('k3'),
+    # 'y': sp.Symbol('y'),
+    # 'z': sp.Symbol('z')
+}
+
+FUNC_TO_OP = {
+    sp.Add: 'add',
+    sp.Mul: 'mul',
+    sp.Pow: 'pow',
+
+    sp.log: 'ln',
+    sp.sqrt: 'sqrt',
+    sp.exp: 'exp',
+    sp.Abs: 'abs',
+    # 'abs': (lambda a: sp.abs(a), 1),
+    # 'sign': (lambda a: sp.sign(a), 1),
+    # Trigonometric Functions
+    sp.sin: 'sin',
+    sp.cos: 'cos',
+    sp.tan: 'tan',
+    sp.cot: 'cot',
+    sp.sec: 'sec',
+    sp.csc: 'csc',
+    # Trigonometric Inverses
+    sp.asin: 'asin',
+    sp.acos: 'acos',
+    sp.atan: 'atan',
+    sp.acot: 'acot',
+    sp.asec: 'asec',
+    sp.acsc: 'acsc',
+    # Hyperbolic
+    # sp.cosh: 'cosh',
+    # sp.sinh: 'sinh',
+    # sp.tanh: 'tanh'
+}
+
+def sympy_func_to_op(f):
+    if f in FUNC_TO_OP:
+        return FUNC_TO_OP[f]
+    else:
+        raise DecodeError(f"Op not found {f}")
+    return str(f)
+
+def isint(s):
+    try:
+        int(s)
+        return True
+    except ValueError:
+        return False
+
+def reverse_iter_prefix(prefs):
+    n = len(prefs) - 1
+    # currnum = 0
+    # currpow = 1
+    currnum = []
+    while n >= 0:
+        if isint(prefs[n]) or prefs[n] in ["e", "+", "-", "."]:
+            currnum += prefs[n]
+           # currnum += currpow * int(prefs[n])
+           # currpow *= 10
+        elif prefs[n][:3] == "INT":
+            parsedint = int("".join(reversed(currnum)))
+            if prefs[n][3] == "+":
+                yield parsedint
+            else:
+                yield -parsedint
+            currnum = []
+            # currpow = 1
+        elif prefs[n][:5] == "FLOAT":
+            parsedfloat = float("".join(reversed(currnum)))
+            if prefs[n][5] == "+":
+                yield parsedfloat
+            else:
+                yield -parsedfloat
+            currnum = []
+        else:
+            yield prefs[n]
+        n -= 1
+
+def parse_prefix_to_sympy(prefs):
+    stack = []
+    for val in reverse_iter_prefix(prefs):
+        # print(stack, val)
+        if val in OPERATORS:
+            spop, numops = OPERATORS[val]
+            operands = [stack.pop() for i in range(numops)]
+            expr = spop(*operands)
+            stack.append(expr)
+        elif val in CONSTANTS:
+            stack.append(CONSTANTS[val])
+        elif val in VARIABLES:
+            stack.append(VARIABLES[val])
+        elif type(val) == int or type(val) == float:
+            stack.append(val)
+        elif val == "(" or val == ")":
+            # Simply ignore brackets
+            continue
+        else:
+            raise DecodeError(f"{val} invalid")
+    
+    if len(stack) != 1:
+        raise DecodeError(f"Stack not empty, invalid expression: {prefs} || {stack}")
+    expr = stack.pop()
+    if not sympy_expr_ok(expr):
+        raise DecodeError("Complex or infinite expression")
+    return expr
+
+def parse_postfix_to_sympy(prefs):
+    stack = []
+    postfix = reversed(list(reverse_iter_prefix(prefs)))
+    for val in postfix:
+        if val in OPERATORS:
+            spop, numops = OPERATORS[val]
+            operands = [stack.pop() for i in range(numops)]
+            expr = spop(*operands)
+            stack.append(expr)
+        elif val in CONSTANTS:
+            stack.append(CONSTANTS[val])
+        elif val in VARIABLES:
+            stack.append(VARIABLES[val])
+        elif type(val) == int or type(val) == float:
+            stack.append(val)
+        elif val == "(" or val == ")":
+            # Simply ignore brackets
+            continue
+        else:
+            raise DecodeError(f"{val} invalid")
+    
+    if len(stack) != 1:
+        raise DecodeError(f"Stack not empty, invalid expression: {prefs} || {stack}")
+    expr = stack.pop()
+    if not sympy_expr_ok(expr):
+        raise DecodeError("Complex or infinite expression")
+    return expr
+
+
+def parse_prefix_to_tree(prefs):
+    tree = nx.DiGraph()
+    stack = []
+    newidx = len(prefs)
+    for nidx, val in enumerate(reverse_iter_prefix(prefs)):
+        tree.add_node(nidx, label=val)
+        if val in OPERATORS:
+            _, numops = OPERATORS[val]
+            childs = [stack.pop() for i in range(numops)]
+            if val in {"pow", "sub", "div"}:
+                # Ordered children
+                tree.add_node(newidx, label="lhs")
+                tree.add_node(newidx+1, label="rhs")
+                tree.add_edge(nidx, newidx)
+                tree.add_edge(nidx, newidx+1)
+                tree.add_edge(newidx, childs[0])
+                tree.add_edge(newidx+1, childs[1])
+                newidx += 2
+            else:
+                for c in childs:
+                    tree.add_edge(nidx, c)
+        elif val in CONSTANTS or val in VARIABLES or type(val) == int:
+            pass
+        else:
+            raise DecodeError(f"Val {val} invalid")
+        stack.append(nidx)
+
+    if len(stack) != 1:
+        raise DecodeError(f"Stack not empty, invalid expression: {prefs} || {stack}")
+
+    return tree, stack.pop() # Root node
+
+def sympy_to_dag(expression, csuf=""):
+    dag = nx.DiGraph()
+    seen = {}
+    nitr = it.count()
+
+    def _dfs(node):
+        children = []
+        for child in node.args:
+            if child in seen:
+                cid = seen[child]
+            else:
+                cid = _dfs(child)
+            children.append(cid)
+
+        nid = next(nitr)
+        dag.add_node(nid, expr=node)
+        seen[node] = nid
+        for cid in children:
+            dag.add_edge(nid, cid)
+        return nid
+
+    _dfs(expression)
+    for node in dag.nodes:
+        if len(dag.adj[node]) == 0:
+            e = dag.nodes[node]["expr"]
+            if isinstance(e, sp.Integer):
+                dag.nodes[node]["label"] = f"{e}.0{csuf}"
+            elif isinstance(e, sp.Rational):
+                dag.nodes[node]["label"] = f"{e.p}.0{csuf}/{e.q}.0{csuf}"
+            elif isinstance(e, sp.Float):
+                dag.nodes[node]["label"] = f"{float(e)}{csuf}"
+            else:
+                dag.nodes[node]["label"] = str(e)
+        else:
+            dag.nodes[node]["label"] = sympy_func_to_op(dag.nodes[node]["expr"].func)
+
+    return dag
+
+def sympy_to_prefix(expr):
+    trav = []
+
+    def _pre(node):
+        nonlocal trav
+        if isinstance(node, sp.Rational):
+            if node.q != 1:
+                trav.append("div") 
+                _pre(node.p)
+                _pre(node.q)
+            else:
+                _pre(node.p)
+        elif isinstance(node, sp.Integer) or isinstance(node, int):
+            v = int(node)
+            if v >= 0:
+                trav.append("INT+")
+                trav.extend(list(str(v)))
+            else:
+                trav.append("INT-")
+                trav.extend(list(str(-v)))
+        elif isinstance(node, sp.Symbol):
+            trav.append(str(node))
+        elif isinstance(node, sp.Mul):
+            mulargs = []
+            divargs = []
+            children = node.args
+            for child in children:
+                if isinstance(child, sp.Pow) and \
+                        isinstance(child.args[1], sp.Integer) and child.args[1] == -1:
+                    divargs.append(child.args[0])
+                else:
+                    mulargs.append(child)
+            if len(divargs) > 0:
+                trav.append("div")
+            if len(mulargs) == 0:
+                trav.append("INT+")
+                trav.append("1")
+            # Insert numerator
+            for i, child in enumerate(mulargs):
+                if i < len(mulargs) - 1:
+                    trav.append("mul")
+                _pre(child)
+            # Insert denominator
+            for i, child in enumerate(divargs):
+                if i < len(divargs) - 1:
+                    trav.append("mul")
+                _pre(child)
+        elif isinstance(node, sp.Add):
+            addargs = []
+            subargs = []
+            children = node.args
+            for child in children:
+                if isinstance(child, sp.Mul) and len(child.args) == 2 and \
+                        isinstance(child.args[1], sp.Integer) and child.args[1] == -1:
+                    subargs.append(child.args[0])
+                elif isinstance(child, sp.Mul) and len(child.args) == 2 and \
+                        isinstance(child.args[0], sp.Integer) and child.args[0] == -1:
+                    subargs.append(child.args[1])
+                else:
+                    addargs.append(child)
+            if len(subargs) > 0:
+                trav.append("sub")
+            if len(addargs) == 0:
+                trav.append("INT+")
+                trav.append("0")
+            # Insert numerator
+            for i, child in enumerate(addargs):
+                if i < len(addargs) - 1:
+                    trav.append("add")
+                _pre(child)
+            # Insert denominator
+            for i, child in enumerate(subargs):
+                if i < len(subargs) - 1:
+                    trav.append("add")
+                _pre(child)
+        elif isinstance(node, sp.Float):
+            rep = sp.nsimplify(node, tolerance=1e-7)
+            if isinstance(rep, sp.Integer):
+                _pre(rep)
+            elif isinstance(rep, sp.Rational) and rep.q <= 16:
+                _pre(rep)
+            else:
+                raise DecodeError(f"Float {node} encountered while generating")
+                # trav.append(str(node))
+        elif node == sp.E or node == sp.pi:
+            # Transcendental constants
+            trav.append(str(node))
+        else:
+            op = sympy_func_to_op(node.func)
+            children = node.args
+            for i, child in enumerate(children):
+                # Insert op repeatedly to maintain binary tree
+                if i == 0 or i < len(children) - 1:
+                    trav.append(op)
+                _pre(child)
+    _pre(expr)
+    return trav
+
+def constant_fold(expr):
+    q = [expr]
+    cidx = 0
+    subsmap = {}
+    constmap = {}
+
+    isconst = lambda e: not any(c.is_symbol for c in e.atoms())
+
+    while len(q) > 0:
+        curr_expr = q.pop(0)
+        if isinstance(curr_expr, sp.Number) or isconst(curr_expr):
+            const_expr = curr_expr.evalf()
+            rep = sp.nsimplify(const_expr, [sp.E, sp.pi], tolerance=1e-7)
+            if isinstance(rep, sp.Integer) or \
+                    (isinstance(rep, sp.Rational) and rep.q <= 16) or \
+                    rep == sp.E or rep == sp.pi:
+                subsmap[curr_expr] = rep
+            else:
+                val = float(const_expr)
+                found = False
+                for c in constmap:
+                    if abs(val - constmap[c]) < 1e-7:
+                        subsmap[curr_expr] = sp.Symbol(c)
+                        found = True
+                    elif abs(1/val - constmap[c]) < 1e-7:
+                        subsmap[curr_expr] = 1/sp.Symbol(c)
+                        found = True
+                    elif abs(-val - constmap[c]) < 1e-7:
+                        subsmap[curr_expr] = -sp.Symbol(c)
+                        found = True
+                    elif abs(-1/val - constmap[c]) < 1e-7:
+                        subsmap[curr_expr] = -1/sp.Symbol(c)
+                        found = True
+                if not found:
+                    subsmap[curr_expr] = sp.Symbol(f"k{cidx}")
+                    constmap[f"k{cidx}"] = val
+                    cidx += 1
+        else:
+            for child in curr_expr.args:
+                q.append(child)
+
+    return expr.subs(subsmap), constmap
+
+
+# For testing only
+if __name__ == "__main__":
+    prefs = "add mul INT- 1 x mul pow ln INT+ 4 INT- 1 add x mul INT- 1 pow x INT+ 5".split(" ")
+    exp = parse_prefix_to_sympy(prefs)
+    exp = sp.simplify(exp)
+    print(exp)
+    print(constant_fold(exp))
+
+    # prefs = "mul x mul pow cos INT+ 4 INT- 3 pow ln INT+ 3 INT- 6".split(" ")
+    # exp = parse_prefix_to_sympy(prefs)
+    # print(exp)
+    # dag = sympy_to_dag(exp)
+
+    # exp = sp.parse_expr("(((((x0) + ((x0) - ((-((x0) + (x0))) / ((x0) + (x0))))) * k0) - (-((x0) + (x0)))) / (-((x0) + (x0)))) * ((-((((x0) + ((x0) - ((-((x0) + (x0))) / ((x0) + (x0))))) * k0) - ((x0) + ((x0) - ((-((x0) + (x0))) / ((x0) + (x0))))))) * ((((x0) + ((x0) - ((-((x0) + (x0))) / ((x0) + (x0))))) * k0) - ((x0) + ((x0) - ((-((x0) + (x0))) / ((x0) + (x0)))))))", evaluate=False)
+    # # print(sympy_to_prefix(exp))
+
+    # simp = sp.simplify(exp)
+    # pre = sympy_to_prefix(simp)
+    # print(pre)
+    # repars = parse_prefix_to_sympy(pre)
+    # print(simp)
+    # print(repars)

remend/plot_loss.py

@@ -0,0 +1,60 @@
+from matplotlib import pyplot as plt
+import json
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser("Plot loss for the training log")
+    parser.add_argument("-t", "--trainlog", required=True, help="Training log file")
+    parser.add_argument("-l", "--loss", help="Loss plot to save (optional)")
+    parser.add_argument("--log-scale", default=False, action="store_true", help="Log scale")
+    parser.add_argument("-P", "--no-plot", default=True, action="store_false", help="Don't open matplotlib figure")
+    args = parser.parse_args()
+
+    train_inner_upd, train_inner_loss = [], []
+    train_upd, train_loss = [], []
+    val_upd, val_loss = [], []
+
+    with open(args.trainlog, "r") as tl:
+        for line in tl:
+            # Filter out json
+            if line[0] != "{":
+                continue
+            try:
+                data = json.loads(line.strip())
+            except:
+                continue
+            if "loss" in data:
+                loss = float(data["loss"])
+                upd = int(data["num_updates"])
+                if len(train_inner_upd) == 0 or train_inner_upd[-1] < upd:
+                    train_inner_upd.append(upd)
+                    train_inner_loss.append(loss)
+            if "valid_loss" in data:
+                loss = float(data["valid_loss"])
+                upd = int(data["valid_num_updates"])
+                if len(val_upd) == 0 or val_upd[-1] < upd:
+                    val_upd.append(upd)
+                    val_loss.append(loss)
+            if "train_loss" in data:
+                loss = float(data["train_loss"])
+                upd = int(data["train_num_updates"])
+                if len(train_upd) == 0 or train_upd[-1] < upd:
+                    train_upd.append(upd)
+                    train_loss.append(loss)
+
+    plt.figure()
+    plt.plot(train_upd, train_loss, "r")
+    plt.plot(val_upd, val_loss, "b")
+    if len(train_inner_upd) > 0:
+        plt.plot(train_inner_upd, train_inner_loss, "r", alpha=0.3)
+    plt.legend(["train", "valid"])
+    if args.log_scale:
+        plt.yscale("log")
+    elif min(min(train_loss), min(val_loss)) < 1:
+        plt.ylim((0, 1))
+    plt.xlabel("Updates")
+    plt.ylabel("Loss")
+    if args.loss:
+        plt.savefig(args.loss)
+    if args.no_plot:
+        plt.show()

remend/preprocess_remaqe.py

@@ -0,0 +1,102 @@
+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:]):
+            # Constant
+            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()

remend/util.py

@@ -0,0 +1,21 @@
+from contextlib import contextmanager
+import signal
+import sympy as sp
+
+def timeout_handler(signum, frame):
+    raise TimeoutError("Block timed out")
+@contextmanager
+def timeout(duration):
+    signal.signal(signal.SIGALRM, timeout_handler)
+    signal.alarm(duration)
+    try:
+        yield
+    finally:
+        signal.alarm(0)
+
+class DecodeError(Exception):
+    pass
+
+def sympy_expr_ok(expr):
+    atoms = expr.atoms()
+    return not (sp.I in atoms or sp.oo in atoms or sp.zoo in atoms or sp.nan in atoms)