Merge pull request #180 from timholy/teh/inference

Some inferrability and precompile improvements

Author: chriselrod

src/add_compute.jl

@@ -57,7 +57,7 @@ function add_parent!(
         opp = getop(ls, var, elementbytes)
         # if var === :kern_1_1
         #     @show operations(ls) ls.preamble_symsym
-        # end 
+        # end
         # @show var opp first(operations(ls)) opp === first(operations(ls))
         if iscompute(opp) && instruction(opp).instr === :identity && length(loopdependencies(opp)) < position && isone(length(parents(opp))) && name(opp) === name(first(parents(opp)))
             first(parents(opp))
@@ -168,7 +168,7 @@ function add_reduction_update_parent!(
     # if !isouterreduction && !isreductzero(parent, ls, reduct_zero)
     add_reduct_instruct = !isouterreduction && !isconstant(parent)
     if add_reduct_instruct
-        # We add 
+        # We add
         reductcombine = reduction_scalar_combine(instrclass)
         # reductcombine = :identity
         reductsym = gensym(:reduction)
@@ -222,7 +222,10 @@ function add_compute!(
     # instr = instruction(first(ex.args))::Symbol
     instr = instruction!(ls, first(ex.args))::Instruction
     args = @view(ex.args[2:end])
-    (instr.instr === :(^) && length(args) == 2 && (args[2] isa Number)) && return add_pow!(ls, var, args[1], args[2], elementbytes, position)
+    if instr.instr === :(^) && length(args) == 2
+        arg2 = args[2]
+        arg2 isa Number && return add_pow!(ls, var, args[1], arg2, elementbytes, position)
+    end
     vparents = Operation[]
     deps = Symbol[]
     reduceddeps = Symbol[]
@@ -232,7 +235,7 @@ function add_compute!(
         if var === arg
             reduction_ind = ind
             # add_reduction!(vparents, deps, reduceddeps, ls, arg, elementbytes)
-            getop(ls, arg, elementbytes)
+            getop(ls, arg::Symbol, elementbytes)   # weird that this needs annotation
         elseif arg isa Expr
             isref, argref = tryrefconvert(ls, arg, elementbytes, varname(mpref))
             if isref
@@ -303,7 +306,7 @@ end
 
 # adds x ^ (p::Real)
 function add_pow!(
-    ls::LoopSet, var::Symbol, x, p::Real, elementbytes::Int, position::Int
+    ls::LoopSet, var::Symbol, @nospecialize(x), p::Real, elementbytes::Int, position::Int
 )
     xop::Operation = if x isa Expr
         add_operation!(ls, gensym(:xpow), x, elementbytes, position)

src/add_loads.jl

@@ -2,13 +2,13 @@ function maybeaddref!(ls::LoopSet, op)
     ref = op.ref
     id = findfirst(r -> r == ref, ls.refs_aliasing_syms)
     # try to CSE
-    if isnothing(id)
+    if id === nothing
         push!(ls.syms_aliasing_refs, name(op))
         push!(ls.refs_aliasing_syms, ref)
         0
     else
         id
-    end    
+    end
 end
 
 function add_load!(ls::LoopSet, op::Operation, actualarray::Bool = true, broadcast::Bool = false)

src/constructors.jl

@@ -63,7 +63,7 @@ function LoopSet(q::Expr, mod::Symbol = :Main)
     resize!(ls.loop_order, num_loops(ls))
     ls
 end
-LoopSet(q::Expr, m::Module) = LoopSet(macroexpand(m, q), Symbol(m))
+LoopSet(q::Expr, m::Module) = LoopSet(macroexpand(m, q)::Expr, Symbol(m))
 
 """
     @avx
@@ -143,6 +143,7 @@ use their `parent`. Triangular loops aren't yet supported.
 """
 macro avx(q)
     q = macroexpand(__module__, q)
+    isa(q, Expr) || return q
     q2 = if q.head === :for
         setup_call(LoopSet(q, __module__), q)
     else# assume broadcast

src/graphs.jl

@@ -136,7 +136,7 @@ function looprange(loop::Loop, incr::Int, mangledname)
 end
 function terminatecondition(
     loop::Loop, us::UnrollSpecification, n::Int, mangledname::Symbol, inclmask::Bool, UF::Int = unrollfactor(us, n)
-) 
+)
     if !isvectorized(us, n)
         looprange(loop, UF, mangledname)
     elseif inclmask
@@ -345,7 +345,7 @@ function LoopSet(mod::Symbol)
         Tuple{Int,Symbol}[],
         Tuple{Int,Int}[],
         Tuple{Int,Float64}[],
-        Int[],Int[],
+        Tuple{Int,NumberType}[],Tuple{Int,Symbol}[],
         Symbol[], Symbol[], Symbol[],
         ArrayReferenceMeta[],
         Matrix{Float64}(undef, 4, 2),
@@ -363,7 +363,7 @@ function cacheunrolled!(ls::LoopSet, u₁loop, u₂loop, vectorized)
         for opp ∈ parents(op)
             push!(children(opp), op)
         end
-    end    
+    end
 end
 
 num_loops(ls::LoopSet) = length(ls.loops)
@@ -404,7 +404,7 @@ function getop(ls::LoopSet, var::Symbol, deps, elementbytes::Int)
 end
 getop(ls::LoopSet, i::Int) = ls.operations[i]
 
-# """ 
+# """
 # Returns an operation with the same name as `s`.
 # """
 # function getoperation(ls::LoopSet, s::Symbol)
@@ -486,15 +486,16 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
     itersym = (looprange.args[1])::Symbol
     r = looprange.args[2]
     if isexpr(r, :call)
+        r = r::Expr        # julia#37342
         f = first(r.args)
         loop::Loop = if f === :(:)
             lower = r.args[2]
             upper = r.args[3]
             lii::Bool = lower isa Integer
-            liiv::Int = lii ? convert(Int, lower) : 1
+            liiv::Int = lii ? convert(Int, lower::Integer) : 1
             uii::Bool = upper isa Integer
             if lii & uii # both are integers
-                Loop(itersym, liiv, convert(Int, upper))
+                Loop(itersym, liiv, convert(Int, upper::Integer)::Int)
             elseif lii # only lower bound is an integer
                 if upper isa Symbol
                     Loop(itersym, liiv, upper)
@@ -504,7 +505,7 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
                     Loop(itersym, liiv, add_loop_bound!(ls, itersym, upper, true))
                 end
             elseif uii # only upper bound is an integer
-                uiiv = convert(Int, upper)
+                uiiv = convert(Int, upper::Integer)::Int
                 Loop(itersym, add_loop_bound!(ls, itersym, lower, false), uiiv)
             else # neither are integers
                 L = add_loop_bound!(ls, itersym, lower, false)
@@ -544,16 +545,16 @@ end
 function register_loop!(ls::LoopSet, looprange::Expr)
     if looprange.head === :block # multiple loops
         for lr ∈ looprange.args
-            register_single_loop!(ls, lr)
+            register_single_loop!(ls, lr::Expr)
         end
     else
         @assert looprange.head === :(=)
         register_single_loop!(ls, looprange)
     end
 end
 function add_loop!(ls::LoopSet, q::Expr, elementbytes::Int)
-    register_loop!(ls, q.args[1])
-    body = q.args[2]
+    register_loop!(ls, q.args[1]::Expr)
+    body = q.args[2]::Expr
     position = length(ls.loopsymbols)
     if body.head === :block
         add_block!(ls, body, elementbytes, position)
@@ -685,10 +686,8 @@ function add_operation!(
     end
 end
 
-function prepare_rhs_for_storage!(ls::LoopSet, RHS::Symbol, array, rawindices, elementbytes::Int, position::Int)
-    add_store!(ls, RHS, array, rawindices, elementbytes)
-end
-function prepare_rhs_for_storage!(ls::LoopSet, RHS::Expr, array, rawindices, elementbytes::Int, position::Int)
+function prepare_rhs_for_storage!(ls::LoopSet, RHS::Union{Symbol,Expr}, array, rawindices, elementbytes::Int, position::Int)
+    RHS isa Symbol && return add_store!(ls, RHS, array, rawindices, elementbytes)
     mpref = array_reference_meta!(ls, array, rawindices, elementbytes)
     cachedparents = copy(mpref.parents)
     ref = mpref.mref.ref
@@ -706,9 +705,9 @@ function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int, position::Int)
         finex = first(ex.args)::Symbol
         if finex === :setindex!
             array, rawindices = ref_from_setindex!(ls, ex)
-            prepare_rhs_for_storage!(ls, ex.args[3], array, rawindices, elementbytes, position)
+            prepare_rhs_for_storage!(ls, ex.args[3]::Union{Symbol,Expr}, array, rawindices, elementbytes, position)
         else
-            throw("Function $finex not recognized.")
+            error("Function $finex not recognized.")
         end
     elseif ex.head === :(=)
         LHS = ex.args[1]
@@ -729,7 +728,7 @@ function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int, position::Int)
                     array, rawindices = ref_from_expr!(ls, LHS)
                     prepare_rhs_for_storage!(ls, RHS, array, rawindices, elementbytes, position)
                 else
-                    add_store_ref!(ls, RHS, LHS, elementbytes)
+                    add_store_ref!(ls, RHS, LHS, elementbytes)  # is this necessary? (Extension API?)
                 end
             elseif LHS.head === :tuple
                 @assert length(LHS.args) ≤ 9 "Functions returning more than 9 values aren't currently supported."

src/precompile.jl

@@ -11,6 +11,8 @@ function _precompile_()
     Base.precompile(Tuple{typeof(LoopVectorization.add_compute!),LoopVectorization.LoopSet,Symbol,Expr,Int64,Int64,Nothing})
     Base.precompile(Tuple{typeof(LoopVectorization.add_constant!),LoopVectorization.LoopSet,Float64,Vector{Symbol},Symbol,Int64})
     Base.precompile(Tuple{typeof(LoopVectorization.add_if!),LoopVectorization.LoopSet,Symbol,Expr,Int64,Int64,LoopVectorization.ArrayReferenceMetaPosition})
+    Base.precompile(Tuple{typeof(avx_body), LoopSet, Tuple{Int8,Int8,Int8,Int}})
+    Base.precompile(Tuple{typeof(avx_loopset), Vector{Instruction}, Vector{OperationStruct}, Vector{ArrayRefStruct}, Core.SimpleVector, Core.SimpleVector, Core.SimpleVector, Any})
     Base.precompile(Tuple{typeof(LoopVectorization.add_mref!),LoopVectorization.LoopSet,LoopVectorization.ArrayReferenceMeta,Int64,Type{VectorizationBase.StridedBitPointer{2, 1, 0, (1, 2), Tuple{ArrayInterface.StaticInt{1}, Int64}, Tuple{ArrayInterface.StaticInt{1}, ArrayInterface.StaticInt{1}}}},Symbol})
     Base.precompile(Tuple{typeof(LoopVectorization.add_mref!),LoopVectorization.LoopSet,LoopVectorization.ArrayReferenceMeta,Int64,Type{VectorizationBase.StridedPointer{Float32, 2, 1, 0, (1, 2), Tuple{ArrayInterface.StaticInt{4}, ArrayInterface.StaticInt{292}}, Tuple{ArrayInterface.StaticInt{1}, ArrayInterface.StaticInt{1}}}},Symbol})
     Base.precompile(Tuple{typeof(LoopVectorization.add_mref!),LoopVectorization.LoopSet,LoopVectorization.ArrayReferenceMeta,Int64,Type{VectorizationBase.StridedPointer{Float32, 2, 1, 0, (1, 2), Tuple{ArrayInterface.StaticInt{4}, ArrayInterface.StaticInt{308}}, Tuple{ArrayInterface.StaticInt{1}, ArrayInterface.StaticInt{1}}}},Symbol})

src/reconstruct_loopset.jl

@@ -138,13 +138,13 @@ function add_mref!(
         return pushvarg!(ls, ar, i, name)
     end
     lic = copy(li);
-    inds = getindices(ar); indsc = copy(inds); 
+    inds = getindices(ar); indsc = copy(inds);
     offsets = ar.ref.offsets; offsetsc = copy(offsets);
 
     # must now sort array's inds, and stack pointer's
     tmpsp = gensym(name)
     pushvarg!(ls, ar, i, tmpsp)
-    # pushpreamble!(ls, 
+    # pushpreamble!(ls,
     strd_tup = Expr(:tuple)
     offsets_tup = Expr(:tuple)
     for (i, p) ∈ enumerate(sp)
@@ -417,7 +417,8 @@ function sizeofeltypes(v, num_arrays)::Int
     # sizeof(T)
 end
 
-function avx_loopset(instr, ops, arf, AM, LPSYM, LB, @nospecialize(vargs))
+function avx_loopset(instr::Vector{Instruction}, ops::Vector{OperationStruct}, arf::Vector{ArrayRefStruct},
+                     AM::Core.SimpleVector, LPSYM::Core.SimpleVector, LB::Core.SimpleVector, @nospecialize(vargs))
     ls = LoopSet(:LoopVectorization)
     num_arrays = length(arf)
     elementbytes = sizeofeltypes(vargs, num_arrays)
@@ -436,7 +437,7 @@ function avx_loopset(instr, ops, arf, AM, LPSYM, LB, @nospecialize(vargs))
     num_params = extract_external_functions!(ls, num_params)
     ls
 end
-function avx_body(ls, UNROLL)
+function avx_body(ls::LoopSet, UNROLL::Tuple{Int8,Int8,Int8,Int})
     inline, u₁, u₂, W = UNROLL
     ls.vector_width[] = W
     q = iszero(u₁) ? lower_and_split_loops(ls, inline % Int) : lower(ls, u₁ % Int, u₂ % Int, inline % Int)
@@ -448,7 +449,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, @nospecialize(vargs))
+function _avx_loopset(OPSsv::Core.SimpleVector, ARFsv::Core.SimpleVector, AMsv::Core.SimpleVector, LPSYMsv::Core.SimpleVector, LBsv::Core.SimpleVector, @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 ]

src/vectorizationbase_compat/contract_pass.jl

@@ -48,8 +48,8 @@ end
 function recursive_muladd_search!(call, argv, cnmul::Bool = false, csub::Bool = false)
     length(argv) < 3 && return length(call.args) == 4, cnmul, csub
     fun = first(argv)
-    isadd = fun === :+ || fun === :vadd! || fun === :vadd || fun == :(Base.FastMath.add_fast)
-    issub = fun === :- || fun === :vsub! || fun === :vsub || fun == :(Base.FastMath.sub_fast)
+    isadd = fun === :+ || fun === :vadd! || fun === :vadd || (fun == :(Base.FastMath.add_fast))::Bool
+    issub = fun === :- || fun === :vsub! || fun === :vsub || (fun == :(Base.FastMath.sub_fast))::Bool
     if !(isadd | issub)
         return length(call.args) == 4, cnmul, csub
     end
@@ -60,7 +60,7 @@ function recursive_muladd_search!(call, argv, cnmul::Bool = false, csub::Bool =
             exa = ex.args
             f = first(exa)
             exav = @view(exa[2:end])
-            if f === :* || f === :vmul! || f === :vmul || f == :(Base.FastMath.mul_fast)
+            if f === :* || f === :vmul! || f === :vmul || (f == :(Base.FastMath.mul_fast))::Bool
                 a, b = mulexpr(exav)
                 call.args[2] = a
                 call.args[3] = b
@@ -107,28 +107,27 @@ function recursive_muladd_search!(call, argv, cnmul::Bool = false, csub::Bool =
                         csub = false
                     end
                     return true, cnmul, csub
-                end                
+                end
             end
         end
     end
     length(call.args) == 4, cnmul, csub
 end
-                          
-function capture_muladd(ex::Expr, mod, LHS = nothing)
+
+function capture_muladd(ex::Expr, mod)
     call = Expr(:call, Symbol(""), Symbol(""), Symbol(""))
     found, nmul, sub = recursive_muladd_search!(call, ex.args)
     found || return ex
     # a, b, c = call.args[2], call.args[3], call.args[4]
     # call.args[2], call.args[3], call.args[4] = c, a, b
-    clobber = false#call.args[4] == LHS
     f = if nmul && sub
-        clobber ? :vfnmsub231 : :vfnmsub
+        :vfnmsub
     elseif nmul
-        clobber ? :vfnmadd231 : :vfnmadd
+        :vfnmadd
     elseif sub
-        clobber ? :vfmsub231 : :vfmsub
+        :vfmsub
     else
-        clobber ? :vfmadd231 : :vfmadd
+        :vfmadd
     end
     if mod === nothing
         call.args[1] = f
@@ -139,7 +138,7 @@ function capture_muladd(ex::Expr, mod, LHS = nothing)
 end
 
 contract_pass!(::Any, ::Any) = nothing
-function contract!(expr::Expr, ex::Expr, i::Int, mod = nothing)
+function contract!(expr::Expr, ex::Expr, i::Int, mod)
     # if ex.head === :call
         # expr.args[i] = capture_muladd(ex, mod)
     if ex.head === :(+=)
@@ -163,7 +162,7 @@ function contract!(expr::Expr, ex::Expr, i::Int, mod = nothing)
         RHS = ex.args[2]
         # @show ex
         if RHS isa Expr && RHS.head === :call
-            ex.args[2] = capture_muladd(RHS, mod, ex.args[1])
+            ex.args[2] = capture_muladd(RHS, mod)
         end
     end
     contract_pass!(expr.args[i], mod)