KernelAbstractions 0.10

Project.toml

@@ -1,7 +1,7 @@
 name = "oneAPI"
 uuid = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
-authors = ["Tim Besard <tim.besard@gmail.com>", "Alexis Montoison", "Michel Schanen <michel.schanen@gmail.com>"]
 version = "2.9.1"
+authors = ["Tim Besard <tim.besard@gmail.com>", "Alexis Montoison", "Michel Schanen <michel.schanen@gmail.com>"]
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -13,6 +13,7 @@ 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"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -32,6 +33,12 @@ oneAPI_Level_Zero_Headers_jll = "f4bc562b-d309-54f8-9efb-476e56f0410d"
 oneAPI_Level_Zero_Loader_jll = "13eca655-d68d-5b81-8367-6d99d727ab01"
 oneAPI_Support_jll = "b049733a-a71d-5ed3-8eba-7d323ac00b36"
 
+[weakdeps]
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+
+[extensions]
+KernelAbstractionsExt = "KernelAbstractions"
+
 [compat]
 AbstractFFTs = "1.5.0"
 AcceleratedKernels = "0.3.1, 0.4"
@@ -41,7 +48,8 @@ ExprTools = "0.1"
 GPUArrays = "11.2.1"
 GPUCompiler = "2 - 2.2.1"
 GPUToolbox = "0.1, 0.2, 0.3, 1, 3"
-KernelAbstractions = "0.9.39"
+KernelAbstractions = "0.10"
+KernelInterface = "0.1"
 LLVM = "6, 7, 8, 9"
 NEO_jll = "=26.18.38308"
 PrecompileTools = "1"

ext/KernelAbstractionsExt.jl

@@ -0,0 +1,138 @@
+module KernelAbstractionsExt
+
+using ..oneAPI
+using ..oneAPI: @device_override, SPIRVIntrinsics, method_table
+
+import KernelAbstractions as KA
+
+import StaticArrays
+
+import Adapt
+
+
+Adapt.adapt_storage(::KA.CPU, a::oneArray) = convert(Array, a)
+
+## Kernel Launch
+
+function KA.mkcontext(kernel::KA.Kernel{oneAPIBackend}, _ndrange, iterspace)
+    KA.CompilerMetadata{KA.ndrange(kernel), KA.DynamicCheck}(_ndrange, iterspace)
+end
+function KA.mkcontext(kernel::KA.Kernel{oneAPIBackend}, I, _ndrange, iterspace,
+                      ::Dynamic) where Dynamic
+    KA.CompilerMetadata{KA.ndrange(kernel), Dynamic}(I, _ndrange, iterspace)
+end
+
+function KA.launch_config(kernel::KA.Kernel{oneAPIBackend}, 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{oneAPIBackend})(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
+        # TODO: maxthreads
+        # maxthreads = prod(KA.get(KA.workgroupsize(obj)))
+    else
+        # maxthreads = nothing
+    end
+
+    kernel = @oneapi launch = false always_inline = backend.always_inline obj.f(ctx, args...)
+
+    # figure out the optimal workgroupsize automatically
+    if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing
+        items = oneAPI.launch_configuration(kernel)
+
+        if backend.prefer_blocks
+            # Prefer blocks over threads:
+            # Reducing the workgroup size (items) increases the number of workgroups (blocks).
+            # We use a simple heuristic here since we lack full occupancy info (max_blocks) from launch_configuration.
+
+            # If the total range is large enough, full workgroups are fine.
+            # If the range is small, we might want to reduce 'items' to create more blocks to fill the GPU.
+            # (Simplified logic compared to CUDA.jl which uses explicit occupancy calculators)
+            total_items = prod(ndrange)
+            if total_items < items * 16 # Heuristic factor
+                # Force at least a few blocks if possible by reducing items per block
+                target_blocks = 16 # Target at least 16 blocks
+                items = max(1, min(items, cld(total_items, target_blocks)))
+            end
+        end
+
+        workgroupsize = threads_to_workgroupsize(items, ndrange)
+        iterspace, dynamic = KA.partition(obj, ndrange, workgroupsize)
+        ctx = KA.mkcontext(obj, ndrange, iterspace)
+    end
+
+    groups = length(KA.blocks(iterspace))
+    items = length(KA.workitems(iterspace))
+
+    if groups == 0
+        return nothing
+    end
+
+    # Launch kernel
+    kernel(ctx, args...; items, groups)
+
+    return nothing
+end
+
+
+## Indexing Functions
+
+@device_override @inline function KA.__validindex(ctx)
+    if KA.__dynamic_checkbounds(ctx)
+        I = @inbounds KA.expand(KA.__iterspace(ctx), get_group_id(), get_local_id())
+        return I in KA.__ndrange(ctx)
+    else
+        return true
+    end
+end
+
+
+## Scratch Memory
+
+@device_override @inline function KA.Scratchpad(ctx, ::Type{T}, ::Val{Dims}) where {T, Dims}
+    StaticArrays.MArray{KA.__size(Dims), T}(undef)
+end
+
+## Other
+
+Adapt.adapt_storage(to::KA.ConstAdaptor, a::oneDeviceArray) = Base.Experimental.Const(a)
+
+KA.argconvert(::KA.Kernel{oneAPIBackend}, arg) = kernel_convert(arg)
+
+end

src/compiler/compilation.jl

@@ -1,6 +1,9 @@
 ## gpucompiler interface implementation
 
-struct oneAPICompilerParams <: AbstractCompilerParams end
+Base.@kwdef struct oneAPICompilerParams <: AbstractCompilerParams
+    sub_group_size::Union{Nothing,Int} = nothing
+end
+
 const oneAPICompilerConfig = CompilerConfig{SPIRVCompilerTarget, oneAPICompilerParams}
 const oneAPICompilerJob = CompilerJob{SPIRVCompilerTarget,oneAPICompilerParams}
 
@@ -48,6 +51,11 @@ function GPUCompiler.finish_module!(job::oneAPICompilerJob, mod::LLVM.Module,
                    Tuple{CompilerJob{SPIRVCompilerTarget}, typeof(mod), typeof(entry)},
                    job, mod, entry)
 
+    # Set the subgroup size
+    if job.config.params.sub_group_size !== nothing
+        metadata(entry)["intel_reqd_sub_group_size"] = MDNode([ConstantInt(Int32(job.config.params.sub_group_size))])
+    end
+
     # OpenCL 2.0
     push!(metadata(mod)["opencl.ocl.version"],
           MDNode([ConstantInt(Int32(2)),
@@ -323,11 +331,16 @@ function _driver_supports_bfloat16_spirv(dev=device())
     end
 end
 
-@noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false, kwargs...)
+@noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false, sub_group_size=32, kwargs...)
     properties = oneL0.module_properties(dev)
     supports_fp16 = properties.fp16flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 == oneL0.ZE_DEVICE_MODULE_FLAG_FP16
     supports_fp64 = properties.fp64flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP64 == oneL0.ZE_DEVICE_MODULE_FLAG_FP64
 
+    if sub_group_size ∉ oneL0.compute_properties(dev).subGroupSizes
+        @error("$sub_group_size is not a valid sub-group size for this device.")
+    end
+
+
     # SPIR-V codegen path. The Aurora LTS NEO/IGC runtime only accepts SPIR-V from the
     # Khronos translator; the rolling stack uses the LLVM SPIR-V back-end. GPUCompiler picks
     # the tool from the target's `backend` field and loads the JLL lazily, so both can be
@@ -360,7 +373,7 @@ end
 
     # create GPUCompiler objects
     target = SPIRVCompilerTarget(; backend, extensions = extensions_str, supports_fp16, supports_fp64, supports_bfloat16, kwargs...)
-    params = oneAPICompilerParams()
+    params = oneAPICompilerParams(; sub_group_size)
     CompilerConfig(target, params; kernel, name, always_inline)
 end
 

src/compiler/execution.jl

@@ -4,7 +4,7 @@ export @oneapi, zefunction, kernel_convert
 ## high-level @oneapi interface
 
 const MACRO_KWARGS = [:launch]
-const COMPILER_KWARGS = [:kernel, :name, :always_inline]
+const COMPILER_KWARGS = [:kernel, :name, :always_inline, :sub_group_size]
 const LAUNCH_KWARGS = [:groups, :items, :queue]
 
 """

src/oneAPI.jl

@@ -12,6 +12,7 @@ using SpecialFunctions
 import Preferences
 
 import KernelAbstractions: KernelAbstractions
+import KernelInterface
 
 using LLVM
 using LLVM.Interop
@@ -76,12 +77,18 @@ include("gpuarrays.jl")
 include("random.jl")
 include("utils.jl")
 
-include("oneAPIKernels.jl")
+# KernelAbstractions
+include("oneAPIKernelsOld.jl")
 import .oneAPIKernels: oneAPIBackend
+export oneAPIBackend
+
+# KernelInterface
+include("oneAPIKernels.jl")
+import .oneAPIInterface
+
 include("accumulate.jl")
 include("sorting.jl")
 include("indexing.jl")
-export oneAPIBackend
 
 # precompilation workload (warms up the SPIR-V compilation pipeline)
 include("compiler/precompile.jl")