Add support for broadcasting arrays where one of the sizes other than the first is 1 (will still crash if the first dimension is 1; use a LowDimArray to work around this).

Author: chriselrod

Manifest.toml

@@ -61,15 +61,15 @@ uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [[SIMDPirates]]
 deps = ["MacroTools", "VectorizationBase"]
-git-tree-sha1 = "5d8212a28fd747bb5f77fe8b8f8d21b4024548d3"
+git-tree-sha1 = "bdb86180981859208e759adab67e2bdd8de55c64"
 uuid = "21efa798-c60a-11e8-04d3-e1a92915a26a"
-version = "0.3.2"
+version = "0.3.3"
 
 [[SLEEFPirates]]
 deps = ["SIMDPirates", "VectorizationBase"]
-git-tree-sha1 = "740604e0e5bb739488da2c2c25b5f3d71517a490"
+git-tree-sha1 = "b7c597af915f4425adf1ec64248a49ac23605fa5"
 uuid = "476501e8-09a2-5ece-8869-fb82de89a1fa"
-version = "0.3.1"
+version = "0.3.2"
 
 [[Serialization]]
 uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
@@ -83,6 +83,6 @@ uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [[VectorizationBase]]
 deps = ["CpuId", "LinearAlgebra"]
-git-tree-sha1 = "4c587c04a6daa03370023e8ea4c431fcae5f7371"
+git-tree-sha1 = "07785cdd42d94d3b5b13e1ba67d3f8feee4e670c"
 uuid = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
-version = "0.2.4"
+version = "0.2.5"

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.6.3"
+version = "0.6.4"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

benchmark/driver.jl

@@ -47,7 +47,7 @@ AtmulB_bench = fetch(AtmulB_future)
 AmulBt_bench = fetch(AmulBt_future)
 Atmulvb_bench = fetch(Atmulvb_future)
 
-v = 4
+v = 6
 const PICTURES = "/home/chriselrod/Pictures"
 save(joinpath(PICTURES, "bench_gemm_v$v.png"), plot(gemm_bench));
 save(joinpath(PICTURES, "bench_AtmulB_v$v.png"), plot(AtmulB_bench));
@@ -63,6 +63,9 @@ save(joinpath(PICTURES, "bench_random_access_v$v.png"), plot(randomaccess_bench)
 save(joinpath(PICTURES, "bench_AmulBt_v$v.png"), plot(AmulBt_bench));
 save(joinpath(PICTURES, "bench_Atmulvb_v$v.png"), plot(Atmulvb_bench));
 
+
+
+
 plot(gemm_bench)
 plot(AtmulB_bench)
 plot(dot_bench)

src/LoopVectorization.jl

@@ -3,7 +3,7 @@ module LoopVectorization
 using VectorizationBase, SIMDPirates, SLEEFPirates, MacroTools, Parameters
 using VectorizationBase: REGISTER_SIZE, REGISTER_COUNT, extract_data, num_vector_load_expr,
     mask, masktable, pick_vector_width_val, valmul, valrem, valmuladd, valadd, valsub, _MM,
-    maybestaticlength, maybestaticsize, staticm1, subsetview, vzero,
+    maybestaticlength, maybestaticsize, staticm1, subsetview, vzero, stridedpointer_for_broadcast,
     Static, StaticUnitRange, StaticLowerUnitRange, StaticUpperUnitRange,
     PackedStridedPointer, SparseStridedPointer, RowMajorStridedPointer, StaticStridedPointer, StaticStridedStruct
 using SIMDPirates: VECTOR_SYMBOLS, evadd, evmul, vrange, reduced_add, reduced_prod, reduce_to_add, reduce_to_prod#,

src/add_compute.jl

@@ -86,6 +86,11 @@ function update_reduction_status!(parentvec::Vector{Operation}, deps::Vector{Sym
         end
     end
 end
+function add_compute!(ls::LoopSet, op::Operation)
+    @assert iscompute(op)
+    pushop!(ls, child, name(op))
+end
+
 function add_reduction_update_parent!(
     vparents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet,
     parent::Operation, instr::Symbol, directdependency::Bool, elementbytes::Int

src/add_loads.jl

@@ -1,3 +1,18 @@
+function add_load!(ls::LoopSet, op::Operation, actualarray::Bool = true, broadcast::Bool = false)
+    @assert isload(op)
+    ref = op.ref
+    id = findfirst(r -> r == ref, ls.refs_aliasing_syms)
+    # try to CSE
+    if id === nothing
+        push!(ls.syms_aliasing_refs, name(op))
+        push!(ls.refs_aliasing_syms, ref)
+    else
+        opp = ls.opdict[ls.syms_aliasing_refs[id]] # throw an error if not found.
+        return isstore(opp) ? getop(ls, first(parents(opp))) : opp
+    end    
+    add_vptr!(ls, op.ref.ref.array, vptr(op.ref), actualarray, broadcast)
+    pushop!(ls, op, name(op))
+end
 
 function add_load!(
     ls::LoopSet, var::Symbol, array::Symbol, rawindices, elementbytes::Int
@@ -10,24 +25,13 @@ function add_load!(
 )
     length(mpref.loopdependencies) == 0 && return add_constant!(ls, var, mpref, elementbytes)
     ref = mpref.mref
-    # try to CSE
-    id = findfirst(r -> r == ref, ls.refs_aliasing_syms)
-    if id === nothing
-        push!(ls.syms_aliasing_refs, var)
-        push!(ls.refs_aliasing_syms, ref)
-    else
-        opp = getop(ls, ls.syms_aliasing_refs[id], elementbytes)
-        return isstore(opp) ? getop(ls, first(parents(opp))) : opp
-    end
-    # else, don't
     op = Operation( ls, var, elementbytes, :getindex, memload, mpref )
-    add_vptr!(ls, op)
-    pushop!(ls, op, var)
+    add_load!(ls, op, true, false)
 end
 
 # for use with broadcasting
 function add_simple_load!(
-    ls::LoopSet, var::Symbol, ref::ArrayReference, elementbytes::Int, actualarray::Bool = true
+    ls::LoopSet, var::Symbol, ref::ArrayReference, elementbytes::Int, actualarray::Bool = true, broadcast::Bool = false
 )
     loopdeps = Symbol[s for s ∈ ref.indices]
     mref = ArrayReferenceMeta(
@@ -38,7 +42,7 @@ function add_simple_load!(
         :getindex, memload, loopdeps,
         NODEPENDENCY, NOPARENTS, mref
     )
-    add_vptr!(ls, op.ref.ref.array, vptr(op.ref), actualarray)
+    add_vptr!(ls, op.ref.ref.array, vptr(op.ref), actualarray, broadcast)
     pushop!(ls, op, var)
 end
 function add_load_ref!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int)

src/add_stores.jl

@@ -8,10 +8,14 @@ function cse_store!(ls::LoopSet, op::Operation)
     ls.opdict[op.variable] = op
     op
 end
-function add_store!(ls::LoopSet, op::Operation)
-    nops = length(ls.operations)
+function add_store!(ls::LoopSet, op::Operation, add_pvar::Bool = name(first(parents(op))) ∉ ls.syms_aliasing_refs)
+    @assert isstore(op)
+    if add_pvar
+        push!(ls.syms_aliasing_refs, name(first(parents(op))))
+        push!(ls.refs_aliasing_syms, op.ref)
+    end
     id = op.identifier
-    id == nops ? add_unique_store!(ls, op) : cse_store!(ls, op)
+    id == length(operations(ls)) ? add_unique_store!(ls, op) : cse_store!(ls, op)
 end
 function add_copystore!(
     ls::LoopSet, parent::Operation, mpref::ArrayReferenceMetaPosition, elementbytes::Int
@@ -20,6 +24,7 @@ function add_copystore!(
     add_store!(ls, name(op), mpref, elementbytes, op)
 end
 
+
 function add_store!(
     ls::LoopSet, var::Symbol, mpref::ArrayReferenceMetaPosition, elementbytes::Int, parent = getop(ls, var, mpref.loopdependencies, elementbytes)
 )
@@ -29,10 +34,7 @@ function add_store!(
     reduceddeps = mpref.reduceddeps
     pvar = name(parent)
     id = length(ls.operations)
-    if pvar ∉ ls.syms_aliasing_refs
-        push!(ls.syms_aliasing_refs, pvar)
-        push!(ls.refs_aliasing_syms, mpref.mref)
-    else
+    if pvar ∈ ls.syms_aliasing_refs
         # try to cse store, by replacing the previous one
         ref = mpref.mref.ref
         for opp ∈ operations(ls)
@@ -44,10 +46,13 @@ function add_store!(
                 # @show ref opp.ref.ref
             end
         end
+        add_pvar = false
+    else
+        add_pvar = true
     end
     pushparent!(parents, ldref, reduceddeps, parent)
     op = Operation( id, name(mpref), elementbytes, :setindex!, memstore, mpref )#loopdependencies, reduceddeps, parents, mpref.mref )
-    add_store!(ls, op)
+    add_store!(ls, op, add_pvar)
 end
 
 function add_store!(

src/broadcast.jl

@@ -101,21 +101,21 @@ function add_broadcast!(
 ) where {D,T,N}
     fulldims = Symbol[loopsyms[n] for n ∈ 1:N if D[n]]
     ref = ArrayReference(bcname, fulldims)
-    add_simple_load!(ls, destname, ref, elementbytes )::Operation
+    add_simple_load!(ls, destname, ref, elementbytes, true, false )::Operation
 end
 function add_broadcast_adjoint_array!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol}, ::Type{A}, elementbytes::Int
 ) where {T,N,A<:AbstractArray{T,N}}
     parent = gensym(:parent)
     pushpreamble!(ls, Expr(:(=), parent, Expr(:call, :parent, bcname)))
     ref = ArrayReference(parent, Symbol[loopsyms[N + 1 - n] for n ∈ 1:N])
-    add_simple_load!( ls, destname, ref, elementbytes )::Operation    
+    add_simple_load!( ls, destname, ref, elementbytes, true, true )::Operation    
 end
 function add_broadcast_adjoint_array!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol}, ::Type{<:AbstractVector}, elementbytes::Int
 )
     ref = ArrayReference(bcname, Symbol[loopsyms[2]])
-    add_simple_load!( ls, destname, ref, elementbytes )
+    add_simple_load!( ls, destname, ref, elementbytes, true, true )
 end
 function add_broadcast!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
@@ -133,7 +133,7 @@ function add_broadcast!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
     ::Type{<:AbstractArray{T,N}}, elementbytes::Int
 ) where {T,N}
-    add_simple_load!(ls, destname, ArrayReference(bcname, @view(loopsyms[1:N])), elementbytes)
+    add_simple_load!(ls, destname, ArrayReference(bcname, @view(loopsyms[1:N])), elementbytes, true, true)
 end
 function add_broadcast!(
     ls::LoopSet, ::Symbol, bcname::Symbol, loopsyms::Vector{Symbol}, ::Type{T}, elementbytes::Int
@@ -147,7 +147,7 @@ function add_broadcast!(
     inds = Vector{Symbol}(undef, N+1)
     inds[1] = Symbol("##DISCONTIGUOUSSUBARRAY##")
     inds[2:end] .= @view(loopsyms[1:N])
-    add_simple_load!(ls, destname, ArrayReference(bcname, inds), elementbytes)
+    add_simple_load!(ls, destname, ArrayReference(bcname, inds), elementbytes, true, true)
 end
 function add_broadcast!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},

src/condense_loopset.jl

@@ -185,6 +185,48 @@ end
 @inline array_wrapper(A::Adjoint) = Adjoint
 @inline array_wrapper(A::SubArray) = A.indices
 
+
+
+@generated function __avx__!(
+    ::Val{UT}, ::Type{OPS}, ::Type{ARF}, ::Type{AM}, lb::LB,
+    ::Val{AR}, ::Val{D}, ::Val{IND}, subsetvals, arraydescript, vargs::Vararg{<:Any,N}
+) where {UT, OPS, ARF, AM, LB, N, AR, D, IND}
+    num_vptrs = length(ARF.parameters)::Int
+    vptrs = [gensym(:vptr) for _ ∈ 1:num_vptrs]
+    call = Expr(:call, lv(:_avx_!), Val{UT}(), OPS, ARF, AM, :lb)
+    for n ∈ 1:num_vptrs
+        push!(call.args, vptrs[n])
+    end
+    q = Expr(:block)
+    j = 0
+    assigned_names = Vector{Symbol}(undef, length(AR))
+    num_arrays = 0
+    for i ∈ eachindex(AR)
+        ari = (AR[i])::Int
+        ind = (IND[i])::Union{Nothing,Int}
+        LHS = ind === nothing ? gensym() : vptrs[ind]
+        assigned_names[i] = LHS
+        d = (D[i])::Union{Nothing,Int}
+        if d === nothing # stridedpointer
+            if ari == -1
+                RHS = Expr(:call, :LoopValue)
+            else
+                num_arrays += 1
+                RHS = Expr(:call, lv(:stridedpointer), Expr(:ref, :vargs, ari), Expr(:ref, :arraydescript, ari))
+            end
+        else #subsetview
+            j += 1
+            RHS = Expr(:call, :subsetview, assigned_names[ari], Expr(:call, Expr(:curly, :Val, d)), Expr(:ref, :subsetvals, j))
+        end
+        push!(q.args, Expr(:(=), LHS, RHS))
+    end
+    for n ∈ num_arrays+1:N
+        push!(call.args, Expr(:ref, :vargs, n))
+    end
+    push!(q.args, call)
+    Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__, @__FILE__), q)
+end
+
 # Try to condense in type stable manner
 function generate_call(ls::LoopSet, IUT)
     operation_descriptions = Expr(:curly, :Tuple)
@@ -200,20 +242,77 @@ function generate_call(ls::LoopSet, IUT)
     foreach(ref -> push!(arrayref_descriptions.args, ArrayRefStruct(ls, ref, arraysymbolinds)), ls.refs_aliasing_syms)
     argmeta = argmeta_and_consts_description(ls, arraysymbolinds)
     loop_bounds = loop_boundaries(ls)
-
-    q = Expr(:call, lv(:_avx_!), Expr(:call, Expr(:curly, :Val, IUT)), operation_descriptions, arrayref_descriptions, argmeta, loop_bounds)
-    foreach(ref -> push!(q.args, vptr(ref)), ls.refs_aliasing_syms)
+    inline, U, T = IUT
+    if inline
+        q = Expr(:call, lv(:_avx_!), Expr(:call, Expr(:curly, :Val, (U,T))), operation_descriptions, arrayref_descriptions, argmeta, loop_bounds)
+        foreach(ref -> push!(q.args, vptr(ref)), ls.refs_aliasing_syms)
+    else
+        arraydescript = Expr(:tuple)
+        q = Expr(:call, lv(:__avx__!), Expr(:call, Expr(:curly, :Val, (U,T))), operation_descriptions, arrayref_descriptions, argmeta, loop_bounds, arraydescript)
+        for array ∈ ls.includedactualarrays
+            push!(q.args, Expr(:call, lv(:unwrap_array), array))
+            push!(arraydescript.args, Expr(:call, lv(:array_wrapper), array))
+        end
+    end
     foreach(is -> push!(q.args, last(is)), ls.preamble_symsym)
     append!(q.args, arraysymbolinds)
     add_reassigned_syms!(q, ls)
     add_external_functions!(q, ls)
     q
 end
 
-function setup_call_noinline(ls::LoopSet, inline = Int8(2), U = zero(Int8), T = zero(Int8))
-    call = generate_call(ls, (inline,U,T))
+function setup_call_noinline(ls::LoopSet, U = zero(Int8), T = zero(Int8))
+    call = generate_call(ls, (false,U,T))
     hasouterreductions = length(ls.outer_reductions) > 0
-    q = ls.preamble
+    q = Expr(:block)
+    vptrarrays = Expr(:tuple)
+    vptrsubsetvals = Expr(:tuple)
+    vptrsubsetdims = Expr(:tuple)
+    vptrindices = Expr(:tuple)
+    stridedpointerLHS = Symbol[]
+    loopvalueLHS = Symbol[]
+    for ex ∈ ls.preamble.args
+        # vptrcalls = Expr(:tuple)
+        if ex isa Expr && ex.head === :(=) && length(ex.args) == 2
+            if ex.args[2] isa Expr && ex.args[2].head === :call
+                gr = first(ex.args[2].args)
+                if gr == lv(:stridedpointer)
+                    array = ex.args[2].args[2]
+                    arrayid = findfirst(a -> a === array, ls.includedactualarrays)
+                    if arrayid isa Int
+                        push!(vptrarrays.args, arrayid)
+                    else
+                        @assert array ∈ loopvalueLHS
+                        push!(vptrarrays.args, -1)
+                    end
+                    push!(vptrsubsetdims.args, nothing)
+                    vp = first(ex.args)::Symbol
+                    push!(stridedpointerLHS, vp)
+                    push!(vptrindices.args, findfirst(a -> vptr(a) == vp, ls.refs_aliasing_syms))
+                elseif gr == lv(:subsetview)
+                    array = ex.args[2].args[2]
+                    vptrarrayid = findfirst(a -> a === array, stridedpointerLHS)#::Int
+                    if vptrarrayid === nothing
+                        @show array, stridedpointerLHS
+                        @assert vptrarrayid isa Int
+                    end
+                    push!(vptrarrays.args, vptrarrayid::Int)
+                    push!(vptrsubsetdims.args, ex.args[2].args[3].args[1].args[2])
+                    push!(vptrsubsetvals.args, ex.args[2].args[4])
+                    vp = first(ex.args)::Symbol
+                    push!(stridedpointerLHS, vp)
+                    push!(vptrindices.args, findfirst(a -> vptr(a) == vp, ls.refs_aliasing_syms))
+                end
+            elseif ex.args[2] == LoopValue()
+                push!(loopvalueLHS, first(ex.args))
+            end
+        end
+        push!(q.args, ex)
+    end
+    insert!(call.args, 7, Expr(:call, Expr(:curly, :Val, vptrarrays)))
+    insert!(call.args, 8, Expr(:call, Expr(:curly, :Val, vptrsubsetdims)))
+    insert!(call.args, 9, Expr(:call, Expr(:curly, :Val, vptrindices)))
+    insert!(call.args, 10, vptrsubsetvals)
     if hasouterreductions
         outer_reducts = Expr(:local)
         for or ∈ ls.outer_reductions
@@ -227,8 +326,6 @@ function setup_call_noinline(ls::LoopSet, inline = Int8(2), U = zero(Int8), T =
         retv = loopset_return_value(ls, Val(false))
         call = Expr(:(=), retv, call)
         push!(q.args, gc_preserve(ls, call))
-        push!(q.args, Expr(:return, retv))
-        q = Expr(:block, Expr(:(=), retv, Expr(:call, Expr(:(->), Expr(:tuple, ls.includedactualarrays...), q), ls.includedactualarrays...)))
         for or ∈ ls.outer_reductions
             op = ls.operations[or]
             var = name(op)
@@ -239,13 +336,11 @@ function setup_call_noinline(ls::LoopSet, inline = Int8(2), U = zero(Int8), T =
         end
     else
         push!(q.args, gc_preserve(ls, call))
-        push!(q.args, Expr(:return, :nothing))
-        q = Expr(:call, Expr(:(->), Expr(:tuple, ls.includedactualarrays...), q), ls.includedactualarrays...)
     end
     q
 end
-function setup_call_inline(ls::LoopSet, inline = Int8(2), U = zero(Int8), T = zero(Int8))
-    call = generate_call(ls, (inline,U,T))
+function setup_call_inline(ls::LoopSet, U = zero(Int8), T = zero(Int8))
+    call = generate_call(ls, (true,U,T))
     hasouterreductions = length(ls.outer_reductions) > 0
     if hasouterreductions
         retv = loopset_return_value(ls, Val(false))
@@ -273,9 +368,9 @@ function setup_call(ls::LoopSet, inline = Int8(2), U = zero(Int8), T = zero(Int8
     # Creating an anonymous function and calling it also achieves the outlining, while still
     # inlining the generated function into the loop preamble.
     if inline == Int8(2)
-        setup_call_inline(ls, Int8(2), U, T)
+        setup_call_inline(ls, U, T)
     else
-        setup_call_noinline(ls, Int8(2), U, T)
+        setup_call_noinline(ls, U, T)
     end
 end