[Temp] Try out KernelIntrinsics

Author: christiangnrd

src/accumulate.jl

@@ -15,16 +15,16 @@
 function partial_scan(op::Function, output::AbstractArray{T}, input::AbstractArray,
                       Rdim, Rpre, Rpost, Rother, neutral, init,
                       ::Val{inclusive}=Val(true)) where {T, inclusive}
-    threads = blockDim().x
-    thread = threadIdx().x
-    block = blockIdx().x
+    threads = KI.get_local_size().x
+    thread = KI.get_local_id().x
+    block = KI.get_group_id().x
 
     temp = CuDynamicSharedArray(T, (2*threads,))
 
     # iterate the main dimension using threads and the first block dimension
-    i = (blockIdx().x-1i32) * blockDim().x + threadIdx().x
+    i = (KI.get_group_id().x-1i32) * KI.get_local_size().x + KI.get_local_id().x
     # iterate the other dimensions using the remaining block dimensions
-    j = (blockIdx().z-1i32) * gridDim().y + blockIdx().y
+    j = (KI.get_group_id().z-1i32) * KI.get_num_groups().y + KI.get_group_id().y
 
     if j > length(Rother)
         return
@@ -47,7 +47,7 @@ function partial_scan(op::Function, output::AbstractArray{T}, input::AbstractArr
     offset = 1
     d = threads>>1
     while d > 0
-        sync_threads()
+        KI.barrier()
         @inbounds if thread <= d
             ai = offset * (2*thread-1)
             bi = offset * (2*thread)
@@ -66,7 +66,7 @@ function partial_scan(op::Function, output::AbstractArray{T}, input::AbstractArr
     d = 1
     while d < threads
         offset >>= 1
-        sync_threads()
+        KI.barrier()
         @inbounds if thread <= d
             ai = offset * (2*thread-1)
             bi = offset * (2*thread)
@@ -78,7 +78,7 @@ function partial_scan(op::Function, output::AbstractArray{T}, input::AbstractArr
         d *= 2
     end
 
-    sync_threads()
+    KI.barrier()
 
     # write results to device memory
     @inbounds if i <= length(Rdim)
@@ -100,14 +100,14 @@ end
 function aggregate_partial_scan(op::Function, output::AbstractArray,
                                 aggregates::AbstractArray, Rdim, Rpre, Rpost, Rother,
                                 init)
-    threads = blockDim().x
-    thread = threadIdx().x
-    block = blockIdx().x
+    threads = KI.get_local_size().x
+    thread = KI.get_local_id().x
+    block = KI.get_group_id().x
 
     # iterate the main dimension using threads and the first block dimension
-    i = (blockIdx().x-1i32) * blockDim().x + threadIdx().x
+    i = (KI.get_group_id().x-1i32) * KI.get_local_size().x + KI.get_local_id().x
     # iterate the other dimensions using the remaining block dimensions
-    j = (blockIdx().z-1i32) * gridDim().y + blockIdx().y
+    j = (KI.get_group_id().z-1i32) * KI.get_num_groups().y + KI.get_group_id().y
 
     @inbounds if i <= length(Rdim) && j <= length(Rother)
         I = Rother[j]

src/device/random.jl

@@ -72,9 +72,9 @@ end
     elseif field === :ctr1
         @inbounds global_random_counters()[warpId]
     elseif field === :ctr2
-        blockId = blockIdx().x + (blockIdx().y - 1i32) * gridDim().x +
-                                 (blockIdx().z - 1i32) * gridDim().x * gridDim().y
-        globalId = threadId + (blockId - 1i32) * (blockDim().x * blockDim().y * blockDim().z)
+        globalId = KI.get_global_id().x +
+                   (KI.get_global_id().y - 1i32) * KI.get_global_size().x +
+                   (KI.get_global_id().z - 1i32) * KI.get_global_size().x * KI.get_global_size().y
         globalId%UInt32
     end::UInt32
 end

src/indexing.jl

@@ -33,7 +33,7 @@ function Base.findall(bools::AnyCuArray{Bool})
     if n > 0
         ## COV_EXCL_START
         function kernel(ys::CuDeviceArray, bools, indices)
-            i = threadIdx().x + (blockIdx().x - 1i32) * blockDim().x
+            i = KI.get_local_id().x + (KI.get_group_id().x - 1i32) * KI.get_local_size().x
 
             @inbounds if i <= length(bools) && bools[i]
                 i′ = CartesianIndices(bools)[i]

src/mapreduce.jl

@@ -19,9 +19,9 @@ end
 @inline function reduce_block(op, val::T, neutral, shuffle::Val{true}) where T
     # shared mem for partial sums
     assume(warpsize() == 32)
-    shared = CuStaticSharedArray(T, 32)
+    shared = KI.localmemory(T, 32)
 
-    wid, lane = fldmod1(threadIdx().x, warpsize())
+    wid, lane = fldmod1(KI.get_local_id().x, warpsize())
 
     # each warp performs partial reduction
     val = reduce_warp(op, val)
@@ -32,10 +32,10 @@ end
     end
 
     # wait for all partial reductions
-    sync_threads()
+    KI.barrier()
 
     # read from shared memory only if that warp existed
-    val = if threadIdx().x <= fld1(blockDim().x, warpsize())
+    val = if KI.get_local_id().x <= fld1(KI.get_local_size().x, warpsize())
          @inbounds shared[lane]
     else
         neutral
@@ -49,8 +49,8 @@ end
     return val
 end
 @inline function reduce_block(op, val::T, neutral, shuffle::Val{false}) where T
-    threads = blockDim().x
-    thread = threadIdx().x
+    threads = KI.get_local_size().x
+    thread = KI.get_local_id().x
 
     # shared mem for a complete reduction
     shared = CuDynamicSharedArray(T, (threads,))
@@ -59,7 +59,7 @@ end
     # perform a reduction
     d = 1
     while d < threads
-        sync_threads()
+        KI.barrier()
         index = 2 * d * (thread-1) + 1
         @inbounds if index <= threads
             other_val = if index + d <= threads
@@ -92,10 +92,10 @@ function partial_mapreduce_grid(f, op, neutral, Rreduce, Rother, shuffle, R::Abs
 
     # decompose the 1D hardware indices into separate ones for reduction (across threads
     # and possibly blocks if it doesn't fit) and other elements (remaining blocks)
-    threadIdx_reduce = threadIdx().x
-    blockDim_reduce = blockDim().x
-    blockIdx_reduce, blockIdx_other = fldmod1(blockIdx().x, length(Rother))
-    gridDim_reduce = gridDim().x ÷ length(Rother)
+    threadIdx_reduce = KI.get_local_id().x
+    blockDim_reduce = KI.get_local_size().x
+    blockIdx_reduce, blockIdx_other = fldmod1(KI.get_group_id().x, length(Rother))
+    gridDim_reduce = KI.get_num_groups().x ÷ length(Rother)
 
     # block-based indexing into the values outside of the reduction dimension
     # (that means we can safely synchronize threads within this block)
@@ -134,7 +134,7 @@ function partial_mapreduce_grid(f, op, neutral, Rreduce, Rother, shuffle, R::Abs
 end
 
 function serial_mapreduce_kernel(f, op, neutral, Rreduce, Rother, R, As)
-    grid_idx = threadIdx().x + (blockIdx().x - 1i32) * blockDim().x
+    grid_idx = KI.get_local_id().x + (KI.get_group_id().x - 1i32) * KI.get_local_size().x
     @inbounds if grid_idx <= length(Rother)
         Iother = Rother[grid_idx]
 
@@ -160,14 +160,14 @@ end
 
 # factored out for use in tests
 function serial_mapreduce_threshold(dev)
-    max_concurrency = attribute(dev, DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK) *
-                      attribute(dev, DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
+    max_concurrency = KI.max_work_group_size(CUDABackend()) * KI.multiprocessor_count(CUDABackend())
     return max_concurrency
 end
 
 function GPUArrays.mapreducedim!(f::F, op::OP, R::AnyCuArray{T},
                                  A::Union{AbstractArray,Broadcast.Broadcasted};
                                  init=nothing) where {F, OP, T}
+    backend = CUDABackend()
     if !isa(A, Broadcast.Broadcasted)
         # XXX: Base.axes isn't defined anymore for Broadcasted, breaking this check
         Base.check_reducedims(R, A)
@@ -201,10 +201,13 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::AnyCuArray{T},
     # If `Rother` is large enough, then a naive loop is more efficient than partial reductions.
     if length(Rother) >= serial_mapreduce_threshold(dev)
         args = (f, op, init, Rreduce, Rother, R, A)
+        # kernel = KI.KIKernel(backend, serial_mapreduce_kernel, args...)
         kernel = @cuda launch=false serial_mapreduce_kernel(args...)
+        # kernel_config = launch_configuration(kernel.kern.fun)
         kernel_config = launch_configuration(kernel.fun)
         threads = kernel_config.threads
         blocks = cld(length(Rother), threads)
+        # kernel(args...; workgroupsize=threads, numworkgroups=blocks)
         kernel(args...; threads, blocks)
         return R
     end
@@ -228,8 +231,10 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::AnyCuArray{T},
     # we might not be able to launch all those threads to reduce each slice in one go.
     # that's why each threads also loops across their inputs, processing multiple values
     # so that we can span the entire reduction dimension using a single thread block.
+    # kernel = KI.KIKernel(backend, partial_mapreduce_grid, f, op, init, Rreduce, Rother, Val(shuffle), R, A)
     kernel = @cuda launch=false partial_mapreduce_grid(f, op, init, Rreduce, Rother, Val(shuffle), R, A)
     compute_shmem(threads) = shuffle ? 0 : threads*sizeof(T)
+    # kernel_config = launch_configuration(kernel.kern.fun; shmem=compute_shmem∘compute_threads)
     kernel_config = launch_configuration(kernel.fun; shmem=compute_shmem∘compute_threads)
     reduce_threads = compute_threads(kernel_config.threads)
     reduce_shmem = compute_shmem(reduce_threads)
@@ -255,6 +260,7 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::AnyCuArray{T},
     # perform the actual reduction
     if reduce_blocks == 1
         # we can cover the dimensions to reduce using a single block
+        # kernel(f, op, init, Rreduce, Rother, Val(shuffle), R, A; workgroupsize=threads, numworkgroups=blocks, shmem)
         kernel(f, op, init, Rreduce, Rother, Val(shuffle), R, A; threads, blocks, shmem)
     else
         # TODO: provide a version that atomically reduces from different blocks
@@ -265,7 +271,9 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::AnyCuArray{T},
         # NOTE: we can't use the previously-compiled kernel, or its launch configuration,
         #       since the type of `partial` might not match the original output container
         #       (e.g. if that was a view).
+        # partial_kernel = KI.KIKernel(backend, partial_mapreduce_grid, f, op, init, Rreduce, Rother, Val(shuffle), partial, A)
         partial_kernel = @cuda launch=false partial_mapreduce_grid(f, op, init, Rreduce, Rother, Val(shuffle), partial, A)
+        # partial_kernel_config = launch_configuration(partial_kernel.kern.fun; shmem=compute_shmem∘compute_threads)
         partial_kernel_config = launch_configuration(partial_kernel.fun; shmem=compute_shmem∘compute_threads)
         partial_reduce_threads = compute_threads(partial_kernel_config.threads)
         partial_reduce_shmem = compute_shmem(partial_reduce_threads)
@@ -286,7 +294,8 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::AnyCuArray{T},
         end
 
         partial_kernel(f, op, init, Rreduce, Rother, Val(shuffle), partial, A;
-                       threads=partial_threads, blocks=partial_blocks, shmem=partial_shmem)
+                    threads=partial_threads, blocks=partial_blocks, shmem=partial_shmem)
+                    # workgroupsize=partial_threads, numworkgroups=partial_blocks, shmem=partial_shmem)
 
         GPUArrays.mapreducedim!(identity, op, R, partial; init)
     end