causal-conv1d / benchmarks /benchmark.py
Ashiedu's picture
Add files using upload-large-folder tool
45eef7f verified
Raw
History Blame Contribute Delete
464 Bytes
import torch
from kernels.benchmark import Benchmark
class CausalConv1dBenchmark(Benchmark):
def setup(self):
self.size = 1024
self.input = torch.randn(self.size, self.size, device=self.device)
self.out = torch.empty_like(self.input)
def benchmark_base(self):
self.kernel.causal_conv1d(self.out, self.input)
def verify_base(self) -> torch.Tensor:
# Reference implementation
return self.input.clone()