atomic vectorization can't handle predicates #8844
Description
Currently, atomic vectorization (i.e. VectorReduce nodes) bails out in the presence of if statements, which happen e.g. due to a GuardWithIf split on the RVar. Here's a failing example from @alexreinking
#include "Halide.h"
using namespace Halide;
int main(int argc, char **argv) {
ImageParam mat(Float(32), 2, "mat");
mat.dim(0).set_min(0).set_extent(mat.dim(0).extent() / 4 * 4);
mat.dim(1).set_min(0).set_stride(mat.dim(0).extent());
ImageParam vec(Float(32), 1, "vec");
vec.dim(0).set_bounds(0, mat.dim(0).extent());
Func mv{"mv"};
Var x{"x"};
// RDom r(0, vec.dim(0).extent() / 4 * 4); <- works with this, because no tail
RDom r(0, vec.dim(0).extent());
mv(x) += mat(r, x) * vec(r);
Func out = mv.in();
RVar ro{"ro"}, ri{"ri"};
Var u{"u"};
out.output_buffer().dim(0).set_bounds(0, mat.dim(1).extent() / 4 * 4);
out.vectorize(x, 4);
auto intm = mv.update().split(r, ro, ri, 4, TailStrategy::Predicate).rfactor(ri, u);
intm.compute_at(out, x).reorder_storage(u, x).vectorize(u).unroll(x);
intm.update().reorder(x, u, ro).vectorize(u).unroll(x);
mv.update().atomic().vectorize(ri, 4);
mv.bound_extent(x, 4);
out.compile_jit();
return 0;
}
I think what needs to happen is this:
- PredicateStoreLoads in VectorizeLoops.cpp should be a prepass on the scalar IR that converts IfThenElse nodes with conditions that depend on a to-be-vectorized loop var into scalar predicates on the enclosed loads and stores. It can walk inside atomic nodes.
- The IfThenElse visitor in VectorSubs mutator gets simpler, because it doesn't have to introduce predication
- The Atomic visitor in VectorSubs should handle predicates on the store node by masking the RHS using the identity of the corresponding reduction operator, and adding a VectorReduce with the Or operator on the original predicate to get a new reduced predicate for the Store itself.