mode=['short'],
sample_func=cross_product
)
@torch.jit.script
def torch_matmul(a, b, iterations):
# type: (Tensor, Tensor, int)
result = torch.jit.annotate(torch.Tensor, None)
for _ in range(iterations):
result = torch.matmul(a, b)
return result
@benchmark_core.register_test
def test_matmul():
generate_pt_test(
[long_config, short_config],
map_pt_config_matmul,
[('matmul', torch_matmul)]
)
generate_c2_test(
[long_config, short_config],
map_c2_config_matmul,
[('matmul', 'MatMul')],
)
if __name__ == "__main__":
benchmark_runner.main()
mode=['short'],
sample_func=cross_product
)
@torch.jit.script
def torch_add(a, b, iterations):
# type: (Tensor, Tensor, int)
result = torch.jit.annotate(torch.Tensor, None)
for _ in range(iterations):
result = torch.add(a, b)
return result
@benchmark_core.register_test
def test_add():
generate_pt_test(
[long_config, short_config],
map_pt_config_add,
[('add', torch_add)]
)
generate_c2_test(
[long_config, short_config],
map_c2_config_add,
[('add', 'Add')],
)
if __name__ == "__main__":
benchmark_runner.main()