ampere_research/pytorch/run.py
2024-12-05 12:20:05 -05:00

120 lines
4.0 KiB
Python
Executable File

#! /bin/python3
import data_stat
from data_stat import Stat, Cpu
import argparse
import os, sys
import subprocess, signal
import json
import time
parser = argparse.ArgumentParser()
parser.add_argument('cpu', choices=[x.name.lower() for x in Cpu])
parser.add_argument('matrix_file')
parser.add_argument('iterations', type=int)
parser.add_argument('baseline_time_s', type=int)
parser.add_argument('baseline_delay_s', type=int)
parser.add_argument('--perf', action='store_true')
parser.add_argument('--power', action='store_true')
parser.add_argument('-d', '--debug', action='store_true')
args = parser.parse_args()
args.cpu = Cpu[args.cpu.upper()]
program = {
Cpu.ALTRA: [
'apptainer', 'run', 'pytorch-altra.sif', '-c',
'numactl --cpunodebind=0 --membind=0 '
+ f'python spmv.py {args.matrix_file} {args.iterations}']
}
perf = ['perf', 'stat']
perf_args = {
Cpu.ALTRA: [
['-d', '-d'],
['-M', 'branch_misprediction_ratio'],
['-M', 'dtlb_walk_ratio,itlb_walk_ratio'],
['-M', 'l1d_cache_miss_ratio,l1i_cache_miss_ratio'],
['-M', 'l2_cache_miss_ratio,l2_tlb_miss_ratio,ll_cache_read_miss_ratio']]
}
def baseline_power(baseline_time_s: int) -> list:
power_process = subprocess.Popen(['./power.py', '-s', str(baseline_time_s)],
stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)
return [float(x) for x in power_process.communicate()[0].strip().split('\n')]
def run_program(program: list[str]) -> tuple[dict, str]:
process = subprocess.run(program,
stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
if args.debug:
print(process.stdout)
print(process.stderr)
return (json.loads(process.stdout), process.stderr)
result = dict()
result[Stat.CPU.value] = args.cpu.value
result[Stat.ITERATIONS.value] = args.iterations
result[Stat.BASELINE_TIME_S.value] = args.baseline_time_s
result[Stat.BASELINE_DELAY_S.value] = args.baseline_delay_s
if args.power:
time.sleep(args.baseline_delay_s)
result[Stat.POWER_BEFORE.value] = baseline_power(args.baseline_time_s)
if args.debug:
print(result)
run_program(program[args.cpu]) # Warmup
power_process = subprocess.Popen(['./power.py'],
stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)
result = result | run_program(program[args.cpu])[0]
power_process.send_signal(signal.SIGINT)
if args.debug:
print(result)
result[Stat.POWER.value] = [float(x) for x in power_process.communicate()[0].strip().split('\n')]
# Riemann Sum
from math import ceil
result[Stat.JOULES.value] = sum(result[Stat.POWER.value][-ceil(result[Stat.TIME_S.value]):-1]) + (result[Stat.POWER.value][-1] * (result[Stat.TIME_S.value] % 1))
if args.debug:
print(result)
#print(len(result['power']))
#print(sum(result['power']) / len(result['power']))
time.sleep(args.baseline_delay_s)
result[Stat.POWER_AFTER.value] = baseline_power(args.baseline_time_s)
if args.debug:
print(result)
if args.perf:
for perf_arg in perf_args[args.cpu]:
output = run_program(perf + perf_arg + program[args.cpu])[1]
print(output, file=sys.stderr)
result = result | data_stat.parse_output(output, args.cpu)
if args.debug:
print(result)
result = result | data_stat.derive_stats(result)
if args.debug:
print(result)
print(json.dumps(result))
#arch = subprocess.run(['uname', '-m'], stdout=subprocess.PIPE, text=True).stdout.strip()
#baseline = subprocess.run(
# ['./power.sh', args.baseline_time_s],
# stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
#print(baseline)
#for line in baseline.stdout.split('\n'):
# print("line")
# print(line)
#os.path.basename(args.matrix_file)
#subprocess.run(
# ['apptainer', 'run', 'pytorch-altra.sif', '-c',
# f'"numactl --cpunodebind=0 --membind=0 python spmv.py {args.matrix_file} {args.iterations}"'
# ],
# stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)