fix: #395, check if loop body is a block and warp it if not (#396)

* fix:#395, check if loop body is a block and warp it if not

* refactor:  test head with isexpr, use `@turbo`

* Bump version to 0.12.107

* fix: missing `!`

Author: wangl-cc

Project.toml

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

src/constructors.jl

@@ -156,9 +156,24 @@ function process_args(
   end
   inline, check_empty, u₁, u₂, v, threads, warncheckarg
 end
+# check if the body of loop is a block, if not convert it to a block issue#395
+function check_loopbody!(q)
+  if q isa Expr && q.head == :for
+    if !Meta.isexpr(q.args[2], :block)
+      q.args[2] = Expr(:block, q.args[2])
+    else
+      for arg in q.args[2].args
+        check_loopbody!(arg) # check recursively for inner loop
+      end
+    end
+  end
+  return q
+end
+
 function turbo_macro(mod, src, q, args...)
   q = macroexpand(mod, q)
   if q.head === :for
+    check_loopbody!(q)
     ls = LoopSet(q, mod)
     inline, check_empty, u₁, u₂, v, threads, warncheckarg = process_args(args)
     esc(setup_call(ls, q, src, inline, check_empty, u₁, u₂, v, threads, warncheckarg))

test/grouptests.jl

@@ -8,6 +8,7 @@ const START_TIME = time()
 
   @time if LOOPVECTORIZATION_TEST == "all" || LOOPVECTORIZATION_TEST == "part1"
     @time include("broadcast.jl")
+    @time include("parsing_inputs.jl")
   end
 
   @time if LOOPVECTORIZATION_TEST == "all" || LOOPVECTORIZATION_TEST == "part2"

test/parsing_inputs.jl

@@ -0,0 +1,50 @@
+using LoopVectorization, Test, ArrayInterface
+
+# macros for generate loops whose body is not a block
+macro gen_loop_issue395(ex)
+  sym, ind = ex.args
+  loop_body = :(ret[$ind] = $sym[$ind])
+  loop = Expr(:for, :($ind = axes($sym, 1)), loop_body)
+  return esc(:(@turbo $loop))
+end
+macro gen_single_loop(B, A)
+  loop_body = :($B[i] = $A[i])
+  loop = Expr(:for, :(i = indices(($B, $A), 1)), loop_body)
+  return esc(:(@turbo $loop))
+end
+macro gen_nest_loop(C, A, B)
+  loop_body = :($C[i, j] = $A[i] * $B[j])
+  loop_head = Expr(:block, :(j = indices(($C, $B), (2, 1))), :(i = indices(($C, $A), 1)))
+  loop = Expr(:for, loop_head, loop_body)
+  return esc(:(@turbo $loop))
+end
+macro gen_A_mul_B(C, A, B)
+  inner_body = :(Cji += $A[j, k] * $B[k, i])
+  inner_loop = Expr(:for, :(k = indices(($A, $B), (2, 1))), inner_body)
+  loop = :(
+    for i in indices(($C, $B), 2), j in indices(($C, $A), 1)
+      Cji = zero(eltype($C))
+      $inner_loop
+      $C[j, i] = Cji
+    end
+  )
+  return esc(:(@turbo $loop))
+end
+
+@testset "check_block, #395" begin
+  A = rand(4)
+  B = rand(4)
+  C = rand(4, 4)
+  D = zeros(4)
+  E = zeros(4, 4)
+  F = zeros(4, 4)
+  ret = zeros(4)
+  @gen_single_loop D A
+  @gen_loop_issue395 B[i]
+  @gen_nest_loop E A B
+  @gen_A_mul_B F C E
+  @test D == A
+  @test ret == B
+  @test E == A * B'
+  @test F == C * E
+end