introduce runtime representation of broadcast fusion

fix #21094 fix #22060 fix #22053 replaces #22063
JuliaLang · Sep 13, 2017 · 535fefa · 535fefa
1 parent fce0a3c
commit 535fefa
Show file tree

Hide file tree

Showing 6 changed files with 349 additions and 129 deletions.
diff --git a/base/broadcast.jl b/base/broadcast.jl
@@ -609,4 +609,229 @@ macro __dot__(x)
     esc(__dot__(x))
 end
 
+############################################################
+## The parser turns dotted calls into the equivalent Fusion expression.
+## Effectively, this turns the Expr tree into a runtime AST,
+## for a limited subset of expression types.
+#
+## For example, in the expression:
+# d = sin.((a .+ (b .* c))...)
+## The kernel becomes
+# d' = Fusion{3}(
+#        FusionApply(
+#          sin,
+#          ( FusionCall(
+#              +,
+#              ( FusionArg{1}(),
+#                FusionCall(
+#                  *,
+#                  ( FusionArg{2}(),
+#                    FusionArg{3}() )), )), )),
+#        (:a, :b, :c))
+## and then the final expansion becomes:
+# d = broadcast(d', a, b, c)
+
+struct Fusion{N, vararg#=::Bool=#, T}
+    f::T
+    # Debugging Metadata:
+    # names::NTuple{N, Symbol}
+    # source::LineNumberNode
+    function Fusion{N, vararg}(f) where {N, vararg}
+        return new{N, vararg::Bool, typeof(f)}(f)
+    end
+end
+
+struct FusionArg{N}
+end
+
+struct FusionConstant{T}
+    c::T
+    function FusionConstant(c) where {}
+        return new{typeof(c)}(c)
+    end
+end
+
+struct FusionCall{F, Args<:Tuple}
+    f::F
+    args::Args
+    function FusionCall(f, args::Tuple) where {}
+        return new{typeof(f), typeof(args)}(f, args)
+    end
+end
+
+struct FusionApply{N, F, Args<:NTuple{N, Any}}
+    f::F
+    args::Args
+    function FusionApply(f, args::NTuple{N, Any}) where {N}
+        return new{N, typeof(f), typeof(args)}(f, args)
+    end
+end
+
+function kw_to_vec(kws::Vector{Any})
+    kwargs = Vector{Any}(2 * length(kws))
+    for i in 1:2:length(kws)
+        kw = kws[i]::Tuple{Any, Any}
+        kwargs[i] = getfield(kw, 1)
+        kwargs[i + 1] = getfield(kw, 2)
+    end
+    return kwargs
+end
+
+struct FusionKWCall{F, Args<:Tuple}
+    f::F
+    args::Args
+    kwargs::Vector{Any}
+    function FusionKWCall(f, args::Tuple; kwargs...) where {}
+        return new{typeof(f), typeof(args)}(f, args, kw_to_vec(kwargs))
+    end
+end
+
+struct FusionKWApply{F, Args<:Tuple}
+    f::F
+    args::Args
+    kwargs::Vector{Any}
+    function FusionKWApply(f, args::Tuple; kwargs...)  where {}
+        return new{typeof(f), typeof(args)}(f, args, kw_to_vec(kwargs))
+    end
+end
+
+function tuplehead(t::Tuple, N::Val)
+    return ntuple(i -> t[i], N)
+end
+@generated function tupletail(t::NTuple{M, Any}, ::Val{N}) where {N, M}
+    # alternative, non-generated versions,
+    # enable when inference is improved:
+    #tupletail(t, Nreq) = ntuple(i -> t[i + Nreq], length(t) - Nreq)
+    #tupletail(t, Nreq) = t[(Nreq + 1):end]
+    args = Any[ :(getfield(t, $i)) for i in (N + 1):M ]
+    tpl = Expr(:tuple)
+    tpl.args = args
+    return tpl
+end
+
+@inline (f::Fusion{N, false})(args::Vararg{Any, N}) where {N} = f.f(args...)
+function (f::Fusion{Nreq, true})(args::Vararg{Any, M}) where {Nreq, M}
+    M >= Nreq || throw(MethodError(f, args))
+    fargs = tuplehead(args, Val(Nreq))
+    vararg = tupletail(args, Val(Nreq))
+    return f.f((fargs..., vararg)...)
+end
+@inline (f::FusionArg{N})(args...) where {N} = args[N]
+@inline (f::FusionConstant)(args...) = f.c
+@inline (f::FusionCall)(args...) = f.f(map(a -> a(args...), f.args)...)
+# TODO: calling _apply on map _apply is not handled by inference
+# for now, we unroll some cases and generate others, to help it out
+#@inline (f::FusionApply)(args...) = Core._apply(f.f, map(a -> a(args...), f.args)...)
+@inline (f::FusionApply{0})(args...) = f.f()
+@inline (f::FusionApply{1})(args...) = f.f(f.args[1](args...)...)
+@inline (f::FusionApply{2})(args...) = f.f(f.args[1](args...)..., f.args[2](args...)...)
+@inline (f::FusionApply{3})(args...) = f.f(f.args[1](args...)..., f.args[2](args...)..., f.args[3](args...)...)
+@generated function (f::FusionApply{N})(args...) where {N}
+    fargs = Any[ :(getfield(f.args, $i)(args...)) for i in 1:N ]
+    return Expr(:call, GlobalRef(Core, :_apply), :(f.f), fargs...)
+end
+@inline function (f::FusionKWCall)(args...)
+    fargs = map(a -> a(args...), f.args)
+    # return f.f(args...; kwargs...)
+    if isempty(f.kwargs)
+        return f.f(fargs...)
+    else
+        return Core.kwfunc(f.f)(f.kwargs, f.f, fargs...)
+    end
+end
+@inline function (f::FusionKWApply)(args...)
+    fargs = map(a -> a(args...), f.args)
+    # return Core._apply(f.f, args...; kwargs...)
+    if isempty(f.kwargs)
+        return Core._apply(f.f, fargs...)
+    else
+        return Core._apply(Core.kwfunc(f.f), (f.kwargs,), (f.f,), fargs...)
+    end
+end
+
+function Base.show(io::IO, f::Fusion{N, vararg}) where {N, vararg}
+    nargs = (vararg ? N + 1 : N)
+    names = String[ "a_$i" for i in 1:nargs ] # f.names
+    print(io, "(")
+    join(io, names, ", ")
+    vararg && print(io, "...")
+    print(io, ") -> ")
+    show_fusion(io, f.f, names)
+end
+
+function show_fusion(io::IO, f::FusionArg{N}, names) where N
+    print(io, names[N])
+    nothing
+end
+
+function show_fusion(io::IO, f::FusionConstant{N}, names) where N
+    print(io, f.c)
+    nothing
+end
+
+function show_fusion(io::IO, f::FusionCall, names)
+    Base.show(io, f.f)
+    print(io, '(')
+    first = true
+    for i in f.args
+        first || print(io, ", ")
+        first = false
+        show_fusion(io, i, names)
+    end
+    print(io, ')')
+    nothing
+end
+
+function show_fusion(io::IO, f::FusionApply, names)
+    print(io, "Core._apply(")
+    Base.show(io, f.f)
+    for i in f.args
+        print(io, ", ")
+        show_fusion(io, i, names)
+    end
+    print(io, ')')
+    nothing
+end
+
+function show_fusion(io::IO, f::FusionKWCall, names)
+    Base.show(io, f.f)
+    print(io, '(')
+    first = true
+    for i in f.args
+        first || print(io, ", ")
+        first = false
+        show_fusion(io, i, names)
+    end
+    print(io, "; ")
+    first = true
+    for i in 1:2:length(f.kwargs)
+        first || print(io, ", ")
+        first = false
+        print(io, f.kwargs[i])
+        print(io, "=")
+    end
+    print(io, ')')
+    nothing
+end
+
+
+function show_fusion(io::IO, f::FusionKWApply, names)
+    print(io, "Core._apply(")
+    Base.show(io, f.f)
+    for i in f.args
+        print(io, ", ")
+        show_fusion(io, i, names)
+    end
+    print(io, "; #=kwargs=#...)")
+    nothing
+end
+
+
+function show_fusion(io::IO, @nospecialize(f), names)
+    print(io, "#= unexpected expression ")
+    show(io, f)
+    print(io, " =#")
+    nothing
+end
+
 end # module
diff --git a/base/inference.jl b/base/inference.jl
@@ -4,7 +4,7 @@ import Core: _apply, svec, apply_type, Builtin, IntrinsicFunction, MethodInstanc
 
 #### parameters limiting potentially-infinite types ####
 const MAX_TYPEUNION_LEN = 3
-const MAX_TYPE_DEPTH = 8
+const MAX_TYPE_DEPTH = 10
 const TUPLE_COMPLEXITY_LIMIT_DEPTH = 3
 
 const MAX_INLINE_CONST_SIZE = 256