Fix for tiled-only operations, adding broadcast support.

Author: chriselrod

Manifest.toml

@@ -61,15 +61,13 @@ uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [[SIMDPirates]]
 deps = ["MacroTools", "VectorizationBase"]
-git-tree-sha1 = "3e45c76dfcc349ff208a955e1ce6e92b1be6d15e"
-repo-rev = "master"
-repo-url = "https://github.com/chriselrod/SIMDPirates.jl"
+git-tree-sha1 = "296cae2ccd6e4766aad669e748c1248fb99ab69c"
 uuid = "21efa798-c60a-11e8-04d3-e1a92915a26a"
 version = "0.1.0"
 
 [[SLEEFPirates]]
 deps = ["SIMDPirates", "VectorizationBase"]
-git-tree-sha1 = "ba032bbcc7038853867119f4cac383a0051b62a8"
+git-tree-sha1 = "01ff5ddb2fe743e93a6d80b072a02cceb90592bf"
 repo-rev = "master"
 repo-url = "https://github.com/chriselrod/SLEEFPirates.jl"
 uuid = "476501e8-09a2-5ece-8869-fb82de89a1fa"
@@ -87,8 +85,6 @@ uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [[VectorizationBase]]
 deps = ["CpuId", "LinearAlgebra"]
-git-tree-sha1 = "1cc48a9bce5c18f2f70fa16cc5b2b39b39332a9e"
-repo-rev = "master"
-repo-url = "https://github.com/chriselrod/VectorizationBase.jl"
+git-tree-sha1 = "30dd7fd08829bfa0fa6c57bf84a7daeac2e9462b"
 uuid = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
-version = "0.1.0"
+version = "0.1.1"

Project.toml

@@ -4,12 +4,21 @@ authors = ["Chris Elrod <[email protected]>"]
 version = "0.1.0"
 
 [deps]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 SIMDPirates = "21efa798-c60a-11e8-04d3-e1a92915a26a"
 SLEEFPirates = "476501e8-09a2-5ece-8869-fb82de89a1fa"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
+[compat]
+MacroTools = "0.5"
+Parameters = "0.12.0"
+SIMDPirates = "0.1.0"
+SLEEFPirates = "0.1.0"
+VectorizationBase = "0.1.2"
+julia = "1.0.0"
+
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 

src/LoopVectorization.jl

@@ -3,9 +3,13 @@ module LoopVectorization
 using VectorizationBase, SIMDPirates, SLEEFPirates, MacroTools, Parameters
 using VectorizationBase: REGISTER_SIZE, extract_data, num_vector_load_expr, mask
 using SIMDPirates: VECTOR_SYMBOLS, evadd, evmul, vrange, reduced_add, reduced_prod
+using Base.Broadcast: Broadcasted, DefaultArrayStyle
+using LinearAlgebra: Adjoint
 using MacroTools: prewalk, postwalk
 
-export vectorizable, @vectorize, @vvectorize, @avx
+
+export LowDimArray, stridedpointer, vectorizable,
+    @vectorize, @vvectorize, @avx, ∗
 
 function isdense end #
 
@@ -896,6 +900,7 @@ end
 include("costs.jl")
 include("operations.jl")
 include("graphs.jl")
+include("broadcast.jl")
 include("determinestrategy.jl")
 include("lowering.jl")
 include("constructors.jl")

src/broadcast.jl

@@ -0,0 +1,129 @@
+struct Product{A,B}
+    a::A
+    b::B
+end
+
+@inline ∗(a::A, b::B) where {A,B} = Product{A,B}(a, b)
+@inline Base.Broadcast.Broadcasted(::typeof(∗), a::A, b::B) where {A, B} = Product{A,B}(a, b)
+# 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 = 8
+) where {T,A,B}
+    K = gensym(:K)
+    mA = gensym(:Aₘₖ)
+    mB = gensym(:Bₖₙ)
+    pushpreamble!(ls, Expr(:(=), mA, Expr(:(.), bcname, QuoteNode(:a))))
+    pushpreamble!(ls, Expr(:(=), mB, Expr(:(.), bcname, QuoteNode(:b))))
+    pushpreamble!(ls, Expr(:(=), K, Expr(:call, :size, mB, 1)))
+
+    k = gensym(:k)
+    ls.loops[k] = Loop(k, K)
+    m = loopsyms[1]; n = loopsyms[2];
+    # load A
+    # loadA = add_load!(ls, gensym(:A), productref(A, mA, m, k), elementbytes)
+    loadA = add_broadcast!(ls, gensym(:A), mA, [m,k], A, elementbytes)
+    # load B
+    loadB = add_broadcast!(ls, gensym(:B), mB, [k,n], B, elementbytes)
+    # set Cₘₙ = 0
+    setC = add_constant!(ls, 0.0, Symbol[m, k], mC, elementbytes)
+    # compute Cₘₙ += Aₘₖ * Bₖₙ
+    reductop = Operation(
+        ls, mC, elementbytes, :vmuladd, compute, Symbol[m, k, n], Symbol[k], Operation[loadA, loadB, setC]
+    )
+    pushop!(ls, reductop, mC)    
+end
+
+struct LowDimArray{D,T,N,A<:DenseArray{T,N}} <: DenseArray{T,N}
+    data::A
+end
+@inline Base.pointer(A::LowDimArray) = pointer(A)
+function LowDimArray{D}(data::A) where {D,T,N,A <: AbstractArray{T,N}}
+    LowDimArray{D,T,N,A}(data)
+end
+function add_broadcast!(
+    ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
+    ::Type{<:LowDimArray{D,T,N}}, elementbytes::Int = 8
+) where {D,T,N}
+    fulldims = Union{Symbol,Int}[loopsyms[n] for n ∈ 1:N if D[n]]
+    ref = ArrayReference(bcname, fulldims, Ref{Bool}(false))
+    add_load!(ls, destname, ref, elementbytes)::Operation
+end
+function add_broadcast!(
+    ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
+    ::Type{Adjoint{T,A}}, elementbytes::Int = 8
+) where {T, N, A <: AbstractArray{T,N}}
+    ref = ArrayReference(bcname, Union{Symbol,Int}[loopsyms[N + 1 - n] for n ∈ 1:N], Ref{Bool}(false))
+    add_load!( ls, destname, ref, elementbytes )::Operation
+end
+function add_broadcast!(
+    ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
+    ::Type{Adjoint{T,V}}, elementbytes::Int = 8
+) where {T, V <: AbstractVector{T}}
+    ref = ArrayReference(bcname, Union{Symbol,Int}[loopsyms[2]], Ref{Bool}(false))
+    add_load!( ls, destname, ref, elementbytes )
+end
+function add_broadcast!(
+    ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
+    ::Type{<:AbstractArray{T,N}}, elementbytes::Int = 8
+) where {T,N}
+    add_load!(ls, destname, ArrayReference(bcname, @view(loopsyms[1:N]), Ref{Bool}(false)), elementbytes)
+end
+function add_broadcast!(
+    ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
+    ::Type{Broadcasted{DefaultArrayStyle{N},Nothing,F,A}},
+    elementbytes::Int = 8
+) where {N,F,A}
+    instr = get(FUNCTIONSYMBOLS, F) do
+        f = gensym(:f)
+        pushpreamble!(ls, Expr(:(=), f, Expr(:(.), bcname, QuoteNode(:f))))
+        f
+    end
+    args = A.parameters
+    Nargs = length(args)
+    bcargs = Expr(:(.), bcname, QuoteNode(:args))
+    # this is the var name in the loop
+    parents = Operation[]
+    deps = Symbol[]
+    reduceddeps = Symbol[]
+    for (i,arg) ∈ enumerate(args)
+        argname = gensym(:arg)
+        pushpreamble!(ls, Expr(:(=), argname, Expr(:ref, bcargs, i)))
+        # dynamic dispatch
+        parent = add_broadcast!(ls, gensym(:temp), argname, loopsyms, arg)::Operation
+        pushparent!(parents, deps, reduceddeps, parent)
+    end
+    op = Operation(
+        length(operations(ls)), destname, elementbytes, instr, compute, deps, reduceddeps, parents
+    )
+    pushop!(ls, op, destname)
+end
+
+# size of dest determines loops
+# @generated
+function vmaterialize!(
+    dest::AbstractArray{T,N}, bc::BC
+# ) where {T, N, BC <: Broadcasted}
+) where {N, T, BC <: Broadcasted}
+    # we have an N dimensional loop.
+    # need to construct the LoopSet
+    loopsyms = [gensym(:n) for n ∈ 1:N]
+    ls = LoopSet()
+    sizes = Expr(:tuple,)
+    for (n,itersym) ∈ enumerate(loopsyms)
+        Nsym = gensym(:N)
+        ls.loops[itersym] = Loop(itersym, Nsym)
+        push!(sizes.args, Nsym)
+    end
+    pushpreamble!(ls, Expr(:(=), sizes, Expr(:call, :size, :dest)))
+    add_broadcast!(ls, :dest, :bc, loopsyms, BC)
+    add_store!(ls, :dest, ArrayReference(:dest, loopsyms, Ref{Bool}(false)))
+    resize!(ls.loop_order, num_loops(ls)) # num_loops may be greater than N, eg Product
+    # lower(ls)
+    ls
+end
+
+function vmaterialize(bc::Broadcasted)
+    ElType = Base.Broadcast.combine_eltypes(bc.f, bc.args)
+    vmaterialize!(similar(bc, ElType), bc)
+end

src/constructors.jl

@@ -6,6 +6,40 @@ function Base.copyto!(ls::LoopSet, q::Expr)
     add_loop!(ls, q)
 end
 
+function add_ci_call!(q::Expr, f, args, syms, i)
+    call = Expr(:call, f)
+    for arg ∈ @view(args[2:end])
+        if arg isa Core.SSAValue
+            push!(call.args, syms[arg.id])
+        else
+            push!(call.args, arg)
+        end
+    end
+    push!(q.args, Expr(:(=), syms[i], call))
+end
+
+function substitute_broadcast(q::Expr)
+    ci = first(Meta.lower(LoopVectorization, q).args).code
+    nargs = length(ci)-1
+    ex = Expr(:block,)
+    syms = [gensym() for _ ∈ 1:nargs]
+    for n ∈ 1:nargs
+        ciₙ = ci[n]
+        ciₙargs = ciₙ.args
+        f = first(ciₙargs)
+        if ciₙ.head === :(=)
+            push!(ex.args, Expr(:(=), f, syms[((ciₙargs[2])::Core.SSAValue).id]))
+        elseif f === GlobalRef(Base, :materialize!)
+            add_ci_call!(ex, lv(:vmaterialize!), ciₙargs, syms, n)
+        elseif f === GlobalRef(Base, :materialize)
+            add_ci_call!(ex, lv(:vmaterialize), ciₙargs, syms, n)
+        else
+            add_ci_call!(ex, f, ciₙargs, syms, n)
+        end
+    end
+    ex
+end
+
 function LoopSet(q::Expr)
     q = SIMDPirates.contract_pass(q)
     ls = LoopSet()
@@ -15,16 +49,13 @@ function LoopSet(q::Expr)
 end
 
 macro avx(q)
-    esc(lower(LoopSet(q)))
+    q2 = if q.head === :for
+        lower(LoopSet(q))
+    else# assume broadcast
+        substitute_broadcast(q)
+    end
+    esc(q2)
 end
 
-#=
-@generated function vmaterialize(
-    dest::AbstractArray{T,N}, bc::BC
-) where {T,N,BC <: Base.Broadcast.Broadcasted}
-    # we have an N dimensional loop.
-    # need to construct the LoopSet
 
-end
-=#
 

src/costs.jl

@@ -71,6 +71,10 @@ const COST = Dict{Symbol,InstructionCost}(
     :(-) => InstructionCost(4,0.5),
     :(*) => InstructionCost(4,0.5),
     :(/) => InstructionCost(13,4.0,-2.0),
+    :vadd => InstructionCost(4,0.5),
+    :vsub => InstructionCost(4,0.5),
+    :vmul => InstructionCost(4,0.5),
+    :vdiv => InstructionCost(13,4.0,-2.0),
     :(==) => InstructionCost(1, 0.5),
     :isequal => InstructionCost(1, 0.5),
     :(&) => InstructionCost(1, 0.5),
@@ -104,6 +108,9 @@ const CORRESPONDING_REDUCTION = Dict{Symbol,Symbol}(
     :(+) => :vsum,
     :(-) => :vsum,
     :(*) => :vprod,
+    :vadd => :vsum,
+    :vsub => :vsum,
+    :vmul => :vprod,
     :(&) => :vall,
     :(|) => :vany,
     :muladd => :vsum,
@@ -194,3 +201,56 @@ function callfun(f::Symbol)
 end
 
 
+
+const FUNCTIONSYMBOLS = Dict{Type{<:Function},Symbol}(
+    typeof(+) => :(+),
+    typeof(SIMDPirates.vadd) => :(+),
+    typeof(Base.FastMath.add_fast) => :(+),
+    typeof(-) => :(-),
+    typeof(SIMDPirates.vsub) => :(-),
+    typeof(Base.FastMath.sub_fast) => :(-),
+    typeof(*) => :(*),
+    typeof(SIMDPirates.vmul) => :(*),
+    typeof(Base.FastMath.mul_fast) => :(*),
+    typeof(/) => :(/),
+    typeof(SIMDPirates.vdiv) => :(/),
+    typeof(Base.FastMath.div_fast) => :(/),
+    typeof(==) => :(==),
+    typeof(isequal) => :isequal,
+    typeof(&) => :(&),
+    typeof(|) => :(|),
+    typeof(>) => :(>),
+    typeof(<) => :(<),
+    typeof(>=) => :(>=),
+    typeof(<=) => :(<=),
+    typeof(inv) => :inv,
+    typeof(muladd) => :muladd,
+    typeof(fma) => :fma,
+    typeof(SIMDPirates.vmuladd) => :vmuladd,
+    typeof(SIMDPirates.vfma) => :vfma,
+    typeof(SIMDPirates.vfmadd) => :vfmadd,
+    typeof(SIMDPirates.vfmsub) => :vfmsub,
+    typeof(SIMDPirates.vfnmadd) => :vfnmadd,
+    typeof(SIMDPirates.vfnmsub) => :vfnmsub,
+    typeof(sqrt) => :sqrt,
+    typeof(Base.FastMath.sqrt_fast) => :sqrt,
+    typeof(SIMDPirates.vsqrt) => :sqrt,
+    typeof(log) => :log,
+    typeof(Base.FastMath.log_fast) => :log,
+    typeof(SIMDPirates.vlog) => :log,
+    typeof(SLEEFPirates.log) => :log,
+    typeof(exp) => :exp,
+    typeof(Base.FastMath.exp_fast) => :exp,
+    typeof(SIMDPirates.vexp) => :exp,
+    typeof(SLEEFPirates.exp) => :exp,
+    typeof(sin) => :sin,
+    typeof(Base.FastMath.sin_fast) => :sin,
+    typeof(SLEEFPirates.sin) => :sin,
+    typeof(cos) => :cos,
+    typeof(Base.FastMath.cos_fast) => :cos,
+    typeof(SLEEFPirates.cos) => :cos,
+    typeof(sincos) => :sincos,
+    typeof(Base.FastMath.sincos_fast) => :sincos,
+    typeof(SLEEFPirates.sincos) => :sincos
+)
+