File size: 5,155 Bytes
02c783d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | # Modifications Copyright(C)[2025] Advanced Micro Devices, Inc. All rights reserved.
# https://github.com/thunlp/TritonBench - Apache License 2.0
import os
import json
import argparse
from tb_eval.constants import NATIVE_PERF_GOLD_ROOT
# golden_metrics_folder = TBG_PERF_GOLD_ROOT
_MODULE_DIR = os.path.dirname(os.path.abspath(__file__))
def write_file(input_folder_path, results_path, golden_metrics_folder):
golden_metrics_list = os.listdir(golden_metrics_folder)
tmp_dir = os.path.join(results_path, 'tmp')
logs_dir = os.path.join(results_path, 'logs')
assert len(golden_metrics_list) > 0, f"golden_metrics_list: {golden_metrics_list} is empty in {golden_metrics_folder}"
if os.path.exists(results_path):
os.system(f'rm -rf {results_path}')
os.mkdir(results_path)
print(f"results_path: {results_path}")
if os.path.exists(tmp_dir):
os.system(f'rm -rf {tmp_dir}')
os.mkdir(tmp_dir)
print(f"tmp_dir: {tmp_dir}")
if os.path.exists(f'{logs_dir}'):
os.system(f'rm -rf {logs_dir}')
os.mkdir(logs_dir)
print(f"logs_dir: {logs_dir}")
tab = ' ' * 4
with open(f'{_MODULE_DIR}/performance_utils.py', 'r') as f:
performance_utils = f.readlines()
# performance_utils = performance_utils.replace('folder_path = "/home/lishangzhan/triton/bench_performance/results"', f'folder_path = "{results_path}"')
performance_utils_lines = []
for line in performance_utils:
if 'folder_path = ' in line:
line = tab * 2 + f'folder_path = "{results_path}"\n'
performance_utils_lines.append(line)
performance_utils = "".join(performance_utils_lines)
with open(f'{tmp_dir}/performance_utils.py', 'w') as f:
f.write(performance_utils)
input_file_list = os.listdir(input_folder_path)
print("input_file_list:", input_file_list)
for file in input_file_list:
if file[-3:] == ".py":
op = file[:-3]
perf_file_name = op + "_perf.py"
assert perf_file_name in golden_metrics_list, f"{perf_file_name} not in {golden_metrics_list}"
with open(os.path.join(golden_metrics_folder, perf_file_name), "r") as f:
# golden_metrics = f.read()
lines = f.readlines()
# print(lines)
updated_lines = []
for line in lines:
if "from TritonBench_v1." in line:
fn_name = line.split(" ")[-1].strip()
updated_lines.append(
f"import tb_eval.data.TritonBench.data.TritonBench_G_v1.{op}.{fn_name} as {fn_name}_ref"
)
if line == "sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))\n":
updated_lines.append(f"sys.path.append('{input_folder_path}')\n")
line = line.replace("from TritonBench_v1.", "from ")
line = line.replace("op_perf.get_do_bench_config()", "op_perf.get_do_bench_config(warmup=100, rep=1000)")
line = line.replace('folder_path = "/home/lishangzhan/triton/bench_performance/results"', f'folder_path = "{results_path}"')
updated_lines.append(line)
golden_metrics = "".join(updated_lines)
golden_metrics_lines = golden_metrics.split("\n")
flag = False
for i in range(len(golden_metrics_lines)):
if "input_tensor = self.to_cuda(input_tensor_)" in golden_metrics_lines[i]:
index_1 = i
if "results.append(result)" in golden_metrics_lines[i]:
index_2 = i + 1
flag = True
if flag:
for i in range(index_1, index_2):
golden_metrics_lines[i] = tab + golden_metrics_lines[i]
golden_metrics_lines.insert(index_1, tab*3 + "try:")
golden_metrics_lines.insert(index_2 + 1, tab*3 + "except Exception as e:")
golden_metrics_lines.insert(index_2 + 2, tab*4 + 'print(f"Failed to run benchmark for input tensor. Error: {e}")')
golden_metrics = "\n".join(golden_metrics_lines)
with open(os.path.join(tmp_dir, perf_file_name), "w") as f:
f.write(golden_metrics)
def parse_args():
parser = argparse.ArgumentParser(description='write_file')
parser.add_argument('--input_folder_path', type=str, help='input_folder_path')
parser.add_argument('--result_folder_path', type=str, help='result_folder_path')
parser.add_argument('--golden_metrics_folder', type=str, default=NATIVE_PERF_GOLD_ROOT, help='golden_metrics_folder_path')
args = parser.parse_args()
return args
if __name__ == "__main__":
args = parse_args()
input_folder_path = args.input_folder_path
results_path = args.result_folder_path
golden_metrics_folder = args.golden_metrics_folder
write_file(input_folder_path, results_path, golden_metrics_folder)
|