Merge #39

39: fix nested unroll macros r=vchuravy a=vchuravy

bors r+

Co-authored-by: Valentin Churavy <v.churavy@gmail.com>

Author: bors[bot]

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/KernelAbstractions.jl

@@ -115,13 +115,34 @@ end
 
 """
    @synchronize()
+
+After a `@synchronize` statement all read and writes to global and local memory
+from each thread in the workgroup are visible in from all other threads in the
+workgroup.
 """
 macro synchronize()
     quote
         $__synchronize()
     end
 end
 
+"""
+   @synchronize(cond)
+
+After a `@synchronize` statement all read and writes to global and local memory
+from each thread in the workgroup are visible in from all other threads in the
+workgroup. `cond` is not allowed to have any visible sideffects.
+
+# Platform differences
+  - `GPU`: This synchronization will only occur if the `cond` evaluates.
+  - `CPU`: This synchronization will always occur.
+"""
+macro synchronize(cond)
+    quote
+        $(esc(cond)) && $__synchronize()
+    end
+end
+
 """
    @index
 

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))

src/macros.jl

@@ -103,7 +103,13 @@ struct WorkgroupLoop
     private :: Vector{Any}
 end
 
-is_sync(expr) = @capture(expr, @synchronize)
+is_sync(expr) = @capture(expr, @synchronize() | @synchronize(a_))
+
+function is_scope_construct(expr::Expr)
+    expr.head === :block # ||
+    # expr.head === :let
+end
+
 function find_sync(stmt)
     result = false
     postwalk(stmt) do expr
@@ -131,7 +137,7 @@ function split(stmts,
             push!(new_stmts, emit(loop))
             allocations = Any[]
             current     = Any[]
-            @capture(stmt, @synchronize) && continue
+            is_sync(stmt) && continue
 
             # Recurse into scope constructs
             # TODO: This currently implements hard scoping
@@ -140,7 +146,7 @@ function split(stmts,
             recurse(x) = x
             function recurse(expr::Expr)
                 expr = unblock(expr)
-                if any(is_sync, expr.args)
+                if is_scope_construct(expr) && any(is_sync, expr.args)
                     new_args = unblock(split(expr.args, deepcopy(indicies), deepcopy(private)))
                     return Expr(expr.head, new_args...)
                 else 

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
@@ -8,16 +9,38 @@ using KernelAbstractions.Extras
 end
 
 @kernel function kernel_unroll!(a, ::Val{N}) where N
-  @unroll for i in 1:N
-    @inbounds a[i] = i
+  let M = N+5
+    @unroll for i in 6:M
+      @inbounds a[i-5] = i
+    end
+    @synchronize
+  end
+end
+
+# Check that nested `@unroll` doesn't throw a syntax error
+@kernel function kernel_unroll2!(A)
+  @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]
+    @synchronize(m % 2 == 0)
   end
 end
 
 let
   a = zeros(5)
   kernel! = kernel_unroll!(CPU(), 1, 1)
-  event = kernel!(a)
-  wait(event)
-  event = kernel!(a, Val(5))
-  wait(event)
+  wait(kernel!(a))
+  wait(kernel!(a, Val(5)))
+  kernel2! = kernel_unroll2!(CPU(), 1, 1)
+  wait(kernel2!(a))
 end