Merge pull request #175 from JuliaGPU/vc/fancy_private

Fancy private

Author: vchuravy

src/KernelAbstractions.jl

@@ -154,6 +154,16 @@ macro private(T, dims)
     end
 end
 
+"""
+    @private mem = 1
+
+Creates a private local of `mem` per item in the workgroup. This can be safely used
+across [`@synchronize`](@ref) statements.
+"""
+macro private(expr)
+    expr
+end
+
 """
     @uniform expr
 

src/backends/cpu.jl

@@ -241,23 +241,9 @@ end
 
 ###
 # CPU implementation of scratch memory
-# - private memory for each workitem
-# - memory allocated as a MArray with size `Dims + WorkgroupSize`
+# - memory allocated as a MArray with size `Dims`
 ###
-struct ScratchArray{N, D}
-    data::D
-    ScratchArray{N}(data::D) where {N, D} = new{N, D}(data)
-end
-Base.eltype(a::ScratchArray) = eltype(a.data)
 
 @inline function Cassette.overdub(ctx::CPUCtx, ::typeof(Scratchpad), ::Type{T}, ::Val{Dims}) where {T, Dims}
-    return ScratchArray{length(Dims)}(MArray{__size((Dims..., __groupsize(ctx.metadata)...)), T}(undef))
-end
-
-Base.@propagate_inbounds function Base.getindex(A::ScratchArray, I...)
-    return A.data[I...]
-end
-
-Base.@propagate_inbounds function Base.setindex!(A::ScratchArray, val, I...)
-    A.data[I...] = val
+    return MArray{__size(Dims), T}(undef)
 end

src/macros.jl

@@ -117,7 +117,8 @@ struct WorkgroupLoop
     indicies :: Vector{Any}
     stmts :: Vector{Any}
     allocations :: Vector{Any}
-    private :: Vector{Any}
+    private_allocations :: Vector{Any}
+    private :: Set{Symbol}
 end
 
 is_sync(expr) = @capture(expr, @synchronize() | @synchronize(a_))
@@ -138,21 +139,23 @@ end
 
 # TODO proper handling of LineInfo
 function split(stmts,
-               indicies = Any[], private=Any[])
+               indicies = Any[], private = Set{Symbol}())
     # 1. Split the code into blocks separated by `@synchronize`
     # 2. Aggregate `@index` expressions
     # 3. Hoist allocations
     # 4. Hoist uniforms
 
     current     = Any[]
     allocations = Any[]
+    private_allocations = Any[]
     new_stmts   = Any[]
     for stmt in stmts
         has_sync = find_sync(stmt)
         if has_sync
-            loop = WorkgroupLoop(deepcopy(indicies), current, allocations, deepcopy(private))
+            loop = WorkgroupLoop(deepcopy(indicies), current, allocations, private_allocations, deepcopy(private))
             push!(new_stmts, emit(loop))
             allocations = Any[]
+            private_allocations = Any[]
             current     = Any[]
             is_sync(stmt) && continue
 
@@ -177,15 +180,26 @@ function split(stmts,
         if @capture(stmt, @uniform x_)
             push!(allocations, stmt)
             continue
+        elseif @capture(stmt, @private lhs_ = rhs_)
+            push!(private, lhs)
+            push!(private_allocations, :($lhs = $rhs))
+            continue
         elseif @capture(stmt, lhs_ = rhs_ | (vs__, lhs_ = rhs_))
             if @capture(rhs, @index(args__))
                 push!(indicies, stmt)
                 continue
             elseif @capture(rhs, @localmem(args__) | @uniform(args__) )
                 push!(allocations, stmt)
                 continue
-            elseif @capture(rhs, @private(args__))
-                push!(allocations, stmt)
+            elseif @capture(rhs, @private(T_, dims_))
+                # Implement the legacy `mem = @private T dims` as
+                # @private mem = Scratchpad(T, Val(dims))
+
+                if dims isa Integer
+                    dims = (dims,)
+                end
+                alloc = :($Scratchpad($T, Val($dims)))
+                push!(private_allocations, :($lhs = $alloc))
                 push!(private, lhs)
                 continue
             end
@@ -196,7 +210,7 @@ function split(stmts,
 
     # everything since the last `@synchronize`
     if !isempty(current)
-        loop = WorkgroupLoop(deepcopy(indicies), current, allocations, deepcopy(private))
+        loop = WorkgroupLoop(deepcopy(indicies), current, allocations, private_allocations, deepcopy(private))
         push!(new_stmts, emit(loop))
     end
     return new_stmts
@@ -212,13 +226,34 @@ function emit(loop)
     end
     stmts = Any[]
     append!(stmts, loop.allocations)
+
+    # private_allocations turn into lhs = ntuple(i->rhs, length(__workitems_iterspace()))
+    N = gensym(:N)
+    push!(stmts, :($N = length($__workitems_iterspace())))
+
+    for stmt in loop.private_allocations
+        if @capture(stmt, lhs_ = rhs_)
+            push!(stmts, :($lhs = ntuple(_->$rhs, $N)))
+        else
+            error("@private $stmt not an assignment")
+        end
+    end
+
     # don't emit empty loops
     if !(isempty(loop.stmts) || all(s->s isa LineNumberNode, loop.stmts))
         body = Expr(:block, loop.stmts...)
         body = postwalk(body) do expr
-            if @capture(expr, A_[i__])
+            if @capture(expr, lhs_ = rhs_)
+                if lhs in loop.private
+                    error("Can't assign to variables marked private")
+                end
+            elseif @capture(expr, A_[i__])
                 if A in loop.private
-                    return :($A[$(i...), $(idx).I...])
+                    return :($A[$__index_Local_Linear($(idx))][$(i...)])
+                end
+            elseif expr isa Symbol
+                if expr in loop.private
+                    return :($expr[$__index_Local_Linear($(idx))])
                 end
             end
             return expr