121 lines
2.8 KiB
Python
121 lines
2.8 KiB
Python
import torch
|
|
import sys
|
|
import torch.utils.benchmark as benchmark
|
|
import numpy as np
|
|
import json
|
|
import polars as pl
|
|
import dataclasses
|
|
|
|
from config import Statistics, Configuration
|
|
|
|
# device = torch.device("cuda:0")
|
|
device = torch.device("cpu")
|
|
|
|
# ITERATIONS = 1_000_000
|
|
ITERATIONS = 1_000
|
|
|
|
|
|
def run_gemv_bench(workload, level):
|
|
match level:
|
|
case "X1":
|
|
ROWS = 16
|
|
case "X2":
|
|
ROWS = 32
|
|
case "X3":
|
|
ROWS = 64
|
|
case "X4":
|
|
ROWS = 128
|
|
|
|
COLUMNS = 128
|
|
|
|
matrix = torch.rand(
|
|
(ROWS, COLUMNS),
|
|
dtype=torch.float16,
|
|
device=device,
|
|
)
|
|
input_vector = torch.rand(COLUMNS, dtype=torch.float16, device=device)
|
|
|
|
timer = benchmark.Timer(
|
|
stmt="gemv(matrix, input_vector)",
|
|
setup="from workloads import gemv",
|
|
globals={"input_vector": input_vector, "matrix": matrix},
|
|
)
|
|
|
|
return int(timer.timeit(ITERATIONS).mean * 1e12)
|
|
|
|
|
|
def run_gemv_layers_bench(workload, level):
|
|
LESS_ITERATIONS = int(ITERATIONS / 100)
|
|
|
|
match level:
|
|
case "X1":
|
|
DIMENSIONS = 128
|
|
case "X2":
|
|
DIMENSIONS = 256
|
|
case "X3":
|
|
DIMENSIONS = 512
|
|
case "X4":
|
|
DIMENSIONS = 1024
|
|
|
|
matrix = torch.rand(
|
|
(DIMENSIONS, DIMENSIONS),
|
|
dtype=torch.float16,
|
|
device=device,
|
|
)
|
|
input_vector = torch.rand(DIMENSIONS, dtype=torch.float16, device=device)
|
|
|
|
timer = benchmark.Timer(
|
|
stmt="gemv_layers(matrix, input_vector)",
|
|
setup="from workloads import gemv_layers",
|
|
globals={"input_vector": input_vector, "matrix": matrix},
|
|
)
|
|
|
|
return int(timer.timeit(LESS_ITERATIONS).mean * 1e12)
|
|
|
|
|
|
def run_vector_bench(workload, level):
|
|
match level:
|
|
case "X1":
|
|
ROWS = 256
|
|
case "X2":
|
|
ROWS = 512
|
|
case "X3":
|
|
ROWS = 1024
|
|
case "X4":
|
|
ROWS = 2048
|
|
|
|
vector_a = torch.rand(ROWS, dtype=torch.float16, device=device)
|
|
vector_b = torch.rand(ROWS, dtype=torch.float16, device=device)
|
|
|
|
timer = benchmark.Timer(
|
|
stmt=f"{workload}(vector_a, vector_b)",
|
|
setup=f"from workloads import {workload}",
|
|
globals={"vector_a": vector_a, "vector_b": vector_b},
|
|
)
|
|
|
|
return int(timer.timeit(ITERATIONS).mean * 1e12)
|
|
|
|
|
|
workloads = [
|
|
("vadd", run_vector_bench),
|
|
("vmul", run_vector_bench),
|
|
("haxpy", run_vector_bench),
|
|
("gemv", run_gemv_bench),
|
|
("gemv_layers", run_gemv_layers_bench),
|
|
]
|
|
|
|
levels = ["X1", "X2", "X3", "X4"]
|
|
|
|
results: list[dict] = []
|
|
|
|
for workload, workload_callback in workloads:
|
|
for level in levels:
|
|
runtime = workload_callback(workload, level)
|
|
element = {"workload": workload, "level": level, "runtime": runtime}
|
|
results.append(element)
|
|
print(element)
|
|
|
|
|
|
df = pl.DataFrame(results)
|
|
df.write_csv("rocm_results.csv")
|