Improve support for AbstractVector{<:Bool} and BitVecOrMat, resolves #91. Add an optimization to favor unrolling in direction of static offsets.

Author: chriselrod

Project.toml

@@ -13,14 +13,15 @@ VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
 OffsetArrays = "1"
-SIMDPirates = "0.7.11"
+SIMDPirates = "0.7.12"
 SLEEFPirates = "0.4.4"
 UnPack = "0"
-VectorizationBase = "0.10"
+VectorizationBase = "0.10.1"
 julia = "1.1"
 
 [extras]
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Random", "Test"]

src/broadcast.jl

@@ -244,6 +244,7 @@ end
     q = lower(ls)
     push!(q.args, :dest)
     pushfirst!(q.args, Expr(:meta,:inline))
+    # @show q
     q
     # ls
 end

src/condense_loopset.jl

@@ -35,7 +35,7 @@ function ArrayRefStruct(ls::LoopSet, mref::ArrayReferenceMeta, arraysymbolinds::
         index_types <<= 8
         indices <<= 8
         offsets <<= 8
-        offsets |= offv[n]
+        offsets |= (offv[n] % UInt8)
         if mref.loopedindex[n]
             index_types |= LoopIndex
             indices |= getloopid(ls, ind)

src/costs.jl

@@ -100,6 +100,7 @@ const OPAQUE_INSTRUCTION = InstructionCost(50, 50.0, -1.0, VectorizationBase.REG
 #    hand, should indicate how many registers we're keeping live for the sake of eventually storing.
 const COST = Dict{Instruction,InstructionCost}(
     Instruction(:getindex) => InstructionCost(-3.0,0.5,3,1),
+    Instruction(:conditionalload) => InstructionCost(-3.0,0.5,3,1),
     Instruction(:setindex!) => InstructionCost(-3.0,1.0,3,0),
     Instruction(:conditionalstore!) => InstructionCost(-3.0,1.0,3,0),
     Instruction(:zero) => InstructionCost(1,0.5),

src/determinestrategy.jl

@@ -458,7 +458,53 @@ function isoptranslation(ls::LoopSet, op::Operation, u1::Symbol, u2::Symbol, vec
     end
     istranslation, translationplus
 end
-function convolution_cost_factor(ls::LoopSet, op::Operation, u1::Symbol, u2::Symbol, v::Symbol)
+function maxnegativeoffset(ls::LoopSet, op::Operation, u::Symbol)
+    opmref = op.ref
+    opref = opmref.ref
+    mno = typemin(Int)
+    id = 0
+    for opp ∈ operations(ls)
+        opp === op && continue
+        oppmref = opp.ref
+        oppref = oppmref.ref
+        sameref(opref, oppref) || continue
+        opinds = getindicesonly(op)
+        oppinds = getindicesonly(opp)
+        opoffs = opref.offsets
+        oppoffs = oppref.offsets
+        # oploopi = opmref.loopedindex
+        # opploopi = oppmref.loopedindex
+        mnonew = typemin(Int)
+        for i ∈ eachindex(opinds)
+            if opinds[i] !== oppinds[i]
+                mnonew = 1
+                break
+            end
+            if opinds[i] === u
+                mnonew = (opoffs[i] - oppoffs[i])
+            elseif opoffs[i] != oppoffs[i]
+                mnonew = 1
+                break
+            end
+        end
+        if mno < mnonew < 0
+            mno = mnonew
+            id = identifier(opp)
+        end
+    end
+    mno, id
+end
+function maxnegativeoffset(ls::LoopSet, op::Operation, u1::Symbol, u2::Symbol, v::Symbol)
+    mno = typemin(Int)
+    if u1 !== v
+        mno = first(maxnegativeoffset(ls, op, u1))
+    end
+    if u2 !== v
+        mno = max(mno, first(maxnegativeoffset(ls, op, u2)))
+    end
+    mno
+end
+function loadelimination_cost_factor(ls::LoopSet, op::Operation, u1::Symbol, u2::Symbol, v::Symbol)
     if first(isoptranslation(ls, op, u1, u2, v))
         for loop ∈ ls.loops
             # If another loop is short, assume that LLVM will unroll it, in which case
@@ -473,7 +519,12 @@ function convolution_cost_factor(ls::LoopSet, op::Operation, u1::Symbol, u2::Sym
         end
         (0.25, VectorizationBase.REGISTER_COUNT == 32 ? 1.2 : 1.0)
     else
-        (1.0, 1.0)
+        offset = maxnegativeoffset(ls, op, u1, u2, v)
+        if -5 < offset < 0
+            (-0.25offset, 1.0)
+        else
+            (1.0, 1.0)
+        end
     end
 end
 # Just tile outer two loops?
@@ -542,7 +593,7 @@ function evaluate_cost_tile(
         rt, lat, rp = cost(ls, op, vectorized, Wshift, size_T)
         # @show op rt, lat, rp
         if isload(op)
-            factor1, factor2 = convolution_cost_factor(ls, op, unrolled, tiled, vectorized)
+            factor1, factor2 = loadelimination_cost_factor(ls, op, unrolled, tiled, vectorized)
             rt *= factor1; rp *= factor2;
         end
         # @show isunrolled, istiled, op rt, lat, rp

src/lower_load.jl

@@ -16,8 +16,8 @@ function lower_load_scalar!(
         end
     else
         condop = last(parents(op))
-        condvar = variable_name(condop, suffix)
-        condunrolled = any(isequal(unrolled), loopdependencies(condop))
+        condvar = tiled ∈ loopdependencies(condop) ? variable_name(condop, suffix) : variable_name(condop, nothing)
+        condunrolled = unrolled ∈ loopdependencies(condop)
         for u ∈ umin:U-1
             condsym = condunrolled ? Symbol(condvar, u) : condvar
             varname = varassignname(var, u, isunrolled)
@@ -32,7 +32,7 @@ end
 function pushvectorload!(
     q::Expr, op::Operation, var::Symbol, td::UnrollArgs, U::Int, W::Symbol, vectorized::Symbol, mask
 )
-    @unpack u, unrolled, suffix = td
+    @unpack u, unrolled, tiled, suffix = td
     ptr = refname(op)
     vecnotunrolled = vectorized !== unrolled
     name, mo = name_memoffset(var, op, td, W, vecnotunrolled)
@@ -43,8 +43,8 @@ function pushvectorload!(
     if iscondstore
         condop = last(parents(op))
         # @show condop
-        condsym = variable_name(condop, suffix)
-        condsym = any(isequal(unrolled), loopdependencies(condop)) ? Symbol(condsym, u) : condsym        
+        condsym = tiled ∈ loopdependencies(condop) ? variable_name(condop, suffix) : variable_name(condop, nothing)
+        condsym = unrolled ∈ loopdependencies(condop) ? Symbol(condsym, u) : condsym        
         if vectorized ∈ loopdependencies(condop)
             if maskend
                 push!(instrcall.args, Expr(:call, :&, condsym, mask))
@@ -99,6 +99,23 @@ function lower_load!(
                 push!(q.args, Expr(:(=), Symbol(varnew, u), Symbol(varold, u + 1)))
             end
             umin = U - 1
+        elseif tiled !== vectorized
+            mno, id = maxnegativeoffset(ls, op, tiled)
+            if -suffix < mno < 0
+                varnew = variable_name(op, suffix)
+                varold = variable_name(operations(ls)[id], suffix + mno)
+                opold = operations(ls)[id]
+                if unrolled ∈ loopdependencies(op)
+                    for u ∈ 0:U-1
+                        push!(q.args, Expr(:(=), Symbol(varnew, u), Symbol(varold, u)))
+                    end
+                else
+                    push!(q.args, Expr(:(=), varnew, varold))
+                end
+                return
+            else
+                umin = 0
+            end
         else
             umin = 0
         end

src/lower_memory_common.jl

@@ -23,19 +23,23 @@ function symbolind(ind::Symbol, op::Operation, td::UnrollArgs)
     end
     unrolled ∈ loopdependencies(parent) ? Symbol(pvar, u) : pvar
 end
+
+addoffset(ex, offset::Integer) = iszero(offset) ? ex : Expr(:call, :+, ex, convert(Int, offset))
+
 function mem_offset(op::Operation, td::UnrollArgs)
     # @assert accesses_memory(op) "Computing memory offset only makes sense for operations that access memory."
     ret = Expr(:tuple)
-    indices = getindices(op)
+    indices = getindicesonly(op)
+    offsets = getoffsets(op)
     loopedindex = op.ref.loopedindex
-    start = (first(indices) === Symbol("##DISCONTIGUOUSSUBARRAY##")) + 1
-    for (n,ind) ∈ enumerate(@view(indices[start:end]))
-        if ind isa Int
-            push!(ret.args, ind)
+    for (n,ind) ∈ enumerate(indices)
+        offset = offsets[n]
+        if ind isa Int # impossible
+            push!(ret.args, ind + offset)
         elseif loopedindex[n]
-            push!(ret.args, ind)
+            push!(ret.args, addoffset(ind, offset))
         else
-            push!(ret.args, symbolind(ind, op, td))
+            push!(ret.args, addoffset(symbolind(ind, op, td), offset))
         end
     end
     ret
@@ -45,22 +49,23 @@ function mem_offset_u(op::Operation, td::UnrollArgs)
     @unpack unrolled, u = td
     incr = u
     ret = Expr(:tuple)
-    indices = getindices(op)
+    indices = getindicesonly(op)
+    offsets = getoffsets(op)
     loopedindex = op.ref.loopedindex
     if incr == 0
         return mem_offset(op, td)
         # append_inds!(ret, indices, loopedindex)
     else
-        start = (first(indices) === Symbol("##DISCONTIGUOUSSUBARRAY##")) + 1
-        for (n,ind) ∈ enumerate(@view(indices[start:end]))
+        for (n,ind) ∈ enumerate(indices)
+            offset = offsets[n]
             if ind isa Int
                 push!(ret.args, ind)
             elseif ind === unrolled
-                push!(ret.args, Expr(:call, :+, ind, incr))
+                push!(ret.args, Expr(:call, :+, ind, incr + offset))
             elseif loopedindex[n]
-                push!(ret.args, ind)
+                push!(ret.args, addoffset(ind, offset))
             else
-                push!(ret.args, symbolind(ind, op, td))
+                push!(ret.args, addoffset(symbolind(ind, op, td), offset))
             end
         end
     end
@@ -71,22 +76,35 @@ function mem_offset_u(op::Operation, td::UnrollArgs, mul::Symbol)
     @unpack unrolled, u = td
     incr = u
     ret = Expr(:tuple)
-    indices = getindices(op)
+    indices = getindicesonly(op)
+    offsets = getoffsets(op)
     loopedindex = op.ref.loopedindex
     if incr == 0
         return mem_offset(op, td)
         # append_inds!(ret, indices, loopedindex)
     else
-        start = (first(indices) === Symbol("##DISCONTIGUOUSSUBARRAY##")) + 1
-        for (n,ind) ∈ enumerate(@view(indices[start:end]))
-            if ind isa Int
-                push!(ret.args, ind)
+        for (n,ind) ∈ enumerate(indices)
+            offset = offsets[n]
+            if ind isa Int # impossible
+                push!(ret.args, ind + offset)
             elseif ind === unrolled
-                push!(ret.args, Expr(:call, :+, ind, Expr(:call, lv(:valmul), mul, incr)))
+                if isone(incr)
+                    if iszero(offset)
+                        push!(ret.args, Expr(:call, lv(:valadd), mul, ind))
+                    else
+                        push!(ret.args, Expr(:call, :+, ind, Expr(:call, lv(:valadd), mul, convert(Int, offset))))
+                    end
+                else
+                    if iszero(offset)
+                        push!(ret.args, Expr(:call, lv(:valmuladd), mul, incr, ind))
+                    else
+                        push!(ret.args, Expr(:call, :+, ind, Expr(:call, lv(:valmuladd), mul, incr, convert(Int, offset))))
+                    end
+                end
             elseif loopedindex[n]
-                push!(ret.args, ind)
+                push!(ret.args, addoffset(ind, offset))
             else
-                push!(ret.args, symbolind(ind, op, td))
+                push!(ret.args, addoffset(symbolind(ind, op, td), offset))
             end
         end
     end

src/memory_ops_common.jl

@@ -65,11 +65,51 @@ function subset_vptr!(ls::LoopSet, vptr::Symbol, indnum::Int, ind, previndices,
     pushpreamble!(ls, Expr(:(=), subsetvptr, Expr(:call, lv(:subsetview), vptr, valcall, indm1)))
     subsetvptr
 end
+
+function addoffset!(ls, indices, offsets, loopedindex, loopdependencies, ind, offset)
+    if typemin(Int8) ≤ offset ≤ typemax(Int8)
+        push!(indices, ind);
+        push!(offsets, offset % Int8)
+        push!(loopedindex, true)
+        push!(loopdependencies, ind)
+        true
+    else
+        false
+    end
+end
+
+function checkforoffset!(
+    ls::LoopSet, indices::Vector{Symbol}, offsets::Vector{Int8}, loopedindex::Vector{Bool}, loopdependencies::Vector{Symbol}, ind::Expr
+)
+    ind.head === :call || return false
+    f = first(ind.args)
+    (((f === :+) || (f === :-)) && (length(ind.args) == 3)) || return false
+    factor = f === :+ ? 1 : -1
+    arg1 = ind.args[2]
+    arg2 = ind.args[3]
+    if arg1 isa Integer
+        if arg2 isa Symbol && arg2 ∈ ls.loopsymbols
+            addoffset!(ls, indices, offsets, loopedindex, loopdependencies, arg2, arg1 * factor)
+        else
+            false
+        end
+    elseif arg2 isa Integer
+        if arg1 isa Symbol && arg1 ∈ ls.loopsymbols
+            addoffset!(ls, indices, offsets, loopedindex, loopdependencies, arg1, arg2 * factor)
+        else
+            false
+        end
+    else
+        false
+    end        
+end
+
 const DISCONTIGUOUS = Symbol("##DISCONTIGUOUSSUBARRAY##")
 function array_reference_meta!(ls::LoopSet, array::Symbol, rawindices, elementbytes::Int, var::Union{Nothing,Symbol} = nothing)
     vptrarray = vptr(array)
     add_vptr!(ls, array, vptrarray) # now, subset
     indices = Symbol[]
+    offsets = Int8[]
     loopedindex = Bool[]
     parents = Operation[]
     loopdependencies = Symbol[]
@@ -82,21 +122,26 @@ function array_reference_meta!(ls::LoopSet, array::Symbol, rawindices, elementby
             length(indices) == 0 && push!(indices, DISCONTIGUOUS)
         elseif ind isa Expr
             #FIXME: position (in loopnest) wont be length(ls.loopsymbols) in general
-            parent = add_operation!(ls, gensym(:indexpr), ind, elementbytes, length(ls.loopsymbols))
-            pushparent!(parents, loopdependencies, reduceddeps, parent)
-            # var = get(ls.opdict, ind, nothing)
-            push!(indices, name(parent)); ninds += 1
-            push!(loopedindex, false)
+            if !checkforoffset!(ls, indices, offsets, loopedindex, loopdependencies, ind)
+                parent = add_operation!(ls, gensym(:indexpr), ind, elementbytes, length(ls.loopsymbols))
+                pushparent!(parents, loopdependencies, reduceddeps, parent)
+                push!(indices, name(parent)); 
+                push!(offsets, zero(Int8))
+                push!(loopedindex, false)
+            end
+            ninds += 1
         elseif ind isa Symbol
             if ind ∈ loopset
                 push!(indices, ind); ninds += 1
+                push!(offsets, zero(Int8))
                 push!(loopedindex, true)
                 push!(loopdependencies, ind)
             else
                 indop = get(ls.opdict, ind, nothing)
                 if indop !== nothing  && !isconstant(indop)
                     pushparent!(parents, loopdependencies, reduceddeps, indop)
                     push!(indices, name(indop)); ninds += 1
+                    push!(offsets, zero(Int8))
                     push!(loopedindex, false)
                 else
                     vptrarray = subset_vptr!(ls, vptrarray, ninds, ind, indices, loopedindex)
@@ -108,7 +153,7 @@ function array_reference_meta!(ls::LoopSet, array::Symbol, rawindices, elementby
         end
     end
     # (length(parents) != 0 && first(indices) !== Symbol("##DISCONTIGUOUSSUBARRAY##")) && pushfirst!(indices, Symbol("##DISCONTIGUOUSSUBARRAY##"))
-    mref = ArrayReferenceMeta(ArrayReference( array, indices ), loopedindex, vptrarray)
+    mref = ArrayReferenceMeta(ArrayReference( array, indices, offsets ), loopedindex, vptrarray)
     ArrayReferenceMetaPosition(mref, parents, loopdependencies, reduceddeps, isnothing(var) ? Symbol("") : var )
 end
 function tryrefconvert(ls::LoopSet, ex::Expr, elementbytes::Int, var::Union{Nothing,Symbol} = nothing)::Tuple{Bool,ArrayReferenceMetaPosition}