let broadcast scoping

Author: chriselrod

Project.toml

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

src/LoopVectorization.jl

@@ -5,8 +5,12 @@ using Static: StaticInt, gt, static, Zero, One, reduce_tup
 using VectorizationBase,
   SLEEFPirates, UnPack, OffsetArrays, ArrayInterfaceOffsetArrays, ArrayInterfaceStaticArrays
 using LayoutPointers:
-  AbstractStridedPointer, StridedPointer, StridedBitPointer, grouped_strided_pointer,
-  stridedpointer_preserve, GroupedStridedPointers
+  AbstractStridedPointer,
+  StridedPointer,
+  StridedBitPointer,
+  grouped_strided_pointer,
+  stridedpointer_preserve,
+  GroupedStridedPointers
 import LayoutPointers
 
 using SIMDTypes: NativeTypes
@@ -114,8 +118,7 @@ using Base.Meta: isexpr
 using DocStringExtensions
 import LinearAlgebra # for check_args
 
-using Base:
-  unsafe_trunc
+using Base: unsafe_trunc
 
 using Base.FastMath:
   add_fast,

src/broadcast.jl

@@ -1,7 +1,7 @@
-@generated function append_true(::Val{D},::Val{N}) where {D,N}
+@generated function append_true(::Val{D}, ::Val{N}) where {D,N}
   length(D) == N && return D
   t = Expr(:tuple)
-  for d = D
+  for d in D
     push!(t.args, d)
   end
   for n = length(D)+1:N
@@ -12,14 +12,16 @@ end
 struct LowDimArray{D,T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
   data::A
   function LowDimArray{D}(data::A) where {D,T,N,A<:AbstractArray{T,N}}
-    new{append_true(Val{D}(),Val{N}()),T,N,A}(data)
+    new{append_true(Val{D}(), Val{N}()),T,N,A}(data)
   end
   function LowDimArray{D,T,N,A}(data::A) where {D,T,N,A<:AbstractArray{T,N}}
-    new{append_true(Val{D}(),Val{N}()),T,N,A}(data)
+    new{append_true(Val{D}(), Val{N}()),T,N,A}(data)
   end
 end
-function LowDimArray{D0}(data::LowDimArray{D1,T,N,A}) where {D0,T,N,D1,A<:AbstractArray{T,N}}
-  LowDimArray{map(|,D0,D1),T,N,A}(parent(data))
+function LowDimArray{D0}(
+  data::LowDimArray{D1,T,N,A},
+) where {D0,T,N,D1,A<:AbstractArray{T,N}}
+  LowDimArray{map(|, D0, D1),T,N,A}(parent(data))
 end
 Base.@propagate_inbounds Base.getindex(
   A::LowDimArray,
@@ -115,7 +117,9 @@ end
   end
   Expr(:block, Expr(:meta, :inline), staticexpr(Cnew))
 end
-function ArrayInterface.contiguous_axis(::Type{LowDimArrayForBroadcast{D,T,N,A}}) where {D,T,N,A}
+function ArrayInterface.contiguous_axis(
+  ::Type{LowDimArrayForBroadcast{D,T,N,A}},
+) where {D,T,N,A}
   ArrayInterface.contiguous_axis(A)
 end
 @inline function ArrayInterface.stride_rank(
@@ -180,8 +184,8 @@ function _strides_expr(@nospecialize(s), @nospecialize(x), R::Vector{Int}, D::Ve
   use_stride_acc = true
   stride_acc::Int = 1
   if is_column_major(R)
-  # elseif is_row_major(R)
-  #   Nrange = reverse(Nrange)
+    # elseif is_row_major(R)
+    #   Nrange = reverse(Nrange)
   else # not worth my time optimizing this case at the moment...
     # will write something generic stride-rank agnostic eventually
     use_stride_acc = false
@@ -323,14 +327,8 @@ function add_broadcast!(
   mA = gensym!(ls, "Aₘₖ")
   mB = gensym!(ls, "Bₖₙ")
   gf = GlobalRef(Core, :getfield)
-  pushprepreamble!(
-    ls,
-    Expr(:(=), mA, Expr(:(.), bcname, QuoteNode(:a))),
-  )
-  pushprepreamble!(
-    ls,
-    Expr(:(=), mB, Expr(:(.), bcname, QuoteNode(:b))),
-  )
+  pushprepreamble!(ls, Expr(:(=), mA, Expr(:(.), bcname, QuoteNode(:a))))
+  pushprepreamble!(ls, Expr(:(=), mB, Expr(:(.), bcname, QuoteNode(:b))))
   pushprepreamble!(ls, Expr(:(=), Klen, Expr(:call, gf, Expr(:call, :size, mB), 1, false)))
   pushpreamble!(ls, Expr(:(=), Krange, Expr(:call, :(:), staticexpr(1), Klen)))
   k = gensym!(ls, "k")
@@ -432,7 +430,7 @@ function add_broadcast!(
     pushprepreamble!(ls, Expr(:(=), bcname2, Expr(:call, forbroadcast, lda)))
     ArrayReference(bcname2, fulldims)
   end
-  
+
   loadop = add_simple_load!(ls, destname, ref, elementbytes, true)::Operation
   doaddref!(ls, loadop)
 end
@@ -486,10 +484,7 @@ function add_broadcast!(
   gf = GlobalRef(Core, :getfield)
   for (i, arg) ∈ enumerate(args)
     argname = gensym!(ls, "arg")
-    pushprepreamble!(
-      ls,
-      Expr(:(=), argname, Expr(:call, gf, bcargs, i, false)),
-    )
+    pushprepreamble!(ls, Expr(:(=), argname, Expr(:call, gf, bcargs, i, false)))
     # dynamic dispatch
     parent = add_broadcast!(
       ls,
@@ -542,7 +537,7 @@ end
   bc::BC,
   ::Val{Mod},
   ::Val{UNROLL},
-  ::Val{dontbc}
+  ::Val{dontbc},
 ) where {T<:NativeTypes,N,BC<:Union{Broadcasted,Product},Mod,UNROLL,dontbc}
   # 2 + 1
   # we have an N dimensional loop.
@@ -580,8 +575,16 @@ end
   bc::BC,
   ::Val{Mod},
   ::Val{UNROLL},
-  ::Val{dontbc}
-) where {T<:NativeTypes,N,A<:AbstractArray{T,N},BC<:Union{Broadcasted,Product},Mod,UNROLL,dontbc}
+  ::Val{dontbc},
+) where {
+  T<:NativeTypes,
+  N,
+  A<:AbstractArray{T,N},
+  BC<:Union{Broadcasted,Product},
+  Mod,
+  UNROLL,
+  dontbc,
+}
   # we have an N dimensional loop.
   # need to construct the LoopSet
   ls = LoopSet(Mod)
@@ -626,14 +629,14 @@ end
   bc::Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(identity),Tuple{T2}},
   ::Val{Mod},
   ::Val{UNROLL},
-  ::Val{dontbc}
+  ::Val{dontbc},
 ) where {T<:NativeTypes,N,T2<:Number,Mod,UNROLL,dontbc}
   inline, u₁, u₂, v, isbroadcast, W, rs, rc, cls, threads, warncheckarg, safe = UNROLL
   quote
     $(Expr(:meta, :inline))
     arg = T(first(bc.args))
     @turbo inline = $inline unroll = ($u₁, $u₂) thread = $threads vectorize = $v for i ∈
-      eachindex(
+                                                                                     eachindex(
       dest,
     )
       dest[i] = arg
@@ -646,7 +649,7 @@ end
   bc::Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(identity),Tuple{T2}},
   ::Val{Mod},
   ::Val{UNROLL},
-  ::Val{dontbc}
+  ::Val{dontbc},
 ) where {T<:NativeTypes,N,A<:AbstractArray{T,N},T2<:Number,Mod,UNROLL,dontbc}
   inline, u₁, u₂, v, isbroadcast, W, rs, rc, cls, threads, warncheckarg, safe = UNROLL
   quote
@@ -680,4 +683,4 @@ end
 end
 
 # vmaterialize!(dest, bc, ::Val, ::Val, ::StaticInt, ::StaticInt, ::StaticInt) =
-  # Base.Broadcast.materialize!(dest, bc)
+# Base.Broadcast.materialize!(dest, bc)

src/codegen/lower_threads.jl

@@ -157,10 +157,10 @@ end
 scale_cost(c) = @fastmath c * (Sys.ARCH === :x86_64 ? 0.0225 : 0.005625)
 scale_cost(c, looplen) = scale_cost(@fastmath c / looplen)
 @inline function choose_num_threads(
-    C::T,
-    NT::UInt,
-    x::Base.BitInteger,
-  ) where {T<:Union{Float32,Float64}}
+  C::T,
+  NT::UInt,
+  x::Base.BitInteger,
+) where {T<:Union{Float32,Float64}}
   _choose_num_threads(scale_cost(T(C)), NT, x)
 end
 @inline function _choose_num_threads(
@@ -529,12 +529,8 @@ function thread_one_loops_expr(
         end
       else# eliminate undef var errors that the compiler should be able to figure out are unreachable, but doesn't
         var"#torelease#tuple#" = (zero(UInt),)
-        var"#threads#tuple#" = (
-          PolyesterWeave.UnsignedIteratorEarlyStop(
-            zero(UInt),
-            0x00000000,
-          ),
-        )
+        var"#threads#tuple#" =
+          (PolyesterWeave.UnsignedIteratorEarlyStop(zero(UInt), 0x00000000),)
       end
     end
     var"#avx#call#args#" = $avxcall_args
@@ -808,10 +804,8 @@ function thread_two_loops_expr(
         end
       else# eliminate undef var errors that the compiler should be able to figure out are unreachable, but doesn't
         var"#torelease#tuple#" = (zero(UInt),)
-        var"#threads#tuple#" = PolyesterWeave.UnsignedIteratorEarlyStop(
-          zero(UInt),
-          0x00000000,
-        )
+        var"#threads#tuple#" =
+          PolyesterWeave.UnsignedIteratorEarlyStop(zero(UInt), 0x00000000)
       end
     end
     # @show $lastboundexpr

src/codegen/lowering.jl

@@ -543,7 +543,11 @@ end
   end
 end
 @inline function of_same_size(::Type{T}, ::Type{S}) where {T,S}
-  of_same_size(T, S, VectorizationBase.register_size() ÷ VectorizationBase.simd_integer_register_size())
+  of_same_size(
+    T,
+    S,
+    VectorizationBase.register_size() ÷ VectorizationBase.simd_integer_register_size(),
+  )
 end
 function outer_reduction_zero(
   op::Operation,

src/condense_loopset.jl

@@ -187,8 +187,8 @@ function parents_uint(op::Operation)
   N = length(opv)
   @assert N ≤ 32
   p0 = parents_uint(view(opv, 1:min(8, N)))
-  p1 = N > 8 ? parents_uint(view(opv, 9:min(16,N))) : zero(p0)
-  p2 = N > 16 ? parents_uint(view(opv, 17:min(24,N))) : zero(p0)
+  p1 = N > 8 ? parents_uint(view(opv, 9:min(16, N))) : zero(p0)
+  p2 = N > 16 ? parents_uint(view(opv, 17:min(24, N))) : zero(p0)
   p3 = N > 24 ? parents_uint(view(opv, 25:N)) : zero(p0)
   p0, p1, p2, p3
 end
@@ -361,8 +361,12 @@ end
 val(x) = Expr(:call, Expr(:curly, :Val, x))
 
 @inline gespf1(x, i) = gesp(x, i)
-@inline gespf1(x::StridedPointer{T,1}, i::Tuple{I}) where {T,I<:Union{Integer,StaticInt}} = gesp(x, i)
-@inline gespf1(x::StridedBitPointer{T,1}, i::Tuple{I}) where {T,I<:Union{Integer,StaticInt}} = gesp(x, i)
+@inline gespf1(x::StridedPointer{T,1}, i::Tuple{I}) where {T,I<:Union{Integer,StaticInt}} =
+  gesp(x, i)
+@inline gespf1(
+  x::StridedBitPointer{T,1},
+  i::Tuple{I},
+) where {T,I<:Union{Integer,StaticInt}} = gesp(x, i)
 @inline gespf1(x::StridedPointer{T,1}, i::Tuple{Zero}) where {T} = x
 @inline gespf1(x::StridedBitPointer{T,1}, i::Tuple{Zero}) where {T} = x
 @generated function gespf1(
@@ -571,7 +575,7 @@ end
     register_size(),
     available_registers(),
     lv_max_num_threads(),
-    cache_linesize()
+    cache_linesize(),
   )
 end
 function find_samename_constparent(op::Operation, opname::Symbol)
@@ -924,7 +928,7 @@ can_turbo(::typeof(Base.literal_pow), ::Val{3}) = true
 can_turbo(::typeof(Base.FastMath.pow_fast), ::Val{2}) = true
 
 for f ∈ (convert, reinterpret, trunc, unsafe_trunc, round, ceil, floor)
-    @eval can_turbo(::typeof($f), ::Val{2}) = true
+  @eval can_turbo(::typeof($f), ::Val{2}) = true
 end
 
 """

src/constructors.jl

@@ -56,7 +56,7 @@ function substitute_broadcast(
 )
   ci = first(Meta.lower(LoopVectorization, q).args).code
   nargs = length(ci) - 1
-  ex = Expr(:block)
+  lb = Expr(:block)
   syms = Vector{Symbol}(undef, nargs)
   configarg = (inline, u₁, u₂, v, true, threads, warncheckarg, safe)
   unroll_param_tup = Expr(:call, lv(:avx_config_val), :(Val{$configarg}()), staticexpr(0))
@@ -70,17 +70,21 @@ function substitute_broadcast(
       ciₙargs = ciₙ.args
       f = first(ciₙargs)
       if ciₙ.head === :(=)
-        push!(ex.args, Expr(:(=), f, syms[((ciₙargs[2])::Core.SSAValue).id]))
+        push!(lb.args, Expr(:(=), f, syms[((ciₙargs[2])::Core.SSAValue).id]))
       elseif isglobalref(f, Base, :materialize!)
-        add_ci_call!(ex, lv(:vmaterialize!), ciₙargs, syms, n, unroll_param_tup, mod)
+        add_ci_call!(lb, lv(:vmaterialize!), ciₙargs, syms, n, unroll_param_tup, mod)
       elseif isglobalref(f, Base, :materialize)
-        add_ci_call!(ex, lv(:vmaterialize), ciₙargs, syms, n, unroll_param_tup, mod)
+        add_ci_call!(lb, lv(:vmaterialize), ciₙargs, syms, n, unroll_param_tup, mod)
       else
-        add_ci_call!(ex, f, ciₙargs, syms, n)
+        add_ci_call!(lb, f, ciₙargs, syms, n)
       end
     end
   end
-  esc(ex)
+  ret::Expr = pop!(lb.args)
+  if Meta.isexpr(ret, :(=), 2)
+    ret = ret.args[2]
+  end
+  esc(Expr(:let, lb, Expr(:block, ret)))
 end
 
 

src/getconstindexes.jl

@@ -37,7 +37,8 @@ const EXTRACTFUNS = (
   :thirtysixth,
   :thirtyseventh,
   :thirtyeighth,
-  :last)
+  :last,
+)
 
 for (i, f) ∈ enumerate(EXTRACTFUNS)
   (i == 1 || i == length(EXTRACTFUNS)) && continue

src/modeling/costs.jl

@@ -266,7 +266,7 @@ const COST = Dict{Symbol,InstructionCost}(
   :significand => InstructionCost(8, 1.0),
 )
 
-for f = EXTRACTFUNS
+for f in EXTRACTFUNS
   COST[f] = InstructionCost(0, 0.0, 0.0, 0)
 end
 
@@ -351,7 +351,8 @@ end
   VectorizationBase.fmap(ier, VectorizationBase.data(a), VectorizationBase.data(b)),
 )
 
-@inline (iec::IfElseCollapser)(a) = VectorizationBase.contract(IfElseOp(iec.f), a, StaticInt{1}())
+@inline (iec::IfElseCollapser)(a) =
+  VectorizationBase.contract(IfElseOp(iec.f), a, StaticInt{1}())
 @inline (iec::IfElseCollapser)(a, ::StaticInt{C}) where {C} =
   VectorizationBase.contract(IfElseOp(iec.f), a, StaticInt{C}())
 

src/modeling/graphs.jl

@@ -509,12 +509,7 @@ end
 available_registers() =
   ifelse(has_opmask_registers(), register_count(), register_count() - One())
 function set_hw!(ls::LoopSet)
-  set_hw!(
-    ls,
-    Int(register_size()),
-    Int(available_registers()),
-    Int(cache_linesize())
-  )
+  set_hw!(ls, Int(register_size()), Int(available_registers()), Int(cache_linesize()))
 end
 reg_size(ls::LoopSet) = ls.register_size
 reg_count(ls::LoopSet) = ls.register_count
@@ -924,7 +919,7 @@ function add_block!(ls::LoopSet, ex::Expr, elementbytes::Int, position::Int)
 end
 function makestatic!(expr)
   expr isa Expr || return expr
-  for i = eachindex(expr.args)
+  for i in eachindex(expr.args)
     ex = expr.args[i]
     if ex isa Int
       expr.args[i] = staticexpr(ex)
@@ -1481,7 +1476,7 @@ function add_operation!(
     add_comparison!(ls, LHS_sym, RHS, elementbytes, position)
   else
     throw(LoopError("Expression not recognized.", RHS))
-  end  
+  end
 end
 
 function prepare_rhs_for_storage!(