Fix unrolled stepranges

Author: chriselrod

Project.toml

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

src/codegen/lower_compute.jl

@@ -62,43 +62,44 @@ function parent_unroll_status(op::Operation, u₁loop::Symbol, u₂loop::Symbol,
 end
 
 function _add_loopvalue!(ex::Expr, loopval::Symbol, vloop::Loop, u::Int)
-    vloopsym = vloop.itersymbol
-    if loopval === vloopsym
-        if iszero(u)
-            push!(ex.args, _MMind(loopval, step(vloop)))
-        else
-            mm = _MMind(loopval, step(vloop))
-            if isone(u)
-                push!(ex.args, Expr(:call, lv(:vadd_nsw), VECTORWIDTHSYMBOL, mm))
-            else
-                push!(ex.args, Expr(:call, lv(:vadd_nsw), Expr(:call, lv(:vmul_nsw), VECTORWIDTHSYMBOL, u), mm))
-            end
-        end
-    elseif u == 0
-        push!(ex.args, loopval)
+  vloopsym = vloop.itersymbol
+  if loopval === vloopsym
+    if iszero(u)
+      push!(ex.args, _MMind(loopval, step(vloop)))
     else
-        push!(ex.args, Expr(:call, lv(:vadd_nsw), loopval, staticexpr(u)))
+      vstep = step(vloop)
+      mm = _MMind(loopval, vstep)
+      if isone(u) & isone(vstep)
+        push!(ex.args, Expr(:call, lv(:vadd_nsw), VECTORWIDTHSYMBOL, mm))
+      else
+        push!(ex.args, Expr(:call, lv(:vadd_nsw), mulexpr(VECTORWIDTHSYMBOL, u, vstep), mm))
+      end
     end
+  elseif u == 0
+    push!(ex.args, loopval)
+  else
+    push!(ex.args, Expr(:call, lv(:vadd_nsw), loopval, staticexpr(u)))
+  end
 end
 function add_loopvalue!(instrcall::Expr, loopval, ua::UnrollArgs, u₁::Int)
-    @unpack u₁loopsym, u₂loopsym, vloopsym, vloop, suffix = ua
-    if loopval === u₁loopsym #parentsunrolled[n]
-        if isone(u₁)
-            _add_loopvalue!(instrcall, loopval, vloop, 0)
-        else
-            t = Expr(:tuple)
-            for u ∈ 0:u₁-1
-                _add_loopvalue!(t, loopval, vloop, u)
-            end
-            push!(instrcall.args, Expr(:call, lv(:VecUnroll), t))
-        end
-    elseif suffix > 0 && loopval === u₂loopsym
-        _add_loopvalue!(instrcall, loopval, vloop, suffix)
-    elseif loopval === vloopsym
-        push!(instrcall.args, _MMind(loopval, step(vloop)))
+  @unpack u₁loopsym, u₂loopsym, vloopsym, vloop, suffix = ua
+  if loopval === u₁loopsym #parentsunrolled[n]
+    if isone(u₁)
+      _add_loopvalue!(instrcall, loopval, vloop, 0)
     else
-        push!(instrcall.args, loopval)
+      t = Expr(:tuple)
+      for u ∈ 0:u₁-1
+        _add_loopvalue!(t, loopval, vloop, u)
+      end
+      push!(instrcall.args, Expr(:call, lv(:VecUnroll), t))
     end
+  elseif suffix > 0 && loopval === u₂loopsym
+    _add_loopvalue!(instrcall, loopval, vloop, suffix)
+  elseif loopval === vloopsym
+    push!(instrcall.args, _MMind(loopval, step(vloop)))
+  else
+    push!(instrcall.args, loopval)
+  end
 end
 
 vecunrolllen(::Type{VecUnroll{N,W,T,V}}) where {N,W,T,V} = (N::Int + 1)

src/codegen/lowering.jl

@@ -412,31 +412,31 @@ function unrollremcomparison(ls::LoopSet, loop::Loop, UFt::Int, n::Int, nisvecto
     end
 end
 function loopvarremcomparison(loop::Loop, UFt::Int, nisvectorized::Bool, remfirst::Bool)
-    loopsym = loop.itersymbol
-    loopstep = loop.step
-    if nisvectorized
-        offset = mulexpr(VECTORWIDTHSYMBOL, UFt, loopstep)
-        itercount = subexpr(last(loop), offset)
-        Expr(:call, GlobalRef(Base,:>), loopsym, itercount)
-    elseif remfirst # requires `isstaticloop(loop)`
-        Expr(:call, GlobalRef(Base,:<), loopsym, gethint(first(loop)) + UFt*gethint(loopstep) - 1)
-    elseif isknown(last(loop))
-        if isknown(loopstep)
-            Expr(:call, GlobalRef(Base,:>), loopsym, gethint(last(loop)) - UFt*gethint(loopstep))
-        elseif isone(UFt)
-            Expr(:call, GlobalRef(Base,:>), loopsym, subexpr(gethint(last(loop)), getsym(loopstep)))
-        else
-            Expr(:call, GlobalRef(Base,:>), loopsym, subexpr(gethint(last(loop)),  mulexpr(getsym(loopstep), UFt)))
-        end
+  loopsym = loop.itersymbol
+  loopstep = loop.step
+  if nisvectorized
+    offset = mulexpr(VECTORWIDTHSYMBOL, UFt, loopstep)
+    itercount = subexpr(last(loop), offset)
+    Expr(:call, GlobalRef(Base,:>), loopsym, itercount)
+  elseif remfirst # requires `isstaticloop(loop)`
+    Expr(:call, GlobalRef(Base,:<), loopsym, gethint(first(loop)) + UFt*gethint(loopstep) - 1)
+  elseif isknown(last(loop))
+    if isknown(loopstep)
+      Expr(:call, GlobalRef(Base,:>), loopsym, gethint(last(loop)) - UFt*gethint(loopstep))
+    elseif isone(UFt)
+      Expr(:call, GlobalRef(Base,:>), loopsym, subexpr(gethint(last(loop)), getsym(loopstep)))
     else
-        if isknown(loopstep)
-            Expr(:call, GlobalRef(Base,:>), loopsym, Expr(:call, lv(:vsub_nsw), getsym(last(loop)), UFt*gethint(loopstep)))
-        elseif isone(UFt)
-            Expr(:call, GlobalRef(Base,:>), loopsym, Expr(:call, lv(:vsub_nsw), getsym(last(loop)), getsym(loopstep)))
-        else
-            Expr(:call, GlobalRef(Base,:>), loopsym, Expr(:call, lv(:vsub_nsw), getsym(last(loop)), mulexpr(getsym(loopstep), UFt)))
-        end
+      Expr(:call, GlobalRef(Base,:>), loopsym, subexpr(gethint(last(loop)),  mulexpr(getsym(loopstep), UFt)))
     end
+  else
+    if isknown(loopstep)
+      Expr(:call, GlobalRef(Base,:>), loopsym, Expr(:call, lv(:vsub_nsw), getsym(last(loop)), UFt*gethint(loopstep)))
+    elseif isone(UFt)
+      Expr(:call, GlobalRef(Base,:>), loopsym, Expr(:call, lv(:vsub_nsw), getsym(last(loop)), getsym(loopstep)))
+    else
+      Expr(:call, GlobalRef(Base,:>), loopsym, Expr(:call, lv(:vsub_nsw), getsym(last(loop)), mulexpr(getsym(loopstep), UFt)))
+    end
+  end
 end
 function pointerremcomparison(ls::LoopSet, termind::Int, UFt::Int, n::Int, nisvectorized::Bool, remfirst::Bool, loop::Loop)
   lssm = ls.lssm