swap to ArrayInterface.axes

Author: chriselrod

Project.toml

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

src/LoopVectorization.jl

@@ -41,7 +41,7 @@ using Base.FastMath: add_fast, sub_fast, mul_fast, div_fast, inv_fast, abs2_fast
 using SLEEFPirates: log_fast, log2_fast, log10_fast, pow, sin_fast, cos_fast, sincos_fast
 
 using ArrayInterface
-using ArrayInterface: OptionallyStaticUnitRange, OptionallyStaticRange, Zero, One, StaticBool, True, False, reduce_tup, indices, UpTri, LoTri, strides, offsets, size, StrideIndex
+using ArrayInterface: OptionallyStaticUnitRange, OptionallyStaticRange, Zero, One, StaticBool, True, False, reduce_tup, indices, UpTri, LoTri, strides, offsets, size, axes, StrideIndex
 using StrideArraysCore: CloseOpen, PtrArray
 # @static if VERSION ≥ v"1.6.0-rc1" #TODO: delete `else` when dropping 1.5 support
     # using ArrayInterface: static_step

src/modeling/graphs.jl

@@ -971,6 +971,7 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
     elseif f === :indices || (isexpr(f, :(.), 2) && (f.args[2] === QuoteNode(:indices)) && ((f.args[1] === :ArrayInterface) || (f.args[1] === :LoopVectorization)))
       indices_loop!(ls, r, itersym)
     else
+      (f === :axes) && (r.args[1] = lv(:axes))
       misc_loop!(ls, r, itersym, (f === :eachindex) | (f === :axes))
     end
   elseif isa(r, Symbol)