Merge #38

38: add ntuple index type r=vchuravy a=vchuravy

bors r+

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

Author: bors[bot]

examples/matmul.jl

@@ -6,15 +6,15 @@ end
 
 # Simple kernel for matrix multiplication
 @kernel function matmul_kernel!(a, b, c)
-    cI = @index(Global, Cartesian)
+    i, j = @index(Global, NTuple)
 
     # creating a temporary sum variable for matrix multiplication
     tmp_sum = zero(eltype(c))
-    for i = 1:size(a)[2]
-        tmp_sum += a[cI[1],i] * b[i,cI[2]]
+    for k = 1:size(a)[2]
+        tmp_sum += a[i,k] * b[k, j]
     end
 
-    c[cI] = tmp_sum
+    c[i,j] = tmp_sum
 end
 
 # Creating a wrapper kernel for launching with error checks

examples/naive_transpose.jl

@@ -5,8 +5,7 @@ if CUDAapi.has_cuda_gpu()
 end
 
 @kernel function naive_transpose_kernel!(a, b)
-  I = @index(Global, Cartesian)
-  i, j = Tuple(I)
+  i, j = @index(Global, NTuple)
   @inbounds b[i, j] = a[j, i]
 end
 

examples/performance.jl

@@ -9,8 +9,7 @@ using CUDAnative
 using CUDAnative.NVTX
 
 @kernel function transpose_kernel_naive!(b, a)
-    I = @index(Global, Cartesian)
-    i, j = I.I 
+    i, j = @index(Global, NTuple)
     @inbounds b[i, j] = a[j, i]
 end
 

src/KernelAbstractions.jl

@@ -138,7 +138,8 @@ A cartesian index is a general N-dimensional index that is derived from the iter
 # 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.
+  - `Cartesian`: Produces a `CartesianIndex{N}` that can be used to index into memory.
+  - `NTuple`: Produces a `NTuple{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.
 
@@ -149,6 +150,7 @@ If the index kind is not provided it defaults to `Linear`, this is suspect to ch
 @index(Global, Cartesian)
 @index(Local, Cartesian)
 @index(Group, Linear)
+@index(Local, NTuple)
 @index(Global)
 ```
 """
@@ -158,7 +160,10 @@ macro index(locale, args...)
     end
 
     if length(args) >= 1
-        if args[1] === :Cartesian || args[1] === :Linear
+
+        if args[1] === :Cartesian ||
+           args[1] === :Linear ||
+           args[1] === :NTuple
             indexkind = args[1]
             args = args[2:end]
         else
@@ -184,6 +189,10 @@ function __index_Local_Cartesian end
 function __index_Group_Cartesian end
 function __index_Global_Cartesian end
 
+__index_Local_NTuple(I...) = Tuple(__index_Local_Cartesian(I...))
+__index_Group_NTuple(I...) = Tuple(__index_Group_Cartesian(I...))
+__index_Global_NTuple(I...) = Tuple(__index_Global_Cartesian(I...))
+
 struct ConstAdaptor end
 
 Adapt.adapt_storage(to::ConstAdaptor, a::Array) = Base.Experimental.Const(a)