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
|
#include "Halide.h"
using namespace Halide;
namespace {
class Complex {
Tuple t;
public:
Complex(Expr real, Expr imag)
: t(real, imag) {
}
Complex(Tuple tup)
: t(tup) {
}
Complex(FuncRef f)
: t(Tuple(f)) {
}
Expr real() const {
return t[0];
}
Expr imag() const {
return t[1];
}
operator Tuple() const {
return t;
}
};
// Define the usual complex arithmetic
Complex operator+(const Complex &a, const Complex &b) {
return Complex(a.real() + b.real(), a.imag() + b.imag());
}
Complex operator*(const Complex &a, const Complex &b) {
return Complex(a.real() * b.real() - a.imag() * b.imag(),
a.real() * b.imag() + a.imag() * b.real());
}
Complex conjugate(const Complex &a) {
return Complex(a.real(), -a.imag());
}
Expr magnitude(Complex a) {
return (a * conjugate(a)).real();
}
class Mandelbrot : public Generator<Mandelbrot> {
public:
Input<float> x_min{"x_min"};
Input<float> x_max{"x_max"};
Input<float> y_min{"y_min"};
Input<float> y_max{"y_max"};
Input<float> c_real{"c_real"};
Input<float> c_imag{"c_imag"};
Input<int> iters{"iters"};
Input<int> w{"w"};
Input<int> h{"h"};
Output<Buffer<int32_t, 2>> count{"count"};
void generate() {
Var x, y, z;
Complex initial(lerp(x_min, x_max, cast<float>(x) / w),
lerp(y_min, y_max, cast<float>(y) / h));
Complex c(c_real, c_imag);
mandelbrot(x, y, z) = initial;
RDom t(1, iters);
Complex current = mandelbrot(x, y, t - 1);
mandelbrot(x, y, t) = current * current + c;
// How many iterations until something escapes a circle of radius 2?
Tuple escape = argmin(magnitude(mandelbrot(x, y, t)) < 4);
// If it never escapes, use the value 0
count(x, y) = select(escape[1], 0, escape[0]);
Var xi, yi, xo, yo;
mandelbrot.compute_at(count, xo);
count.tile(x, y, xo, yo, xi, yi, 8, 8).parallel(yo).vectorize(xi, 4).unroll(xi).unroll(yi, 2);
}
private:
// Declared as a member variable to verify that Funcs-as-members won't cause
// spurious "Invalid Param name: __user_context" errors (Issue #561)
Func mandelbrot{"mandelbrot"};
};
} // namespace
HALIDE_REGISTER_GENERATOR(Mandelbrot, mandelbrot)
|
