Apply masks to more selfop-funs

Author: chriselrod

Project.toml

@@ -30,5 +30,5 @@ Static = "0.2, 0.3"
 StrideArraysCore = "0.1.12"
 ThreadingUtilities = "0.4.5"
 UnPack = "1"
-VectorizationBase = "0.20.27"
+VectorizationBase = "0.20.31"
 julia = "1.5"

src/codegen/lower_compute.jl

@@ -234,6 +234,7 @@ end
             push!(q.args, :($gf(vargs, $k, false)))
         end
         return Expr(:block, Expr(:meta, :inline), q)
+        # return Expr(:block, Expr(:meta, :inline), :(@show($q)))
     end
     if Sreduced
         M = N
@@ -272,6 +273,7 @@ end
         push!(t.args, :($gf(dd, $m, false)))
     end
     push!(q.args, :(VecUnroll($t)))
+    # push!(q.args, :(@show(VecUnroll($t))))
     q
 end
 
@@ -560,7 +562,8 @@ function lower_compute!(
             end
         end
     end
-    selfdepreduce = ifelse(((!u₁unrolledsym) & isu₁unrolled(op)) & (u₁ > 1), selfdep, 0)
+  selfdepreduce = ifelse(((!u₁unrolledsym) & isu₁unrolled(op)) & (u₁ > 1), selfdep, 0)
+  # @show selfdepreduce, selfdep, maskreduct, op
     if maskreduct
         ifelsefunc = if us.u₁ == 1
             :ifelse # don't need to be fancy
@@ -577,7 +580,7 @@ function lower_compute!(
                 insert!(instrcall.args, 4, staticexpr(u₁))
                 insert!(instrcall.args, 5, staticexpr(selfdepreduce))
             end
-        elseif all(in(loopdependencies(op)), reduceddeps) || any(opp -> mangledvar(opp) === mangledvar(op), parents_op)
+        elseif all(in(loopdependencies(op)), reduceddeps) || selfdep ≠ 0#any(opp -> mangledvar(opp) === mangledvar(op), parents_op)
             # Here, we are evaluating the function, and then `ifelse`-ing it with `hasf == false`.
             # That means we still need to adjust the `instrcall` in case we're reducing/accumulating across the unroll
             if ifelsefunc ≡ :ifelse # ifelse means it's unrolled by 1, no need
@@ -593,8 +596,8 @@ function lower_compute!(
                 push!(q.args, Expr(:(=), varsym, Expr(:call, lv(:ifelse), MASKSYMBOL, instrcall, selfopname)))
             end
             return
-        elseif selfdep != 0
-            make_partial_map!(instrcall, selfopname, u₁, selfdepreduce)
+        # elseif selfdep != 0
+        #   make_partial_map!(instrcall, selfopname, u₁, selfdepreduce)
         end
     elseif selfdep != 0 && (dopartialmap ||
         (isouterreduct && (opunrolled) && (u₁ < us.u₁)) ||

src/modeling/costs.jl

@@ -269,7 +269,9 @@ const COST = Dict{Symbol,InstructionCost}(
   :convert => InstructionCost(4,0.5),
   :vpermilps177 => InstructionCost(1, 1.0),
   :vmovsldup => InstructionCost(1, 1.0),
-  :vmovshdup => InstructionCost(1, 1.0)
+  :vmovshdup => InstructionCost(1, 1.0),
+  :exponent => InstructionCost(8, 1.0),
+  :significand => InstructionCost(8, 1.0)
 )
 
 # # @inline prefetch0(x::Ptr, i) = VectorizationBase.prefetch(x, Val{3}(), Val{0}())
@@ -596,7 +598,7 @@ const FUNCTIONSYMBOLS = IdDict{Type{<:Function},Instruction}(
     typeof(ifelse) => :ifelse,
     typeof(identity) => :identity,
     typeof(conj) => :identity,#conj,
-    typeof(÷) => :div_fast
+    typeof(÷) => :vdiv_fast
     # typeof(zero) => :zero,
     # typeof(one) => :one,
     # typeof(axes) => :axes,

src/reconstruct_loopset.jl

@@ -711,7 +711,7 @@ Execute an `@turbo` block. The block's code is represented via the arguments:
 @aggressive_constprop @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#")))
   # return @show avx_body(ls, var"#UNROLL#")

test/outer_reductions.jl

@@ -54,6 +54,48 @@ function test_awmean(::Type{T}) where {T}
     end
 end
 
+function logℒ_fast(α, β, t, c, x)
+    eα = abs(α)
+    n, k = size(x)
+  
+    (n == length(t) == length(c) && length(β) == k + 1) || throw(DimensionMismatch())
+    s = zero(typeof(α))
+    @inbounds for i in 1:n
+      xb = sum(@inbounds(x[i, j] * β[j+1]) for j in 1:k) + β[1]
+      ti = t[i]
+      s += (1 - (c[i] == ti)) * (log(eα) + (eα - 1) * log(ti) + xb) - ti^eα * exp(xb)
+    end
+    return s
+end
+function logℒ_fast_turbo(α, β, t, c, x)
+    eα = abs(α)
+    n, k = size(x)
+  
+    (n == length(t) == length(c) && length(β) == k + 1) || throw(DimensionMismatch())
+    s = zero(typeof(α))
+    @turbo for i in 1:n
+      xb0 = 0.0
+      for j in 1:k
+        xb0 += x[i,j] * β[j+1]
+      end
+      xb = xb0 + β[1]
+      ti = t[i]
+      s += (1 - (c[i] == ti)) * (log(eα) + (eα - 1) * log(ti) + xb) - ti^eα * exp(xb)
+    end
+    return s
+end
+
+function test_logℒ(n, k)
+  t = rand(n)
+  c = copy(t);
+  b = rand(n) .> 0.5
+  c[b] .= rand.();
+  x = rand(n,k)
+  α = 2.85
+  β = rand(k+1)
+  @test logℒ_fast(α, β, t, c, x) ≈ logℒ_fast_turbo(α, β, t, c, x)
+end
+
 function not_an_outer_reduct!(r, N::Int, x = 2.0, y= nothing) # there was a bug where this was classified as one
   @turbo for i ∈ eachindex(r)
     acc = y === nothing ? x : r[i]
@@ -69,5 +111,8 @@ end
     test_awmean(T)
   end
   @test all(==(7.4), not_an_outer_reduct!(Vector{Float64}(undef, 5), 17, 7.4))
+  for n ∈ 1:20, k ∈ 1:5
+    test_logℒ(n,k)
+  end
 end