Adding "algebraic expansions" ?

[marius311]

I've been toying with adding an "algebraic expansion" feature to StaticArrays and was curious for feedback / ideas. Sorry if this is just noise, feel free to close, but I thought people working on this package would be both interested and knowledgeable on this topic. So, by algebraic expansion I mean, suppose I'm doing the matrix operation,

(x::RowVector{SVector}) * (A::SMatrix) * (y::SVector)

Say these are all dimension 2, then it would be really cool if somehow we could evaluate this most efficiently, which is just to do,

x[1]*A[1,1]*y[1] + x[1]*A[1,2]*y[2] + x[2]*A[2,1]*y[1] + x[2]*A[2,2]*y[2]

as opposed to the way its currently done which uses temporary arrays.

If you define a new operator, say ⨳, and define that broadcasting over this operator does the aforementioed expansion, i.e. above you would write @. x' ⨳ A ⨳ y, then in theory you could overload broadcast,

@generated function broadcast(func, args::StaticArray...)

    # here you know the AST of func via Base.uncompressed_ast(func) 
    # and the types of all the arguments
    # that's everything needed to do the algebraic expansion
    
    new_func_expr = algebraic_expand(func, args...)
    
    quote
        # here somehow eval new_func_expr and apply it to args
    end

end

By making it generated, it should be pretty easy to maintain type-stability. The problem, however, is how to evaluate my new_func_expr (which is an Expr) and use it inside the generated body. Turns out, Julia works pretty hard to prevent me from doing this :) (for good reason, I imagine) Here's three failed attempts on a simplified example,

@generated function foo()
    ex = :(x->x^2)
    quote
        ($ex)(3)
    end
end
foo() # -> generated function body not pure

@generated function foo()
    ex = :(x->x^2)
    quote
        eval($(QuoteNode(ex)))(3)
    end
end
foo() # -> world age / applicable method may be too new


@generated function foo()
    ex = :(x->x^2)
    quote
        $(eval(ex))(3)
    end
end
foo() # -> eval cannot be used in a generated function

This is basically where I am at thinking about this. I'm curious if anyone has other ideas how to get around this? Or maybe how to write a non-@generated version that could be type-stable?

[marius311]

Ah sorry, I realize now that in reducing the thing I'm really after to the above example, I actually created something that has a trivial solution, that is, just interpolate new_func_expr directly into the generated function body with the right variable names.

The reason for the seemingly convoluted ways I was trying to go about it above is because really I want the algebraically expanded expression to also be broadcast over, and since you can't currently use loop fusion in the generated body, I was trying to get around that.

I suppose this means it'd be quite doable to have this type of feature in StaticArrays.

[c42f]

I suspect we already generate very good code for the original example so TBH I'm not really sure what's actionable here.

But if you've got specific benchmarks where we should be doing better, please do share :-)