Merge branch 'master' of https://github.com/chriselrod/LoopVectorization.jl

Author: chriselrod

src/broadcast.jl

@@ -60,8 +60,9 @@ the final operation.
 # TODO: Need to make this handle A or B being (1 or 2)-D broadcast objects.
 function add_broadcast!(
     ls::LoopSet, mC::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
-    ::Type{Product{A,B}}, elementbytes::Int
-) where {A, B}
+    @nospecialize(prod::Type{<:Product}), elementbytes::Int
+)
+    A, B = prod.parameters
     K = gensym(:K)
     mA = gensym(:Aₘₖ)
     mB = gensym(:Bₖₙ)
@@ -126,9 +127,10 @@ function LowDimArray{D}(data::A) where {D,T,N,A <: AbstractArray{T,N}}
 end
 function add_broadcast!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
-    ::Type{<:LowDimArray{D,T,N}}, elementbytes::Int
-) where {D,T,N}
-    fulldims = Symbol[loopsyms[n] for n ∈ 1:N if D[n]]
+    @nospecialize(LDA::Type{<:LowDimArray}), elementbytes::Int
+)
+    D,T,N::Int,_ = LDA.parameters
+    fulldims = Symbol[loopsyms[n] for n ∈ 1:N if D[n]::Bool]
     ref = ArrayReference(bcname, fulldims)
     add_simple_load!(ls, destname, ref, elementbytes, true, false )::Operation
 end
@@ -185,11 +187,13 @@ function add_broadcast!(
     inds[2:end] .= @view(loopsyms[1:N])
     add_simple_load!(ls, destname, ArrayReference(bcname, inds), elementbytes, true, true)
 end
+BroadcastedArray{S<:Broadcast.AbstractArrayStyle,F,A} = Broadcasted{S,Nothing,F,A}
 function add_broadcast!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
-    ::Type{Broadcasted{S,Nothing,F,A}},
+    @nospecialize(B::Type{<:BroadcastedArray}),
     elementbytes::Int
-) where {N,S<:Base.Broadcast.AbstractArrayStyle{N},F,A}
+)
+    S,_,F,A = B.parameters
     instr = get(FUNCTIONSYMBOLS, F) do
         f = gensym(:func)
         pushpreamble!(ls, Expr(:(=), f, Expr(:(.), bcname, QuoteNode(:f))))

src/constructors.jl

@@ -6,7 +6,7 @@ function Base.copyto!(ls::LoopSet, q::Expr)
     add_loop!(ls, q, 8)
 end
 
-function add_ci_call!(q::Expr, f, args, syms, i, mod = nothing)
+function add_ci_call!(q::Expr, @nospecialize(f), args, syms, i, mod = nothing)
     call = if f isa Core.SSAValue
         Expr(:call, syms[f.id])
     else

src/precompile.jl

@@ -1,3 +1,32 @@
 function _precompile_()
     ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
+
+    precompile(Tuple{Type{LoopVectorization.LoopSet},Expr})
+    precompile(Tuple{typeof(Base.mapreduce_impl),typeof(LoopVectorization.elsize),typeof(max),Array{LoopVectorization.Operation,1},Int64,Int64})
+    precompile(Tuple{typeof(LoopVectorization._avx_loopset),Core.SimpleVector,Core.SimpleVector,Core.SimpleVector,Core.SimpleVector,Core.SimpleVector,Any})
+    precompile(Tuple{typeof(LoopVectorization.add_broadcast!),LoopVectorization.LoopSet,Symbol,Symbol,Array{Symbol,1},Type{Array{Bool,1}},Int64})
+    precompile(Tuple{typeof(LoopVectorization.add_ci_call!),Expr,Any,Array{Any,1},Array{Symbol,1},Int64,Symbol})
+    precompile(Tuple{typeof(LoopVectorization.add_ci_call!),Expr,Any,Array{Any,1},Array{Symbol,1},Int64})
+    precompile(Tuple{typeof(LoopVectorization.add_constant!),LoopVectorization.LoopSet,Float64,Array{Symbol,1},Symbol,Int64})
+    precompile(Tuple{typeof(LoopVectorization.add_parent!),Array{LoopVectorization.Operation,1},Array{Symbol,1},Array{Symbol,1},LoopVectorization.LoopSet,Int64,Int64,Int64})
+    precompile(Tuple{typeof(LoopVectorization.avx_body),LoopVectorization.LoopSet,Tuple{Int8,Int8}})
+    precompile(Tuple{typeof(LoopVectorization.avx_loopset),Array{LoopVectorization.Instruction,1},Array{LoopVectorization.OperationStruct,1},Array{LoopVectorization.ArrayRefStruct,1},Core.SimpleVector,Core.SimpleVector,Core.SimpleVector,Any})
+    precompile(Tuple{typeof(LoopVectorization.cost_vec_buf),LoopVectorization.LoopSet})
+    precompile(Tuple{typeof(LoopVectorization.evaluate_cost_tile),LoopVectorization.LoopSet,Array{Symbol,1},Symbol,Symbol,Symbol})
+    precompile(Tuple{typeof(LoopVectorization.evaluate_cost_unroll),LoopVectorization.LoopSet,Array{Symbol,1},Symbol,Float64})
+    precompile(Tuple{typeof(LoopVectorization.lower_block),LoopVectorization.LoopSet,LoopVectorization.UnrollSpecification,Int64,Nothing,Int64})
+    precompile(Tuple{typeof(LoopVectorization.lower_block),LoopVectorization.LoopSet,LoopVectorization.UnrollSpecification,Int64,Symbol,Int64})
+    precompile(Tuple{typeof(LoopVectorization.lower_compute!),Expr,LoopVectorization.Operation,Symbol,Symbol,Symbol,Symbol,Int64,Int64,Nothing,Bool})
+    precompile(Tuple{typeof(LoopVectorization.lower_compute!),Expr,LoopVectorization.Operation,Symbol,Symbol,Symbol,Symbol,Int64,Int64,Symbol,Bool})
+    precompile(Tuple{typeof(LoopVectorization.lower_compute!),Expr,LoopVectorization.Operation,Symbol,Symbol,Symbol,Symbol,Int64,Nothing,Nothing,Bool})
+    precompile(Tuple{typeof(LoopVectorization.lower_compute!),Expr,LoopVectorization.Operation,Symbol,Symbol,Symbol,Symbol,Int64,Nothing,Symbol,Bool})
+    precompile(Tuple{typeof(LoopVectorization.lower_load!),Expr,LoopVectorization.Operation,Symbol,LoopVectorization.LoopSet,Symbol,Symbol,Int64,Int64,Nothing})
+    precompile(Tuple{typeof(LoopVectorization.lower_load_scalar!),Expr,LoopVectorization.Operation,Symbol,Symbol,Symbol,Symbol,Int64,Nothing,Int64})
+    precompile(Tuple{typeof(LoopVectorization.reg_pres_buf),LoopVectorization.LoopSet})
+    precompile(Tuple{typeof(LoopVectorization.setup_call),LoopVectorization.LoopSet})
+    precompile(Tuple{typeof(LoopVectorization.solve_tilesize),SubArray{Float64,1,Array{Float64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true},SubArray{Float64,1,Array{Float64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true},Int64,Int64})
+    precompile(Tuple{typeof(LoopVectorization.substitute_broadcast),Expr,Symbol})
+    precompile(Tuple{typeof(LoopVectorization.vmap_quote),Int64,Type{Float32}})
+    precompile(Tuple{typeof(println),Base.GenericIOBuffer{Array{UInt8,1}},Array{LoopVectorization.Operation,1}})
+    precompile(Tuple{typeof(resize!),LoopVectorization.LoopOrder,Int64})
 end

src/reconstruct_loopset.jl

@@ -45,7 +45,8 @@ function add_loops!(ls::LoopSet, LPSYM, LB)
         end
     end
 end
-function add_loops!(ls::LoopSet, i::Int, sym::Symbol, l::Type{CartesianIndices{N,T}}) where {N,T}
+function add_loops!(ls::LoopSet, i::Int, sym::Symbol, @nospecialize(l::Type{<:CartesianIndices}))
+    N, T = l.parameters
     ssym = String(sym)
     for k = N:-1:1
         axisexpr = Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), Expr(:ref, Expr(:., Expr(:ref, :lb, i), QuoteNode(:indices)), k))
@@ -375,7 +376,7 @@ function sizeofeltypes(v, num_arrays)::Int
     sizeof(T)
 end
 
-function avx_loopset(instr, ops, arf, AM, LPSYM, LB, vargs)
+function avx_loopset(instr, ops, arf, AM, LPSYM, LB, @nospecialize(vargs))
     ls = LoopSet(:LoopVectorization)
     num_arrays = length(arf)
     elementbytes = sizeofeltypes(vargs, num_arrays)
@@ -407,7 +408,7 @@ function _avx_loopset_debug(::Type{OPS}, ::Type{ARF}, ::Type{AM}, ::Type{LPSYM},
     @show OPS ARF AM LPSYM LB vargs
     _avx_loopset(OPS.parameters, ARF.parameters, AM.parameters, LPSYM.parameters, LB.parameters, typeof.(vargs))
 end
-function _avx_loopset(OPSsv, ARFsv, AMsv, LPSYMsv, LBsv, vargs)
+function _avx_loopset(OPSsv, ARFsv, AMsv, LPSYMsv, LBsv, @nospecialize(vargs))
     nops = length(OPSsv) ÷ 3
     instr = Instruction[Instruction(OPSsv[3i+1], OPSsv[3i+2]) for i ∈ 0:nops-1]
     ops = OperationStruct[ OPSsv[3i] for i ∈ 1:nops ]

utils/generate_precompiles.jl

@@ -0,0 +1,15 @@
+using LoopVectorization, SnoopCompile
+pkgdir = dirname(@__DIR__)
+tinf = @snoopi tmin=0.01 include(joinpath(pkgdir, "test", "runtests.jl"))
+pc = SnoopCompile.parcel(tinf; blacklist=["vmaterialize", "vmaterialize!"])
+pcs = pc[:LoopVectorization]
+open(joinpath(pkgdir, "src", "precompile.jl"), "w") do io
+    println(io, """
+    function _precompile_()
+        ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
+    """)
+    for stmt in sort(pcs)
+        println(io, "    ", stmt)
+    end
+    println(io, "end")
+end