Add warn_check_args=true option to @turbo so that it can emit a warning when LoopVectorization.check_args fails. Resolves #296

Author: chriselrod

src/broadcast.jl

@@ -398,7 +398,7 @@ end
     # need to construct the LoopSet
     # @show typeof(dest)
     ls = LoopSet(Mod)
-    inline, u₁, u₂, isbroadcast, W, rs, rc, cls, l1, l2, l3, threads = UNROLL
+    inline, u₁, u₂, isbroadcast, W, rs, rc, cls, l1, l2, l3, threads, warncheckarg = UNROLL
     set_hw!(ls, rs, rc, cls, l1, l2, l3)
     ls.isbroadcast = isbroadcast # maybe set `false` in a DiffEq-like `@..` macro
     loopsyms = [gensym!(ls, "n") for n ∈ 1:N]
@@ -409,7 +409,7 @@ end
     doaddref!(ls, storeop)
     resize!(ls.loop_order, num_loops(ls)) # num_loops may be greater than N, eg Product
   # return ls
-  sc = setup_call(ls, :(Base.Broadcast.materialize!(dest, bc)), LineNumberNode(0), inline, false, u₁, u₂, threads%Int)
+    sc = setup_call(ls, :(Base.Broadcast.materialize!(dest, bc)), LineNumberNode(0), inline, false, u₁, u₂, threads%Int, warncheckarg)
   # return sc
     Expr(:block, Expr(:meta,:inline), sc, :dest)
 end
@@ -419,7 +419,7 @@ end
     # we have an N dimensional loop.
     # need to construct the LoopSet
     ls = LoopSet(Mod)
-    inline, u₁, u₂, isbroadcast, W, rs, rc, cls, l1, l2, l3, threads = UNROLL
+    inline, u₁, u₂, isbroadcast, W, rs, rc, cls, l1, l2, l3, threads, warncheckarg = UNROLL
     set_hw!(ls, rs, rc, cls, l1, l2, l3)
     ls.isbroadcast = isbroadcast # maybe set `false` in a DiffEq-like `@..` macro
     loopsyms = [gensym!(ls, "n") for n ∈ 1:N]
@@ -430,7 +430,7 @@ end
     storeop = add_simple_store!(ls, :dest, ArrayReference(:dest, reverse(loopsyms)), elementbytes)
     doaddref!(ls, storeop)
     resize!(ls.loop_order, num_loops(ls)) # num_loops may be greater than N, eg Product
-    Expr(:block, Expr(:meta,:inline), setup_call(ls, :(Base.Broadcast.materialize!(dest′, bc)), LineNumberNode(0), inline, false, u₁, u₂, threads%Int), :dest′)
+    Expr(:block, Expr(:meta,:inline), setup_call(ls, :(Base.Broadcast.materialize!(dest′, bc)), LineNumberNode(0), inline, false, u₁, u₂, threads%Int, warncheckarg), :dest′)
 end
 # these are marked `@inline` so the `@turbo` itself can choose whether or not to inline.
 @generated function vmaterialize!(

src/condense_loopset.jl

@@ -525,10 +525,11 @@ end
     ::Val{CNFARG}, ::StaticInt{W}, ::StaticInt{RS}, ::StaticInt{AR}, ::StaticInt{NT},
     ::StaticInt{CLS}, ::StaticInt{L1}, ::StaticInt{L2}, ::StaticInt{L3}
 ) where {CNFARG,W,RS,AR,CLS,L1,L2,L3,NT}
-    inline,u₁,u₂,BROADCAST,thread = CNFARG
-    nt = min(thread % UInt, NT % UInt)
-    t = Expr(:tuple, inline, u₁, u₂, BROADCAST, W, RS, AR, CLS, L1,L2,L3, nt)
-    Expr(:call, Expr(:curly, :Val, t))
+  inline,u₁,u₂,BROADCAST,thread = CNFARG
+  nt = min(thread % UInt, NT % UInt)
+  t = Expr(:tuple, inline, u₁, u₂, BROADCAST, W, RS, AR, CLS, L1,L2,L3, nt)
+  length(CNFARG) == 6 && push!(t.args, last(CNFARG))
+  Expr(:call, Expr(:curly, :Val, t))
 end
 @inline function avx_config_val(
     ::Val{CNFARG}, ::StaticInt{W}
@@ -787,18 +788,20 @@ function setup_call_debug(ls::LoopSet)
   generate_call(ls, (false,zero(Int8),zero(Int8)), zero(UInt), true)
 end
 function setup_call(
-    ls::LoopSet, q::Expr, source::LineNumberNode, inline::Bool, check_empty::Bool, u₁::Int8, u₂::Int8, thread::Int
+  ls::LoopSet, q::Expr, source::LineNumberNode, inline::Bool, check_empty::Bool, u₁::Int8, u₂::Int8, thread::Int, warncheckarg::Bool
 )
-    # We outline/inline at the macro level by creating/not creating an anonymous function.
-    # The old API instead was based on inlining or not inline the generated function, but
-    # the generated function must be inlined into the initial loop preamble for performance reasons.
-    # Creating an anonymous function and calling it also achieves the outlining, while still
-    # inlining the generated function into the loop preamble.
-    lnns = extract_all_lnns(q)
-    pushfirst!(lnns, source)
-    call = generate_call(ls, (inline, u₁, u₂), thread%UInt, false)
-    call = check_empty ? check_if_empty(ls, call) : call
-    pushprepreamble!(ls, Expr(:if, check_args_call(ls), call, make_crashy(make_fast(q))))
-    prepend_lnns!(ls.prepreamble, lnns)
-    return ls.prepreamble
+  # We outline/inline at the macro level by creating/not creating an anonymous function.
+  # The old API instead was based on inlining or not inline the generated function, but
+  # the generated function must be inlined into the initial loop preamble for performance reasons.
+  # Creating an anonymous function and calling it also achieves the outlining, while still
+  # inlining the generated function into the loop preamble.
+  lnns = extract_all_lnns(q)
+  pushfirst!(lnns, source)
+  call = generate_call(ls, (inline, u₁, u₂), thread%UInt, false)
+  call = check_empty ? check_if_empty(ls, call) : call
+  argfailure = make_crashy(make_fast(q))
+  warncheckarg && (argfailure = Expr(:block, :(@warn "`LoopVectorization.check_args` on your inputs failed; running fallback `@inbounds @fastmath` loop instead."  maxlog=1), argfailure))
+  pushprepreamble!(ls, Expr(:if, check_args_call(ls), call, argfailure))
+  prepend_lnns!(ls.prepreamble, lnns)
+  return ls.prepreamble
 end

src/constructors.jl

@@ -35,12 +35,12 @@ function add_ci_call!(q::Expr, @nospecialize(f), args, syms, i, valarg = nothing
     push!(q.args, Expr(:(=), syms[i], call))
 end
 
-function substitute_broadcast(q::Expr, mod::Symbol, inline, u₁, u₂, threads)
+function substitute_broadcast(q::Expr, mod::Symbol, inline, u₁, u₂, threads, warncheckarg)
     ci = first(Meta.lower(LoopVectorization, q).args).code
     nargs = length(ci)-1
     ex = Expr(:block)
     syms = [gensym() for _ ∈ 1:nargs]
-    configarg = (inline,u₁,u₂,true,threads)
+    configarg = (inline,u₁,u₂,true,threads,warncheckarg)
     unroll_param_tup = Expr(:call, lv(:avx_config_val), :(Val{$configarg}()), staticexpr(0))
     for n ∈ 1:nargs
         ciₙ = ci[n]
@@ -75,7 +75,7 @@ function loopset(q::Expr) # for interactive use only
     ls
 end
 
-function check_macro_kwarg(arg, inline::Bool, check_empty::Bool, u₁::Int8, u₂::Int8, threads::Int)
+function check_macro_kwarg(arg, inline::Bool, check_empty::Bool, u₁::Int8, u₂::Int8, threads::Int, warncheckarg::Bool)
     ((arg.head === :(=)) && (length(arg.args) == 2)) || throw(ArgumentError("macro kwarg should be of the form `argname = value`."))
     kw = (arg.args[1])::Symbol
     value = (arg.args[2])
@@ -100,27 +100,29 @@ function check_macro_kwarg(arg, inline::Bool, check_empty::Bool, u₁::Int8, u
         else
             throw(ArgumentError("Don't know how to process argument in `thread=$value`."))
         end
+    elseif kw === :warn_check_args
+        warncheckarg = value::Bool
     else
         throw(ArgumentError("Received unrecognized keyword argument $kw. Recognized arguments include:\n`inline`, `unroll`, `check_empty`, and `thread`."))
     end
-    inline, check_empty, u₁, u₂, threads
+    inline, check_empty, u₁, u₂, threads, warncheckarg
 end
-function process_args(args; inline = false, check_empty = false, u₁ = zero(Int8), u₂ = zero(Int8), threads = 1)
+function process_args(args; inline = false, check_empty = false, u₁ = zero(Int8), u₂ = zero(Int8), threads = 1, warncheckarg = false)
     for arg ∈ args
-        inline, check_empty, u₁, u₂, threads = check_macro_kwarg(arg, inline, check_empty, u₁, u₂, threads)
+        inline, check_empty, u₁, u₂, threads, warncheckarg = check_macro_kwarg(arg, inline, check_empty, u₁, u₂, threads, warncheckarg)
     end
-    inline, check_empty, u₁, u₂, threads
+    inline, check_empty, u₁, u₂, threads, warncheckarg
 end
 function turbo_macro(mod, src, q, args...)
     q = macroexpand(mod, q)
 
     if q.head === :for
         ls = LoopSet(q, mod)
-        inline, check_empty, u₁, u₂, threads = process_args(args)
-        esc(setup_call(ls, q, src, inline, check_empty, u₁, u₂, threads))
+        inline, check_empty, u₁, u₂, threads, warncheckarg = process_args(args)
+        esc(setup_call(ls, q, src, inline, check_empty, u₁, u₂, threads, warncheckarg))
     else
-        inline, check_empty, u₁, u₂, threads = process_args(args, inline=true)
-        substitute_broadcast(q, Symbol(mod), inline, u₁, u₂, threads)
+        inline, check_empty, u₁, u₂, threads, warncheckarg = process_args(args, inline=true)
+        substitute_broadcast(q, Symbol(mod), inline, u₁, u₂, threads, warncheckarg)
     end
 end
 """
@@ -209,6 +211,10 @@ and `@fastmath` is generated. Note that `VectorizationBase` provides functions s
 ignore `@fastmath`, preserving IEEE semantics both within `@turbo` 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.
+
+Setting the keyword argument `warn_check_args=true` (e.g. `@turbo warn_check_args=true for ...`) in a loop or
+broadcast statement will cause it to warn once if `LoopVectorization.check_args` fails and the fallback
+loop is executed instead of the LoopVectorization-optimized loop.
 """
 macro turbo(args...)
     turbo_macro(__module__, __source__, last(args), Base.front(args)...)
@@ -250,7 +256,7 @@ end
 macro _turbo(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)
+  inline, check_empty, u₁, u₂ = check_macro_kwarg(arg, false, false, zero(Int8), zero(Int8), 1, false)
   ls = LoopSet(q, __module__)
   set_hw!(ls)
   def_outer_reduct_types!(ls)

src/modeling/graphs.jl

@@ -989,12 +989,13 @@ end
 
 
 function maybe_const_compute!(ls::LoopSet, LHS::Symbol, op::Operation, elementbytes::Int, position::Int)
-    # return op
-    if iscompute(op) && iszero(length(loopdependencies(op)))
-        ls.opdict[LHS] = add_constant!(ls, LHS, ls.loopsymbols[1:position], gensym!(ls, instruction(op).instr), elementbytes, :numericconstant)
-    else
-        op
-    end
+  # return op
+  if iscompute(op) && iszero(length(loopdependencies(op)))
+    ls.opdict[LHS] = add_constant!(ls, LHS, ls.loopsymbols[1:position], gensym!(ls, instruction(op).instr), elementbytes, :numericconstant)
+  else
+    # op.dependencies = ls.loopsymbols[1:position]
+    op
+  end
 end
 strip_op_linenumber_nodes(q::Expr) = only(filter(x -> !isa(x, LineNumberNode), q.args))
 

src/parse/add_compute.jl

@@ -183,10 +183,12 @@ function add_reduction_update_parent!(
   parent::Operation, instr::Instruction, reduction_ind::Int, elementbytes::Int
 )
   var = name(parent)
+  # isouterreduction = iszero(length(loopdependencies(parent))) && (parent.instruction === LOOPCONSTANT)
   isouterreduction = parent.instruction === LOOPCONSTANT
   # @show instr, vparents, parent, reduction_ind
   # if parent is not an outer reduction...
   # if !isouterreduction && !isreductzero(parent, ls, reduct_zero)
+  # @show isouterreduction, parent, length(loopdependencies(parent))
   add_reduct_instruct = !isouterreduction && !isconstant(parent)
   if add_reduct_instruct
     if instr.instr === :ifelse
@@ -226,7 +228,7 @@ function add_reduction_update_parent!(
   update_reduction_status!(vparents, reduceddeps, name(reductinit))
   # this is the op added by add_compute
   op = Operation(length(operations(ls)), reductsym, elementbytes, instr, compute, deps, reduceddeps, vparents)
-  parent.instruction === LOOPCONSTANT && push!(ls.outer_reductions, identifier(op))
+  isouterreduction && push!(ls.outer_reductions, identifier(op))
   opout = pushop!(ls, op, var) # note this overwrites the entry in the operations dict, but not the vector
   # isouterreduction || iszero(length(reduceddeps)) && return opout
   # return opout

src/parse/add_constants.jl

@@ -62,7 +62,7 @@ end
 # assignedsym will be assigned to value within the preamble
 function add_constant!(
     ls::LoopSet, value::Symbol, deps::Vector{Symbol}, assignedsym::Symbol, elementbytes::Int, f::Symbol = Symbol("")
-)
+  )
     retop = get(ls.opdict, value, nothing)
     if retop === nothing
         op = Operation(length(operations(ls)), assignedsym, elementbytes, Instruction(f, value), constant, deps, NODEPENDENCY, NOPARENTS)

src/parse/add_ifelse.jl

@@ -5,37 +5,49 @@
 negateop!(ls::LoopSet, condop::Operation, elementbytes::Int) = add_compute!(ls, gensym!(ls, "negated#mask"), :~, [condop], elementbytes)
 
 function add_if!(ls::LoopSet, LHS::Symbol, RHS::Expr, elementbytes::Int, position::Int, mpref::Union{Nothing,ArrayReferenceMetaPosition} = nothing)
-    # for now, just simple 1-liners
-    @assert length(RHS.args) == 3 "if statements without an else cannot be assigned to a variable."
-    condition = first(RHS.args)
-    condop = if condition isa Symbol
-        getop(ls, condition, elementbytes)
-    elseif mpref === nothing
-        add_operation!(ls, gensym!(ls, "mask"), condition, elementbytes, position)
-    else
-        add_operation!(ls, gensym!(ls, "mask"), condition, mpref, elementbytes, position)
+  # for now, just simple 1-liners
+  @assert length(RHS.args) == 3 "if statements without an else cannot be assigned to a variable."
+  condition = first(RHS.args)
+  condop = if condition isa Symbol
+    getop(ls, condition, elementbytes)
+  elseif mpref === nothing
+    add_operation!(ls, gensym!(ls, "mask"), condition, elementbytes, position)
+  else
+    add_operation!(ls, gensym!(ls, "mask"), condition, mpref, elementbytes, position)
+  end
+  iftrue = RHS.args[2]
+  if iftrue isa Expr
+    trueop = add_operation!(ls, gensym!(ls, "iftrue"), iftrue, elementbytes, position)
+    if iftrue.head === :ref && all(ld -> ld ∈ loopdependencies(trueop), loopdependencies(condop)) && !search_tree(parents(condop), trueop)
+      trueop.instruction = Instruction(:conditionalload)
+      push!(parents(trueop), condop)
     end
-    iftrue = RHS.args[2]
-    if iftrue isa Expr
-        trueop = add_operation!(ls, gensym!(ls, "iftrue"), iftrue, elementbytes, position)
-        if iftrue.head === :ref && all(ld -> ld ∈ loopdependencies(trueop), loopdependencies(condop)) && !search_tree(parents(condop), trueop)
-            trueop.instruction = Instruction(:conditionalload)
-            push!(parents(trueop), condop)
-        end
-    else
-        trueop = getop(ls, iftrue, elementbytes)
+  else
+    trueop = getop(ls, iftrue, elementbytes)
+  end
+  iffalse = RHS.args[3]
+  if trueop.instruction === Instruction(:conditionalload)
+    if ((iffalse isa Number) && (iffalse == 0)) || (Meta.isexpr(iffalse, :call, 2) && (iffalse.args[1] === :zero))
+      trueop.variable = LHS
+      ls.opdict[LHS] = trueop
+      return trueop
     end
-    iffalse = RHS.args[3]
-    if iffalse isa Expr
-        falseop = add_operation!(ls, gensym!(ls, "iffalse"), iffalse, elementbytes, position)
-        if iffalse.head === :ref && all(ld -> ld ∈ loopdependencies(falseop), loopdependencies(condop)) && !search_tree(parents(condop), falseop)
-            falseop.instruction = Instruction(:conditionalload)
-            push!(parents(falseop), negateop!(ls, condop, elementbytes))
-        end
-    else
-        falseop = getop(ls, iffalse, elementbytes)
+  end
+  if iffalse isa Expr
+    falseop = add_operation!(ls, gensym!(ls, "iffalse"), iffalse, elementbytes, position)
+    if iffalse.head === :ref && all(ld -> ld ∈ loopdependencies(falseop), loopdependencies(condop)) && !search_tree(parents(condop), falseop)
+      falseop.instruction = Instruction(:conditionalload)
+      push!(parents(falseop), negateop!(ls, condop, elementbytes))
+      if (any(==(identifier(trueop)), Iterators.map(first, ls.preamble_zeros)))
+        falseop.variable = LHS
+        ls.opdict[LHS] = falseop
+        return falseop
+      end
     end
-    add_compute_ifelse!(ls, LHS, condop, trueop, falseop, elementbytes)
+  else
+    falseop = getop(ls, iffalse, elementbytes)
+  end
+  add_compute_ifelse!(ls, LHS, condop, trueop, falseop, elementbytes)
 end
 
 function add_andblock!(ls::LoopSet, condop::Operation, LHS, rhsop::Operation, elementbytes::Int, position::Int)

test/fallback.jl

@@ -11,7 +11,7 @@
     end
     function msdavx(x)
         s = zero(eltype(x))
-        @turbo for i in eachindex(x)
+        @turbo warn_check_args=true for i in eachindex(x)
             s = muladd(x[i], x[i], s) # Avoids fastmath in fallback loop.
         end
         s
@@ -33,6 +33,8 @@
     @test @inferred !LoopVectorization.check_args(['a'])
     @test @inferred !LoopVectorization.check_args(Diagonal(x))
 
+    @test_nowarn msdavx(x)
+    @test_logs (:warn,"`LoopVectorization.check_args` on your inputs failed; running fallback `@inbounds @fastmath` loop instead.") msdavx(FallbackArrayWrapper(x))
     @test msdavx(FallbackArrayWrapper(x)) == 1e18
     @test msd(x) == msdavx(FallbackArrayWrapper(x))
     @test msdavx(x) != msdavx(FallbackArrayWrapper(x))