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; }