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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
|
#include "tensorflow/core/framework/types.pb.h"
#include <gtest/gtest.h>
#include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/lib/random/simple_philox.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/tensor.h"
namespace tensorflow {
random::PhiloxRandom philox(1, 1);
random::SimplePhilox rnd(&philox);
void Sparsify(Tensor* t, float sparsity) {
const int64 N = t->NumElements();
CHECK_LE(sparsity, 1);
if (sparsity <= 0) return;
auto flat = t->flat<float>();
static const uint32 K = 10000;
for (int64 i = 0; i < N; ++i) {
if (rnd.Uniform(K) < sparsity * K) {
flat(i) = 0;
}
}
}
Node* SparseMatMulNode(Graph* g, Node* in0, Node* in1, bool transpose_a,
bool transpose_b, bool a_sparse, bool b_sparse) {
Node* ret;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "SparseMatMul")
.Input(in0)
.Input(in1)
.Attr("transpose_a", transpose_a)
.Attr("transpose_b", transpose_b)
.Attr("a_is_sparse", a_sparse)
.Attr("b_is_sparse", b_sparse)
.Finalize(g, &ret));
return ret;
}
static Graph* SparseMatMulHelper(Graph* g, int m, int n, int d, float sparsity,
bool transpose_a, bool transpose_b,
bool a_sparse, bool b_sparse) {
a_sparse = a_sparse && (sparsity > 0);
b_sparse = b_sparse && (sparsity > 0);
auto left_shape = transpose_a ? TensorShape({d, m}) : TensorShape({m, d});
Tensor left(DataTypeToEnum<float>::value, left_shape);
left.flat<float>().setRandom();
if (a_sparse) {
Sparsify(&left, sparsity);
}
auto right_shape = transpose_b ? TensorShape({n, d}) : TensorShape({d, n});
Tensor right(DataTypeToEnum<float>::value, right_shape);
right.flat<float>().setRandom();
if (b_sparse) {
Sparsify(&right, sparsity);
}
SparseMatMulNode(g, test::graph::Constant(g, left),
test::graph::Constant(g, right), transpose_a, transpose_b,
a_sparse, b_sparse);
return g;
}
static Graph* SparseMatMul(int m, int n, int d, float sparsity,
bool transpose_a, bool transpose_b) {
Graph* g = new Graph(OpRegistry::Global());
return SparseMatMulHelper(g, m, n, d, sparsity, transpose_a, transpose_b,
true, false);
}
static Graph* MultiSparseMatMul(int m, int n, int d, float sparsity_a,
float sparsity_b) {
Graph* g = new Graph(OpRegistry::Global());
if (sparsity_a == 0 && sparsity_b > 0) {
SparseMatMulHelper(g, m, n, d, sparsity_a, false, false, false, false);
SparseMatMulHelper(g, n, d, m, sparsity_b, true, true, true, false);
SparseMatMulHelper(g, m, d, n, sparsity_b, false, false, true, false);
} else {
SparseMatMulHelper(g, m, n, d, sparsity_a, false, true, true, false);
SparseMatMulHelper(g, d, n, m, sparsity_a, true, false, true, true);
SparseMatMulHelper(g, m, d, n, sparsity_b, false, false, true, false);
}
return g;
}
#define BM_SPARSE(M, K, N, S) \
static void BM_Sparse##_##M##_##K##_##N##_##S(int iters) { \
testing::ItemsProcessed(static_cast<int64>(iters) * M * K * N * 2); \
std::string label = strings::Printf("%d_%d_%d_%0.2f", M, K, N, S / 100.0); \
testing::SetLabel(label); \
test::Benchmark("cpu", SparseMatMul(M, N, K, S / 100.0, false, false)) \
.Run(iters); \
} \
BENCHMARK(BM_Sparse##_##M##_##K##_##N##_##S);
BM_SPARSE(2048, 2048, 2048, 0);
BM_SPARSE(2048, 2048, 2048, 1);
BM_SPARSE(2048, 2048, 2048, 85);
BM_SPARSE(1024, 1024, 1024, 0);
BM_SPARSE(1024, 1024, 1024, 1);
BM_SPARSE(1024, 1024, 1024, 85);
BM_SPARSE(256, 256, 256, 1);
BM_SPARSE(512, 512, 512, 1);
#define BM_SPARSE_MULTI(M, K, N, S1, S2) \
static void BM_Sparse_Multi##_##M##_##K##_##N##_##S1##_##S2(int iters) { \
testing::ItemsProcessed(static_cast<int64>(iters) * M * K * N * 2 * 3); \
std::string label = strings::Printf("%d_%d_%d_%0.2f_%0.2f", M, K, N, \
S1 / 100.0, S2 / 100.0); \
testing::SetLabel(label); \
test::Benchmark("cpu", MultiSparseMatMul(M, N, K, S1 / 100.0, S2 / 100.0)) \
.Run(iters); \
} \
BENCHMARK(BM_Sparse_Multi##_##M##_##K##_##N##_##S1##_##S2);
BM_SPARSE_MULTI(512, 2140, 4096, 0, 82);
BM_SPARSE_MULTI(512, 4096, 2048, 83, 83);
#define BM_SPARSE_TR(M, K, N, S, TA, TB) \
static void BM_Sparse##_##M##_##K##_##N##_##S##_##TA##_##TB(int iters) { \
testing::ItemsProcessed(static_cast<int64>(iters) * M * K * N * 2); \
std::string label = \
strings::Printf("%d_%d_%d_%d_%d_%0.2f", M, K, N, TA, TB, S / 100.0); \
testing::SetLabel(label); \
test::Benchmark("cpu", SparseMatMul(M, N, K, S / 100.0, TA, TB)) \
.Run(iters); \
} \
BENCHMARK(BM_Sparse##_##M##_##K##_##N##_##S##_##TA##_##TB);
BM_SPARSE_TR(2048, 2048, 2048, 1, true, false);
BM_SPARSE_TR(2048, 2048, 2048, 1, false, true);
BM_SPARSE_TR(2048, 2048, 2048, 1, true, true);
} // end namespace tensorflow
|
