33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from src.common import DataType, Op
|
|
from src.pytorch.base import TorchBase
|
|
|
|
|
|
class TorchMatmulBench(TorchBase):
|
|
def __init__(self, output_path: Path):
|
|
super().__init__(output_path, Op.MATMUL)
|
|
|
|
def experiment(self, experiment_args: tuple[int, int], length: int, dtype: torch.dtype, device: torch.device):
|
|
shape_1, shape_2 = experiment_args
|
|
tensor_1 = torch.ones(shape_1, dtype=dtype, device=device, requires_grad=False)
|
|
tensor_2 = torch.ones(shape_2, dtype=dtype, device=device, requires_grad=False)
|
|
|
|
for _ in range(length):
|
|
_ = tensor_1 @ tensor_2
|
|
|
|
def name(self, experiment_args: tuple[int, int]) -> str:
|
|
shape_1, shape_2 = experiment_args
|
|
return f'{shape_1[0]}x{shape_1[1]} @ {shape_2[0]}x{shape_2[1]}'
|
|
|
|
def mop(self, experiment_args: tuple[int, int]) -> float:
|
|
shape_1, shape_2 = experiment_args
|
|
return (shape_1[0] * shape_2[1] / 1000_000) * 2 * (shape_1[1] - 1)
|
|
|
|
def run(self,
|
|
experiment_args: list[tuple[tuple[int, int], tuple[int, int]]],
|
|
experiment_count: int,
|
|
data_type: DataType):
|
|
super().run(experiment_args, experiment_count, data_type)
|