Updates to fix transposed/adjoint broadcasting.

Author: chriselrod

Manifest.toml

@@ -61,15 +61,13 @@ uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [[SIMDPirates]]
 deps = ["MacroTools", "VectorizationBase"]
-git-tree-sha1 = "296cae2ccd6e4766aad669e748c1248fb99ab69c"
+git-tree-sha1 = "5f7fc8a48d9806817bb3f8a2ef9793398744e8a6"
 uuid = "21efa798-c60a-11e8-04d3-e1a92915a26a"
-version = "0.1.0"
+version = "0.1.1"
 
 [[SLEEFPirates]]
 deps = ["SIMDPirates", "VectorizationBase"]
 git-tree-sha1 = "1c5b6827da87a12bdb7a4c893f44c3adbce3389d"
-repo-rev = "master"
-repo-url = "https://github.com/chriselrod/SLEEFPirates.jl"
 uuid = "476501e8-09a2-5ece-8869-fb82de89a1fa"
 version = "0.1.1"
 

Project.toml

@@ -14,7 +14,7 @@ VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 [compat]
 MacroTools = "0.5"
 Parameters = "0.12.0"
-SIMDPirates = "0.1.0"
+SIMDPirates = "0.1.1"
 SLEEFPirates = "0.1.1"
 VectorizationBase = "0.1.3"
 julia = "1.3.0"

README.md

@@ -115,10 +115,10 @@ In the future, I would like it to also model the cost of memory movement in the
 Until then, performance will degrade rapidly compared to BLAS as the size of the matrices increase. The advantage of the `@avx` macro, however, is that it is general. Not every operation is supported by BLAS.
 
 For example, what if `A` were the outter product of two vectors?
-```julia
+<!-- ```julia -->
 
 
-```
+<!-- ``` -->
 
 
 Another example, a straightforward operation expressed well via broadcasting:
@@ -137,6 +137,8 @@ d2 = @avx @. a + B * c′;
 can be optimized in a similar manner to BLAS, albeit to a much smaller degree because the naive version already benefits from vectorization (unlike the naive BLAS).
 
 
+<!-- You can also use `\ast` to for a lazy matrix multiplication that can fuse with broadcasts. `.\ast` behaves similarly, to allow it's arguments to -->
+
 
 
 Originally, LoopVectorization only provided a simple, dumb, transform on a single loop using the `@vectorize` macro. This transformation took element type and unroll factor arguments, performing no analysis of the loop, simply applying the specified arguments.

src/LoopVectorization.jl

@@ -4,7 +4,7 @@ using VectorizationBase, SIMDPirates, SLEEFPirates, MacroTools, Parameters
 using VectorizationBase: REGISTER_SIZE, REGISTER_COUNT, extract_data, num_vector_load_expr, mask, masktable, pick_vector_width_val, valmul, valrem, valmuladd
 using SIMDPirates: VECTOR_SYMBOLS, evadd, evmul, vrange, reduced_add, reduced_prod
 using Base.Broadcast: Broadcasted, DefaultArrayStyle
-using LinearAlgebra: Adjoint
+using LinearAlgebra: Adjoint, Transpose
 using MacroTools: prewalk, postwalk
 
 

src/broadcast.jl

@@ -95,19 +95,31 @@ function add_broadcast!(
     ref = ArrayReference(bcname, fulldims, Ref{Bool}(false))
     add_load!(ls, destname, ref, elementbytes)::Operation
 end
+function add_broadcast_adjoint_array!(
+    ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol}, ::Type{A}, elementbytes::Int = 8
+) where {T,N,A<:AbstractArray{T,N}}
+    parent = gensym(:parent)
+    pushpreamble!(ls, Expr(:(=), parent, Expr(:call, :parent, bcname)))
+    ref = ArrayReference(parent, 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_adjoint_array!(
+    ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol}, ::Type{<:AbstractVector}, elementbytes::Int = 8
+)
+    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{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
+) where {T, A <: AbstractArray{T}}
+    add_broadcast_adjoint_array!( ls, destname, bcname, loopsyms, A, elementbytes )
 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 )
+    ::Type{Transpose{T,A}}, elementbytes::Int = 8
+) where {T, A <: AbstractArray{T}}
+    add_broadcast_adjoint_array!( ls, destname, bcname, loopsyms, A, elementbytes )
 end
 function add_broadcast!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},

src/lowering.jl

@@ -88,8 +88,10 @@ function lower_load_unrolled!(
     else
         sn = findfirst(x -> x === unrolled, loopdependencies(op))::Int
         ustrides = Expr(:call, lv(:vmul), Expr(:call, :stride, ptr, sn), Expr(:call, lv(:vrange), W))
+        ustride = gensym(:ustride)
+        push!(q.args, Expr(:(=), ustride, ustrides))
         for u ∈ 0:U-1
-            instrcall = Expr(:call, lv(:gather), ptr, Expr(:call,lv(:vadd),mem_offset(op, u, W, unrolled),ustrides))
+            instrcall = Expr(:call, lv(:gather), ptr, mem_offset(op, u, W, unrolled), ustride)
             if mask !== nothing && u == U - 1
                 push!(instrcall.args, mask)
             end
@@ -207,7 +209,7 @@ function lower_store_unrolled!(
         sn = findfirst(x -> x === unrolled, loopdependencies(op))::Int
         ustrides = Expr(:call, lv(:vmul), Expr(:call, :stride, ptr, sn), Expr(:call, lv(:vrange), W))
         for u ∈ 0:U-1
-            instrcall = Expr(:call, lv(:scatter!), ptr, Symbol("##",var,:_,u), Expr(:call,lv(:vadd),mem_offset(op,u,W,unrolled),ustrides))
+            instrcall = Expr(:call, lv(:scatter!), ptr, mem_offset(op,u,W,unrolled), ustrides, Symbol("##",var,:_,u))
             if mask !== nothing && u == U - 1
                 push!(instrcall.args, mask)
             end

test/runtests.jl

@@ -1,5 +1,6 @@
 using Test
 using LoopVectorization
+using LinearAlgebra
 
 
 @testset "LoopVectorization.jl" begin
@@ -356,6 +357,16 @@ end
         D1 = C .+ A * B;
         D2 = @avx C .+ A ∗ B;
         @test D1 ≈ D2
+
+        D3 = exp.(B')
+        D4 = @avx exp.(B')
+        @test D3 ≈ D4
+
+        fill!(D3, -1e3); fill!(D4, 9e9)
+        Bt = Transpose(B)
+        @. D3 = exp(Bt)
+        @avx @. D4 = exp(Bt)
+        @test D3 ≈ D4
     end
 end