A simpler DSL? (Pass code literal to _avx_!)

[timholy]

It occurred to me that one thing making it a bit hard to debug problems is that in the final generated code, the names of the variables are mostly unrelated to the original code. Have you considered passing the raw expression as a type parameter to _avx_!? I'm envisioning something like translating

ex = :(for i = 1:n
    s += a[i]
end)

into

struct TypedExpr{head,args} end
struct TailTuple{head,tail} end
TypedExpr{:for,TailTuple{TypedExpr{:(=),TailTuple{:i,TypedExpr{:call,TailTuple{:colon,TailTuple{1,:n}}}...

This is basically just an Expr but one that could be passed as a type-parameter.

If you like this idea I can work on it.

[chriselrod]

I agree that the variable name mangling is a problem. It forces me to do a lot of untangling to figure out what's going on. I hate bugs that show up with @avx but not @_avx (this macro immediately lowers the LoopSet, rather than converting it to a type to call _avx_!), because that means I can't avoid the gensyms by using LoopSet(expr) to diagnose the problem; @avx_debug is a nightmare in comparison, mostly because of the gensym issue (that all the objects it references have to be defined in scope is also inconvenient).

I think two reasonable solutions are:

1. Returning to Expr-as-type, to preserve the original symbol names. Here was @MasonProtter's extremely clean implementation of such an idea. Ignore all the functions with gensym in their names; they don't actually call gensym.
2. Wherever reconstruct_loopset.jl uses gensym, add that symbol to the appropriate object in condense_loopset, and use it instead of the gensym call.

The reason I favored the "condense -> reconstruct" model was to support differentiating loops in a (Bayesian) modeling DSL, and have it be fast*. I'd call this a work in progress, but it's been about a month since I was last able to spend any time on it. I hope to be able to focus on this again soon to actually have "progress".
The idea was to, when someone defines a model, create the forward and reverse pass LoopSets -- which should actually be surprisingly straightforward -- and then convert them into @_avx_! calls so that they can later be compiled with type information when the function is called.
But I could probably delay the creation of the forward and reverse pass LoopSets until the compilation of the @generated logdensity functions, and dump the expression directly. This would require that they manage type information themselves so that they can pass it on to the LoopSets then (and some API for providing the information). I might eventually want to do that anyway, if I wish to make it more aggressively optimizing. But I also have some time pressure to actually produce working code.

*Because MLIR was originally created as part of TensorFlow, they must have this motivation too? I really should find the time to play around with TensorFlow Probability. Maybe it already is what I've been dreaming of.

---

An additional note on the original implementation of the Expr-as-types idea, is that it did not pass on any of the type information to the LoopSet. Meaning feature-parity would require more work than in that example. Although there aren't too many features yet, mostly "condense" -> "reconstruct" just provides info on:

1. Column vs row major (are the arrays transposed?)
2. Statically known size information (if they support the maybestatic* interface).
3. (I'd also call solving the need for managing types in the probability packages, and then the need for an API to provide the information to LoopSets a feature.)

This doesn't sound like too much to re-implement.

I do think the "reconstruct_loopset.jl" is simpler to work on and extend than add_*.jl, and that providing some separation is nice itself.
Whether you think adding CartesianIndexing support to the existing "condense -> reconstruct", or to the assembly of LoopSets from expressions is something to consider.

Some code is also needed to figure out what the call to _avx_! should look like, and I found having a LoopSet object around to answer questions about returned reductions convenient.

[timholy]

I'm 100% on board with the idea of pushing forward with the current design. Getting something that you can use to persuade others of its value will provide the stability and permission you need to keep working on this.

In the longer run (when your deadlines are not so pressing) I'd argue we should think about a refactor. The crux of the issue is that currently

• the expression is only available to the macro, not the generated function
• the types are only available to the generated function, not the macro

In my opinion, a lot of what makes some things a bit difficult here is that you have to do quite a lot of analysis with only half the information available. To me it seems likely that a fair bit of the complexity (e.g., all the CartesianIndex offsets due to needing to expand symbol names into tuples) will just disappear if we can get both types of information in the same place. This would be true if we wired this in as a compiler pass in Julia itself, but in the shorter term I think encoding the Expr as a type is by far the easiest way to make these gymnastics prettier. I don't think compile time will be that big a deal if we can prevent recompilation of the internal methods of LoopVectorization (#76, likely to require improvements to Julia itself), other than for the compilation of the functions this produces.