Halide::Func convolution_layer(Halide::Func input, Halide::Func weights,
Halide::Func bias, int filter_size, int input_layers, int pool_size) {
// Convolution
Halide::Func convolution;
Halide::Var x, y, z, w;
Halide::RDom r(0, filter_size, 0, filter_size, 0, input_layers);
convolution(x, y, z, w) = 0.0f;
convolution(x, y, z, w) += weights(r.x, r.y, r.z, z) *
input(x + r.x, y + r.y, r.z, w);
// Max pool
Halide::Func subsample;
Halide::RDom s(0, pool_size, 0, pool_size);
subsample(x, y, z, w) = 0.0f;
subsample(x, y, z, w) = Halide::max(convolution(pool_size * x + s.x,
pool_size * y + s.y, z, w), subsample(x, y, z, w));
// Non-linear bias
Halide::Func biased;
biased(x, y, z, w) = tanh(subsample(x, y, z, w) + bias(z, 0));
Halide::Var x_inner, x_outer, y_inner, y_outer;
biased.parallel(w);
biased.tile(x, y, x_outer, y_outer, x_inner, y_inner, VECTORS, 2);
biased.vectorize(x_inner);
biased.unroll(y_inner);
return biased;
}
