Improved check_args, make some literals static for calls like axes, and bump compat requirements on VectorizationBase/ArrayInterface to force using their inference improvements.

Author: chriselrod

Project.toml

@@ -14,13 +14,13 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
-ArrayInterface = "2.14"
+ArrayInterface = "2.14.1"
 DocStringExtensions = "0.8"
 IfElse = "0.1"
 OffsetArrays = "1"
 SLEEFPirates = "0.6"
 UnPack = "0,1"
-VectorizationBase = "0.13"
+VectorizationBase = "0.13.2"
 julia = "1.5"
 
 [extras]

src/condense_loopset.jl

@@ -255,28 +255,13 @@ It returns true for `AbstractArray{T}`s when `check_type(T) == true` and the arr
 To provide support for a custom array type, ensure that `check_args` returns true, either through overloading it or subtyping `DenseArray`.
 Additionally, define `pointer` and `stride` methods.
 """
-@inline check_args(A::SubArray{T,N,P,I}) where {T,N,P,I<:Tuple{Vararg{Union{Int,Colon,AbstractRange}}}} = check_args(parent(A))
-@inline check_args(A::OffsetArray) = check_args(parent(A))
-@inline check_args(A::Adjoint) = check_args(parent(A))
-@inline check_args(A::Transpose) = check_args(parent(A))
-@inline check_args(A::PermutedDimsArray) = check_args(parent(A))
-@inline check_args(A::StridedArray) = check_type(eltype(A))
-@inline check_args(A::AbstractRange) = check_type(eltype(A))
-@inline check_args(A::BitVector) = true
-@inline check_args(A::BitMatrix) = true
-@inline function check_args(A::AbstractArray)
-    M = parentmodule(typeof(A))
-    if parent(A) === A # SparseMatrix, StaticArray, etc
-        false
-    elseif M === Base || M === Core || M ===LinearAlgebra
-        # reshapes which aren't StridedArrays, plus UpperTriangular, etc.
-        false
-    else
-        check_args(parent(A)) # PermutedDimsArray, NamedDimsArray
-    end
+@inline function check_args(A::AbstractArray{T}) where {T}
+    check_type(T) && ArrayInterface.device(A) === ArrayInterface.CPUPointer()
 end
-@inline check_args(A::VectorizationBase.AbstractStridedPointer{T}) where {T} = check_type(T)
-@inline check_args(A, Bs...) = check_args(A) && check_args(Bs...)
+@inline check_args(A::BitVector) = true
+@inline check_args(A::BitArray) = iszero(size(A,1) & 7)
+@inline check_args(_) = false
+@inline check_args(A, B, C::Vararg{Any,K}) where {K} = check_args(A) && check_args(B, C...)
 """
     check_type(::Type{T}) where {T}
 

src/graphs.jl

@@ -458,6 +458,16 @@ function add_loop_bound!(ls::LoopSet, itersym::Symbol, bound, upper::Bool = true
     pushprepreamble!(ls, Expr(:(=), N, bound))
     N
 end
+function static_literals!(q::Expr)
+    for (i,ex) ∈ enumerate(q.args)
+        if ex isa Number
+            q.args[i] = staticexpr(ex)
+        elseif ex isa Expr
+            static_literals!(ex)
+        end
+    end
+    q
+end
 
 """
 This function creates a loop, while switching from 1 to 0 based indices
@@ -503,7 +513,7 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
             end
         else
             N = gensym("loop" * string(itersym))
-            pushprepreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticrange), r)))
+            pushprepreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticrange), static_literals!(r))))
             L = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticfirst), N), false)
             U = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticlast), N), true)
             Loop(itersym, L, U)