support multiple stores on the same line, as well as storing symbols into multiple locations

Author: chriselrod

Project.toml

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

src/modeling/graphs.jl

@@ -1192,6 +1192,44 @@ function prepare_rhs_for_storage!(ls::LoopSet, RHS::Union{Symbol,Expr}, array, r
   return op
 end
 
+function unpack_tuple!(ls::LoopSet, LHS::Expr, RHS, elementbytes::Int, position::Int)
+  if Meta.isexpr(RHS, :tuple)
+    for i ∈ eachindex(LHS.args)
+      add_assignment!(ls, LHS.args[i], RHS.args[i], elementbytes, position)
+    end
+    return last(operations(ls)) # FIXME: dummy
+  end
+  @assert length(LHS.args) ≤ 14 "Functions returning more than 9 values aren't currently supported."
+  lhstemp = gensym!(ls, "lhstuple")
+  vparents = Operation[maybe_const_compute!(ls, lhstemp, add_operation!(ls, lhstemp, RHS, elementbytes, position), elementbytes, position)]
+  unpack_tuple!(ls, LHS, vparents, elementbytes, position)
+end
+
+function unpack_tuple!(ls::LoopSet, LHS::Expr, vparents::Vector{Operation}, elementbytes::Int, position::Int)
+  for i ∈ eachindex(LHS.args)
+    f = (:first,:second,:third,:fourth,:fifth,:sixth,:seventh,:eighth,:ninth,:tenth,:eleventh,:twelfth,:thirteenth,:last)[i]
+    lhsi = LHS.args[i]
+    if lhsi isa Symbol
+      add_compute!(ls, lhsi, f, vparents, elementbytes)
+      continue
+    elseif lhsi isa Expr
+      if lhsi.head === :ref
+        tempunpacksym = gensym!(ls, "tempunpack")
+        add_compute!(ls, tempunpacksym, f, vparents, elementbytes)
+        add_store_ref!(ls, tempunpacksym, lhsi, elementbytes)
+        continue
+      elseif lhsi.head === :tuple
+        lhstemp = gensym!(ls, "lhstuple")
+        op = add_compute!(ls, lhstemp, f, vparents, elementbytes)
+        unpack_tuple!(ls, lhsi, [op], elementbytes, position)
+        continue
+      end
+    end
+    throw(LoopError("Unpacking the above expression in the left hand side was not understood/supported.", lhsi))
+  end
+  first(vparents)
+end
+
 function add_assignment!(ls::LoopSet, LHS, RHS, elementbytes::Int, position::Int)
   if LHS isa Symbol
     if RHS isa Expr
@@ -1212,29 +1250,7 @@ function add_assignment!(ls::LoopSet, LHS, RHS, elementbytes::Int, position::Int
         add_store_ref!(ls, RHS, LHS, elementbytes)  # is this necessary? (Extension API?)
       end
     elseif LHS.head === :tuple
-      if RHS.head === :tuple
-        for i ∈ eachindex(LHS.args)
-          add_assignment!(ls, LHS.args[i], RHS.args[i], elementbytes, position)
-        end
-        return last(operations(ls)) # FIXME: dummy
-      end
-      @assert length(LHS.args) ≤ 14 "Functions returning more than 9 values aren't currently supported."
-      lhstemp = gensym!(ls, "lhstuple")
-      vparents = Operation[maybe_const_compute!(ls, lhstemp, add_operation!(ls, lhstemp, RHS, elementbytes, position), elementbytes, position)]
-      for i ∈ eachindex(LHS.args)
-        f = (:first,:second,:third,:fourth,:fifth,:sixth,:seventh,:eighth,:ninth,:tenth,:eleventh,:twelfth,:thirteenth,:last)[i]
-        lhsi = LHS.args[i]
-        if lhsi isa Symbol
-          add_compute!(ls, lhsi, f, vparents, elementbytes)
-        elseif lhsi isa Expr && lhsi.head === :ref
-          tempunpacksym = gensym!(ls, "tempunpack")
-          add_compute!(ls, tempunpacksym, f, vparents, elementbytes)
-          add_store_ref!(ls, tempunpacksym, lhsi, elementbytes)
-        else
-          throw(LoopError("Unpacking the above expression in the left hand side was not understood/supported.", lhsi))
-        end
-      end
-      first(vparents)
+      unpack_tuple!(ls, LHS, RHS, elementbytes, position)
     else
       throw(LoopError("LHS not understood; only `:ref`s and `:tuple`s are currently supported.", LHS))
     end

test/multiassignments.jl

@@ -20,9 +20,28 @@ function multiassign_turbo!(y, x)
 end
 multiassign_turbo(x) = multiassign_turbo!(similar(x, length(x)-3), x)
 
+multistorefunc(x) = exp(x), sincos(x)
+function multistore!(x, y, z, a)
+  @inbounds for i ∈ eachindex(x, y, z, a)
+    x[i], (y[i], z[i]) = multistorefunc(a[i])
+  end
+end
+function multistore_turbo!(x, y, z, a)
+  @turbo for i ∈ eachindex(x, y, z, a)
+    x[i], (y[i], z[i]) = multistorefunc(a[i])
+  end
+end
+
 @testset "Multiple assignments" begin
   @show @__LINE__
   x = rand(111);
-  @test multiassign(x) ≈ multiassign_turbo(x)  
+  @test multiassign(x) ≈ multiassign_turbo(x)
+  a0 = similar(x); b0 = similar(x); c0 = similar(x);
+  a1 = similar(x); b1 = similar(x); c1 = similar(x);
+  multistore!(a0, b0, c0, x)
+  multistore_turbo!(a1, b1, c1, x)
+  @test a0 ≈ a1
+  @test b0 ≈ b1
+  @test c0 ≈ c1
 end