import json
import os

import numpy as np
import pandas as pd
import torch
import torch_npu
from openpyxl import Workbook

def write_excel(file, excel_info):
result_xlsx = Workbook()
del result_xlsx["Sheet"]
for sheet in excel_info.keys():
sheet_xlsx = result_xlsx.create_sheet(sheet, excel_info[sheet]["ids"])
sheet_xlsx.append(excel_info[sheet]["titles"])
for data in excel_info[sheet]["datas"]:
colum_data = []
for title in excel_info[sheet]["titles"]:
if title in data.keys():
colum_data.append(data[title])
else:
colum_data.append(None)
sheet_xlsx.append(colum_data)
if os.path.exists(file):
os.remove(file)
result_xlsx.save(file)
def find_file_by_name_fuzz_csv(directory, fuzzname):
results = []
for root, dirs, files in os.walk(directory):
for file_name in files:
if fuzzname in file_name and "csv" in file_name:
results.append(os.path.join(root, file_name))
return results

def read_csv(file):
csv_info = {"titles": [], "datas": []}
if os.path.exists(file):
result_df = pd.read_csv(file, header=0)
if result_df is not None:
for index, row in result_df.iterrows():
csv_info["titles"] = list(row.keys())
data = {}
for key in row.keys():
data[key] = row[key]
csv_info["datas"].append(data)
return csv_info

def get_perf_info_by_name(path, times, names):
run_times = 0
avg_duration = -1
max_aic_mte2_ratio = max_aiv_mte2_ratio = max_aic_mte3_ratio = max_aiv_mte3_ratio= -1
avg_aic_mte2_time = avg_aiv_mte2_time = avg_aic_mte3_time = avg_aiv_mte3_time = -1
api_name_list = []
collect_all = False
if names is None:
collect_all = True
else:
if isinstance(names, list):
api_name_list = names
else:
api_name_list.append(names)
durations = []
aic_mte2_ratios = []
aic_mte3_ratios = []
aic_mte2_times = []
aic_mte3_times = []

aiv_mte2_ratios = []
aiv_mte3_ratios = []
aiv_mte2_times = []
aiv_mte3_times = []
aicore_time = []
if os.path.exists(path):
    kernel_details_files = find_file_by_name_fuzz_csv(path, "kernel_details")
    if len(kernel_details_files) > 0:
        kernel_details = read_csv(kernel_details_files[0])
        for data in kernel_details["datas"]:
            if ("Type" in data.keys() and data["Type"] in api_name_list) or collect_all:
                run_times += 1
                if "Duration(us)" in data.keys():
                    durations.append(data["Duration(us)"])
                if "aic_mte2_time(us)" in data.keys():
                    aic_mte2_times.append(data["aic_mte2_time(us)"])
                if "aic_mte2_ratio" in data.keys():
                    aic_mte2_ratios.append(data["aic_mte2_ratio"])
                if "aiv_mte2_time(us)" in data.keys():
                    aiv_mte2_times.append(data["aiv_mte2_time(us)"])
                if "aiv_mte2_ratio" in data.keys():
                    aiv_mte2_ratios.append(data["aiv_mte2_ratio"])

                if "aic_mte3_time(us)" in data.keys():
                    aic_mte3_times.append(data["aic_mte3_time(us)"])
                if "aic_mte3_ratio" in data.keys():
                    aic_mte3_ratios.append(data["aic_mte3_ratio"])
                if "aiv_mte3_time(us)" in data.keys():
                    aiv_mte3_times.append(data["aiv_mte3_time(us)"])
                if "aiv_mte3_ratio" in data.keys():
                    aiv_mte3_ratios.append(data["aiv_mte3_ratio"])

    print(f"path:{path} names:{names} run_times:{run_times} times:{times}")
    if len(durations) > 0:
        avg_duration = sum(durations) / times
    if len(aic_mte2_ratios) > 0:
        max_aic_mte2_ratio = max(aic_mte2_ratios)
    if len(aiv_mte2_ratios) > 0:
        max_aiv_mte2_ratio = max(aiv_mte2_ratios)
    if len(aic_mte3_ratios) > 0:
        max_aic_mte3_ratio = max(aic_mte3_ratios)
    if len(aiv_mte3_ratios) > 0:
        max_aiv_mte3_ratio = max(aiv_mte3_ratios)

    if len(aic_mte2_times) > 0:
        avg_aic_mte2_time = sum(aic_mte2_times) / times
    if len(aiv_mte2_times) > 0:
        avg_aiv_mte2_time = sum(aiv_mte2_times) / times
    if len(aic_mte3_times) > 0:
        avg_aic_mte3_time = sum(aic_mte3_times) / times
    if len(aiv_mte3_times) > 0:
        avg_aiv_mte3_time = sum(aiv_mte3_times) / times
    print(f"avg_duration:{avg_duration} max_aic_mte2_ratio:{max_aic_mte2_ratio} max_aiv_mte2_ratio:{max_aiv_mte2_ratio} max_aic_mte3_ratio:{max_aic_mte3_ratio} max_aiv_mte3_ratio:{max_aiv_mte3_ratio}")
    print(f"avg_aic_mte2_time:{avg_aic_mte2_time} avg_aiv_mte2_time:{avg_aiv_mte2_time} avg_aic_mte3_time:{avg_aic_mte3_time} avg_aiv_mte3_time:{avg_aiv_mte3_time}")
return avg_duration, max_aic_mte2_ratio, max_aiv_mte2_ratio,max_aic_mte3_ratio,max_aiv_mte3_ratio,avg_aic_mte2_time,avg_aiv_mte2_time,avg_aic_mte3_time,avg_aiv_mte3_time

def find_file_by_name_fuzz(directory, fuzzname):
results = []
for root, dirs, files in os.walk(directory):
for file_name in files:
if fuzzname in file_name:
results.append(os.path.join(root, file_name))
return results

def write_aclnn_perf_data(prof_path, dest_csv, case_name):
api_statistic_files = find_file_by_name_fuzz(prof_path, "api_statistic")
op_statistic_files = find_file_by_name_fuzz(prof_path, "op_statistic")
op_summary_files = find_file_by_name_fuzz(prof_path, "kernel_details")
op_l2cache_files = find_file_by_name_fuzz(prof_path, "l2_cache")
print(f"[write_perf_data]api_statistic_files:{api_statistic_files} op_statistic_files:{op_statistic_files} op_summary_files:{op_summary_files} op_l2cache_files:{op_l2cache_files}")
api_statistic_json = ""
if len(api_statistic_files) < 1:
print(f"[write_perf_data] not api_statistic_files")
else:
api_statistic_json = json.dumps(read_csv(api_statistic_files[0]))
api_statistic_json = api_statistic_json.replace("\"", "'").replace(""", "'").replace("''", "'").replace("\n", " ")

op_statistic_json = ""
if len(op_statistic_files) < 1:
    print(f"[write_perf_data] not op_statistic_files")
else:
    op_statistic_json = json.dumps(read_csv(op_statistic_files[0]))
op_statistic_json = op_statistic_json.replace("\\\"", "'").replace("\"", "'").replace("''", "'").replace("\n", " ")

op_summary_json = ""
if len(op_summary_files) < 1:
    print(f"[write_perf_data] not op_summary_files")
else:
    op_summary_json = json.dumps(read_csv(op_summary_files[0]))
op_summary_json = op_summary_json.replace("\\\"", "'").replace("\"", "'").replace("''", "'").replace("\n", " ")

op_l2cache_json = ""
if len(op_l2cache_files) < 1:
    print(f"[write_perf_data] not op_l2cache_files")
else:
    op_l2cache_json = json.dumps(read_csv(op_l2cache_files[0]))
op_l2cache_json = op_l2cache_json.replace("\\\"", "'").replace("\"", "'").replace("''", "'").replace("\n", " ")

write_head = False
if not os.path.exists(dest_csv):
    write_head = True
with open(dest_csv, "a+") as file:
    if write_head:
        file.write("case_name,api_statistic,op_statistic,op_summary,op_l2cache\n")
    file.write(f"{case_name}," + f"\"{api_statistic_json}\"" + "," + f"\"{op_statistic_json}\"" + "," + f"\"{op_summary_json}\"" + "," + f"\"{op_l2cache_json}\"" +"\n")

def add_robustness_noise(data, seed, noise_ratio=0.001, scale=1000):
shape = data.shape
total_elements = data.numel()
print(f"[add_robustness_noise]total_elements: {total_elements}")
num_noisy_points = int(total_elements * noise_ratio)

if num_noisy_points == 0:
    print(f"[add_robustness_noise] warning: {noise_ratio} * {total_elements} = 0, no need inject noisy.")
    return data
np.random.seed(seed)
indices = np.random.choice(total_elements, size=num_noisy_points, replace=False)
print(f"[add_robustness_noise]inject noisy indices: {indices}")
flat_data = data.flatten()
flat_data[indices] *= scale
return flat_data.reshape(shape)

def gen_tensor_data(shape,dtype,seed,prama=0,pramb=1,mode=0,robust=False):
print(f"[gen_tensor_data]shape:{shape} dtype:{dtype} seed:{seed} prama:{prama} pramb:{pramb} mode:{mode} robust:{robust}")
if mode == 0:
#default
tensor = torch.randn(shape, generator=torch.Generator().manual_seed(seed)).to(dtype)
elif mode == 1:
#normal 正态分布 prama
np.random.seed(seed)
tensor = torch.from_numpy(np.random.normal(prama, pramb, size=shape)).to(dtype)
elif mode == 2:
# uniform 均匀分布
np.random.seed(seed)
tensor = torch.from_numpy(np.random.uniform(low=prama, high=pramb, size=shape)).to(dtype)
else:
tensor = torch.empty(shape).to(dtype)
if robust:
tensor = add_robustness_noise(tensor, seed, noise_ratio=0.001, scale=1000)
return tensor

def gen_single_input_tensor_data(range_type,manual_seed,shape,input_dtype,robustness=False,name=None):
print(f"[gen_fa_input_tensor_data][info]run gen data range_type:{range_type} manual_seed:{manual_seed}")
if range_type == 0:
# 默认的torch正态分布
prama = 0
pramb = 0
mode = 0
robustness = False
elif range_type == 1:
# numpy 均匀分布 [-0.001, 0.001]
mode = 2
prama = -0.001
pramb = 0.001
elif range_type == 2:
# numpy 均匀分布 [-5, 5]
mode = 2
prama = -5
pramb = 5
elif range_type == 3:
# numpy 正态分布
mode = 1
np.random.seed(manual_seed)
prama = np.random.uniform(-100, 100)
np.random.seed(manual_seed)
pramb = np.random.uniform(1, 25)
elif range_type == 4:
# numpy 均匀分布 [-2^-10, 2^-10]
mode = 2
prama = -2e-10
pramb = 2e-10
elif range_type == 5:
# numpy 均匀分布 [-2^-20, 2^-20]
mode = 2
prama = -2e-20
pramb = 2e-20
elif range_type == 6:
# numpy 均匀分布 [0, 1]
mode = 2
prama = 0
pramb = 1
elif range_type == 7:
# numpy 正态分布
mode = 1
np.random.seed(manual_seed)
prama = np.random.uniform(-5, 5)
np.random.seed(manual_seed)
pramb = np.random.uniform(0.1, 2)
else:
prama = 0
pramb = 0
mode = 0
robustness = False

tensor= gen_tensor_data(shape, input_dtype, manual_seed, prama=prama, pramb=pramb, mode=mode,robust=robustness)
print(f"[gen_data]name:{name} shape:{tensor.shape} dtype:{tensor.dtype}")
return tensor

def generate_shapes(type_size=2):
#1维size分区间1-256M:8M间隔, 256-64G:128M间隔 ，用脚本生成1维shape，类型是bfloat16(type_size=2)
shapes = []
# 区间1: 1-256MB, 间隔8MB
for size_mb in range(0, 256, 8):
size_bytes = size_mb * 1024 * 1024
num_elements = size_bytes // type_size
# Find a valid 1D shape (just the size itself)
if num_elements == 0:
num_elements = 1024 * 1024 // type_size
shapes.append((num_elements,))
# 区间2: 256MB-64GB, 间隔256MB
for size_mb in range(256, 64 * 1024 + 1, 256):
size_bytes = size_mb * 1024 * 1024
num_elements = size_bytes // type_size
shapes.append((num_elements,))
return shapes

def write_csv(file,head,line):
if not os.path.exists(file):
write_head = True
else:
write_head = False
with open(file, "a+") as f:
if write_head:
f.write(f"{head}\n")
f.write(f"{line}\n")

def run_index_select(input_shape, dim, casename, device, times=5, l2cache=True):
print('start_run_index_select')
kernel_names = "GatherV2"
src = torch.randn(input_shape, device=device, dtype=torch.bfloat16)
index = torch.randint(0, input_shape[dim], input_shape, device=device)
size_in = (src.numel() * src.element_size() + index.numel() * index.element_size()) / 1024 / 1024
current_dir = os.getcwd()
prof_path = os.path.join(current_dir, "index_select_prof")
case_prof = os.path.join(prof_path, casename)
os.system(f"rm -rf {case_prof}")
experimental_config = torch_npu.profiler._ExperimentalConfig(
profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
l2_cache=True
)
torch.npu.synchronize()
with torch_npu.profiler.profile(
activities=[
torch_npu.profiler.ProfilerActivity.NPU,
torch_npu.profiler.ProfilerActivity.CPU
],
schedule=torch_npu.profiler.schedule(wait=1 if l2cache else 0, warmup=1 if l2cache else 0, active=times, repeat=1, skip_first=1 if l2cache else 0),
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(case_prof),
record_shapes=True,
profile_memory=True,
with_stack=True,
experimental_config=experimental_config,
with_modules=True,
) as prof:
for i in range(times + (3 if l2cache else 0)):
output = torch.ops.aten.index_select(src, dim, index)
torch.npu.synchronize()
prof.step()
size_out = output.numel() * output.element_size() / 1024 / 1024
if os.path.exists(case_prof):
avg_duration, max_aic_mte2_ratio, max_aiv_mte2_ratio, max_aic_mte3_ratio, max_aiv_mte3_ratio, avg_aic_mte2_time, avg_aiv_mte2_time, avg_aic_mte3_time, avg_aiv_mte3_time = get_perf_info_by_name(case_prof, times,kernel_names)
write_aclnn_perf_data(case_prof, f"index_select_profiler.csv", casename)
return avg_duration, max_aic_mte2_ratio, max_aiv_mte2_ratio, max_aic_mte3_ratio, max_aiv_mte3_ratio, avg_aic_mte2_time, avg_aiv_mte2_time, avg_aic_mte3_time, avg_aiv_mte3_time, size_in, size_out
else:
print('error：profiler failed!')
return None, None, None, None, None, None, None, None, None, None, None

device_id = 0
device = torch.device(f'npu:{device_id}')
torch.npu.set_device(device)
soc_version = torch_npu.npu.get_device_name(device)

#run_index_select
result_file = f"index_select_{soc_version}.xlsx"
excel_info = {}
excel_info["index_select"] = {"ids":0,"titles":["case_name","size(in)(M)","size(out)(M)","avg_duration","max_aic_mte2_ratio","max_aiv_mte2_ratio","max_aic_mte3_ratio","max_aiv_mte3_ratio","avg_aic_mte2_time","avg_aiv_mte2_time","avg_aic_mte3_time","avg_aiv_mte3_time","size(in)/avg_aiv_mte2_time","size(out)/avg_aiv_mte3_time"],"datas":[]}
src_shapes = generate_shapes()
print(len(src_shapes))
dim = 0
for shape in src_shapes:
case_name = "index_select-" + str(shape).replace(",)", "").replace(",]", "").replace("[", "").replace("]", "").replace(",", "_").replace("(", "").replace(")", "")
try:
avg_duration, max_aic_mte2_ratio, max_aiv_mte2_ratio, max_aic_mte3_ratio, max_aiv_mte3_ratio, avg_aic_mte2_time, avg_aiv_mte2_time, avg_aic_mte3_time, avg_aiv_mte3_time, size_in, size_out = run_index_select(shape, dim, case_name, device=device)
in_mte2_bw = size_in / avg_aiv_mte2_time if avg_aiv_mte2_time and avg_aiv_mte2_time > 0 else -1
out_mte3_bw = size_out / avg_aiv_mte3_time if avg_aiv_mte3_time and avg_aiv_mte3_time > 0 else -1
excel_info["index_select"]["datas"].append(
{"case_name": case_name,"size(in)(M)":size_in,"size(out)(M)":size_out,"avg_duration": avg_duration, "max_aic_mte2_ratio": max_aic_mte2_ratio,
"max_aiv_mte2_ratio": max_aiv_mte2_ratio, "max_aic_mte3_ratio": max_aic_mte3_ratio,
"max_aiv_mte3_ratio": max_aiv_mte3_ratio, "avg_aic_mte2_time": avg_aic_mte2_time,
"avg_aiv_mte2_time": avg_aiv_mte2_time, "avg_aic_mte3_time": avg_aic_mte3_time,
"avg_aiv_mte3_time": avg_aiv_mte3_time, "size(in)/avg_aiv_mte2_time": in_mte2_bw,
"size(out)/avg_aiv_mte3_time": out_mte3_bw})
write_excel(result_file, excel_info)
except Exception as e:
print(e)
continue
write_excel(result_file, excel_info)