Support chained comparisons, fixes #313 (tested), fix key error, which should fix #312 (untested)

Author: chriselrod

Project.toml

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

src/modeling/graphs.jl

@@ -1050,9 +1050,25 @@ strip_op_linenumber_nodes(q::Expr) = only(filter(x -> !isa(x, LineNumberNode), q
 function add_operation!(ls::LoopSet, LHS::Symbol, RHS::Symbol, elementbytes::Int, position::Int)
   add_constant!(ls, RHS, ls.loopsymbols[1:position], LHS, elementbytes)
 end
+function add_comparison!(ls::LoopSet, LHS::Symbol, RHS::Expr, elementbytes::Int, position::Int)
+  Nargs = length(RHS.args)
+  @assert (Nargs ≥ 5) & isodd(Nargs)
+  p1 = add_assignment!(ls, gensym!(ls, "leftcmp"), RHS.args[1], elementbytes, position)::Operation
+  p2 = add_assignment!(ls, gensym!(ls, "middlecmp"), RHS.args[3], elementbytes, position)::Operation
+  cmpname = Nargs == 3 ? LHS : gensym!(ls, "cmp")
+  cmp = add_compute!(ls, cmpname, RHS.args[2], Operation[p1, p2], elementbytes)::Operation
+  for i ∈ 5:2:Nargs
+    pnew = add_assignment!(ls, gensym!(ls, "rightcmp"), RHS.args[i], elementbytes, position)::Operation
+    cmpchain = add_compute!(ls, gensym!(ls, "cmpchain"), RHS.args[i-1], Operation[p2, pnew], elementbytes)::Operation
+    cmpname = Nargs == i ? LHS : gensym!(ls, "cmp")
+    cmp = add_compute!(ls, cmpname, :&, [cmp, cmpchain], elementbytes)::Operation
+    p2 = pnew
+  end
+  return cmp
+end
 function add_operation!(
   ls::LoopSet, LHS::Symbol, RHS::Expr, elementbytes::Int, position::Int
-  )
+)
   if RHS.head === :ref
     add_load_ref!(ls, LHS, RHS, elementbytes)
   elseif RHS.head === :call
@@ -1083,6 +1099,8 @@ function add_operation!(
     # op = add_constant!(ls, c, ls.loopsymbols[1:position], LHS, elementbytes, :numericconstant)
     # pushpreamble!(ls, op, c)
     # op
+  elseif Meta.isexpr(RHS, :comparison)
+    add_comparison!(ls, LHS, RHS, elementbytes, position)
   else
     throw(LoopError("Expression not recognized.", RHS))
   end
@@ -1125,6 +1143,8 @@ function add_operation!(
     # op = add_constant!(ls, c, ls.loopsymbols[1:position], LHS_sym, elementbytes, :numericconstant)
     # pushpreamble!(ls, op, c)
     # op
+  elseif Meta.isexpr(RHS, :comparison, 5)
+    add_comparison!(ls, LHS, RHS, elementbytes, position)
   else
     throw(LoopError("Expression not recognized.", RHS))
   end
@@ -1141,6 +1161,7 @@ function prepare_rhs_for_storage!(ls::LoopSet, RHS::Union{Symbol,Expr}, array, r
   mpref.parents = cachedparents
   op = add_store!(ls, mpref, elementbytes)
   ls.syms_aliasing_refs[findfirst(==(mpref.mref), ls.refs_aliasing_syms)] = lrhs
+  ls.opdict[lrhs] = op
   return op
 end
 

src/parse/add_loads.jl

@@ -32,9 +32,7 @@ function add_load!(ls::LoopSet, op::Operation, actualarray::Bool = true)
         end
       end
     end
-    if allmatch
-      return isstore(opp) ? getop(ls, first(parents(opp))) : opp
-    end
+    allmatch && isstore(opp) ? first(parents(opp)) : opp
   end
   add_vptr!(ls, op.ref.ref.array, vptr(op), actualarray)
   pushop!(ls, op, name(op))

src/parse/add_stores.jl

@@ -1,20 +1,20 @@
 function add_unique_store!(ls::LoopSet, op::Operation)
-    add_vptr!(ls, op)
-    pushop!(ls, op, name(op.ref))
+  add_vptr!(ls, op)
+  pushop!(ls, op, name(op.ref))
 end
 function cse_store!(ls::LoopSet, op::Operation)
-    id = identifier(op)
-    ls.operations[id] = op
-    ls.opdict[op.variable] = op
+  id = identifier(op)
+  ls.operations[id] = op
+  ls.opdict[op.variable] = op
 end
 function add_store!(ls::LoopSet, op::Operation, add_pvar::Bool = !any(r -> r == op.ref, ls.refs_aliasing_syms))
-    @assert isstore(op)
-    if add_pvar
-        push!(ls.syms_aliasing_refs, name(first(parents(op))))
-        push!(ls.refs_aliasing_syms, op.ref)
-    end
-    id = op.identifier
-    id == length(operations(ls)) ? add_unique_store!(ls, op) : cse_store!(ls, op)
+  @assert isstore(op)
+  if add_pvar
+    push!(ls.syms_aliasing_refs, name(first(parents(op))))
+    push!(ls.refs_aliasing_syms, op.ref)
+  end
+  id = op.identifier
+  id == length(operations(ls)) ? add_unique_store!(ls, op) : cse_store!(ls, op)
 end
 function add_copystore!(
     ls::LoopSet, parent::Operation, mpref::ArrayReferenceMetaPosition, elementbytes::Int
@@ -89,7 +89,7 @@ function add_conditional_store!(ls::LoopSet, LHS, condop::Operation, storeop::Op
     id = length(ls.operations)
     @assert pvar ∉ ls.syms_aliasing_refs
     # if pvar ∉ ls.syms_aliasing_refs
-    # FIXME properly handle CSE of conditional stores.
+  # FIXME properly handle CSE of conditional stores.
     push!(ls.syms_aliasing_refs, pvar)
     push!(ls.refs_aliasing_syms, mref)
     storeparents = [storeop, condop]

test/ifelsemasks.jl

@@ -428,6 +428,46 @@ T = Float32
         return w
     end
 
+
+  function mwe_v(dest, src, lut)
+    # for i in eachindex(src)
+    LoopVectorization.@turbo for i in eachindex(src)
+      v = src[i]
+      s = 128
+      s2 = s | 64
+      s = ifelse(lut[s2] <= v, s2, s)
+      s2 = s | 32
+      s = ifelse(lut[s2] <= v, s2, s)
+      dest[i] = s
+    end
+    return dest
+  end
+
+  function mwe_s(dest, src, lut)
+    for i in eachindex(src)
+      # LoopVectorization.@turbo for i in eachindex(src)
+      v = src[i]
+      s = 128
+      s2 = s | 64
+      s = ifelse(lut[s2] <= v, s2, s)
+      s2 = s | 32
+      s = ifelse(lut[s2] <= v, s2, s)
+      dest[i] = s
+    end
+    return dest
+  end
+
+  function turbocomparison!(m)
+    @turbo for i in eachindex(m)
+      m[i] = ifelse(0<m[i]<0.5, 0.0, m[i])
+    end
+  end
+  function turbocomparison!(m, y)
+    @turbo for i in eachindex(m)
+      m[i] = ifelse(0<m[i]<y[i]<0.5, 0.0, m[i])
+    end
+  end
+
     N = 117
     for T ∈ (Float32, Float64, Int32, Int64)
         @show T, @__LINE__
@@ -612,5 +652,21 @@ T = Float32
         @test barycentric_weight3(X) ≈ bX
         @test barycentric_weight4(X) ≈ bX
     end
+
+  let
+    lut = let x = cumsum(rand(Float32, 256)./128); x[end] = Inf; x end;
+    src = rand(Float32, N);
+
+    @test mwe_v(Vector{Int}(undef, N), src, lut) == mwe_s(Vector{Int}(undef, N), src, lut)
+    @test mwe_v(Vector{Int32}(undef, N), src, lut) == mwe_s(Vector{Int32}(undef, N), src, lut)
+  end
+
+  let m = rand(25, 25), y = rand(25,25), baseline5 = (@. ifelse(0 < m < 0.5, 0.0, m)), baseline7 = @. ifelse(0 < y < m < 0.5, 0.0, y)
+    turbocomparison!(y, m)
+    @test y == baseline7
+    turbocomparison!(m)
+    @test m == baseline5
+  end
+
 end