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
|
from . import benchmark
class PoolingBench(benchmark.Benchmark):
def __init__(self, case, mode, device, dtype, kernel_size, N, C, H, W):
super().__init__(mode, device)
self.case = case
self.kernel_size = kernel_size
self.N = N
self.C = C
self.H = H
self.W = W
self.data = self.rand(
[N, C, H, W], device=device, dtype=dtype, requires_grad=self.requires_grad
)
def forward(self):
if self.case == "maxpool":
y = self.max_pool2d(self.data, self.kernel_size, stride=1)
elif self.case == "avgpool":
y = self.avg_pool2d(self.data, self.kernel_size, stride=1)
return y
def config(self):
return [self.kernel_size, self.N, self.C, self.H, self.W]
def memory_workload(self):
if self.mode == "fwd":
sol_count = 1 + 1
algorithmic_count = 1 + 1
else:
sol_count = (1 + 1) + (1 + 1)
algorithmic_count = (1 + 1) + (2 + 1)
buffer_size = self.N * self.C * self.H * self.W
return {
"sol": buffer_size * sol_count,
"algorithmic": buffer_size * algorithmic_count,
}
@staticmethod
def default_configs():
return [[3, 16, 32, 256, 256]]
class MaxPoolBench(PoolingBench):
def __init__(self, *args):
super().__init__("maxpool", *args)
@staticmethod
def module():
return "maxpool"
class AvgPoolBench(PoolingBench):
def __init__(self, *args):
super().__init__("avgpool", *args)
@staticmethod
def module():
return "avgpool"
benchmark.register_benchmark_class(MaxPoolBench)
benchmark.register_benchmark_class(AvgPoolBench)
|
