fix nested unroll macros

Author: vchuravy

examples/Project.toml

@@ -4,4 +4,5 @@ CUDAdrv = "c5f51814-7f29-56b8-a69c-e4d8f6be1fde"
 CUDAnative = "be33ccc6-a3ff-5ff2-a52e-74243cff1e17"
 CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/extras/loopinfo.jl

@@ -1,6 +1,6 @@
 module LoopInfo
 
-const HAS_LOOPINFO_EXPR = VERSION >= v"1.2.0-DEV.462"
+using MacroTools
 export @unroll
 
 ##
@@ -20,13 +20,16 @@ module MD
 end
 
 function loopinfo(expr, nodes...)
-    if expr.head != :for
+    if @capture(expr, for i_ in iter_ body__ end)
+        return quote
+            for $i in $iter
+                $(body...)
+                $(Expr(:loopinfo, nodes...))
+            end
+        end
+    else
         error("Syntax error: loopinfo needs a for loop")
     end
-    if HAS_LOOPINFO_EXPR
-        push!(expr.args[2].args, Expr(:loopinfo, nodes...))
-    end
-    return expr
 end
 
 """
@@ -48,6 +51,7 @@ if it is safe to do so.
 """
 macro unroll(N, expr)
     if !(N isa Integer)
+        @debug "@unroll macro inputs" N expr
         error("Syntax error: `@unroll N expr` needs a constant integer N")
     end
     expr = loopinfo(expr, MD.unroll_count(N))

test/unroll.jl

@@ -1,5 +1,6 @@
 using KernelAbstractions
 using KernelAbstractions.Extras
+using StaticArrays
 
 @kernel function kernel_unroll!(a)
   @unroll for i in 1:5
@@ -13,6 +14,25 @@ end
   end
 end
 
+# Check that nested `@unroll` doesn't throw a syntax error
+@kernel function kernel_unroll!(a, ::Val{N}) where N
+  @uniform begin
+    a = MVector{3, Float64}(1, 2, 3)
+    b = MVector{3, Float64}(3, 2, 1)
+    c = MMatrix{3, 3, Float64}(undef)
+  end
+  I = @index(Global)
+  @inbounds for m in 1:3
+    @unroll for j = 1:3
+      @unroll for i = 1:3
+        c[1, j] = m * a[1] * b[j]
+      end
+    end
+    a[I] = c[1, 1]
+    m % 2 == 0 && @synchronize
+  end
+end
+
 let
   a = zeros(5)
   kernel! = kernel_unroll!(CPU(), 1, 1)