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
|
import time
import numpy as np
import torch
"""Microbenchmarks for Tensor repeat operator. Supports PyTorch."""
input_shapes = (
(4, 4, 1),
(16, 1, 32),
(64, 64, 1, 1),
(8, 256, 128),
(1, 64, 128, 32),
(512, 512),
)
repeats = (
(1, 1, 1, 64),
(1, 4, 1, 2),
(1, 2, 2, 15),
(1, 1, 3, 2),
(128, 1, 8, 1),
(1, 1, 2, 16),
)
NUM_WARMUP_ITERS = 5
NUM_BENCHMARK_ITERS = 10
DTYPE_TO_BYTES = {"float": 4}
def generate_data_for_repeat():
input_tensors = [torch.randn(*input_shape) for input_shape in input_shapes]
total_num_elements = 0
for input_tensor, repeat in zip(input_tensors, repeats):
total_num_elements += input_tensor.numel()
total_num_elements += input_tensor.numel() * np.prod(repeat)
return input_tensors, (total_num_elements * DTYPE_TO_BYTES["float"])
input_tensors, total_bytes = generate_data_for_repeat()
BYTES_TO_MB = 1.0 / 1000.0 / 1000.0
def pt_repeat(input_tensor, repeat):
return input_tensor.repeat(repeat)
def pt_repeat_n_times(niters):
for _ in range(niters):
for input_tensor, repeat in zip(input_tensors, repeats):
pt_repeat(input_tensor, repeat)
if __name__ == "__main__":
# Warm up runs.
pt_repeat_n_times(NUM_WARMUP_ITERS)
s = time.time()
pt_repeat_n_times(NUM_BENCHMARK_ITERS)
total_time_s = time.time() - s
total_time_per_iter_s = total_time_s / NUM_BENCHMARK_ITERS
achieved_bandwidth = (total_bytes * BYTES_TO_MB) / total_time_per_iter_s
print(f"Time:{total_time_per_iter_s} Achieved Bandwidth:{achieved_bandwidth} MB/s")
|
