Add loopdeps and reduced children to new parents when removing ifelse.

Author: chriselrod

src/codegen/lower_constant.jl

@@ -103,7 +103,7 @@ function lower_constant!(
   !opu₂ && suffix > 0 && return
   instr = instruction(op)
   instr.mod === GLOBALCONSTANT && return
-  constsym = instr.instr
+  constsym = constantopname(op)# instr.instr
   reducedchildvectorized = vloopsym ∈ reducedchildren(op)
   if reducedchildvectorized || isvectorized(op) || vloopsym ∈ reduceddependencies(op) || should_broadcast_op(op)
     # call = Expr(:call, lv(:vbroadcast), W, Expr(:call, lv(:maybeconvert), typeT, constsym))

src/condense_loopset.jl

@@ -483,12 +483,26 @@ function split_ifelse!(
   falseops = operations(lsfalse)
   true_op = parents(true_ops[k])[2]
   falseop = parents_op[3]
+  true_op.dependencies = loopdependencies(op)
+  falseop.dependencies = loopdependencies(op)
+  true_op.reduced_children = reducedchildren(op)
+  falseop.reduced_children = reducedchildren(op)
   condop_count = 0
   for i ∈ eachindex(falseops)
     fop = falseops[i]
     parents_false = parents(fop)
     for (j,opp) ∈ enumerate(parents_false)
       if opp === op # then ops[i]'s jth parent is the ifelse
+        # These reduction to scalar instructions are added for non-outer reductions initialized with non-constant ops
+        # So we check if now 
+        # if (j == 2) && (Base.sym_in(instruction(fop).instr, (:reduced_add, :reduced_prod, :reduced_max, :reduced_min, :reduced_all, :reduced_any)))
+          # if isconstantop(true_op)
+          #   (true_ops[i]).instruction = Instruction(:identity)
+          # end
+          # if isconstantop(falseop)
+          #   fop.instruction = Instruction(:identity)
+          # end
+        # end
         parents(true_ops[i])[j] = true_op
         parents_false[j] = falseop
       end
@@ -513,12 +527,10 @@ end
 function generate_call_split(
   ls::LoopSet, preserve::Vector{Symbol}, shouldindbyind::Vector{Bool}, roots::Vector{Bool}, extra_args::Expr, inlineu₁u₂::Tuple{Bool,Int8,Int8}, thread::UInt, debug::Bool
 )
-  if !debug
-    for (k,op) ∈ enumerate(operations(ls))
-      parents_op = parents(op)
-      if (iscompute(op) && (instruction(op).instr === :ifelse)) && (length(parents_op) == 3) && isconstantop(first(parents_op))
-        return split_ifelse!(ls, preserve, shouldindbyind, roots, extra_args, k, inlineu₁u₂, thread, debug)
-      end
+  for (k,op) ∈ enumerate(operations(ls))
+    parents_op = parents(op)
+    if (iscompute(op) && (instruction(op).instr === :ifelse)) && (length(parents_op) == 3) && isconstantop(first(parents_op))
+      return split_ifelse!(ls, preserve, shouldindbyind, roots, extra_args, k, inlineu₁u₂, thread, debug)
     end
   end
   return generate_call_types(ls, preserve, shouldindbyind, roots, extra_args, inlineu₁u₂, thread, debug)

src/modeling/costs.jl

@@ -474,6 +474,23 @@ function reduction_zero(x::Float64)
         throw("Reduction not found.")
     end
 end
+function reduction_zero_class(x::Symbol)::Float64
+  if x === :one
+    MULTIPLICATIVE_IN_REDUCTIONS
+  elseif x === :typemin
+    MAX
+  elseif x === :typemax
+    MIN
+  elseif x === :max_mask
+    ALL
+  elseif x === :zero_mask
+    ANY
+  elseif x === :zero#sorted last, as should go into preamble_zeros
+    ADDITIVE_IN_REDUCTIONS
+  else
+    throw("Reduction not found.")
+  end
+end
 reduction_zero(x) = reduction_zero(reduction_instruction_class(x))
 
 function isreductcombineinstr(instr::Symbol)

src/modeling/graphs.jl

@@ -1028,13 +1028,13 @@ function add_operation!(
         f = first(RHS.args)
         if f === :getindex
             add_load_getindex!(ls, LHS, RHS, elementbytes)
-        elseif f === :zero || f === :one
+        elseif f isa Symbol && Base.sym_in(f, (:zero, :one, :typemin, :typemax))
             c = gensym!(ls, f)
             op = add_constant!(ls, c, ls.loopsymbols[1:position], LHS, elementbytes, :numericconstant)
             if f === :zero
                 push!(ls.preamble_zeros, (identifier(op), IntOrFloat))
             else
-                push!(ls.preamble_funcofeltypes, (identifier(op), MULTIPLICATIVE_IN_REDUCTIONS))
+                push!(ls.preamble_funcofeltypes, (identifier(op), reduction_zero_class(f)))
             end
             op
         else
@@ -1070,14 +1070,14 @@ function add_operation!(
         f = first(RHS.args)
         if f === :getindex
             add_load!(ls, LHS_sym, LHS_ref, elementbytes)
-        elseif f === :zero || f === :one
+        elseif f isa Symbol && Base.sym_in(f, (:zero, :one, :typemin, :typemax))
             c = gensym!(ls, f)
             op = add_constant!(ls, c, ls.loopsymbols[1:position], LHS_sym, elementbytes, :numericconstant)
             # op = add_constant!(ls, c, Symbol[], LHS_sym, elementbytes, :numericconstant)
             if f === :zero
                 push!(ls.preamble_zeros, (identifier(op), IntOrFloat))
             else
-                push!(ls.preamble_funcofeltypes, (identifier(op), MULTIPLICATIVE_IN_REDUCTIONS))
+                push!(ls.preamble_funcofeltypes, (identifier(op), reduction_zero_class(f)))
             end
             op
         else

test/offsetarrays.jl

@@ -196,15 +196,19 @@ using LoopVectorization: Static
        end
        out
     end
-    function avxgeneric2!(out, A, kern)
-        @avx for I in CartesianIndices(out)
-            tmp = zero(eltype(out))
-            for J in CartesianIndices(kern)
-                tmp += A[I+J]*kern[J]
-            end
-            out[I] = tmp
+    function avxgeneric2!(out, A, kern, keep = nothing)
+      @avx for I in CartesianIndices(out)
+        tmp = if keep === nothing
+          zero(eltype(out))
+        else
+          out[I]
         end
-        out
+        for J in CartesianIndices(kern)
+          tmp += A[I+J]*kern[J]
+        end
+        out[I] = tmp
+      end
+      out
     end
     function pparent(x) # go through nested parents
         px = parent(x)

test/runtests.jl

@@ -1,6 +1,6 @@
 include("testsetup.jl")
 
-import InteractiveUtils
+import InteractiveUtils, Aqua
 
 InteractiveUtils.versioninfo(stdout; verbose = true)
 
@@ -12,7 +12,7 @@ const START_TIME = time()
 
 @time @testset "LoopVectorization.jl" begin
 
-    # @time Aqua.test_all(LoopVectorization)
+    @time Aqua.test_all(LoopVectorization)
     # @test isempty(detect_unbound_args(LoopVectorization))
 
     @time include("printmethods.jl")

test/testsetup.jl

@@ -1,4 +1,4 @@
-using Test, Aqua
+using Test
 using LoopVectorization
 
 using LinearAlgebra