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
|
import operator_benchmark as op_bench
import benchmark_caffe2 as op_bench_c2
from benchmark_caffe2 import Caffe2BenchmarkBase # noqa: F401
from caffe2.python import core
"""Microbenchmarks for MatMul operator"""
# Configs for C2 Matmul operator
mm_long_configs = op_bench.cross_product_configs(
M=[8, 64, 128],
N=range(2, 10, 3),
K=[2 ** x for x in range(0, 3)],
trans_a=[True, False],
trans_b=[True, False],
tags=["long"]
)
mm_short_configs = op_bench.config_list(
attrs=[
[128, 128, 128, False, True],
[1024, 1024, 256, True, False],
[8192, 8192, 1024, True, False],
],
attr_names=["M", "N", "K", "trans_a", "trans_b"],
tags=["short"],
)
class MatMulBenchmark(op_bench_c2.Caffe2BenchmarkBase):
def init(self, M, N, K, trans_a, trans_b):
self.input_one = self.tensor([N, M]) if trans_a else self.tensor([M, N])
self.input_two = self.tensor([K, N]) if trans_b else self.tensor([N, K])
self.args = {'trans_a': trans_a, 'trans_b': trans_b}
self.output = self.tensor([M, K])
self.set_module_name("matmul")
def forward(self):
op = core.CreateOperator(
"MatMul", [self.input_one, self.input_two], self.output, **self.args
)
return op
op_bench_c2.generate_c2_test(mm_long_configs + mm_short_configs, MatMulBenchmark)
if __name__ == "__main__":
op_bench.benchmark_runner.main()
|
