ifelse reduct fixes

Author: chriselrod

Project.toml

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

src/codegen/lower_threads.jl

@@ -213,25 +213,38 @@ end
 end
 
 function outer_reduct_combine_expressions(ls::LoopSet, retv)
-    gf = GlobalRef(Core, :getfield)
-    q = Expr(:block, :(var"#load#thread#ret#" = $gf(ThreadingUtilities.load(var"#thread#ptr#", typeof($retv), $(reg_size(ls))),2,false)))
-    # push!(q.args, :(@show var"#load#thread#ret#"))
-    for (i,or) ∈ enumerate(ls.outer_reductions)
-        op = ls.operations[or]
-        var = name(op)
-        mvar = mangledvar(op)
-        instr = instruction(op)
-        out = Symbol(mvar, "##onevec##")
-        instrcall = Expr(:call, lv(reduce_to_onevecunroll(instr)))
-        push!(instrcall.args, Expr(:call, lv(:vecmemaybe), out))
-        if length(ls.outer_reductions) > 1
-            push!(instrcall.args, Expr(:call, lv(:vecmemaybe), Expr(:call, gf, Symbol("#load#thread#ret#"), i, false)))
-        else
-            push!(instrcall.args, Expr(:call, lv(:vecmemaybe), Symbol("#load#thread#ret#")))
+  gf = GlobalRef(Core, :getfield)
+  q = Expr(:block, :(var"#load#thread#ret#" = $gf(ThreadingUtilities.load(var"#thread#ptr#", typeof($retv), $(reg_size(ls))),2,false)))
+  # push!(q.args, :(@show var"#load#thread#ret#"))
+  for (i,or) ∈ enumerate(ls.outer_reductions)
+    op = ls.operations[or]
+    var = name(op)
+    mvar = mangledvar(op)
+    instr = instruction(op)
+    out = Symbol(mvar, "##onevec##")
+    instrcall::Expr = if instr.instr ≢ :ifelse
+      Expr(:call, lv(reduce_to_onevecunroll(instr)))
+    else
+      reductexpr::Expr = let ls = ls  #FIXME: this should be tested...
+        ifelse_reduction(:IfElseReduced, op) do opv
+          @assert length(ls.outer_reductions) > 1
+          j = findfirst(==(identifier(opv)), ls.outer_reductions)
+          otherarg = Expr(:call, lv(:vecmemaybe), Expr(:call, GlobalRef(Core,:getfield), Symbol("#load#thread#ret#"), j, false))
+          Expr(:call, lv(:vecmemaybe), Symbol(mangledvar(opv), "##onevec##")), (otherarg,)
         end
-        push!(q.args, Expr(:(=), out, Expr(:call, :data, instrcall)))
+      end
+      Expr(:call, reductexpr)
     end
-    q
+    push!(instrcall.args, Expr(:call, lv(:vecmemaybe), out))
+    if length(ls.outer_reductions) > 1
+      push!(instrcall.args, Expr(:call, lv(:vecmemaybe), Expr(:call, gf, Symbol("#load#thread#ret#"), i, false)))
+    else
+      push!(instrcall.args, Expr(:call, lv(:vecmemaybe), Symbol("#load#thread#ret#")))
+    end
+    push!(q.args, Expr(:(=), out, Expr(:call, :data, instrcall)))
+    # push!(q.args, Expr(:(=), out, :(@show $data($instrcall))))
+  end
+  q
 end
 
 function thread_loop_summary!(ls::LoopSet, ua::UnrollArgs, threadedloop::Loop, issecondthreadloop::Bool)

src/reconstruct_loopset.jl

@@ -713,7 +713,7 @@ Execute an `@turbo` block. The block's code is represented via the arguments:
 @generated function _turbo_!(
   ::Val{var"#UNROLL#"}, ::Val{var"#OPS#"}, ::Val{var"#ARF#"}, ::Val{var"#AM#"}, ::Val{var"#LPSYM#"}, ::Val{Tuple{var"#LB#",var"#V#"}}, var"#flattened#var#arguments#"::Vararg{Any,var"#num#vargs#"}
 ) where {var"#UNROLL#", var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#", var"#V#", var"#num#vargs#"}
-  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 = _turbo_loopset(var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#".parameters, var"#V#".parameters, var"#UNROLL#")
   pushfirst!(ls.preamble.args, :(var"#lv#tuple#args#" = reassemble_tuple(Tuple{var"#LB#",var"#V#"}, var"#flattened#var#arguments#")))
   post = hoist_constant_memory_accesses!(ls)

test/ifelsemasks.jl

@@ -468,6 +468,14 @@ T = Float32
     end
   end
 
+  function absmax_tturbo(a) # LV threaded
+    result=zero(eltype(a))
+    @tturbo for i in 1:length(a)
+      abs(a[i])>result && (result=abs(a[i]))
+    end
+    result
+  end
+
   function findminturbo(x)
     indmin = 0
     minval = typemax(eltype(x))
@@ -478,6 +486,16 @@ T = Float32
     end
     minval, indmin
   end
+  function findmintturbo(x)
+    indmin = 0
+    minval = typemax(eltype(x))
+    @tturbo for i ∈ eachindex(x)
+      newmin = x[i] < minval
+      minval = newmin ? x[i] : minval
+      indmin = newmin ?   i  : indmin
+    end
+    minval, indmin
+  end
   function findminturbo_u4(x)
     indmin = 0
     minval = typemax(eltype(x))
@@ -497,9 +515,21 @@ T = Float32
       mv, mi = findminturbo(a)
       mv2, mi2 = findminturbo_u4(a)
       @test mv == a[mi] == minimum(a) == mv2 == a[mi2]
+      for n in 1000:1000:10_000
+        x = rand(-T(100):T(100), n);
+        @test absmax_tturbo(x) == mapreduce(abs, max, x)
+        mv, mi = findmintturbo(x)
+        @test mv == x[mi] == minimum(x)
+      end
     else
       a = rand(T, N); b = rand(T, N);
       @test findmin(a) == findminturbo(a) == findminturbo_u4(a)
+      for n in 1000:1000:10_000
+        x = randn(T, n);
+        @test absmax_tturbo(x) == mapreduce(abs, max, x)
+        mv, mi = findmintturbo(x)
+        @test mv == x[mi] == minimum(x)
+      end
     end;
     c1 = similar(a); c2 = similar(a);
 

test/outer_reductions.jl

@@ -1,3 +1,4 @@
+
 function reg_term(omega, B = size(omega,2); alpha=0.01)
   reg = 0.0
   for b in 1:B
@@ -199,5 +200,6 @@ end
   end
   omega = rand(87,87);
   @test reg_term(omega) ≈ reg_term_turbo(omega)
+
 end