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
|
#include <benchmark/benchmark.h>
#include <ATen/code_template.h>
#include <torch/csrc/jit/ir/ir.h>
#include <torch/csrc/jit/ir/irparser.h>
#include <torch/csrc/jit/tensorexpr/kernel.h>
using namespace torch::jit;
using namespace torch::jit::tensorexpr;
static const std::string kernel_static_shapes_template = R"IR(
graph(%0 : Float(${dim}, strides=[1], device=cpu),
%1 : Float(${dim}, strides=[1], device=cpu)):
%2 : Float(${dim}, strides=[1]) = aten::mul(%0, %1)
%4 : Float(${dim}, strides=[1]) = aten::mul(%0, %2)
return (%4))IR";
static const std::string kernel_symbolic_shapes = R"IR(
graph(%0 : Float(SS(-2), strides=[1], device=cpu),
%1 : Float(SS(-2), strides=[1], device=cpu),
%SS_2 : int):
%2 : Float(SS(-2), strides=[1]) = aten::mul(%0, %1)
%4 : Float(SS(-2), strides=[1]) = aten::mul(%0, %2)
return (%4))IR";
class KernelBench : public benchmark::Fixture {
public:
void Eager(benchmark::State& state) {
auto dim = state.range(0);
auto a = at::rand({dim}, at::TensorOptions(at::kCPU).dtype(at::kFloat));
auto b = at::rand({dim}, at::TensorOptions(at::kCPU).dtype(at::kFloat));
for (auto _ : state) {
auto o = at::mul(a, at::mul(a, b));
}
}
void GraphWithStaticShapes(benchmark::State& state) {
auto dim = state.range(0);
auto graph = std::make_shared<Graph>();
at::jit::TemplateEnv env;
env.d("dim", dim);
const auto kernel_static_shapes =
format(kernel_static_shapes_template, env);
parseIR(kernel_static_shapes, &*graph);
TensorExprKernel k(graph);
auto a = at::rand({dim}, at::TensorOptions(at::kCPU).dtype(at::kFloat));
auto b = at::rand({dim}, at::TensorOptions(at::kCPU).dtype(at::kFloat));
std::vector<at::Tensor> inputs = {a, b};
for (auto _ : state) {
std::vector<IValue> stack = at::fmap<at::IValue>(inputs);
k.run(stack);
}
}
void GraphWithSymbolicShapes(benchmark::State& state) {
auto dim = state.range(0);
auto graph = std::make_shared<Graph>();
parseIR(kernel_symbolic_shapes, &*graph);
std::vector<torch::jit::StrideInput> input_desc = {
torch::jit::StrideInput::TENSOR_CONT};
std::unordered_map<
const torch::jit::Value*,
std::vector<torch::jit::StrideInput>>
symbolic_strides;
symbolic_strides[graph->inputs().at(0)] = input_desc;
symbolic_strides[graph->inputs().at(1)] = input_desc;
symbolic_strides[graph->outputs().at(0)] = input_desc;
std::vector<int64_t> symbolic_shape_inputs = {-2};
TensorExprKernel k(
graph, {}, symbolic_shape_inputs, false, symbolic_strides);
auto a = at::rand({dim}, at::TensorOptions(at::kCPU).dtype(at::kFloat));
auto b = at::rand({dim}, at::TensorOptions(at::kCPU).dtype(at::kFloat));
std::vector<at::Tensor> inputs = {a, b};
for (auto _ : state) {
std::vector<IValue> stack = at::fmap<at::IValue>(inputs);
stack.push_back(dim);
k.run(stack);
}
}
};
BENCHMARK_DEFINE_F(KernelBench, Eager)(benchmark::State& state) {
Eager(state);
}
BENCHMARK_DEFINE_F(KernelBench, StaticShapes)(benchmark::State& state) {
GraphWithStaticShapes(state);
}
BENCHMARK_DEFINE_F(KernelBench, SymbolicShapes)(benchmark::State& state) {
GraphWithSymbolicShapes(state);
}
BENCHMARK_REGISTER_F(KernelBench, Eager)->Range(32, 2048);
BENCHMARK_REGISTER_F(KernelBench, StaticShapes)->Range(32, 2048);
BENCHMARK_REGISTER_F(KernelBench, SymbolicShapes)->Range(32, 2048);
|
