[Question] How to vectorize update definitions of a function?

[mcleary]

This is just a question on scheduling update definitions of a function, so if there is a better place to ask this just let me know, but here we go.

I was experimenting with different schedules for a simple dot product between two input functions, here is the code:

Var x("x"), y("y");
Func in1("in1"), in2("in2");
RDom r(0, 4, "rdom");
in1(x) = x;
in2(x) = x;

Func g("g");
g(x) = 0;
g(x) += in1(4 * x + r) * in2(4 * x + r);

g.update().atomic().vectorize(r);

Now, when I use inline reductions, I can't specify a schedule for the update definition since it adds an anonymous function to do the sum and I get an error:

g(x) = sum(in1(4 * x + r) * in2(4 * x + r));
g.update().atomic().vectorize(r);

// Error: Call to update with index larger than last defined update stage for Func "g".

This error makes sense since the update definition is not part of g itself. Now, looking into the simd_op_check test I found this piece of code that does the trick

class HasInlineReduction : public Internal::IRVisitor {
    using Internal::IRVisitor::visit;
    void visit(const Internal::Call *op) override {
        if (op->call_type == Internal::Call::Halide) {
            Internal::Function f(op->func);
            if (f.has_update_definition()) {
                inline_reduction = f;
                result = true;
            }
        }
        IRVisitor::visit(op);
    }

public:
    Internal::Function inline_reduction;
    bool result = false;
} has_inline_reduction;
e.accept(&has_inline_reduction);

// Define a vectorized Halide::Func that uses the pattern.
Halide::Func f(name);
f(x, y) = e;
f.bound(x, 0, W).vectorize(x, vector_width);
f.compute_root();

Now, I could use this, but I was wondering if there is a better way to access the update definitions when using the inline reductions.