extract_triton_kernels.py

import sys

filename = sys.argv[1]
with open(filename, 'r') as f:
    lines = f.readlines()

def extract_info(line):
    line = line.split()
    name = line[0].strip()
    self_gpu_time = line[6].strip()
    num_of_calls = int(line[10].strip())
    if self_gpu_time.endswith('ms'):
        self_gpu_time = float(self_gpu_time.replace('ms', ''))
    elif self_gpu_time.endswith('us'):
        self_gpu_time = float(self_gpu_time.replace('us', '')) / 1000.0
    elif self_gpu_time.endswith('s'):
        self_gpu_time = float(self_gpu_time.replace('s', '')) * 1000.0
    print(f'{name}, {self_gpu_time}, {num_of_calls}')
    return [name, self_gpu_time, num_of_calls]

total_timems = 0
triton_kernel_info = []
aten_kernel_info = []
for line in lines:
    line = line.strip().replace('XPU Triton kernel:', '')
    if 'triton_' in line:
        triton_kernel_info.append(extract_info(line))
    elif 'aten::' in line:
        aten_kernel_info.append(extract_info(line))

from openpyxl import Workbook
workbook = Workbook()
sheet = workbook.active
sheet.title = 'gpu_triton'
sheet.append(['opname', 'self_gpu_timems', 'calls'])
for item in triton_kernel_info:
    sheet.append(item)

sheet2 = workbook.create_sheet(title='gpu_aten')
sheet2.append(['opname', 'self_gpu_timems', 'calls'])
for item in aten_kernel_info:
    sheet2.append(item)

workbook.save(f'{filename}.xlsx')