KernelAbstractions 0.10

CUDACore/Project.toml

@@ -14,7 +14,7 @@ ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
-KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LLVMLoopInfo = "8b046642-f1f6-4319-8d3c-209ddc03c586"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
@@ -34,11 +34,13 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [extensions]
 ChainRulesCoreExt = "ChainRulesCore"
 EnzymeCoreExt = "EnzymeCore"
+KernelAbstractionsExt = "KernelAbstractions"
 SpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
@@ -56,7 +58,8 @@ ExprTools = "0.1"
 GPUArrays = "11.5.4"
 GPUCompiler = "2.4"
 GPUToolbox = "3"
-KernelAbstractions = "0.9.38"
+KernelAbstractions = "0.9.38, 0.10"
+KernelInterface = "0.1.0"
 LLVM = "9.6"
 LLVMLoopInfo = "1"
 LazyArtifacts = "1"

CUDACore/ext/KernelAbstractionsExt.jl

@@ -0,0 +1,130 @@
+module KernelAbstractionsExt
+
+using CUDACore
+using CUDACore: @device_override, default_memory, UnifiedMemory, GPUArrays
+
+import KernelAbstractions as KA
+
+
+import StaticArrays
+
+import Adapt
+
+# TODO: Move AbstractGPUSparseArray stuff out
+Adapt.adapt_storage(::KA.CPU, a::Union{CuArray,GPUArrays.AbstractGPUSparseArray}) = Adapt.adapt(Array, a)
+
+## kernel launch
+
+function KA.mkcontext(kernel::KA.Kernel{CUDABackend}, _ndrange, iterspace)
+    KA.CompilerMetadata{KA.ndrange(kernel), KA.DynamicCheck}(_ndrange, iterspace)
+end
+
+function KA.launch_config(kernel::KA.Kernel{CUDABackend}, ndrange, workgroupsize)
+    if ndrange isa Integer
+        ndrange = (ndrange,)
+    end
+    if workgroupsize isa Integer
+        workgroupsize = (workgroupsize, )
+    end
+
+    # partition checked that the ndrange's agreed
+    if KA.ndrange(kernel) <: KA.StaticSize
+        ndrange = nothing
+    end
+
+    iterspace, dynamic = if KA.workgroupsize(kernel) <: KA.DynamicSize &&
+        workgroupsize === nothing
+        # use ndrange as preliminary workgroupsize for autotuning
+        KA.partition(kernel, ndrange, ndrange)
+    else
+        KA.partition(kernel, ndrange, workgroupsize)
+    end
+
+    return ndrange, workgroupsize, iterspace, dynamic
+end
+
+function threads_to_workgroupsize(threads, ndrange)
+    total = Ref(1)
+    return map(ndrange) do n
+        x = min(div(threads, total[]), n)
+        total[] *= x
+        return x
+    end
+end
+
+function (obj::KA.Kernel{CUDABackend})(args...; ndrange=nothing, workgroupsize=nothing)
+    backend = KA.backend(obj)
+
+    ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, workgroupsize)
+    # this might not be the final context, since we may tune the workgroupsize
+    ctx = KA.mkcontext(obj, ndrange, iterspace)
+
+    # If the kernel is statically sized we can tell the compiler about that
+    if KA.workgroupsize(obj) <: KA.StaticSize
+        maxthreads = prod(KA.get(KA.workgroupsize(obj)))
+    else
+        maxthreads = nothing
+    end
+
+    call = CUDACore.kernel_call(obj.f, (ctx, args...))
+    kernel = CUDACore.kernel_compile(call; always_inline=backend.always_inline, maxthreads)
+
+    # figure out the optimal workgroupsize automatically
+    if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing
+        config = CUDACore.launch_configuration(kernel.fun; max_threads=prod(ndrange))
+        if backend.prefer_blocks
+            # Prefer blocks over threads
+            threads = min(prod(ndrange), config.threads)
+            # XXX: Some kernels performs much better with all blocks active
+            cu_blocks = max(cld(prod(ndrange), threads), config.blocks)
+            threads = cld(prod(ndrange), cu_blocks)
+        else
+            threads = config.threads
+        end
+
+        workgroupsize = threads_to_workgroupsize(threads, ndrange)
+        iterspace, dynamic = KA.partition(obj, ndrange, workgroupsize)
+        ctx = KA.mkcontext(obj, ndrange, iterspace)
+        call = CUDACore.rebind(call, ctx, 1)
+    end
+
+    blocks = length(KA.blocks(iterspace))
+    threads = length(KA.workitems(iterspace))
+
+    if blocks == 0
+        return nothing
+    end
+
+    # Launch kernel
+    CUDACore.kernel_launch(kernel, call; threads, blocks)
+
+    return nothing
+end
+
+## indexing
+
+## COV_EXCL_START
+@device_override @inline function KA.__validindex(ctx)
+    if KA.__dynamic_checkbounds(ctx)
+        I = @inbounds KA.expand(KA.__iterspace(ctx), blockIdx().x, threadIdx().x)
+        return I in KA.__ndrange(ctx)
+    else
+        return true
+    end
+end
+
+## shared and scratch memory
+
+@device_override @inline function KA.Scratchpad(ctx, ::Type{T}, ::Val{Dims}) where {T, Dims}
+    StaticArrays.MArray{KA.__size(Dims), T}(undef)
+end
+
+## COV_EXCL_STOP
+
+## other
+
+Adapt.adapt_storage(to::KA.ConstAdaptor, a::CuDeviceArray) = Base.Experimental.Const(a)
+
+KA.argconvert(k::KA.Kernel{CUDABackend}, arg) = cudaconvert(arg)
+
+end

CUDACore/src/CUDACore.jl

@@ -17,7 +17,7 @@ using LLVM
 using LLVM.Interop
 using Core: LLVMPtr
 
-import KernelAbstractions
+import KernelInterface as KI
 
 using Adapt: Adapt, adapt, WrappedArray
 
@@ -123,7 +123,7 @@ include("sorting.jl")
 include("complex.jl")
 include("library_types.jl")
 
-# KernelAbstractions
+# KernelInterface
 include("CUDAKernels.jl")
 import .CUDAKernels: CUDABackend
 export CUDABackend