Add limited ifelse reduction support

Author: chriselrod

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.12.65"
+version = "0.12.66"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -19,7 +19,7 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
-ArrayInterface = "3.1.9"
+ArrayInterface = "3.1.9 - 3.1.23"
 DocStringExtensions = "0.8"
 IfElse = "0.1"
 OffsetArrays = "1.4.1"
@@ -30,5 +30,5 @@ Static = "0.2, 0.3"
 StrideArraysCore = "0.1.12"
 ThreadingUtilities = "0.4.5"
 UnPack = "1"
-VectorizationBase = "0.20.32"
+VectorizationBase = "0.20.36"
 julia = "1.5"

src/codegen/lower_compute.jl

@@ -303,7 +303,7 @@ function parent_op_name!(
         parent = Symbol(parent, suffix_, '_', u)
       elseif u₂max > 1
         t = Expr(:tuple)
-        reduction = Expr(:call, GlobalRef(ArrayInterface, :reduce_tup), reduce_to_onevecunroll(instruction(opp)), t)
+        reduction = Expr(:call, GlobalRef(ArrayInterface, :reduce_tup), reduce_to_onevecunroll(opp), t)
         for u₂ ∈ 0:u₂max-1
           push!(t.args, Symbol(parent, u₂, "__", u))
         end
@@ -388,7 +388,14 @@ function reduce_parent!(q::Expr, ls::LoopSet, op::Operation, opp::Operation, par
     return parent
   end
   newp = gensym(parent)
-  push!(q.args, Expr(:(=), newp, Expr(:call, lv(reduction_to_scalar(reduct_class)), parent)))
+  if instruction(op).instr ≢ :ifelse
+    push!(q.args, Expr(:(=), newp, Expr(:call, lv(reduction_to_scalar(reduct_class)), parent)))#IfElseReducer
+  else
+    reductexpr = ifelse_reduction(:IfElseReducer,op) do opv
+      throw(LoopError("Does not support storing mirrored ifelse-reductions yet"))
+    end
+    push!(q.args, Expr(:(=), newp, Expr(:call, reductexpr, parent)))
+  end
   newp
 end
 function lower_compute!(
@@ -439,7 +446,7 @@ function lower_compute!(
             else
                 newpname = Symbol(newparentname, '_', u₁)
                 push!(q.args, Expr(:(=), newpname, Symbol(parentname, '_', u₁)))
-                reduce_expr!(q, newparentname, instruction(newparentop), u₁, -1, true, false)
+                reduce_expr!(q, newparentname, newparentop, u₁, -1, true, false)
                 push!(q.args, Expr(:(=), Symbol(newparentname, '_', 1), Symbol(newparentname, "##onevec##")))
             end
         end
@@ -542,7 +549,7 @@ function lower_compute!(
           # elseif parents_u₂syms[n] & (!u₂unrolledsym)
             #&& (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.
-          reduced_u₂ = reduce_expr_u₂(mangledvar(opp), instruction(opp), u₂max, Symbol("__", u₁))#ureduct(ls))
+            reduced_u₂ = reduce_expr_u₂(mangledvar(opp), opp, u₂max, Symbol("__", u₁))#ureduct(ls))
             reducedparentname = gensym!(ls, "reducedop")
             push!(q.args, Expr(:(=), reducedparentname, reduced_u₂))
             reduced_u₂ = reduce_parent!(q, ls, op, opp, reducedparentname)

src/codegen/lower_store.jl

@@ -5,51 +5,55 @@ function opisreduced(op::Operation)
   false
 end
 function storeinstr_preprend(op::Operation, vloopsym::Symbol)
-    # defaultstoreop = :vstore!
-    # defaultstoreop = :vnoaliasstore!
-    isvectorized(op) && return Symbol("")
-    vloopsym ∉ reduceddependencies(op) && return Symbol("")
-    # vectorized is not a loopdep, but is a reduced dep
-    opp::Operation = first(parents(op))
-    # while vectorized ∉ loopdependencies(opp)
-    while ((!isvectorized(opp)) || opisreduced(opp))
-        oppold = opp
-        for oppp ∈ parents(opp)
-            if vloopsym ∈ reduceddependencies(oppp)
-                @assert opp !== oppp "More than one parent is a reduction over the vectorized variable."
-                opp = oppp
-            end
-        end
-        @assert opp !== oppold "Failed to find any parents "
+  # defaultstoreop = :vstore!
+  # defaultstoreop = :vnoaliasstore!
+  isvectorized(op) && return Symbol("")
+  vloopsym ∉ reduceddependencies(op) && return Symbol("")
+  # vectorized is not a loopdep, but is a reduced dep
+  opp::Operation = first(parents(op))
+  # while vectorized ∉ loopdependencies(opp)
+  while ((!isvectorized(opp)) || opisreduced(opp))
+    oppold = opp
+    for oppp ∈ parents(opp)
+      if vloopsym ∈ reduceddependencies(oppp)
+        @assert opp !== oppp "More than one parent is a reduction over the vectorized variable."
+        opp = oppp
+      end
     end
-    reduction_to_scalar(reduction_instruction_class(instruction(opp)))
+    @assert opp !== oppold "Failed to find any parents "
+  end
+  instr = instruction(opp).instr
+  instr === :ifelse && return Symbol("") #throw(LoopError("ifelse not yet supported for inner reductions"))
+  return reduction_to_scalar(instr)
+  # if instr ≢ :ifelse
+  #   Expr(:(.), LoopVectorization, QuoteNode(reduction_to_scalar(instr))) # :IfElseReducer
+  # else
+  #   ifelse_reduction(:IfElseReducer,opp) do opv
+  #     throw(LoopError("Does not support storing mirrored ifelse-reductions yet"))
+  #   end
+  # end
 end
 
-function reduce_expr_u₂(toreduct::Symbol, instr::Instruction, u₂::Int, suffix::Symbol)
-    t = Expr(:tuple)
-    for u ∈ 0:u₂-1
-        push!(t.args, Symbol(toreduct, u, suffix))
-    end
-    Expr(:call, lv(:reduce_tup), reduce_to_onevecunroll(instr), t)
+function reduce_expr_u₂(toreduct::Symbol, op::Operation, u₂::Int, suffix::Symbol)
+  t = Expr(:tuple)
+  for u ∈ 0:u₂-1
+    push!(t.args, Symbol(toreduct, u, suffix))
+  end
+  Expr(:call, lv(:reduce_tup), reduce_to_onevecunroll(op), t)
 end
-function reduce_expr!(q::Expr, toreduct::Symbol, instr::Instruction, u₁::Int, u₂::Int, isu₁unrolled::Bool, isu₂unrolled::Bool)
-    if isu₂unrolled# u₂ != -1
-        _toreduct = Symbol(toreduct, 0)
-        push!(q.args, Expr(:(=), _toreduct, reduce_expr_u₂(toreduct, instr, u₂, Symbol(""))))
-    else#if u₂ == -1
-        _toreduct = Symbol(toreduct, '_', u₁)
-    # else
-        # _toreduct = Symbol(toreduct, 0)
-    end
-    # @show toreduct, _toreduct, u₁, u₂, isu₁unrolled, isu₂unrolled
-    if (u₁ == 1) | (~isu₁unrolled)
-        push!(q.args, Expr(:(=), Symbol(toreduct, "##onevec##"), _toreduct))
-    else
-        push!(q.args, Expr(:(=), Symbol(toreduct, "##onevec##"), Expr(:call, lv(reduction_to_single_vector(instr)), _toreduct)))
-        # push!(q.args, :(@show $_toreduct))
-        # push!(q.args, Expr(:(=), Symbol(toreduct, "##onevec##"), :(@show $(Expr(:call, lv(reduction_to_single_vector(instr)), _toreduct)))))
-    end
-    nothing
+function reduce_expr!(q::Expr, toreduct::Symbol, op::Operation, u₁::Int, u₂::Int, isu₁unrolled::Bool, isu₂unrolled::Bool)
+  if isu₂unrolled# u₂ != -1
+    _toreduct = Symbol(toreduct, 0)
+    push!(q.args, Expr(:(=), _toreduct, reduce_expr_u₂(toreduct, op, u₂, Symbol(""))))
+  else#if u₂ == -1
+    _toreduct = Symbol(toreduct, '_', u₁)
+  end
+  if (u₁ == 1) | (~isu₁unrolled)
+    push!(q.args, Expr(:(=), Symbol(toreduct, "##onevec##"), _toreduct))
+  else
+    push!(q.args, Expr(:(=), Symbol(toreduct, "##onevec##"), Expr(:call, reduction_to_single_vector(op), _toreduct)))
+  end
+  nothing
 end
 
 function lower_store_collection!(
@@ -278,7 +282,7 @@ function lower_tiled_store!(blockq::Expr, op::Operation, ls::LoopSet, ua::Unroll
     end
     opp = first(parents(op))
     if (opp.instruction.instr === reductfunc) && isone(length(parents(opp)))
-        throw("Operation $opp's instruction is $reductfunc, shouldn't be able to reach here.")
+        throw(LoopError("Operation $opp's instruction is $reductfunc, shouldn't be able to reach here."))
         # opp = only(parents(opp))
     end
     isu₁, isu₂ = isunrolled_sym(opp, u₁loopsym, u₂loopsym, vloopsym, ls)#, u₂)