WIP: support CartesianIndex

Author: timholy

src/add_loads.jl

@@ -24,7 +24,6 @@ function add_load!(
     ls::LoopSet, var::Symbol, mpref::ArrayReferenceMetaPosition, elementbytes::Int
 )
     length(mpref.loopdependencies) == 0 && return add_constant!(ls, var, mpref, elementbytes)
-    ref = mpref.mref
     op = Operation( ls, var, elementbytes, :getindex, memload, mpref )
     add_load!(ls, op, true, false)
 end

src/condense_loopset.jl

@@ -4,6 +4,11 @@
 Base.:|(u::Unsigned, it::IndexType) = u | UInt8(it)
 Base.:(==)(u::Unsigned, it::IndexType) = (u % UInt8) == UInt8(it)
 
+"""
+`ArrayRefStruct` stores a representation of an array-reference expression such as `A[i,j]`.
+It supports array-references with up to 8 indexes, where the data for each consecutive index is packed into corresponding 8-bit fields
+of `index_types` (storing the enum `IndexType`), `indices` (the `id` for each index symbol), and `offsets` (currently unused).
+"""
 struct ArrayRefStruct
     index_types::UInt64
     indices::UInt64
@@ -392,4 +397,3 @@ function setup_call(ls::LoopSet, inline = Int8(2), U = zero(Int8), T = zero(Int8
         setup_call_noinline(ls, U, T)
     end
 end
-

src/determinestrategy.jl

@@ -21,7 +21,7 @@ function findparent(ls::LoopSet, s::Symbol)#opdict isn't filled when reconstruct
 end
 function unitstride(ls::LoopSet, op::Operation, s::Symbol)
     inds = getindices(op)
-    li = op.ref.loopedindex
+    li, lookup = op.ref.loopedindex, op.ref.indexlookup
     # The first index is allowed to be indexed by `s`
     fi = first(inds)
     if fi === Symbol("##DISCONTIGUOUSSUBARRAY##")
@@ -32,7 +32,7 @@ function unitstride(ls::LoopSet, op::Operation, s::Symbol)
         indexappearences(parent, s) > 1 && return false
     end
     for i ∈ 2:length(inds)
-        if li[i]
+        if li[lookup[i]]
             s === inds[i] && return false
         else
             parent = findparent(ls, inds[i])
@@ -344,7 +344,7 @@ function maybedemotesize(T::Int, N::Int, U::Int, Uloop::Loop, maxTbase::Int)
 end
 function solve_tilesize(
     ls::LoopSet, unrolled::Symbol, tiled::Symbol,
-    cost_vec::AbstractVector{Float64}, 
+    cost_vec::AbstractVector{Float64},
     reg_pressure::AbstractVector{Int},
     W::Int, vectorized::Symbol
 )
@@ -436,7 +436,7 @@ function evaluate_cost_tile(
     # Need to check if fusion is possible
     size_T = biggest_type_size(ls)
     W, Wshift = VectorizationBase.pick_vector_width_shift(length(ls, vectorized), size_T)::Tuple{Int,Int}
-    # costs = 
+    # costs =
     # cost_mat[1] / ( unrolled * tiled)
     # cost_mat[2] / ( tiled)
     # cost_mat[3] / ( unrolled)
@@ -569,7 +569,7 @@ function choose_unroll_order(ls::LoopSet, lowest_cost::Float64 = Inf)
         iter = iterate(lo, state)
         iter === nothing && return best_order, best_vec, lowest_cost
         new_order, state = iter
-    end    
+    end
 end
 function choose_tile(ls::LoopSet)
     lo = LoopOrders(ls)
@@ -627,4 +627,3 @@ function register_pressure(ls::LoopSet)
         tU * tT * rp[1] + tU * rp[2] + rp[3] + rp[4]
     end
 end
-

src/graphs.jl

@@ -174,6 +174,7 @@ Base.@propagate_inbounds Base.getindex(lo::LoopOrder, i...) = lo.oporder[LinearI
 # O(N) search is faster at small sizes
 struct LoopSet
     loopsymbols::Vector{Symbol}
+    loopsymbol_offsets::Vector{Int}  # symbol loopsymbols[i] corresponds to loops[lso[i]+1:lso[i+1]] (CartesianIndex handling)
     loops::Vector{Loop}
     opdict::Dict{Symbol,Operation}
     operations::Vector{Operation} # Split them to make it easier to iterate over just a subset
@@ -281,7 +282,7 @@ includesarray(ls::LoopSet, array::Symbol) = array ∈ ls.includedarrays
 
 function LoopSet(mod::Symbol)# = :LoopVectorization)
     LoopSet(
-        Symbol[], Loop[],
+        Symbol[], [0], Loop[],
         Dict{Symbol,Operation}(),
         Operation[],
         Int[],

src/lower_load.jl

@@ -8,7 +8,7 @@ function pushvectorload!(q::Expr, op::Operation, var::Symbol, td::UnrollArgs, U:
     end
     push!(q.args, Expr(:(=), name, instrcall))
 end
-function lower_load_scalar!( 
+function lower_load_scalar!(
     q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
     suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
 )
@@ -60,6 +60,3 @@ function lower_load!(
         lower_load_scalar!(q, op, vectorized, W, unrolled, tiled, U, suffix, mask)
     end
 end
-
-
-

src/lower_memory_common.jl

@@ -27,12 +27,12 @@ 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)
-    loopedindex = op.ref.loopedindex
+    loopedindex, indexlookup = op.ref.loopedindex, op.ref.indexlookup
     start = (first(indices) === Symbol("##DISCONTIGUOUSSUBARRAY##")) + 1
     for (n,ind) ∈ enumerate(@view(indices[start:end]))
         if ind isa Int
             push!(ret.args, ind)
-        elseif loopedindex[n]
+        elseif loopedindex[indexlookup[n]]
             push!(ret.args, ind)
         else
             push!(ret.args, symbolind(ind, op, td))
@@ -46,7 +46,7 @@ function mem_offset_u(op::Operation, td::UnrollArgs)
     incr = u
     ret = Expr(:tuple)
     indices = getindices(op)
-    loopedindex = op.ref.loopedindex
+    loopedindex, indexlookup = op.ref.loopedindex, op.ref.indexlookup
     if incr == 0
         return mem_offset(op, td)
         # append_inds!(ret, indices, loopedindex)
@@ -57,7 +57,7 @@ function mem_offset_u(op::Operation, td::UnrollArgs)
                 push!(ret.args, ind)
             elseif ind === unrolled
                 push!(ret.args, Expr(:call, :+, ind, incr))
-            elseif loopedindex[n]
+            elseif loopedindex[indexlookup[n]]
                 push!(ret.args, ind)
             else
                 push!(ret.args, symbolind(ind, op, td))
@@ -117,4 +117,3 @@ function name_memoffset(var::Symbol, op::Operation, td::UnrollArgs, W::Symbol, v
     end
     name, mo
 end
-

src/memory_ops_common.jl

@@ -64,7 +64,7 @@ function array_reference_meta!(ls::LoopSet, array::Symbol, rawindices, elementby
             else
                 indop = get(ls.opdict, ind, nothing)
                 if indop !== nothing  && !isconstant(indop)
-                    pushparent!(parents, loopdependencies, reduceddeps, parent)
+                    pushparent!(parents, loopdependencies, reduceddeps, parent)   # FIXME where does `parent` come from?
                     # var = get(ls.opdict, ind, nothing)
                     push!(indices, name(parent)); ninds += 1
                     push!(loopedindex, false)

src/operations.jl

@@ -20,10 +20,11 @@ struct ArrayReferenceMeta
     ref::ArrayReference
     loopedindex::Vector{Bool}
     ptr::Symbol
+    indexlookup::Vector{Int}
 end
-function ArrayReferenceMeta(ref::ArrayReference, loopedindex, ptr = vptr(ref))
+function ArrayReferenceMeta(ref::ArrayReference, loopedindex, ptr = vptr(ref), indexlookup = [i for i in 1:length(loopedindex)])
     ArrayReferenceMeta(
-        ref, loopedindex, ptr
+        ref, loopedindex, ptr, indexlookup
     )
 end
 # function Base.hash(x::ArrayReference, h::UInt)
@@ -174,7 +175,7 @@ These names will be further processed if op is tiled and/or unrolled.
     if tiled ∈ loopdependencies(op) # `suffix` is tilenumber
         mvar = Symbol(op, suffix, :_)
     end
-    if unrolled ∈ loopdependencies(op) # `u` is unroll number 
+    if unrolled ∈ loopdependencies(op) # `u` is unroll number
         mvar = Symbol(op, u)
     end
 ```
@@ -240,6 +241,3 @@ getindices(op::Operation) = op.ref.ref.indices
     # # access stride info?
     # op.numerical_metadata[symposition(op,sym)]
 # end
-
-
-

src/reconstruct_loopset.jl

@@ -25,12 +25,31 @@ function Loop(ls, l, ::Type{StaticUnitRange{L,U}}) where {L,U}
     Loop(gensym(:n), L, U, Symbol(""), Symbol(""), true, true)::Loop
 end
 
+function Loop(ls::LoopSet, l::Int, k::Int, ::Type{<:CartesianIndices{N}}) where N
+    start = gensym(:loopstart); stop = gensym(:loopstop)
+    axisexpr = Expr(:ref, Expr(:., Expr(:ref, :lb, l), QuoteNode(:indices)), k)
+    pushpreamble!(ls, Expr(:(=), start, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), Expr(:(.), axisexpr, QuoteNode(:start)))))
+    pushpreamble!(ls, Expr(:(=), stop, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), Expr(:(.), axisexpr, QuoteNode(:stop)))))
+    Loop(gensym(:n), 0, 1024, start, stop, false, false)::Loop
+end
+
 function add_loops!(ls::LoopSet, LB)
-    loopsyms = [gensym(:n) for _ ∈ eachindex(LB)]
     for (i,l) ∈ enumerate(LB)
-        add_loop!(ls, Loop(ls, i, l)::Loop)
+        if l<:CartesianIndices
+            add_loops!(ls, i, l)
+        else
+            add_loop!(ls, Loop(ls, i, l)::Loop)
+            push!(ls.loopsymbol_offsets, ls.loopsymbol_offsets[end]+1)
+        end
     end
 end
+function add_loops!(ls, i, l::Type{<:CartesianIndices{N}}) where N
+    for k = N:-1:1
+        add_loop!(ls, Loop(ls, i, k, l)::Loop)
+    end
+    push!(ls.loopsymbol_offsets, ls.loopsymbol_offsets[end]+N)
+end
+
 function ArrayReferenceMeta(
     ls::LoopSet, ar::ArrayRefStruct, arraysymbolinds::Vector{Symbol}, opsymbols::Vector{Symbol},
     array::Symbol, vp::Symbol
@@ -39,21 +58,28 @@ function ArrayReferenceMeta(
     indices = ar.indices
     offsets = ar.offsets
     ni = filled_8byte_chunks(index_types)
-    index_vec = Vector{Symbol}(undef, ni)
+    index_vec = Symbol[]
     offset_vec = Vector{Int8}(undef, ni)
     loopedindex = fill(false, ni)
+    indexlookup = Int[]
     while index_types != zero(UInt64)
         ind = indices % UInt8
-        symind = if index_types == LoopIndex
+        if index_types == LoopIndex
+            for inda in ls.loopsymbol_offsets[ind]+1:ls.loopsymbol_offsets[ind+1]
+                pushfirst!(index_vec, ls.loopsymbols[inda])
+                pushfirst!(indexlookup, ni)
+            end
             loopedindex[ni] = true
-            ls.loopsymbols[ind]
-        elseif index_types == ComputedIndex
-            opsymbols[ind]
         else
-            @assert index_types == SymbolicIndex
-            arraysymbolinds[ind] 
+            symind = if index_types == ComputedIndex
+                opsymbols[ind]
+            else
+                @assert index_types == SymbolicIndex
+                arraysymbolinds[ind]
+            end
+            pushfirst!(index_vec, symind)
+            pushfirst!(indexlookup, ni)
         end
-        index_vec[ni] = symind
         offset_vec[ni] = offsets % Int8
         index_types >>>= 8
         indices >>>= 8
@@ -62,7 +88,7 @@ function ArrayReferenceMeta(
     end
     ArrayReferenceMeta(
         ArrayReference(array, index_vec, offset_vec),
-        loopedindex, vp
+        loopedindex, vp, indexlookup
     )
 end
 
@@ -134,14 +160,16 @@ function process_metadata!(ls::LoopSet, AM, num_arrays::Int)::Vector{Symbol}
     arraysymbolinds
 end
 function parents_symvec(ls::LoopSet, u::Unsigned)
-    i = filled_4byte_chunks(u)
-    loops = Vector{Symbol}(undef, i)
+    loops = Symbol[]
+    offsets = ls.loopsymbol_offsets
     while u != zero(u)
-        loops[i] = getloopsym(ls, ( u % UInt8 ) & 0x0f )
-        i -= 1
+        idx = ( u % UInt8 ) & 0x0f
+        for j = offsets[idx]+1:offsets[idx+1]
+            push!(loops, getloopsym(ls, j))
+        end
         u >>= 4
     end
-    loops
+    return reverse!(loops)
 end
 loopdependencies(ls::LoopSet, os::OperationStruct) = parents_symvec(ls, os.loopdeps)
 reduceddependencies(ls::LoopSet, os::OperationStruct) = parents_symvec(ls, os.reduceddeps)
@@ -227,7 +255,7 @@ function avx_loopset(instr, ops, arf, AM, LB, vargs)
     num_arrays = length(arf)
     elementbytes = sizeofeltypes(vargs, num_arrays)
     add_loops!(ls, LB)
-    resize!(ls.loop_order, length(LB))
+    resize!(ls.loop_order, ls.loopsymbol_offsets[end])
     arraysymbolinds = process_metadata!(ls, AM, length(arf))
     opsymbols = [gensym(:op) for _ ∈ eachindex(ops)]
     mrefs = create_mrefs!(ls, arf, arraysymbolinds, opsymbols, vargs)
@@ -249,7 +277,7 @@ function _avx_loopset_debug(::Type{OPS}, ::Type{ARF}, ::Type{AM}, ::Type{LB}, va
     @show OPS ARF AM LB vargs
     _avx_loopset(OPS.parameters, ARF.parameters, AM.parameters, LB.parameters, typeof.(vargs))
 end
-function _avx_loopset(OPSsv, ARFsv, AMsv, LBsv, vargs) where {UT, OPS, ARF, AM, LB}
+function _avx_loopset(OPSsv, ARFsv, AMsv, LBsv, 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 ]

test/offsetarrays.jl

@@ -83,7 +83,7 @@ using Test
     end
 
 
-    
+
     struct SizedOffsetMatrix{T,LR,UR,LC,RC} <: AbstractMatrix{T}
         data::Matrix{T}
     end
@@ -141,25 +141,37 @@ using Test
     # lsuq = LoopVectorization.LoopSet(macroexpand(Base, uq));
     # LoopVectorization.choose_order(lsuq)
 
-    
+
+    function avxgeneric!(out, A, kern, R=CartesianIndices(out), z=zero(eltype(out)))
+       Rk = CartesianIndices(kern)
+       @avx for I in R
+           tmp = z
+           for J in Rk
+               tmp += A[I+J]*kern[J]
+           end
+           out[I] = tmp
+       end
+       out
+   end
+
     for T ∈ (Float32, Float64)
         @show T, @__LINE__
         A = rand(T, 100, 100);
         kern = OffsetArray(rand(T, 3, 3), -1:1, -1:1);
         skern = SizedOffsetMatrix{T,-1,1,-1,1}(parent(kern));
         out1 = OffsetArray(similar(A, size(A).-2), 1, 1);   # stay away from the edges of A
-        out2 = similar(out1); out3 = similar(out1);
+        out2 = similar(out1); out3 = similar(out1); out4 = similar(out1)
 
         old2d!(out1, A, kern);
         avx2d!(out2, A, kern);
         @test out1 ≈ out2
-        
+
         avx2douter!(out3, A, kern);
         @test out1 ≈ out3
 
         fill!(out2, NaN); avx2d!(out2, A, skern);
         @test out1 ≈ out2
-        
+
         fill!(out3, NaN); avx2douter!(out3, A, skern);
         @test out1 ≈ out3
 
@@ -168,11 +180,12 @@ using Test
 
         fill!(out3, NaN); avx2dunrolled2x2!(out3, A, skern);
         @test out1 ≈ out3
-        
+
         fill!(out3, NaN); avx2dunrolled3x3!(out3, A, skern);
         @test out1 ≈ out3
-end
 
-    
-end
+        @test_broken avxgeneric!(out4, A, kern) ≈ out1
+    end
 
+
+end