float_quantize_fast.py (7268B)
1 #!/usr/bin/env python3 2 """Ablate IEEE bit packing, width caching and bounded positive shifts.""" 3 import argparse 4 import json 5 from pathlib import Path 6 import random 7 import statistics 8 9 from float_quantize import ROOT, BIAS, cases, near_cases, oracle, encoded 10 from float_multiply import Runner, sha 11 from base_workloads import encode as nat 12 from http_headers import TRUE, pair, tree_list, expect 13 14 VARIANTS=['control','packed','cachedWidths','optimized','selected'] 15 16 17 def main(): 18 p=argparse.ArgumentParser(description=__doc__) 19 p.add_argument('tricu',type=Path) 20 p.add_argument('--directory',required=True,type=Path) 21 args=p.parse_args() 22 root=args.directory.resolve();root.mkdir(parents=True,exist_ok=False) 23 work=root/'workspace';work.mkdir() 24 sources={'driver.tri':'bench/float-quantize-fast.tri','quantize.tri':'bench/float-quantize.tri', 25 'pipelines.tri':'bench/float-pipelines.tri','kernels.tri':'bench/float-significands.tri', 26 'layouts.tri':'bench/float-layouts.tri','native_math.tri':'lib/math/native.tri'} 27 for name,path in sources.items(): (work/name).write_bytes((ROOT/path).read_bytes()) 28 report=dict(complete=False,runtime_sha256=sha(args.tricu),compiler_sha256=sha(ROOT/'bootstrap/build-stable.arboricx'), 29 harness_sha256=sha(Path(__file__)),oracle_sha256=sha(ROOT/'test/float_quantize.py'), 30 sources={n:sha(ROOT/p) for n,p in sources.items()},validation=[],demand=[],samples=[],summaries=[]) 31 runner=Runner(args.tricu,root,1800,memory_mib=2048) 32 try: 33 tool=runner.imported(ROOT/'bootstrap/build-stable.arboricx','compiler') 34 bundle=root/'fast.arboricx' 35 runner.run('compile',['compile','--file',work/'driver.tri','--toolchain',tool,'--tool-store',runner.store, 36 '--store',runner.store,'--output',bundle,'--compile-timeout','240'],270) 37 pin=runner.imported(bundle,'fast');report['bundle_sha256']=sha(bundle) 38 def run(label,name,raw,wanted,count,measure=False): 39 file=root/(label+'.input.tree');file.write_bytes(raw) 40 actual=runner.invoke(label,pin,name,['--arg-tree-file',file]+(['--profile-full-result'] if measure else [])) 41 expect(actual,wanted,label) 42 if not measure: 43 report['validation'].append(dict(label=label,cases=count,input_sha256=sha(file))) 44 return 45 points={v['profile']:v for line in (root/(label+'.log')).read_text().splitlines() 46 if line.startswith('{') for v in [json.loads(line)]} 47 a,b=points['invoke-begin-1'],points['force-end-1'] 48 return dict(wall_ns=b['wall_ns']-a['wall_ns'],allocated_bytes=b['allocated_bytes']-a['allocated_bytes']) 49 def quant_material(rows): 50 return tree_list([pair(nat(s),nat(n)) for n,s in rows]),tree_list([encoded(oracle(n,s)) for n,s in rows]) 51 rows=cases() 52 # Exercise every normalization distance, including exact subnormal clamps. 53 rows += [(2**w-1,s) for w in range(1,54) for s in [1230,1231,1256,2252,3275]] 54 report['quantizer_cases_per_variant']=len(rows) 55 for offset in range(0,len(rows),24): 56 batch=rows[offset:offset+24];raw,wanted=quant_material(batch) 57 for name in VARIANTS: run(f'check-{name}-{offset}',name,raw,wanted,len(batch)) 58 pack_rows=[(e,2**52+f) for e in [1,2,3,1022,1023,2045,2046] 59 for f in [0,1,2,2**52-1]+[2**k for k in range(52)]] 60 shift_rows=[(k,n) for k in range(64) for n in [1,2,3,2**52,2**53-1,2**105+1]] 61 for name,rows,wants in [('packs',pack_rows,[(e << 52)+(m % 2**52) for e,m in pack_rows]), 62 ('shifts',shift_rows,[n << k for k,n in shift_rows]), 63 ('selectShifts',shift_rows,[n << k for k,n in shift_rows])]: 64 report[name+'_cases']=len(rows) 65 for offset in range(0,len(rows),32): 66 batch=rows[offset:offset+32] 67 run(f'check-{name}-{offset}',name,tree_list([pair(nat(a),nat(b)) for a,b in batch]), 68 tree_list([nat(n) for n in wants[offset:offset+32]]),len(batch)) 69 def near_material(rows,sign): 70 raw=tree_list([pair(nat(s),pair(nat(g),pair(nat(a),nat(b)))) for s,g,a,b in rows]) 71 wanted=tree_list([encoded(oracle((a << g)+sign*b,s)) for s,g,a,b in rows]) 72 return raw,wanted 73 near=near_cases();report['near_cases_per_worker']=len(near) 74 for suffix,sign in [('Adds',1),('Subtracts',-1)]: 75 for offset in range(0,len(near),24): 76 batch=near[offset:offset+24];raw,wanted=near_material(batch,sign) 77 for mode in ['Generic','Splice','Select']: 78 name='near'+mode+suffix 79 run(f'check-{name}-{offset}',name,raw,wanted,len(batch)) 80 for name,wanted in [('zeroPacking',encoded((0,False,False))),('zeroCached',encoded((0,False,False))), 81 ('zeroFast',encoded((0,False,False))),('packLowDemand',TRUE),('zeroShiftDemand',TRUE),('zeroSelective',encoded((0,False,False))),('zeroSelectiveShift',TRUE)]: 82 actual=runner.invoke('demand-'+name,pin,name,['--arg','0']);expect(actual,wanted,name) 83 report['demand'].append(dict(name=name,outcome='expected-result')) 84 rng=random.Random(107453) 85 groups={ 86 'normal-106':[(rng.randrange(2**105,2**106),BIAS-104) for _ in range(16)], 87 'cancellation':[(i+1,BIAS-52) for i in range(16)], 88 'subnormal-midpoints':[(2**54+i,BIAS-1077) for i in range(16)], 89 'tiny-to-zero':[(i+1,BIAS-1126) for i in range(16)], 90 'overflow-midpoint':[((2**54-1)*8+i%5-2,BIAS+967) for i in range(16)]} 91 jobs=[] 92 for group,rows in groups.items(): 93 raw,wanted=quant_material(rows) 94 jobs += [(group,name,raw,wanted,len(rows)) for name in VARIANTS] 95 for gap in [1,27,54]: 96 rows=[(BIAS-52-gap,gap,2**52+i,2**53-1-i) for i in range(16)] 97 for suffix,sign in [('Adds',1),('Subtracts',-1)]: 98 raw,wanted=near_material(rows,sign) 99 jobs += [(f'near-{gap}-{suffix}','near'+mode+suffix,raw,wanted,len(rows)) for mode in ['Generic','Splice','Select']] 100 for rep in range(5): 101 schedule=list(jobs);rng.shuffle(schedule) 102 for group,name,raw,wanted,size in schedule: 103 sample=run(f'measure-{group}-{name}-{rep}',name,raw,wanted,size,True) 104 report['samples'].append(dict(group=group,name=name,size=size,repeat=rep,warmup=rep==0,**sample)) 105 for group,name,_,_,size in jobs: 106 samples=[s for s in report['samples'] if s['group']==group and s['name']==name and not s['warmup']] 107 row=dict(group=group,name=name,size=size) 108 for key in ['wall_ns','allocated_bytes']: 109 row[key]=statistics.median(s[key] for s in samples) 110 row[key+'_min']=min(s[key] for s in samples);row[key+'_max']=max(s[key] for s in samples) 111 report['summaries'].append(row);print(json.dumps(row),flush=True) 112 report['complete']=True 113 finally: 114 (root/'report.json').write_text(json.dumps(report,indent=2)+'\n') 115 116 117 if __name__=='__main__': main()