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
|
#include "Halide.h"
#include "halide_benchmark.h"
#include <cstdio>
using namespace Halide;
using namespace Halide::Tools;
void simple_version(float *A, float *B, float *C, int width, int stride) {
for (int iy = 0; iy < width; iy++) {
for (int ix = 0; ix < width; ix++) {
float *cc = C + iy * stride + ix;
*cc = 0.0f;
for (int ik = 0; ik < width; ik++) {
*cc = *cc + A[iy * stride + ik] * B[ik * stride + ix];
}
}
}
}
int main(int argc, char **argv) {
Target target = get_jit_target_from_environment();
if (target.arch == Target::WebAssembly) {
printf("[SKIP] Performance tests are meaningless and/or misleading under WebAssembly interpreter.\n");
return 0;
}
const int matrix_size = 992;
ImageParam A(type_of<float>(), 2);
ImageParam B(type_of<float>(), 2);
Var x("x"), xi("xi"), xo("xo"), y("y"), yo("yo"), yi("yi"), yii("yii"), xii("xii");
Func matrix_mul("matrix_mul");
RDom k(0, matrix_size);
RVar ki;
matrix_mul(x, y) += A(k, y) * B(x, k);
Func out;
out(x, y) = matrix_mul(x, y);
Var xy;
out.tile(x, y, xi, yi, 24, 32)
.fuse(x, y, xy)
.parallel(xy)
.split(yi, yi, yii, 4)
.vectorize(xi, 8)
.unroll(xi)
.unroll(yii);
matrix_mul.compute_at(out, yi)
.vectorize(x, 8)
.unroll(y);
matrix_mul.update(0)
.reorder(x, y, k)
.vectorize(x, 8)
.unroll(x)
.unroll(y)
.unroll(k, 2);
out
.bound(x, 0, matrix_size)
.bound(y, 0, matrix_size);
out.compile_jit();
Buffer<float> mat_A(matrix_size, matrix_size);
Buffer<float> mat_B(matrix_size, matrix_size);
Buffer<float> output(matrix_size, matrix_size);
// init randomly
for (int iy = 0; iy < matrix_size; iy++) {
for (int ix = 0; ix < matrix_size; ix++) {
mat_A(ix, iy) = (rand() % 256) / 256.0f;
mat_B(ix, iy) = (rand() % 256) / 256.0f;
}
}
A.set(mat_A);
B.set(mat_B);
out.realize(output);
double t = benchmark([&]() {
out.realize(output);
});
// check results
Buffer<float> output_ref(matrix_size, matrix_size);
Buffer<float> output_halide(matrix_size, matrix_size);
simple_version(mat_A.data(), mat_B.data(), output_ref.data(), mat_A.width(), mat_A.stride(1));
out.realize(output_halide);
bool halide_correct = true;
for (int iy = 0; iy < matrix_size && halide_correct; iy++) {
for (int ix = 0; ix < matrix_size; ix++) {
halide_correct = halide_correct && (std::abs(output_ref(ix, iy) - output_halide(ix, iy)) < 0.001f);
}
}
if (halide_correct) {
printf("Halide results - OK\n");
} else {
printf("Halide results - FAIL\n");
return 1;
}
// Uncomment to see the generated assembly.
/*
{
Target t("host-no_asserts-no_runtime-no_bounds_query");
out.compile_to_assembly("/dev/stdout", matrix_mul.infer_arguments(), t);
}
*/
float gflops = 2.0f * matrix_size * matrix_size * matrix_size / 1e9f;
printf("Halide: %fms, %f GFLOP/s\n\n", t * 1e3, (gflops / t));
printf("Success!\n");
return 0;
}
|
