test += in more places, handle vectorized -> non-vectorized

Author: chriselrod

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.12.9"
+version = "0.12.10"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/codegen/lower_compute.jl

@@ -333,7 +333,22 @@ function getu₁forreduct(ls::LoopSet, op::Operation, u₁::Int)
         return getu₁full(ls, u₁)
     end    
 end
-
+isidentityop(op::Operation) = iscompute(op) && (instruction(op).instr === :identity) && (length(parents(op)) == 1)
+function reduce_parent!(q::Expr, ls::LoopSet, op::Operation, opp::Operation, parent::Symbol)
+    isvectorized(op) && return parent
+    dependent_outer_reducts(ls, op) && return parent
+    if isvectorized(opp)
+        oppt = opp
+    elseif isidentityop(opp)
+        oppt = only(parents(opp))
+        isvectorized(oppt) || return parent
+    else
+        return parent
+    end
+    newp = gensym(parent)
+    push!(q.args, Expr(:(=), newp, Expr(:call, lv(reduction_to_scalar(oppt.instruction)), parent)))
+    newp
+end
 function lower_compute!(
     q::Expr, op::Operation, ls::LoopSet, ua::UnrollArgs, mask::Bool
 )
@@ -468,25 +483,26 @@ function lower_compute!(
             add_loopvalue!(instrcall, loopval, ua, u₁)
         elseif name(opp) === name(op)
             selfdep = n
-            if ((isvectorized(first(parents_op)) && !isvectorized(op)) && !dependent_outer_reducts(ls, op)) ||
+            if ((isvectorized(opp) && !isvectorized(op)) && !dependent_outer_reducts(ls, op)) ||
                 (parents_u₁syms[n] != u₁unrolledsym) || (parents_u₂syms[n] != u₂unrolledsym)
 
                 selfopname, uₚ = parent_op_name(ls, parents_op, n, modsuffix, suffix_, parents_u₁syms, parents_u₂syms, u₁, opisvectorized, tiledouterreduction)
                 # if (uₚ ≠ 0) & (uₚ ≠ u₁)
                 #     dopartialmap = true
                 # end
+                # @show selfopname, instr
                 push!(instrcall.args, selfopname)
             else
                 push!(instrcall.args, varsym)
             end
         elseif ((!isu₂unrolled(op)) & isu₂unrolled(opp)) && (isouterreduction(ls, opp) != -1)
             # this checks if the parent is u₂ unrolled but this operation is not, in which case we need to reduce it.
-            push!(instrcall.args, reduce_expr_u₂(mangledvar(opp), instruction(opp), ureduct(ls)))
+            reduced_u₂ = reduce_expr_u₂(mangledvar(opp), instruction(opp), ureduct(ls))
+            reduced_u₂ = reduce_parent!(q, ls, op, opp, reduced_u₂)
+            push!(instrcall.args, reduced_u₂)
         else
             parent, uₚ = parent_op_name(ls, parents_op, n, modsuffix, suffix_, parents_u₁syms, parents_u₂syms, u₁, opisvectorized, tiledouterreduction)
-            # if name(op) === Symbol("##op#9536")
-            #     @show parent
-            # end
+            parent = reduce_parent!(q, ls, op, opp, parent)
             if (selfdep == 0) && search_tree(parents(opp), name(op))
                 selfdep = n
                 push!(instrcall.args, parent)
@@ -495,8 +511,6 @@ function lower_compute!(
             else
                 push!(instrcall.args, parent)
             end
-
-            # @show uₚ, u₁, op
         end
     end
     selfdepreduce = ifelse(((!u₁unrolledsym) & isu₁unrolled(op)) & (u₁ > 1), selfdep, 0)

src/reconstruct_loopset.jl

@@ -623,7 +623,7 @@ Execute an `@avx` block. The block's code is represented via the arguments:
 @generated function _avx_!(
     ::Val{UNROLL}, ::Val{OPS}, ::Val{ARF}, ::Val{AM}, ::Val{LPSYM}, var"#lv#tuple#args#"::Tuple{LB,V}
 ) where {UNROLL, OPS, ARF, AM, LPSYM, LB, V}
-    1 + 1 # Irrelevant line you can comment out/in to force recompilation...
+    # 1 + 1 # Irrelevant line you can comment out/in to force recompilation...
     ls = _avx_loopset(OPS, ARF, AM, LPSYM, LB.parameters, V.parameters, UNROLL)
     # return @show avx_body(ls, UNROLL)
     if last(UNROLL) > 1

test/gemv.jl

@@ -31,7 +31,7 @@ using Test
             for j ∈ eachindex(x)
                 yᵢ += A[i,j] * x[j]
             end
-            y[i] = yᵢ
+            y[i] += yᵢ
         end
     end
     function mygemvavx_range!(y, A, x)
@@ -246,8 +246,12 @@ using Test
         # y1 = Vector{TC}(undef, M); y2 = similar(y1);
 
         mygemv!(y1, A, x);
-        mygemvavx!(y2, A, x);
+        fill!(y2, 0); mygemvavx!(y2, A, x);
         @test y1full ≈ y2full
+        mygemvavx!(y2, A, x);
+        @test y1 .* 2 ≈ y2
+        mygemvavx!(y2, A, x);
+        @test y1 .* 3 ≈ y2
         fill!(y2, -9999); mygemv_avx!(y2, A, x);
         @test y1full ≈ y2full
         fill!(y2, -9999);
@@ -263,12 +267,12 @@ using Test
                 fill!(y2, -9999); mygemv_avx!(y2, Abit, x);
                 @test y2 ≈ Abit * x
             end
-            fill!(y2, -9999); mygemvavx!(y2, Abit, x);
+            fill!(y2, 0); mygemvavx!(y2, Abit, x);
             @test y2 ≈ Abit * x
             xbit = x .> 0.5;
             fill!(y2, -9999); mygemv_avx!(y2, A, xbit);
             @test y2 ≈ A * xbit
-            fill!(y2, -9999); mygemvavx!(y2, A, xbit);
+            fill!(y2, 0); mygemvavx!(y2, A, xbit);
             @test y2 ≈ A * xbit
         end
 

test/shuffleloadstores.jl

@@ -209,15 +209,17 @@ end
             @test Cc1 ≈ Cc2# ≈ Cc3
         end
     end
-    M = 100
-    G = 50
-    J = 50
-    H = 300
+    if VERSION ≥ v"1.6.0-rc1"
+        M = 100
+        G = 50
+        J = 50
+        H = 300
 
-    A = Matrix(Tridiagonal(rand(G-1,G-1)));
-    B = rand(Complex{Float64}, 2*J+1, G-1, H+1, M+1);
-    ϕ = rand(Complex{Float64}, 2*J+1, G+1, H+1, M+1);
-    @test issue209(M, G, J, H, A, B, ϕ) ≈ issue209_noavx(M, G, J, H, A, B, ϕ)
+        A = Matrix(Tridiagonal(rand(G-1,G-1)));
+        B = rand(Complex{Float64}, 2*J+1, G-1, H+1, M+1);
+        ϕ = rand(Complex{Float64}, 2*J+1, G+1, H+1, M+1);
+        @test issue209(M, G, J, H, A, B, ϕ) ≈ issue209_noavx(M, G, J, H, A, B, ϕ)
+    end
 end