forked from ccf-ai-infra/TileOPs-Metax
29 lines
817 B
Python
29 lines
817 B
Python
"""Workload definitions for elementwise op workloads with custom generators."""
|
|
|
|
import torch
|
|
|
|
from workloads.workload_base import WorkloadBase
|
|
|
|
|
|
class ReluTest(WorkloadBase):
|
|
|
|
def __init__(self, n_total: int, dtype: torch.dtype):
|
|
self.n_total = n_total
|
|
self.dtype = dtype
|
|
|
|
def gen_inputs(self) -> tuple[torch.Tensor]:
|
|
x = torch.randn(self.n_total, dtype=self.dtype, device="cuda")
|
|
return (x,)
|
|
|
|
|
|
class AddSameShapeTest(WorkloadBase):
|
|
|
|
def __init__(self, n_total: int, dtype: torch.dtype):
|
|
self.n_total = n_total
|
|
self.dtype = dtype
|
|
|
|
def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor]:
|
|
a = torch.randn(self.n_total, dtype=self.dtype, device="cuda")
|
|
b = torch.randn(self.n_total, dtype=self.dtype, device="cuda")
|
|
return a, b
|