Failing gracefully

[stev47]

It might be worth thinking about falling back to traditional loops if there are unrecognized expressions while issuing a warning about not optimizing instead of a hard error.

Feel free to close if this doesn't align with your goals: errors are fine too if the expectance is that @avx should fail if unable to optimize.

[chriselrod]

Hi, thanks for the issue.

The two times problems can be noticed are:

1. Macro-expand time.
2. When compiling the generated function.

Macro-expansion happens immediately upon defining a function. It'll create a LoopSet object that describes the loop, and then summarizes it, creating a call to the generated function _avx_! (which is passed the summary, and all necessary variables, as arguments).

I think I'd prefer errors immediately at macro-expansion time, if possible.
If I wanted fall-back behavior, I could just delete the @avx. You could use @test_broken in your test suite to monitor and find out if/when it starts to work (in which case you would replace @test_broken with @test).

When compiling the generated function, I would like to at least have type-based fall-back behavior.
This could allow either using a black or whitelist to make it easier to write generic code. If there are types involved that it can make fast, it will. Otherwise, they'll run as written.

[baggepinnen]

I've been enjoying some awesome 2x performance improvements from sticking an @avx in front of some lines of code. Unfortunately, this appears to have lost me the ability to pass in Double64 to my function, or differentiate it with ForwardDiff, which is also need to be able to do. Would there be a way for me to keep the @avx and have it fall back to standard code if the types involved are "exotic"? I could potentially maintain two different functions, that that would lead to a ton of code duplication.

When compiling the generated function, I would like to at least have type-based fall-back behavior.
This could allow either using a black or whitelist to make it easier to write generic code. If there are types involved that it can make fast, it will. Otherwise, they'll run as written.

This sounds like it would be exactly what I need!

[chriselrod]

The major obstacle to this is that the original loop expression is lost, and the LoopSet object doesn't preserve the original structure. It represents loops as a set of operations that depend on one another, external arrays, and specific loop indices.

Perhaps this is another argument for
#74

As a stop gap, we could revive a modification of this code as a stop-gap.
@_avx_! could then be given the original loop expression, and check to confirm that the element type of each AbstractArray is in

julia> Union{Bool, Base.HWReal}
Union{Bool, Float32, Float64, Int16, Int32, Int64, Int8, UInt16, UInt32, UInt64, UInt8}

If not, it'll return the original quote.

PRs welcome if anyone wants to try their hand at implementing this. ;)

BTW, you can see some examples of using it with ForwardDiff here:
#93

[baggepinnen]

I'm not sure if I understand what the @_avx macro is doing. Is it inserting some info about "I want to @avx this" rather than avx:ing it immediately?

I was contemplating if a macro could be put at the entire function definition to from the following

function f(a::T, b::Vector{T}) where T

generate two methods, the one above, and additionally this one

function f(a::T, b::Vector{T}) where T <: Base.HWReal

where only the last one would contain the @avx features and the other one would be completely free of @avx. That would make compilation time faster for the @avx-free method I guess, but I'm not sure how to make the macro on the function definition "see" the code before it's transformed by the inner @avx annotations.

[chriselrod]

The file I linked no longer exists, and is no longer how things work.

The macro now:

julia> @macroexpand @avx for i ∈ eachindex(a,b)
               s += a[i]*b[i]
           end
quote
    var"##loopi#1406" = LoopVectorization.maybestaticrange(eachindex(a, b))
    var"##i_loop_lower_bound#1407" = LoopVectorization.maybestaticfirst(var"##loopi#1406")
    var"##i_loop_upper_bound#1408" = LoopVectorization.maybestaticlast(var"##loopi#1406")
    var"##vptr##_a" = LoopVectorization.stridedpointer(a)
    var"##vptr##_b" = LoopVectorization.stridedpointer(b)
    local var"##s_0"
    begin
        $(Expr(:gc_preserve, :(var"##s_0" = LoopVectorization._avx_!(Val{(0, 0)}(), Tuple{:LoopVectorization, :getindex, LoopVectorization.OperationStruct(0x0000000000000001, 0x0000000000000000, 0x0000000000000000, 0x0000000000000000, LoopVectorization.memload, 0x01, 0x01), :LoopVectorization, :getindex, LoopVectorization.OperationStruct(0x0000000000000001, 0x0000000000000000, 0x0000000000000000, 0x0000000000000000, LoopVectorization.memload, 0x02, 0x02), :LoopVectorization, :LOOPCONSTANTINSTRUCTION, LoopVectorization.OperationStruct(0x0000000000000000, 0x0000000000000000, 0x0000000000000001, 0x0000000000000000, LoopVectorization.constant, 0x00, 0x03), :LoopVectorization, :vfmadd_fast, LoopVectorization.OperationStruct(0x0000000000000001, 0x0000000000000001, 0x0000000000000000, 0x0000000000010203, LoopVectorization.compute, 0x00, 0x03)}, Tuple{LoopVectorization.ArrayRefStruct{:a,Symbol("##vptr##_a")}(0x0000000000000001, 0x0000000000000001, 0x0000000000000000), LoopVectorization.ArrayRefStruct{:b,Symbol("##vptr##_b")}(0x0000000000000001, 0x0000000000000001, 0x0000000000000000)}, Tuple{0, Tuple{4}, Tuple{3}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}, Tuple{:i}, (var"##i_loop_lower_bound#1407":var"##i_loop_upper_bound#1408",), var"##vptr##_a", var"##vptr##_b", s)), :a, :b))
    end
    s = LoopVectorization.reduced_add(var"##s_0", s)
end

It generates a call to _avx_!, along with a summary of what _avx_! needs to know about the original loop. This is how it uses compile-time information (instead of just macro-expansion-time info) to optimize the loops.

but I'm not sure how to make the macro on the function definition "see" the code before it's transformed by the inner @avx annotations.

Outer macros expand first, so this shouldn't be a problem.

julia> a = rand(43); b = rand(43);

julia> s = 0.0;

julia> macro showex(x)
          esc(@show x)
       end
@showex (macro with 1 method)

julia> @showex @avx for i ∈ eachindex(a,b)
               s += a[i]*b[i]
           end
x = :(#= REPL[23]:1 =# @avx for i = eachindex(a, b)
          #= REPL[23]:2 =#
          s += a[i] * b[i]
      end)
11.942335582598114

LoopVectorization actually has to call macroexpand so that it can avoid parsing them, and so that things like Base.Cartesian.@nexprs can work as intended for unrolling loops.

I'd definitely accept a macro like that as a PR. Although I think it should be Union{Bool, Base.HWReal}, because it should be support Bools and I have a few tests for it.

[baggepinnen]

EDIT: hold on a second, I had made a mistake

[baggepinnen]

This seems to be working. Right now it only handles a single type parameter, but that can be improved if you find the approach okay

using MacroTools, LoopVectorization
macro maybeavx(ex)
    def = splitdef(ex)
    haskey(def,:whereparams) || throw(ArgumentError("Found no type parameter in the function definition"))
    length(def[:whereparams]) == 1 || throw(ArgumentError("Only a single type parameter supported"))
    T = def[:whereparams][1] # This is the type parameter

    # Remove @avx annotations
    ex_noavx = MacroTools.prewalk(ex) do x
        if x isa Expr && x.head === :macrocall
            return x.args[3] # This returns the expression inside the macrocall
        end
        x
    end
    quote
        # Define method without @avx annotations
        $(esc(ex_noavx))

        # Define the method that retains the @avx annotations
        $(esc(def[:name]))($(esc.(def[:args])...)) where $(esc(T)) <: Union{Bool, Base.HWReal} = $(esc(def[:body]))
    end
end


@maybeavx function f(a::T, b::Vector{T}) where T
    @avx a .+ b
end


using DoubleFloats
julia> f(1.0, ones(3))
3-element Array{Float64,1}:
 2.0
 2.0
 2.0

julia> f(Double64(1.0), Double64.(ones(3)))
3-element Array{Double64,1}:
 2.0
 2.0
 2.0

julia> @code_lowered f(1.0, ones(3))
CodeInfo(

1 ─      266 = Base.broadcasted(Main.:+, a, b)
│   %2 = 266
│   %3 = Core.apply_type(Main.Val, :Main)
│   %4 = (%3)()
│   %5 = LoopVectorization.vmaterialize(%2, %4)
│        267 = %5
└──      return %5
)

julia> @code_lowered f(Double64(1.0), Double64.(ones(3)))
CodeInfo(
1 ─ %1 = Base.broadcasted(Main.:+, a, b)
│   %2 = Base.materialize(%1)
└──      return %2
)

[mcabbott]

Such a macro sounds like a neat idea. I wonder if it shouldn't insert the two options as if/else in the body of the function, to avoid strange interactions with dispatch.

Perhaps too complicated, but I also wonder whether it should check that the arrays aren't sparse, or CuArrays, too. Something like this, with a,b extracted from the LoopSet, and check_types recursively calling parent and to find something like Array{<:HWReal}:

@maybeavx function f(x, s)
    a, b = x .+ 1, x .- 1
    @avx for i ∈ eachindex(a,b)
        s += a[i]*b[i]
    end
    ...
# becomes:

function f(x, s)
    a, b = x .+ 1, x .- 1
    if check_types(a, b) 
        @avx for i ∈ eachindex(a,b)
            s += a[i]*b[i]
        end
   else
       for i ∈ eachindex(a,b)
            ...

Edit: if you want if right in front of @avx, then there's little need for a second macro on the function. The two alternatives could simply be what @avx returns.

[chriselrod]

@baggepinnen , I liked your approach until:

Edit: if you want if right in front of @avx, then there's little need for a second macro on the function. The two alternatives could simply be what @avx returns.

Ha ha, that's a lot simpler than what I had suggested.
If we wanted to replicate making it optional (like @maybeavx), @avx could have a key word argument (such as force=false), that you can set to true to remove the check if for some reason you thought it would work with some non-hw types.

A modification of the @maybeavx macro that returns two functions with different names (one with, and one without) @avx would be useful for testing to make sure the function is still returning the correct results.
The test suite, for example, is filled with such nearly-duplicate function definitions.

[chriselrod]

Perhaps something like this:

check_valid_args() = true
check_valid_args(::Any) = false
check_valid_args(::Type{T}) where {T <: Union{Base.HWReal, Bool}} = true
check_valid_args(::T) where {T <: Number} = check_valid_args(T)
check_valid_args(::StridedArray{T}) where {T} = check_valid_args(T)
check_valid_args(::AbstractRange{T}) where {T} = check_valid_args(T)
check_valid_args(a, b, args...) = check_valid_args(a) && check_valid_args(b) && check_valid_args(args....)

This would let you to opt in an array type MyArray via either defining:

struct MyArray{T,N} <: DenseArray{T,N}
...
end

or

LoopVectorization.check_valid_args(::MyArray{T}) where {T} = LoopVectorization.check_valid_args(T)

And to opt in your own element type, MyElementType, similarly via

LoopVectorization.check_valid_args(::Type{MyElementType}) = true

Some additional definitions:

using LinearAlgebra
check_valid_args(A::Adjoint) = check_valid_args(parent(A))
check_valid_args(A::Transpose) = check_valid_args(parent(A))
using OffsetArrays
check_valid_args(A::OffsetArray) = check_valid_args(parent(A))

Anything else I'm missing?
I don't think I should make a generic AbstractArray fallback with parent, because odds are that the wrapper type will need explicit support if it didn't subtype DenseArray.

[mcabbott]

My suggestion for how to check applicability was going be something ... not so different to yours, but allowing for wrappers such as those of NamedDims, which don't do anything which this has to know about, but are nice to pass along.

[chriselrod]

My problem with that generic definition is

I don't think I should make a generic AbstractArray fallback with parent, because odds are that the wrapper type will need explicit support if it didn't subtype DenseArray

For example:

struct MyOffsetArray{T,N} <: AbstractArray{T,N}
    data::Array{T,N}
    offsets::NTuple{N,Int}
end
Base.@propagate_inbounds function Base.getindex(A::MyOffsetArray{T,N}, i::Varargs{Int,N}) where {T,N}
    A.data[(i .+ A.offsets)...]
end
Base.parent(A::MyOffsetArray) = A.data
# Other methods, including `Base.strides` and `Base.pointer`/`Base.unsafe_convert(Ptr{T}, ...)`

But maybe that's okay. Those array types that don't work correctly:

julia> pointer(A')
ERROR: conversion to pointer not defined for Adjoint{Float64,Array{Float64,2}}
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] unsafe_convert(::Type{Ptr{Float64}}, ::Adjoint{Float64,Array{Float64,2}}) at ./pointer.jl:67
 [3] pointer(::Adjoint{Float64,Array{Float64,2}}) at ./abstractarray.jl:933
 [4] top-level scope at REPL[27]:1

julia> pointer(OffsetArray(A,1,1))
ERROR: conversion to pointer not defined for OffsetArray{Float64,2,Array{Float64,2}}
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] unsafe_convert(::Type{Ptr{Float64}}, ::OffsetArray{Float64,2,Array{Float64,2}}) at ./pointer.jl:67
 [3] pointer(::OffsetArray{Float64,2,Array{Float64,2}}) at ./abstractarray.jl:933
 [4] top-level scope at REPL[28]:1

julia> pointer(A)
Ptr{Float64} @0x000055bfbb939f40

Don't define some of the needed methods, meaning they'll get errors and at least learn about it.

EDIT:
It's a question of how far on the opt-in vs opt-out spectrum we should be.
When it comes to things like NamedArray, it is unfortunate that it'd be type piracy for someone else to make it work.

Would be great if there were a acts_like_a_normal_strided_array trait that everyone implemented.

[mcabbott]

I guess the clear criteria are:

• false for CuArrays (which we can't see) and sparse arrays
• false for reshapes which aren't StridedArrays, and for Diagonal, UpperTriangular etc (these wrappers we could see)
• true for wrappers PermutedDimsArray, Transpose. And Adjoint, since it will only contain real numbers, and Make transposes of StridedArrays strided JuliaLang/julia#29135 will define strides & pointer.

As you say, after that it depends how optimistic you want to be about weird wrappers. Besides OffsetArrays (for which one could test axes(A)), are there any defined by packages which break things? NamedDimsArrays, and AxisArrays & friends, and NNlib.batched_transpose are examples of things which don't.

My PR here https://github.com/FluxML/NNlib.jl/pull/191/files has a go at making an is_strided which is optimistic about foreign wrappers. (But should probably reject OffsetArrays.)

(Re traits, there is https://github.com/JuliaMatrices/ArrayLayouts.jl, and I quite like JuliaLang/julia#30432 but it hasn't been merged.)

[chriselrod]

Okay, I'm fine with being more optimistic and using the recursive typeof(parent(A)) === typeof(A) trick.

I'll work on this tonight.

[chriselrod]

Got around to it about a day late, but I did add it yesterday, using mcabbot's implementation.

Example from the unit tests, wrapper type that does not subtype DenseArray and does not define parent:

using LoopVectorization
struct FallbackArrayWrapper{T,N} <: AbstractArray{T,N}
    data::Array{T,N}
end
Base.size(A::FallbackArrayWrapper) = size(A.data)
Base.@propagate_inbounds Base.getindex(A::FallbackArrayWrapper, i::Vararg{Int, N}) where {N} = getindex(A.data, i...)
Base.@propagate_inbounds Base.setindex!(A::FallbackArrayWrapper, v, i::Vararg{Int, N}) where {N} = setindex!(A.data, v, i...)
Base.IndexStyle(::Type{<:FallbackArrayWrapper}) = IndexLinear()

function msd(x)
    s = zero(eltype(x))
    for i in eachindex(x)
        s += x[i] * x[i]
    end
    s
end
function msdavx(x)
    s = zero(eltype(x))
    @avx for i in eachindex(x)
        s += x[i] * x[i]
    end
    s
end

x = fill(sqrt(63.0), 128);
x[1] = 1e9;

This yields

julia> msd(x)
1.0e18

julia> msdavx(x)
1.0000000000000081e18

julia> msdavx(FallbackArrayWrapper(x))
1.0e18

The test is based on the fact that 1e18 + 63 == 1e18, so if it doesn't get unrolled, no matter how many times we add 63, we still have 1e18 at the end. If it does get unrolled however, each 63 that doesn't share a lane with 1e18 avoids getting eaten, and gets added on at the end.

Although, this begs the question:
Perhaps the fall back definition should use one or both of @inbounds and @fastmath (in which case I'd have to modify the above test)?