Skip to content
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

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion base/coreimg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,9 @@ include("options.jl")
typealias Cint Int32
typealias Csize_t UInt
include("promotion.jl")
include("expr.jl")
include("tuple.jl")
include("range.jl")
include("expr.jl")
include("error.jl")

# core numeric operations & types
Expand Down
4 changes: 1 addition & 3 deletions base/deprecated.jl
Original file line number Diff line number Diff line change
Expand Up @@ -403,8 +403,6 @@ end

@deprecate with_env(f::Function, key::AbstractString, val) withenv(f, key=>val)

@deprecate ntuple(n::Integer, f::Function) ntuple(f, n)

# 0.4 discontinued functions

@noinline function subtypetree(x::DataType, level=-1)
Expand Down Expand Up @@ -841,4 +839,4 @@ for f in (:remotecall, :remotecall_fetch, :remotecall_wait)
@deprecate ($f)(w::Worker, f::Function, args...) ($f)(f, w::Worker, args...)
@deprecate ($f)(id::Integer, f::Function, args...) ($f)(f, id::Integer, args...)
end
end
end
3 changes: 3 additions & 0 deletions base/inference.jl
Original file line number Diff line number Diff line change
Expand Up @@ -458,6 +458,9 @@ const apply_type_tfunc = function (A, args...)
end
add_tfunc(apply_type, 1, IInf, apply_type_tfunc)

## external tfunc additions ##
add_tfunc(ntuple, 2, 2, _ntuple_tfunc)

function tuple_tfunc(argtype::ANY)
if isa(argtype,DataType) && argtype.name === Tuple.name
p = map(x->(isType(x) && !isa(x.parameters[1],TypeVar) ? typeof(x.parameters[1]) : x),
Expand Down
2 changes: 1 addition & 1 deletion base/sysimg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -35,9 +35,9 @@ include("options.jl")

# core operations & types
include("promotion.jl")
include("expr.jl")
include("tuple.jl")
include("range.jl")
include("expr.jl")
include("error.jl")

# core numeric operations & types
Expand Down
46 changes: 30 additions & 16 deletions base/tuple.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)]...)
Copy link
Member

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...

Copy link
Member Author

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))

Copy link
Member

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.

Copy link
Member Author

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.

Copy link
Member

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.

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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does all this have to go here rather than in inference.jl?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, that's part of the strategy here. there are two ntuple functions (Core.Inference and Base), and inference needs to know about both.


# 0 argument function
map(f) = f()
Expand Down
16 changes: 9 additions & 7 deletions test/tuple.jl
Original file line number Diff line number Diff line change
Expand Up @@ -145,12 +145,14 @@ foo(x, y, z) = x + y + z
@test any((true,true,false)) === true
@test any((true,true,true)) === true

@test @inferred(ntuple(Base.Abs2Fun(), Val{0})) == ()
@test @inferred(ntuple(Base.Abs2Fun(), Val{2})) == (1, 4)
@test @inferred(ntuple(Base.Abs2Fun(), Val{3})) == (1, 4, 9)
@test @inferred(ntuple(Base.Abs2Fun(), Val{4})) == (1, 4, 9, 16)
@test @inferred(ntuple(Base.Abs2Fun(), Val{5})) == (1, 4, 9, 16, 25)
@test @inferred(ntuple(Base.Abs2Fun(), Val{6})) == (1, 4, 9, 16, 25, 36)
ntuple_val{T}(x, ::Type{Val{T}}) = ntuple(x, T)
@test @inferred(ntuple_val(Base.Abs2Fun(), Val{0})) == ()
@test @inferred(ntuple_val(Base.Abs2Fun(), Val{2})) == (1, 4)
@test @inferred(ntuple_val(Base.Abs2Fun(), Val{3})) == (1, 4, 9)
@test @inferred(ntuple_val(Base.Abs2Fun(), Val{4})) == (1, 4, 9, 16)
@test @inferred(ntuple_val(Base.Abs2Fun(), Val{5})) == (1, 4, 9, 16, 25)
@test @inferred(ntuple_val(Base.Abs2Fun(), Val{6})) == (1, 4, 9, 16, 25, 36)

# issue #12854
@test_throws TypeError ntuple(identity, Val{1:2})
@test_throws TypeError ntuple(identity, Val{1})
@test_throws TypeError ntuple_val(identity, Val{1:2})