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

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 q.args[2].head != :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(:(@avx $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(:(@avx $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(:(@avx $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(:(@avx $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