Staged Functions

[Keno]

I've been nagging @JeffBezanson about staged functions for a while. For context, @JeffBezanson explained the basic idea on the mailing list here: https://groups.google.com/forum/#!searchin/julia-dev/staged$20functions/julia-dev/eSbNplQpTXQ/WMIUIRhBrxcJ. There is also a proof of concept implementation (which is unfortunately way too slow for my purposes) in examples/staged.jl (though it didn't work for me when I just tried it). I need them now, so I want to go ahead and implement them with proper compiler support, but before that I want to discuss some semantics right here.

Syntax wise I think it's fine to maintain the @staged macro but make it efficient, so that you would write, e.g. (yes the following example is a terrible use case, but I want to make some semantics clear).

@staged function foo{T<:Real}(::Type{T},::Type{T})
quote
$T
end
end

which would evaluate the above when called, e.g.

julia> foo(1,1) # This compiles a specialized method of foo for Int64, Int64
Int64

Some questions that remain are:

• Are you allowed to mix staged/non-staged methods? E.g. would I be allowed to define a foo(::Int64,::Int64) method?
• What happens when your functions are not constant on methods (i.e. do not depend solely on the inputs you pass in). I would be strongly inclined to make this undefined behavior.
• Implementation

I would appreciate any ideas, feedback, etc. Also @JeffBezanson, feel free to take this over at any time.

cc @vtjnash @StefanKarpinski @timholy (I imagine this would make cartesian simpler and straightforward)

[JeffBezanson]

Probably the same feature as #3706

My guess is that the right approach is to invoke the method right before type inference if it's marked as staged, and pass its result to type inference instead. Could potentially be done with just a few lines of code.

If the staged function is not pure, then yes, the behavior has to be undefined.

[Keno]

Yes, probably the same feature, but I think it might be slightly cleaner in a function context. Yes, I agree that this shouldn't be too hard to do. Essentially all it has to do is do specialization slightly differently and have a little bit of magic in type inference to actually run the function, but other than that it shouldn't be too bad.

[JeffBezanson]

I believe staged methods can be inserted in method tables as usual, alongside normal definitions. The only oddity is how to handle type inference with approximate types. Most likely, staged methods would need to be prepared to handle abstract types.

I'd prefer staged methods to use normal signatures, i.e.

@staged function foo{T<:Real}(::T, ::T)
quote
$T
end
end

Needing to wrap everything in Type{ } is boilerplatey, and using normal signatures will make it easier to understand applicability. You can see that foo is meant to be called on two Reals; only the author of foo needs to know that the arguments it gets will actually be types and not values.

[Keno]

That seems reasonable to me.

[Keno]

I'd be fine with having the staged functions be passed abstract types as long as I'm allowed to error out and demand better type info :)

[JeffBezanson]

I suppose that's fine with me; I will catch the error and return Any :)

[Keno]

Not sure that's what I want though, since the function wouldn't actually generate any meaningful code. Shouldn't that be some sort of compile time error?

[JeffBezanson]

That would just lead to a run time dispatch, which would call inference on the staged method again at run time with actual argument types.

[Keno]

Ah, I see. That makes sense then.

[toivoh]

Thanks for moving on this!

[timholy]

While it's similar to #3706, it's identical to #5395 (particularly the description here). So yes, I agree with Keno that this would be useful and makes more sense in a functional context 😄. In fact this has been currently on my mind too: I was going to ask if @Keno, @JeffBezanson, and/or @vtjnash were going to be around during the JuliaCon HackFest and if so whether I could steal some of your time to fill in some of the gaps in what I need to know to implement this. If there's even more widespread interest in this, so much the better. I think this will be a pretty impactful language feature, and almost essential for delivering a really great multidimensional array experience.

To the "stagedfunction"/"latefunction" method-matching rules for deciding whether a previously-compiled version matches or whether to compile a new one, I would also suggest the following extra tweak : in addition to matching on types, any function-valued argument also has to match in value. That will then also fix "function-valued argument inlining" (#3440) in a way that also permits the choice not to inline function-valued arguments (for example, in cases where "method explosion" is a concern).

[JeffBezanson]

Ah yes, I thought there was another @timholy issue about this, but didn't find it :)

Doing a sprint on this at the hackfest is a good idea.

At this point I'm planning to do function specialization on values in general.

[lindahua]

The staged functions together with function value specialization would make it much nicer to write generic functions such as reduce, reducedim, and mapreduce, etc.

[cdsousa]

Is "function specialization on values in general" something that enables doing

sin(π) = 0.0

?
Or am I completely misunderstanding this thread?

[JeffBezanson]

Yes, although that's technically possible now since pi is of type MathConst{:π}.

[cdsousa]

Hum, ok, I didn't know, thanks. The question could had been sin(2π) = 0.0 then.

[ivarne]

Try sinipi(2)

[JeffBezanson]

Yes, but defining methods on specific floating point values is probably a really bad idea. That requires checking for special values on every call to sin, which in practice means sin would always be dynamically dispatched. Realistically, the compiler will hardly ever be able to prove the value of a floating-point number, so such definitions are less useful than dispatching on, say, booleans.

[quinnj]

I think finally getting enums ironed out would also be a great here since a usual use case is having an enum arg to a function and branching code based on the specific enum value.

[cdsousa]

Ok I see, if it happens, function specialization on values must not be used like pattern matching of other languages. Thanks to both.

[carnaval]

If we do get value specialization through singleton types, a simple enum implementation could be :
typealias MyEnum Union(Single{:a}, Single{:b}, Single{:c})
thus staying compatible with code expecting symbols. Of course then some of your enums can be of non-empty intersection which might be a problem.

[mlubin]

This will be very useful in JuMP, although we actually need a further generalization. Ideally we'd like to write staged functions where the inputs are typed expression trees, or some transformation thereof. Is this feasible?

[JeffBezanson]

Greedy, greedy!
Actually macros and staged functions together are sufficient for that. The macro can generate a call to a staged function, passing the expression trees and, separately, all of its leaves. Then it will have access to everything, with a bit of work.

[mlubin]

As long as it's technically possible. So we'll have varargs staged functions?

[JeffBezanson]

Sure.