Merge pull request #431 from JuliaGPU/tb/launch_config

Roll our own launch configuration

Author: maleadt

lib/level-zero/driver.jl

@@ -100,15 +100,18 @@ function ipc_properties(drv::ZeDriver)
     )
 end
 
-# FIXME: throws ZE_RESULT_ERROR_UNSUPPORTED_FEATURE
 function extension_properties(drv::ZeDriver)
     count_ref = Ref{UInt32}(0)
     zeDriverGetExtensionProperties(drv, count_ref, C_NULL)
 
     all_props = Vector{ze_driver_extension_properties_t}(undef, count_ref[])
     zeDriverGetExtensionProperties(drv, count_ref, all_props)
 
-    return [(name=String([props.name[1:findfirst(isequal(0), props.name)-1]...]),
-             version=Int(props.version),
-            ) for props in all_props[1:count_ref[]]]
+    extensions = Dict{String,VersionNumber}()
+    for prop in all_props[1:count_ref[]]
+        name = String(UInt8[prop.name[1:findfirst(isequal(0), prop.name)-1]...])
+        version = unmake_version(prop.version)
+        extensions[name] = version
+    end
+    return extensions
 end

lib/level-zero/memory.jl

@@ -50,11 +50,11 @@ function device_alloc(ctx::ZeContext, dev::ZeDevice, bytesize::Integer,
         flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE,
     ))
     GC.@preserve relaxed_allocation_ref begin
-        desc_ref = if bytesize > properties(dev).maxMemAllocSize
-            pNext = Base.unsafe_convert(Ptr{Cvoid}, relaxed_allocation_ref)
-            Ref(ze_device_mem_alloc_desc_t(; flags, ordinal, pNext))
+        if bytesize > properties(dev).maxMemAllocSize
+            desc_ref = Ref(ze_device_mem_alloc_desc_t(; flags, ordinal))
+            link_extensions(desc_ref, relaxed_allocation_ref)
         else
-            Ref(ze_device_mem_alloc_desc_t(; flags, ordinal))
+            desc_ref = Ref(ze_device_mem_alloc_desc_t(; flags, ordinal))
         end
 
         ptr_ref = Ref{Ptr{Cvoid}}()

lib/level-zero/module.jl

@@ -234,6 +234,16 @@ export properties
 
 function properties(kernel::ZeKernel)
     props_ref = Ref(ze_kernel_properties_t())
+    preferred_group_size_props_ref = Ref(ze_kernel_preferred_group_size_properties_t())
+    link_extensions(props_ref, preferred_group_size_props_ref)
+    if haskey(oneL0.extension_properties(kernel.mod.context.driver),
+              "ZE_extension_kernel_max_group_size_properties")
+        # TODO: memoize
+        max_group_size_props_ref = Ref(ze_kernel_max_group_size_properties_ext_t())
+        link_extensions(preferred_group_size_props_ref, max_group_size_props_ref)
+    else
+        max_group_size_props_ref = nothing
+    end
     zeKernelGetProperties(kernel, props_ref)
 
     props = props_ref[]
@@ -251,6 +261,9 @@ function properties(kernel::ZeKernel)
         spillMemSize=Int(props.spillMemSize),
         kernel_uuid=Base.UUID(reinterpret(UInt128, [props.uuid.kid...])[1]),
         module_uuid=Base.UUID(reinterpret(UInt128, [props.uuid.mid...])[1]),
+        preferredGroupSize=Int(preferred_group_size_props_ref[].preferredMultiple),
+        maxGroupSize=max_group_size_props_ref === nothing ? missing :
+            Int(max_group_size_props_ref[].maxGroupSize)
     )
 end
 

lib/level-zero/oneL0.jl

@@ -69,6 +69,20 @@ zeroinit(::Type{_ze_native_kernel_uuid_t}) =
 zeroinit(::Type{ze_kernel_uuid_t}) =
     ze_kernel_uuid_t(ntuple(_->zero(UInt8), 16), ntuple(_->zero(UInt8), 16))
 
+# link extension objects in pNext
+function link_extensions(refs...)
+    length(refs) >= 2 || return
+    for (parent, child) in zip(refs[1:end-1], refs[2:end])
+        pNext = Base.unsafe_convert(Ptr{Cvoid}, child)
+        typ = eltype(parent)
+        @assert fieldnames(typ)[2] == :pNext
+        field = Base.unsafe_convert(Ptr{Cvoid}, parent) + fieldoffset(typ, 2)
+        field = convert(Ptr{Ptr{Cvoid}}, field)
+        unsafe_store!(field, pNext)
+    end
+    return
+end
+
 # core wrappers
 include("error.jl")
 include("common.jl")

res/patches/.dummy

@@ -161,6 +161,33 @@ struct HostKernel{F,TT} <: AbstractKernel{F,TT}
     fun::ZeKernel
 end
 
+function launch_configuration(kernel::HostKernel{F,TT}) where {F,TT}
+    # XXX: have the user pass in a global size to clamp against
+    #      maxGroupSizeX/Y/Z?
+
+    # XXX: shrink until a multiple of preferredGroupSize?
+
+    # once the MAX_GROUP_SIZE extension is implemented, we can use it here
+    kernel_props = oneL0.properties(kernel.fun)
+    if kernel_props.maxGroupSize !== missing
+        return kernel_props.maxGroupSize
+    end
+
+    # otherwise, we'd use `zeKernelSuggestGroupSize` but it's been observed
+    # to return really bad configs (JuliaGPU/oneAPI.jl#430)
+
+    # so instead, calculate it ourselves based on the device properties
+    dev = kernel.fun.mod.device
+    compute_props = oneL0.compute_properties(dev)
+    max_size = compute_props.maxTotalGroupSize
+    ## when the kernel uses many registers (which we can't query without
+    ## extensions that landed _after_ MAX_GROUP_SIZE, so don't bother)
+    ## the groupsize should be halved
+    group_size = max_size ÷ 2
+
+    return group_size
+end
+
 
 ## host-side API
 

src/compiler/execution.jl

@@ -16,9 +16,7 @@ struct oneKernelContext <: AbstractKernelContext end
                                              elements::Int, elements_per_thread::Int) where {F,N}
     kernel = @oneapi launch=false f(oneKernelContext(), args...)
 
-    items = suggest_groupsize(kernel.fun, elements).x
-    # XXX: the z dimension of the suggested group size is often non-zero.
-    #      preserve this in GPUArrays?
+    items = launch_configuration(kernel)
     # XXX: how many groups is a good number? the API doesn't tell us.
     #      measured on a low-end IGP, 32 blocks seems like a good sweet spot.
     #      note that this only matters for grid-stride kernels, like broadcast.

src/gpuarrays.jl

@@ -146,7 +146,7 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::oneWrappedArray{T},
     kernel_args = kernel_convert.(args)
     kernel_tt = Tuple{Core.Typeof.(kernel_args)...}
     kernel = zefunction(partial_mapreduce_device, kernel_tt)
-    reduce_items = compute_items(suggest_groupsize(kernel.fun, wanted_items).x)
+    reduce_items = launch_configuration(kernel)
 
     # how many groups should we launch?
     #

src/mapreduce.jl

@@ -90,8 +90,7 @@ function (obj::KA.Kernel{oneAPIBackend})(args...; ndrange=nothing, workgroupsize
 
     # figure out the optimal workgroupsize automatically
     if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing
-        items = oneAPI.suggest_groupsize(kernel.fun, prod(ndrange)).x
-        # XXX: the z dimension of the suggested group size is often non-zero. use this?
+        items = oneAPI.launch_configuration(kernel)
         workgroupsize = threads_to_workgroupsize(items, ndrange)
         iterspace, dynamic = KA.partition(obj, ndrange, workgroupsize)
         ctx = KA.mkcontext(obj, ndrange, iterspace)

src/oneAPIKernels.jl

@@ -20,7 +20,7 @@ api_version(drv)
 
 properties(drv)
 ipc_properties(drv)
-#extension_properties(drv)
+extension_properties(drv)
 
 end