Fix for i in rng

Author: timholy

src/LoopVectorization.jl

@@ -12,6 +12,7 @@ using SIMDPirates: VECTOR_SYMBOLS, evadd, evmul, vrange, reduced_add, reduced_pr
     vmullog2, vmullog10, vdivlog2, vdivlog10, vmullog2add!, vmullog10add!, vdivlog2add!, vdivlog10add!, vfmaddaddone
 using Base.Broadcast: Broadcasted, DefaultArrayStyle
 using LinearAlgebra: Adjoint, Transpose
+using Base.Meta: isexpr
 
 const SUPPORTED_TYPES = Union{Float16,Float32,Float64,Integer}
 

src/graphs.jl

@@ -369,47 +369,55 @@ This function creates a loop, while switching from 1 to 0 based indices
 """
 function register_single_loop!(ls::LoopSet, looprange::Expr)
     itersym = (looprange.args[1])::Symbol
-    r = (looprange.args[2])::Expr
-    @assert r.head === :call
-    f = first(r.args)
-    loop::Loop = if f === :(:)
-        lower = r.args[2]
-        upper = r.args[3]
-        lii::Bool = lower isa Integer
-        liiv::Int = lii ? (convert(Int, lower)-1) : 0
-        uii::Bool = upper isa Integer
-        if lii & uii # both are integers
-            Loop(itersym, liiv, convert(Int, upper))
-        elseif lii # only lower bound is an integer
-            if upper isa Symbol
-                Loop(itersym, liiv, upper)
-            elseif upper isa Expr
-                Loop(itersym, liiv, add_loop_bound!(ls, itersym, upper, true))
-            else
-                Loop(itersym, liiv, add_loop_bound!(ls, itersym, upper, true))
+    r = looprange.args[2]
+    if isexpr(r, :call)
+        f = first(r.args)
+        loop::Loop = if f === :(:)
+            lower = r.args[2]
+            upper = r.args[3]
+            lii::Bool = lower isa Integer
+            liiv::Int = lii ? (convert(Int, lower)-1) : 0
+            uii::Bool = upper isa Integer
+            if lii & uii # both are integers
+                Loop(itersym, liiv, convert(Int, upper))
+            elseif lii # only lower bound is an integer
+                if upper isa Symbol
+                    Loop(itersym, liiv, upper)
+                elseif upper isa Expr
+                    Loop(itersym, liiv, add_loop_bound!(ls, itersym, upper, true))
+                else
+                    Loop(itersym, liiv, add_loop_bound!(ls, itersym, upper, true))
+                end
+            elseif uii # only upper bound is an integer
+                uiiv = convert(Int, upper)
+                Loop(itersym, add_loop_bound!(ls, itersym, lower, false), uiiv)
+            else # neither are integers
+                L = add_loop_bound!(ls, itersym, lower, false)
+                U = add_loop_bound!(ls, itersym, upper, true)
+                Loop(itersym, L, U)
             end
-        elseif uii # only upper bound is an integer
-            uiiv = convert(Int, upper)
-            Loop(itersym, add_loop_bound!(ls, itersym, lower, false), uiiv)
-        else # neither are integers
-            L = add_loop_bound!(ls, itersym, lower, false)
-            U = add_loop_bound!(ls, itersym, upper, true)
-            Loop(itersym, L, U)
-        end
-    elseif f === :eachindex
-        N = gensym(Symbol(:loop, itersym))
-        pushpreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticlength), r.args[2])))
-        Loop(itersym, 0, N)
-    elseif f === :OneTo || f == Expr(:(.), :Base, QuoteNode(:OneTo))
-        otN = r.args[2]
-        if otN isa Integer
-            Loop(itersym, 0, otN)
-        else
-            otN isa Expr && maybestatic!(otN)
+        elseif f === :eachindex
             N = gensym(Symbol(:loop, itersym))
-            pushpreamble!(ls, Expr(:(=), N, otN))
+            pushpreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticlength), r.args[2])))
             Loop(itersym, 0, N)
+        elseif f === :OneTo || f == Expr(:(.), :Base, QuoteNode(:OneTo))
+            otN = r.args[2]
+            if otN isa Integer
+                Loop(itersym, 0, otN)
+            else
+                otN isa Expr && maybestatic!(otN)
+                N = gensym(Symbol(:loop, itersym))
+                pushpreamble!(ls, Expr(:(=), N, otN))
+                Loop(itersym, 0, N)
+            end
+        else
+            throw("Unrecognized loop range type: $r.")
         end
+    elseif isa(r, Symbol)
+        # Treat similar to `eachindex`
+        N = gensym(Symbol(:loop, itersym))
+        pushpreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticlength), r)))
+        loop = Loop(itersym, 0, N)
     else
         throw("Unrecognized loop range type: $r.")
     end

test/dot.jl

@@ -1,3 +1,6 @@
+using LoopVectorization
+using Test
+
 @testset "dot" begin
     dotq = :(for i ∈ eachindex(a,b)
              s += a[i]*b[i]
@@ -46,6 +49,14 @@
         end
         s
     end
+    function myselfdotavx_range(a)
+        s = zero(eltype(a))
+        rng = axes(a, 1)
+        @avx for i ∈ rng
+            s += a[i]*a[i]
+        end
+        s
+    end
     function myselfdot_avx(a)
         s = zero(eltype(a))
         @_avx for i ∈ eachindex(a)
@@ -190,6 +201,7 @@
         @test dot_unroll3avx_inline(a,b) ≈ s
         s = myselfdot(a)
         @test myselfdotavx(a) ≈ s
+        @test myselfdotavx_range(a) ≈ s
         @test myselfdot_avx(a) ≈ s
         @test myselfdotavx(a) ≈ s
 

test/gemv.jl

@@ -1,3 +1,6 @@
+using LoopVectorization
+using Test
+
 @testset "GEMV" begin
     gemvq = :(for i ∈ eachindex(y)
               yᵢ = 0.0
@@ -27,6 +30,16 @@
             y[i] = yᵢ
         end
     end
+    function mygemvavx_range!(y, A, x)
+        rng1, rng2 = axes(A)
+        @avx for i ∈ rng1
+            yᵢ = zero(eltype(y))
+            for j ∈ rng2
+                yᵢ += A[i,j] * x[j]
+            end
+            y[i] = yᵢ
+        end
+    end
     q = :(for i ∈ eachindex(y)
           yᵢ = zero(eltype(y))
           for j ∈ eachindex(x)
@@ -150,6 +163,9 @@
         @test y1 ≈ y2
         fill!(y2, -999.9); mygemv_avx!(y2, A, x)
         @test y1 ≈ y2
+        fill!(y2, -999.9)
+        mygemvavx_range!(y2, A, x)
+        @test y1 ≈ y2
 
         B = rand(R, N, N);
         G1 = Matrix{TC}(undef, N, 1);