| | import json |
| | import os |
| | import ast |
| | import subprocess |
| | from random import randint |
| | from tqdm import tqdm |
| | import signal |
| | from multiprocessing import Pool, Lock, Value |
| | from dataloaders.ProblemState import ProblemState |
| | from dataloaders.TB_eval.utils import code_call_exec_success_allclose, code_kernel_profiling |
| | import re |
| | from tb_eval.evaluators.interface import get_evaluators |
| |
|
| | class TritonBench: |
| | def __init__(self, |
| | statis_path, |
| | py_folder, |
| | instruction_path, |
| | golden_metrics, |
| | py_interpreter, |
| | perf_ref_folder, |
| | perf_G_path, |
| | result_path=None |
| | ): |
| | self.statis_path = statis_path |
| | self.py_folder = py_folder |
| | self.instruction_path = instruction_path |
| | self.golden_metrics_folder = golden_metrics |
| | self.py_interpreter = py_interpreter |
| | self.perf_ref_folder = perf_ref_folder |
| | self.perf_G_path = perf_G_path |
| | self.result_path = result_path |
| |
|
| | self.problem_states = self.load_ps(result_path) |
| | self.evaluator = get_evaluators['tbg']() |
| | |
| | def load_ps(self, path): |
| | problem_states = [] |
| | if path is None: |
| | with open(self.instruction_path, "r", encoding='utf-8') as file: |
| | instructions = json.load(file) |
| | statis_data = json.loads(open(self.statis_path, 'r', encoding='utf-8').read()) |
| |
|
| | for line in instructions: |
| | instruction = line["instruction"] |
| | label = line["output"] |
| |
|
| | |
| | g = label.replace("<|im_end|>", "").replace("<|EOT|>", "") |
| | tmp = False |
| | for item in statis_data: |
| | if g in item["output"]: |
| | file = item["file"] |
| | tmp = item |
| | break |
| | if tmp: |
| | statis_data.remove(tmp) |
| | elif g[50:220] == 'as tl\n\nif triton.__version__ >= "2.1.0":\n @triton.jit\n def _fwd_kernel(\n Q, K, V, sm_scale, B_Start_Loc, B_Seqlen, # B_LOC 内部记录每个batch 输入的真实位置， B_SEQ_len 记录': |
| | file = "context_attn_nopad.py" |
| | path = os.path.join(self.py_folder, file) |
| | assert os.path.exists(path), f"{file} not exist!" |
| | test_code = open(path, "r", encoding="utf-8").read().split("#"*146)[-1] |
| | assert "def test_" in test_code, "" |
| |
|
| | problemstate = ProblemState(instruction=instruction, |
| | label=label, |
| | test_code=test_code, |
| | filename=file, |
| | ) |
| | |
| | problem_states.append( |
| | problemstate |
| | ) |
| | else: |
| | with open(path, 'r', encoding='utf-8') as file: |
| | for line in file.readlines(): |
| | content = json.loads(line) |
| | problem_state = ProblemState(instruction=content["instruction"], |
| | label=content["label"], |
| | filename=content["filename"], |
| | ) |
| | if "test_code" in content: |
| | problem_state.test_code = content["test_code"] |
| | if "predict" in content: |
| | problem_state.solution = content["predict"] |
| | problem_states.append(problem_state) |
| | return problem_states |
| |
|
| | def __len__(self): |
| | return len(self.problem_states) |
| | |
| | def write_file(self, file_path): |
| | os.makedirs(os.path.dirname(file_path), exist_ok=True) |
| | with open(file_path, 'w') as f: |
| | for ps in self.problem_states: |
| | output = { |
| | "instruction": ps.instruction, |
| | "label": ps.label, |
| | "filename": ps.filename, |
| | } |
| | if ps.test_code: |
| | output["test_code"] = ps.test_code |
| | if ps.solution: |
| | output["predict"] = ps.solution |
| | output["speedup"] = ps.speedup |
| | |
| | f.write(json.dumps(output) + "\n") |
| | |
| | def test_opt_correctness(self, code, filename, tmp_dir, exe_dir="pass_exe",gpu_id=0): |
| | """ |
| | Runs a given Python script on a specified GPU. |
| | """ |
| | pass_call, pass_exe, speedup, call_stdout, call_stderr = self.evaluator(code, tmp_dir, exe_dir, filename, atol=1e-3, rtol=1e-3, custom_tests_path=None) |
| |
|
| | return pass_call, pass_exe, speedup, call_stdout, call_stderr |
| | |
| | def test_kernel_profiling(self, code, filename, tmp_dir, save_scripts=True, exe_dir="pass_exe", target_gpu=None, timeout=20*60): |
| | os.makedirs(exe_dir, exist_ok=True) |
| | profile_status, stdout_profile, stderr_profile, stdout_analyze = code_kernel_profiling(code=code, fname=filename, temp_root=tmp_dir, py_folder=self.py_folder, target_gpu=target_gpu, timeout=timeout) |
| | pass_prfiler = False |
| | if "True" in str(profile_status): |
| | pass_prfiler=True |
| | |
| | return pass_prfiler, stdout_profile, stderr_profile, stdout_analyze |
| | |
| | |
| |
|
