File size: 3,858 Bytes
a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a a7f6a13 f30427a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | import torch
import torch.nn.functional as F
from kernels.benchmark import Benchmark
class ReluBenchmark(Benchmark):
"""Same workloads/names as upstream kernels-community/relu (fp32
base/large) so results are directly comparable across the two repos."""
seed: int = 42
def setup(self):
self.x = torch.randn(1024, 1024, device=self.device, dtype=torch.float32)
self.out = torch.empty_like(self.x)
def benchmark_base(self):
self.out = self.kernel.relu(self.x)
def verify_base(self) -> torch.Tensor:
return F.relu(self.x)
def setup_large(self):
self.x = torch.randn(4096, 4096, device=self.device, dtype=torch.float32)
self.out = torch.empty_like(self.x)
def benchmark_large(self):
self.out = self.kernel.relu(self.x)
def verify_large(self) -> torch.Tensor:
return F.relu(self.x)
class ReluDtypeBenchmark(Benchmark):
"""Dtypes this build ships beyond upstream's fp32-only: fp16, bf16, int8.
base = 1024^2, large = 4096^2, matching ReluBenchmark's shapes.
Uses the kernel's out= form with a preallocated output: the allocating
form (`self.out = kernel.relu(x)`) keeps the previous output alive at
allocation time, so the caching allocator alternates two output blocks
and the effective footprint grows by a full buffer — at L2-boundary
working sets that alone pushes the timed loop off the L2 cliff while
the single-shot reference timing stays on it (torch.relu itself
measures ~0.33x vs its own single-shot under the same loop). The out=
form keeps one live output, matching the reference's footprint."""
seed: int = 42
def _float_setup(self, side, dtype):
self.x = torch.randn(side, side, device=self.device, dtype=dtype)
self.out = torch.empty_like(self.x)
self.ref_out = torch.empty_like(self.x)
def setup_fp16_base(self):
self._float_setup(1024, torch.float16)
def setup_fp16_large(self):
self._float_setup(4096, torch.float16)
def setup_bf16_base(self):
self._float_setup(1024, torch.bfloat16)
def setup_bf16_large(self):
self._float_setup(4096, torch.bfloat16)
def _int8_setup(self, side):
self.x = torch.randint(
-128, 128, (side, side), device=self.device, dtype=torch.int8
)
self.out = torch.empty_like(self.x)
self.ref_out = torch.empty_like(self.x)
def setup_int8_base(self):
self._int8_setup(1024)
def setup_int8_large(self):
self._int8_setup(4096)
def benchmark_fp16_base(self):
self.kernel.relu(self.x, out=self.out)
def benchmark_fp16_large(self):
self.kernel.relu(self.x, out=self.out)
def benchmark_bf16_base(self):
self.kernel.relu(self.x, out=self.out)
def benchmark_bf16_large(self):
self.kernel.relu(self.x, out=self.out)
def benchmark_int8_base(self):
self.kernel.relu(self.x, out=self.out)
def benchmark_int8_large(self):
self.kernel.relu(self.x, out=self.out)
def verify_fp16_base(self) -> torch.Tensor:
torch.clamp(self.x, min=0, out=self.ref_out)
return self.ref_out
def verify_fp16_large(self) -> torch.Tensor:
torch.clamp(self.x, min=0, out=self.ref_out)
return self.ref_out
def verify_bf16_base(self) -> torch.Tensor:
torch.clamp(self.x, min=0, out=self.ref_out)
return self.ref_out
def verify_bf16_large(self) -> torch.Tensor:
torch.clamp(self.x, min=0, out=self.ref_out)
return self.ref_out
def verify_int8_base(self) -> torch.Tensor:
torch.clamp(self.x, min=0, out=self.ref_out)
return self.ref_out
def verify_int8_large(self) -> torch.Tensor:
torch.clamp(self.x, min=0, out=self.ref_out)
return self.ref_out
|