Fixed a bug where loops of the form a[i] = f(a[i]) weren't handled correctly.

Author: chriselrod

src/determinestrategy.jl

@@ -114,7 +114,12 @@ function determine_unroll_factor(
 
     # The strategy is to use an unroll factor of 1, unless there appears to be loop carried dependencies (ie, num_reductions > 0)
     # The assumption here is that unrolling provides no real benefit, unless it is needed to enable OOO execution by breaking up these dependency chains
-    num_reductions = sum(isreduction, operations(ls))
+    num_reductions = 0#sum(isreduction, operations(ls))
+    for op ∈ operations(ls)
+        if isreduction(op) & iscompute(op)
+            num_reductions += 1
+        end
+    end
     # @show num_reductions
     if iszero(num_reductions) # the 4 is a hack, based on the idea that there is some cost to moving through columns
         return length(order) == 1 ? 1 : 4

src/graphs.jl

@@ -88,6 +88,13 @@ struct LoopSet
     # ref_to_sym_aliases::Dict{ArrayReference,Symbol}
 end
 
+# function op_to_ref(ls::LoopSet, op::Operation)
+    # s = op.variable
+    # id = findfirst(ls.syms_aliasing_regs)
+    # @assert id !== nothing
+    # ls.refs_aliasing_syms[id]
+# end
+
 function includesarray(ls::LoopSet, array::Symbol)
     for (a,i) ∈ ls.includedarrays
         a === array && return i
@@ -235,7 +242,7 @@ function add_load!(
         :getindex, memload, loopdependencies(ref),
         NODEPENDENCY, NOPARENTS, ref
     )
-    add_vptr!(ls, indexed, identifier(op))
+    add_vptr!(ls, ref.array, identifier(op))
     pushop!(ls, op, var)
 end
 function add_load_ref!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
@@ -311,12 +318,16 @@ function maybe_cse_load!(ls::LoopSet, expr::Expr, elementbytes::Int = 8)
     else
         return add_operation!(ls, gensym(:temporary), expr, elementbytes)
     end
-    ref = ArrayReference( ex.args[1+offset], @view(ex.args[2+offset:end]) )::ArrayReference
+    ref = ArrayReference(
+        expr.args[1+offset],
+        @view(expr.args[2+offset:end]),
+        Ref(false)
+    )::ArrayReference
     id = findfirst(r -> r == ref, ls.refs_aliasing_syms)
     if id === nothing
-        add_load!( ls, gensym(:temporary), array, args, elementbytes )
+        add_load!( ls, gensym(:temporary), ref, elementbytes )
     else
-        ls.syms_aliasing_refs[id]
+        getop(ls, ls.syms_aliasing_refs[id])        
     end
     # id = includesarray(ls, array)
     # if id > 0
@@ -371,7 +382,7 @@ function add_compute!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8,
     # op = Operation( length(operations(ls)), var, elementbytes, instr, compute )
     reduction = false
     for arg ∈ args
-        if arg === var
+        if var === arg
             reduction = true
             add_reduction!(parents, deps, reduceddeps, ls, arg, elementbytes)
         elseif ref == arg
@@ -389,18 +400,20 @@ function add_compute!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8,
     end
 end
 function add_store!(
-    ls::LoopSet, indexed::Symbol, var::Symbol, indices::AbstractVector, elementbytes::Int = 8
+    ls::LoopSet, var::Symbol, ref::ArrayReference, elementbytes::Int = 8
 )
     parent = getop(ls, var)
-    op = Operation( length(operations(ls)), indexed, elementbytes, :setindex!, memstore, indices, reduceddependencies(parent), [parent] )
-    add_vptr!(ls, indexed, identifier(op))
+    op = Operation( length(operations(ls)), ref.array, elementbytes, :setindex!, memstore, loopdependencies(ref), reduceddependencies(parent), [parent], ref )
+    add_vptr!(ls, ref.array, identifier(op))
     pushop!(ls, op, var)
 end
 function add_store_ref!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
-    add_store!(ls, ex.args[1], var, @view(ex.args[2:end]), elementbytes)
+    ref = ref_from_ref(ex)
+    add_store!(ls, var, ref, elementbytes)
 end
 function add_store_setindex!(ls::LoopSet, ex::Expr, elementbytes::Int = 8)
-    add_store!(ls, ex.args[2], ex.args[3], @view(ex.args[4:end]), elementbytes)
+    ref = ref_from_setindex(ex)
+    add_store!(ls, var, ref, elementbytes)
 end
 # add operation assigns X to var
 function add_operation!(
@@ -418,6 +431,21 @@ function add_operation!(
         throw("Expression not recognized:\n$x")
     end
 end
+function add_operation!(
+    ls::LoopSet, LHS_sym::Symbol, RHS::Expr, LHS_ref::ArrayReference, elementbytes::Int = 8
+)
+    if RHS.head === :ref# || (RHS.head === :call && first(RHS.args) === :getindex)
+        add_load!(ls, LHS_sym, LHS_ref, elementbytes)
+    elseif RHS.head === :call
+        if first(RHS.args) === :getindex
+            add_load!(ls, LHS_sym, LHS_ref, elementbytes)
+        else
+            add_compute!(ls, LHS_sym, RHS, elementbytes, LHS_ref)
+        end
+    else
+        throw("Expression not recognized:\n$x")
+    end
+end
 function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int = 8)
     if ex.head === :call
         finex = first(ex.args)::Symbol
@@ -446,7 +474,9 @@ function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int = 8)
                 ref = ArrayReference(LHS)
                 id = findfirst(r -> r == ref, ls.refs_aliasing_syms)
                 lrhs = id === nothing ? gensym(:RHS) : ls.syms_aliasing_refs[id]
-                add_operation!(ls, lrhs, RHS, elementbytes, ref)
+                # we pass ref, so it can compare references within RHS, and realize
+                # they equal lrhs
+                add_operation!(ls, lrhs, RHS, ref, elementbytes)
             end
             add_store_ref!(ls, lrhs, LHS, elementbytes)
         else

src/lowering.jl

@@ -5,12 +5,13 @@
 function mem_offset(op::Operation, incr::Int = 0)
     @assert accesses_memory(op) "Computing memory offset only makes sense for operations that access memory."
     ret = Expr(:tuple, )
-    deps = op.dependencies
+    deps = op.ref.ref
     if incr == 0
         append!(ret.args, deps)
     else
-        push!(ret.args, Expr(:call, :+, first(deps), incr))
-        for n ∈ 2:length(deps)
+        dep = first(deps)
+        push!(ret.args, dep isa Symbol ? Expr(:call, :+, dep, incr) : dep + incr)
+         for n ∈ 2:length(deps)
             push!(ret.args, deps[n])
         end
     end
@@ -19,7 +20,7 @@ end
 function mem_offset(op::Operation, incr::Int, unrolled::Symbol)
     @assert accesses_memory(op) "Computing memory offset only makes sense for operations that access memory."
     ret = Expr(:tuple, )
-    deps = op.dependencies
+    deps = op.ref.ref
     if incr == 0
         append!(ret.args, deps)
     else
@@ -55,7 +56,7 @@ function lower_load_scalar!(
     if suffix !== nothing
         var = Symbol(var, :_, suffix)
     end
-    ptr = Symbol("##vptr##_", first(op.reduced_deps))
+    ptr = refname(op)
     push!(q.args, Expr(:(=), Symbol("##", var), Expr(:call, lv(:load),  ptr, mem_offset(op))))
     nothing
 end
@@ -69,7 +70,7 @@ function lower_load_unrolled!(
     if suffix !== nothing
         var = Symbol(var, :_, suffix)
     end
-    ptr = Symbol("##vptr##_", first(op.reduced_deps))
+    ptr = refname(op)
     val = Expr(:call, Expr(:curly, :Val, W))
     if first(loopdependencies(op)) === unrolled # vload
         for u ∈ 0:U-1
@@ -157,7 +158,7 @@ function lower_store_reduction!(
     if suffix !== nothing
         var = Symbol(var, :_, suffix)
     end
-    ptr = Symbol("##vptr##_", op.variable)
+    ptr = refname(op)
     # need to find out reduction type
     instr = first(parents(op)).instruction
     reduce_expr!(q, var, instr, U) # assigns reduction to storevar
@@ -174,7 +175,7 @@ function lower_store_scalar!(
     if suffix !== nothing
         var = Symbol(var, :_, suffix)
     end
-    ptr = Symbol("##vptr##_", op.variable)
+    ptr = refname(op)
     push!(q.args, Expr(:call, lv(:store!), ptr, Symbol("##", var), mem_offset(op)))
     nothing
 end
@@ -188,7 +189,7 @@ function lower_store_unrolled!(
     if suffix !== nothing
         var = Symbol(var, :_, suffix)
     end
-    ptr = Symbol("##vptr##_", op.variable)
+    ptr = refname(op)
     if first(loopdependencies(op)) === unrolled # vstore!
         for u ∈ 0:U-1
             instrcall = Expr(:call,lv(:vstore!), ptr, Symbol("##",var,:_,u), mem_offset(op, u*W))

src/operations.jl

@@ -23,12 +23,12 @@ end
 
 Base.:(==)(x::ArrayReference, y::ArrayReference) = isequal(x, y)
 
-function ref_from_ref(ex::Expr)
-    ArrayReference( ex.args[1], @view(ex.args[2:end]), Ref(false) )
-end
-function ref_from_getindex(ex::Expr)
-    ArrayReference( ex.args[2], @view(ex.args[3:end]), Ref(false) )
+function ref_from_expr(ex, offset1::Int = 0, offset2 = 0)
+    ArrayReference( ex.args[1 + offset1], @view(ex.args[2 + offset2:end]), Ref(false) )
 end
+ref_from_ref(ex::Expr) = ref_from_expr(ex, 0, 0)
+ref_from_getindex(ex::Expr) = ref_from_expr(ex, 1, 1)
+ref_from_setindex(ex::Expr) = ref_from_expr(ex, 1, 2)
 function ArrayReference(ex::Expr)
     ex.head === :ref ? ref_from_ref(ex) : ref_from_getindex(ex)
 end
@@ -46,7 +46,7 @@ Base.:(==)(x::ArrayReference, y) = false
 
 
 # Avoid memory allocations by accessing this
-const NOTAREFERENCE = ArrayReference(Symbol(""), Union{Symbol,Int}[])
+const NOTAREFERENCE = ArrayReference(Symbol(""), Union{Symbol,Int}[], Ref(false))
 
 @enum OperationType begin
     constant
@@ -128,6 +128,9 @@ identifier(op::Operation) = op.identifier + 1
 name(op::Operation) = op.variable
 instruction(op::Operation) = op.instruction
 
+refname(op::Operation) = Symbol("##vptr##_", op.ref.array)
+
+
 """
 Returns `0` if the op is the declaration of the constant outerreduction variable.
 Returns `n`, where `n` is the constant declarations's index among parents(op), if op is an outter reduction.

test/runtests.jl

@@ -77,6 +77,7 @@ function mygemmavx!(C, A, B)
 end
 
 M, K, N = rand(70:81, 3);
+M, K, N = 72, 75, 71;
 C = Matrix{Float64}(undef, M, N); A = randn(M, K); B = randn(K, N);
 C2 = similar(C);
 mygemmavx!(C, A, B)
@@ -328,6 +329,7 @@ lscolsum = LoopVectorization.LoopSet(colsumq);
 lscolsum
 lscolsum.operations
 
+LoopVectorization.choose_order(lscolsum)
 @test LoopVectorization.choose_order(lscolsum) == (Symbol[:j,:i], 4, -1)
 
 function mycolsum!(x, A)
@@ -370,6 +372,7 @@ varq = :(for j ∈ eachindex(s²), i ∈ 1:size(A,2)
          s²[j] += δ*δ
          end)
 lsvar = LoopVectorization.LoopSet(varq);
+LoopVectorization.choose_order(lsvar)
 @test LoopVectorization.choose_order(lsvar) == (Symbol[:j,:i], 4, -1)
 
 function myvar!(s², A, x̄)