spmd -> brrr, and sometimes init reductions with loads rather than zeros when that load would otherwise be added at the end

Author: chriselrod

src/LoopVectorization.jl

@@ -51,7 +51,7 @@ using Requires
 
 
 export LowDimArray, stridedpointer, indices,
-    @avx, @avxt, @spmd, @spmdt, *ˡ, _avx_!,
+    @avx, @avxt, @brrr, @tbrrr, *ˡ, _avx_!,
     vmap, vmap!, vmapt, vmapt!, vmapnt, vmapnt!, vmapntt, vmapntt!,
     tanh_fast, sigmoid_fast,
     vfilter, vfilter!, vmapreduce, vreduce

src/codegen/lower_compute.jl

@@ -525,7 +525,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), ureduct(ls))
+            reduced_u₂ = reduce_expr_u₂(mangledvar(opp), instruction(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_constant.jl

@@ -42,13 +42,11 @@ function typeof_sym(ls::LoopSet, op::Operation, zerotyp::NumberType)
         ELTYPESYMBOL
     end
 end
-# function in_reduced_children(op::Operation, s::Symbol)
-# end
 
 function lower_zero!(
     q::Expr, op::Operation, ls::LoopSet, ua::UnrollArgs, zerotyp::NumberType = zerotype(ls, op)
 )
-  @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, u₂max, suffix = ua
+  @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, vloop, u₂max, suffix = ua
   mvar, opu₁, opu₂ = variable_name_and_unrolled(op, u₁loopsym, u₂loopsym, vloopsym, suffix, ls)
   !opu₂ && suffix > 0 && return
   # TODO: for u₁, needs to consider if reducedchildren are u₁-unrolled

src/codegen/lower_store.jl

@@ -19,17 +19,17 @@ function storeinstr_preprend(op::Operation, vloopsym::Symbol)
     reduction_to_scalar(reduction_instruction_class(instruction(opp)))
 end
 
-function reduce_expr_u₂(toreduct::Symbol, instr::Instruction, u₂::Int)
+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))
+        push!(t.args, Symbol(toreduct, u, suffix))
     end
     Expr(:call, lv(:reduce_tup), reduce_to_onevecunroll(instr), 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₂)))
+        push!(q.args, Expr(:(=), _toreduct, reduce_expr_u₂(toreduct, instr, u₂, Symbol(""))))
     else#if u₂ == -1
         _toreduct = Symbol(toreduct, '_', u₁)
     # else

src/codegen/operation_evaluation_order.jl

@@ -56,40 +56,94 @@ function set_upstream_family!(adal::Vector{T}, op::Operation, val::T, ld::Vector
         set_upstream_family!(adal, opp, val, ld, id)
     end
 end
-
+function search_for_reductinit!(op::Operation, opswap::Operation, var::Symbol, loopdeps::Vector{Symbol})
+  for (i,opp) ∈ enumerate(parents(op))
+    if (name(opp) === var) && (length(reduceddependencies(opp)) == 0) && (length(loopdependencies(opp)) == length(loopdeps)) && (length(children(opp)) == 1)
+      if all(in(loopdeps), loopdependencies(opp))
+        parents(op)[i] = opswap
+        return opp
+      end
+    end
+    opcheck = search_for_reductinit!(opp, opswap, var, loopdeps)
+    opcheck === opp || return opcheck
+  end
+  return op
+end
 function addoptoorder!(
     ls::LoopSet, included_vars::Vector{Bool}, place_after_loop::Vector{Bool}, op::Operation,
     loopsym::Symbol, _n::Int, u₁loop::Symbol, u₂loop::Symbol, vectorized::Symbol, u₂max::Int
 )
-    lo = ls.loop_order
-    id = identifier(op)
-    included_vars[id] || return nothing
-    loopsym ∈ loopdependencies(op) || return nothing
-    for opp ∈ parents(op) # ensure parents are added first
-        addoptoorder!(ls, included_vars, place_after_loop, opp, loopsym, _n, u₁loop, u₂loop, vectorized, u₂max)
-    end
-    included_vars[id] || return nothing
-    included_vars[id] = false
-    isunrolled = (isu₁unrolled(op)) + 1
-    istiled = isu₂unrolled(op) + 1
-    # optype = Int(op.node_type) + 1
-    after_loop = place_after_loop[id] + 1
-    if !isloopvalue(op)
-        isnopidentity(ls, op, u₁loop, u₂loop, vectorized, u₂max) || push!(lo[isunrolled,istiled,after_loop,_n], op)
-        # if istiled
-        #     isnopidentity(ls, op, u₁loop, u₂loop, vectorized, u₂max, u₂max) || push!(lo[isunrolled,2,after_loop,_n], op)
-        # else
-        #     isnopidentity(ls, op, u₁loop, u₂loop, vectorized, u₂max, nothing) || push!(lo[isunrolled,1,after_loop,_n], op)
-        # end
+  lo = ls.loop_order
+  id = identifier(op)
+  included_vars[id] || return nothing
+  loopsym ∈ loopdependencies(op) || return nothing
+  for opp ∈ parents(op) # ensure parents are added first
+    addoptoorder!(ls, included_vars, place_after_loop, opp, loopsym, _n, u₁loop, u₂loop, vectorized, u₂max)
+  end
+  included_vars[id] || return nothing
+  included_vars[id] = false
+  isunrolled = (isu₁unrolled(op)) + 1
+  istiled = isu₂unrolled(op) + 1
+  # optype = Int(op.node_type) + 1
+  after_loop = place_after_loop[id] + 1
+  if !isloopvalue(op)
+    isnopidentity(ls, op, u₁loop, u₂loop, vectorized, u₂max) || push!(lo[isunrolled,istiled,after_loop,_n], op)
+    # if istiled
+    #     isnopidentity(ls, op, u₁loop, u₂loop, vectorized, u₂max, u₂max) || push!(lo[isunrolled,2,after_loop,_n], op)
+    # else
+    #     isnopidentity(ls, op, u₁loop, u₂loop, vectorized, u₂max, nothing) || push!(lo[isunrolled,1,after_loop,_n], op)
+    # end
+  end
+  # @show op, after_loop
+  # isloopvalue(op) || push!(lo[isunrolled,istiled,after_loop,_n], op)
+  # all(opp -> iszero(length(reduceddependencies(opp))), parents(op)) &&
+  set_upstream_family!(place_after_loop, op, false, loopdependencies(op), identifier(op)) # parents that have already been included are not moved, so no need to check included_vars to filter
+  nothing
+end
+function replace_reduct_init!(ls::LoopSet, op::Operation, opsub::Operation, opcheck::Operation)
+  deleteat!(parents(op), 2)
+  op.variable = opcheck.variable
+  opsub.variable = opcheck.variable
+  op.mangledvariable = opcheck.mangledvariable
+  opsub.mangledvariable = opcheck.mangledvariable
+  op.instruction = instruction(:identity)
+  fill_children!(ls)
+end
+function nounrollreduction(op::Operation, u₁loop::Symbol, u₂loop::Symbol, vectorized::Symbol)
+  reduceddeps = reduceddependencies(op)
+  (vectorized ∉ reduceddeps) &&
+    (u₁loop ∉ reduceddeps) &&
+    (u₂loop ∉ reduceddeps)
+end
+function load_short_static_reduction_first!(ls::LoopSet, u₁loop::Symbol, u₂loop::Symbol, vectorized::Symbol)
+  for op ∈ operations(ls)
+    iscompute(op) || continue
+    length(reduceddependencies(op)) == 0 && continue
+    if (instruction(op).instr === :reduced_add)
+      vecloop = getloop(ls, vectorized)
+      if isstaticloop(vecloop) && (length(vecloop) ≤ 16) && nounrollreduction(op, u₁loop, u₂loop, vectorized)
+        opsub = parents(op)[2]
+        length(children(opsub)) == 1 || continue
+        opsearch = parents(op)[1]
+        opcheck = search_for_reductinit!(opsearch, opsub, name(opsearch), loopdependencies(op))
+        opcheck === opsearch || replace_reduct_init!(ls, op, opsub, opcheck)
+        
+      end
+    elseif (instruction(op).instr === :add_fast) && (instruction(first(parents(op))).instr === :identity)
+      vecloop = getloop(ls, vectorized)
+      if isstaticloop(vecloop) && (length(vecloop) ≤ 16) && nounrollreduction(op, u₁loop, u₂loop, vectorized)
+        opsub = parents(op)[2]
+        ((length(reduceddependencies(opsub)) == 0) & (length(children(opsub)) == 1)) || continue
+        opsearch = parents(op)[1]
+        opcheck = search_for_reductinit!(opsearch, opsub, name(opsearch), loopdependencies(op))
+        opcheck === opsearch || replace_reduct_init!(ls, op, opsub, opcheck)
+      end
     end
-    # @show op, after_loop
-    # isloopvalue(op) || push!(lo[isunrolled,istiled,after_loop,_n], op)
-    # all(opp -> iszero(length(reduceddependencies(opp))), parents(op)) &&
-    set_upstream_family!(place_after_loop, op, false, loopdependencies(op), identifier(op)) # parents that have already been included are not moved, so no need to check included_vars to filter
-    nothing
+  end
 end
 
 function fillorder!(ls::LoopSet, order::Vector{Symbol}, u₁loop::Symbol, u₂loop::Symbol, u₂max::Int, vectorized::Symbol)
+  load_short_static_reduction_first!(ls, u₁loop, u₂loop, vectorized)
     lo = ls.loop_order
     resize!(lo, length(ls.loopsymbols))
     ro = lo.loopnames # reverse order; will have same order as lo

src/constructors.jl

@@ -111,7 +111,7 @@ function process_args(args; inline = false, check_empty = false, u₁ = zero(Int
     end
     inline, check_empty, u₁, u₂, threads
 end
-function spmd_macro(mod, src, q, args...)
+function brrr_macro(mod, src, q, args...)
     q = macroexpand(mod, q)
 
     if q.head === :for
@@ -124,12 +124,12 @@ function spmd_macro(mod, src, q, args...)
     end
 end
 """
-    @spmd
+    @brrr
 
 Annotate a `for` loop, or a set of nested `for` loops whose bounds are constant across iterations, to optimize the computation. For example:
 
     function AmulB!(C, A, B)
-        @spmd for m ∈ 1:size(A,1), n ∈ 1:size(B,2)
+        @brrr for m ∈ 1:size(A,1), n ∈ 1:size(B,2)
             Cₘₙ = zero(eltype(C))
             for k ∈ 1:size(A,2)
                 Cₘₙ += A[m,k] * B[k,n]
@@ -148,23 +148,23 @@ julia> using LoopVectorization
 
 julia> a = rand(100);
 
-julia> b = @spmd exp.(2 .* a);
+julia> b = @brrr exp.(2 .* a);
 
 julia> c = similar(b);
 
-julia> @spmd @. c = exp(2a);
+julia> @brrr @. c = exp(2a);
 
 julia> b ≈ c
 true
 ```
 
 # Extended help
 
-Advanced users can customize the implementation of the `@spmd`-annotated block
+Advanced users can customize the implementation of the `@brrr`-annotated block
 using keyword arguments:
 
 ```
-@spmd inline=false unroll=2 body
+@brrr inline=false unroll=2 body
 ```
 
 where `body` is the code of the block (e.g., `for ... end`).
@@ -197,42 +197,42 @@ but it applies to the loop ordering and unrolling that will be chosen by LoopVec
 `uᵢ=0` (the default) indicates that LoopVectorization should pick its own value,
 and `uᵢ=-1` disables unrolling for the correspond loop.
 
-The `@spmd` macro also checks the array arguments using `LoopVectorization.check_args` to try and determine
+The `@brrr` macro also checks the array arguments using `LoopVectorization.check_args` to try and determine
 if they are compatible with the macro. If `check_args` returns false, a fall back loop annotated with `@inbounds`
 and `@fastmath` is generated. Note that `VectorizationBase` provides functions such as `vadd` and `vmul` that will
-ignore `@fastmath`, preserving IEEE semantics both within `@spmd` and `@fastmath`.
+ignore `@fastmath`, preserving IEEE semantics both within `@brrr` and `@fastmath`.
 `check_args` currently returns false for some wrapper types like `LinearAlgebra.UpperTriangular`, requiring you to
 use their `parent`. Triangular loops aren't yet supported.
 """
-macro spmd(args...)
-    spmd_macro(__module__, __source__, last(args), Base.front(args)...)
+macro brrr(args...)
+    brrr_macro(__module__, __source__, last(args), Base.front(args)...)
 end
 """
-Equivalent to `@spmd`, except it adds `thread=true` as the first keyword argument.
+Equivalent to `@brrr`, except it adds `thread=true` as the first keyword argument.
 Note that later arguments take precendence.
 
-Meant for convenience, as `@spmdt` is shorter than `@spmd thread=true`.
+Meant for convenience, as `@tbrrr` is shorter than `@brrr thread=true`.
 """
-macro spmdt(args...)
-    spmd_macro(__module__, __source__, last(args), :(thread=true), Base.front(args)...)
+macro tbrrr(args...)
+    brrr_macro(__module__, __source__, last(args), :(thread=true), Base.front(args)...)
 end
 
 """
-    @_spmd
+    @_brrr
 
-This macro mostly exists for debugging/testing purposes. It does not support many of the use cases of [`@spmd`](@ref).
+This macro mostly exists for debugging/testing purposes. It does not support many of the use cases of [`@brrr`](@ref).
 It emits loops directly, rather than punting to an `@generated` function, meaning it doesn't have access to type
 information when generating code or analyzing the loops, often leading to bad performance.
 
-This macro accepts the `inline` and `unroll` keyword arguments like `@spmd`, but ignores the `check_empty` argument.
+This macro accepts the `inline` and `unroll` keyword arguments like `@brrr`, but ignores the `check_empty` argument.
 """
-macro _spmd(q)
+macro _brrr(q)
     q = macroexpand(__module__, q)
     ls = LoopSet(q, __module__)
     set_hw!(ls)
     esc(Expr(:block, ls.prepreamble, lower_and_split_loops(ls, -1)))
 end
-macro _spmd(arg, q)
+macro _brrr(arg, q)
     @assert q.head === :for
     q = macroexpand(__module__, q)
     inline, check_empty, u₁, u₂ = check_macro_kwarg(arg, false, false, zero(Int8), zero(Int8), 1)
@@ -241,14 +241,14 @@ macro _spmd(arg, q)
     esc(Expr(:block, ls.prepreamble, lower(ls, u₁ % Int, u₂ % Int, -1)))
 end
 
-macro spmd_debug(q)
+macro brrr_debug(q)
     q = macroexpand(__module__, q)
     ls = LoopSet(q, __module__)
     esc(LoopVectorization.setup_call_debug(ls))
 end
 
 # define aliases
-const var"@avx" = var"@spmd"
-const var"@avxt" = var"@spmdt"
-const var"@avx_debug" = var"@spmd_debug"
+const var"@avx" = var"@brrr"
+const var"@avxt" = var"@tbrrr"
+const var"@avx_debug" = var"@brrr_debug"