-
-
Notifications
You must be signed in to change notification settings - Fork 5.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
RFC: better ntuple implementation without @generated
#13439
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -29,29 +29,43 @@ eltype{T,_}(::Type{NTuple{_,T}}) = T | |
|
||
## mapping ## | ||
|
||
ntuple(f::Function, n::Integer) = | ||
const ntuple = function(f, n::Integer) | ||
n<=0 ? () : | ||
n==1 ? (f(1),) : | ||
n==2 ? (f(1),f(2),) : | ||
n==3 ? (f(1),f(2),f(3),) : | ||
n==4 ? (f(1),f(2),f(3),f(4),) : | ||
n==5 ? (f(1),f(2),f(3),f(4),f(5),) : | ||
tuple(ntuple(f,n-5)..., f(n-4), f(n-3), f(n-2), f(n-1), f(n)) | ||
|
||
ntuple(f, ::Type{Val{0}}) = () | ||
ntuple(f, ::Type{Val{1}}) = (f(1),) | ||
ntuple(f, ::Type{Val{2}}) = (f(1),f(2)) | ||
ntuple(f, ::Type{Val{3}}) = (f(1),f(2),f(3)) | ||
ntuple(f, ::Type{Val{4}}) = (f(1),f(2),f(3),f(4)) | ||
ntuple(f, ::Type{Val{5}}) = (f(1),f(2),f(3),f(4),f(5)) | ||
@generated function ntuple{N}(f, ::Type{Val{N}}) | ||
if !isa(N,Int) | ||
:(throw(TypeError(:ntuple, "", Int, $(QuoteNode(N))))) | ||
tuple([f(i) for i = 1:Int(n)]...) | ||
end | ||
|
||
const _ntuple_tfunc_gensym = gensym() | ||
const _ntuple_tfunc = function(argtypes, argexprs, vtypes, sv) | ||
n = false | ||
if isa(argexprs[2], Integer) | ||
try | ||
n = Int(argexprs[2]) | ||
end | ||
end | ||
|
||
Core.Inference.setindex!(vtypes, | ||
Core.Inference.call(Core.Inference.VarState, Int, false), | ||
_ntuple_tfunc_gensym) | ||
F = Core.Inference.abstract_eval_call( | ||
Expr(:call, argexprs[1], _ntuple_tfunc_gensym), | ||
vtypes, sv) | ||
Core.Inference.delete!(vtypes, | ||
_ntuple_tfunc_gensym) | ||
|
||
if isa(n, Int) | ||
return NTuple{n, F} | ||
elseif F === Any | ||
return Tuple | ||
else | ||
M = N-5 | ||
:(tuple(ntuple(f, Val{$M})..., f($N-4), f($N-3), f($N-2), f($N-1), f($N))) | ||
return Tuple{Vararg{F}} | ||
end | ||
end | ||
if isdefined(Core, :Inference) && isdefined(Core.Inference, :add_tfunc) | ||
Core.Inference.add_tfunc(ntuple, 2, 2, _ntuple_tfunc) | ||
end | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Does all this have to go here rather than in There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. yes, that's part of the strategy here. there are two |
||
|
||
# 0 argument function | ||
map(f) = f() | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Would be awesome if we didn't have that splatting penalty...
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if
ntuple
doesn't get inlined (say, because there are too many cases or f is unknown), the cost of splatting is the easy part. ideally, this will eventually be accessible as a direct iterator:(f(i) for i = 1:Int(n))
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
From the compiler standpoint it may be the easy part, but it's the one part that you can't do anything about from the Julia side short of writing a
@generated
function. (Inlining can be forced with@inline
,f
could be defined as a Functor.) I expect you noticed #13359. We're talking factors of 100 in terms of performance.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
i haven't looked into the exact issues in #13359 to see which point in codegen is pessimizing its assumptions for that example. but this is calling
tuple
, which is a more like syntax than an actual function call. the tfunc below then tries to do a better job of computing the type information that will be needed by the callee to limit the loss of type information in the caller where it could cause a much greater slowdown.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sounds like a big step in the right direction, then.