#!/usr/bin/env python3

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import subprocess
import sys, os

if len(sys.argv) <= 1:
	exit(1)

subprocess.run(['rocprof', '--obj-tracking', 'on', '--basenames', 'on', '--timestamp', 'on', '-o', '/tmp/result.csv'] + sys.argv[1:])
print('\nROCm program finished.\n')

funcs = dict()

with open('/tmp/result.csv', 'r') as fp:
	line = fp.readline() # Skip Header Title
	while True:
		line = fp.readline()
		if not line:
			break
		line = line.strip()
		if not line:
			continue
		line = line.split(',')
		funcName = line[1]
		if len(funcName) > 2 and funcName[0] == funcName[-1] and funcName[0] == '"':
			funcName = funcName[1:-1]
		if 'CLS' in os.environ:
			funcName = funcName.split('_')[0]
		if funcName not in funcs:
			funcs[funcName] = {
				"SharedMem": line[8],
				"VGPRs": line[10],
				"SGPRs": line[11],
				"InitCost": int(line[-2]) - int(line[-3]),
				"PostCount": 1,
				"PostCost": 0
			}
		else:
			funcs[funcName]["PostCost"] += int(line[-2]) - int(line[-3])
			funcs[funcName]["PostCount"] += 1

ordered = []
full_gpu_cost = 0

for key in funcs:
	ordered.append([funcs[key]["PostCount"], funcs[key]["PostCost"], key])
	full_gpu_cost += funcs[key]["PostCost"]


ordered = sorted(ordered, key=lambda item: -item[1])

for item in ordered:
	print(' %8s%% %12d  %12.2f(ms)  %12.4f(ms)  [%s]' % (str(item[1] * 100.0 / full_gpu_cost)[0:5], item[0], item[1] * 1e-6, item[1] * 1e-6 / item[0], item[2] if len(item[2]) < 60 else item[2]))

