add group index

Author: vchuravy

src/KernelAbstractions.jl

@@ -113,13 +113,38 @@ macro synchronize()
 end
 
 """
-   @index(Global)
-   @index(Local)
-   @index(Global, Cartesian)
+   @index
+
+The `@index` macro can be used to give you the index of a workitem within a kernel
+function. It supports both the production of a linear index or a cartesian index.
+A cartesian index is a general N-dimensional index that is derived from the iteration space.
+
+# Index granularity
+
+  - `Global`: Used to access global memory.
+  - `Group`: The index of the `workgroup`.
+  - `Local`: The within `workgroup` index.
+
+# Index kind
+
+  - `Linear`: Produces an `Int64` that can be used to linearly index into memory.
+  - `Global`: Produces a `CartesianIndex{N}` that can be used to index into memory.
+
+If the index kind is not provided it defaults to `Linear`, this is suspect to change.
+
+# Examples
+
+```julia
+@index(Global, Linear)
+@index(Global, Cartesian)
+@index(Local, Cartesian)
+@index(Group, Linear)
+@index(Global)
+```
 """
 macro index(locale, args...)
-    if !(locale === :Global || locale === :Local)
-        error("@index requires as first argument either :Global or :Local")
+    if !(locale === :Global || locale === :Local || locale === :Group)
+        error("@index requires as first argument either :Global, :Local or :Group")
     end
 
     if length(args) >= 1
@@ -142,9 +167,11 @@ end
 ###
 
 function __index_Local_Linear end
+function __index_Group_Linear end
 function __index_Global_Linear end
 
 function __index_Local_Cartesian end
+function __index_Group_Cartesian end
 function __index_Global_Cartesian end
 
 struct ConstAdaptor end

src/backends/cpu.jl

@@ -80,6 +80,11 @@ end
     return @inbounds LinearIndices(indices)[idx]
 end
 
+@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Group_Linear), idx::CartesianIndex)
+    indices = blocks(__iterspace(ctx.metadata))
+    return @inbounds LinearIndices(indices)[__groupindex(ctx.metadata)]
+end
+
 @inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Global_Linear), idx::CartesianIndex)
     I = @inbounds expand(__iterspace(ctx.metadata), __groupindex(ctx.metadata), idx)
     @inbounds LinearIndices(__ndrange(ctx.metadata))[I]
@@ -89,6 +94,10 @@ end
     return idx
 end
 
+@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Group_Cartesian), idx::CartesianIndex)
+    __groupindex(ctx.metadata)
+end
+
 @inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Global_Cartesian), idx::CartesianIndex)
     return @inbounds expand(__iterspace(ctx.metadata), __groupindex(ctx.metadata), idx)
 end

src/backends/cuda.jl

@@ -116,6 +116,10 @@ end
     return CUDAnative.threadIdx().x
 end
 
+@inline function Cassette.overdub(ctx::CUDACtx, ::typeof(__index_Group_Linear))
+    return CUDAnative.blockIdx().x
+end
+
 @inline function Cassette.overdub(ctx::CUDACtx, ::typeof(__index_Global_Linear))
     I =  @inbounds expand(__iterspace(ctx.metadata), CUDAnative.blockIdx().x, CUDAnative.threadIdx().x)
     # TODO: This is unfortunate, can we get the linear index cheaper
@@ -126,6 +130,10 @@ end
     @inbounds workitems(__iterspace(ctx.metadata))[CUDAnative.threadIdx().x]
 end
 
+@inline function Cassette.overdub(ctx::CUDACtx, ::typeof(__index_Group_Cartesian))
+    @inbounds blocks(__iterspace(ctx.metadata))[CUDAnative.blockIdx().x]
+end
+
 @inline function Cassette.overdub(ctx::CUDACtx, ::typeof(__index_Global_Cartesian))
     return @inbounds expand(__iterspace(ctx.metadata), CUDAnative.blockIdx().x, CUDAnative.threadIdx().x)
 end

test/test.jl

@@ -44,15 +44,31 @@ end
 end
 @kernel function index_linear_local(A)
        I  = @index(Global, Linear)
-       li = @index(Local, Linear)
-       A[I] = li
+       i = @index(Local, Linear)
+       A[I] = i
+end
+@kernel function index_linear_group(A)
+       I  = @index(Global, Linear)
+       i = @index(Group, Linear)
+       A[I] = i
 end
 @kernel function index_cartesian_global(A)
        I = @index(Global, Cartesian)
        A[I] = I
 end
+@kernel function index_cartesian_local(A)
+       I = @index(Global, Cartesian)
+       i = @index(Local, Cartesian)
+       A[I] = i
+end
+@kernel function index_cartesian_group(A)
+       I = @index(Global, Cartesian)
+       i = @index(Group, Cartesian)
+       A[I] = i
+end
 
 function indextest(backend, ArrayT)
+    # TODO: add test for _group and _local_cartesian
     A = ArrayT{Int}(undef, 16, 16)
     wait(index_linear_global(backend, 8)(A, ndrange=length(A)))
     @test all(A .== LinearIndices(A))