mirror of
https://github.com/jessicanataliagta/PSPRecomp
synced 2026-09-26 16:49:34 -04:00
Jessica_Natalia
62 lines
3.1 KiB
Python
62 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Conservative basic-block FPR register cache for generated AOT."""
|
|
from __future__ import annotations
|
|
import argparse, collections, pathlib, re
|
|
from dataclasses import dataclass
|
|
LABEL_RE = re.compile(r"(?m)^L_[0-9A-F]+:\n")
|
|
FPR_RE = re.compile(r"ctx\.fpr\[([0-9]|[12][0-9]|3[01])\]")
|
|
@dataclass
|
|
class Stats:
|
|
blocks_cached:int=0; registers_cached:int=0; occurrences_replaced:int=0; dirty_registers:int=0
|
|
def add(self,o):
|
|
self.blocks_cached+=o.blocks_cached; self.registers_cached+=o.registers_cached
|
|
self.occurrences_replaced+=o.occurrences_replaced; self.dirty_registers+=o.dirty_registers
|
|
|
|
def assign(line,name):
|
|
return re.search(rf"(?<![=!<>])\b{re.escape(name)}\s*=(?!=)",line) is not None
|
|
|
|
def block_transform(label,block,threshold):
|
|
st=Stats()
|
|
if 'rt.' in block or 'ctx.execute_' in block or 'ctx.set_fpr_bits' in block or 'ctx.fpr_bits(' in block:
|
|
return label+block,st
|
|
c=collections.Counter(FPR_RE.findall(block)); sel=sorted(int(r) for r,n in c.items() if n>=threshold)
|
|
if not sel:return label+block,st
|
|
dirty={r for r in sel if re.search(rf'ctx\.fpr\[{r}\]\s*=(?!=)',block)}
|
|
for r in sel:block=block.replace(f'ctx.fpr[{r}]',f'f{r}')
|
|
out=[*(f' float f{r} = ctx.fpr[{r}];\n' for r in sel)]
|
|
active=set()
|
|
def flush():
|
|
nonlocal active
|
|
if active:
|
|
out.extend(f' ctx.fpr[{r}] = f{r};\n' for r in sorted(active)); active=set()
|
|
for line in block.splitlines(True):
|
|
if 'if (branch_taken)' in line and active:flush()
|
|
for r in dirty:
|
|
if assign(line,f'f{r}'):active.add(r)
|
|
if ('goto ' in line or 'return;' in line) and active:flush()
|
|
out.append(line)
|
|
flush()
|
|
st.blocks_cached=1;st.registers_cached=len(sel);st.occurrences_replaced=sum(c[str(r)] for r in sel);st.dirty_registers=len(dirty)
|
|
return label+'{\n'+''.join(out)+'}\n',st
|
|
|
|
def transform(text,threshold=3):
|
|
ms=list(LABEL_RE.finditer(text)); total=Stats()
|
|
if not ms:return text,total
|
|
out=[];last=0
|
|
for i,m in enumerate(ms):
|
|
end=ms[i+1].start() if i+1<len(ms) else text.find('\n}\n\nvoid ',m.end())
|
|
if end<0:end=len(text)
|
|
out.append(text[last:m.start()]);t,st=block_transform(text[m.start():m.end()],text[m.end():end],threshold);out.append(t);total.add(st);last=end
|
|
out.append(text[last:]);return ''.join(out),total
|
|
|
|
def main():
|
|
ap=argparse.ArgumentParser();ap.add_argument('paths',nargs='+',type=pathlib.Path);ap.add_argument('--threshold',type=int,default=3);ap.add_argument('--check',action='store_true');a=ap.parse_args()
|
|
fs=[]
|
|
for p in a.paths:fs.extend(sorted(p.glob('generated_unit_*.cpp')) if p.is_dir() else [p])
|
|
total=Stats()
|
|
for p in fs:
|
|
s=p.read_text();t,st=transform(s,a.threshold);total.add(st)
|
|
if not a.check and t!=s:p.write_text(t,newline='\n')
|
|
print(f'FPR block cache: files={len(fs)} blocks={total.blocks_cached} locals={total.registers_cached} occurrences={total.occurrences_replaced} dirty={total.dirty_registers} threshold={a.threshold} check={int(a.check)}')
|
|
if __name__=='__main__':main()
|